The Teacher Is a Direction, Not a Destination: Extrapolating RL-Induced Representation Residuals in On-Policy Distillation

The Teacher Is a Direction, Not a Destination. Extrapolating RL-Induced Representation Residuals in On-Policy Distillation. By Hao Li and six others, from the University of Science and Technology of China, Rutgers University, and two other institutions. Published on the twenty-ninth of September. Its code repository is a placeholder: a readme, and nothing to run yet.

On-policy distillation trains a student model to match the teacher's next-token distributions, on trajectories the student generates itself rather than on text the teacher wrote. It has produced substantial empirical gains. Plain matching, though, puts a ceiling on the student: at best it becomes the teacher. Generalized variants go further and let the student surpass the teacher, by extrapolating an implicit reward in output space. Here the authors find a snag. The language-model head dampens the change unevenly across directions. Much of what changed in the teacher's hidden states reaches the logits with only a small fraction of its weight. Meanwhile, the sampled-token log-probability ratios that output-space extrapolation depends on inject noise, the extrapolation amplifies that noise, and training becomes unstable. Two things go wrong at once, then. The useful signal arrives weakened at the output, and what does arrive is noisy, with the noise made louder by the very step meant to go past the teacher.

The starting observation is about reinforcement learning. Training a model with reinforcement learning shifts its internal representations relative to the base checkpoint it started from, and the direction of that shift can be measured at every layer. If the useful change lives in representation space, then extrapolate there rather than at the output, before the language-model head has had a chance to attenuate it. That is RIDE, short for RL-Induced Direction Extrapolation.

At every layer and every token position, RIDE computes the residual between the teacher and the teacher's own checkpoint from before reinforcement learning. That residual is the direction. The student's hidden states are then regressed toward targets placed beyond the teacher, displaced further along the same residual.

The objective has a clean reading. Conditioned on a sampled trajectory, this regression is equivalent to maximizing a linear directional reward, defined by the residual, under a quadratic penalty centered at the teacher. Put plainly, the reward pulls the student along the direction reinforcement learning opened up, and the penalty keeps it from straying too far from the teacher. The equivalence makes both forces explicit.

The evaluation uses four pairs of base model and reinforcement-learned teacher, spanning different scales, architectures and pre-training lineages. On every pair, RIDE approaches or exceeds the teacher that was trained with reinforcement learning, and it is the only method whose mean does so. It also consistently beats output-space extrapolation. That alternative actually degrades the student whenever the teacher sits close to its base. For a method whose whole aim is to let the student go past its teacher, that contrast is the relevant one: the output-space route can leave the student worse off, while RIDE reaches the teacher or goes beyond it.

Instead of a target to be copied, the teacher becomes a heading: the difference between it and its pre-reinforcement self says which way is forward, and RIDE keeps the student moving that way, one layer and one token at a time.