LLM을 위한 하이브리드 정책 증류

대형 언어 모델을 경량화하는 지식 증류 과정에서 발생하는 성능 저하와 학습 불안정성 문제를 해결합니다. 기존의 여러 증류 기법을 하나의 통합된 시각으로 재정의하고, 상황에 맞는 최적의 가중치를 부여함으로써 적은 비용으로도 고성능의 소형 모델을 제작할 수 있게 합니다.

지식 증류는 거대한 교사 모델의 지식을 작은 학생 모델에게 전달하는 과정이다. 기존 방식인 Forward KL은 교사 모델이 가능성이 있다고 보는 모든 선택지를 학생이 다 배우게 하려다 보니 학생 모델의 용량을 초과해 출력이 뭉개지는 현상이 발생한다. 반면 Reverse KL은 학생이 잘 아는 부분에만 집중하게 하여 선명한 결과를 내지만, 교사와 학생의 격차가 크면 학습이 매우 불안정해지는 단점이 있다.

HPD는 이 두 방식의 장점을 토큰 단위에서 조절한다. 학생 모델이 정답 토큰의 확률을 너무 낮게 예측하면(Under-estimation) Forward KL의 원리를 이용해 확률을 끌어올리고, 반대로 엉뚱한 토큰을 너무 높게 예측하면(Over-estimation) Reverse KL의 원리를 이용해 해당 확률을 억제하고 정답 쪽으로 에너지를 재분배한다.

결과적으로 학생 모델은 교사 모델의 전체적인 분포를 놓치지 않으면서도, 자신이 감당할 수 있는 범위 내에서 가장 정확한 답을 찾는 능력을 갖추게 된다. 이는 특히 수학 문제 풀이처럼 긴 사고 과정이 필요한 작업에서 학생 모델이 길을 잃지 않고 정답에 도달하도록 돕는 결정적인 역할을 한다.

기존의 지식 증류 목적 함수들을 L(θ) = -E[w(at|st) log qθ(at|st)] 형태의 통합 수식으로 정의한다. 여기서 w(at|st)는 각 토큰에 부여되는 가중치이며, qθ는 학생 모델의 정책이다. [입력된 상태 st와 선택된 행동 at에 대해] → [교사와 학생의 확률 분포 차이를 계산하여 가중치 w를 결정하고] → [이 가중치를 로그 가능도에 곱해 손실 함수를 구성하며] → [이 값이 작아지도록 모델을 학습시켜 교사의 분포를 효과적으로 모사한다].

HPD는 두 가지 가중치 wt(전문가 토큰용)와 wt(학생 샘플링 토큰용)를 사용한다. 전문가 토큰 가중치 wt는 k1(Reverse KL의 근사치)이 양수일 때 p(a*t|st) + k1을 부여하여 학습을 강화하고, k1이 음수일 때는 k1만 사용하여 과도한 학습을 방지한다. [교사가 정답이라 한 토큰의 확률 차이를 입력으로] → [학생이 정답을 과소평가하면 가중치를 높이고 과대평가하면 낮추는 연산을 수행해] → [최종 전문가 가중치를 얻고] → [모델이 정답에 더 확신을 갖게 유도한다].

학생이 직접 샘플링한 비전문가 토큰에 대해서는 wt 가중치를 적용한다. 학생이 샘플링한 토큰이 교사의 의도와 다를 경우(k'1 < 0), 해당 토큰의 확률을 억제하는 음수 가중치를 부여한다. [학생이 스스로 생성한 토큰을 입력으로] → [교사의 확률 분포와 대조하여 부적절함을 수치화하고] → [음수 가중치를 결과로 산출하여] → [잘못된 생성 패턴을 억제하고 확률 질량을 정답 토큰으로 재분배한다].

수학적 추론 벤치마크에서 Qwen 2.5 3B 모델에 HPD를 적용한 결과, 평균 점수가 28.25에서 39.83으로 약 41% 향상되었다. LLaMA 3 3B 모델의 경우 19.43에서 34.56으로 77.9%라는 비약적인 성능 향상을 기록하며 기존 SFT나 RKLD 기반 증류 방식보다 월등한 성과를 보였다.

대화(Personalization) 및 코드 생성 작업에서도 HPD는 일관되게 우수한 성능을 나타냈다. AlpacaEval2 LC 기준 Qwen 2.5 7B에서 1.5B로의 증류 실험에서 HPD는 13.75%를 기록하여 SFT(12.74%)나 JSD(13.48%)를 앞질렀으며, 특히 다회차 대화의 문맥 유지 능력을 측정하는 MT-Bench에서 가장 높은 점수를 획득했다.

Ablation Study를 통해 학생 샘플링(Student Sampling)과 강화 연산(Reinforce Operation)의 중요성을 확인했다. 학생 샘플링을 제거했을 때는 모델이 교사의 분포에만 갇혀 성능 개선이 조기에 멈췄으며, 강화 연산을 제거했을 때는 KL Loss의 감소 속도가 현저히 느려져 학습 효율이 떨어짐이 입증됐다.

HPD는 Forward KL(FKLD)과 Reverse KL(RKLD)의 상호보완적인 유도 편향(Inductive Bias)을 활용한다. FKLD는 교사의 모든 모드를 덮으려는 성질이 있어 분포를 부드럽게 만들고, RKLD는 주요 모드에 집중하여 선명한 생성을 돕는다. HPD는 이를 토큰 수준에서 마스킹 메커니즘을 통해 결합하여, 학생이 이미 잘 알고 있는 부분에서는 불필요한 FKLD 업데이트를 차단하고 부족한 부분에만 집중하게 한다.

수학적으로는 K1 추정기(Estimator)를 보상 신호로 활용한다. K1 = qθ(a|s) * [log p(a|s) - log qθ(a|s)]로 정의되며, 이는 RKLD의 비편향 추정치 역할을 한다. HPD는 이 K1 값을 기반으로 전문가 토큰과 학생 샘플링 토큰에 대해 비대칭적인 가중치를 할당함으로써, 온-폴리시(On-policy) 증류의 높은 계산 비용 없이도 오프-폴리시(Off-policy) 데이터에서 온-폴리시와 유사한 정렬 효과를 얻는다.

학습 과정에서는 2k step 정도의 짧은 파인튜닝만으로도 유의미한 성능 향상을 이끌어낸다. 이는 HPD가 제공하는 토큰 수준의 세밀한 가중치 신호가 단순한 원-핫 타겟(One-hot target)을 사용하는 SFT보다 훨씬 밀도 높은 정보를 제공하기 때문이다. 또한, 테스트 시점의 엔트로피 붕괴(Entropy Collapse)를 방지하여 모델이 더 다양하고 일관된 답변을 생성하도록 돕는다.

교사 모델과 학생 모델이 동일한 토크나이저를 공유해야 한다는 전제가 필요하다. 또한 전체 KL 다이버전스를 직접 계산하는 대신 추정기를 사용하므로, 추정 오차에 따른 잠재적 불안정성이 존재할 수 있으나 실험적으로는 안정성이 확인되었다.