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

Cleaned-text code points [0, 15479). Bibliography omitted here; source.md includes it.
# DMax: Aggressive Parallel Decoding for dLLMs

###### Abstract

We present DMax, a new paradigm for efficient diffusion language models (dLLMs). It mitigates error accumulation in parallel decoding, enabling aggressive decoding parallelism while preserving generation quality. Unlike conventional masked dLLMs that decode through a binary mask-to-token transition, DMax reformulates decoding as a progressive self-refinement from mask embeddings to token embeddings. At the core of our approach is On-Policy Uniform Training, a novel training strategy that efficiently unifies masked and uniform dLLMs, equipping the model to recover clean tokens from both masked inputs and its own erroneous predictions. Building on this foundation, we further propose Soft Parallel Decoding. We represent each intermediate decoding state as an interpolation between the predicted token embedding and the mask embedding, enabling iterative self-revising in embedding space. Extensive experiments across a variety of benchmarks demonstrate the effectiveness of DMax. Compared with the original LLaDA-2.0-mini, our method improves TPF on GSM8K from 2.04 to 5.47 while preserving accuracy. On MBPP, it increases TPF from 2.71 to 5.86 while maintaining comparable performance. On two H200 GPUs, our model achieves an average of 1,338 TPS at batch size 1. Code is available at: https://github.com/czg1225/DMax

## 1 Introduction

Recently, Diffusion Language Models (dLLMs) [ 93 , 95 , 98 , 42 , 56 , 108 ] have emerged as a compelling alternative to the long-standing dominance of Autoregressive Language Models (AR-LLM) [ 1 , 6 , 25 ] in text generation. The primary allure of dLLMs lies in their capacity for parallel decoding, which holds great promise for improving inference efficiency

Despite this promise, the practical decoding parallelism of existing dLLMs [ 58 , 109 , 92 , 17 , 10 , 90 ] remains limited, as their performance drops sharply under aggressive parallel decoding. Some prior work has attempted to improve this trade-off through improved decoding [ 37 , 81 , 41 , 27 , 8 , 32 , 34 , 89 ] or distillation strategies [ 16 , 62 , 101 , 38 ] . Nevertheless, these methods do not address the fundamental bottleneck underlying parallel decoding in current dLLM paradigm: error accumulation.

In current mask-based dLLMs, decoding is a binary, one-way mask-to-token process. Once a masked position is decoded into a token, that token is fixed and propagated as context to subsequent decoding steps, with no opportunity for revision. 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 [ 39 , 11 , 43 ] , dLLMs lack a mechanism to recover from incorrect predictions, which fundamentally restricts their performance under highly parallel decoding. Addressing this challenge requires a new dLLM paradigm with an intrinsic capability to revise its own predictions during decoding.

Building on this insight, we propose DMax, a novel paradigm that reformulates the binary mask-to-token decoding process into a self-revising transformation in the embedding space. Central to our approach is On-Policy Uniform Training (OPUT), a training recipe that efficiently extends a pretrained masked diffusion language model into a self-corrective uniform diffusion language model while preserving its original mask denoising capability. Unlike conventional uniform diffusion training that constructs noisy sequences by randomly sampling tokens from the vocabulary, OPUT samples noisy inputs on-policy from the model’s own predictive distribution. This substantially bridges the train-inference gap and enables the model to effectively learn to correct its own potential prediction errors. Building upon OPUT, we further present Soft Parallel Decoding (SPD) for inference. Instead of treating decoded tokens as discrete and irrevocable commitments, SPD represents each intermediate decoding state as a hybrid soft embedding, formed by interpolating between the predicted token embedding and the mask embedding according to the model’s prediction confidence. This simple design provides the model with confidence priors from previous steps, enabling more robust self-correction.

Using LLaDA-2.0-mini [ 10 ] , a state-of-the-art open-source dLLM, as the base model, we validate the effectiveness of our method across multiple widely used benchmarks. On the mathematical reasoning benchmark GSM8K [ 20 ] , our method increases tokens per forward (TPF) from 2.04 to 5.48 with only minimal accuracy degradation relative to the original model. On the code generation benchmark MBPP [ 5 ] , it improves TPF from 2.71 to 5.86 while maintaining comparable performance.

In summary, we propose DMax, a novel paradigm that enables highly parallel decoding for dLLMs while preserving strong performance. Our central idea is to mitigate the error accumulation issue caused by the conventional one-way mask-to-token decoding. To realize this, we introduce two key designs: on-policy uniform training and soft parallel decoding. Extensive experiments demonstrate the effectiveness and superiority of our approach. This work establishes a new strong baseline for future research on parallel decoding in dLLMs.

## 2 Preliminaries

We begin by briefly reviewing the diffusion language modeling paradigms, and then highlight the central challenge for highly parallel decoding and introduce our key motivation.

Masked Diffusion Language Models (MDLMs). MDLMs [ 70 , 4 , 66 , 105 , 51 ] formulate text generation as a discrete denoising process over token sequences, where clean tokens are progressively replaced by a special [MASK] symbol during corruption. Let x 0 = ( x 0 1 , … , x 0 L ) ∈ 𝒱 L x_{0}=(x_{0}^{1},\dots,x_{0}^{L})\in\mathcal{V}^{L} denote a clean sequence of length L L , where 𝒱 \mathcal{V} is the vocabulary. Given a corrupted sequence x t x_{t} at noise level t ∈ [ 0 , 1 ] t\in[0,1] , the denoising model is trained to recover the original tokens only at masked positions. The standard MDLM objective is ℒ MDLM  ( θ ) = − 𝔼 x 0 , t , x t  [ 1 t  ∑ i = 1 L 𝟏  ( x t i = [MASK] )  log ⁡ p θ  ( x 0 i ∣ x t ) ] . \mathcal{L}_{\mathrm{MDLM}}(\theta)=-\mathbb{E}_{x_{0},\,t,\,x_{t}}\left[\frac{1}{t}\sum_{i=1}^{L}\mathbf{1}(x_{t}^{i}=\texttt{[MASK]})\log p_{\theta}(x_{0}^{i}\mid x_{t})\right]. (1) At inference time, MDLMs start from a fully masked sequence and iteratively decode masked positions in parallel, with an optional remasking step to enable further refinement.

Uniform Diffusion Language Models (UDLMs). UDLMs [ 68 , 67 , 69 ] generalize the corruption process by replacing tokens with uniformly sampled vocabulary tokens rather than a dedicated [MASK] symbol. As a result, the model is trained to recover clean tokens from arbitrary noisy token inputs, instead of only from masked positions. A standard UDLM training objective is ℒ UDLM  ( θ ) = − 𝔼 x 0 , t , x t  [ ∑ i = 1 L log ⁡ p θ  ( x 0 i ∣ x t ) ] . \mathcal{L}_{\mathrm{UDLM}}(\theta)=-\mathbb{E}_{x_{0},\,t,\,x_{t}}\left[\sum_{i=1}^{L}\log p_{\theta}(x_{0}^{i}\mid x_{t})\right]. (2) During inference, UDLMs typically start from a fully noisy sequence sampled uniformly from the vocabulary and iteratively update all positions.

Error accumulation in MDLMs. Existing dLLMs based on the MDLM paradigm degrade sharply under highly parallel decoding, which limits the practical speedup. The main reason behind it is error accumulation. MDLM decoding follows a binary mask-to-token process: each position is either a mask token or a committed token. Once a masked position is decoded, its prediction is treated as fixed context for subsequent steps. Early mistakes cannot be revised, and instead propagate through later denoising steps as erroneous context.

UDLMs as a Promising Solution. In contrast, UDLMs are trained to denoise from arbitrary vocabulary tokens rather than only from [MASK] , so all positions can be re-evaluated at every decoding step. This token-to-token denoising mechanism naturally enables self-correction and improves robustness to prediction errors. However, UDLM decoding typically starts from a fully random sequence, which makes denoising harder and leads to very unstable generation.

Unify the Strengths of MDLMs and UDLMs. Motivated by this trade-off, we propose to unify the strengths of both paradigms. Specifically, we retain a fully masked sequence as the initialization of UDLM decoding to preserve stability, while continuing to re-predict all tokens that have been decoded from [MASK] at every subsequent step. This design combines the stable initialization with the self-revising capability, enabling a more robust parallel decoding process.

## 3 Methodology

### 3.1 On-Policy Uniform Training

A practical way to achieve this goal is to extend a pretrained MDLM into a UDLM. Accordingly, our first objective is to endow a pretrained MDLM with the self-revision capability of UDLMs while preserving its original mask denoising ability.

Extending MDLM toward UDLM is Nontrivial. This is nontrivial because the training objective of UDLMs differs substantially from that of MDLMs and is considerably harder to optimize. In the standard UDLM training paradigm, a clean sequence is first corrupted by randomly selecting a subset of positions and replacing the selected tokens with tokens sampled uniformly from the vocabulary. The resulting noisy sequence is then used as model input, and the model is trained to recover the original clean sequence.

However, this training strategy is often unstable in practice and tends to yield suboptimal performance. A key reason is that uniformly sampled tokens lie far outside the natural language manifold, producing highly unnatural corrupted inputs. As a result, the model must spend substantial capacity merely learning to map these corrupted sequences back toward plausible language, rather than directly acquiring effective language modeling and self-correction behaviors. More importantly, this corruption process introduces a severe train–inference mismatch. Unlike conventional UDLMs, our paradigm first predicts tokens from masked positions in parallel and then iteratively refines its own predictions. Consequently, the noisy sequences encountered at inference time are sampled from the model’s own output distribution rather than from a uniform vocabulary distribution. This mismatch hinders self-correction and leads to ineffective training.

On-Policy Uniform Training. To address these issues, we propose On-Policy Uniform Training (OPUT), a simple yet effective method for equipping MDLMs with self-corrective denoising capability. The core idea is to construct training inputs using noisy sequences sampled on-policy from the model’s own predictive distribution, rather than from a uniform vocabulary distribution, thereby bridging the train–inference gap. The overview of the training procedure is shown in Figure 2 .

Training Procedure. Let M θ M_{\theta} denote a pretrained diffusion language model built on the MDLM paradigm, parameterized by θ \theta . We further adapt M θ M_{\theta} on a training dataset 𝒟 \mathcal{D} of clean sequences x 0 = ( x 0 1 , … , x 0 L ) x_{0}=(x_{0}^{1},\dots,x_{0}^{L}) . At each training iteration, we first sample a corruption level t ∼ Uniform ⁡ ( t l , t h ) t\sim\mathrm{Uniform}(t_{l},t_{h}) , where t l t_{l} and t h t_{h} denote the lower and upper bounds of the noise level, respectively. Given a clean sequence x 0 ∼ 𝒟 x_{0}\sim\mathcal{D} , we construct a masked noisy sequence x t ( m ) x_{t}^{(m)} by independently replacing each token with [MASK] with probability t t .

We feed x t ( m ) x_{t}^{(m)} into M θ M_{\theta} and predict all masked positions in parallel. By sampling from the model’s predictive distribution at masked positions, we obtain a predicted noisy sequence x t ( p ) x_{t}^{(p)} , defined as x t ( p ) , i = { x t ( m ) , i , if  x t ( m ) , i ≠ [MASK] , x ^ i , x ^ i ∼ p θ ( ⋅ ∣ x t ( m ) ) , if  x t ( m ) , i = [MASK] . x_{t}^{(p),i}=\begin{cases}x_{t}^{(m),i},&\text{if }x_{t}^{(m),i}\neq\texttt{[MASK]},\\[4.0pt] \hat{x}^{i},\quad\hat{x}^{i}\sim p_{\theta}(\cdot\mid x_{t}^{(m)}),&\text{if }x_{t}^{(m),i}=\texttt{[MASK]}.\end{cases} (3) Importantly, x t ( p ) x_{t}^{(p)} is sampled using the current model parameters at each iteration, making this a strictly on-policy rollout process.

Next, we perform two forward passes, using the masked noisy sequence x t ( m ) x_{t}^{(m)} and the predicted noisy sequence x t ( p ) x_{t}^{(p)} as inputs, respectively: p θ ( m ) ( ⋅ ∣ x t ( m ) ) = M θ ( x t ( m ) ) , p θ ( p ) ( ⋅ ∣ x t ( p ) ) = M θ ( x t ( p ) ) . p_{\theta}^{(m)}(\cdot\mid x_{t}^{(m)})=M_{\theta}(x_{t}^{(m)}),\qquad p_{\theta}^{(p)}(\cdot\mid x_{t}^{(p)})=M_{\theta}(x_{t}^{(p)}). (4) We then supervise both outputs against the original clean sequence x 0 x_{0} using cross-entropy loss over all token positions, regardless of whether a position is masked: ℒ mask = − ∑ i = 1 L log p θ ( m ) ( x 0 i ∣ x t ( m ) ) , ℒ pred = − ∑ i = 1 L log p θ ( p ) ( x 0 i ∣ x t ( p ) ) . \mathcal{L}_{\mathrm{mask}}=-\sum_{i=1}^{L}\log p_{\theta}^{(m)}(x_{0}^{i}\mid x_{t}^{(m)}),\qquad\mathcal{L}_{\mathrm{pred}}=-\sum_{i=1}^{L}\log p_{\theta}^{(p)}(x_{0}^{i}\mid x_{t}^{(p)}). (5) The final training objective is ℒ on  -  policy = ℒ mask + ℒ pred , \mathcal{L}_{\mathrm{on\mbox{-}policy}}=\mathcal{L}_{\mathrm{mask}}+\mathcal{L}_{\mathrm{pred}}, (6) By reducing the train–inference mismatch, the proposed OPUT strategy enables a pretrained MDLM to efficiently learn self-correction through limited post-training, while retaining its original mask denoising ability. As a result, the model can correct self-generated errors and effectively mitigate error accumulation under highly parallel decoding. On LLaDA-2.0-mini, our method improves GSM8K accuracy from 78 % 78\% to 90 % 90\% under confidence-threshold decoding with a threshold of 0.5 0.5 , while also delivering faster decoding.

### 3.2 Soft Parallel Decoding

Although OPUT substantially mitigates error accumulation, it still struggles when many erroneous predictions arise simultaneously within a block. When many positions are decoded in parallel, correlated errors can appear at once, making them difficult to fully correct through iterative refinement. For example, for OPUT-trained LLaDA-2.0-mini, if we decode all masked positions in a block at once using a confidence threshold of 0 0 , and then iteratively refine them, the accuracy on GSM8K drops to only 68 % 68\% .

Soft Parallel Decoding. To further enhance self-revising in iterative refinement, we propose soft parallel decoding. The central idea is to preserve predictive uncertainty from earlier iterations and explicitly propagate it to later refinement steps. Concretely, instead of treating intermediate decoding states as discrete tokens, we represent each decoded token as a soft embedding interpolated between the predicted token embedding and the mask embedding. Because the mask embedding naturally encodes maximal uncertainty, this interpolation serves as an explicit carrier of uncertainty across iterations. This enables the model to better distinguish confident predictions from unreliable ones, allowing it to focus on refining low-confidence tokens while avoiding interference from noisy signals.