기술 요약: 온-폴리시 증류를 통한 데이터 효율적 자기회귀-확산 언어 모델 (On-Policy Distillation을 통한 Data-Efficient Autoregressive-to-Diffusion Language Models)

1. 문제 정의

본 논문은 사전 학습된 자기회귀 언어 모델(ARLM)을 확산 언어 모델(DLM)로 변환하는 과정에서 발생하는 고유한 과제들을 다룹니다. 최근 연구들은 인과적 어텐션(causal attention)을 양방향 어텐션(bidirectional attention)으로 교체하고 확산 목적 함수(diffusion objective)로 재학습함으로써 ARLM을 DLM으로 변환하려고 시도해 왔으나, 저자들은 학습 효율과 성능을 제한하는 두 가지 결정적인 "분포 변화(distribution shifts)"를 식별했습니다:

1. 지식 유지 불일치(Knowledge-Retention Mismatch): 다음 토큰 예측 목적 함수(ARLM)에서 확산 목적 함수(DLM)로 전환하는 과정은 기존 ARLM 사전 학습 중에 습득한 방대한 지식을 약화시키거나 유실할 수 있습니다.
2. 학습-추론 상태 불일치(Train-Inference State Mismatch): 표준 DLM은 균등 무작위 마스킹(uniform random masking, 순방향 과정)으로 오염된 시퀀스 상에서 학습됩니다. 그러나 추론 시, 모델은 신뢰도 기반 디코딩 휴리스틱에 의해 유도되는 역방향 언마스킹 궤적(reverse unmasking trajectories)을 통해 시퀀스를 생성합니다. 추론 중에 마주치는 상태는 학습 중에 경험하는 무작위 상태와 크게 다르며, 이는 최적화되지 않은 학습으로 이어집니다.

기존의 변환 방법들은 이를 사전 학습의 연장선으로 취급하며, 이러한 불일치를 명시적으로 해결하지 못하거나 DLM을 처음부터 사전 학습하기 위해 천문학적인 양의 데이터(수조 개의 토큰)를 요구하는 경우가 많습니다.

2. 방법론: 온-폴리시 확산 언어 모델 (OPDLM)

저자들은 **온-폴리시 증류(On-Policy Distillation, OPD)**를 사용하여 ARLM을 DLM으로 변환하는 프레임워크인 OPDLM을 제안합니다. 이 접근 방식은 변환 과정을 사전 학습의 문제가 아닌 사후 학습(post-training) 문제로 재정의합니다.

핵심 메커니즘

OPDLM은 학생 모델이 스스로 학습 데이터를 생성하고, 원래의 ARLM이 동결된(frozen) 교사 역할을 수행하는 "자기-OPD(self-OPD)" 전략을 채택합니다.

• 학생 모델 (Student Model): 원래 ARLM 가중치로부터 초기화되었으며, 블록 단위 인과적(양방향) 어텐션을 사용하도록 수정된 ARLM 체크포인트입니다.
• 교사 모델 (Teacher Model): 동결된 원래의 ARLM입니다.
• 학습 궤적 (On-Policy Training Trajectory): 학생 DLM은 무작위 마스크를 샘플링하는 대신, 자신의 역방향 디코딩 궤적($\tau \sim \Gamma^{rev}_{\theta}$)을 샘플링합니다. 이 궤적은 모델이 추론 중에 실제로 마주치는 상태를 나타냅니다.
• 증류 목적 함수 (Distillation Objective):
  • 샘플링된 역방향 궤적으로부터 부분적으로 노이즈가 제거된 상태 $x_t$를 선택합니다.
  • 블록 내 각 마스킹된 토큰 위치 $i$에 대해, 학생의 예측 분포 $p_{\theta}$가 교사의 분포를 따르도록 학습합니다.
  • 결정적으로, 교사의 분포는 터미널 시퀀스 $x_0$(해당 궤적의 최종 생성 출력)의 **언마스킹된 인과적 접두사(unmasked causal prefix)**에 조건화됩니다. 구체적으로, 타겟은 $p_{ARLM}(\cdot | x_{<b}^0, x_{<i}^0)$입니다.
  • 손실 함수는 마스킹된 위치에 대한 교사의 토큰 수준 분포와 학생의 분포 사이의 KL 발산(Kullback-Leibler divergence)입니다.

커리큘럼 학습 (Curriculum Learning)

변환 직후의 불안정성(구조적 변화인 인과적 어텐션에서 블록 단위 양방향 어텐션으로의 전환으로 인해 학생 모델이 낮은 품질의 궤적을 생성할 수 있음)을 해결하기 위해, 저자들은 **롤아웃 길이 커리큘럼(Rollout-Length Curriculum)**을 도입합니다. 온-폴리시 생성 길이를 최소($L_{min}$)에서 최대($L_{max}$)까지 웜업 단계에 걸쳐 점진적으로 증가시켜, 학생 모델이 긴 시퀀스를 다루기 전에 일관된 컨텍스트 상에서 교사를 모방하는 법을 배우도록 보장합니다.

3. 주요 기여

1. ARLM-to-DLM을 사후 학습으로 정립: 본 논문은 DLM을 처음부터 사전 학습하는 데 드는 막대한 계산 비용 없이도 DLM을 얻을 수 있음을 입증하여, 변환을 효율적인 사후 학습 절차로 포지셔닝했습니다.
2. 분포 변화 해결: 학생 자신의 역방향 궤적(온-폴리시)에서 학습함으로써 학습-추론 상태 불일치를 제거합니다. 또한 원래의 ARLM을 교사로 사용함으로써 지식 유지를 보장합니다.
3. 극단적인 데이터 효율성: OPDLM은 기존 DLM 베이스라인(예: LLaDA, Dream, SDAR)에 비해 15배에서 7,000배 적은 학습 토큰을 사용하여 경쟁력 있는 성능을 달로합니다.
4. 범용성 및 태스크 특화 적용 가능성: 이 프레임워크는 별도의 DLM 사전 학습 단계 없이도 범용 DLM과 태스크 특화 모델(예: 수학용)을 모두 생성할 수 있음을 보여줍니다.

4. 실험 결과

저자들은 일반 지식, 수학, 코딩 벤치마크에 대해 LLaDA, Dream, SDAR, Fast-dLLM-v2와 같은 베이스라인을 대상으로 OPDLM(4B 및 8B 규모)을 평가했습니다.

• 성능 대 효율성: OPDLM-8B는 0.066B 개의 학습 토큰과 4.2 × 10¹⁸ FLOPs만을 사용하여 AIME-24(정확도 14.7%) 및 GPQA-Diamond(36.1%)와 같은 벤치마크에서 강력한 결과를 달성했습니다. 반면, 베이스라인들은 수천억 개의 토큰이 필요했습니다 (예: LLaDA-8B는 1,500B 개의 토큰 사용).
• 절제 연구 (Ablation Studies):
  • 온-폴리시 vs 오프-폴리시: 표준 무작위 마스킹(오프-폴리시)으로 교체했을 때 성능이 일관되게 하락하였으며, 이는 모델 자신의 추론 상태에서 학습하는 것이 성공의 주요 동력임을 확인시켜 줍니다.
  • 지식 유지: OPDLM은 변환 과정에서 명시적으로 목표하지 않았던 제로샷 능력, 즉 확장된 사고(extended thinking)( <thought> 태그 내의 추론 과정) 및 다국어 생성 능력을 유지했습니다.
• 추론 효율성: 모델은 멀티 토큰 디코딩을 지원합니다. 디코딩 신뢰도 임계값과 블록 크기를 조정함으로써, 저자들은 약간의 정확도 비용을 감수하면서 처리량(throughput)이 단계당 2~3 토큰을 초과할 수 있음을 입증했습니다.
• 태스크 특화 적응: MATH 데이터셋에 엄격하게 학습되었을 때, OPDLM-MATH는 표준 DLM 베이스라인을 능가했으며, 별도의 DLM 사전 학습이나 검증 가능한 보상을 사용하는 강화 학습(RL) 없이도 전문화된 RL 학습 모델(TraDo)과 대등한 성능을 보였습니다.

5. 의의 및 주장

본 논문은 OPDLM이 DLM 학습의 효율성에 있어 새로운 파레토 최적선(Pareto frontier)을 구축한다고 주장합니다. 온-폴리시 증류를 활용함으로써, ARLM을 DLM으로 변환하는 작업이 최소한의 계산 자원으로 가능함을 보여주며, 이를 통해 처음부터 사전 학습하는 "과도한 비용" 없이도 확산 기반 텍스트 생성을 가능하게 합니다.

저자들은 이 접근 방식이 커뮤니티가 미래의 ARLM을 효율적으로 확산 모델로 적응시킬 수 있게 해준다고 강조합니다. 즉, 원본 모델의 풍부한 지식을 보존하면서 동시에 확산 모델의 병렬 디코딩 이점을 얻을 수 있게 합니다. 이 연구는 성숙한 자기회귀 모델과 잠재력이 큰 확산 모델 사이의 간극을 메우는 실질적인 경로를 제시합니다.