교사를 넘어선 학습: 보상 외삽을 활용한 일반화된 온-폴리시 증류

Learning beyond Teacher: Generalized On-Policy Distillation with Reward Extrapolation

학생이 생성한 궤적에 대해 학생 모델을 교사 모델의 로짓 분포와 일치시키는 온-폴리시 증류(On-policy distillation, OPD)는 학생 모델의 성능을 향상시키는 데 있어 강력한 경험적 이점을 입증했으며, 오프-폴리시 증류 및 강화학습(RL) 패러다임보다 뛰어난 성능을 자주 보여준다. 본 연구에서는 먼저 OPD가 보상 함수와 KL 정규화가 항상 동일한 가중치를 가지며 참조 모델이 어떤 모델이든 될 수 있는 밀집(dense) KL-제한 강화학습의 특수한 경우임을 이론적으로 증명한다. 그런 다음, 우리는 유연한 참조 모델과 KL 정규화 대비 보상 항의 상대적 가중치를 제어하는 보상 스케일링 계수를 도입하여 표준 OPD 목적 함수를 확장하는 일반화된 온-폴리시 증류(Generalized On-Policy Distillation, G-OPD) 프레임워크를 제안한다. 수학적 추론 및 코드 생성 작업에 대한 포괄적인 실험을 통해 우리는 다음과 같은 두 가지 새로운 통찰을 도출한다: (1) 보상 스케일링 계수를 1보다 크게 설정하는 것(즉, 보상 외삽), 우리가 ExOPD라고 명명한 이 방법은 다양한 교사-학생 크기 쌍 조합 전반에 걸쳐 표준 OPD보다 일관되게 향상된 성능을 보여준다. 특히 동일한 학생 모델에 도메인 특화 RL을 적용하여 얻은 다양한 도메인 전문가의 지식을 원래 학생 모델로 다시 병합하는 설정에서, ExOPD는 학생 모델이 교사 모델의 성능 한계를 넘어서 도메인 교사 모델들을 능가할 수 있도록 해준다. (2) ExOPD를 바탕으로, 강-약 증류(strong-to-weak distillation) 설정(즉, 더 큰 교사 모델로부터 더 작은 학생 모델로 증류하는 것)에서 RL 적용 전 교사 모델의 기본 모델(base model)을 참조 모델로 선택하여 보상 보정을 수행하는 것이 더 정확한 보상 신호를 산출하고 증류 성능을 더욱 향상시킨다는 점을 추가로 발견했다. 그러나 이러한 선택은 교사 모델의 RL 적용 이전 버전에 대한 접근을 전제로 하며 더 많은 계산적 오버헤드를 발생시킨다. 우리는 본 연구가 향후 OPD 연구에 새로운 통찰력을 제공하기를 기대한다.

On-policy distillation (OPD), which aligns the student with the teacher's logit distribution on student-generated trajectories, has demonstrated strong empirical gains in improving student performance and often outperforms off-policy distillation and reinforcement learning (RL) paradigms. In this work, we first theoretically show that OPD is a special case of dense KL-constrained RL where the reward function and the KL regularization are always weighted equally and the reference model can by any model. Then, we propose the Generalized On-Policy Distillation (G-OPD) framework, which extends the standard OPD objective by introducing a flexible reference model and a reward scaling factor that controls the relative weight of the reward term against the KL regularization. Through comprehensive experiments on math reasoning and code generation tasks, we derive two novel insights: (1) Setting the reward scaling factor to be greater than 1 (i.e., reward extrapolation), which we term ExOPD, consistently improves over standard OPD across a range of teacher-student size pairings. In particular, in the setting where we merge the knowledge from different domain experts, obtained by applying domain-specific RL to the same student model, back into the original student, ExOPD enables the student to even surpass the teacher's performance boundary and outperform the domain teachers. (2) Building on ExOPD, we further find that in the strong-to-weak distillation setting (i.e., distilling a smaller student from a larger teacher), performing reward correction by choosing the reference model as the teacher's base model before RL yields a more accurate reward signal and further improves distillation performance. However, this choice assumes access to the teacher's pre-RL variant and incurs more computational overhead. We hope our work offers new insights for future research on OPD.