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

DMax: Aggressive Parallel Decoding for dLLMs

본 논문에서는 효율적인 디퓨전 언어 모델(dLLM)을 위한 새로운 패러다임인 DMax를 제시합니다. DMax는 병렬 디코딩 과정에서의 오류 누적을 완화하여, 생성 품질을 유지하면서 공격적인 병렬 디코딩을 가능하게 합니다. 기존의 마스크 기반 dLLM과는 달리, DMax는 디코딩을 이진 마스크-토큰 변환 과정이 아닌, 마스크 임베딩에서 토큰 임베딩으로의 점진적인 자체 개선 과정으로 재구성합니다. 저희 접근 방식의 핵심은 On-Policy Uniform Training이라는 새로운 학습 전략으로, 이를 통해 마스크 기반 dLLM과 균일 기반 dLLM을 효율적으로 통합하고, 모델이 마스크 입력과 자체의 오류 예측으로부터 모두 정확한 토큰을 복구할 수 있도록 합니다. 이러한 기반 위에, 저희는 Soft Parallel Decoding을 추가적으로 제안합니다. 각 중간 디코딩 상태를 예측된 토큰 임베딩과 마스크 임베딩 사이의 보간값으로 표현하여, 임베딩 공간에서 반복적인 자체 수정이 가능하도록 합니다. 다양한 벤치마크를 사용한 광범위한 실험 결과는 DMax의 효과를 입증합니다. 원본 LLaDA-2.0-mini와 비교하여, 저희 방법은 GSM8K에서 TPF를 2.04에서 5.47로 향상시키면서 정확도를 유지합니다. MBPP에서는 TPF를 2.71에서 5.86으로 증가시키면서도 유사한 성능을 유지합니다. 두 개의 H200 GPU에서 저희 모델은 배치 크기 1에서 평균 1,338 TPS를 달성합니다. 코드: https://github.com/czg1225/DMax

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