💬 NLP

Hybrid Policy Distillation for LLMs

이 논문은 기존 지식 증류 방법들을 토큰 수준의 재가중 로그 가능도 목적함수로 통합하여 해석하고, 정방향 및 역방향 KL 발산을 결합하고 오프-정책 데이터와 경량 온-정책 샘플링을 융합한 '하이브리드 정책 증류 (HPD)'를 제안하여 다양한 LLM 작업에서 최적화 안정성과 성능을 향상시켰음을 보여줍니다.

원저자: Wenhong Zhu, Ruobing Xie, Rui Wang, Pengfei Liu

게시일 2026-04-23

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Wenhong Zhu, Ruobing Xie, Rui Wang, Pengfei Liu

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

🎓 비유: 거장 요리사 (선생님) vs. 초보 요리사 (학생)

상상해 보세요. 세계 최고의 요리사 (거대 AI, Teacher) 가 있습니다. 이 분은 어떤 재료를 넣어도 완벽한 요리를 만들어내죠. 하지만 이 분은 너무 비싸고, 주방이 너무 커서 모든 식당에서 쓸 수 없습니다.

그래서 우리는 **초보 요리사 (작은 AI, Student)**를 키워서 그 분의 실력을 대신하게 하려고 합니다. 이때 중요한 건 **"어떻게 가르칠 것인가"**입니다.

기존의 방법들은 두 가지 극단으로 나뉩니다:

단순 암기 (SFT): "이 요리는 소금 1 스푼, 후추 1 스푼!"이라고 정답만 외우게 하는 방식입니다.
- 문제: 학생은 정답만 외우지, 왜 그 양을 넣었는지 이해하지 못합니다. 만약 재료가 조금 부족하면 당황해서 요리를 망칩니다. (너무 딱딱함)
완전 모방 (KL Divergence): 거장 요리사가 만든 모든 레시피와 맛의 미세한 차이까지 다 분석하게 하는 방식입니다.
- 문제: 학생이 감당할 수 없는 양의 정보를 받아서 머리가 터질 수도 있고 (계산 비용 과다), 혹은 모든 맛을 다 섞으려다 보니 요리의 맛이 뭉개져서 평범해집니다. (너무 흐릿함)

💡 새로운 방법: HPD (하이브리드 정책 증류)

이 논문에서 제안한 **HPD(Hybrid Policy Distillation)**는 이 두 가지의 장점을 섞은 '현명한 코치' 같은 방법입니다.

1. "맞는 답"과 "틀린 답"을 동시에 가르칩니다.

기존 방법들은 보통 '정답'만 보여주거나, '학생이 만든 답'만 교정했습니다. 하지만 HPD 코치는 이렇게 합니다:

정답 (Expert Token): 거장 요리사의 레시피를 보여주며 "이건 이렇게 해야 해!"라고 가르칩니다. (Forward KL: 다양한 가능성을 열어둠)
실수 (Student Sampling): 학생이 직접 요리를 해보게 한 뒤, "이건 너무 짜네!"라고 지적합니다. (Reverse KL: 잘못된 길을 막음)

2. "무엇을 강조할지"를 상황에 따라 바꿉니다.

이게 HPD 의 가장 멋진 부분입니다. 코치는 학생의 상태를 보고 **가중치 (중요도)**를 실시간으로 조절합니다.

학생이 정답을 너무 낮게 평가할 때 (예: "소금 1 스푼이 아니라 2 스푼이겠지?"): "아니야, 소금 1 스푼이 정답이야!"라고 강하게 가르칩니다.
학생이 잘못된 답을 너무 높게 평가할 때 (예: "소금 10 스푼이 좋겠어!"): "그건 너무 짜! 그 맛을 줄이고 정답인 소금 1 스푼 쪽으로 집중해!"라고 억제합니다.

즉, 정답을 더 잘 찾게 유도하면서도, 학생이 엉뚱한 길로 빠지지 않게 막는 두 마리 토끼를 다 잡는 방식입니다.

🚀 왜 이 방법이 특별한가요?

계산이 훨씬 가볍습니다: 거장 요리사의 모든 레시피를 다 분석할 필요 없이, 중요한 부분만 골라 가르쳐도 됩니다. (기존 방법보다 빠르고 저렴함)
안정적입니다: 학생이 혼란스러워하지 않고, 차근차근 실력을 키워갑니다.
다양한 분야에서 잘 작동합니다:
- 수학 문제 풀이: 복잡한 논리도 잘 따라갑니다.
- 대화: 자연스러운 말투를 유지합니다.
- 코딩: 정확한 프로그램을 작성합니다.

📊 실제 결과 (시험 성적)

실험 결과, HPD 를 쓴 작은 모델들은 기존 방법들보다 훨씬 좋은 성적을 냈습니다.

특히 **30 억 파라미터 (작은 모델)**가 기존 70 억~80 억 파라미터 (큰 모델) 못지않은 수학 실력을 보여주기도 했습니다.
마치 초보 요리사가 코치의 현명한 지도를 받아, 거장 요리사의 맛을 거의 그대로 재현해낸 것과 같습니다.

🏁 결론

이 논문은 "AI 를 작게 만들 때, 단순히 정답만 외우게 하거나, 모든 것을 다 가르치려 하지 말고, 학생이 어디를 잘못 알고 있는지 파악해서 정답을 강조하고 오답을 억제하는 '하이브리드' 방식으로 가르쳐야 한다"는 것을 증명했습니다.

이 방법은 앞으로 우리가 더 작고 빠르면서도 똑똑한 AI 를 만드는 데 큰 도움이 될 것입니다.

1. 문제 정의 (Problem)

LLM 을 경량화하기 위한 지식 증류는 성공적인 배포에 필수적이지만, 기존 방법들은 다음과 같은 한계를 가집니다:

발산 방향의 상충 (Divergence Trade-off):
- Forward KL (FKLD): 교사의 분포를 모두 커버하려 하지만, 학생 모델의 용량이 부족할 경우 예측이 과도하게 평활화 (over-smoothed) 되어 생성 품질이 떨어집니다.
- Reverse KL (RKLD): 교사의 주요 모드 (high-probability modes) 에 집중하지만, 학생과 교사의 격차가 크면 학습이 불안정해지고 저확률이지만 유효한 출력을 무시할 수 있습니다.
데이터 및 최적화 전략의 복잡성:
- Off-policy (SFT 등): 외부 데이터에 의존하지만 훈련 - 추론 행동 불일치 (train-inference mismatch) 가 발생할 수 있습니다.
- On-policy (OPD 등): 학생이 생성한 데이터를 사용하지만, 계산 비용이 매우 크고 분포 이동 (distribution shift) 으로 인한 불안정성이 있습니다.
현재의 과제: 이 세 가지 요소 (발산 방향, 최적화 전략, 데이터 regime) 를 어떻게 조화롭게 결합하여 안정적이고 효율적인 증류를 수행할 것인가?

2. 방법론: 하이브리드 정책 증류 (HPD)

저자들은 증류 과정을 토큰 수준의 가중 로그 가능도 (reweighted log-likelihood) 문제로 재정의하고, 이를 기반으로 HPD 를 제안합니다.

핵심 아이디어

HPD 는 Forward KL 과 Reverse KL 의 상호 보완적 이점을 통합하며, 오프-폴리시 (off-policy) 데이터와 경량 온-폴리시 (on-policy) 샘플링을 동시에 활용합니다.

통일된 관점 (Unified View):
- SFT, FKLD, RKLD 를 모두 가중치 $w(a_t|s_t)$ 가 적용된 로그 가능도 손실 함수로 통합하여 설명합니다.
- 양의 가중치는 해당 토큰의 확률을 높이고, 음의 가중치는 확률을 억제하여 다른 토큰으로 확률 질을 재분배합니다.
HPD 의 알고리즘적 구조:
- 전문가 토큰 (Expert Token, $a^*_t$ ): 교사 모델이 생성한 정답 토큰.
- 샘플링된 토큰 (Sampled Token, $a_t$ ): 학생 모델이 현재 정책으로 샘플링한 비전문가 토큰.
- 가중치 계산 (Negative K1 Estimator 활용):
  - 전문가 토큰: 학생이 전문가 토큰을 과소평가 ( $k_1 > 0$ ) 하면 Forward KL 가중치를 적용하여 보정합니다. 반대로 과대평가 ( $k_1 \le 0$ ) 하면 가중치를 억제하여 불필요한 학습을 방지합니다.
  - 샘플링된 토큰: 학생이 비전문가 토큰을 과대평가 ( $k'_1 < 0$ ) 하면, 해당 토큰의 확률을 억제 (음의 가중치) 하고 그 확률 질을 전문가 토큰으로 재분배합니다.
- 강화 (Reinforcement): 학생이 비합리적인 샘플링을 할 때, 전문가 토큰의 가중치를 추가로 강화 (Forward KL 가중치 2 배 등) 하여 학습을 가속화합니다.
손실 함수:
- 오프라인 데이터의 전문가 토큰과 학생이 샘플링한 토큰에 대해 각각 계산된 가중치를 적용한 재가중 로그 가능도 손실을 최소화합니다.
- 장점: 전체 시퀀스 롤아웃 (full rollout) 을 수행하지 않고도, 토큰 단위의 샘플링만으로 온-폴리시 효과를 얻어 계산 비용을 크게 절감합니다.

3. 주요 기여 (Key Contributions)

통일된 증류 관점 제시: 발산 방향, 최적화 전략, 데이터 regime 을 통합하여 기존 증류 방법들을 재해석하고 새로운 접근법을 체계적으로 개발할 수 있는 기반을 마련했습니다.
HPD 알고리즘 제안: Forward/Reverse KL 의 이점을 하이브리드화하고, 오프-폴리시 데이터와 경량 온-폴리시 샘플링을 결합하여 SFT 의 효율성과 OPD 의 안정성을 동시에 달성했습니다.
실험적 검증: 수학 추론 (긴 생성), 대화 (짧은 생성), 코드 생성 등 다양한 태스크와 모델 규모 (1B~8B) 에서 HPD 가 기존 방법 (SFT, SeqKD, RKLD, JSD 등) 보다 우수한 성능과 안정성을 보임을 입증했습니다.

4. 실험 결과 (Results)

저자는 Qwen 2.5 와 LLaMA 3 계열 모델을 사용하여 다양한 벤치마크에서 HPD 를 검증했습니다.

수학 추론 (Off-policy):
- HPD 는 모든 베이스라인을 압도했습니다. 특히 3B 모델이 7B/8B 교사 모델의 추론 능력을 크게 따라잡는 성과를 거두었습니다 (예: Qwen 2.5 3B 성능 41% 향상, LLaMA 3 3B 성능 77.9% 향상).
- SFT 는 학습 중 엔트로피가 급격히 감소 (overfitting) 하는 반면, HPD 는 안정적인 엔트로피를 유지하며 교사 분포에 점진적으로 정렬됩니다.
대화 및 개인화 (Off-policy):
- MT-Bench, AlpacaEval 등 대화 벤치마크에서 HPD 가 가장 높은 점수를 기록했으며, 특히 멀티턴 대화의 일관성과 문맥 이해도가 뛰어났습니다.
코드 생성 (Off-policy):
- HumanEval 및 MBPP 벤치마크에서 HPD 가 가장 높은 평균 점수를 달성하며, 다른 방법들보다 더 안정적이고 견고한 성능을 보였습니다.
온-폴리시 데이터 (On-policy) 및 DPO 초기화:
- HPD 로 초기화된 모델은 이후 OPD 나 DPO (Direct Preference Optimization) 를 수행할 때 더 빠른 수렴과 더 높은 최종 성능을 보여주었습니다. 이는 HPD 가 강력한 초기화 (initialization) 역할을 함을 의미합니다.
학습 동역학:
- HPD 는 SFT 대비 교사 - 학생 간 KL 발산 거리를 더 효과적으로 줄이며, 테스트 시 엔트로피 붕괴를 방지하여 추론 행동의 일관성을 유지합니다.

5. 의의 및 결론 (Significance)

효율성과 성능의 균형: HPD 는 계산 비용이 큰 온-폴리시 학습 (full rollouts) 없이도, 토큰 수준의 경량 샘플링을 통해 온-폴리시 증류의 이점을 얻습니다. 이는 실제 산업 환경에서의 적용 가능성을 높입니다.
안정적인 증류 프레임워크: Forward 와 Reverse KL 의 상충되는 특성을 동적으로 조절하여 학습 불안정성을 해결하고, 모델의 용량 한계 내에서 최적의 성능을 끌어냅니다.
미래 지향성: HPD 는 LLM 의 전처리 (pre-training) 나 중간 학습 (mid-training) 단계에서도 적용 가능할 것으로 기대되며, 향후 더 강력한 소형 모델 개발의 핵심 기술로 자리 잡을 수 있습니다.

결론적으로, 이 논문은 LLM 증류의 복잡한 설계 공간을 단순화하고, **하이브리드 정책 증류 (HPD)**를 통해 안정성, 효율성, 성능을 모두 잡은 새로운 표준을 제시했습니다.