교사 학습을 넘어선 학습: 보상 외삽을 통한 일반화된 온-정책 지식 증류

초록

온-정책 증류(OPD)는 학생 모델이 생성한 궤적에서 교사 모델의 로짓 분포와 학생 모델을 정렬하는 방식으로, 학생 모델 성능 향상에 있어 강력한 실증적 효과를 보였으며 오프-정책 증류나 강화학습(RL) 패러다임을 종종 능가합니다. 본 연구에서는 먼저 OPD가 밀집 KL 제약 강화학습의 특수한 경우임을 이론적으로 보입니다. 여기서 보상 함수와 KL 정규화는 항상 동일한 가중치를 가지며, 참조 모델은 어떤 모델이든 될 수 있습니다. 이후 우리는 표준 OPD 목적 함수를 유연한 참조 모델과 KL 정규화 대비 보상 항의 상대적 가중치를 조절하는 보상 스케일링 인자를 도입하여 확장하는 일반화 온-정책 증류(G-OPD) 프레임워크를 제안합니다. 수학적 추론 및 코드 생성 과제에 대한 포괄적 실험을 통해 두 가지 새로운 통찰을 도출했습니다: (1) 보상 스케일링 인자를 1보다 크게 설정하는 것(즉, 보상 외삽)은 ExOPD라고 명명하며, 다양한 교사-학생 모델 크기 조합에서 표준 OPD보다 consistently 향상된 성능을 보입니다. 특히, 동일한 학생 모델에 도메인 특화 RL을 적용하여 얻은 서로 다른 분야 전문가들의 지식을 원래 학생 모델로 다시 융합하는 설정에서 ExOPD는 학생 모델이 교사 모델의 성능 한계를 넘어서고 도메인 교사들을 능가하는 성과를 가능하게 합니다. (2) ExOPD를 기반으로, 강한 모델에서 약한 모델로의 증류 설정(즉, 더 큰 교사 모델로부터 더 작은 학생 모델을 증류)에서 참조 모델을 RL 적용 전 교사 모델의 기본 모델로 선택하여 보상 보정을 수행하면 더 정확한 보상 신호를 얻어 증류 성능을 추가로 향상시킴을 발견했습니다. 그러나 이 선택은 교사 모델의 RL 적용 전 변형에 대한 접근을 전제로 하며 더 많은 계산 오버헤드를 수반합니다. 우리의 연구가 OPD에 대한 향후 연구에 새로운 통찰을 제공하기를 바랍니다.

English

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.

교사 학습을 넘어선 학습: 보상 외삽을 통한 일반화된 온-정책 지식 증류

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

초록

Support