DMax: Making Parallel Decoding Actually Work for Diffusion Language Models
TL;DR: Diffusion language models (dLLMs) can decode in parallel in theory, but the hard mask→token two-step design of existing methods collapses as parallelism increases. A team from the National University of Singapore moves decoding from discrete jumps to gradual refinement in embedding space: first teaching the model to correct its own mistakes (OPUT), then giving decoding states a gradual uncertainty profile (SPD). On LLaDA-2.0-mini, TPF for math reasoning jumps from 2.04 to 5.48, and from 2.71 to 5.86 for code generation, with almost no accuracy loss. Two H200 GPUs reach 1,338 tokens/second.
1. The Promise and Reality of Diffusion Models
Diffusion language models (dLLMs, such as 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 reality, accuracy collapses as parallelism grows.
The problem lies in the decoding mechanism. Existing masked dLLMs (like LLaDA) decode in a binary fashion:
1. All positions start as [MASK]
2. The model predicts a token for every masked position
3. Based on confidence, some masks are "promoted" to fixed tokens
4. The rest stay masked and are re-guessed next round
The issue: once a token is promoted, it is locked in. If the model guessed wrong at that step, all subsequent iterations build on the error. The higher the parallelism, the more tokens get promoted at once, and the more likely errors explode simultaneously. Like a snowball — when LLaDA pushes TPF to 6 on GSM8K, accuracy falls from 80%+ to 15%.
This is not the diffusion model's fault — it's the decoding strategy's fault.
2. DMax's Core Insight: Decoding Should Slide, Not Jump
DMax moves decoding from discrete token space to embedding space.
- Traditional: mask → token (hard switch, irreversible)
- DMax: mask → hybrid embedding → token (gradual, correctable)
- Traditional: a position is 100% mask or 100% token. Once decided, it can't change.
- SPD: every 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 position confidences exceed 0.9 (τ_acc)
- τ_dec (decoding threshold): 0.5 (math) / 0.65 (coder) — what counts as "high confidence"
- τ_acc (acceptance threshold): 0.9 — minimum confidence for a block to be "submitted"
- MATH500: DMax 71.6%, LLaDA crashes to 15.2%
- MBPP: DMax 79.2%, LLaDA crashes to 2.3%
- Data: self-distillation. LLaDA-2.0-mini generates its own answers as training targets. 0.7M math samples + 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 in different iterations to avoid 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: Teaching the Model to Correct Itself
The training problem: standard uniform diffusion training samples random noise tokens from the vocabulary as training inputs. But at inference time, the "noise" the model faces isn't random — it's its own wrong guesses from the previous round. The training and inference distributions are misaligned.
OPUT (On-Policy Uniform Training) constructs noise from the model's own predictions: it samples from the model's top-k prediction distribution and feeds the samples back as training inputs. The model learns to recover the correct token both from [MASK] embeddings and from embeddings of its own wrong predictions.
Effect: the model builds a mapping from both mask embeddings and self-predicted token embeddings to the correct answer. This mapping is the foundation of SPD — without it, hybrid embeddings are meaningless.
A brutal ablation confirms this: applying SPD directly to LLaDA without OPUT causes catastrophic collapse. OPUT is not optional; it's mandatory.
2.2 SPD: Leaving an Escape Hatch in Embedding Space
SPD (Soft Parallel Decoding) centers on the Hybrid Embedding:
where \(c\) is the model's prediction confidence for the position. High confidence → the embedding looks more like a token, so the model keeps it. Low confidence → it looks more like a mask, so the model knows to revise it.
Difference from the traditional approach:
Benefit: at the start of each iteration, the model receives confidence signals from the previous step — it knows which positions are stable and which are still uncertain. High-confidence positions don't waste attention; low-confidence positions get concentrated correction. Like an editor marking uncertain sentences lightly and only reviewing those next round.
3. Decoding Pipeline: Block-wise Semi-Autoregressive
DMax is not fully parallel (contextual coherence would suffer). Instead: parallel within blocks, sequential across blocks:
1. Text is split into 32-token blocks 2. Within each block: all positions start masked → 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 criteria:
In-block selection strategy: each round, scan left to right and promote the longest contiguous high-confidence prefix to tokens. At the first low-confidence position, everything to the right stays masked. This ensures the "uncertain future" doesn't disturb the "determined 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) measures how many tokens are determined per forward pass. 2.04 means roughly 2 tokens per step; 5.48 means 5.5 — more than double the speed.
4.2 The More Important Result: Behavior Under Extreme Parallelism
At TPF 6.5 (very aggressive parallelism):
This isn't "DMax is a bit faster" — it's that DMax remains 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 holistically evaluate the parallelism-accuracy trade-off. DMax substantially outperforms the base model and all baselines on every benchmark, showing the advantage isn't from threshold tuning but from a fundamentally more robust framework.
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. A plug-and-play training recipe.
6. Ablations: Who Contributes the Gains?
The paper runs careful controls:
1. Traditional uniform diffusion training: performance actually drops. Random noise mismatches inference-time noise, so the model oscillates during parallel decoding. 2. OPUT without 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.