On-Policy Distillation의 다양한 얼굴: 함정, 메커니즘, 그리고 개선

OPD/OPSD는 대형언어모델의 추가 학습 시나리오에서 교사 정보를 토큰 수준으로 활용하는 강력한 방법이다. 그러나 교사-학생 간 분포 불일치, Top-K 역KL의 편향적 그래디언트, instance-specific PI의 한계 등으로 안정성과 효과가 저하될 수 있다. 이 연구는 이러한 실패 메커니즘을 정밀 분석하고, stop-gradient Top-K KL, RLVR 기반 교사, SFT 안정화 등 구체적 해결책을 제시한다. 결과적으로 PI의 구조가 공유 잠재 규칙으로 작동하는 경우 OPSD가 효과적일 수 있으며, 반대로 PI가 인스턴스 특이적일 때 OPSD의 한계가 명확해진다.

단계1: OPD는 교사(policy_T)로부터 학생(policy_θ)이 샘플링한 on-policy 시퀀스에 토큰 수준의 감독을 받도록 한다. OPD에서 교사는 외부 모델이 될 수 있으며 PI가 optional하다. OPSD는 학생 자체를 교사로 삼되 PI를 통해 교사를 조건화한다. 이때 토큰-수준 KL 손실을 통해 교사-학생 분포를 맞추는데, 역KL은 모드-추구적으로 작동하여 교사의 선호 분기를 보존하려 한다. 단계2: TopK 역KL은 전체 어휘를 사용하는 경우와 다르게, TopK 집합의 확률 mass를 부분적으로만 비교하기 때문에 편향이 생길 수 있다. stop-gradient를 도입하면 이 편향이 제거되거나 감소한다. 또한 OPSD는 PI-conditioned 교사들의 모든 instance에서의 공통된 PI-free 정책으로 수렴하려 하지만, PI가 instance-specific일 때 이 수렴은 성능 저하를 낳는다. 단계3: RLVR로 교사를 트레이닝 데이터 분포에 맞춰 개선하면 교사-학생 분포 간 차이를 줄일 수 있다. SFT를 통해 학생의 출력 포맷과 길이 안정성을 보장하면 이후 OPD의 효과를 더욱 높일 수 있다. 단계4: PI의 구조에 따라 OPSD의 효과가 달라지므로, 시스템 프롬프트 내부화나 alignment와 같이 공유 잠재 규칙이 있는 PI 구조에서는 OPSD가 유효할 수 있다.

단락 1: 전체 접근 방식과 핵심 아이디어 → OPD/OPSD의 기본 정의와 Loss 설계(Full-vocabulary KL, Reverse KL, Forward KL, Sampled-token KL) 및 TopK 근사. [Pattern: 입력 x, y<t → 교사·학생의 로짓 분포를 비교하는 손실을 계산한다.]

단락 1: 수학 추론에서 OPSD는 PI 조건부에서도 안정적 개선을 보이지 못했고, Qwen3-1.7B 학생의 경우 Math500/AIME24/AIME25에서 성능 저하를 보였다. OPD는 초기 개선을 보이다가 이후 롤아웃 길이가 증가하며 반복 토큰이 늘고, step 700~1000에서 언어 반복 및 응답 길이 증가로 정확도가 크게 하락했다. 단, RLVR 어댑티드 교사는 블렌드된 설정에서 성능이 개선될 수 있음을 시사했다. 단락 2: Alignment 및 System Prompt Internalization에서 OPSD는 초기 학습에서 빠르게 개선되지만 최종 성능은 교사에 의해 제한된다. System prompt internalization은 응답 길이를 크게 줄이고, reasoning 정리에서 OPSd가 길이 관리를 도와 효율을 높일 수 있음을 보였다. 단락 3: PI 구조의 중요성 분석에서 PI가 문제-특정일 때 OPSD의 효과가 감소하는 경향을 확인했다. 단계4: 제시된 안정화 기법들(stop-gradient Top-K KL, RLVR 교사, SFT 안정화)이 OPD/OPSD의 안정성을 크게 개선한다는 점이 확인되었다.

단락1: OPD/OPSD의 기본 아키텍처는 x 입력에 대해 학생 정책 πθ와 교사 정책 πT를 두고, y<t 단계에서 토큰 stograd(πT)와 비교하는 로스 ℓt를 적용한다. Full-vocabulary KL은 pθ(v|x,y<t)와 pT(v|x,y<t,I) 간 뒤집힌 KL을 사용한다. 역KL은 모드-추구적 성향을 가지며, forward KL은 모드-커버링 특성을 가진다. 분모는 로그 우도 및 로그 확률 차이를 포함하는 형태로 나타난다.

본 연구는 소수의 모델 계열/스케일에서의 실험에 기초하며, 더 큰 규모에서의 행동은 다를 수 있다.