DMax: dLLM을 위한 공격적 병렬 디코딩

DMax: Aggressive Parallel Decoding for dLLMs

April 9, 2026

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

cs.AI

초록

우리는 효율적인 확산 언어 모델(dLLM)을 위한 새로운 패러다임인 DMax를 제안한다. 이는 병렬 디코딩에서의 오류 누적을 완화하여 생성 품질을 유지하면서도 공격적인 디코딩 병렬화를 가능하게 한다. 이진 마스크-토큰 전환을 통해 디코딩하는 기존의 마스크 dLLM과 달리, DMax는 디코딩을 마스크 임베딩에서 토큰 임베딩으로의 점진적 자기 정제 과정으로 재구성한다. 우리 접근법의 핵심은 On-Policy Uniform Training이라는 새로운 훈련 전략으로, 마스크 dLLM과 균일 dLLM을 효율적으로 통일하여 모델이 마스크된 입력과 자체 오류 예측 모두에서 깨끗한 토큰을 복원할 수 있도록 한다. 이를 기반으로 우리는 추가로 Soft Parallel Decoding을 제안한다. 각 중간 디코딩 상태를 예측된 토큰 임베딩과 마스크 임베딩 사이의 보간으로 표현하여 임베딩 공간에서 반복적인 자기 수정이 가능하게 한다. 다양한 벤치마크에서의 광범위한 실험을 통해 DMax의 효과성을 입증하였다. 기존 LLaMA-2.0-mini 대비 우리 방법은 GSM8K에서 정확도를 유지하면서 TPF를 2.04에서 5.47로 향상시켰다. MBPP에서는 비슷한 성능을 유지하면서 TPF를 2.71에서 5.86으로 증가시켰다. H200 GPU 2개에서 우리 모델은 배치 크기 1 기준 평균 1,338 TPS를 달성했다. 코드는 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