Scaling Diffusion Language Models via Adaptation from Autoregressive Models

Diffusion Language Models (DLMs) have emerged as a promising new paradigm for text generative modeling, potentially addressing limitations of autoregressive (AR) models. However, current DLMs have been studied at a smaller scale compared to their AR counterparts and lack fair comparison on language modeling benchmarks. Additionally, training diffusion models from scratch at scale remains challenging. Given the prevalence of open-source AR language models, we propose adapting these models to build text diffusion models. We demonstrate connections between AR and diffusion modeling objectives and introduce a simple continual pre-training approach for training diffusion models. Through systematic evaluation on language modeling, reasoning, and commonsense benchmarks, we show that we can convert AR models ranging from 127M to 7B parameters (GPT2 and LLaMA) into diffusion models DiffuGPT and DiffuLLaMA, using less than 200B tokens for training. Our experimental results reveal that these models outperform earlier DLMs and are competitive with their AR counterparts. We release a suite of DLMs (127M-355M-7B) capable of generating fluent text, performing in-context learning, filling in the middle without prompt re-ordering, and following instructions https://github.com/HKUNLP/DiffuLLaMA.

한 줄 요약

GPT2와 LLaMA를 기반으로 확장된 DiffuGPT와 DiffuLLaMA는 200B 토큰 미만으로 훈련되어 기존 DLM을 능가하고 AR 모델과 경쟁한다.

핵심 기여도

GPT2와 LLaMA(127M~7B 파라미터)를 기반으로 DiffuGPT와 DiffuLLaMA를 생성하여 DLM의 규모를 확장.
200B 토큰 미만의 훈련 데이터로 기존 DLM보다 우수한 성능 달성.
DiffuLLaMA는 코드 생성, 인컨텍스트 학습, 인필링 등에서 뛰어난 성능을 보임.
7B 규모의 DLM을 성공적으로 구축하여 기존 1B 이하 DLM과 비교해 확장성을 입증.

핵심 아이디어

기존 DLM은 AR 모델 대비 규모가 작고 훈련이 어려운 문제가 있었다. 본 연구는 기존 오픈소스 AR 모델(GPT2, LLaMA)을 기반으로 DLM을 확장하는 새로운 접근법을 제안한다. AR 모델과 DLM의 목표 간 연결성을 분석하고, 주요 차이점인 **주의력 마스크**(attention mask)와 **시퀀스 예측 방식**을 해결하기 위해 **마스크 점진 해제**(attention mask annealing)와 **시프트 연산**(shift operation)을 도입했다. 이는 AR 모델의 인과 마스크(causal mask)를 DLM의 양방향 마스크로 전환하는 데 기여한다. 또한, **연속 확산**(continuous diffusion, CD)과 **이산 확산**(discrete diffusion, DD) 손실 함수 간의 차이를 분석하여, DD 손실이 AR 손실과 더 잘 정렬된다는 점을 밝혔다.

기술적 접근법

**모델 기반**: GPT2와 LLaMA2를 기반으로 DiffuGPT와 DiffuLLaMA를 생성.
**확산 프로세스**: 연속 시간 이산 확산(continuous-time discrete diffusion)을 정의하고, AR과 DLM의 목표 간 동등성을 수립(§3.1, §3.2).
**적응 전략**:
**주의력 마스크 점진 해제**(attention mask annealing)로 AR 모델의 인과 마스크를 해제.
**시프트 연산**(shift operation)을 통해 AR 모델의 구조를 유지하면서 DLM의 양방향 예측을 가능하게 함.
**훈련 데이터**: FineWeb과 SlimPajama를 사용, 200B 토큰 미만으로 훈련.
**샘플링 알고리즘**: 확산 모델의 샘플링 과정을 정의(§3.4).

주요 결과

**DiffuGPT**: GPT2 기반으로 훈련된 DLM이 GPT2 AR 모델보다 대부분의 작업에서 우수한 성능을 보임.
**DiffuLLaMA**: 7B 파라미터 규모로 훈련된 DLM이 기존 DLM(SED, Plaid 1B)을 능가.
**인필링 성능**: 256 단계의 확산 샘플링으로 1024 토큰을 생성하며, AR 모델과 유사한 속도를 보임.
**수학 추론**: GSM8K에서 DiffuLLaMA는 61.8%의 정확도를 달성(AR 훈련 대비 +12.1%).
**인컨텍스트 학습**: 4 샷 수준에서 정확도가 상승하며, CoT 형식은 성능 저하를 초래(지시 학습 부족 때문).

의의 및 한계

**의의**: 기존 DLM의 규모와 훈련 데이터 부족 문제를 해결하며, AR 모델과의 공정한 비교 가능성을 열었다. DiffuLLaMA는 인컨텍스트 학습, 코드 생성, 인필링 등에서 뛰어난 성능을 보여 DLM의 잠재력을 입증.
**한계**: CoT 형식의 지시 학습이 부족하여 수학 추론 성능이 제한됨. 또한, 확산 모델의 샘플링 과정은 여전히 AR 모델보다 느리고, 확산 단계 수가 증가할수록 계산 비용이 증가함. 지시 학습(instruction tuning)은 향후 연구 주제로 남음.

실용적 활용

DiffuLLaMA는 텍스트 생성, 인컨텍스트 학습, 코드 생성, 인필링 등 다양한 NLP 작업에 활용 가능하다. 특히, **중간 토큰 보완**(middle infilling)이나 **지시 없이 생성**(unconditional generation)이 필요한 시스템에서 AR 모델 대비 유연한 제어가 가능하다. 산업적으로는 **대화형 AI**, **코드 생성 도구**, **문서 자동 작성 시스템** 등에 적용 가능하며, 기존 AR 모델을 기반으로 확장할 수 있어 연구 및 개발 비용을 절감할 수 있다.