On the Optimal Teacher for On-Policy (Self-)Distillation

On-policy distillation has become an attractive way to post-train language models. Unlike RL approaches that use sparse outcome rewards, OPD provides dense supervision along the student’s own trajectories, which can enable greater sample and compute efficiency.

Recently On-Policy Self-Distillation (OPSD) has become especially appealing. Instead of requiring a stronger teacher, one can use the same base model as both the student and the teacher. While the student policy is only given the default task environment, the teacher is provided with privileged information. This privileged information can come in the form of hints, feedback, verifier traces, or reference-side context.

The hope of OPSD is that this privileged information is enough to create a strong teacher for the student to imitate. In the following, we will derive the optimal form that this teacher would take.

The Optimal Teacher Setting

It is difficult to say anything theoretical on techniques that merely prompt models. We therefore consider a cleaner formulation where we can train a policy to behave as an optimal teacher for a given student.

Let the student policy be

and let the teacher policy be

Here \(x\) is the task input, \(c\) is privileged information available only to the teacher during training, and \(\tau\) is a trajectory. The object we want to characterize is the optimal teacher,

The OPD Optimality Condition

The student update induced by OPD is determined by the teacher-student log-ratio. For a fixed student \(\pi_S\), OPD gives the policy-gradient direction

The corresponding plain RL update is

So, if we define the optimal teacher as the teacher whose OPD update exactly matches the RL update, the necessary and sufficient condition is

This is the exact condition. It only requires equality after projection onto the student score function.

A Pointwise Sufficient Teacher

It suffices to satisfy the stronger pointwise condition

The constant \(\log Z(x,c)\) does not affect the policy-gradient update because

Exponentiating gives

So the pointwise construction gives one clean optimal teacher: the student distribution exponentially tilted toward high-advantage trajectories.

Training Toward The Optimal Teacher

Now we ask how to train a teacher toward this pointwise optimal form.

Take any candidate teacher \(\pi_T\). Since the target distribution \(\pi_T^\star\) is now defined, the most direct objective is to minimize the distance from the candidate teacher to the optimal teacher:

Expanding this KL gives the teacher-training objective

The above objective indicates that the optimal teacher should not simply maximize reward. It should maximize reward while staying close to the student distribution.

This gives a precise form to the intuition behind pedagogical teaching: the teacher should only move away from the student when the advantage is large enough to justify that movement. High-reward trajectories are useful only insofar as they can be connected back to the student’s current policy. In this view, the teacher-student KL is not an extra regularizer added after the fact; it is the cost of producing a learning signal that remains compatible with OPD.

A Practical Algorithm

The above objective suggests an iterative algorithm: alternately train the teacher and the student until convergence.

At iteration \(k\), freeze the current student \(\pi_S^{(k)}\) and train a privileged teacher against it:

Here \(\widehat{A}(x,\tau)\) is whatever practical advantage estimate is available. For example, it could be a centered verifier reward.

Then freeze the teacher and update the student using the OPD signal induced by that teacher:

The student policy is then updated in the policy-gradient direction

In pseudocode:

At convergence, the teacher should no longer need to move far from the student to produce high-advantage trajectories. In that sense, the student has absorbed the part of the privileged teacher that is useful under the original task environment.