KL for a KL: On-Policy Distillation with Control Variate Baseline

Motivation

• Known: OPD는 sparse reward에 의존하는 RLVR 대비 dense token-level reverse KL 신호를 이용해 빠른 post-training을 가능하게 하며, Qwen3·GLM-5·Nemotron-Cascade2·DeepSeek-V4 등 산업 규모 post-training에서 효과가 입증되었다. 표준 구현은 single-sample Monte Carlo estimator를 사용하는 policy-gradient RL 형태로 최적화된다.
• Gap: OPD는 single-sample estimator의 높은 gradient variance로 인해 학습이 불안정하며, 기존 안정화 기법인 full-vocabulary reverse KL(OPDfull-V)은 계산 비용이 크고, top-k reverse KL(OPDtop-k)은 gradient를 편향시키면서도 성능 향상은 미미하다는 한계가 있다.
• Why: OPD는 대형 언어모델의 수학·과학 추론 post-training에서 RLVR을 대체할 만큼 중요한 방법으로 부상했으나, 안정적 학습을 위한 원칙적이고 저비용인 레시피가 부재했다는 점에서 이를 해결하는 것은 실무적, 이론적으로 중요하다.
• Approach: OPD를 policy-gradient RL로 재해석하여 RL 문헌의 control variate baseline(value function)을 도입해 gradient variance를 줄이면서도 불편성(unbiasedness)을 유지하는 vOPD를 제안한다.

Achievement

1. closed-form value function 도출: OPD의 value function이 student와 teacher 간 per-token negative reverse KL과 동일함을 이론적으로 보이고, 이를 이미 계산된 forward pass 로짓에서 추가 critic이나 추가 rollout 없이 바로 얻을 수 있음을 증명했다.
2. unbiased variance reduction 실현: baseline이 샘플링된 토큰에 독립적이기만 하면 gradient의 unbiasedness가 유지된다는 RL 이론(§3.1, §A.1)을 활용해, single-sample estimator를 유지하면서도 분산을 낮추는 방법을 제시했다.
3. top-k 근사 baseline 제안: baseline 계산을 top-k student 토큰으로 근사해도 unbiasedness가 손상되지 않으며(샘플 토큰과 무관하므로), k값 선택이 성능에 거의 영향을 주지 않음을 실험적으로 확인했다(§4.3, Fig 2).
4. 경험적 성능 개선: Qwen3 및 Olmo-3 계열 4개 모델, MATH500·Minerva Math·AMC23·AIME24/25·SciKnowEval·GPQA-Diamond 6개 벤치마크에서 vOPD가 vanilla OPD 대비 평균 최대 +3%, MATH500에서 최대 +6.2%의 절대 정확도 향상을 보였다.
5. 효율성 검증: OPDtop-k를 큰 폭으로 능가하고 가장 비용이 큰 OPDfull-V와 동등한 성능을 내면서 wall-clock time을 최대 57.7% 절감했다(Fig 3).
6. 안정성 검증: 학습 중 gradient norm이 일관되게 낮아짐을 보이고, vOPD가 불안정을 유발하는 reward 토큰에 대한 regularizer로 작동함을 확인했다(Fig 4, Fig 5).

How

• OPD를 policy-gradient RL 형태로 정식화하고(Eq. 2-3), reward를 rt(ct, yt) = log πT(yt|ct) - log πθ(yt|ct)로 정의
• 기존 안정화 기법인 OPDfull-V(전체 vocabulary에 대한 KL, Eq. 4-5)와 OPDtop-k(top-k support로 제한된 KL, Eq. 6-7)의 비용/편향 트레이드오프를 분석
• RL의 control variate baseline 이론(Eq. 8)을 도입하여 advantage at(ct, yt) = rt(ct, yt) - bt(ct)로 재정의하고, baseline이 action에 독립이면 unbiased하고 variance를 줄임을 이론적으로 증명(§A.1, §A.2)
• 표준 baseline인 value function이 OPD에서 per-token negative reverse KL의 closed form을 가짐을 도출하고, 이를 forward pass에서 추가 연산 없이 계산
• baseline을 top-k student 토큰으로 근사하여 추가 비용 절감(vOPDtop-k)
• Qwen3, Olmo-3 계열 4개 모델과 6개 reasoning benchmark에서 vOPD, OPD, OPDfull-V, OPDtop-k를 비교 평가
• gradient norm과 reward-token 상관관계 분석을 통해 안정성 메커니즘을 실증적으로 검증

Originality

• OPD를 policy-gradient RL의 일반적 프레임워크로 명시적으로 재해석하고, RL의 표준 분산 감소 기법(control variate baseline)을 지식 증류(knowledge distillation) 맥락에 처음으로 체계적으로 적용
• OPD의 value function이 student-teacher 간 negative reverse KL이라는 closed-form 표현을 발견하여, 별도 critic 없이 baseline을 "공짜로" 얻을 수 있다는 통찰 제시
• top-k 근사가 action-independent하다는 성질을 이용해 unbiasedness를 훼손하지 않으면서 계산 비용을 낮추는 방법을 제안함으로써 기존 OPDtop-k의 편향 문제를 우회

Limitation & Further Study

• 실험이 수학·과학 reasoning 도메인과 특정 모델 계열(Qwen3, Olmo-3)에 한정되어 있어, 코드 생성이나 대화형 태스크 등 다른 도메인에 대한 일반화 가능성은 추가 검증이 필요
• top-k 근사의 이론적 unbiasedness는 baseline이 샘플 토큰에 독립적이라는 가정에 의존하는데, 매우 작은 k에서의 근사 품질 저하나 극단적 분포에서의 안정성은 충분히 탐구되지 않음
• teacher 모델의 품질이나 student-teacher 크기 격차가 매우 클 때 value function 추정의 신뢰성에 대한 분석이 부족하며, 후속 연구로 다양한 teacher-student 조합 및 non-reasoning 도메인으로의 확장이 필요

Evaluation

총평: OPD의 불안정성 문제를 RL의 고전적 control variate baseline 이론으로 우아하게 해결하며, closed-form value function이라는 실용적이고 저비용인 해법을 제시한 점에서 이론적 통찰과 실무적 효율성을 동시에 갖춘 견고한 연구이다. 다만 워크숍 논문 특성상 실험 범위가 제한적이므로 더 넓은 도메인과 모델 규모에서의 검증이 후속 과제로 남는다.