싱가포르 국립대, 확산 언어 모델(dLLM)의 병렬 디코딩을 가속화하는 새로운 패러다임 'DMax' 발표
National University of Singapore Presents "DMax": A New Paradigm For Diffusion Language Models (dLLMs) Enabling Aggressive Parallel Decoding.
핵심 요약
DMax는 확산 언어 모델의 병렬 디코딩 시 발생하는 오류 누적을 스스로 수정하는 방식으로 해결하여, 생성 품질을 유지하면서도 처리 속도를 대폭 향상함.
- 오류 누적 해결 — 디코딩 과정을 자기 수정 과정으로 재정의하여 잘못된 예측을 실시간으로 보정함.
- 병렬 디코딩 가속 — 마스크 임베딩에서 토큰 임베딩으로의 점진적 정제를 통해 더 공격적인 병렬 생성을 가능하게 함.
- 학습 전략 혁신 — 온-정책 균일 학습을 통해 모델이 마스크 입력과 자신의 오류로부터 올바른 토큰을 복구하도록 훈련함.
- 성능 검증 완료 — GSM8K 및 MBPP 벤치마크에서 정확도를 유지하면서도 기존 모델 대비 월등한 처리 속도를 기록함.
##TL;DR:
DMax는 디코딩을 점진적인 자기 수정 과정으로 재구성하여 오류 누적을 영리하게 완화하며, 모델이 생성 중에 스스로 잘못된 예측을 수정할 수 있도록 합니다.
##Abstract:
우리는 효율적인 확산 언어 모델(dLLM)을 위한 새로운 패러다임인 DMax를 제시합니다. 이는 병렬 디코딩에서의 오류 누적을 완화하여, 생성 품질을 유지하면서도 공격적인 디코딩 병렬화를 가능하게 합니다. 이진 마스크-투-토큰 전환을 통해 디코딩하는 기존의 마스크 dLLM과 달리, DMax는 디코딩을 마스크 임베딩에서 토큰 임베딩으로의 점진적인 자기 수정 과정으로 재구성합니다.
우리 접근 방식의 핵심은 마스크 dLLM과 균일 dLLM을 효율적으로 통합하는 새로운 학습 전략인 온-정책 균일 학습(On-Policy Uniform Training)으로, 모델이 마스크 입력과 자신의 잘못된 예측 모두로부터 올바른 토큰을 복구할 수 있도록 합니다. 이 기반 위에 우리는 소프트 병렬 디코딩(Soft Parallel Decoding)을 추가로 제안합니다. 우리는 각 중간 디코딩 상태를 예측된 토큰 임베딩과 마스크 임베딩 사이의 보간으로 표현하여, 임베딩 공간에서 반복적인 자기 수정을 가능하게 합니다.
다양한 벤치마크에 걸친 광범위한 실험은 DMax의 효율성을 입증합니다. 기존 LLaDA-2.0-mini와 비교했을 때, 우리 방법은 정확도를 유지하면서 GSM8K의 TPF를 2.04에서 5.47로 향상시킵니다. MBPP에서는 성능을 유지하면서 TPF를 2.71에서 5.86으로 증가시킵니다. H200 GPU 2대에서 우리 모델은 배치 사이즈 1 기준으로 평균 1,338 TPS를 달성합니다.
##Layman's Explanation:
핵심 아이디어는 확산 언어 모델이 여러 토큰을 동시에 채울 수 있기 때문에 일반 LLM보다 더 빠르게 텍스트를 생성할 수 있어야 한다는 것입니다. 하지만 실제로는 초기 단계의 잘못된 추측이 눈덩이처럼 불어나기 때문에 그 속도 이점이 제한됩니다. 모델이 잘못된 토큰을 선택하면 그 토큰이 다음 단계의 컨텍스트 일부가 되므로, 디코딩이 너무 공격적일 경우 품질이 빠르게 무너질 수 있습니다. DMax가 하는 일은 모델이 자신의 실수로부터 회복할 수 있는 더 나은 방법을 제공하는 것입니다. 마스크 슬롯에서 최종 토큰으로 가는 경직된 일방향 경로 대신, 모델이 중간 추측을 확정하기 전에 계속해서 정제할 수 있도록 합니다.
이 논문의 두 가지 주요 아이디어는 꽤 직관적입니다. 첫째, 모델은 자신의 불완전한 예측을 바탕으로 학습되므로 추론 시점에 실제로 발생할 오류를 정리하는 방법을 배웁니다. 둘째, 디코딩 중에 모든 추측을 즉시 최종적인 것으로 취급하는 대신 더 부드러운 중간 표현을 사용하여 불확실성을 유지하고 수정을 더 쉽게 만듭니다. 결과적으로 DMax는 품질 저하 없이 훨씬 더 많은 병렬 디코딩을 추진합니다. 논문의 수학 및 코딩 벤치마크에서 정확도를 원래 모델과 가깝게 유지하면서 큰 속도 향상을 얻었으며, 일부 낮은 병렬 설정에서는 정확도가 약간 향상되기도 했습니다. 따라서 주요 결론은 단순히 '더 빠른 확산 LLM'이 아니라, 공격적인 병렬 디코딩을 실제로 실용화할 수 있을 만큼 스스로를 잘 수정할 수 있는 확산 LLM이라는 점입니다.
######Link to the Paper: https://arxiv.org/pdf/2604.08302
