DMax: Aggressive Parallel Decoding for dLLMs

ArXiv: 2604.08302

Pitch

DMax unlocks the latent speed potential of diffusion language models by solving the error accumulation problem that causes quality collapse under aggressive parallel decoding. Through On-Policy Uniform Training and Soft Parallel Decoding, the method enables models to iteratively self-correct their predictions in embedding space—transforming the binary, irreversible mask-to-token process into a self-revising transformation. This achieves nearly 3x speedup (2.04→5.48 TPF on GSM8K) while maintaining accuracy, proving that dLLMs can deliver on their promise of efficient parallel generation without sacrificing quality.

---

1. Executive Summary

This paper presents DMax, a new paradigm for diffusion language models (dLLMs) that enables aggressive parallel decoding while maintaining generation quality. The core innovation is reformulating the binary mask-to-token decoding process into a self-revising transformation in the embedding space, achieved through two key mechanisms: On-Policy Uniform Training (OPUT), which trains models to correct their own prediction errors, and Soft Parallel Decoding (SPD), which represents intermediate decoding states as interpolations between predicted token embeddings and mask embeddings. On GSM8K, DMax improves tokens per forward (TPF) from 2.04 to 5.48 while preserving 92.1% accuracy (vs. 92.6% baseline); on MBPP, it increases TPF from 2.71 to 5.86 while maintaining comparable performance, demonstrating that dLLMs can achieve near-autoregressive-level parallelism without sacrificing quality.

2. Context and Motivation

The Core Problem: Error Accumulation Limits Parallel Decoding in dLLMs

Diffusion language models (dLLMs) have emerged as a promising alternative to autoregressive language models (AR-LLMs) due to their inherent capacity for parallel decoding—the ability to generate multiple tokens simultaneously rather than sequentially. This parallelism theoretically enables significant inference speedups, as the model can predict multiple positions in a single forward pass rather than requiring one forward pass per token.

However, the paper identifies a fundamental bottleneck that prevents existing dLLMs from realizing this potential: error accumulation. In current masked diffusion language models (MDLMs), decoding follows a binary, one-way mask-to-token process. Each position is either a [MASK] token or a committed token. Once a masked position is decoded, that prediction becomes fixed context for all subsequent decoding steps—it cannot be revised. Under highly parallel decoding, erroneous predictions are inevitable. Once such errors are committed, they contaminate future predictions and trigger cascading error accumulation, ultimately leading to semantic collapse. Unlike speculative decoding in autoregressive models (which can verify and reject incorrect predictions), dLLMs lack any mechanism to recover from incorrect predictions.

This creates a severe tradeoff: dLLMs can either decode conservatively (high accuracy, low parallelism/speed) or aggressively (high parallelism/speed, but accuracy collapses). The paper demonstrates this with a concrete example: when decoding all masked positions at once with an aggressive threshold on GSM8K, accuracy drops to only 68%, making aggressive parallelism impractical for real-world deployment.

Why This Problem Matters

The practical implications are significant for several reasons:

1. Inference efficiency is critical for LLM deployment. Autoregressive models require sequential token generation, creating a fundamental bottleneck. If dLLMs could achieve true parallel decoding at scale, they could dramatically reduce inference costs and latency—key concerns for production systems.

2. The promise of dLLMs remains unfulfilled. While dLLMs have shown theoretical advantages (parallel decoding, bidirectional context, better handling of certain tasks), their practical decoding parallelism has been severely limited by this error accumulation problem. Prior work like LLaDA and Dream scaled MDLMs to billion-parameter regimes but still faced this fundamental limitation.

3. Self-correction is a missing capability. Unlike humans, who can revise and refine their reasoning, current dLLMs commit irrevocably to early predictions. This limits their applicability to tasks requiring iterative refinement or error correction.

Prior Approaches and Their Limitations

The paper surveys several categories of prior work:

Improved decoding strategies. Methods like Hierarchical Decoding use divide-and-conquer procedures to improve parallel decoding. The paper shows (Table 1) that this improves TPF modestly (e.g., 2.04 → 2.44 on GSM8K) but at the cost of accuracy (92.6% → 91.6%). These methods don't address the root cause—they work around error accumulation rather than solving it.

Distillation strategies. Approaches like dParallel use certainty-forcing distillation to accelerate confidence convergence. dParallel-SFT improves TPF to 2.79 on GSM8K with minimal accuracy loss, but the gains remain limited because the underlying error accumulation problem persists.

Uniform Diffusion Language Models (UDLMs). UDLMs train models to recover clean tokens from arbitrary vocabulary tokens (not just [MASK]), enabling token-to-token denoising that naturally allows self-correction. However, UDLM decoding typically starts from fully random sequences, which makes denoising harder and leads to unstable generation. The paper's experiments (Table 1) show that naively applying uniform diffusion training actually degrades performance (GSM8K accuracy drops from 92.6% to 68.7%).

Soft embedding approaches. Methods like Soft-Mask (SM) and EvoToken introduce soft embeddings into decoding but haven't translated this into improved decoding efficiency—they address representation but not the self-correction mechanism.

How This Paper Positions Itself

The paper's central insight is that existing dLLMs lack self-correction capability because their training paradigm doesn't match their inference behavior. At inference, errors come from the model's own predictions, but training uses either masked inputs (MDLMs) or uniformly random tokens (UDLMs)—neither reflects the actual error distribution the model encounters.

DMax addresses this through a unified paradigm that combines:

1. On-Policy Uniform Training (OPUT): Train the model to recover clean tokens from its own predictions, not random tokens—bridging the train-inference gap.

2. Soft Parallel Decoding (SPD): Instead of committing to discrete tokens, represent intermediate states as soft embeddings that preserve uncertainty information across decoding iterations.

The paper positions this as "unifying the strengths of MDLMs and UDLMs"—maintaining stable masked initialization (from MDLMs) while enabling token-to-token self-revision (from UDLMs).

3. Technical Approach

3.1 Reader Orientation

DMax is a training-and-inference framework that transforms diffusion language models from systems that commit irrevocably to predictions into systems that can iteratively refine their own outputs in the embedding space, enabling aggressive parallel decoding without accuracy degradation.

3.2 Big-Picture Architecture (Diagram in Words)

The system has three major components:

1. Base Model (LLaDA-2.0-mini): A pretrained masked diffusion language model that serves as the initialization point. The model predicts tokens at masked positions given a partially masked sequence.

2. On-Policy Uniform Training (OPUT) Module: A training procedure that extends the base MDLM to handle self-predicted tokens as input (not just masks). During training, it constructs two types of noisy inputs—masked sequences (from MDLM training) and on-policy predicted sequences (from the model's own outputs)—and trains the model to recover the original clean sequence from both.

3. Soft Parallel Decoding (SPD) Inference Engine: A decoding procedure that replaces discrete token commitments with "hybrid embeddings"—interpolations between predicted token embeddings and mask embeddings weighted by prediction confidence. These soft embeddings serve as model inputs, preserving uncertainty across iterations.

The data flow: A prompt enters → the block-wise semi-autoregressive process begins with fully masked blocks → at each step, the model predicts tokens → predictions above a confidence threshold become "token positions" represented as hybrid embeddings (not discrete tokens) → positions below threshold remain masked → the model re-predicts all positions using these soft/hard inputs → iteration continues until convergence (predictions stabilize or all positions exceed acceptance threshold) → final tokens are committed.

3.3 Roadmap for the Deep Dive

We will explain the components in the following order:

1. Background on diffusion language model paradigms (MDLMs vs. UDLMs), to establish the theoretical foundation and understand what DMax builds upon.

2. The On-Policy Uniform Training procedure, since this is the training innovation that equips the model with self-correction capability—it must come before we can understand how inference exploits this capability.

3. The Soft Parallel Decoding procedure, the inference mechanism that leverages the trained model's self-correction ability through embedding-space representations.

4. The relationship between OPUT and SPD, explaining why OPUT is a prerequisite for SPD and how they synergize.

5. Training data construction and hyperparameters, the practical details for replication.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a methods paper that introduces a new training-inference paradigm for diffusion language models, with the core idea being: reformulate decoding as a self-revising transformation in embedding space rather than an irrevocable mask-to-token process.

---

Background: MDLMs and UDLMs

Masked Diffusion Language Models (MDLMs) treat text generation as a discrete denoising process. Given a clean token sequence \(x_0 = (x_0^1, \ldots, x_0^L)\) of length \(L\) with vocabulary \(\mathcal{V}\), the corruption process replaces clean tokens with a special [MASK] symbol. At noise level \(t \in [0, 1]\), each token is independently masked with probability \(t\). The MDLM training objective is:

\[\mathcal{L}_{\text{MDLM}}(\theta) = -\mathbb{E}_{x_0, t, x_t}\left[\frac{1}{t}\sum_{i=1}^{L} \mathbb{1}(x_t^i = \text{[MASK]}) \log p_\theta(x_0^i | x_t)\right]\]

In plain language: the model is trained to predict the original tokens only at positions that are masked. Unmasked positions are treated as fixed context. At inference, MDLMs start from a fully masked sequence and iteratively decode masked positions in parallel.

Uniform Diffusion Language Models (UDLMs) generalize this by replacing tokens with uniformly sampled vocabulary tokens rather than [MASK]. The UDLM objective is:

\[\mathcal{L}_{\text{UDLM}}(\theta) = -\mathbb{E}_{x_0, t, x_t}\left[\sum_{i=1}^{L} \log p_\theta(x_0^i | x_t)\right]\]

Here, the model must recover clean tokens from arbitrary noisy token inputs at all positions, not just masked ones. This enables token-to-token denoising where all positions can be re-evaluated at every step—naturally enabling self-correction. However, UDLM decoding typically starts from fully random sequences, leading to unstable generation.

The key difference: MDLMs have stable masked initialization but cannot revise committed tokens; UDLMs can revise tokens but have unstable random initialization. DMax unifies both: masked initialization (stability) + token-to-token refinement (self-correction).

---

On-Policy Uniform Training (OPUT)

The training procedure addresses a critical mismatch: standard UDLM training corrupts sequences with uniformly random tokens, but at inference, noisy inputs come from the model's own predictions. Random vocabulary tokens lie far outside the natural language manifold, forcing the model to waste capacity learning to map unnatural sequences back to plausible language rather than learning useful self-correction behaviors.

The core idea: Construct training inputs by sampling noisy sequences from the model's own predictive distribution (on-policy), not from uniform random sampling.

Training procedure step-by-step:

1. Sample a corruption level: At each training iteration, sample \(t \sim \text{Uniform}(t_l, t_h)\), where \(t_l\) and \(t_h\) are lower and upper bounds of the noise level. In the paper's implementation, a fixed mask ratio of 0.75 is used.

2. Construct masked noisy sequence: Given a clean sequence \(x_0\) from training data, construct \(x_t^{(m)}\) by independently replacing each token with [MASK] with probability \(t\): \(\(x_t^{(m),i} = \begin{cases} \text{[MASK]} & \text{with probability } t \\ x_0^i & \text{otherwise} \end{cases}\)\)

3. Generate on-policy predictions: Feed \(x_t^{(m)}\) into the model \(M_\theta\) and sample predictions at masked positions from the model's predictive distribution: \(\(x_t^{(p),i} = \begin{cases} x_t^{(m),i} & \text{if } x_t^{(m),i} \neq \text{[MASK]} \\ \hat{x}^i \sim p_\theta(\cdot | x_t^{(m)}) & \text{if } x_t^{(m),i} = \text{[MASK]} \end{cases}\)\)

This is a strictly on-policy rollout—the predictions are sampled using the current model parameters at each iteration.

1. Two forward passes: Compute model outputs for both sequences:
2. \(p_\theta^{(m)}(\cdot | x_t^{(m)}) = M_\theta(x_t^{(m)})\) — the model's prediction given masked input

3. \(p_\theta^{(p)}(\cdot | x_t^{(p)}) = M_\theta(x_t^{(p)})\) — the model's prediction given on-policy predicted input

4. Supervise both outputs against clean sequence: Use cross-entropy loss over all positions (regardless of mask status): \(\(\mathcal{L}_{\text{mask}} = -\sum_{i=1}^{L} \log p_\theta^{(m)}(x_0^i | x_t^{(m)})\)\) \(\(\mathcal{L}_{\text{pred}} = -\sum_{i=1}^{L} \log p_\theta^{(p)}(x_0^i | x_t^{(p)})\)\)

5. Combined objective: \(\(\mathcal{L}_{\text{on-policy}} = \mathcal{L}_{\text{mask}} + \mathcal{L}_{\text{pred}}\)\)

Why this works: By training the model to recover clean tokens from its own predictions (the \(\mathcal{L}_{\text{pred}}\) term), OPUT bridges the train-inference gap. The model learns a consistent mapping from both mask embeddings and self-predicted token embeddings toward the correct output. This equips it with self-correction capability while preserving the original mask denoising ability (the \(\mathcal{L}_{\text{mask}}\) term).

Implementation detail: To avoid extra memory overhead, the masked noisy sequence and predicted noisy sequence are optimized in separate iterations within the same epoch rather than jointly in a single iteration.

---

Soft Parallel Decoding (SPD)

Even with OPUT, error accumulation remains challenging when many erroneous predictions arise simultaneously. The paper shows that decoding all positions with threshold 0 (maximally aggressive) still drops GSM8K accuracy to 68% with OPUT alone. Soft Parallel Decoding addresses this by preserving predictive uncertainty across iterations.

The core idea: Instead of treating decoded tokens as discrete, irrevocable commitments, represent each intermediate decoding state as a "hybrid embedding"—an interpolation between the predicted token embedding and the mask embedding, weighted by prediction confidence. The mask embedding naturally encodes maximal uncertainty, so this interpolation serves as an explicit carrier of uncertainty across iterations.

Decoding procedure step-by-step:

The process follows a block-wise semi-autoregressive scheme with block size 32. For each block:

1. Initialize: All positions start as mask positions with mask embeddings as input: \(\(h_j^{(t)} = e_{\text{mask}}, \quad j \in \mathcal{M}^{(t)}\)\) where \(\mathcal{M}^{(t)}\) is the set of mask positions at step \(t\).

2. Model forward pass: Predict distributions over all positions: \(\(p_j(\cdot) = p_\theta(\cdot | \{h_j\}_{j \in B}), \quad \forall j \in B\)\) where \(B\) is the set of all block positions.

3. Compute top-1 predictions and confidences: \(\(\hat{y}_j = \arg\max_y p_j(y), \quad c_j = p_j(\hat{y}_j)\)\)

4. Promote positions to "token status": Use an aggressive confidence threshold \(\tau_{\text{dec}}\). Importantly, the paper uses a contiguous prefix rule: scan the masked region left-to-right and promote only the longest contiguous prefix whose confidence exceeds \(\tau_{\text{dec}}\). Once a position with confidence below \(\tau_{\text{dec}}\) is encountered, all positions to its right remain masked. This keeps the masked region contiguous and prevents unreliable future tokens on the right from interfering with mask predictions on the left.

If no position satisfies the criterion, the leftmost mask position is still promoted to ensure progress.

1. Construct hybrid embeddings for token positions: For each token position \(j \in \mathcal{T}^{(t)}\) (the set of promoted positions), compute the hybrid embedding. Let \(y_j^{(t-1)}\) be the top-1 prediction and \(\pi_j^{(t-1)}\) its probability from the previous step. Assign remaining probability mass to the mask embedding: \(\(\pi_{j,\text{mask}}^{(t-1)} = 1 - \pi_j^{(t-1)}\)\)

The unnormalized hybrid embedding is: \(\(\tilde{h}_j^{(t)} = \pi_j^{(t-1)} e_{y_j^{(t-1)}} + \pi_{j,\text{mask}}^{(t-1)} e_{\text{mask}}\)\)

1. Renormalize: To prevent norm collapse from adding high-dimensional embeddings, renormalize to match the probability-weighted sum of component norms: \(\(h_j^{(t)} = \frac{\tilde{h}_j^{(t)}}{\|\tilde{h}_j^{(t)}\|_2} \left( \pi_j^{(t-1)} \|e_{y_j^{(t-1)}}\|_2 + \pi_{j,\text{mask}}^{(t-1)} \|e_{\text{mask}}\|_2 \right)\)\)

2. Check for convergence: A block is considered converged when either:

3. Consistency criterion: Top-1 predictions at all positions remain unchanged for two consecutive steps

4. Confidence criterion: Confidence of every position exceeds acceptance threshold \(\tau_{\text{acc}} = 0.9\)

5. Commit and move to next block: Once converged, commit all token positions and proceed to the next block.

Why this works: The hybrid embedding serves as a soft intermediate state that explicitly carries forward uncertainty. High-confidence predictions are closer to the predicted token embedding; low-confidence predictions are closer to the mask embedding. This gives the model explicit uncertainty priors from previous steps, enabling more robust self-correction.

Key design choice—contiguous prefix rule: By keeping masked positions as a contiguous left-anchored region, the model avoids interference from unreliable future tokens. The paper's ablations show this improves performance (Table 3).

---

The Critical Relationship: OPUT as a Prerequisite for SPD

A crucial finding: Soft Parallel Decoding must be used with OPUT-trained models. Applying SPD to a standard diffusion language model without OPUT causes "catastrophic performance collapse" (Table 3 shows 0.0% accuracy across all thresholds).

The reason: OPUT trains the model to recover correct targets from both mask embeddings and self-predicted token embeddings, making interpolation between them a meaningful and effective input for denoising. Without OPUT, the model has no concept of how to interpret intermediate points between mask and token embeddings—it was never trained on such inputs.

This is demonstrated empirically in Table 3, where the row "Hybrid Embedding" checked but "On-Policy Rollout" unchecked (applying SPD to original model) yields 0% accuracy. In contrast, OPUT alone (without SPD) achieves 68.2% accuracy at threshold 0, and OPUT + SPD achieves 90.4%.

---

Training Data Construction

The paper uses self-distillation: all training data is generated by the base model itself, not from external high-quality sources.

Math reasoning data (0.7M samples): - Prompts from: GSM8K trainset, PRM12K, subset of Numina-Math, subset of OpenThoughts - Responses generated by LLaDA-2.0-mini with confidence threshold 0.95, block size 32, max length 2048 - Incomplete generations discarded

Code generation data (1.0M samples): - Prompts from: subset of OpenCodeInstruct - Same generation procedure

Training hyperparameters: - Base model: LLaDA-2.0-mini - Fixed mask ratio: 0.75 - Full-parameter fine-tuning for 2 epochs - Batch size: 8 - Initial learning rate: \(2 \times 10^{-6}\) - Schedule: Cosine learning rate schedule - Block size: 32 (block-diffusion setting) - Hardware: 8 H200 GPUs

Two trained models: - DMax-Math: For mathematical reasoning tasks - DMax-Coder: For code generation tasks

---

Inference Configuration

• Block size: 32 (semi-autoregressive block diffusion)
• Acceptance threshold: \(\tau_{\text{acc}} = 0.9\)
• Decoding threshold \(\tau_{\text{dec}}\): 0.5 for DMax-Math, 0.65 for DMax-Coder
• Evaluation framework: dInFer
• Hardware: 2 H200 GPUs with tensor parallelism
• Max generation length: 2048 tokens
• Evaluation: Zero-shot, batch size 1

4. Key Insights and Innovations

Innovation 1: On-Policy Training to Bridge the Train-Inference Gap

The most conceptually significant contribution is identifying and addressing the train-inference mismatch in existing approaches. Standard UDLM training constructs noisy sequences by uniformly sampling tokens from the vocabulary. But at inference, noisy inputs come from the model's own predictions—not random tokens. These on-policy predictions have a very different distribution: they lie closer to the language manifold but contain structured errors characteristic of the model's failure modes.

By constructing training inputs on-policy (sampling from the model's own predictions), OPUT ensures the model learns to correct the actual types of errors it will encounter at inference. This is analogous to the insight behind on-policy RL algorithms like PPO—training on the distribution you'll act on improves policy quality.

The empirical validation is striking: Table 1 shows that conventional uniform diffusion training actually degrades GSM8K accuracy from 92.6% to 68.7% with an AUP score of 0, while OPUT achieves 92.1% with dramatically improved TPF. This demonstrates that simply enabling self-correction isn't enough—the training distribution must match the inference distribution.

Innovation 2: Soft Embedding Representations for Iterative Refinement

The Soft Parallel Decoding mechanism is a clean, theoretically motivated solution to preserving uncertainty across decoding iterations. Rather than committing to discrete tokens (which lose all uncertainty information), SPD represents intermediate states as points on a continuous spectrum between "fully confident prediction" and "maximal uncertainty" (mask embedding).

The key insight is that the mask embedding naturally encodes maximal uncertainty in the model's representation space—it's the symbol the model has been trained to associate with "unknown." Interpolating toward the mask embedding for low-confidence predictions gives the model an explicit signal: "this prediction might be wrong, treat it skeptically."

The renormalization step (Equation 10) is a subtle but important design choice—without it, adding high-dimensional embeddings distorts magnitudes and causes representation collapse.

Innovation 3: Contiguous Prefix Rule for Structured Decoding

The paper introduces a simple but effective rule: keep the masked region as a contiguous prefix. When promoting positions from mask to token status, only the longest contiguous prefix (scanning left-to-right) whose confidence exceeds the threshold is promoted.

This design prevents a practical failure mode: if low-confidence predictions are scattered throughout the block, they can interfere with mask predictions elsewhere. By ensuring the masked region remains contiguous and left-anchored, the model avoids contamination from unreliable "future" tokens (positions to the right of the current decoding focus).

Table 3's ablation confirms this: adding the contiguous prefix rule improves accuracy from 91.3% to 92.1% at threshold 0.5 (comparing row with "Contiguous Prefix" unchecked vs. checked).

Innovation 4: Training-Inference Synergy (OPUT as Prerequisite for SPD)

The paper makes a critical discovery: OPUT and SPD are not independent improvements—they must be combined. Table 3 unambiguously shows that applying SPD to an untrained model causes complete collapse (0% accuracy). The explanation reveals a deep dependency:

• OPUT trains the model on inputs that interpolate between mask and token embeddings (by using both \(x_t^{(m)}\) and \(x_t^{(p)}\))
• This creates a model that can interpret and denoise from hybrid inputs
• SPD then exploits this capability by constructing such hybrid inputs at inference

Without OPUT, the model has never seen inputs that blend mask and token embeddings, so SPD's hybrid embeddings fall outside its training distribution, causing collapse. This training-inference synergy is a non-obvious finding that the paper's ablations carefully establish.

5. Experimental Analysis

Evaluation Methodology

Benchmarks: - Mathematical reasoning: GSM8K, MATH500, Minerva-Algebra, ASDIV - Code generation: HumanEval-Instruct, MBPP-Instruct - All evaluations use chain-of-thought prompting for math benchmarks

Metrics: - TPF (Tokens Per Forward): Average number of tokens decoded per forward pass—higher is better (more parallelism) - TPS (Tokens Per Second): Actual throughput measured on hardware - Accuracy: Task-specific accuracy (correct answer for math, pass@1 for code) - AUP Score: Area under the accuracy-parallelism curve—comprehensive measure of parallel decoding performance

Baselines: 1. LLaDA-2.0-mini: Base model with default confidence-threshold decoding (threshold 0.95) 2. Hierarchical Decoding: Advanced inference strategy with divide-and-conquer procedure (low threshold 0.2) 3. dParallel-SFT (LLaDA-2.0-mini-CAP): Model incorporating certainty-forcing loss for improved parallelism 4. Uniform Diffusion Training: Naive extension continuing training with standard UDLM objective (random token corruption)

Experimental setup: - Zero-shot evaluation, batch size 1 - Hardware: 2 H200 GPUs with tensor parallelism - Framework: dInFer - Max generation length: 2048 tokens - Decoding thresholds: \(\tau_{\text{dec}} = 0.5\) for DMax-Math, \(\tau_{\text{dec}} = 0.65\) for DMax-Coder - Acceptance threshold: \(\tau_{\text{acc}} = 0.9\)

---

Main Quantitative Results

Table 1: Main Comparison (Math & Reasoning Benchmarks)

Benchmark	Method	TPF	TPS	Accuracy	AUP Score
GSM8K	LLaDA-2.0-mini	2.04	512	92.6%	340
GSM8K	Hierarchical Decoding	2.44	577	91.6%	357
GSM8K	dParallel SFT	2.79	721	92.3%	395
GSM8K	Uniform Diffusion Training	2.26	493	68.7%	0
GSM8K	DMax-Math	5.48	1258	92.1%	557
MATH500	LLaDA-2.0-mini	2.58	626	75.8%	257
MATH500	DMax-Math	5.94	1286	75.4%	507
Minerva-Algebra	LLaDA-2.0-mini	3.01	755	91.4%	363
Minerva-Algebra	DMax-Math	7.03	1492	91.5%	658

Table 1: Main Comparison (Code Generation Benchmarks)

Benchmark	Method	TPF	TPS	Accuracy	AUP Score
HumanEval-Instruct	LLaDA-2.0-mini	4.38	1044	84.2%	369
HumanEval-Instruct	DMax-Coder	7.36	1557	83.5%	637
MBPP-Instruct	LLaDA-2.0-mini	2.71	662	80.6%	276
MBPP-Instruct	DMax-Coder	5.86	1264	79.2%	482

Key findings:

1. Dramatic TPF improvement: DMax nearly triples TPF on average (from 2.8 to 6.2) while preserving accuracy. On GSM8K, TPF improves from 2.04 to 5.48 (2.7×); on Minerva-Algebra, from 3.01 to 7.03 (2.3×).

2. Accuracy preservation: Accuracy drops are minimal—GSM8K: 92.6% → 92.1% (-0.5%); MATH500: 75.8% → 75.4% (-0.4%); HumanEval: 84.2% → 83.5% (-0.7%). In some cases, accuracy slightly improves (Minerva-Algebra: 91.4% → 91.5%).

3. AUP Score dominance: DMax consistently achieves the highest AUP scores, indicating superior accuracy-parallelism trade-offs across all threshold settings. On GSM8K, DMax achieves 557 vs. the next best 395 (dParallel SFT).

4. Baseline failures: Uniform Diffusion Training catastrophically fails—GSM8K accuracy drops to 68.7% with AUP score of 0, validating the paper's claim that naive uniform training creates a train-inference mismatch.

---

Ablation Studies

Table 3: Ablation on Training and Inference Strategies

The paper conducts a comprehensive ablation on GSM8K across three decoding thresholds (0.95, 0.5, 0.0):

Training	Inference	τdec = 0.95	τdec = 0.50	τdec = 0.0
Original model	Standard	2.04 TPF, 92.6%	4.47 TPF, 78.0%	7.86 TPF, 0.9%
+ SPD only	Hybrid Emb.	1.04 TPF, 0.0%	1.73 TPF, 0.0%	5.39 TPF, 0.0%
+ OPUT only	Standard	2.95 TPF, 92.6%	5.14 TPF, 90.1%	5.89 TPF, 68.2%
+ OPUT	Contiguous	2.85 TPF, 93.0%	5.28 TPF, 91.3%	5.98 TPF, 69.6%
+ OPUT	Hybrid Emb.	3.25 TPF, 92.8%	5.64 TPF, 91.4%	6.01 TPF, 90.4%
+ OPUT	Both	3.00 TPF, 93.3%	5.48 TPF, 92.1%	6.01 TPF, 90.4%

Critical findings:

1. SPD without OPUT causes collapse: The second row shows 0% accuracy across all thresholds when applying hybrid embeddings to the original model—this is catastrophic failure, validating that OPUT is a prerequisite.

2. OPUT alone enables self-correction: With just OPUT (third row), accuracy at τdec = 0.5 improves from 78.0% to 90.1%—a 12.1 percentage point gain. At τdec = 0, accuracy improves from 0.9% to 68.2%.

3. SPD completes the picture: Adding hybrid embeddings (fifth row) further improves τdec = 0 accuracy from 68.2% to 90.4%—demonstrating SPD's value for extreme parallelism.

4. Both components work synergistically: The full system (last row) maintains 90.4% accuracy even at τdec = 0 (maximally aggressive), compared to 0.9% for the baseline—a 100× improvement in the accuracy-parallelism tradeoff at the extreme.

Table 4: Convergence Criteria Ablation

Consistency	Confidence	GSM8K TPF	GSM8K Acc.	MBPP TPF	MBPP Acc.
✓		5.13	92.1%	5.16	79.9%
	✓	2.28	92.2%	3.36	80.1%
✓	✓	5.48	92.1%	5.86	79.2%

The consistency criterion (predictions unchanged for two consecutive steps) is the primary convergence signal. Adding the confidence criterion (all positions > 0.9) can improve TPF by allowing earlier termination, saving the final forward pass. Accuracy is unaffected by the choice.

---

Accuracy-TPF Trade-off Curves

Figure 4 presents curves showing accuracy vs. TPF across different decoding thresholds:

• GSM8K: At ~6.5 TPF, DMax maintains 92%+ accuracy while the original model drops to ~40%
• MATH500: At ~6.5 TPF, DMax achieves 71.6% vs. original's 15.2%
• MBPP: At similar TPF (~5.8), DMax achieves 79.2% vs. original's 2.3%

These curves demonstrate that DMax dramatically shifts the Pareto frontier for the accuracy-parallelism tradeoff.

---

Performance at Low Parallelism

Table 2 shows DMax improves accuracy even in low-parallelism regimes:

Benchmark	LLaDA TPF	LLaDA Acc.	DMax TPF	DMax Acc.
GSM8K	2.04	92.6%	3.54	93.4% (+0.8%)
MATH500	2.58	75.8%	3.45	78.0% (+2.2%)
HumanEval	4.38	84.2%	4.58	87.2% (+3.0%)
MBPP	2.71	80.6%	3.58	83.4% (+2.8%)

The gains (0.8%–3.0%) come from iterative re-evaluation of earlier predictions, allowing recovery from reasoning errors that would otherwise persist.

---

Assessment: Do the Experiments Support the Claims?

Claim 1: DMax enables aggressive parallel decoding while preserving accuracy. Strongly supported. Across all benchmarks, DMax achieves 2–3× higher TPF with minimal accuracy loss. The AUP scores comprehensively demonstrate superior trade-offs.

Claim 2: OPUT bridges the train-inference gap. Strongly supported. The comparison with uniform diffusion training (Table 1) shows that naive uniform training fails catastrophically while OPUT succeeds—same task, different training distribution. The ablations further validate that OPUT alone (without SPD) substantially improves accuracy at aggressive thresholds.

Claim 3: SPD requires OPUT. Strongly supported. The ablation (Table 3, row 2) shows SPD alone causes complete collapse (0% accuracy). This is a critical finding with clear mechanistic explanation.

Claim 4: Self-distillation with model-generated data suffices. Supported but not extensively analyzed. The paper uses only model-generated responses as training data without external supervision, achieving strong results. However, there's no comparison against externally supervised training, so it's unclear whether external data would help further.

Limitations in experimental design: - All experiments use LLaDA-2.0-mini as the base model—generalization to other dLLM architectures is not tested - Training data is self-distilled (model-generated), which may limit capabilities to what the base model already knows—no exploration of external supervision - Only zero-shot evaluation is reported; few-shot or fine-tuned performance is not analyzed - Latency/TPS at different batch sizes beyond batch size 1 is not explored

6. Limitations and Trade-offs

Assumption: The Base Model Must Already Possess Sufficient Capability

DMax operates by enabling self-correction of prediction errors, but it cannot create capabilities the base model lacks entirely. The self-distillation training data is generated by LLaDA-2.0-mini itself—the model learns to recover correct answers from its own predictions, but those correct answers must exist in the model's output distribution in the first place.

This creates an implicit ceiling: DMax can improve efficiency and reduce error accumulation, but it cannot teach the model new reasoning patterns or factual knowledge beyond what the base model can already produce. The paper acknowledges that training data is "obtained from the model's own generations" (Section 4.1), but does not explore whether external high-quality supervision could push capabilities further.

The empirical results hint at this limitation: accuracy improvements at low parallelism are modest (0.8%–3.0% across benchmarks in Table 2). These gains come from error correction, not capability expansion. For tasks where the base model's pass@1 is already high, DMax provides efficiency benefits; for tasks where the base model fundamentally fails, DMax offers no solution.

Single Architecture Dependency

All experiments use LLaDA-2.0-mini as the base model. The paper claims this model is "a state-of-the-art open-source diffusion language model" (Section 4.1), but provides no evidence that the approach transfers to other dLLM architectures such as Dream, Mercury, or different scaling variants of LLaDA.

This matters because OPUT's effectiveness depends on specific properties of the base model's embedding space. The hybrid embedding construction assumes that linear interpolation between token embeddings and mask embeddings produces semantically meaningful intermediate representations. This assumption may not hold universally—different architectures may have different embedding geometries. The paper's ablation shows that SPD fails catastrophically without OPUT, but does not test whether OPUT transfers to other architectures.

Additionally, the base model is relatively small (implied by the "-mini" designation and the training scale of 8 GPUs). Whether the approach scales to larger models (LLaDA's 100B variant mentioned in the related work) remains untested. Larger models may have different error profiles—perhaps fewer single-step errors but more complex multi-step reasoning failures—which could change the optimal training strategy.

Computational Overhead Not Fully Analyzed

The paper reports impressive TPS improvements (e.g., GSM8K: 512 → 1258 TPS), but several computational aspects remain unclear:

Memory overhead of hybrid embeddings: The renormalization step (Equation 10) requires computing norms of both the predicted token embedding and the mask embedding, then scaling. This adds operations per token position per decoding step. The paper does not quantify this overhead relative to standard decoding.

Iteration count variability: SPD iterates until convergence, but the number of iterations varies per block. The paper reports TPF (tokens per forward) but not the total number of forward passes. Under aggressive thresholds, more iterations may be needed, potentially offsetting parallelism gains.

Latency vs. throughput trade-off: All evaluations use batch size 1, optimizing for latency. For throughput-oriented deployments (large batch sizes), the memory cost of maintaining hybrid embeddings across many sequences could become significant. The paper does not explore batched inference.

Threshold Sensitivity

The approach requires setting multiple thresholds: decoding threshold \(\tau_{\text{dec}}\), acceptance threshold \(\tau_{\text{acc}}\), and the mask ratio during training (fixed at 0.75). The paper uses different \(\tau_{\text{dec}}\) values for math (0.5) and code (0.65), suggesting the optimal settings are task-dependent.

Figure 4 shows that DMax's accuracy-TPF curve is flatter than the baseline's, indicating reduced sensitivity to threshold choice—a strength. However, the paper does not provide guidance on how to select these thresholds for new tasks or domains. Practitioners would need to tune thresholds on a validation set, adding deployment complexity.

Limited Task Coverage

The evaluation covers mathematical reasoning (4 benchmarks) and code generation (2 benchmarks). These are structured generation tasks with well-defined correctness criteria. The paper does not test:

• Open-ended generation (story writing, dialogue): Error accumulation may manifest differently when there's no single correct answer
• Long-form reasoning (document writing, multi-step planning): The block-wise approach may struggle with coherence across many blocks
• Factual QA: Self-correction for factual errors may require different training strategies

The contiguous prefix rule specifically assumes a left-to-right structure where positions to the right depend on positions to the left. This may not generalize to tasks requiring bidirectional reasoning or non-sequential dependencies.

Convergence Can Be Slow in Extreme Cases

The ablation in Table 4 reveals that relying solely on the confidence criterion for convergence (without consistency checking) results in much lower TPF (2.28 vs. 5.48 on GSM8K). This suggests that confidence alone is an unreliable convergence signal—the model may have high confidence in incorrect predictions.

The consistency criterion (predictions unchanged for two consecutive steps) is more reliable but requires additional forward passes. In the worst case, if the model oscillates between predictions, convergence could take many iterations. The paper does not analyze convergence speed distributions or worst-case behavior.

The 0% AUP Score for Uniform Diffusion Training Requires Investigation

Table 1 reports an AUP score of 0 for the Uniform Diffusion Training baseline across all benchmarks. This is unusual—it suggests that increasing parallelism (adjusting thresholds) provides no accuracy benefit whatsoever, or that the accuracy-parallelism curve is completely flat.

The paper attributes this to "unstable oscillations within each block" (Section 4.2), but does not provide detailed analysis. Understanding why naive uniform training fails so completely—and whether this failure is inevitable or could be mitigated with different hyperparameters—would strengthen the paper's claims about the necessity of on-policy training.

No Comparison to Autoregressive Alternatives

The paper compares against other dLLM approaches but does not benchmark against autoregressive models with speculative decoding or other inference acceleration techniques. This makes it difficult to assess the practical value proposition: is a dLLM with DMax faster than an AR-LLM with speculative decoding on the same hardware?

The reported throughput (~1250 TPS on GSM8K at batch size 1) is impressive, but without an AR-LLM comparison, practitioners cannot make informed deployment decisions. The related work mentions speculative decoding (Leviathan et al., 2023; Cai et al., 2024; Li et al., 2024) but does not include them in experiments.

Potential Distribution Shift at Inference Time

OPUT trains on self-generated data with a specific generation configuration: confidence threshold 0.95, block size 32, max length 2048. At inference, different thresholds and generation lengths may produce outputs that diverge from the training distribution.

The paper uses thresholds 0.5 (math) and 0.65 (code) at inference—lower than the 0.95 used for training data generation. This means the model encounters more aggressive decoding during inference than during training, potentially creating a new distribution shift. The results suggest this works, but the mismatch between training and inference threshold regimes is not analyzed.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper makes a fundamental contribution to diffusion language model research by identifying and solving the train-inference mismatch that plagued prior approaches to enabling self-correction. The key insight—that models must be trained on their own prediction errors, not random noise—reframes how researchers should think about extending dLLMs beyond their original training paradigm.

Prior to this work, the field faced a frustrating dilemma: UDLMs theoretically enabled self-correction but failed catastrophically in practice (as Table 1's Uniform Diffusion Training results starkly demonstrate). This failure may have led researchers to conclude that self-correction was inherently incompatible with dLLMs or that much more complex training schemes were needed. DMax shows that the solution is conceptually simple—sample training data from the model's own distribution—but critically important.

The demonstration that OPUT is a prerequisite for SPD is equally significant. This reveals that training and inference modifications cannot be designed independently; they must co-evolve. Attempting to improve inference without corresponding training changes (applying SPD to an untrained model) causes complete collapse. This finding should guide future research: inference-time improvements must be validated with compatible training procedures, not bolted onto existing models.

Perhaps most importantly, DMax demonstrates that dLLMs can achieve near-autoregressive-level parallelism without sacrificing quality. The accuracy-TPF curves in Figure 4 show a dramatically shifted Pareto frontier: at similar accuracy levels, DMax achieves 2–3× higher TPF. This validates the long-held promise of dLLMs as efficient alternatives to AR-LLMs, potentially reinvigorating research investment in this direction.

Follow-Up Research This Work Enables or Suggests

1. External supervision integration. The paper uses only self-generated training data, which limits capabilities to what the base model already knows. A natural extension is combining OPUT with external high-quality data (e.g., human-written solutions, outputs from stronger models). The on-policy rollout mechanism could be applied to predictions that need to match external targets, potentially enabling capability gains alongside efficiency gains.

2. Architecture transfer and scaling studies. Testing OPUT and SPD on other dLLM architectures (Dream, Mercury, different LLaDA scales) would reveal which components are architecture-specific versus universally applicable. Scaling to larger models (LLaDA's 100B variant) is particularly important—larger models may have different error distributions, potentially changing optimal training strategies.

3. Dynamic threshold adaptation. Currently, thresholds are fixed per task. Future work could develop adaptive threshold policies that adjust based on block-level confidence statistics or uncertainty estimates. For example, a block with uniformly low confidence might warrant a higher decoding threshold (more conservative), while high-confidence blocks could use lower thresholds (more aggressive).

4. Alternative embedding interpolation schemes. The hybrid embedding uses linear interpolation weighted by confidence. Other interpolation schemes might be more effective: - Learned interpolation weights (a small network predicting the mixing coefficient) - Non-linear interpolation paths in embedding space - Multiple intermediate points (not just token and mask endpoints)

5. Combination with other acceleration techniques. DMax focuses on parallel decoding but could be combined with: - KV caching optimizations (the paper cites dKV-cache but does not combine it) - Sparse attention mechanisms (Sparsed, cited in related work) - Quantization for further speedups

6. Theoretical analysis of convergence. The empirical results show effective self-correction, but theoretical analysis is lacking. Under what conditions does OPUT guarantee improvement? What are the convergence properties of SPD? Formal analysis could inform hyperparameter selection and predict failure modes.

7. Extension to multimodal tasks. The related work mentions multimodal dLLMs (LLaDA-V, Dream-VL, MMaDA). Extending OPUT and SPD to vision-language tasks—where self-correction might help align text and image reasoning—could demonstrate broader applicability.

Practical Applications and Downstream Use Cases

Latency-sensitive deployments. The dramatic TPS improvements (e.g., 512 → 1258 on GSM8K) make DMax attractive for real-time applications where response latency matters: interactive coding assistants, mathematical tutoring systems, or conversational agents. The ability to decode multiple tokens in parallel without quality degradation directly translates to better user experience.

Edge deployment scenarios. Smaller models with efficient parallel decoding could enable running capable dLLMs on resource-constrained devices (laptops, phones) where autoregressive models' sequential generation would be too slow. The paper's batch-size-1 results are particularly relevant here.

High-throughput batch inference. While the paper focuses on batch size 1, the efficiency gains should compound in batch processing scenarios (evaluating many examples, generating training data). Organizations running large-scale inference workloads could see significant cost reductions.

Self-improvement pipelines. The self-distillation approach used for training data construction suggests a recursive application: use DMax to generate high-quality outputs, fine-tune with OPUT, use the improved model to generate better data, and repeat. This could enable automated capability improvement without human supervision.

When to Prefer This Method Over Alternatives

Prefer DMax when: - The base model already has reasonable accuracy on the target task (DMax improves efficiency, not capability) - Latency or throughput is a primary concern - The task involves structured generation (reasoning, code) where error accumulation is the main failure mode - You can afford the training overhead of OPUT (full-parameter fine-tuning for 2 epochs) - Batch size 1 or small batch inference is the target deployment scenario

Prefer standard AR-LLMs with speculative decoding when: - The task requires capabilities beyond what dLLMs currently offer - Batched inference at large batch sizes is needed (untested regime for DMax) - You want mature tooling and deployment infrastructure (dLLM support is less developed)

Prefer standard MDLM decoding when: - Accuracy at conservative thresholds is the only metric that matters - Training compute is unavailable for OPUT fine-tuning - The task requires maximum reliability (e.g., medical or legal applications) and parallelism is not needed

Reproduction and Integration Guidance

Training OPUT from scratch: 1. Start with a pretrained MDLM (LLaDA-2.0-mini in the paper) 2. Generate training data using the base model with high confidence threshold (0.95) to ensure quality 3. For each training iteration: - Sample corruption level \(t\) (paper uses fixed 0.75) - Construct masked sequence \(x_t^{(m)}\) - Generate on-policy predictions \(x_t^{(p)}\) by sampling from the model - Compute losses \(\mathcal{L}_{\text{mask}}\) and \(\mathcal{L}_{\text{pred}}\) separately - Update with combined loss 4. Train for 2 epochs with learning rate \(2 \times 10^{-6}\), cosine schedule

Implementing Soft Parallel Decoding: 1. Use block size 32 for semi-autoregressive generation 2. For each block: - Initialize all positions with mask embeddings - At each step, compute model predictions and confidences - Promote the longest contiguous prefix exceeding \(\tau_{\text{dec}}\) - For promoted positions, construct hybrid embeddings per Equations 8–10 - Check convergence: predictions unchanged for 2 steps OR all confidences > \(\tau_{\text{acc}} = 0.9\) 3. Commit block and move to next

Critical integration note: Do not attempt to use SPD without OPUT training—Table 3 shows this causes complete collapse (0% accuracy). The two components must be deployed together.

Hyperparameter recommendations: - Training mask ratio: 0.75 - Decoding threshold \(\tau_{\text{dec}}\): 0.5 (math), 0.65 (code) — may need tuning for other tasks - Acceptance threshold \(\tau_{\text{acc}}\): 0.9 - Block size: 32 - Use both consistency and confidence convergence criteria for best TPF