PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2609.36484v1cs.LG

TL;DR

Output-space extrapolation can underrepresent RL-induced changes in hidden states and amplify sampled-token log-ratio noise. RIDE extrapolates the teacher-to-pre-RL residual directly in representation space; across four base/teacher pairs, it approaches or exceeds the RL-trained teacher and outperforms output-space extrapolation.

  • Problem

    Output-space extrapolation measures RL-induced change after the language-model head attenuates it anisotropically, while sampled-token log-ratios are length-biased and can destabilize training when scaled.

    Source: §1 · p. 2

  • Method

    RIDE regresses student hidden states toward layerwise targets displaced beyond the teacher along the teacher-to-pre-RL-checkpoint residual on student-generated contexts.

    Source: §1 · p. 2

  • Results

    Across four base/RL-teacher pairs, RIDE approaches or exceeds the teacher and outperforms OPRD and output-space baselines on mean Avg@16 across AIME24, AIME25, and AIMO.

    Source: §1 · p. 2

  • Takeaways & Limitations

    In the same-initialization setting, the layerwise teacher-to-base hidden-state difference measures what RL changed on the student’s visited contexts.

    Source: §3 · p. 4

    RIDE transfers biases or errors introduced by RL, which could in principle be amplified because its targets extend beyond the teacher.

    Source: ETHICS STATEMENT · p. 10

Abstract

from arXiv · show

On-policy distillation (OPD) trains a student to match the teacher's next-token distributions on the student's own trajectories and has yielded substantial empirical gains. Generalized variants allow the student to surpass the teacher by extrapolating an implicit reward in output space. The language-model head, however, attenuates this change anisotropically: much of the change encoded in the teacher's hidden states reaches the logits at a small fraction of its weight, and the sampled-token log-probability ratios on which output-space extrapolation relies inject noise that the extrapolation amplifies, making training unstable. We observe that reinforcement learning (RL) shifts a model's internal representations relative to its base checkpoint, and that the direction of this shift can be measured at every layer. Motivated by this observation, we propose RIDE (RL-Induced Direction Extrapolation), which extrapolates the RL-induced change directly in representation space: at every layer and token position, RIDE computes the residual between the teacher and its pre-RL checkpoint and regresses the student's hidden states toward targets displaced beyond the teacher along this residual. 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, which makes explicit how the objective moves the student along the RL-induced direction while limiting its deviation from the teacher. Across four base/RL-teacher pairs spanning different scales, architectures, and pre-training lineages, RIDE approaches or exceeds the RL-trained teacher on every pair and is the only method whose mean does so, and it consistently outperforms output-space extrapolation, which degrades the student whenever the teacher is close to its base. Project page: https://github.com/xixixixixxxx/RIDE.

1 INTRODUCTION

On-policy distillation transfers an RL teacher’s gains on student-generated trajectories, and RIDE treats the teacher-to-base representation change as a direction the student can continue along. It extrapolates hidden-state targets along that residual and approaches or exceeds the RL-trained teacher across four model pairs.

1 INTRODUCTION

  • On-policy distillation transfers teacher supervision on student-sampled contexts, providing a dense signal and reducing exposure bias.A common use is transferring reinforcement-learning gains back to a student sharing the teacher’s initialization.Source: §1 · p. 1
  • 79.8% of the teacher-to-base residual’s hidden-state energy lies in the head’s 512 weakest directions, versus 30.0% of its centered-logit energy.This attenuation leaves much of the residual weakly supervised in output space and earlier layers unconstrained.Source: Figure 1 · p. 2, §1 · p. 2
  • The teacher-to-reference log-probability ratio can act as an implicit reward, so up-weighting it enables students to outperform the teacher.The paper frames an RL-trained teacher as both a destination and a direction for a same-checkpoint student.Source: §1 · p. 2
  • RIDE regresses student hidden states toward layerwise targets displaced beyond the teacher along the teacher-to-pre-RL-checkpoint residual.Conditioned on a rollout, this is equivalent to maximizing a linear directional reward under a quadratic penalty around the teacher.Source: §1 · pp. 2–3
  • Across four base/RL-teacher pairs, RIDE approaches or exceeds the teacher and outperforms OPRD and output-space baselines on mean Avg@16 across AIME24, AIME25, and AIMO.Output-space extrapolation falls below its teacher on every pair.Source: §1 · p. 2

2 RELATED WORK

Related work spans output-focused on-policy distillation, methods for learning beyond the teacher, and representation-level distillation. RIDE combines the latter two lines by extrapolating hidden-state residuals from the pre-RL checkpoint to the teacher.

On-policy distillation

  • On-policy distillation variants supervise outputs, including approaches that remove teacher logits, reformulate targets for stability, analyze failure modes, or use privileged self-teaching.On-policy distillation samples from the student and supervises visited prefixes.Source: §2 · p. 3

Learning beyond the teacher

Representation-level distillation

  • Representation-level methods include layerwise hidden-state regression in OPRD and hidden-state transition-geometry matching in PHF (Yang et al., 2026a; Li et al., 2026d).RIDE instead uses the pre-RL-to-teacher residual to place targets beyond the teacher, drawing on structured RL updates and behavioral directions in hidden space (Mukherjee et al., 2025; Zhu et al., 2025; Shenfeld et al., 2026; Park et al., 2024a; Turner et al., 2023; Zou et al., 2023; Arditi et al., 2024).Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries set up on-policy distillation as teacher matching on student-generated prefixes and relate it to KL-regularized learning with an implicit reward. They then introduce representation matching and the shared-initialization setting in which teacher-to-base differences measure RL-induced change.

3 PRELIMINARIES

  • On-policy distillation samples student responses and minimizes reverse KL to the teacher on visited prefixes; sampled-token estimates yield a tokenwise update.Relative to a reference policy, the objective is KL-regularized RL with an implicit log-probability-ratio reward.Source: §3 · pp. 3–4
  • Target matching compares a model quantity at each position with a target, using a discrepancy measure and stopping gradients through the target.Representation matching uses hidden states across selected layers and masked positions with dimension-normalized squared error.Source: §3 · p. 4
  • On-policy feature distillation uses teacher hidden states as targets and has a deterministic gradient conditioned on the rollout.Teacher matching treats the teacher as the endpoint of learning.Source: §3 · p. 4
  • With a teacher obtained by RL from a base checkpoint and a student initialized at that base, the shared representational origin makes their layerwise difference on common prefixes a measure of RL’s change.This same-initialization setting occurs when distilling an RL result back into its base model or merging RL-trained experts that share a base.Source: §3 · p. 4

4 METHOD

RIDE treats the teacher’s displacement from its pre-RL checkpoint as a direction the student can continue along, rather than treating the teacher as the endpoint. It extrapolates targets in hidden-state space and interprets the resulting regression as directional reward maximization under a quadratic proximity penalty.

4.1 FROM TEACHER MATCHING TO RESIDUAL EXTRAPOLATION

  • RIDE displaces the teacher target along the RL-induced change, with λ = 1 matching the teacher and λ > 1 placing the target beyond it.The direction is recomputed for each rollout and prefix, rather than applied as a fixed offset.Source: §4.1 · p. 4

4.2 RL-INDUCED REPRESENTATION RESIDUAL AND THE RIDE OBJECTIVE

  • On each student prefix, RIDE computes deterministic layerwise teacher-minus-base residuals and regresses student states toward extrapolated targets, supervising all layers and final k response positions.The regression uses masked targets; λ = 1 recovers OPRD exactly.Source: §4.2 · p. 5, Figure 2 · p. 5

4.3 WHY THE DISPLACEMENT SHOULD BE MEASURED IN HIDDEN STATES

  • Hidden-state extrapolation preserves residual directions weakened by the head and uses a deterministic conditional gradient, whereas sampled output-space residual noise is amplified with (λ−1)^2.The head constrains weakly the directions it amplifies least; the RIDE gradient has zero conditional variance.Source: §4.3 · pp. 5–6

4.4 OPTIMIZATION INTERPRETATION

  • Conditioned on a rollout, the regression is equivalent to maximizing a linear reward aligned with the RL-induced residual under a quadratic penalty on displacement from the teacher.The reward weight vanishes at λ = 1, recovering OPRD.Source: §4.4 · p. 6

5 EXPERIMENTS

Across four base/RL-teacher pairs, RIDE is the only method whose mean exceeds the teacher on every pair. The experiments test the effect of extrapolation coefficients, analyze learning signals, and ablate the residual direction.

5.1 EXPERIMENTAL SETUP

  • Runs use four model pairs, shared prompts and training settings, student-generated rollouts, and three seeds for the main results; evaluation averages correctness over 16 responses on three math benchmarks.The pairs differ in scale, architecture, vocabulary, and pre-training lineage; the teacher and untouched student are fixed-checkpoint references.Source: §5.1 · pp. 6–7

5.2 MAIN RESULTS

  • RIDE alone exceeds the RL-trained teacher mean on all four pairs, though three margins are within one across-seed standard deviation; its margin over OPRD is 0.97–4.06 points.Output-space extrapolation falls below the teacher on every pair, and RIDE exceeds it by 9.1 points on average.Source: Figure 3 · p. 7, §5.2 · p. 7

5.3 EFFECT OF THE EXTRAPOLATION COEFFICIENT

  • On R1-Distill-1.5B, RIDE peaks at 55.4 Avg@16 with λ=1.25 and every λ in [1.15, 1.35] beats OPRD; ExOPD worsens for every λ>1.The coefficient sweep uses one run per value; for RIDE, λ=1.5 and 2 finish at 52.3 and 52.4, respectively.Source: §5.3 · pp. 7–8

5.4 MECHANISTIC ANALYSIS

  • The 512 weakest head directions contain 79.8% of residual hidden-state energy versus 33.3% under isotropy, and receive 73.0% of representation-loss gradient but 48.6% of output-KL gradient.The centered head attenuates rather than removes the residual; the student update aligns with it at cosine 0.954 and projects 1.69 along it, beyond the target coefficient 1.25.Source: Figure 5 · p. 9, §5.4 · p. 9
  • ExOPD’s conditional variance rises 12.6× from λ=1 to 2, while resampling the next token leaves the RIDE per-position gradient norm unchanged.The extrapolation term dominates ExOPD variance beyond λ≈1.3.Source: §5.4 · p. 9

5.5 ABLATION: THE RESIDUAL DIRECTION

  • Against OPRD’s 54.50, each of four direction controls scores within one point, while RIDE with the RL-induced residual reaches 56.38.The controls change the residual direction while holding displacement magnitude fixed.Source: §5.5 · p. 9

6 CONCLUSION

The conclusion frames the RL-induced representation residual as a direction for extrapolation and reports that RIDE approaches or exceeds its teacher across four pairs. It also identifies the method’s checkpoint and evaluation boundaries, while the appendix derives the implicit-reward form of OPD extrapolation.

6 CONCLUSION

  • RIDE displaces hidden-state targets beyond the teacher along the layerwise teacher–pre-RL-checkpoint difference, and approaches or exceeds the teacher on all four pairs.The method requires the pre-RL checkpoint and a shared representation space; experiments cover mathematical reasoning with one RL recipe.Source: §6 · p. 10
  • Because RIDE supervises hidden-state targets no model has produced, its calibration and safety are not guaranteed to match the teacher and require independent evaluation before deployment.The authors also note that displacing targets beyond the teacher could in principle amplify biases or errors introduced by RL.Source: ETHICS STATEMENT · p. 10, §6 · p. 10

A.1 ON-POLICY DISTILLATION AS IMPLICIT-REWARD REINFORCEMENT LEARNING

  • Weighting OPD’s implicit reward by λ yields a normalized target distribution whose log-probabilities combine the teacher and reference policies affinely.The sampled-token advantage becomes ℓ(v)−(λ−1)ρ(v); the tokenwise update holds sampled prefixes fixed.Source: §A.1 · p. 14

A.2 PROOF OF PROPOSITION 4.1

The appendix shows that ExOPD’s sampled-token variance includes a quadratic-in-λ−1 term, while RIDE’s per-position gradient is invariant to resampling the next token. It also derives RIDE’s objective as a residual-alignment reward balanced by a quadratic penalty.

A.2 PROOF OF PROPOSITION 4.1

  • Because the frozen teacher and base make ρ independent of the student, ExOPD’s variance includes a term growing quadratically in λ−1; it vanishes near the teacher only if RL left the prefix’s next-token distribution unchanged.For RIDE, the per-position loss and gradient depend on the prefix and frozen models, not the sampled token, so conditional variance is zero.Source: §A.2 · p. 14

A.3 REWARD-AND-PENALTY FORM: DERIVATION OF PROPOSITION 4.2

  • RIDE’s objective rewards student displacement aligned with the RL-induced residual and penalizes displacement quadratically around the teacher; λ−1 weights the reward against the penalty.At λ=1 the reward weight vanishes, recovering teacher matching.Source: §A.3 · p. 14

B EXTENDED EXTRAPOLATION-COEFFICIENT SWEEP

The coefficient sweep shows RIDE improving over teacher matching near λ = 1.25, while increasing λ beyond 1 degrades ExOPD and makes its late training less stable. The sweep also relates coefficient choice to response length and validation accuracy.

B EXTENDED EXTRAPOLATION-COEFFICIENT SWEEP

C DIRECTION CONTROLS

The direction controls hold the training protocol, displacement magnitude, and supervised layer and position sets fixed while changing the displacement direction. Their reported comparison is macro Avg@16 across AIME24, AIME25, and AIMO on R1-Distill-1.5B.

C DIRECTION CONTROLS

  • Each direction control changes only the displacement direction while keeping the training protocol, displacement magnitude, and RIDE layer and position sets fixed.The displacement magnitude is |λ − 1| ∥∆(l)t∥.Source: C DIRECTION CONTROLS · p. 16
  • The fixed-magnitude design isolates differences associated with displacement direction under the stated protocol.The table caption identifies macro Avg@16 as the comparison metric.Source: C DIRECTION CONTROLS · p. 16, Table 3 · p. 16
  • The reported outcome is Avg@16 (%) averaged over AIME24, AIME25, and AIMO on the R1-Distill-1.5B pair.The table uses λ = 1.25 unless otherwise stated.Source: Table 3 · p. 16

D TRAINING-SEED SUMMARY

The training-seed summary reports macro Avg@16 means and across-seed standard deviations for OPRD and RIDE, while additional diagnostics examine head attenuation and the shared-head assumption.

D TRAINING-SEED SUMMARY

  • Table 4 summarizes OPRD and RIDE macro Avg@16 means and standard deviations across three independent seeds.Each macro score equally averages AIME24, AIME25, and AIMO; the seeds are 14, 42, and 2027.Source: D TRAINING-SEED SUMMARY · p. 16, Table 4 · p. 17

E ADDITIONAL HEAD DIAGNOSTICS

  • On the R1-Distill-1.5B pair, the residual’s normalized head gain is 0.59, and its energy outside the top-r head subspace exceeds the isotropic expectation at every r.The measurements use 292,864 positions; the centered head is full rank, with singular values from 71.70 to 1.78. The 512 weakest directions carry 73.00% of the representation-loss gradient versus 48.60% of the output-KL gradient at the head input.Source: E ADDITIONAL HEAD DIAGNOSTICS · p. 17
  • The output-space identity assumes the base and student share a head; RL-induced teacher-head drift adds a term absent from the representation-space target.The drift is 1.65% in relative Frobenius norm, and fixed-head versus native-head targets differ by 18.70% in mean relative logit error.Source: E ADDITIONAL HEAD DIAGNOSTICS · p. 17

F ADDITIONAL LEARNING-SIGNAL DIAGNOSTICS

The learning-signal diagnostics compare ExOPD’s coefficient-dependent conditional variance with RIDE’s gradient variation under next-token resampling. ExOPD variance grows at larger coefficients, whereas RIDE has no next-token sampling variance in the measured per-position gradient norm.

F ADDITIONAL LEARNING-SIGNAL DIAGNOSTICS

  • The ExOPD conditional moments are computed over the full 151,936-token vocabulary at 9,152 fixed prefixes, and resampling reproduces them within 0.38%.RIDE gradient norms are measured at the same fixed prefixes with 8,192 next-token resamples per prefix.Source: F ADDITIONAL LEARNING-SIGNAL DIAGNOSTICS · p. 17
  • ExOPD’s conditional variance reaches 0.0067 at λ ≈ 1.14, then rises to 2.80× its λ = 1 variance at λ = 1.50 and 12.60× at λ = 2.The quadratic term (λ − 1)^2 Var[ρ] accounts for essentially all growth beyond λ ≈ 1.30.Source: F ADDITIONAL LEARNING-SIGNAL DIAGNOSTICS · p. 17
  • RIDE’s per-position gradient norm has exactly zero variance across resampled next tokens at every prefix and λ.Its target depends on the prefix and frozen models, not the sampled next token; the between-prefix variance is approximately 4 × 10^3 at every λ.Source: F ADDITIONAL LEARNING-SIGNAL DIAGNOSTICS · pp. 17–18

G IMPLEMENTATION AND DIAGNOSTICS

All methods share the same training and evaluation configuration, while runs differ only in target construction. The section also specifies RIDE’s per-rollout cost and training diagnostics.

G IMPLEMENTATION AND DIAGNOSTICS

  • RIDE adds one frozen base-model forward pass per rollout, with no extra communication and the same transfer volume as OPRD.The teacher worker forms the target in place and transmits one hidden-state tensor.Source: G IMPLEMENTATION AND DIAGNOSTICS · p. 18, Table 6 · p. 18
  • The training and evaluation configuration is shared across methods and all four base/teacher pairs; only target construction varies.Final answers are extracted with the standard boxed parser and graded by exact match.Source: G IMPLEMENTATION AND DIAGNOSTICS · p. 18, Table 5 · p. 18
  • RIDE forms each target using an additional base-model forward pass; with λ = 1 or no base checkpoint, training coincides with OPRD.Masked or beyond-response positions are zero-padded in both hidden-state tensors, leaving the masked loss unaffected.Source: G IMPLEMENTATION AND DIAGNOSTICS · p. 18
  • Because the loss is unnormalized, the loss coefficient is rescaled by λ^-2 to equalize initial loss across λ values.The initial per-position loss is λ^2 times the OPRD loss before this rescaling.Source: G IMPLEMENTATION AND DIAGNOSTICS · pp. 18–19
  • Training diagnostics track hidden-state cosine similarity, residual-norm ratios overall and per layer, and target-norm ratios.The target norm ratio monitors the scale sensitivity associated with changing λ.Source: G IMPLEMENTATION AND DIAGNOSTICS · p. 19
Loading 2609.36484v1…