Zone of Proximal Policy Optimization: Teacher in Prompts, Not Gradients

A new approach to knowledge distillation aims to overcome the limitations of logit imitation and reinforcement learning by integrating teacher guidance directly into prompts rather than gradient-based constraints, specifically targeting the "small-student" regime to improve generalization.

Addressing the Brittleness of Knowledge Distillation

Traditional knowledge distillation focuses on transferring the competence of a large teacher model to a smaller student model. However, this process often proves brittle when the student model is significantly smaller than the teacher. The primary issue arises from forcing the student to imitate the teacher's logits, which tends to concentrate the student's learning on the teacher's sharpest modes. This narrow focus often degrades the student's ability to generalize across benchmark families that extend beyond the initial training corpus.

Limitations of Reinforcement Learning (RL)

To circumvent the pitfalls of logit imitation, Reinforcement Learning (RL) is often employed, allowing the student to train on its own rollouts. While this avoids direct imitation, it introduces a new failure mode: when a student fails every rollout for a specific question, the process yields zero advantage. Consequently, these critical learning opportunities are silently discarded, leaving the student unable to bridge the gap in competence.

The Proposed Solution: Zone of Proximal Policy Optimization

The researchers introduce a method where the teacher acts as a guide within the prompts rather than through gradient-based constraints. By shifting the teacher's role from the objective function to the prompt context, the method seeks to create a "Zone of Proximal Policy Optimization." This approach aims to provide the student with the necessary guidance to succeed in rollouts that would otherwise fail, thereby ensuring a steady stream of positive reinforcement and improved generalization.

Note: Due to the limited nature of the provided source text, specific architectural details and quantitative results of the Zone of Proximal Policy Optimization method are not available in this summary.