Multi-Mask Diffusion Language Models for Few-Step Generation
Sijin Chen, Yinuo Ren, Heyang Zhao, Ziheng Cheng, Quanquan Gu, Lexing Ying
COLM 2026
cs.CL, cs.LG
2026-07-22
MultiMDM sends each clean token to a designated mask, then mixes in the mask set. Continual training from MDLM cuts 4-step entropy-matched GenPPL on a 170M DiT from 721.1 to 558.8.
Masked diffusion language models (MDMs) have a clean training target and a reconstruction job that matches language modeling. Few-step sampling is still awkward. Every forward path collapses to one fully masked state, so terminal entropy is zero and consistency distillation has no random seed to start from. Uniform-state diffusion (USDM) spreads mass over the vocabulary and keeps a stochastic endpoint, but clean tokens and noise look the same, which usually hurts modeling quality and training efficiency.
ByteDance Seed, with Princeton, Stanford, UCLA, and Berkeley, introduce MultiMDM: keep the masking structure, replace the single absorbing state with a set of masks, and give few-step generation somewhere to start.
Each clean token is assigned a designated mask. The forward process has two knobs: αt pushes a visible token toward its designated mask, then βt mixes inside the mask subspace. The terminal law is uniform over the mask set, with entropy L log M, not zero. In reverse, the model can draft by predicting the designated mask before refining it into a clean token. The mask count M is much smaller than the vocabulary; M in [5, 100] is the useful range, and the main runs use 50.
The training objective is a closed-form ELBO with two terms: the same clean-token reconstruction used by MDM, plus an intra-mask identification term. Visible positions contribute zero. Continual training from a pretrained MDM is straightforward: keep the clean-token head, copy the old single-mask embedding onto every new mask, then let them separate. A curriculum turns on reconstruction first, then the identification term.
Few-step distillation stays in discrete state space. One Gumbel draw per coordinate and state is shared across time. Gumbel-max keeps the teacher marginals; the path is deterministic given that noise. The consistency target is the clean-token posterior along the coupled path. In practice an EMA teacher compresses the future-filtration posterior into a Markov predictor that only sees the current sequence.
The backbone is a 170M DiT on two unnamed English corpora (names withheld for data-disclosure rules). Quality is GPT-2 generative perplexity (GenPPL, lower is better), reported both at matched corpus entropy and at temperature τ=1. Pretraining runs 200K steps. The continual recipe is 150K steps of MDM plus 50K steps of MultiMDM.
Corpus A (entropy 5.44), entropy-matched GenPPL:
| Method | 4 steps | 8 steps | 16 steps | 64 steps |
| DUO | 995.3 | 509.4 | 193.6 | 70.9 |
| MDLM | 721.1 | 280.1 | 117.0 | 65.9 |
| CANDI | 666.5 | 237.5 | 126.8 | 69.6 |
| MultiMDM-cont | 558.8 | 219.1 | 102.8 | 59.7 |
At 4 steps, continual MultiMDM moves MDLM from 721.1 to 558.8. On Corpus B, 2-step GenPPL falls from 730.3 (MDLM) and 679.1 (DUO) to 400.0. Training MultiMDM from scratch is mostly competitive, not dominant; the gain comes from MDM-then-multi-mask continual training. M=50 is the sweet spot. M=100 gets worse at large step counts (96.7 vs 59.7 at 64 steps).
LoRA adaptation of LLaDA-8B-Base into MultiLLaDA with M=50 lifts MATH500 at 1 token per step from 27.0 to 30.2. GSM8K, HumanEval, and MBPP improve in most settings; HumanEval at 4 tokens per step drops from 14.3 to 9.5. The 8B run does not include shared-Gumbel distillation or full-parameter finetuning.
If the goal is few-step sampling from an existing MDM, this path does not require a continuous embedding or a switch to uniform noise. Designated masks are draft paper for each token: the endpoint has entropy, and noisy positions stay explicitly marked. The 170M unconditional numbers are clear. The 8B result is a transfer check, not a product claim.
The two English corpora are unnamed, so the data recipe cannot be reproduced from the paper. The headline metric is GPT-2 GenPPL, which tracks sample entropy; matching entropy helps, but it is still not a downstream task. The 8B experiment uses LoRA only, skips shared-Gumbel distillation, and regresses on HumanEval at coarser sampling. The designated mapping is a one-hot column assignment, not a semantic one. There is no standalone Limitations section; these caveats sit in the experimental setup and discussion.