Research map / Image generation / Fast diffusion sampling
Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking
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.
Proposes MDM-Prime, which encodes tokens as base-b sub-token sequences so masked diffusion can pass through intermediate partially-masked states, reducing idle sampling steps and improving text and image generation.
Problem
Masked diffusion models have many idle sampling steps where the token sequence is unchanged between steps, wasting computation and limiting likelihood performance on discrete data.
Approach
Represent each token as a base-b sub-token sequence and mask at the sub-token level, creating intermediate states between masked and unmasked; the decoder uses a joint softmax over valid sub-token sequences with carry-over parameterization, and the embedding concatenates sub-token embeddings.
Key results
- MDLM-Prime with l=6 reaches 15.36 perplexity on OpenWebText, versus 21.52 for the best prior MDM (EDLM-NCE) and 17.54 for ARM.
- MDM-Prime achieves FID 3.26 on CIFAR-10 and 6.98 on ImageNet-32, outperforming the discrete baselines reported in the paper.
- Idle step ratio on OWT falls from 36.77% for MDLM to 0.25% for MDLM-Prime (l=6).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add partial masking with sub-token length ℓ=6 | MDLM · OpenWebText | PPL: ≤22.98 -> ≤15.36 | Helped |
| Add partial masking with sub-token length ℓ=2 | MDM · CIFAR-10 | FID: 4.66 -> 3.26 | Helped |
| Add partial masking with sub-token length ℓ=2 | MDM · CIFAR-10 | IS: 9.09 -> 9.67 | Helped |
| Add partial masking with sub-token length ℓ=2 | MDM · ImageNet-32 | FID: 7.91 -> 6.98 | Helped |
| Add partial masking with sub-token length ℓ=2 | MDM · ImageNet-32 | IS: 11.60 -> 11.65 | Helped |
| Increase sub-token length from ℓ=2 to ℓ=6 (ablation) | MDLM-Prime · OpenWebText | PPL: ≤17.90 -> ≤15.36 | Helped |
| Increase sub-token length from ℓ=6 to ℓ=8 (ablation) | MDLM-Prime · OpenWebText | PPL: ≤15.36 -> ≤15.48 | Hurt |
| Use ℓ=4 instead of ℓ=1 partial masking (ablation) | MDLM · AG News zero-shot | PPL: 61.15 -> 59.44 | Hurt |
| Use ℓ=4 instead of ℓ=1 partial masking (ablation) | MDLM · LAMBADA zero-shot | PPL: 47.52 -> 24.44 | Helped |
| Use ℓ=4 instead of ℓ=1 partial masking (ablation) | MDLM · WikiText zero-shot | PPL: 32.83 -> 23.86 | Helped |
| Use the carry-over parameterization (ablation) | MDLM-Prime · LAMBADA and PubMed zero-shot | PPL | Helped |
| Use an independent rather than joint sub-token distribution (ablation) | MDM-Prime · Synthetic 512×512 sampling distribution | sampling distribution quality | Hurt |
Surprising findings
- Increasing sub-token length l does not monotonically improve performance; performance converges around l>=4 for text.
Limitations
- Performance is not monotonic in sub-token length, and choosing l requires heuristics (ISR elbow point).
- Results on most benchmarks come from single runs without error bars.
- Bit Diffusion-style quantization methods lack tractable likelihood, which the authors contrast with their approach.
Experimental setup
Datasets: OpenWebText CIFAR-10 ImageNet-32 LAMBADA WikiText Penn Treebank LM1B AG News PubMed ArXiv
Benchmarks: OpenWebText perplexity LAMBADA WikiText PTB LM1B AG News PubMed ArXiv CIFAR-10 FID/IS ImageNet-32 FID/IS
Metrics: perplexity FID IS idle step ratio
Baselines: ARM BD3-LM EDLM-coAR SEDD EDLM-NCE MDLM GenMD4 MDM-Mixture PixelCNN D3PM CTDD-DG Tau-LDR
Builds on: MDLM D3PM carry-over parameterization ADM
Closest papers by meaning
- Simplified and Generalized Masked Diffusion for Discrete Data (NeurIPS 2024 · similarity 0.95)
- Learn from Your Mistakes: Self-Correcting Masked Diffusion Models (NeurIPS 2026 · similarity 0.93)
- Simple and Effective Masked Diffusion Language Models (NeurIPS 2024 · similarity 0.92)
- Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling (ICLR 2025 · similarity 0.92)
- Why Masking Diffusion Works: Condition on the Jump Schedule for Improved Discrete Diffusion (NeurIPS 2025 · similarity 0.91)
- Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking (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.