Learning to Foresee: Unveiling the Unlocking Efficiency of On-Policy Distillation

대규모 언어 모델의 포스트-트레이닝에서 On-Policy Distillation(OPD)가 유의미한 속도 이점을 보여 주지만, 그 기저의 파라미터 업데이트 역학은 잘 이해되지 않았다. 본 연구는 OPD의 효율성이 초기 학습 단계에서의 방향 정렬과 모듈 선택의 차별적 업데이트에 의해 유도된다고 관찰하고, 이를 바탕으로 EffOPD를 제안해 학습 시간은 크게 줄이면서 최종 성능은 유지한다. 이는 Dense supervision의 이점이 단순한 업데이트 규모의 증가에 의한 것이 아니라 고효율 방향으로의 빠른 수렴에 의해 좌우된다는 통찰을 제공한다.

출발점: RL과 비교해 OPD의 업데이트는 더 작지만 더 높은 품질의 정보량을 전달한다. 모듈 수준에서 비핵심 업데이트를 억제하고 핵심 모듈에 집중하는 ‘Functional Redundancy Avoidance’가 관찰된다. 업데이트 방향 측면에서, OPD의 주 업데이트 방향은 학습 초기부터 최종 해의 방향과 강하게 정렬되어 있으며, 이는 상위 서브스페이스의 에너지가 빠르게 집중되는 현상으로 나타난다.

중요 기여: Top-k 주 서브스페이스의 방향성이 초기부터 우수하게 작동하며, tail 방향은 큰 노름 비용에 비해 성능 기여가 작다. 이러한 저랭크 집중 현상은 학습 초기에 방향을 잠그고, 이후 학습은 이 방향의 크기만 키우는 방식으로 진행됨을 시사한다.

결과적 시사점: EffOPD는 초기 방향의 안정성을 활용해 외삽(extrapolation)으로 더 빠르게 수렴한다. 1.5B~32B 파라미터 규모에서 평균 3×의 학습 속도 향상을 달성하고, 최종 성능은 유사하게 유지한다. 이로써 OPD의 효율성은 Dense supervision 그 자체보다는 파라미터 업데이트의 방향성 및 공간적 구성이 핵심임을 제시한다.

입력: teacher 모듈러 네트워크와 student 모델(πθ) 간의 온-policy 샘플링. 손실은 reverse KL divergence DKL(pθ(y|x) || p*(y|x))로 정의되며, 토큰 단위로 샘플링된 trajectory에 대해 업데이트 방향을 계산한다. 로컬 근사: zθ(c) = z0(c) + JcΔθ로 선형화하고, DKL은 Fc에 의해 가중된 이차 형태로 근사된다. 전체 목적은 Δθ에 대한 2차 형태의 최소화로, g(Δθ) = AΔθ − b를 따른다. gradient descent는 Δθs+1 = Δθs − η(AΔθs − b)로 진행되며, 해는 Δθ∞ = A−1b로 수렴한다. 모듈 간 상호 결합을 Amn으로 표현하는 블록 행렬 구조를 도출하고, bm = Ec[J⊤c,mFcrc]로 각 모듈의 driving term을 분석한다. RL의 정책그래디언트 업데이트 gRL은 gRL = E[∑t At ∇θ log πθ(yt|ct)]로 나타나며, OPD의 업데이트 방향은 더 집중적이고 저잡음의 업데이트 신호(b, A, Fc를 통한)를 따른다. 차이점은 OPD가 높은 확률의 타깃 방향으로 수렴하는 반면, RL은 고변동성의 방향 탐색을 계속 수행한다.

주요 벤치마크에서 OPD는 RL과 비교해 동일한 성능 향상을 달성하면서 업데이트 노름은 더 작다. Top-1% 서브스페이스가 전체 성능의 95% 이상을 회복할 수 있을 정도로 충분하며, Principal 방향의 집중이 강화된다. Tail 방향은 상대적으로 낮은 한계 기여도와 높은 노름 비용을 동반한다. 1.5B~32B 모델에서 OPD는 더 компакт하고 효과적인 업데이트 패턴을 보여주고, 학습 곡선은 빨리 수렴한다. EffOPD의 도입으로 10단계 수준의 수렴이 가능해지며 vanilla OPD 대비 더 강한 수렴 속도를 달성한다. 실험은 코드 구현과 다양한 데이터셋(MATH500, AIME 등)에서 수행되며, 학습 스케일에 따른 안정성과 속도 이점이 확인된다.

아키텍처: teacher- student 구조의 Transformer 계열 모델에서 OPD(On-Policy Distillation) 학습을 수행한다. 학습 신호는 토큰 단위의 역KL 손실이며, 샘플링은 student-prompt에 의해 온-policy로 진행된다. 수학적 기초는 로컬 근사에 기반하며, 업데이트 벡터 Δθ에 대한 2차 근사식으로 AΔθ−b 형태의 목적 함수와 g(Δθ) = AΔθ−b를 도출한다. A는 ∑c J⊤c Fc Jc의 Gram 행렬이며, Fc는 Teacher 분포의 logit 공간의 피셔 정보 행렬로 정의된다. 업데이트는 학습 초기 단계에서 Dominant 방향에 집중되며, 작은 업데이트 규모로도 최종 해에 가까운 방향 정렬을 달성한다. 모듈 간 상호 영향은 Amn으로 표현되며, bm=Ec[J⊤c,mFc rc]가 작으면 해당 모듈의 업데이트가 억제된다. RL 업데이트와의 차이는 RL이 At(토큰 단위의 Advantage) 기반의 불안정한 방향 탐색에 의존하는 반면, OPD는 rc(teacher-base logit 차이)와 Fc의 가중치를 포함한 고정적이고 저잡음의 신호로 작동한다. 실험 설정: 모델 스케일은 1.5B〜32B, RL 알고리즘으로 PPO/GRPO/DAPO를 포함하며, OPD 학생 모델은 2,048 토큰 입력 길이와 18,432 토큰의 총 예산으로 3개의 트레이닝 에폭을 수행한다. 최소 역전파를 위해 배치 사이즈는 512(FP16)로 설정되며, 프롬프트 길이는 2048, 최대 응답 길이는 16384다.

한계점으로는 다중 턴 에이전트 작업 및 다중 모달 추론과 같은 더 복잡한 설정에서의 적용성 불확실성, 이론적 분석이 지역적(local) 성격을 가지며 글로벌 비선형성 전체를 포착하지 못하는 점이 있다. 또한 EffOPD의 외삽 전략은 초기 방향의 안정성에 의존하므로, 극단적 extrapolation에 의한 성능 저하 가능성이 있다. 향후 연구에서는 distillation objective, evolving on-policy distribution, 및 parameter update의 스펙트럼 진화를 보다 포괄적으로 다루는 이론이 필요하다.