Research map / Image generation / Fast diffusion sampling
Simplified and Generalized Masked Diffusion for Discrete Data
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Simplifies masked (absorbing) discrete diffusion so the continuous-time ELBO is a weighted integral of cross-entropy losses, and extends it to state-dependent masking schedules (GenMD4), achieving strong text and pixel-level likelihoods.
Problem
Masked diffusion models for discrete data have complex formulations and ELBOs that are hard to optimize, leading to suboptimal parameterizations and training objectives that lag behind autoregressive and continuous diffusion models.
Approach
Derive a continuous-time masked diffusion ELBO that reduces to a simple weighted integral of cross-entropy losses over masked tokens, showing that the masking schedule only affects the loss via SNR endpoints and variance; generalize to state-dependent masking schedules optimized with REINFORCE leave-one-out.
Key results
- MD4 small on OpenWebText beats SEDD Absorb reimplementation on all five zero-shot perplexity benchmarks (e.g., WikiText2 34.94 vs 38.94).
- On text8, GenMD4 reaches 1.34 BPC vs SEDD Absorb 1.39 and MD4 1.37.
- MD4 achieves 2.75 BPD on CIFAR-10 (better than best AR result 2.80) and 3.40 BPD on ImageNet 64x64, matching Perceiver AR with 770M params using a 198M model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a cosine schedule instead of a linear schedule with a uniform sampling grid (ablation) | MD4 · ImageNet 64×64 | FID: 70 to 17 | Helped |
| Use a cosine sampling grid with a linear training schedule instead of a uniform grid (ablation) | MD4 · ImageNet 64×64 | FID | Helped |
| Use a state-dependent masking schedule | MD4 · Text8 test | BPC: ≤ 1.37 -> ≤ 1.34 | Helped |
| Use a state-dependent masking schedule | MD4 · OpenWebText validation | perplexity | Helped |
| Use a state-dependent masking schedule | MD4 · zero-shot transfer tasks | perplexity | Hurt |
| Replace the uniform sampling grid with a cosine grid (ablation) | MD4 · text infilling | generation quality | Helped |
| Use the simplified analytic cross-entropy objective instead of the doubly stochastic rate-matrix estimate | masked diffusion · pixel-level image modeling | training variance | Helped |
Surprising findings
- A cosine unmasking discretization grid, applied to a linear-schedule model, improved ImageNet 64x64 FID from about 70 to 17 at T=256 steps.
- Sampling with JAX Gumbel-argmax categorical sampling implicitly truncates small probabilities, acting like top-p sampling and improving sample quality.
Limitations
- Masked diffusion still lags autoregressive models on text8 (AR transformer 1.23 BPC).
- GenMD4 with learned state-dependent schedules overfits to dataset statistics and transfers worse zero-shot.
- Inference with state-dependent schedules is more challenging.
Experimental setup
Datasets: OpenWebText text8 FineWeb-Edu CIFAR-10 ImageNet 64x64 LAMBADA WikiText2 PTB WikiText103 1BW
Benchmarks: OpenWebText zero-shot perplexity (LAMBADA, WikiText2, PTB, WikiText103, 1BW) text8 BPC Hellaswag CIFAR-10 BPD ImageNet 64x64 BPD FID
Metrics: perplexity bits per character bits per dimension zero-shot accuracy FID
Baselines: SEDD Absorb D3PM Plaid GPT-2 BFN ARDM MAC Discrete Flow VDM PixelCNN++ Sparse Transformer Perceiver AR
Builds on: SEDD D3PM VDM MaskGIT
Closest papers by meaning
- Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking (NeurIPS 2025 · similarity 0.95)
- Simple and Effective Masked Diffusion Language Models (NeurIPS 2024 · similarity 0.94)
- Why Masking Diffusion Works: Condition on the Jump Schedule for Improved Discrete Diffusion (NeurIPS 2025 · similarity 0.93)
- Sparsely Supervised Diffusion (NeurIPS 2026 · similarity 0.92)
- Generalized Interpolating Discrete Diffusion (ICML 2025 · similarity 0.92)
- Remasking Discrete Diffusion Models with Inference-Time Scaling (NeurIPS 2025 · similarity 0.91)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.