Multi-Mask Diffusion Language Models for Few-Step Generation
cs.CL, cs.LG
Submitted: 2026-07-22
Updated: 2026-09-02
Comments: 38 pages; Accepted at COLM 2026
Code: https://github.com/ML-GSAI/LLaDA
Project page: http://skylion007.github.io/OpenWebTextCorpus
License: http://creativecommons.org/licenses/by/4.0/
The gist: Masked diffusion models (MDMs) are a promising family of language generators, but achieving high-quality few-step generation remains challenging.
Terminology
Abstract
Masked diffusion models (MDMs) are a promising family of language generators, but achieving high-quality few-step generation remains challenging. In MDMs, all forward trajectories collapse to a single fully masked state, leaving no terminal entropy for consistency-style few-step generation. While recent few-step alternatives based on uniform-state diffusion avoid this degeneracy, it becomes harder to distinguish clean tokens from noise than MDMs, which usually harms modeling quality and training efficiency. In this work, we propose a multi-mask diffusion model (MultiMDM) that preserves the masking structure towards few-step generation. In the forward process, each clean token is first pushed towards a designated mask and then gradually mixes over the mask set. As a result, the backward process has a drafting capability by predicting a designated mask before refining to a clean token. We derive a closed-form ELBO training objective for MultiMDM that supports continual training from pretrained MDMs. In addition, we formulate a purely discrete-state consistency distillation scheme, with a shared-Gumbel coupling to reduce pathwise entropy. Experiments on pretraining and distillation show that MultiMDM provides an effective foundation for principled few-step generation.
Sources
- MDM-Prime-v2: Binary Encoding and Index Shuffling Enable Scaling of Diffusion Language Models
- One Billion Word Benchmark for Measuring Progress in Statistical Language Modeling
- Edit Flows: Flow Matching with Edit Operations
- Beyond Single Tokens: Distilling Discrete Diffusion Models via Discrete MMD
- Any-Order Flexible Length Masked Diffusion
- Flow Map Language Models: One-step Language Modeling via Continuous Denoising
- Training Optimal Large Diffusion Language Models
- CANDI: Hybrid Discrete-Continuous Diffusion Models
- Categorical Flow Maps
- Seed Diffusion: A Large-Scale Diffusion Language Model with High-Speed Inference
- Dream 7B: Diffusion Large Language Models
- Continuously Augmented Discrete Diffusion model for Categorical Generative Modeling
- Coevolutionary Continuous Discrete Diffusion: Make Your Diffusion Language Model a Latent Reasoner
Related papers
- Exploring Solution Divergence and Its Effect on Large Language Model Problem Solving
- Ishigaki-IDS-Bench: A Benchmark for Generating Information Delivery Specification from BIM Information Requirements
- Subliminal Steering: Stronger Encoding of Hidden Signals
- MedStruct-S: A Benchmark for Key Discovery, Key-Conditioned QA and Semi-Structured Extraction from OCR Clinical Reports
- The End of Transformers? On Challenging Attention and the Rise of Sub-Quadratic Architectures
- Untangling the Mechanisms of Misleading Context in Medical Question Answering