DMax: Aggressive Parallel Decoding for dLLMs — part 2/3

Cleaned-text code points [15479, 31061). Bibliography omitted here; source.md includes it.

Decoding Procedure. An overview of the decoding process is shown in Figure 3 . It follows a block-wise semi-autoregressive process. For each block, we partition its positions into two sets: mask positions and token positions . At initialization, all positions in the block are mask positions. At each decoding step, we use an aggressive confidence threshold τ dec \tau_{\mathrm{dec}} to promote some mask positions into token positions. Specifically, we scan the masked region from left to right and promote only its longest contiguous prefix whose confidence exceeds τ dec \tau_{\mathrm{dec}} . Once the first mask position with confidence below τ dec \tau_{\mathrm{dec}} is encountered, all mask positions to its right remain masked. If no mask position satisfies this criterion, we still promote the leftmost mask position to ensure decoding progress. This design keeps the masked region contiguous and prevents unreliable future tokens on the right from interfering with mask predictions on the left.

At decoding step t t , every mask position uses the mask embedding as model input: 𝐡 j ( t ) = 𝐞 mask , j ∈ ℳ ( t ) . \mathbf{h}^{(t)}_{j}=\mathbf{e}_{\mathrm{mask}},\qquad j\in\mathcal{M}^{(t)}. (7) where ℳ ( t ) \mathcal{M}^{(t)} denotes the set of mask positions at step t t .

For each token position j ∈ 𝒯 ( t ) j\in\mathcal{T}^{(t)} , where 𝒯 ( t ) \mathcal{T}^{(t)} is the set of token positions, we construct a hybrid embedding from the top-1 prediction at the previous step t − 1 t-1 as the model input. Let y j ( t − 1 ) y^{(t-1)}_{j} denote the top-1 predicted token at position j j , and let π j ( t − 1 ) \pi^{(t-1)}_{j} be its predicted probability. We assign the remaining probability mass to the mask embedding: π j , mask ( t − 1 ) = 1 − π j ( t − 1 ) . \pi^{(t-1)}_{j,\mathrm{mask}}=1-\pi^{(t-1)}_{j}. (8) The unnormalized hybrid embedding is then 𝐡 ~ j ( t ) = π j ( t − 1 ) ​ 𝐞 ​ ( y j ( t − 1 ) ) + π j , mask ( t − 1 ) ​ 𝐞 mask , j ∈ 𝒯 ( t ) . \tilde{\mathbf{h}}^{(t)}_{j}=\pi^{(t-1)}_{j}\,\mathbf{e}\!\left(y^{(t-1)}_{j}\right)+\pi^{(t-1)}_{j,\mathrm{mask}}\,\mathbf{e}_{\mathrm{mask}},\qquad j\in\mathcal{T}^{(t)}. (9)

Directly adding high-dimensional embeddings may distort their magnitude and lead to norm collapse. To avoid this issue, we renormalize the hybrid embedding so that its norm matches the probability-weighted sum of the component norms: 𝐡 j ( t ) = 𝐡 ~ j ( t ) ‖ 𝐡 ~ j ( t ) ‖ 2 ​ ( π j ( t − 1 ) ​ ‖ 𝐞 ⁡ ( y j ( t − 1 ) ) ‖ 2 + π j , mask ( t − 1 ) ​ ‖ 𝐞 mask ‖ 2 ) . \mathbf{h}^{(t)}_{j}=\frac{\tilde{\mathbf{h}}^{(t)}_{j}}{\left\|\tilde{\mathbf{h}}^{(t)}_{j}\right\|_{2}}\left(\pi^{(t-1)}_{j}\left\|\mathbf{e}\!\left(y^{(t-1)}_{j}\right)\right\|_{2}+\pi^{(t-1)}_{j,\mathrm{mask}}\left\|\mathbf{e}_{\mathrm{mask}}\right\|_{2}\right). (10) This hybrid embedding serves as a soft intermediate state between decoding steps, explicitly carrying forward the uncertainty of previous predictions.

We regard a block as having converged to a stable state if either of the following conditions holds: (1) the top-1 predictions at all positions remain unchanged for two consecutive decoding steps, or (2) the confidence of every position in the block exceeds a high acceptance threshold τ acc \tau_{\mathrm{acc}} . Once a block converges, we commit all token positions in the block according to the final predictions and move on to the next block.

By interpolating prediction-mask embeddings as intermediate states, the model receives an explicit uncertainty prior before every forward pass, leading to substantially more robust parallel decoding. On OPUT-trained LLaDA-2.0-mini, under the highly aggressive setting of τ dec = 0 \tau_{\mathrm{dec}}=0 , soft parallel decoding improves GSM8K accuracy from 68 % 68\% to 90 % 90\% while achieving a higher speedup.

OPUT as a Prerequisite. Notably, soft parallel decoding must be used together with OPUT-trained models. OPUT trains the model to recover the correct target not only from masked inputs, but also from its own sampled predictions. As a result, the model learns a consistent mapping from both mask embeddings and self-predicted token embeddings toward the correct output, which makes interpolation between them meaningful. In contrast, applying soft parallel decoding to a standard diffusion language model without OPUT leads to catastrophic performance collapse.

## 4 Experiments

### 4.1 Experimental Setup

Implementation Details. We build our method on LLaDA-2.0-mini [ 10 ] , a state-of-the-art open-source diffusion language model. During training, we use OPUT with a fixed mask ratio of 0.75. We perform full-parameter fine-tuning for 2 epochs with a batch size of 8, an initial learning rate of 2 × 10 − 6 2\times 10^{-6} , and a cosine learning rate schedule. Training follows the block-diffusion setting with a block size of 32. To avoid extra memory overhead, the masked noisy sequence and the predicted noisy sequence are optimized in separate iterations within the same epoch, rather than jointly in a single iteration. Under this setup, we train two models: DMax-Math for mathematical reasoning and DMax-Coder for code generation tasks. All training runs are conducted on 8 H200 GPUs. At inference time, we adopt the proposed SPD decoding strategy under semi-autoregressive block diffusion with a block size of 32. The acceptance threshold for determining whether a block has converged to a stable state is set to τ acc = 0.9 \tau_{\mathrm{acc}}=0.9 .

Training Data. We construct all training data through self-distillation. Specifically, we take prompts from public datasets and use LLaDA-2.0-mini to generate responses as training targets. For math, prompts are collected from GSM8K trainset [ 20 ] , PRM12K [ 44 ] , a subset of Numina-Math [ 40 ] , and a subset of OpenThoughts [ 26 ] . For code, prompts are drawn from a subset of OpenCodeInstruct [ 2 ] . Responses are generated with a confidence threshold of 0.95, a block size of 32, and a maximum generation length of 2048 tokens. We discard incomplete generations that do not finish within the length budget. This yields 0.7M math samples and 1.0M code samples. Notably, we do not use any external high-quality responses, all supervision is obtained from the model’s own generations.

Evaluation Details. We evaluate our method on multiple benchmarks. For mathematical reasoning, we use GSM8K [ 20 ] , MATH500 [ 44 ] , Minerva-Algebra [ 29 ] , and ASDIV [ 55 ] , and prompt the model to produce chain-of-thought [ 80 ] reasoning. For code generation, we use the instruction versions of HumanEval [ 13 ] and MBPP [ 5 ] . All evaluations are conducted with the dInFer [ 54 ] framework on 2 H200 GPUs using tensor parallelism. Besides TPF, TPS, and accuracy, we also report AUP Score [ 62 ] to measure parallel decoding performance. The generation length for all benchmarks is 2048.

Baselines. We compare our method against four baselines in terms of both decoding efficiency and generation accuracy: (1) LLaDA-2.0-mini, the base model, evaluated with its default confidence-threshold-based parallel decoding strategy using a threshold of 0.95; (2) Hierarchical Decoding, an advanced inference strategy that improves parallel decoding via a divide-and-conquer procedure [ 61 ] . The low threshold is set as 0.2; (3) dParallel-SFT, for which we use the LLaDA-2.0-mini-CAP model [ 10 ] , where the certainty-forcing loss proposed in dParallel [ 16 ] is incorporated into large-scale supervised fine-tuning to improve decoding parallelism; and (4) Uniform Diffusion Training, which continues training the base model using the conventional UDLM objective. In addition to masked noisy sequences, this baseline also replaces tokens with random vocabulary samples to construct uniformly corrupted noisy sequences, while keeping all other training settings identical to those of DMax. During inference, it updates all tokens within a block at every step until convergence.

### 4.2 Experimental Results

Aggressive Parallelism While Preserving Accuracy. As shown in Table 1 , compared with the original LLaDA-2.0-mini, our method substantially increases decoding parallelism, improving the average TPF from 2.8 to 6.2 while preserving the original accuracy. In contrast, the other baselines provide only limited gains in parallel decoding. This advantage is further reflected in the AUP Score, where DMax consistently outperforms both the original model and all baselines by a large margin. These results demonstrate that our paradigm enables a much stronger parallel decoding capability than conventional MDLMs. Moreover, on two H200 GPUs, our model achieves a practical inference throughput of over 1000 tokens per second.

On-Policy Training as the Cornerstone. Table 1 also compares our method with conventional uniform diffusion training. The latter neither improves decoding speed nor preserves model quality, instead causing a noticeable performance drop. We find that this failure stems from the large mismatch between the randomly sampled noisy sequences used in training and the model’s actual decoding trajectories at inference time. Consequently, the model struggles to revise erroneous predictions while unnecessarily perturbing correct ones, resulting in unstable oscillations within each block. By contrast, our on-policy training samples noisy sequences from the model’s own outputs, effectively bridging this train–inference gap and substantially improving self-revision under parallel decoding.

Superior Efficiency–Performance Trade-off. Figure 4 compares the accuracy–TPF trade-off curves of our method and the original model on GSM8K, MATH500, HumanEval, and MBPP. As TPF increases, the original model suffers a sharp accuracy drop, whereas our method maintains stable performance. For instance, on MATH500, at around 6.5 TPF, our method still retains over 71.6% accuracy, while the original model falls to 15.2%. The gap is even larger on code benchmarks: on MBPP, at a similar TPF, our method achieves 79.2%, whereas the original model drops to only 2.3%. This superior trade-off stems from the self-revision capability of our paradigm, which effectively mitigates error accumulation under aggressive parallel decoding.

Improved Performance at Low Parallelism. By enabling dLLMs to revise their own predictions, our method not only mitigates error accumulation under aggressive parallel decoding, but also improves performance in the low-parallelism regime. Through iterative re-evaluation of earlier predictions, the model can recover from reasoning errors that would otherwise remain on the original decoding path. As shown in Table 2 , our method consistently improves accuracy by 0.8%–3.0% across multiple benchmarks at low parallelism. Importantly, these gains are obtained using only the model’s own generated responses as training data, without introducing any external supervision.

## 5 Ablation Study

Ablation Study on Training and Inference Strategies. Table 3 presents a comprehensive ablation study of both our training and inference designs. We compare different combinations of training and decoding strategies on GSM8K under three decoding thresholds, τ dec ∈ { 0.95 , 0.5 , 0.0 } \tau_{\mathrm{dec}}\in\{0.95,0.5,0.0\} . On-policy rollout is the core of our training method. Even with OPUT alone, the model acquires the ability to revise its own errors, yielding substantial accuracy gains over the original model at τ dec = 0.5 \tau_{\mathrm{dec}}=0.5 and 0.0 0.0 . Our proposed SPD further improves robustness when many erroneous predictions emerge simultaneously, allowing the model to remain stable under highly parallel decoding and to preserve strong performance even in the extreme case of τ dec = 0.0 \tau_{\mathrm{dec}}=0.0 . The key ingredient of SPD is to use soft embeddings, rather than discrete tokens, as intermediate decoding states. Maintaining the non-masked region as a contiguous prefix further improves performance. Another important result is that OPUT is a prerequisite for SPD. As shown in Table 3 , directly applying SPD to the original model causes generation to collapse. This is because OPUT trains the model to recover clean tokens from both mask tokens and predicted tokens, making interpolation between their embeddings a meaningful and effective input for denoising.

Ablation Study on Convergence Criteria. We further study in Table 4 how different block-level convergence criteria affect the efficiency–performance trade-off. We consider two criteria: (1) consistency , where decoding is considered converged if the model produces the same top-1 prediction for the block in two consecutive steps; and (2) confidence , where decoding is considered converged if the confidence of every token in the block exceeds 0.9. As shown in Table 4 , consistency serves as the primary convergence signal, with most blocks terminating once this condition is met. Adding the confidence criterion can further improve TPF by allowing decoding to stop before two consecutive identical predictions are observed, thereby saving the final forward pass. Importantly, neither criterion affects the accuracy.

## 6 Related Work

Diffusion Language Models. Diffusion models [ 31 , 71 ] have become dominant in visual generation [ 64 , 60 , 65 , 99 ] , and recent work has explored their application to text generation. Among existing paradigms, masked diffusion language models (MDLMs) [ 70 , 4 , 66 , 105 , 51 ] have emerged as a promising alternative to AR-LLMs by modeling language in discrete space through masked token prediction. Building on this formulation, LLaDA [ 58 ] and Dream [ 92 ] scale MDLMs to the billion-parameter regime with large-scale pretraining, demonstrating their practical potential. LLaDA-2.0 [ 10 ] and LLaDA-MoE [ 110 ] further show that MDLMs can be effectively scaled with mixture-of-experts architectures. Beyond these developments, dLLMs are also attracting increasing attention in reasoning [ 109 , 59 , 86 , 63 , 74 , 57 , 103 ] , multimodal tasks [ 94 , 96 , 90 , 91 , 48 , 82 , 97 , 18 ] , code generation [ 87 , 24 , 21 ] , long-context modeling [ 47 , 28 , 106 ] , and agent [ 104 , 102 ] .

Accelerating Diffusion Language Models. dLLMs are viewed as promising due to their potential for low-cost inference, yet their efficiency remains largely underexplored. Existing efforts improve efficiency from several perspectives. Some methods reduce the cost of each decoding step through techniques including KV caching [ 53 , 49 , 84 , 35 , 45 ] , token dropping [ 15 , 36 , 72 , 85 ] , and sparse attention [ 79 , 19 ] . Others design more effective decoding strategies [ 37 , 81 , 41 , 27 , 8 , 32 , 34 , 89 , 50 , 12 , 77 , 61 , 22 ] to improve generation efficiency. A separate line of work [ 73 , 62 , 101 , 7 , 14 , 33 ] learns better decoding trajectories so that fewer decoding steps are required. dParallel [ 16 ] employs certainty-forcing distillation to accelerate confidence convergence and enable higher parallel decoding. Other methods [ 83 , 17 , 78 , 3 , 52 , 75 , 46 , 23 ] interpolate between diffusion and autoregressive language models to better balance speed and accuracy. [ 9 , 76 , 100 ] implement uniform training, which trains the model to recover clean tokens from random noisy tokens, thereby enabling token correction during generation. SM [ 30 ] and EvoToken [ 107 ] introduce soft embeddings into the decoding process, but neither method translates this design into improved decoding efficiency. Further efforts [ 88 ] leverage compression techniques to construct lightweight dLLMs.

## 7 Conclusion