DMax: Decodificación Paralela Agresiva para dLLMs

DMax: Aggressive Parallel Decoding for dLLMs

April 9, 2026

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

cs.AI

Resumen

Presentamos DMax, un nuevo paradigma para modelos de lenguaje de difusión (dLLMs) eficientes. Mitiga la acumulación de errores en la decodificación paralela, permitiendo un paralelismo de decodificación agresivo mientras preserva la calidad de la generación. A diferencia de los dLLMs enmascarados convencionales que decodifican mediante una transición binaria de máscara a token, DMax reformula la decodificación como un auto-refinamiento progresivo desde *embeddings* de máscara hasta *embeddings* de token. En el núcleo de nuestro enfoque se encuentra el Entrenamiento Uniforme *On-Policy*, una novedosa estrategia de entrenamiento que unifica eficientemente los dLLMs enmascarados y uniformes, equipando al modelo para recuperar tokens limpios tanto a partir de entradas enmascaradas como de sus propias predicciones erróneas. Sobre esta base, proponemos además la Decodificación Paralela Suave. Representamos cada estado intermedio de decodificación como una interpolación entre el *embedding* del token predicho y el *embedding* de la máscara, permitiendo una auto-revisión iterativa en el espacio de *embeddings*. Experimentos exhaustivos en diversos benchmarks demuestran la efectividad de DMax. En comparación con el LLaMA-2.0-mini original, nuestro método mejora el TPF en GSM8K de 2.04 a 5.47 mientras preserva la precisión. En MBPP, aumenta el TPF de 2.71 a 5.86 manteniendo un rendimiento comparable. En dos GPUs H200, nuestro modelo alcanza un promedio de 1,338 TPS con un tamaño de lote de 1. El código está disponible en: 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