온-정책 증류를 통한 데이터 효율적 자기회귀-확산 언어 모델

초록

우리는 자기회귀 모델(ARLM)을 확산 언어 모델(DLM)로 변환하는 과정을 연구한다. 기존 연구는 처음부터 사전 학습을 진행하는 대신, ARLM의 인과적 어텐션을 양방향 어텐션으로 대체한 후 결과 모델을 DLM 목적 함수로 학습시킨다. 그러나 이러한 접근법은 두 가지 분포 이동을 초래한다. 첫째, 다음 토큰 예측 목적 함수에서 DLM 목적 함수로 전환할 때 ARLM이 학습 과정에서 획득한 지식이 소실될 수 있다. 둘째, 표준 DLM은 학습 손실이 추론 시 신뢰도 기반 디코딩에 의해 생성된 궤적이 아닌 무작위 마스킹된 시퀀스에 대해 정의되므로 학습-추론 불일치가 발생한다. 이러한 두 가지 문제를 해결하기 위해, 우리는 ARLM에서 DLM으로의 변환을 위해 정책 내 증류(OPD)를 사용하는 정책 내 확산 언어 모델(OPDLM)을 도입한다. 구체적으로, OPDLM은 자기-OPD를 통해 학습되며, 학생 모델(양방향 어텐션을 갖춘 ARLM)은 자체 궤적을 생성하고, 교사 모델(원래의 고정된 ARLM)은 해당 궤적에 대한 대상 로짓을 제공하여 지식을 증류한다. OPDLM은 정책 내 방식으로 직접 학습함으로써 DLM의 학습-추론 불일치를 제거하며, 원본 모델로부터의 증류를 통해 ARLM의 지식 보존을 강화한다. 실험 결과, OPDLM은 다양한 작업에서 강력한 성능을 보이면서도 15배에서 7,000배 더 적은 학습 토큰만을 필요로 한다. OPDLM은 DLM 사전 학습의 엄청난 비용을 피하고, DLM 변환을 ARLM 후속 학습의 한 형태로 자리매김한다.

English

We study the transformation of autoregressive models (ARLMs) into diffusion language models (DLMs). Rather than pretraining from scratch, prior work replaces the causal attention in ARLMs with bidirectional attention and then trains the resulting model using a DLM objective. However, these approaches incur two distribution shifts. First, transitioning from a next-token prediction objective to a DLM objective can discard knowledge acquired by the ARLM during training. Second, standard DLMs suffer from a train-inference mismatch, as the training loss is defined on randomly masked sequences rather than the trajectories encountered at inference produced by confidence-based decoding. To address both challenges, we introduce an On-Policy Diffusion Language Model (OPDLM) in which On-Policy Distillation (OPD) is employed for ARLM-to-DLM transformation. Specifically, OPDLM is trained via self-OPD, where the student, an ARLM with bidirectional attention, generates its own trajectories, and the teacher, the original frozen ARLM, distills its knowledge by providing target logits on these trajectories. By training directly in an on-policy manner, OPDLM eliminates the train-inference mismatch in DLMs, while distillation from the original model enhances knowledge retention from the ARLM. Empirical results demonstrate that OPDLM requires 15x to 7,000x fewer training tokens with strong performance across a wide variety of tasks. OPDLM avoids the prohibitive cost of DLM pretraining and positions DLM transformation as a form of ARLM post-training.