Source-linked AI summary

Stable On-Policy Distillation through Adaptive Target Reformulation

Ijun Jang, Jewon Yeom, Juan Yeo, Hyunggyu Lim, Taesup Kim

arXiv:2601.07155v3cs.LGcs.AI

TL;DR

On-policy 지식 증류는 불안정한 최적화와 diversity collapse를 유발하는 teacher-student 분포 간극 때문에 어려움을 겪는다. Veto는 logit space에 중간 target을 구성하며, 실험에서 reasoning, code generation, summarization 전반에 걸쳐 supervised fine-tuning 및 기존 on-policy baseline 대비 일관된 성능 향상을 보였다.

  • 문제

    On-policy 지식 증류는 큰 teacher-student 분포 간극을 안정적으로 처리하지 못해 forward-KL gradient explosion과 reverse-KL mode collapse를 유발한다.

  • 방법

    Veto는 유해한 low-confidence update를 억제하고 diversity와 decisiveness 간 trade-off를 조절하는 intermediate logit-space target을 구성해 distillation objective를 재정식화한다.

  • 결과

    Reasoning, code generation, summarization task 전반에서 Veto는 supervised fine-tuning 및 기존 on-policy baseline을 일관되게 능가했으며, on-policy distillation에서 +4.8% 향상된 39.9% accuracy에 도달했다.

  • 시사점 및 한계

    Veto는 forward-KL distillation을 안정화하면서 reverse-KL regime에서 performance와 diversity의 균형을 맞추는 one objective-level reformulation을 제공한다.

Abstract

from arXiv · show

Knowledge distillation (KD) is a widely adopted technique for transferring knowledge from large language models to smaller student models; however, conventional supervised KD often suffers from a distribution mismatch between training and inference. While on-policy KD approaches attempt to mitigate this issue by learning directly from student-generated outputs, they frequently encounter training instabilities because the distributional gap between the novice student and the expert teacher is often too wide to bridge directly. These challenges manifest as pathological gradients in forward KL objectives or diversity collapse in reverse KL regimes. To address these limitations, we propose Veto, an objective-level reformulation that constructs a geometric bridge in the logit space. Unlike prior methods that mix data samples, Veto creates an intermediate target distribution that promotes alignment between the teacher and the student. By introducing a tunable parameter beta, Veto serves as an Adaptive Gradient Veto that stabilizes optimization by suppressing harmful gradients on low-confidence tokens, while simultaneously acting as a Decisiveness Knob to balance reward-driven performance with output diversity. Extensive experiments across various reasoning and generation tasks demonstrate that Veto consistently outperforms supervised fine-tuning and existing on-policy baselines.

1 서론

전통적인 supervised KD는 teacher의 궤적과 student가 생성한 출력이 달라 발생하는 노출 편향을 겪으며, on-policy 방법은 초보 student가 날카로운 expert 분포를 직접 모방할 때 불안정해진다. Veto는 모델 아키텍처를 변경하지 않고 목적 함수 수준의 logit-space 보간을 통해 최적화를 안정화한다.

  • 동기: Supervised KD는 teacher가 제공한 궤적과 student가 생성한 출력이 불일치해 노출 편향(exposure bias)을 겪으며, 이로 인해 특히 장기 생성에서 자기회귀 성능이 저하된다 [Ranzato et al., 2015; Bengio et al., 2015].On-policy KD는 대신 student가 생성한 출력으로 학습해 훈련과 추론 시 동작을 정렬한다.
  • 동기: 데이터 수준 혼합으로는 목적 함수 수준의 간극(objective-level gap)을 해소할 수 없다. 날카로운 teacher 분포가 초보 student에게 가파른 최적화 절벽을 만들기 때문이다.이러한 실패 양상은 모델 아키텍처나 데이터 생성 전략이 아니라 발산 기하에서 비롯된다.
  • 실패 양상: Forward KL은 student가 거의 0의 확률을 부여하는 teacher 선호 토큰에서 gradient explosions를 일으킬 수 있는 반면, reverse KL은 mode collapse와 다양성 손실을 초래할 수 있다.이처럼 상반된 불안정성이 통합된 target 재정식을 요구한다.
  • Veto: Veto는 logit space에서 teacher와 student 분포를 보간해 geometric bridge를 구성하고, 일치하는 부분을 강조하는 동시에 유해한 저신뢰도 업데이트를 거부한다.이 재정식화는 아키텍처를 수정하지 않고 안정성을 향상한다.
  • 기여: 단일 매개변수화 target(single parameterized target)이 forward- 및 reverse-KL 처리를 통합하며, 이론적으로 adaptive gradient suppression 및 entropy-regularized policy gradients와 연결된다.실험은 reasoning, code generation, summarization을 포괄하며 supervised 및 on-policy baseline 대비 일관된 향상을 보고한다.

2 관련 연구

Knowledge distillation은 높은 용량의 teacher에서 효율적인 student로 지식을 전달해 LLM을 압축하지만, supervised 방식은 학습–추론 불일치에 직면한다. On-policy 방식은 student가 생성한 output을 통해 이 불일치를 완화하며, 이후의 objective들은 행동을 개선했지만 초기 optimization 안정성은 대부분 간과한다.

  • Supervised Knowledge Distillation: Knowledge distillation은 높은 용량의 teacher에서 효율적인 student로 지식을 전달하며, supervised 접근법은 고정된 dataset에서 teacher distribution 또는 sequence output을 모방한다.이러한 접근법에는 고정된 data에 대한 distribution-level training과 sequence-level imitation이 포함된다.
  • On-Policy Knowledge Distillation: On-policy KD는 student가 생성한 output에서 학습해 training을 student inference에 맞추며, ImitKD (Lin et al., 2020)의 고정 sequence와 self-generated sequence 혼합에서 GKD (Agarwal et al., 2024)의 완전한 self-generated distillation으로 발전했다.GKD는 self-generated mistake에서 학습하면 performance를 크게 향상할 수 있다고 보고한다.
  • On-Policy Objectives: MiniLLM (Gu et al., 2023)과 f-distill (Wen et al., 2023)은 reverse KL과 f-divergence를 사용해 mode-seeking behavior를 유도하는 반면, SKD (Xu et al., 2025)는 low-quality generation을 교정하기 위해 interleaved sampling을 사용한다.SKD의 interleaved sampling은 feedback quality를 개선하기 위한 것이다.
  • Open Challenge: 이러한 발전에도 불구하고 기존 on-policy 접근법은 optimization objective의 안정성, 특히 early training 중의 안정성을 대부분 간과한다.이 한계는 training distribution이나 sampling strategy의 선택만이 아니라 objective 자체와 관련된다.

3 예비 지식

Knowledge distillation은 고용량 teacher의 예측 능력을 더 작은 student에게 전달하지만, supervised KD는 학습에 고정된 teacher trajectory를 사용하므로 exposure bias를 유발할 수 있다. 반면 on-policy KD는 student policy에서 샘플링해 학습과 추론을 정렬하지만, 초기 출력이 불안정할 때 KL 방향 선택으로 수치적 문제가 발생한다.

  • Supervised KD는 고정된 teacher trajectory에 대한 divergence를 최소화하므로, autoregressive inference 중 student가 직면하는 distribution shift에 노출된다.이러한 exposure bias는 student가 inference 시점의 자체 상태로 학습되지 않기 때문에 발생한다.
  • On-policy KD는 student policy에서 시퀀스를 샘플링해 학습과 추론을 정렬하고, 가능성 높은 test-time 영역에서 발생하는 오류에 대한 피드백을 제공한다 (Agarwal et al., 2024; Xu et al., 2025).
  • Forward KL은 zero-avoiding인 반면 reverse KL은 mode-seeking이며, 초기 student 출력이 매우 불안정할 때 둘 다 on-policy distillation에서 수치적 문제를 일으킬 수 있다.Forward KL은 teacher가 선호할 가능성이 높은 token에 0이 아닌 확률을 부여하도록 유도하는 반면, reverse KL은 확률 질량을 주요 mode에 집중시킨다.

4 방법론

Veto는 teacher의 품질과 student의 confidence를 기하학적으로 연결하는 intermediate logit-space target으로 on-policy distillation을 재구성한다. Adaptive parameter β는 optimization을 안정화하고 결정을 선명하게 하며, diversity를 제어하면서 reverse-KL distillation과 reinforcement learning을 연결한다.

  • 분석 I: Adaptive Veto를 통한 안정성: Forward KL에서 Veto는 student의 uncertainty를 반영해 pathological gradient explosion을 방지하며, β > 0일 때 loss가 0으로 수렴하게 한다.Student-dependent term은 log PS가 발산하는 것보다 빠르게 감쇠하므로, student가 무지한 token에 대한 update를 차단한다.
  • Veto Objective: Veto는 teacher와 직접 일치시키는 대신 geometric bridge로서 intermediate target distribution Q를 구성해 teacher–student distribution mismatch가 심한 문제를 해결한다.Q(y|x) ∝ exp(zT(y|x) + β·zS(y|x))는 Product of Experts를 형성하며, teacher quality와 student confidence가 모두 뒷받침하는 token에만 높은 probability를 부여한다.
  • 분석 I: Adaptive Veto를 통한 안정성: Veto는 implicit temperature scaling T = 1 − β를 통해 student를 teacher 방향으로 선명하게 하며, 0 ≤ β < 1일 때 더 단호한 output을 유도한다.이 sharpening effect는 student가 teacher의 uncertainty를 단순히 재현하는 대신 teacher보다 더 단호해지게 한다.
  • 분석 II: Reinforcement Learning으로의 연결: Reverse KL에서 Veto는 gradient가 scaled entropy regularization을 적용한 REINFORCE와 동등한 Decisiveness Knob로 작동한다.이는 knowledge distillation과 reinforcement learning을 연결하며, β를 mode-seeking behavior와 diversity를 제어하는 변수로 드러낸다.
  • 결정성의 스펙트럼: β가 0에서 1을 향해 증가할수록 objective는 distribution matching에서 reward-seeking으로 이동하고, entropy regularization이 감소해 최고 reward mode에 수렴할 수 있다.중간 β 값은 (1 − β)에 비례하는 diversity budget을 유지하는 반면, β → 1에서는 REINFORCE objective가 된다.

5 실험

수학적 추론, 코드 생성, 대화 요약 전반에서 Veto는 supervised 및 on-policy distillation baseline보다 student 성능을 향상시킨다. Ablation 결과, Veto는 두 KL objective를 모두 안정화하고 다양한 data-generation strategy를 지원하며, adaptive β scheduling과 더 폭넓은 model-family 적용에서 이점을 얻는 것으로 나타난다.

  • 주요 결과: Veto는 supervised student fine-tuning 대비 GSM8K 정확도를 30.7%에서 39.9%로 향상시키며, SKD 및 on-policy KD를 최대 6.3%p 앞선다.평가는 1,319개 instance로 구성된 GSM8K test set에서 answer accuracy를 사용해 수행된다.
  • 주요 결과: Veto는 on-policy KD 대비 HumanEval Pass@1을 22.9에서 29.0으로 (+6.1), Pass@10을 35.3에서 37.7로 (+2.4) 높인다.코드 생성은 HumanEval에서 pass@k metric으로 평가된다.
  • 주요 결과: Veto는 DialogSum에서 가장 높은 win-rate를 달성하며, supervised KD 및 on-policy KD 대비 54.3에서 56.5로 (+2.2), SKD 대비 53.6%에서 56.5%로 향상시킨다.이러한 향상은 student-training example이 1K개에 불과한 경우에도 나타나 제한된 student data에서도 효과적임을 보여준다.
  • Data Generation Setting: supervised KD, SKD, on-policy KD 전반에서 Veto는 mathematical-reasoning accuracy를 향상시키며, supervised KD에서는 34.3%를 기록해 standard baseline보다 0.9%p 높다.Figure 3은 여러 on-policy data-generation strategy에서 Veto를 평가한다.
  • Adaptability of KL Loss: Veto는 forward 및 reverse KL distillation을 모두 개선한다. forward KL에서는 ignorant token의 pathological gradient를 억제하고, reverse-KL accuracy를 2.3% 향상시킨다 (37.9% →40.2%).이 결과는 Veto가 objective-agnostic reformulation으로 기능함을 뒷받침한다. 즉 안정성을 위한 Adaptive Gradient Veto이자 premature mode collapse를 방지하는 Decisiveness Knob으로 작동한다.
  • β Scheduling and Generalizability: 더 높은 β는 optimization stability를 우선시하는 반면, 더 낮은 β는 student가 숙련도를 얻은 뒤 teacher를 더 밀접하게 따른다. linear decay는 fixed scheduling보다 우수하며, initial β = 0.8 best이다.model-family 확장에서도 β = 0.3에서 GSM8K에 대해 Gemma2-9B-IT에서 distill한 Gemma2-2B-IT의 일관된 향상이 보고된다.

6 결론

Veto는 logit space에서 기하학적 target distribution을 구성해 최적화 안정성을 높이는 on-policy knowledge distillation용 objective-level reformulation이다. forward-KL gradient의 병리적 현상을 억제하는 동시에 reverse-KL regime에서 reward-driven performance와 distributional diversity의 균형을 맞춘다.

  • 6 결론: Veto는 teacher–student agreement를 강조하는 logit space의 기하학적 target distribution을 구성해 on-policy KD의 최적화 안정성을 높인다.이 reformulation은 data sample을 혼합하는 대신 objective level에서 작동한다.
  • 6 결론: Veto는 기하학적 target을 사용해 teacher–student agreement를 촉진함으로써 forward KL objective에서 일반적으로 발생하는 병리적 gradient를 억제한다.
  • 6 결론: reverse KL regime에서 Veto는 reward-driven performance와 distributional diversity의 균형을 맞추는 Decisiveness Knob로 작동한다.

A 상세한 수학적 증명 · A.1 정리 1의 증명 (Adaptive Gradient Veto)

이 증명은 표준 forward KD와 Veto의 β > 0 정식화를 대조하며, 토큰에 대한 student probability가 0에 가까워질 때 Veto가 발산하는 거동을 해결함을 보인다. 그 결과의 메커니즘은 student가 알지 못하는 토큰에 대한 업데이트를 효과적으로 거부한다.

  • A.1 정리 1의 증명 (Adaptive Gradient Veto): β = 0인 표준 forward KD는 loss L ≈ −c log P_S(y)를 가지며, P_S(y) → 0일 때 loss와 gradient가 발산한다.
  • A 상세한 수학적 증명: 따라서 이 정리의 adaptive-gradient 해석은 표준 forward-KD의 발산을 허용하는 대신 student probability가 낮은 토큰에 대한 업데이트를 억제하는 것과 연결된다.
  • A.1 정리 1의 증명 (Adaptive Gradient Veto): β > 0인 Veto는 product-of-experts target Q ∝ P_T P_S^β를 loss에 대입한다.
  • A.1 정리 1의 증명 (Adaptive Gradient Veto): 이 증명은 student probability P_S(y)가 위에서 0에 접근할 때 Veto의 극한 거동을 평가한다.
  • A.1 정리 1의 증명 (Adaptive Gradient Veto): L’Hôpital’s rule을 적용하는 것은 logarithmic term에 대한 극한 분석의 일부다.
  • A.1 정리 1의 증명 (Adaptive Gradient Veto): 이 증명은 student가 알지 못하는 토큰에 대한 업데이트를 Veto가 효과적으로 거부한다고 결론짓는다.

A.2 정리 2의 증명 (Sharpening Effect)

증명은 student가 target distribution과 일치할 때 forward KL이 최소화되며, 0 < β < 1에서 이 target이 sharpened teacher distribution을 유도함을 보인다.

  • A.2 정리 2의 증명 (Sharpening Effect): Forward KL divergence D_KL(Q∥P_S)는 student distribution P_S가 target Q와 일치할 때 최소화된다.증명은 Eq. (4)의 target definition을 대입한 뒤 이 최소화 조건을 적용한다.
  • A.2 정리 2의 증명 (Sharpening Effect): 0 < β < 1이면 γ = 1/(1−β) > 1이므로, student는 teacher distribution의 sharpened version으로 수렴한다.주어진 β 범위에서 지수가 정확히 1을 초과하므로, 정리의 sharpening effect가 성립한다.

A.3 Theorem 3 증명 (REINFORCE로의 연결)

증명은 gradient step 동안 teacher distribution Q를 고정된 것으로 취급하는 reverse KL objective에서 시작한다. 이를 재정식화하면 scaled student entropy cost와 teacher log-probability reward signal이 분리된다.

  • reverse KL objective는 J(θ) = D_KL(P_S∥Q)이며, gradient step에서 Q를 고정된 target으로 취급한다.
  • 재정식화된 항 (1−β) log P_S(y)는 scaled entropy cost로 작용하는 반면, log P_T(y)는 reward signal로 기능한다.
Loading 2601.07155v3…