Entropy-Aware On-Policy Distillation of Language Models

이 논문은 고엔트로피 상태의 교사 모델 예측 시 역 KL 발산을 보완하기 위해 순 KL 발산을 결합한 '엔트로피 인식 온-정책 증류 (Entropy-Aware On-Policy Distillation)'를 제안하여, 생성 다양성을 유지하면서 수학 추론 성능을 크게 향상시켰음을 보여줍니다.

핵심

🎓 1. 상황 설정: "선생님의 답을 그대로 외우는 학생"

가상의 상황을 상상해 보세요.

• 선생님 (Teacher): 수학 천재입니다. 문제를 풀 때 확실히 정답인 경우엔 "정답은 5 입니다!"라고 아주 단호하게 말합니다. 하지만 어려운 문제나 여러 가지 풀이가 가능한 상황에서는 "음... 5 일 수도 있고, 6 일 수도 있고, 7 일 수도 있겠네..."라며 여러 가능성을 열어두고 고민합니다.
• 학생 (Student): 이 선생님의 말을 듣고 똑같이 문제를 풀려고 노력합니다.

⚠️ 2. 기존 방식의 문제: "무조건 확신만 믿는 학생"

기존의 AI 학습 방식 (Reverse KL) 은 학생에게 **"선생님이 가장 확신하는 답만 골라라"**라고 가르쳤습니다.

• 상황: 선생님이 "5 일 수도 있고 6 일 수도 있어"라고 고민할 때 (불확실성이 높은 상황), 학생은 "아, 선생님이 5 라고 했으니 5 가 정답이겠지!"라고 무조건 5 로 고정해 버립니다.
• 결과:
  1. 다양성 상실: 학생은 6 이나 7 이라는 다른 가능성도 잊어버리고, 오직 5 만 고집하게 됩니다.
  2. 혼란: 선생님이 진짜로 고민하고 있을 때 (여러 답이 공존할 때), 학생은 그 고민의 '흐름'을 이해하지 못하고 엉뚱한 방향으로 쏠리게 되어 학습이 불안정해집니다.
  3. 비유: 마치 요리 실습에서, 선생님이 "이 요리는 소금 1 티스푼도 좋고 1.5 티스푼도 좋네"라고 말했을 때, 학생이 "선생님이 1 티스푼을 먼저 말했으니 1 티스푼만 넣어야지!"라고 딱딱하게만 따라 하는 것과 같습니다. 결국 창의적인 요리 (다양한 풀이) 가 사라지고 실패할 확률이 높아집니다.

✨ 3. 이 논문의 해결책: "상황에 맞춰 가르치는 '지각 있는' 선생님"

이 논문은 **"지각 있는 온-정책 증류 (Entropy-Aware On-Policy Distillation, EOPD)"**라는 새로운 방법을 제안합니다. 핵심은 **"선생님이 확신할 때는 확실히 배우고, 고민할 때는 고민하는 법까지 배우라"**는 것입니다.

🧠 핵심 아이디어: 두 가지 학습 모드

이 방법은 학생이 문제를 풀 때, 선생님의 '고민 정도 (엔트로피)'를 보고 두 가지 방식을 상황에 맞게 섞어 사용합니다.

1. 선생님이 확신할 때 (낮은 엔트로피):

   • 상황: "정답은 5 입니다!"라고 단호할 때.
   • 학습법: "선생님이 말하는 5 를 정확히 따라 해라." (기존 방식 유지)
   • 효과: 효율적이고 빠르게 핵심을 배웁니다.

2. 선생님이 고민할 때 (높은 엔트로피):

   • 상황: "5 일 수도 있고 6 일 수도 있어..."라고 여러 가능성을 제시할 때.
   • 학습법: "선생님이 고민하는 그 **모든 가능성 (5, 6, 7)**을 다 기억해 둬라. 5 만 고집하지 말고, 선생님이 가진 '의심'과 '다양성'까지 그대로 받아라." (새로운 방식 적용)
   • 효과: 학생이 다양한 해결책을 탐색할 수 있게 되어, 나중에 비슷한 어려운 문제를 만나도 더 유연하게 대처할 수 있습니다.

🎨 비유: "등산 가이드"

• 기존 방식: 가이드가 "저기 길이 하나만 있어"라고 말하면 그 길만 가고, "여러 길이 있을 수 있어"라고 말하면 학생은 당황해서 아무 길도 못 가거나 엉뚱한 길로만 갑니다.
• 새로운 방식 (EOPD): 가이드가 "여기 길이 여러 개 있어"라고 말하면, 학생은 "아, 저기에도 길이 있고, 저기에도 있구나"라고 모든 길을 훑어보며 지도에 그려 넣습니다. 나중에 그 지역을 다시 갈 때, 어떤 길이든 찾아낼 수 있게 됩니다.

📈 4. 실제 성과: "수학 시험에서 더 높은 점수"

이 새로운 방법으로 훈련된 학생들은 실제로 수학 문제 풀이 능력에서 큰 발전을 보였습니다.

• 다양성 유지: 학생이 문제를 풀 때, 다양한 풀이 경로를 시도할 수 있게 되어 (다양한 답을 낼 수 있게 되어) 실수할 확률이 줄어듭니다.
• 성적 향상: 복잡한 수학 문제 (AIME, MATH 등) 를 풀 때, 기존 방식보다 **정답률 (Pass@8)**이 크게 향상되었습니다. 특히 40 억 개 파라미터 (4B) 모델의 경우 정답률이 5% 이상이나 높아졌습니다.

💡 5. 결론: 왜 이 기술이 중요한가요?

이 논문은 **"AI 가 무조건 확신하는 것만 배우는 게 아니라, '모르는 것'이나 '고민하는 것'도 배우게 해야 진짜 똑똑해진다"**는 사실을 증명했습니다.

• 기존: "정답만 외워라." (다양성 파괴)
• 새로운 방식: "선생님이 고민하는 부분까지 함께 고민해라." (다양성 보존 + 효율성 유지)

이 기술은 거대한 AI 모델을 작은 모델로 옮길 때, 작은 모델이 큰 모델의 '지혜'와 '유연함'까지 모두 가져갈 수 있게 해줍니다. 이는 더 작고 빠르면서도 똑똑한 AI 를 만드는 데 중요한 열쇠가 될 것입니다.

기술 요약

1. 문제 정의 (Problem)

대규모 언어 모델 (LLM) 의 능력을 작은 모델로 전달하는 온-폴리시 증류 (On-Policy Distillation, OPD) 는 추론 시 발생하는 분포 불일치 (Distribution Mismatch) 를 해결하기 위해 유망한 접근법입니다. 기존 OPD 는 주로 역 KL 발산 (Reverse KL Divergence) 을 최적화 목표로 사용합니다.

• 역 KL 의 한계: 역 KL 은 '모드 탐색 (Mode-seeking)' 특성을 가집니다. 즉, 교사 (Teacher) 모델이 높은 확률을 가진 토큰에 집중하도록 유도하지만, 교사 분포의 높은 엔트로피 (High Entropy) 영역, 즉 여러 가지 유효한 추론 경로가 존재하는 불확실한 위치에서는 학습 신호가 불안정해집니다.
• 발생하는 문제:
  1. 다양성 붕괴 (Diversity Degradation): 역 KL 은 학생 모델이 교사 모델의 불확실성을 학습하지 못하게 하여, 생성되는 토큰의 엔트로피를 급격히 낮춥니다. 실험 결과, 기존 OPD 는 교사 모델의 고엔트로피 토큰을 18.5% 에서 6.8% 로 크게 감소시켰습니다.
  2. 학습 불안정성: 교사 모델이 불확실한 (고엔트로피) 위치에서 역 KL 기반의 보상 신호는 발산하거나 수렴하지 못하는 불안정한 경향을 보입니다. 이는 복잡한 추론 작업에서 중요한 의사결정 지점 (Key decision points) 에서 학생 모델의 성능 저하를 초래합니다.

2. 방법론 (Methodology)

저자들은 역 KL 의 효율성과 순 KL (Forward KL) 의 다양성 보존 능력을 결합한 엔트로피 인식 온-폴리시 증류 (Entropy-Aware On-Policy Distillation, EOPD) 를 제안합니다.

• 핵심 아이디어: 교사 모델의 토큰 수준 엔트로피 ($H^{te}_t$) 를 기반으로 학습 목적 함수를 동적으로 적응시킵니다.
  • 낮은 엔트로피 영역 (교사가 확신하는 경우): 역 KL (Reverse KL) 을 사용하여 효율적이고 빠른 수렴을 유도합니다.
  • 높은 엔트로피 영역 (교사가 불확실한 경우): 순 KL (Forward KL) 을 활성화하여 교사 모델의 전체 분포 구조와 불확실성을 학생에게 전달합니다.
• 목적 함수:
  $$L^{EOPD}_t = L^{OPD}_t + \mathbb{I}[H^{te}_t > \tau] \cdot L^{FKL}_t$$
  여기서 $\tau$는 엔트로피 임계값이며, $L^{FKL}_t$는 교사 모델의 상위 $k$개 토큰에 대해서만 계산하여 계산 효율성을 유지합니다.
• 알고리즘: PPO(Proximal Policy Optimization) 스타일의 업데이트에 통합되어, 고엔트로피 토큰에서는 모드 커버링 (Mode-covering) 특성을, 저엔트로피 토큰에서는 모드 탐색 특성을 발휘하도록 설계되었습니다.

3. 주요 기여 (Key Contributions)

1. 다양성 저하 및 학습 불안정성에 대한 체계적 분석: 역 KL 기반 증류가 고엔트로피 토큰에서 다양성을 붕괴시키고 학습 신호를 불안정하게 만든다는 것을 정량적으로 증명했습니다.
2. EOPD 프레임워크 제안: 교사 모델의 불확실성을 감지하여 역 KL 과 순 KL 을 선택적으로 적용하는 새로운 증류 전략을 제시했습니다. 이는 계산 오버헤드를 늘리지 않으면서도 교사 모델의 분포 구조를 효과적으로 전달합니다.
3. 추론 벤치마크에서의 성능 향상: 수학적 추론 작업에서 기존 온-폴리시 증류 방법 대비 일관된 성능 개선을 입증했습니다.

4. 실험 결과 (Results)

저자는 Qwen3 시리즈 (0.6B, 1.7B, 4B) 를 학생 모델로, Qwen3-8B 를 교사 모델로 사용하여 6 가지 수학적 추론 벤치마크 (MATH500, AIME24/25, AMC23 등) 에서 실험을 수행했습니다.

• 성능 향상:
  • Pass@8 정확도: Qwen3-4B-Base 모델에서 기존 OPD 대비 +5.05%p 향상, 1.7B 모델에서 +2.39%p, 0.6B 모델에서 +1.37%p 향상.
  • Avg@8 정확도: 모든 모델 크기에서 일관된 개선 (예: 4B 모델에서 +1.80%p).
  • Pass@k 분석: $k$가 증가함에 따라 EOPD 와 OPD 간의 성능 격차가 더 벌어지는데, 이는 EOPD 가 더 다양한 추론 경로를 탐색하여 정답에 도달할 확률을 높였음을 시사합니다.
• 엔트로피 보존: EOPD 는 고엔트로피 영역에서 교사 모델의 확률 질량을 더 잘 보존하며, OPD 대비 교사 분포와의 정렬 (Forward KL) 이 우수함을 확인했습니다.
• 외부 도메인 (Out-of-Domain) 성능: 수학 데이터로만 학습되었음에도 불구하고, GPQA-Diamond, MMLU-Pro 등 일반 추론 및 지시 따르기 벤치마크에서도 OPD 보다 우수한 성능을 보였습니다.
• 비교 실험: 단순 엔트로피 보너스 (Entropy Bonus) 나 어드밴티지 셰이핑 (Advantage Shaping) 과 같은 다른 엔트로피 기반 방법들보다 EOPD 가 교사 모델의 불확실성을 더 정확하게 전달하여 더 높은 성능을 달성했습니다.

5. 의의 및 결론 (Significance)

이 논문은 교사 모델의 불확실성 (Uncertainty) 을 명시적으로 모델링하는 것이 지식 증류의 성공에 필수적임을 보여줍니다.

• 기술적 의의: 기존의 역 KL 중심 접근법이 가진 '다양성 붕괴' 문제를 해결하여, 효율성 (온-폴리시 학습) 과 다양성 (고엔트로피 영역 보존) 사이의 균형을 성공적으로 찾았습니다.
• 실용적 가치: 복잡한 추론 작업 (수학, 논리 등) 에서 여러 가지 유효한 해결책이 존재할 때, 학생 모델이 이를 망각하지 않고 학습할 수 있게 함으로써 더 강력하고 안정적인 소형 언어 모델을 구축할 수 있는 길을 열었습니다.
• 환경적 영향: 대형 모델의 추론 능력을 효율적인 소형 모델로 전달함으로써, 대규모 모델 배포에 따른 계산 및 환경 비용을 줄이는 데 기여할 수 있습니다.

요약하자면, EOPD는 교사 모델이 "무엇을 확신하는지"뿐만 아니라 "어디서 불확실한지"를 인지하고 이에 맞춰 학습 전략을 조정함으로써, 기존 증류 방법의 한계를 극복하고 추론 능력을 효과적으로 전이하는 새로운 패러다임을 제시합니다.