DMax: Decodifica Parallela Aggressiva per dLLM

DMax: Aggressive Parallel Decoding for dLLMs

April 9, 2026

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

cs.AI

Abstract

Presentiamo DMax, un nuovo paradigma per modelli linguistici di diffusione (dLLM) efficienti. Questo approccio mitiga l'accumulo di errori nel decoding parallelo, consentendo un parallelismo di decoding aggressivo preservando al contempo la qualità della generazione. A differenza dei dLLM mascherati convenzionali che decodificano attraverso una transizione binaria da maschera a token, DMax riformula il decoding come un auto-affinamento progressivo da embedding di maschera a embedding di token. Il cuore del nostro approccio è l'Addestramento Uniforme On-Policy, una nuova strategia di training che unisce efficientemente dLLM mascherati e uniformi, dotando il modello della capacità di recuperare token puliti sia da input mascherati che dalle sue stesse previsioni errate. Basandoci su queste fondamenta, proponiamo inoltre il Soft Parallel Decoding. Rappresentiamo ogni stato di decoding intermedio come un'interpolazione tra l'embedding del token previsto e l'embedding della maschera, abilitando un auto-revisione iterativa nello spazio degli embedding. Esperimenti estesi su una varietà di benchmark dimostrano l'efficacia di DMax. Rispetto al LLaMA-2.0-mini originale, il nostro metodo migliora il TPF su GSM8K da 2.04 a 5.47 preservando l'accuratezza. Su MBPP, aumenta il TPF da 2.71 a 5.86 mantenendo prestazioni comparabili. Su due GPU H200, il nostro modello raggiunge una media di 1.338 TPS con batch size 1. Il codice è disponibile all'indirizzo: https://github.com/czg1225/DMax

English

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