TIP: 온폴리시 증류에서의 토큰 중요도 분석

대형 언어 모델의 지식을 작은 모델로 옮기는 지식 증류 과정에서 모든 토큰을 동일하게 학습시키는 것은 비효율적이다. 이 논문은 학생 모델이 확신을 가지고 틀리는 '과잉 확신' 토큰이 가장 중요한 학습 신호를 담고 있음을 발견하여, 학습 데이터의 20%만 사용하고도 메모리 사용량을 50% 이상 절감하면서 성능을 높이는 방법을 제시한다.

On-policy Distillation(OPD)은 학생 모델이 직접 생성한 답변에 대해 교사 모델이 토큰별로 피드백을 주며 학습하는 방식이다. 기존에는 주로 학생 모델이 잘 모르는 부분, 즉 Softmax 출력의 확률 분포가 고르지 않은 'High Entropy' 지점이 중요하다고 여겨졌다. 하지만 이는 학생 모델이 '확신을 가지고 틀리는 경우'를 포착하지 못한다는 한계가 있다.

이 논문은 학생 모델의 상태를 두 가지 지표로 측정한다. 첫째는 학생 모델 스스로의 불확실성인 Entropy이고, 둘째는 교사 모델과의 확률 분포 차이인 KL Divergence이다. 단순히 Entropy만 높을 때 학습하는 것이 아니라, Entropy는 낮지만(확신하지만) Divergence는 높은(교사와 다른) 지점을 찾아내는 것이 핵심이다. 이는 마치 학생이 아예 모르는 문제보다, 정답이라고 확신하며 잘못된 개념을 적용하는 순간을 교정해 줄 때 학습 효과가 가장 큰 것과 같은 원리이다.

결과적으로 '불확실한 지점'과 '확신하며 틀린 지점'을 모두 포함하는 Soft-OR 점수를 통해 토큰을 선별한다. 이 방식을 통해 학습에 불필요한 '이미 알고 있는 쉬운 토큰'들을 대거 제거함으로써 연산 효율과 학습 품질을 동시에 잡을 수 있다.

TIP 분류 체계는 학생 엔트로피 $h_t$와 교사-학생 발산 $\delta_t$를 기준으로 토큰을 네 구역으로 나눈다. Q1(High $h_t$, High $\delta_t$)은 가장 밀도 높은 신호를 가진 영역이고, Q2(High $h_t$, Low $\delta_t$)는 예측 안정화에 기여하며, Q3(Low $h_t$, High $\delta_t$)는 학생이 과잉 확신하는 오류 지점이며, Q4(Low $h_t$, Low $\delta_t$)는 이미 해결된 토큰이다.

핵심 선택 규칙인 Soft-OR 점수 $s_t$는 다음과 같이 계산된다. $s_t = \hat{h}_t + \hat{\delta}_t - \hat{h}_t \cdot \hat{\delta}_t = 1 - (1 - \hat{h}_t)(1 - \hat{\delta}_t)$. 여기서 $\hat{h}_t$와 $\hat{\delta}_t$는 각 배치 내에서 최소-최대 정규화(Min-Max Normalization)된 값이다. [0~1 사이의 정규화된 엔트로피와 발산 값을 입력으로] → [두 지표 중 하나라도 크면 높은 점수가 나오도록 논리적 OR 연산의 부드러운 형태를 적용하여] → [최종 중요도 점수 $s_t$를 얻고] → [이 값이 상위 $\rho$ 비율에 해당하는 토큰들만 학습에 사용한다].

학습 손실 함수는 선택된 토큰 집합 $\mathcal{T}$에 대해서만 계산된다. $\mathcal{L}_{TIP} = \frac{1}{|\mathcal{T}|} \sum_{t \in \mathcal{T}} D_{KL}(P_S(\cdot | c_t) \parallel P_T(\cdot | c_t))$. [선택된 토큰 위치에서의 학생과 교사의 확률 분포를 입력으로] → [KL Divergence를 계산하여] → [평균 손실값을 얻고] → [이 값을 최소화하도록 학생 모델의 가중치를 업데이트한다].

MATH-500 벤치마크에서 Qwen3-8B를 4B로 증류할 때, 전체 토큰의 50%만 엔트로피 기반으로 샘플링해도 정확도가 76.7%에서 78.6%로 상승했다. 특히 제안된 Soft-OR 방식을 적용하여 50%를 선택했을 때는 79.1%로 성능이 더욱 향상되었으며, 이는 모든 토큰을 사용한 베이스라인보다 높은 수치이다. Llama-70B에서 8B로의 증류에서도 50% 선택 시 71.0%에서 74.7%로 성능이 개선되었다.

에이전트 계획 벤치마크인 DeepPlanning에서는 더욱 극적인 결과가 나타났다. Q3(과잉 확신) 토큰 20%만 사용하여 학습했을 때, 14B 교사 모델 기준 정확도가 12.6%를 기록하여 전체 토큰 학습(11.7%)을 능가했다. 이는 에이전트 작업에서 단 한 번의 잘못된 확신이 전체 계획을 망칠 수 있기 때문에, 이러한 오류를 교정하는 Q3 토큰의 학습 밀도가 매우 높음을 시사한다.

효율성 측면에서 Qwen3 모델 학습 시 토큰 50% 유지 시 피크 메모리는 72.0GB에서 38.1GB로 약 47% 감소했으며, 10% 유지 시에는 35.3GB까지 떨어져 최대 58%의 메모리 절감 효과를 확인했다. 이는 제한된 GPU 자원 환경에서 더 큰 모델을 증류할 수 있는 실질적인 가능성을 보여준다.

본 연구는 On-Policy Distillation(OPD)에서 토큰별 그래디언트 신호의 세기를 이론적으로 분석했다. Proposition 1을 통해 오라클 토큰 가중치 $w_t^*$가 모집단 그래디언트의 정렬도($\phi_t$)와 에너지($M_t$)의 비율에 비례함을 보였으며, 이를 통해 Q1 > Q2 > Q3 >> Q4 순의 중요도 서열을 도출했다.

특히 엔트로피 기반 선택의 구조적 결함(Proposition 2)을 수학적으로 증명했다. 엔트로피만 고려하는 규칙은 $h_t \approx 0$인 지점에서 함구하게 되는데, 이는 학생이 확신을 가지고 틀리는 Q3 영역을 완전히 무시하게 만든다. Soft-OR 점수는 이러한 'Blind Spot'을 보완하기 위해 설계되었으며, 추가적인 하이퍼파라미터 없이도 오라클 서열을 잘 추종함을 실험적으로 확인했다.

구현 측면에서는 매 배치마다 토큰별 엔트로피와 KL Divergence를 계산한 후, 배치 내에서 Min-Max 정규화를 수행하여 Soft-OR 점수를 산출한다. 이후 Top-K 선택을 통해 상위 토큰만 역전파(Backpropagation)에 참여시킨다. 이 과정에서 발생하는 추가 연산 비용은 $O(m \log m)$ 수준으로, 전체 모델의 순전파/역전파 연산량에 비해 무시할 수 있는 수준이다.

Q3 영역의 토큰을 감지하기 위해서는 교사 모델의 전체 출력 분포가 필요하므로, 교사 모델의 로짓(Logit)에 접근할 수 없는 API 기반 증류 환경에서는 적용이 어렵다. 또한 배치 단위의 Min-Max 정규화가 이상치(Outlier)에 민감할 수 있다는 잠재적 문제가 있으며, Forward KL이나 JSD 등 다른 거리 함수를 사용할 때도 동일한 사분면 논리가 성립하는지는 추가 연구가 필요하다.