[논문리뷰] Hybrid Policy Distillation for LLMs

링크: 논문 PDF로 바로 열기

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

## 1. Key Terms & Definitions (핵심 용어 및 정의)

• HPD (Hybrid Policy Distillation): Forward KL과 Reverse KL의 장점을 통합하고, 오프라인 데이터와 온폴리시(on-policy) 샘플링을 결합하여 효율적인 지식 증류를 수행하는 제안 방법론입니다.
• FKLD (Forward KL Divergence): 스튜던트 모델이 티처 모델의 모든 모드를 커버하도록 유도하여 모드 커버리지(mode coverage)를 높이는 지표입니다.
• RKLD (Reverse KL Divergence): 스튜던트 모델이 티처 모델의 주요 모드에 집중하게 하여 모드 탐색(mode-seeking) 성능을 향상시키는 지표입니다.
• K1 Estimator: Reverse KL을 효율적으로 근사하기 위해 토큰 수준의 보상 신호를 제공하는 Monte Carlo 근사 기법입니다.
• Reweighted Log-likelihood: SFT, FKLD, RKLD를 단일 목적 함수로 일반화하여 토큰 수준에서 재가중치(reweighting)를 적용하는 최적화 프레임워크입니다.

## 2. Motivation & Problem Statement (연구 배경 및 문제 정의) 본 연구는 LLM 압축 과정에서 발생하는 divergence direction, optimization strategy, data regime 간의 복잡한 상호작용 문제를 해결하고자 합니다. 기존의 SFT는 타겟 토큰에 대한 희소한 감독 신호만을 제공하여 일반화 능력이 제한적이며, FKLD는 과도하게 평활화된(over-smoothed) 예측을, RKLD는 스튜던트-티처 간 격차가 클 때 학습 불안정성을 초래하는 단점이 있습니다. 또한, 기존 On-policy Distillation(OPD)은 높은 컴퓨팅 비용과 티처 측의 분포 이동(distribution shift) 문제를 겪고 있습니다. 이러한 기존 기법들의 한계를 극복하기 위해 저자들은 증류 과정을 reweighted log-likelihood 관점에서 통합 분석하고, 보다 안정적이고 효율적인 HPD를 제안합니다.

## 3. Method & Key Results (제안 방법론 및 핵심 결과) 본 논문은 HPD를 통해 Forward와 Reverse KL의 인덕티브 바이어스를 토큰 수준에서 결합하고, 오프라인 데이터와 함께 가벼운 온폴리시 샘플링을 활용하여 학습 효율성을 극대화합니다. HPD는 K1 추정치를 사용하여 티처-스튜던트 간의 불일치를 측정하고, 토큰 수준의 보상 신호로 재가중치를 계산하여 SFT의 효율성을 유지하면서도 티처의 전체 분포를 학습에 효과적으로 반영합니다. [Table 2]와 [Table 4]에서 확인할 수 있듯이, HPD는 추론 및 코딩 작업 전반에서 기존 SFT, SeqKD, RKLD, JSD 등 다양한 베이스라인 기법을 압도하는 성능을 보였습니다. 특히 Qwen 2.5 3B 모델에서 41.0%, LLaMA 3 3B 모델에서 77.9%의 성능 향상을 기록하며, SFT + OPD 파이프라인과 비교해도 훨씬 적은 비용으로 우월한 성능을 달성하였습니다. [Figure 1]은 HPD가 훈련 과정에서 SFT에서 발생하는 엔트로피 붕괴를 방지하고 KLD 격차를 지속적으로 감소시킴으로써 안정적인 최적화와 성능 향상을 유도함을 입증합니다.

Figure 1 — SFT와 HPD의 학습 동역학 비교를 통한 방법론의 우수성 입증

Table 2 — 추론 작업에서의 오프라인 데이터 증류 성능 비교 결과

Table 4 — 코딩 작업에서의 오프라인 데이터 증류 성능 비교 결과

## 4. Conclusion & Impact (결론 및 시사점) 본 연구는 LLM 지식 증류의 설계 원칙을 reweighted log-likelihood 관점에서 재정립하고, HPD를 통해 증류 효율성과 모델 정렬 능력을 동시에 개선했습니다. HPD는 학습의 안정성을 보장하며, 복잡한 온폴리시 rollout 없이도 우수한 증류 성능을 제공함으로써 산업계와 학계의 LLM 경량화 연구에 실질적인 기여를 합니다. 향후 본 기법은 대규모 LLM 학습 초기 단계나 중간 미세 조정 단계에 적용되어 강력한 초기화 전략으로서 그 효용성이 더욱 확대될 것으로 기대됩니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.