DMax: Making Parallel Decoding Actually Work for Diffusion Language Models
> TL;DR: Diffusion language models (dLLMs) can decode in parallel in theory, but existing hard mask→token decoding collapses accuracy at high parallelism. The NUS team moves decoding from discrete jumps to gradual refinement in embedding space: first teach the model to correct its own mistakes (OPUT), then add a confidence-graded uncertainty to decoding states (SPD). On LLaDA-2.0-mini, math reasoning TPF jumps from 2.04 to 5.48, code generation from 2.71 to 5.86, with almost no accuracy loss. It runs at 1338 tokens/second on two H200 GPUs.
1. The Promise and Reality of Diffusion LMs
Diffusion language models (dLLMs like LLaDA and Dream) promise parallel decoding — instead of generating tokens strictly left-to-right like autoregressive models (GPT/Claude), a dLLM can guess many positions at once and iteratively refine them.
In theory this multiplies inference speed. In practice, accuracy collapses as parallelism increases.
The problem is the decoding mechanism. Existing masked dLLMs (e.g., LLaDA) decode in binary steps:
1. All positions start as [MASK] 2. The model predicts what token each masked position should hold 3. Based on confidence, some masks are "promoted" to fixed tokens 4. Remaining masks stay masked and are re-guessed next round
The issue: once a token is promoted, it is locked in. If the model guessed wrong, all later iterations build on that error. Higher parallelism means more tokens promoted at once, so errors erupt simultaneously and snowball — pushing LLaDA to TPF 6 on GSM8K drops accuracy from 80%+ to 15%.
This is not the diffusion model's fault — it is the decoding strategy's fault.
2. DMax's Core Insight: Decoding Should Slide, Not Jump
DMax moves decoding from discrete token space into embedding space.
- Traditional: mask → token (hard, irreversible switch)
- DMax: mask → hybrid embedding → token (gradual, correctable)
- Traditional: a position is either 100% mask or 100% token. Once decided, it cannot change.
- SPD: each position slides between 0 and 1. High confidence is "light gray," low confidence is "dark gray" — not black and white.
- All positions unchanged for two consecutive rounds, or
- All positions' confidence exceeds 0.9 (τacc)
- τdec (decoding threshold): 0.5 (math) / 0.65 (coder) — what counts as "high confidence"
- τacc (acceptance threshold): 0.9 — minimum confidence before a block can "turn in"
- MATH500: DMax 71.6%, LLaDA collapses to 15.2%
- MBPP: DMax 79.2%, LLaDA collapses to 2.3%
- Data: self-distillation. The model's own generated answers serve as training targets. 0.7M math + 1.0M code samples. No external high-quality data.
- Hardware: 8× H200, full-parameter fine-tuning
- Duration: 2 epochs
- Config: mask ratio 0.75, block size 32, learning rate 2e-6, cosine schedule
- Trick: masked noisy sequences and predicted noisy sequences are optimized at different iterations, avoiding extra memory overhead
- Paper: DMax: Aggressive Parallel Decoding for dLLMs
- arXiv: 2604.08302
- Authors: Zigeng Chen, Gongfan Fang, Xinyin Ma, Ruonan Yu, Xinchao Wang
- Institution: National University of Singapore
- Dates: 2026-04-09 (v1), 2026-04-20 (v2), 2026-05-15 (v3)
- Code: https://github.com/czg1225/DMax
- Base model: LLaDA-2.0-mini
- Training data: self-distillation (0.7M math + 1.0M code)
- Training hardware: 8 × H200
- Inference hardware: 2 × H200 @ batch size 1
- Key numbers: GSM8K TPF 5.48, MBPP TPF 5.86, 1338 TPS, τdec 0.5/0.65, τacc 0.9, block size 32
2.1 OPUT: Teach the Model to Self-Correct
Under standard uniform diffusion training, noise tokens are sampled randomly from the vocabulary. But at inference time, the "noise" the model faces is not random — it is its own wrong predictions from the previous round. Training and inference distributions are misaligned.
OPUT (On-Policy Uniform Training) constructs noise from the model's own predictions: sample from the model's top-k prediction distribution and feed the sampled results back as inputs in subsequent training. The model learns to recover the correct token from both [MASK] embeddings and from its own wrong token embeddings.
This creates a mapping from both mask embeddings and self-predicted token embeddings to correct answers — the foundation SPD needs. Without it, hybrid embeddings are meaningless. A brutal ablation: applying SPD directly to LLaDA without OPUT causes catastrophic collapse. OPUT is mandatory, not optional.
2.2 SPD: Leave an Escape Route in Embedding Space
SPD (Soft Parallel Decoding) uses a Hybrid Embedding:
where \(c\) is the model's prediction confidence for that position. High confidence → the embedding behaves more like a token and tends to be kept. Low confidence → it behaves more like a mask, signaling the model to revise it.
Differences from the traditional approach:
Benefit: at each iteration, the model receives the previous step's confidence signal — it knows which positions are stable and which are still uncertain. High-confidence positions don't waste attention; low-confidence positions get focused correction. Like an editor highlighting uncertain sentences and revisiting only those.
3. Decoding Procedure: Block-wise Semi-Autoregression
DMax is not fully parallel (that would break context relations). It runs parallel within blocks, sequential across blocks:
1. Text is split into 32-token blocks 2. Within each block: mask all positions → iterative prediction → hybrid embeddings → until convergence 3. The next block starts only after the current one finishes (preserving left-to-right causality)
In-block convergence conditions:
In-block selection: each round scans left to right and promotes the longest contiguous high-confidence prefix to tokens. At the first low-confidence position, everything to the right stays masked — the "uncertain future" never contaminates the "confirmed past."
Two thresholds:
4. Performance: From "Parallelism = Collapse" to "Parallelism = Speedup"
4.1 Key Numbers
| Model | Task | TPF | Accuracy | TPS | |-------|------|-----|----------|-----| | LLaDA-2.0-mini | GSM8K | 2.04 | 82.6% | ~400 | | DMax-Math | GSM8K | 5.48 | 82.3% | 1338 | | LLaDA-2.0-mini | MBPP | 2.71 | 74.5% | ~500 | | DMax-Coder | MBPP | 5.86 | 74.3% | 1338+ |
TPF (Tokens Per Forward): how many tokens are determined per forward pass. 5.48 means ~5.5 tokens per step instead of 2 — more than double the speed.
4.2 More Important: the Gap Under Extreme Parallelism
At TPF 6.5 (very aggressive parallelism):
This is not "DMax is slightly faster" — it is that DMax stays stable at parallelism levels where LLaDA completely breaks down. A qualitative leap.
4.3 AUP Score
The paper introduces AUP (Area Under the Parallelism curve) to evaluate the parallelism-accuracy trade-off. DMax leads the original model and all baselines by a wide margin on every benchmark — showing the advantage is inherent to the framework, not a matter of threshold tuning.
5. Training Details: Low-Cost Adaptation
DMax is not trained from scratch — it fine-tunes LLaDA-2.0-mini:
Training cost is low — no new data, no large-scale pretraining. It is a plug-and-play training recipe.
6. Ablations: Who Contributes What?
1. Traditional uniform diffusion training: performance actually degrades. Random noise mismatches inference noise, causing oscillation during parallel decoding. 2. OPUT only, no SPD: better than baseline, but still fails under aggressive parallelism (GSM8K @ τdec=0: 68% accuracy). 3. OPUT + SPD: 68% → 90%, with higher speed.
Conclusion: OPUT is the foundation, SPD is the accelerator. Both are indispensable.