[논문리뷰] Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes

링크: 논문 PDF로 바로 열기

Total number of browse results: 1

The paper is titled "Revisiting On-Policy Distillation: Empirical Failure Modes and Simple Fixes". Authors: Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Yuanheng Zhu, Dongbin Zhao. Affiliations: State Key Laboratory of Multimodal Artificial Intelligence Systems, CASIA; School of Artificial Intelligence, UCAS.

I need to extract 5-8 keywords. From the abstract and introduction, relevant keywords appear to be: On-policy distillation (OPD), large language model (LLM), post-training, sampled-token, variance, local support matching, truncated reverse-KL, top-pp sampling, special-token masking. I will select a concise set.

Now I will structure the summary following the given format.

Part 1: 요약 본문

저자: Yuqian Fu, Haohuan Huang, Kaiwen Jiang, Yuanheng Zhu, Dongbin Zhao et al.

---

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

• On-policy Distillation (OPD) : Student model이 자체 생성한 Rollout에 대해 Teacher feedback을 평가하며 학습하는 방식. Fixed teacher traces 대신 student-generated rollouts를 사용하여 long-horizon reasoning 및 agentic post-training에 특히 유용하다.
• Sampled-token OPD : 각 디코딩 단계에서 student가 샘플링한 단일 토큰의 log-ratio를 통해 업데이트되는 일반적인 OPD 구현 방식.
• Sequence-level Reverse-KL : 전체 시퀀스에 대한 Kullback-Leibler (KL) Divergence를 최소화하는 Objective로, 각 토큰 업데이트가 미래 Reward에 Coupling된다.
• Teacher Top-K Local Support Matching : sampled-token supervision 대신, Teacher 모델의 locally plausible support 내에서 Teacher와 Student 분포를 비교하는 방법론.
• Top-pp Rollout Sampling : Rollout 생성 시 확률 분포의 상위 p%에 해당하는 토큰들 중에서 샘플링하여, 비정상적으로 낮은 확률의 토큰 생성을 줄이고 Teacher signal의 신뢰도를 높이는 기법.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

Large Language Model (LLM)의 Post-training에 있어 On-policy Distillation (OPD)은 student-generated rollouts에 대한 teacher feedback을 활용하기 때문에 매력적이다. 그러나 기존의 일반적인 Sampled-token OPD 방식은 long-horizon 환경에서 여러 가지 문제점을 보인다. 첫째, 분포 매칭을 단일 토큰 신호로 축소시켜 highly imbalanced signal을 초래한다. 대부분의 샘플링된 토큰은 Negative Reward를 받으며, Positive 학습 신호는 소수의 토큰에 집중된다 [cite: 1, Figure 2]. 이는 최적화가 소수의 Locally favorable tokens에 불균형적으로 의존하게 만든다. 둘째, student-generated prefixes에 대한 teacher signal의 신뢰성이 낮아질 수 있다. Student가 Teacher에게는 uncommon한 prefixes로 Rollout될 때, Teacher는 겉보기에 그럴듯한 토큰에 높은 확률을 부여할 수 있지만, 전체 Trajectory는 바람직하지 않은 방향으로 이탈할 수 있다 [cite: 1, Figure 3]. 셋째, Tokenizer 및 special-token mismatch로 인해 one-token comparison이 왜곡될 수 있다 [cite: 1, Figure 5]. Student와 Teacher 모델이 다른 Tokenization을 사용할 경우, 동일한 텍스트도 다르게 분할되어 의미적으로는 정확한데도 Teacher로부터 낮은 확률을 받을 수 있다. 이러한 문제점들은 Sampled-token OPD의 brittle함을 야기하며, long-horizon LLM post-training에서 안정적인 학습을 저해한다.

3. Method & Key Results (제안 방법론 및 핵심 결과)

저자들은 Sampled-token OPD의 문제점을 해결하기 위해 Teacher Top-K Local Support Matching 방법론을 제안한다. 이 방법은 단일 샘플링된 토큰에 대한 비교 대신, 각 Prefix에서 Teacher가 정의한 Locally plausible support set 내에서 Teacher와 Student의 분포를 비교한다. 구체적으로, Student inference policy를 사용하여 Outputs {oi}를 샘플링하고, 각 Prefix ci,t에서 Teacher 모델 q의 상위 K개 확률 토큰으로 구성된 Teacher support set S(ci,t)를 정의한다. 이후 이 Local support 내에서 Teacher 및 Student 분포를 정규화하고, Truncated Reverse-KL Objective를 사용하여 학습을 진행한다. 이 Objective는 Full-vocabulary KL보다 계산 비용이 훨씬 저렴하면서도 단일 토큰 보상보다 더 균형 잡힌 학습 신호를 제공한다.

안정적인 학습을 위해 몇 가지 Practical stabilization choices도 적용되었다. Support-set Renormalization은 Truncated support에서 Objective가 평가될 때 Teacher와 Student 분포의 질량을 직접 비교 가능하게 하여 최적화의 불안정성을 방지한다. Top-pp Rollout Sampling은 Rollout이 Typical continuations에 더 가깝게 유지되도록 하여 Teacher signal의 신뢰성을 높인다. 또한 Special-token Masking은 호환되지 않는 Tokenization conventions으로 인한 False negatives를 줄이는 데 기여한다.

실험은 Single-task math reasoning과 Multi-task agentic-plus-math training 두 가지 설정에서 진행되었다. Single-task math reasoning (Table 1)에서 제안하는 Local support matching 방법은 Sampled-token OPD 대비 우수한 성능을 보였다. Sampled-token OPD가 평균 점수를 28.2 에서 36.4 로 향상시켰고, Special-token masking을 적용한 경우 40.7 까지 개선되었지만, 제안하는 Full method는 평균 41.5 를 달성하여 더 강력한 Distillation signal을 입증했다. 이는 Tokenization artifacts 처리뿐만 아니라, 더 강력한 Local distillation signal 덕분임을 시사한다.

Multi-task agentic-plus-math training (Table 2)에서는 ALFWorld 성능을 유지하면서 Math task에서 상당한 개선을 보였다. Unmasked 버전의 제안 방법은 Math500 점수를 Sampled-token OPD의 76.0 에서 82.0 으로, 평균 Math 점수를 36.6 에서 41.7 로 향상시키면서도 ALFWorld에서는 경쟁력 있는 성능을 유지했다. Masked 버전은 ALFWorld에서 97.7 로 최고 성능을 달성했으나 Math task에서는 일부 이득을 포기했다. 전반적으로 Local support matching이 Long-horizon token-level supervision이 가장 brittle한 영역에서 효과적임을 보여준다.

Training dynamics 분석 결과 (Figure 6, Figure 7, Figure 8), 제안 방법은 Training reward와 Final evaluation 성능을 일관되게 개선했으며, 더 작은 Gradient norms와 낮은 Clipping-boundary fraction을 보이며 더 안정적인 Optimization을 나타냈다. 또한 Sampled tokens에 대한 Teacher-student log-probability gap도 감소하여 Teacher-student alignment가 향상되었음을 확인했다. Ablation study (Table 3, Figure 9)는 Support-set renormalization의 필수적인 역할과 Top-pp rollout sampling의 중요성을 강조하며, 제안 방법의 성능이 여러 설계 선택의 시너지 효과임을 보여준다.

4. Conclusion & Impact (결론 및 시사점)

본 논문은 long-horizon post-training 환경에서 On-policy Distillation (OPD)의 핵심적인 Tradeoff를 재조명했다. Sequence-level coupling objective는 근본적인 목적에 가깝지만 Variance가 매우 높고, Sampled-token OPD는 최적화가 쉽지만 신뢰할 수 없는 Supervision을 제공하여 brittle한 경향이 있다. 저자들은 이 두 극단 사이의 중간 지점을 차지하는 Teacher Top-K Local Support Matching 을 제안한다. 이 방법은 Objective의 Local성을 유지하면서 One-token supervision을 Truncated distribution-level comparison으로 대체한다.

실험 결과, Single-task math reasoning과 Alternating agentic-plus-math training 모두에서 제안하는 방법론이 기존 Sampled-token OPD 대비 최적화 동작과 Downstream 성능을 개선함을 입증했다. 이 연구는 LLM post-training에서 더욱 안정적이고 효율적인 On-policy Distillation 기법을 위한 중요한 기반을 제공하며, 향후 Long-horizon reasoning 및 Agentic LLM 개발에 긍정적인 영향을 미칠 것으로 예상된다. 특히, Teacher matching과 Task success 사이의 여전한 Gap은 Rollout drift 제어 및 Teacher uncertainty 활용에 대한 추가 연구의 필요성을 시사한다.

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