1. 문제 정의 (Problem Definition)

대규모 언어 모델 (LLM) 의 지식 증류 (Knowledge Distillation, KD) 는 모델의 크기와 복잡도가 증가함에 따라 점점 더 중요해지고 있습니다. 그러나 기존 증류 방법론들은 다음과 같은 한계를 가지고 있습니다.

• 행동 모방의 한계: 기존 방법들은 주로 '교사 모델 (Teacher)'의 행동 (출력 분포) 을 모방하는 데 초점을 맞춥니다. 이는 행동 복제 (Behavioral Cloning, BC) 에 해당하며, 교사가 학습했던 원래의 학습 환경 (Original Learning Environment) 과 그 환경이 부여한 보상 함수 (Reward Function) 를 고려하지 않습니다.
• 환경 무시: 역강화학습 (Inverse Reinforcement Learning, IRL) 이론에 따르면, 단순히 행동을 모방하는 것보다 학습자가 교사가 학습했던 환경과 보상 구조를 이해하고 그 환경에서 학습하는 것이 더 효과적입니다. 하지만 기존 KD 는 학생 모델 (Student) 을 교사의 원래 학습 환경에 배치하지 않습니다.
• MiniLLM 의 복잡성: 최근 제안된 MiniLLM 은 역강화학습 프레임워크를 도입하여 보상 함수를 고려하지만, 정책 경사 (Policy Gradient) 방법을 사용하여 훈련 안정성 문제가 발생하고, 기존 지도 학습 (Supervised Learning) 패러다임과 통합하기 어렵다는 단점이 있습니다. 또한 다양한 발산 함수 (Divergence Functions) 나 블랙박스 모델에 적용하는 데 제약이 있습니다.

2. 방법론 (Methodology: X-KD)

저자들은 경험적 지식 증류 (Experiential Knowledge Distillation, X-KD) 라는 새로운 프레임워크를 제안합니다. 이는 베이즈 역강화학습 (Bayesian IRL) 의 일종인 AVRIL (Approximated Variational Reward Imitation Learning) 프레임워크를 기반으로 합니다.

핵심 아이디어

학생 모델이 단순히 교사의 행동을 복사하는 것이 아니라, 교사가 학습했던 원래 환경의 보상 함수를 학습하고, 그 보상 함수 하에서 최적의 정책을 학습하도록 유도합니다.

주요 구성 요소 및 수식

1. 원래 보상 모델링 (Original Reward Modeling, ORM):

   • 교사의 원래 보상 함수 $R_\pi$를 알 수 없으므로, 이를 확률 분포 $p(R|\pi)$로 가정하고 파라미터화된 분포 $q_\phi(R)$로 근사합니다.
   • AVRIL 프레임워크를 사용하여 보상 함수와 정책 (Q-value decoder) 을 동시에 학습합니다.
   • 손실 함수는 행동의 가능도 (Likelihood), 보상 사후분포와 사전분포 간의 KL 발산, 그리고 TD 오차 (Temporal Difference error) 를 보상 분포에 맞추는 항으로 구성됩니다.

2. X-KD 의 일반화된 손실 함수:
   X-KD 는 기존의 증류 손실 (지식 전달) 과 경험적 정규화 항 (환경 학습) 을 결합합니다.

   • 시퀀스 레벨 X-KD (Sequence-level X-KD):
     $$L_{orm} = L_{seq} + L_{ex}$$
     여기서 $L_{seq}$는 교사 정책에서 학생 정책으로의 지식 증류 (Sequence-level KD) 이고, $L_{ex}$는 학생 정책이 교사의 원래 보상 함수와 일관되도록 하는 경험적 항입니다.
   • 지도 학습 X-KD (Supervised X-KD):
     기존 KL 발산 기반의 지도 증류 (Forward-KL) 와 경험적 항을 결합합니다.
     $$L_{xkd} = L_{kd} + L_{ex}$$
   • 일반화된 X-KD (Generalized X-KD):
     GKD (Generalized KD) 와 경험적 항을 결합하여, 온-정책 (on-policy) 데이터와 오프-정책 (offline) 데이터를 모두 활용하고 다양한 발산 함수 (예: $\beta$-JS divergence) 를 지원합니다.
     $$L_{g-xkd} = L_{gkd} + L_{ex}$$

3. 구현 및 유연성:

   • 지도 학습 패러다임 유지: X-KD 는 정책 경사 (Policy Gradient) 를 사용하지 않고, 기존 지도 학습 프레임워크 내에서 정규화 항을 추가하는 방식으로 구현되어 기존 LLM 훈련 파이프라인에 쉽게 통합됩니다.
   • 블랙박스 지원: 교사 모델의 내부 분포에 접근할 수 없는 경우 (블랙박스), 시퀀스 레벨의 X-KD 목적 함수를 사용하여 교사 모델의 출력 응답을 통해 경험적 학습을 수행할 수 있습니다.

3. 주요 기여 (Key Contributions)

1. 새로운 프레임워크 제안: 베이즈 역강화학습을 기반으로 한 X-KD 를 제안하여, 학생 모델이 교사의 행동 모방을 넘어 교사의 원래 학습 환경에서 학습하도록 합니다.
2. 유연성과 단순성: 다양한 증류 방법 (시퀀스 레벨, KL 기반, GKD 등) 과 통합 가능하며, 여전히 지도 학습 패러다임을 유지하여 구현이 간단하고 적용 범위가 넓습니다.
3. 실험적 검증: 여러 자연어 생성 (NLG) 작업에서 기존 방법들 (GKD, MiniLLM 등) 보다 우수한 성능을 보이며, 성능 - 다양성 (Performance-Diversity) 트레이드오프와 데이터 효율성 측면에서도 우위를 입증했습니다.

4. 실험 결과 (Results)

저자들은 추상적 요약 (XSum), 기계 번역 (WMT14 en-de), 산술 추론 (GSM8k) 작업에서 X-KD 를 평가했습니다.

• 성능 향상:
  • XSum (요약): G-XKD 가 GKD 와 MiniLLM 보다 ROUGE 점수에서 일관되게 우위를 보였습니다.
  • WMT14 (번역): 다양한 모델 크기 설정에서 X-KD 가 모든 베이스라인을 능가했습니다. 특히 샘플링 온도가 낮아질 때 (다양성이 높아질 때) G-XKD 가 번역 품질과 다양성을 동시에 잘 유지했습니다.
  • GSM8k (추론): 시퀀스 레벨 및 Forward-KL 증류에 경험적 학습을 도입한 SeqXKD 와 S-XKD 가 정확도 향상을 보였습니다.
• 성능 - 다양성 트레이드오프:
  • X-KD 는 GKD 보다 더 나은 성능 - 다양성 균형을 보여주었습니다. 즉, 다양성을 높여도 성능 저하가 상대적으로 적게 발생하여 더 일반화된 전략을 학습했음을 시사합니다.
• 데이터 효율성:
  • 전체 학습 데이터의 50%~75% 만 사용해도 베이스라인이 전체 데이터로 학습한 것과 유사하거나 더 좋은 성능을 달성했습니다. 이는 X-KD 가 표면적 패턴이 아닌 작업의 근본 원리를 더 효율적으로 포착함을 의미합니다.
• 블랙박스 증류:
  • Llama-2-7b(교사) 와 T5-large(학생) 를 사용한 블랙박스 증류 실험에서, GPT-4 평가를 통해 SeqXKD 가 기존 SeqKD 보다 더 높은 승률 (Win Rate) 을 기록했습니다.

5. 의의 및 결론 (Significance & Conclusion)

이 논문은 대규모 언어 모델의 지식 증류에 있어 학습 환경의 중요성을 재조명했습니다. 단순히 교사의 출력을 복사하는 것을 넘어, 교사가 학습했던 '경험 (Reward Function)'을 학생에게 전달함으로써 더 강력하고 효율적인 모델을 만드는 데 성공했습니다.

• 실용성: 복잡한 정책 경사 기반의 방법론 (MiniLLM 등) 과 달리, 기존 지도 학습 파이프라인에 쉽게 적용 가능하여 실제 엔지니어링 환경에서의 활용도가 높습니다.
• 일반성: 화이트박스뿐만 아니라 블랙박스 모델에도 적용 가능하며, 다양한 발산 함수와 증류 시나리오에 유연하게 대응합니다.

한계점:

• 경험적 가중치 (Experiential weight, $\lambda$) 에 대한 민감도로 인해 하이퍼파라미터 튜닝 비용이 발생할 수 있습니다.
• 현재 연구는 주로 화이트박스 설정에 집중되어 있으며, 블랙박스 설정에 대한 체계적인 연구는 향후 과제로 남았습니다.

요약하자면, X-KD 는 경험적 학습 이론을 LLM 증류에 적용하여, 학생 모델이 교사의 '지식'뿐만 아니라 '학습 경험 (환경과 보상)'까지 습득하도록 함으로써 성능, 다양성, 데이터 효율성을 동시에 개선한 획기적인 접근법입니다.