On-Policy Distillation 분석: 어디에 도움이 되고 어디에 해를 주며 그 이유

On-Policy Distillation은 reasoning 모델 학습에 dense한 토큰 단위 감독을 제공하지만, 어떤 교사 구성과 맥락이 유용한지에 대한 일반화된 지침이 부재하다. 본 연구는 per-token, per-question, per-teacher 차원의 진단 프레임워크를 제시해 토큰별로 교사 신호의 downstream 효과를 정량적으로 평가한다.

1. 문제 정의: On-Policy Distillation은 토큰 단위로 교사 신호를 제공하지만 이 신호가 언제 유효한지 판단하기 어렵다. 교사 분포의 차이가 실제로 좋은 업데이트를 이끄는지 확인하려면 토큰 수준의 downstream 효과를 확인해야 한다. 2) 해결 원리: 각 토큰 위치에서의 ideal gradient를 empirical success 확률로 정의하고, softmax Jacobian을 통해 ∂Lideal/∂zj를 얻는다. 이 방향이 학생의 성공 확률을 최대화하는 방향이다. 3) 주요 발견: Gradients는 대부분의 경우 토큰 단위로 국지적으로 다르게 작용하며, 잘못된 롤아웃에서 교사 신호의 정합도가 더 높다. 또한 컨텍스트 설계와 학생의 용량 간의 상호작용이 커서 하나의 고정된 파이프라인으로 최적 구성을 얻기 힘들다. 4) 기여의 공통점: GKD, 단일-샘플 GKD, MiniLLM 등은 모두 fj − f̄ 형태의 로컬 구조를 가지므로, 코사인 유사도로 비교가 가능하고, alignment를 극대화하는 방향으로 교육 신호를 조정할 수 있다. 5) 한계 및 방향: 신호의 유용성은 토큰별로 크게 다르고, 높은 다변성은 compute 비용 증가와 함께 나타난다. 이 프레임워크는 토큰-수준 진단을 위한 일반 도구로 확장 가능하다.

1. 전체 접근: 학생 롤아웃에서 프레이밍된 생성 트리를 구성하고, 노드 u에서 Su에 속하는 토큰들에 대해 Psucc를 추정한다. 이 성공 확률로 ideal gradient를 정의하고, gideal_u를 계산한다. 2) Distillation gradients: GKD는 ∂LGKD/∂zj = Pjθ(ℓj − ¯ℓ)로, ¯ℓ = KL(πθ ∥ πte). 싱글-샘플 GKD는 E[∇zjℓ]를 동일 방향으로 수렴시키고, MiniLLM은 롤아웃의 reward-to-go를 연결하는 경향을 가진다. Dr. GRPO는 길이 정규화 제거 및 중요도 비율의 마진화를 통해, per-node gradient가 ideal gradient에 근접하도록 만든다. 3) align 정의: Align(u) = cos(gideal_u, gDu)로 정의하며, Su에서만 계산하고, −1~+1의 범위를 가진다. 4) 대규모 계산: 노드의 샘플링이 충분히 이루어지도록 타깃 롤아웃을 이용하고, depth 윈도우를 지수적으로 증가시키며 예산에 따라 분배한다. 5) 다-교사 파이프라인: Phase 1에서 G initial rollouts를 공유 트리에 저장하고, Phase 2에서 각 교사마다 한 번의 forward pass로 Pk_te를 계산하고 Align를 구한다.

1. 메인 벤치마크 결과: gradient alignment는 BoolQ에서 0.6B 모델의 경우 평균 약 +0.027, 1.7B에서도 +0.026으로 나타났고, per-token 분산은 대략 0.83~0.91 수준으로 크다. 2) 잘못된 경로에서의 유의미한 이득: 4.1에서 1.7B의 BoolQ에서 incorrect path에서의 alignment 차이가 가장 크며, Δ = −0.056(p<1e−9). MMLU에서는 차이가 작아도 weighted cosine에서 유의미한 차이가 나타난다(p<1e−5). 3) 학생 용량에 따른 최적 교사 구성이 달라짐: 0.6B에서 Self-distillation(정확한 데모 중심) 계열이 2–3× 더 높은 정합을 보였고, 1.7B에서는 External 교사(Qwen3-8B 등)가 BoolQ에서 가장 높은 정합을 보였다. 4) 경로 내부의 편차와 정합의 예측: 경로 내 분포 간의 divergence(KL, JS, L2)가 정합과 양의 상관을 보이나 상관도는 작다( |ρ|≈0.02–0.04 ). 5) AIME 2025 케이스 스터디: 짧은 추론보다 길고 어려운 수학 문제에서는 잘못된 데모를 포함한 Self-Sum-1C1W가 더 강력한 교사로 작용하는 경우가 나타났고, 요약 컨텍스트는 0.6B에서만 이득을 주고 큰 모델에서는 Raw demos가 더 나을 수 있다.

1. 아키텍처/개념: 생성 트리는 Node u에서 토큰의 next-token 선택과 downstream 성과를 추적하는 트리로 표현한다. Su에 속하는 토큰들에 대해 Psucc를 추정하고, 이에 따라 ideal gradient를 계산한다. 2) 수학적 기반: Lideal(u) = Σk Pkθ succ. ∂Lideal/∂zj = Pjθ Pj succ − Pjθ Psucc̄. 3) 교사-학습 기법 간 차별점: LGKD(u) = KL(πθ ∥ πte)이며 ∂LGKD/∂zj = Pjθ(ℓj − ¯ℓ). Thinking Lab의 Single-sample GKD는 이 방향을 기대값으로 복원한다. MiniLLM은 reward-to-go를 도입해 fk가 경로 의존적이 되도록 한다. 4) 로컬 구조의 공통성: 모든 방법은 ∂L/∂zj = Pjθ(fj − ¯f) 형태를 가지며, gideal와 gD의 코사인 유사도로 비교 가능하다. 5) 구현 세부: Exponential depth windows, 예산 기반 롤아웃 배치, Phase 1의 공유 트리, Phase 2의 per-teacher forward pass를 통한 신호 계산. 6) 한계와 확장: Su의 방문 수가 충분한 토큰에만 Align를 계산하며, 대규모 시스템에서 연산량이 상당하다.

토큰별로 관측가능한 Su의 범위 밖 토큰은 평가 불가이며, targeted rollouts를 위한 계산 비용이 크다. 오프라인 진단이므로 실시간 학습 신호로 바로 적용하기에는 제약이 있다.