1. 연구 배경 및 문제 정의 (Problem)

• 기존 지식 증류 (Knowledge Distillation, KD) 의 한계:
  • 기존 KD 는 대규모 언어 모델 (LLM) 의 최종 출력 (logits) 분포만을 최적화하여 학생 모델을 훈련시킵니다.
  • 이 방식은 교사 모델의 중간 계층에서 일어나는 '사고 과정 (thought process)'을 블랙박스로 취급합니다.
  • 결과적으로 학생 모델은 최종 답변은 맞출 수 있지만, 교사 모델과 다른 추론 경로를 따르게 되어 일반화 능력이 떨어지고, 복잡한 추론 작업에서 실패할 수 있습니다 (그림 1 참조).
• 기존 중간 특징 (Intermediate Feature) 증류의 문제점:
  • 중간 계층의 특징을 직접 전달하려는 시도 (예: MSE, 비대칭 KL 발산) 가 존재하지만, 두 가지 주요 결함이 있습니다.
    1. MSE (평균 제곱 오차): 고차원 언어 임베딩 공간이 이방성 (anisotropic) 이라는 점을 고려하지 않아, 확률적 편차를 제대로 반영하지 못합니다.
    2. 비대칭 KL 발산 (Asymmetric KL Divergence):
       • Forward KL 은 학생이 교사의 모든 가능성을 커버하도록 강요하여 '과도한 확신 (overconfidence)'이나 '평균 찾기 (mean-seeking)' 행동을 유발합니다.
       • Reverse KL 은 저확률 영역을 무시하여 '과소 확신 (underconfidence)'이나 '모드 찾기 (mode-seeking)' 행동을 유발합니다.
  • 이러한 비대칭성은 중간 계층의 불확실성 (uncertainty) 정보를 왜곡시켜 최종 추론 능력을 저해합니다.

2. 제안 방법론: DISTILLLENS (Methodology)

저자들은 DISTILLLENS라는 새로운 증류 프레임워크를 제안하며, 이는 학생과 교사 모델의 진화하는 사고 과정 (evolving thought processes) 을 대칭적으로 정렬하는 것을 목표로 합니다.

• Logit Lens 활용:
  • 중간 계층의 은닉 상태 (hidden states, $h^{(l)}$) 를 모델의 사전 학습된 언디범딩 행렬 ($W_U$) 을 통해 어휘 공간 (vocabulary space) 으로 투영합니다.
  • 이를 통해 각 중간 계층에서 모델이 예측하는 확률 분포 $p^{(l)}(y|x)$를 계산합니다 (식 1, 2).
• 대칭적 발산 목표 (Symmetric Divergence Objective):
  • 투영된 확률 분포 간의 정렬을 위해 **대칭적 발산 (Symmetric Divergence)**을 사용합니다. 주로 **제이슨 - 샨논 발산 (Jensen-Shannon Divergence, JSD)**을 사용합니다.
  • 이론적 근거 (이중 측면 페널티):
    • JSD 는 과신 (overconfidence, $c_\theta \to \infty$) 에 대해 선형적인 페널티를 부과하고, 과소신 (underconfidence, $c_\theta \to 0$) 에 대해서는 유계된 (bounded) 페널티를 부과합니다.
    • 이는 학생 모델이 교사의 확률 분포를 정확히 따라가도록 ($c_\theta \to 1$) 유도하며, 고엔트로피 정보 유로를 보존하고 최종 추론에 필수적인 불확실성 정보를 유지하게 합니다.
• 최적화 목표:
  • 전체 손실 함수는 작업 손실 ($L_{task}$, 예: 최종 로그잇의 KL) 과 중간 정렬 손실 ($L_{inter}$) 의 가중 합입니다.
  • $L_{total} = L_{task} + \lambda \cdot L_{inter}$
  • 학생 모델의 특정 중간 계층을 교사 모델의 해당 계층에 매핑하여 (Layer Mapping) JSD 를 계산하고 역전파합니다.

3. 주요 기여 (Key Contributions)

1. DISTILLLENS 프레임워크 도입: Logit Lens 를 통해 중간 계층의 분포를 투영하고, 대칭적 발산 (JSD) 을 사용하여 LLM 의 사고 과정을 대칭적으로 증류하는 새로운 방법론 제시.
2. 이론적 증명: 대칭적 손실 지형 (loss landscape) 에 대한 이론적 유도를 통해, 이 방법이 과신과 과소신 양쪽을 모두 방지하는 '이중 측면 페널티 (dual-sided penalty)'를 부과함을 증명.
3. 광범위한 실험 검증: GPT-2 및 Llama 아키텍처를 대상으로 한 실험을 통해, 기존 표준 KD 및 특징 전달 (feature-transfer) 기반 방법론보다 다양한 지시 따르기 (instruction-following) 벤치마크에서 일관되게 우월한 성능을 보임을 입증.

4. 실험 결과 (Results)

• 데이터셋 및 모델:
  • 교사: GPT-2-1.5B, Llama-7B
  • 학생: GPT-2-120M/340M, TinyLlama-1.1B
  • 데이터: Databricks-dolly-15k
  • 평가 벤치마크: DollyEval, SelfInst, VicunaEval, S-NI, UnNI
• 성능 비교:
  • GPT-2-120M: DISTILLLENS 는 평균 Rouge-L 점수 21.12를 기록하여, 표준 KD(17.74) 와 SFT(16.56) 를 크게 상회했습니다.
  • GPT-2-340M: 평균 점수 23.72로, 교사 모델 (23.52) 의 성능을 약간 능가했습니다.
  • TinyLlama-1.1B: 평균 점수 25.48로, 기존 최상위 기법 (FKL+RKL, 24.56) 보다 1.6 포인트 이상 향상되었습니다.
• 온-폴리시 (On-policy) vs 오프-폴리시 (Off-policy):
  • 기존 온-폴리시 증류 (MiniLLM 등) 는 노출 편향 (exposure bias) 을 줄이지만 학습 속도가 느립니다.
  • DISTILLLENS 는 오프-폴리시 방식임에도 불구하고 노출 편향을 효과적으로 줄여, 온-폴리시 방법과 유사한 성능을 내면서 학습 속도는 훨씬 빠릅니다.
• 특징 전달 비교:
  • 기존 MSE 기반 특징 증류나 비대칭 KL 기반 방법 (FDD) 보다 DISTILLLENS(JSD/JD) 가 모든 모델 크기와 벤치마크에서 우월한 성능을 보였습니다.
• Ablation Study:
  • 중간 계층의 수를 5 개로 설정했을 때 (교차/interleaved 패턴) 최적의 성능을 발휘했습니다.

5. 의의 및 결론 (Significance)

• 사고 과정의 정렬: DISTILLLENS 는 단순히 정답을 맞추는 것을 넘어, 모델이 어떻게 그 정답에 도달하는지 (중간 추론 단계) 를 학생 모델에게 가르칩니다. 이는 모델의 일반화 능력과 추론 능력을 향상시킵니다.
• 효율성과 성능의 균형: 계산 비용이 증가하는 온-폴리시 학습의 단점을 보완하면서도, 오프-폴리시 학습의 속도 장점을 유지하면서 높은 성능을 달성하는 효율적인 방법을 제공합니다.
• 모듈성: 기존 증류 방법 (FKL, RKL, AKL 등) 과 쉽게 결합하여 성능을 추가로 향상시킬 수 있는 모듈형 접근법입니다.
• 미래 전망: 추론 능력을 갖춘 더 작은 모델 개발과, 서로 다른 아키텍처 간의 지식 증류 (Cross-architecture distillation) 로의 확장을 위한 기반을 마련했습니다.

요약하자면, DISTILLLENS 는 Logit Lens를 통해 모델의 내부 사고 과정을 가시화하고, **대칭적 발산 (JSD)**을 통해 그 과정을 정밀하게 증류함으로써, 기존 지식 증류의 한계를 극복하고 더 강력하고 정확한 언어 모델을 만드는 획기적인 접근법입니다.