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

Hao Li, MeiJia Chen, Weijie Ren, Donghan Li, Zijun Tian, Jingchun Huang, Naibo Wang

RIDE improves on-policy distillation by extrapolating RL-induced hidden-state residuals rather than noisy output-space logits.

Can we improve on-policy distillation by treating the RL-trained teacher's hidden-state shift as a directional vector to be extrapolated, rather than just a destination to be matched?

On-policy distillation (OPD) typically treats the teacher as the final destination, but students can often surpass their teachers by extrapolating the RL-induced change. Existing output-space methods fail because the language-model head attenuates these signals anisotropically and amplifies sampled-token noise. RIDE (RL-Induced Direction Extrapolation) instead computes the layerwise hidden-state residual between the teacher and its pre-RL base checkpoint, then regresses the student toward a target displaced beyond the teacher along this residual. Across four model pairs, RIDE consistently outperforms both standard distillation and output-space extrapolation, approaching or exceeding the performance of the RL-trained teacher itself.

Paper Primer

RIDE treats the RL-induced update as a directional vector in representation space rather than a scalar reward in logit space. By measuring the residual at every layer before the head, the method avoids the head's anisotropic attenuation and provides a deterministic, noise-free learning signal that remains stable even when extrapolating far beyond the teacher.

RIDE consistently outperforms output-space extrapolation and standard on-policy distillation.

Across four base/RL-teacher pairs (R1-Distill-1.5B, Qwen3-4B, Llama-3.2-3B, and Phi-4-mini), RIDE is the only method whose mean performance exceeds the teacher on every pair. RIDE outperforms OPRD by 0.97 to 4.06 points and ExOPD by an average of 9.1 points on the R1-Distill-1.5B pair.

Why is measuring the RL-induced change in hidden states superior to measuring it in output logits?

The language-model head attenuates the RL-induced residual anisotropically, meaning output-space objectives supervise the most important changes only weakly. Furthermore, output-space extrapolation amplifies sampled-token noise, whereas the hidden-state residual provides a deterministic gradient that is invariant to next-token sampling.

What is the role of the pre-RL base checkpoint in this method?

The base checkpoint serves as the representational origin. By comparing the frozen teacher to this base checkpoint on identical student-generated prefixes, RIDE isolates the specific change RL introduced, allowing the student to continue along that learned trajectory.

Researchers can now treat RL-trained teachers as directional guides rather than static targets, using representation-space residuals to push student performance beyond the teacher's own capabilities without the instability of logit-based extrapolation.

Introduction and Motivation

We expose why on‑policy distillation stalls at the teacher’s performance.

On‑policy distillation (OPD) lets a student learn from a teacher by matching the teacher’s next‑token distribution on the student’s own trajectories, but the language‑model head dampens most of the underlying representation change, making extrapolation in output space noisy and unstable.

OPD aligns a student’s hidden‑state predictions with a teacher’s on the exact trajectories the student generates, providing dense, on‑policy supervision that avoids exposure bias.

How does OPD differ from conventional (offline) knowledge distillation?

Offline distillation trains the student on a fixed dataset of teacher‑generated examples, so the student sees contexts the teacher never visited. OPD, by contrast, samples trajectories from the student itself and then asks the teacher to label exactly those states, eliminating the mismatch between training and inference distributions.

The limitation of OPD is the teacher’s performance ceiling.

Notation and Setup

Defines notation, on‑policy distillation, and the base concepts needed for RIDE.

We denote the student policy by $π_θ$ and the frozen teacher by $π_T$, both decoder‑only transformers with $L$ layers, hidden dimension $d$, vocabulary $V$, and a language‑model head $W_{head} \in ℝ^{|V| \times d}$. For a prompt $x∼D_x$, the student generates a token sequence $\hat y=(\hat y_1,…,\hat y_T)∼π_θ( \cdot |x)$, and we write $π( \cdot |x,\hat y_{<t})$ for the next‑token distribution given the generated prefix. The layer‑$l$ residual‑stream state at position $t$ is $h^{(l)}_t \in ℝ^d$.

On‑policy distillation (OPD) trains the student by sampling its own trajectories and minimizing the reverse KL to the teacher, yielding the objective $L_{OPD}(θ)=𝔼_{x,\hat y∼π_θ}[ \sum _{t=1}^{T}D_{KL}(π_θ( \cdot |x,\hat y_{<t})‖π_T( \cdot |x,\hat y_{<t}))]$ (1). In the same‑initialization setting we start from a base checkpoint $π_B$, obtain a teacher $π_T=RL(π_B)$, and initialize the student at $π_{θ0}=π_B$, so the only difference between teacher and student lies in the RL‑induced changes.

The base checkpoint is the pretrained model state before any reinforcement‑learning fine‑tuning, serving as the common starting point for both teacher and student.

The teacher‑to‑base residual captures how the RL fine‑tuning altered the hidden representation at each layer and token position, expressed as the difference between the teacher’s state and the base state.

The RIDE Mechanism

RIDE extrapolates hidden‑state residuals so a student can keep moving along the teacher’s RL‑induced improvement direction.

The student’s policy stalls at the teacher’s checkpoint because the RL‑trained teacher only supplies a destination, not a continuation direction.

The residual is the hidden‑state difference that RL added when turning the base model into the teacher; it tells us exactly how the teacher’s representation moved.

Compute the residual: $\Delta_t = (3,5) - (1,2) = (2,3)$.

The residual shows that RL increased the first component by 2 and the second by 3.

If the student were to copy the teacher exactly, it would stop at $(3,5)$; the residual tells us the direction to keep moving.

The residual isolates the effect of RL alone, removing any confounding influence of the input prefix.

RIDE takes the teacher’s hidden state and pushes it further along the RL‑induced residual, scaling the step with a coefficient $\lambda$.

Compute the target: $\tau = 1.5\cdot(3,5) + (1-1.5)\cdot(1,2) = (4.5,7.5) + (-0.5,-1) = (4.0,6.5)$.

The target lies beyond the teacher along the residual direction (2,3).

The student will be trained to match $(4.0,6.5)$, thereby continuing the RL‑induced improvement.

Scaling $\lambda$ above 1 creates a new point that preserves the residual’s direction while extending its magnitude.

$\lambda$ is a scalar that determines how far the target moves beyond the teacher along the residual direction.

Run a student rollout to obtain prefixes $(x,\hat y_{<t})$.

Forward the same prefixes through the frozen teacher $\pi_{T}$ and base $\pi_{B}$ to collect $h_T^{(\ell)}$ and $h_B^{(\ell)}$ at every layer ℓ.

Compute the residual $\Delta^{(\ell)}_t = h_T^{(\ell)} - h_B^{(\ell)}$.

Form the target $h^{* (\ell)}_t = \lambda\,h_T^{(\ell)} + (1-\lambda)\,h_B^{(\ell)}$.

Apply the masked $L_2$ loss $L_\lambda(\theta)$ between the student hidden state $h_{\theta}^{(\ell)}$ and the stopped‑gradient target $h^{* (\ell)}_t$.

Back‑propagate the loss to update the student parameters $\theta$.

RIDE training step (pseudocode)

**Figure 1.** The language-model head attenuates RL-induced representation changes unevenly. (a) Residual extrapolation before and after projection through the head. (b) Hidden-state and centered-logit energy shares of the teacher-to-base residual across the singular directions of the head. Whiskers show prompt-bootstrap 95% confidence intervals.

**Figure 2.** **Overview of RIDE.** On student-generated prefixes, the frozen base and teacher define layerwise residuals. RIDE extrapolates hidden-state targets along these residuals and trains the student by masked regression toward them. Right: target geometry at a single layer and position.

Empirical Performance

RIDE outperforms the teacher on every evaluated model pair.

RIDE consistently outperforms the teacher across all tested base/RL‑teacher pairs.

Table 1 shows RIDE’s mean Avg@16 scores exceed the teacher’s scores for every pair, and Figure 3 confirms the positive relative difference.

**Table 1.** Main results (Avg@16, %) on four base/RL-teacher pairs. Trained-method entries are means over three independent training seeds; the teacher and untouched student are fixed-checkpoint reference evaluations. Each run uses its own on-policy rollouts; within a pair, all runs share prompts, sampling settings, and optimizer schedule. ExOPD and RIDE use $\lambda = 1.25$ with the pre-RL checkpoint as reference. Bold marks the best distillation method per column.

**Table 4.** Mean ± SD of macro Avg@16 (%) over three independent training runs (seeds 14, 42, and 2027). Means are taken from Table 1; SDs are across training seeds. Teacher scores are fixed-checkpoint references.

**Figure 3.** Avg@16 relative to the RL-trained teacher (points; the horizontal line is the teacher) for every method and base/teacher pair in Table 1. RIDE is the only method whose mean lies above the line on all four pairs; across-seed standard deviations are in Table 4.

RIDE consistently outperforms the teacher across all tested pairs.

Mechanistic Analysis and Controls

Direction‑control ablations show how altering the residual direction impacts Avg@16 performance.

We evaluate how the direction of the displacement vector influences the final Avg@16 score while keeping the displacement magnitude $|\lambda-1|\|\Delta_t\|$, the layer set, and the position set fixed. Each row in Table 3 corresponds to a different way of constructing the target direction.

Using a random direction (same norm) improves Avg@16 by +0.54 over the teacher‑matching baseline.

Avg@16 = 55.04 vs. 54.50 for Teacher matching (OPRD, $\lambda=1$).

Reversing the direction ($\lambda$ = 0.75) hurts performance, dropping Avg@16 by –0.60.

Avg@16 = 53.90 vs. 54.50 for the baseline.

Choosing a mismatched origin ($h_T - h_B'$) yields a modest gain of +0.25.

Avg@16 = 54.75 vs. 54.50 baseline.

Employing a trajectory‑mismatched residual gives a gain of +0.62.

Avg@16 = 55.12 vs. 54.50 baseline.

RIDE (the RL‑induced residual) achieves the highest score, improving Avg@16 by +1.88.

Avg@16 = 56.38 vs. 54.50 baseline.

**Figure 5.** Mechanistic analysis on the R1-Distill-1.5B pair ($\lambda = 1.25$, step 500, final layer). (a) The 512 weakest head directions hold 79.8% of the residual's hidden-state energy versus 33.3% under isotropy. (b) Realized student change versus target. (c) Next-token conditional variance of the ExOPD advantage and the RIDE gradient norm across $\lambda$.

Theoretical Foundations

Formal derivations clarify the implicit‑reward view of OPD and the RIDE loss.

We first rewrite the OPD KL loss to expose an implicit reward that matches the teacher policy.

Weighting the reward by a scalar $\lambda$ yields a family of objectives whose optimal policy mixes the teacher and reference policies.

Substituting the optimal $p^{\star}$ back into the per‑token loss gives the extrapolated advantage used in training.

Proposition 4.1 gives the variance of a linear combination of two random variables, which underpins the analysis of $A_{\lambda}$.

Proposition 4.2 rewrites the RIDE loss as a reward‑penalty quadratic form, making the optimisation landscape explicit.

Setting the gradient of the quadratic form to zero gives the closed‑form optimal student state $h^{\star}=h_{T}+(\lambda-1)\Delta$.

The extended $\lambda$‑sweep (Table 2) shows that coefficients in $[1.15,1.35]$ improve Avg@16 over the baseline $\lambda=1$, while larger $\lambda$ values degrade ExOPD performance and destabilise late‑training behaviour (Figures 6‑7).

Implementation and Diagnostics

Implementation details cover shared configs, per‑rollout costs, target construction, loss scaling, and diagnostics.

Table 5 lists the configuration that all four base/teacher pairs share during training and evaluation. It specifies the rollout dataset (DAPO‑Math‑17K), sampling temperature 1.00, maximum response length 16 384 tokens, 500 optimization steps with AdamW at a peak learning rate of 1 × 10⁻⁵, a 3 % linear warm‑up followed by cosine decay, bf16 precision with FSDP (v1) on four H200 GPUs, and the evaluation metric Avg@16 on the AIME24, AIME25, and AIMO benchmarks.

Target construction proceeds by loading the frozen pre‑RL checkpoint $\pi$B onto the worker that already hosts the teacher, running one additional forward pass of $\pi$B for each rollout, forming the hidden‑state tensor $h$⋆ in place according to Eq. (6), and transmitting that single tensor to the trainer. This mirrors OPRD’s communication volume while adding no extra data beyond the hidden state.

Remark G.1 explains that at initialization the student equals the base model (h$\theta$₀ = hB), so the per‑position loss is $\lambda$²‖$\Delta$‖²₂/d, i.e. $\lambda$² times larger than OPRD. To keep the effective step size comparable across $\lambda$ values, the loss coefficient is multiplied by $\lambda$⁻², which normalizes the initial loss back to the OPRD level for any $\lambda$.

During training we log three diagnostics: the cosine similarity between the base hidden state hB and the teacher hidden state hT, the residual norm ratio ‖$\Delta$‖₂/‖hT‖₂ (both overall and per layer), and the target norm ratio ‖h⋆‖₂/‖hT‖₂. These metrics directly monitor the scale sensitivity described in Remark G.1.

Conclusion

Summarizes RIDE’s empirical gains, ethical considerations, reproducibility steps, and AI‑assisted writing policy.

RIDE extrapolates the hidden‑state residual between a teacher and its pre‑RL checkpoint, yielding students that match or surpass their teachers on four base/teacher pairs while remaining stable where output‑space extrapolation collapses.

RIDE requires a pre‑RL checkpoint, a shared representation space, and a single global $\\lambda$; it was evaluated only on mathematical‑reasoning tasks using one RL recipe.

Ethics Statement: the work distills mathematical‑reasoning ability using publicly released checkpoints and datasets, involves no human subjects, but may amplify teacher biases because RIDE extrapolates beyond the teacher’s behavior.

Reproducibility Statement: the RIDE objective is fully specified by equations (6) and (7) plus Remark G.1, with all hyperparameters listed in Table 5, rollout costs in Table 6, and code, checkpoints, and scripts to be released upon publication.

AI Use Statement: generative AI tools were used only for language polishing and consistency checks; the authors retain responsibility for all research decisions and reported results.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers