PaperlayerView exact revision on arXiv

Source-linked AI summary

DMax: Aggressive Parallel Decoding for dLLMs

Zigeng Chen, Gongfan Fang, Xinyin Ma, Ruonan Yu, Xinchao Wang

arXiv:2604.08302v3cs.LGcs.AI

TL;DR

DMax addresses error accumulation that limits highly parallel decoding in diffusion language models by combining On-Policy Uniform Training with Soft Parallel Decoding for self-correction. Compared with LLaDA-2.0-mini, it raises decoding parallelism while preserving accuracy, including TPF gains on GSM8K and MBPP and over 1,000 tokens per second on two H200 GPUs.

  • Problem

    Existing masked diffusion language models suffer error accumulation under highly parallel decoding because early committed mistakes cannot be revised and propagate as erroneous context.

  • Method

    DMax combines On-Policy Uniform Training, which trains the model on its own predictive errors, with Soft Parallel Decoding, which supports iterative self-revision in embedding space.

  • Results

    DMax substantially increases decoding parallelism while preserving accuracy, raising TPF from 2.04 to 5.48 on GSM8K and from 2.71 to 5.86 on MBPP relative to LLaDA-2.0-mini.

  • Takeaways & Limitations

    DMax provides a stronger parallel decoding capability for dLLMs and also improves accuracy by 0.8%–3.0% across multiple benchmarks at low parallelism.

  • Takeaways & Limitations

    When many positions are decoded in parallel, correlated errors can remain difficult to correct; one OPUT-trained LLaDA-2.0-mini setting reached only 68% GSM8K accuracy.

Abstract

from arXiv · show

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

DMax addresses error accumulation that limits aggressive parallel decoding in diffusion language models. It combines on-policy self-correction training with embedding-space refinement to increase decoding parallelism while preserving performance.

  • Motivation: Existing dLLMs promise efficient parallel decoding, but their practical parallelism drops sharply under aggressive decoding because errors accumulate.In binary mask-to-token decoding, committed mistakes propagate as erroneous context and can trigger cascading errors.
  • DMax: DMax reformulates binary mask-to-token decoding as a self-revising transformation in embedding space.The model can revise predictions during decoding rather than treating decoded tokens as irrevocable commitments.
  • DMax: On-Policy Uniform Training samples noisy inputs from the model’s predictive distribution, teaching it to correct potential prediction errors while retaining mask denoising.This training recipe bridges the train–inference gap more directly than uniform vocabulary corruption.
  • DMax: Soft Parallel Decoding represents intermediate states as hybrid embeddings interpolated between predicted token and mask embeddings according to confidence.The design provides confidence priors from previous steps for more robust self-correction.
  • Results: TPF rises from 2.04 to 5.48 on GSM8K and from 2.71 to 5.86 on MBPP while preserving comparable performance.The GSM8K improvement has only minimal accuracy degradation relative to the original model, while MBPP performance remains comparable.
  • Conclusion: DMax enables highly parallel decoding while preserving strong performance and establishes a baseline for future dLLM parallel-decoding research.The paper attributes this outcome to mitigating error accumulation through on-policy uniform training and soft parallel decoding.

2 Preliminaries

This section reviews masked and uniform diffusion language models and explains why masked models suffer error accumulation during parallel decoding. It motivates combining masked initialization with repeated token re-evaluation.

  • Masked Diffusion Language Models: Masked diffusion language models formulate generation as discrete denoising, progressively replacing clean tokens with [MASK] and recovering tokens at masked positions.Their inference process begins from a fully masked sequence and decodes masked positions in parallel.
  • Uniform Diffusion Language Models: Uniform diffusion language models replace tokens with uniformly sampled vocabulary tokens and train recovery from arbitrary noisy token inputs.At inference, they typically begin from a fully noisy sequence and update all positions iteratively.
  • Error Accumulation: Masked-model decoding degrades under aggressive parallelism because committed early mistakes become fixed context and propagate through later denoising steps.This error accumulation limits practical speedup and prevents early errors from being revised.
  • Trade-off: Uniform models can re-evaluate all positions and thereby support self-correction, but fully random initialization makes denoising harder and generation unstable.This creates a trade-off between masked models’ stable initialization and uniform models’ token-to-token refinement.
  • Unifying the Paradigms: The proposed unification retains fully masked initialization while repeatedly re-predicting decoded tokens at subsequent steps.This combines stable initialization with self-revising capability for more robust parallel decoding.

3 Methodology

DMax extends pretrained masked diffusion models with self-correction through On-Policy Uniform Training and revisable embedding-space decoding. Soft Parallel Decoding preserves uncertainty during block-wise refinement, enabling aggressive parallelism while maintaining performance.

  • On-Policy Uniform Training: OPUT trains a pretrained MDLM on model-generated noisy sequences to learn self-corrective denoising while retaining mask denoising.It bridges the train–inference gap by sampling noisy inputs from the model’s predictive distribution rather than uniformly from the vocabulary.
  • On-Policy Uniform Training: OPUT uses masked and predicted noisy sequences as inputs, supervises both against clean sequences, and combines their losses.The predicted sequence is generated by sampling model predictions at masked positions, and both forward-pass outputs receive cross-entropy supervision over all token positions.
  • Soft Parallel Decoding: OPUT alone struggles when many correlated errors arise simultaneously within a parallel block, with GSM8K accuracy dropping to 68% under τdec = 0.Soft interpolation is introduced to make iterative refinement more robust in this setting.
  • Soft Parallel Decoding: Soft Parallel Decoding represents decoded positions with embeddings interpolated between predicted tokens and the mask embedding, preserving uncertainty for iterative self-revision.Mask positions retain mask embeddings, while token positions use hybrid embeddings derived from prior top-1 predictions and their probabilities.
  • Soft Parallel Decoding: The block-wise decoder promotes the longest contiguous prefix whose confidence exceeds τdec, then refines token positions while keeping later positions masked.A block converges when predictions remain unchanged for two consecutive steps or every position exceeds τacc, after which its tokens are committed.
  • Soft Parallel Decoding: 68% to 90%: soft parallel decoding raises GSM8K accuracy under τdec = 0 while achieving a higher speedup on OPUT-trained LLaDA-2.0-mini.OPUT is required because it teaches the model to map both mask and self-predicted token embeddings toward the correct output.

4 Experiments

DMax is evaluated against the base model and several decoding or training baselines across mathematical reasoning and code-generation benchmarks. It increases parallel decoding while preserving accuracy and improves the efficiency–performance trade-off through on-policy training and self-revision.

  • Experimental Results: Conventional uniform diffusion training neither improves decoding speed nor preserves model quality, whereas on-policy training bridges the train–inference gap and improves self-revision.The baseline uses random vocabulary samples to construct uniformly corrupted noisy sequences.
  • Experimental Results: At around 6.5 TPF on MATH500, DMax retains over 71.6% accuracy versus 15.2% for the original model.Figure 4 compares accuracy–TPF curves on GSM8K, MATH500, HumanEval, and MBPP.
  • Experimental Results: At similar TPF on MBPP, DMax achieves 79.2% accuracy while the original model reaches 2.3%.The reported trade-off advantage is attributed to self-revision mitigating error accumulation under aggressive parallel decoding.

5 Ablation Study

The ablations show that on-policy training and soft parallel decoding jointly support robust aggressive decoding. Convergence criteria can improve efficiency without changing accuracy.

  • Training and Inference Strategies: On-policy rollout is the core training component, while soft parallel decoding further improves robustness when many erroneous predictions emerge simultaneously.Applying soft parallel decoding directly to the original model causes generation to collapse, making OPUT a prerequisite.
  • Training and Inference Strategies: Soft embeddings, rather than discrete tokens, make intermediate decoding states effective for denoising and self-revision.Maintaining the non-masked region as a contiguous prefix further improves performance.
  • Convergence Criteria: Consistency is the primary convergence signal, while adding a confidence criterion can improve TPF by saving the final forward pass.The confidence criterion stops decoding when every token in a block exceeds 0.9 confidence.
  • Convergence Criteria: Neither consistency nor confidence-based convergence affects accuracy.The ablation evaluates block-level convergence at a decoding threshold of 0.5.

6 Related Work

Related work positions diffusion language models as an alternative to autoregressive generation and surveys their development across scaling and applications. Prior acceleration efforts target decoding-step cost or decoding strategy, while aggressive parallelism remains challenging.

  • Diffusion Language Models: Masked diffusion language models model language in discrete space through masked-token prediction and have emerged as an alternative to autoregressive language models.LLaDA and Dream scaled masked diffusion language models to the billion-parameter regime.
  • Diffusion Language Models: LLaDA-2.0 and LLaDA-MoE show that masked diffusion language models can be scaled with mixture-of-experts architectures.Diffusion language models are also being explored for reasoning, multimodal tasks, and code generation.
  • Accelerating Diffusion Language Models: Existing efficiency methods reduce per-step cost through KV caching, token dropping, and sparse attention, or use more effective decoding strategies.The related-work discussion frames dLLM efficiency as promising but still underexplored.

7 Conclusion

DMax is presented as an efficient dLLM paradigm that mitigates error accumulation during parallel decoding. It combines On-Policy Uniform Training and Soft Parallel Decoding, with experiments supporting its effectiveness across diverse benchmarks.

  • DMax mitigates error accumulation to enable aggressive parallel decoding while preserving the original model's accuracy.
  • The approach introduces On-Policy Uniform Training and Soft Parallel Decoding as its two key components.
  • Experiments across diverse benchmarks demonstrate DMax's effectiveness and establish a strong new baseline for parallel decoding in dLLMs.
Loading 2604.08302v3…