d3LLM: 의사 경로 증류를 이용한 초고속 확산형 거대 언어 모델

d3LLM: Ultra-Fast Diffusion LLM using Pseudo-Trajectory Distillation

확산형 거대 언어 모델(dLLM)은 병렬 디코딩 및 임의 순서 생성과 같은, 자기 회귀(AR) 모델보다 뛰어난 기능을 제공합니다. 그러나 이러한 이점을 실제 환경에서 구현하는 것은 쉽지 않으며, dLLM은 본질적으로 정확성과 병렬성 간의 균형 문제를 안고 있습니다. 기존 연구들은 효율성 또는 성능 중 한쪽에 집중하는 경향이 있었지만, 본 연구에서는 이러한 한계를 극복하기 위해 d3LLM(의사 경로 증류를 이용한 확산형 거대 언어 모델)을 제안합니다. d3LLM은 정확성과 병렬성 사이의 균형을 맞추며, (i) 학습 단계에서 모델에게 초기 단계에서 어떤 토큰을 안전하게 디코딩할 수 있는지 가르치는 의사 경로 증류를 도입하여 병렬성을 향상시키고, (ii) 추론 단계에서 엔트로피 기반 다중 블록 디코딩과 KV-캐시 갱신 메커니즘을 사용하여 높은 병렬성을 유지하면서 정확도를 확보합니다. 또한, dLLM의 성능을 보다 정확하게 평가하기 위해 정확성과 병렬성을 동시에 측정하는 새로운 지표인 AUP(Accuracy Under Parallelism)를 제안합니다. 실험 결과, 제안하는 d3LLM은 기존 LLaDA/Dream 모델보다 최대 10배, 자기 회귀 모델보다 5배 빠른 속도를 보였으며, 정확도 저하 없이 이러한 성능 향상을 달성했습니다. 저희의 코드는 https://github.com/hao-ai-lab/d3LLM 에서 확인할 수 있습니다.

Diffusion large language models (dLLMs) offer capabilities beyond those of autoregressive (AR) LLMs, such as parallel decoding and random-order generation. However, realizing these benefits in practice is non-trivial, as dLLMs inherently face an accuracy-parallelism trade-off. Despite increasing interest, existing methods typically focus on only one-side of the coin, targeting either efficiency or performance. To address this limitation, we propose d3LLM (Pseudo-Distilled Diffusion Large Language Model), striking a balance between accuracy and parallelism: (i) during training, we introduce pseudo-trajectory distillation to teach the model which tokens can be decoded confidently at early steps, thereby improving parallelism; (ii) during inference, we employ entropy-based multi-block decoding with a KV-cache refresh mechanism to achieve high parallelism while maintaining accuracy. To better evaluate dLLMs, we also introduce AUP (Accuracy Under Parallelism), a new metric that jointly measures accuracy and parallelism. Experiments demonstrate that our d3LLM achieves up to 10$\times$ speedup over vanilla LLaDA/Dream and 5$\times$ speedup over AR models without much accuracy drop. Our code is available at https://github.com/hao-ai-lab/d3LLM.