A Shared Learning Rate Is Not a Neutral Control in Selective On-Policy Distillation

Chencheng Zhu

cs.LGAug 17, 2026

A Shared Learning Rate Is Not a Neutral Control in Selective On-Policy Distillation

Authors: Chencheng Zhu

Organizations: UNSW Sydney

Abstract

Selective on-policy distillation trains a student only at the token positions a selector scores highest, and the literature compares selectors under a single shared learning rate--a control chosen to be neutral. We show it is not. Under LoRA on GSM8K (Qwen2.5-1.5B student, 7B teacher), across an 8x learning-rate grid, dense supervision is statistically flat (swing 1.8 pp, p=0.26) while every selective arm moves with the rate: 5.4 pp for a random 5% subset, 6.7 pp for a total-variation selector, up to 17.7 pp for a teachability selector. Consequently the dense-versus-selective verdict reads 10.1 pp at lr=1e-4 but 5.1 pp at 5e-5--a 2.0x difference decided by a parameter the protocol treats as scenery--and two of six pairwise significance calls between selectors flip between adjacent rates without any rank inversion. We call this selector-rate entanglement and trace it to selection itself rather than step size: AdamW update magnitudes track the rate to within 2.2% despite 15.5x gradient-norm differences across arms. A preregistered frozen-scoring ablation (selection scored by the initial student; criterion, budget, and on-policy rollouts unchanged; 12 seeds per cell) shows live scoring adds 3.79+/-1.69 pp of rate sensitivity (p=0.035) while the frozen arm remains significantly entangled (p=0.015): the feedback loop aggravates the phenomenon rather than causing it. Under full fine-tuning at the rates this literature actually uses (1e-6 to 1e-5) the pattern grows: dense itself swings 19.8 pp, the selective arm 49.5 pp, and the verdict ranges from a non-significant +3.6 pp at the published operating point to +34 pp (p=0.005) one notch hotter. On MATH-500 the rate dependence does not reproduce under LoRA, scoping that result, while the ~10 pp cost of selective training does. We prescribe reporting the arm x rate matrix, not a shared-rate column, as a precondition for selector comparisons.

Figures & tables

Figure 1: Dense supervision is flat across the learning-rate grid; every selective arm is not. Accuracy change over the untrained student (GSM8K, mean

\pm

seed SD; dots are seeds) versus learning rate. Swings: dense

1.8\,\text{pp}

(

t{=}1.40

, n.s.), random

5\%

subset

5.4

, frozen-scored TV

2.9

(

p{=}0.015

), live-scored TV

6.7

(

p{<}0.001

). The consequence for comparisons: the dense-versus-selective verdict is

10.1\,\text{pp}

\eta{=}10^{-4}

and

5.1\,\text{pp}

5{\cdot}10^{-5}

—a factor of

2.0

decided by a parameter the protocol treats as neutral.

Arm	Loop	$\eta{=}10^{-4}$	$5{\cdot}10^{-5}$	$2.5{\cdot}10^{-5}$	$1.25{\cdot}10^{-5}$	Swing
Full	–	$+1.87{\pm}2.0$ (3)	$+3.64{\pm}1.0$ (3)	$+3.29{\pm}1.3$ (3)	$+2.58$ (1)	$1.8$
Random	open	$-3.64{\pm}3.3$ (3)	$+1.44{\pm}0.7$ (3)	$+1.72{\pm}0.9$ (3)	–	$5.4$
TV-frozen	$\theta_{0}$	$-7.37{\pm}3.2$ (12)	$-4.48{\pm}1.8$ (12)	–	–	$2.9$
TV-live	$\theta_{t}$	$-8.18{\pm}4.1$ (12)	$-1.50{\pm}1.9$ (12)	$-3.21{\pm}2.3$ (3)	$-3.49{\pm}1.5$ (3)	$6.7$

Table 1: Arm

\times

learning-rate matrix (duration-homogeneous). Accuracy change (pp) over the untrained baseline (

64.97\%

, greedy, 1,319 problems), mean

\pm

seed SD; parenthesized counts are seeds. Swing is the difference between the arm’s best and worst cell means. TV-frozen is the frozen-scoring arm of Section 5 .

Arm	Scoring	$\eta{=}10^{-4}$	$\eta{=}5{\cdot}10^{-5}$	$\Delta(\eta)$
Preregistered powered test, $n{=}12$ per cell
TV-frozen	$\theta_{0}$	$-7.37$	$-4.48$	$2.89\pm 1.07$ ( $p{=}0.015$ )
TV-live	$\theta_{t}$	$-8.18$	$-1.50$	$6.68\pm 1.31$ ( $p{<}0.001$ )
interaction				$\mathbf{3.79\pm 1.69}$ ( $p{=}0.035$ )
Original underpowered version
TV-frozen ( $n{=}3$ )	$\theta_{0}$	$-5.99\pm 2.84$	$-3.36\pm 1.08$	$2.63$ ( $p{=}0.25$ )

Table 2: Scoring selection with the live model adds

3.79\,\text{pp}

of rate sensitivity. Accuracy change (pp) over the untrained baseline, mean

\pm

seed SD. Both arms: TV criterion,

5\%

budget, on-policy rollouts, 512 steps/round, selection recomputed every round; the arms differ only in which model scores the selection.

\Delta(\eta)

is the within-arm difference across rates (Welch). The

n{=}12

block is the preregistered powered test (P7); the smaller-

n

block is the original, underpowered version, reported as preregistered.

Figure 2: The loop changes the shape of selection drift, not its magnitude—with a dose–response in the rate. Top-5% Jaccard

J_{t}

between

\theta_{t}

- and

\theta_{0}

-scored selections on the same positions (mean over seeds; bands are seed ranges;

n{=}12

per TV cell at the final round). After one round every arm sits at

J_{1}\approx 0.35

regardless of rate or scoring source. Then they separate: frozen-scored selections keep drifting away (

\to 0.28

–

0.29

, both rates), live-scored ones recover toward the initial selection, with the recovery monotone in the rate (

J_{2}=0.39

0.44

0.44

0.48

from

\eta{=}10^{-4}

down to

1.25{\cdot}10^{-5}

). The frozen and live ranges do not overlap at any seed (

\max 0.314

vs.

\min 0.354

;

t{=}17.6

Arm	Scoring	$\eta{=}10^{-4}$	$5{\cdot}10^{-5}$	$\Delta(\eta)$	$J_{1}$
Ent-live	$\theta_{t}$	$+0.38{\pm}1.2$	$+3.54{\pm}0.1$	$3.16$ ( $p{=}0.043$ )	$0.62$
Ent-frozen	$\theta_{0}$	$-0.48{\pm}2.0$	$+3.82{\pm}0.3$	$4.30$ ( $p{=}0.059$ )	$0.62$
Teach-live	$\theta_{t}$	$-15.54{\pm}3.2$	$-3.87{\pm}1.1$	$11.68$ ( $p{=}0.016$ )	$0.86$
Teach-frozen	$\theta_{0}$	$-22.11{\pm}7.1$	$-4.37{\pm}4.0$	$17.74$ ( $p{=}0.030$ )	—

Table 3: S2 replication. Accuracy change (pp), mean

\pm

seed SD,

n{=}3

per cell.

J_{1}

is the first-round selection drift (smaller = the criterion’s scores move more as the model moves).

Arm	$\eta{=}10^{-4}$	$5{\cdot}10^{-5}$	$\Delta(\eta)$
Full	$48.00{\pm}3.49$ ( $-1.23$ )	$50.56{\pm}0.77$ ( $+1.33$ )	$2.56$ ( $t{=}1.24$ , $p{=}0.33$ )
TV-live	$37.95{\pm}2.69$ ( $-11.28$ )	$39.69{\pm}2.77$ ( $-9.54$ )	$1.74$ ( $t{=}0.78$ , $p{=}0.48$ )
gap	$10.05{\pm}2.55$ ( $p{=}0.019$ )	$10.87{\pm}1.66$ ( $p{=}0.015$ )

Table 4: MATH-500 replication. Greedy accuracy (%) on the 325 numeric-answer problems, mean

\pm

seed SD,

n{=}3

; the untrained student scores

49.23\%

on the same set, and the bracketed values are changes over it.

Figure 3: Under full fine-tuning the asymmetry survives in graded form and grows. Accuracy change over the untrained student (A800 baseline, mean

\pm

seed SD,

n{=}3

) at rates bracketing those the literature publishes with. Dense swings

19.8\,\text{pp}

; the TV-selective arm swings

49.5\,\text{pp}

(

2.5\times

more). The dense-versus-selective verdict is a non-significant

+3.6\,\text{pp}

at the published operating point (

2{\cdot}10^{-6}

) and

+34.0\,\text{pp}

(

p{=}0.005

) one notch hotter.

Arm	$\eta{=}10^{-5}$	$2{\cdot}10^{-6}$	$10^{-6}$	Swing
Full	$-14.30{\pm}1.9$	$+4.22{\pm}0.9$	$+5.53{\pm}1.2$	$19.84$ ( $p{<}0.001$ )
TV	$-48.32{\pm}5.6$	$+0.61{\pm}1.9$	$+1.21{\pm}3.5$	$49.53$ ( $p{=}0.001$ )
gap	$+34.02{\pm}3.4$ ( $p{=}0.005$ )	$+3.61{\pm}1.2$ ( $p{=}0.060$ )	$+4.32{\pm}2.1$ ( $p{=}0.153$ )

Table 5: Full fine-tuning grid. Accuracy change (pp) over the untrained student on the same hardware (

63.46\%

; see Appendix A for why this baseline differs from the LoRA tables’), mean

\pm

seed SD,

n{=}3

Paper	Reported rate	Per-arm tuning?	Adaptation
Xu et al. 2026	$10^{-6}$ / $3{\cdot}10^{-7}$ (by model pair)	no (shared)	full FT
Jiang et al. 2026	$2{\cdot}10^{-5}$ off-policy, $2{\cdot}10^{-6}$ on-policy	no (shared)	full FT
Koo et al. 2025	grid $\{5,1,0.5\}{\cdot}10^{-5}$ (by model size)	no (by size)	LoRA $r{=}16$
Huang et al. 2025	not reported for LLMs	—	LoRA
Wang et al. 2026	not reported	—	not reported
This work	grid $\{10,5,2.5,1.25\}{\cdot}10^{-5}$	yes (the point)	LoRA $r{=}32$

Table 6: Learning-rate practice in selective on-policy distillation. Rates as reported in each paper; “shared” means one value across all selector arms compared in that paper.

Selective Knowledge Distillation

Selective knowledge distillation transfers only the examples, tokens, features, or teacher signals most useful to a student.

See all 49 papers

Sep 24, 2026

Not Every Token Is Worth Distilling: Selective Supervision for Direct-OPD

Direct On-Policy Distillation (Direct-OPD) transfers reinforcement-learning-induced policy improvements from a small model to a larger student by using the token-level log-ratio between post-RL and pre-RL checkpoints as dense supervision on the student's own rollouts. This transfer rewards the policy shift at every state, yet the log-ratio measures only relative change: it can stay fixed even as the probability mass that both checkpoints assign to the student's candidate tokens vanishes. Through an exact construction, we show that the Direct-OPD reward and its update can remain unchanged while the Jensen-Shannon divergence (JSD) and both KL directions between the checkpoints vanish with this mass, and we note that a small JSD bounds how much the teacher's behavior changed. Motivated by this analysis, we propose Selective Supervision for Direct-OPD (S

^2

D-OPD), which ranks student-sampled states by their teacher-reference JSD and masks Direct-OPD supervision at low-divergence states, retaining only the top 10% of states per response. Across two teacher pairs and four student models ranging from 1.7B to 8B parameters, S

^2

D-OPD improves held-out accuracy over dense Direct-OPD on AIME and HMMT benchmarks in seven of eight settings and matches it in the eighth, without extra forward passes.

Yibo Zhao, Zixuan Yang, Yunshi Lan +1

School of Data Science and Engineering East China Normal University · Hugging Face

Selective Knowledge Distillation On-Policy Distillation RL Post-Training Knowledge Distillation

Ask PDF

Jun 21, 2026

On the Position Bias of On-Policy Distillation

On-Policy Distillation (OPD) improves the learning efficiency of standard reinforcement learning through dense, token-level supervision from teachers. In the standard KL objective of OPD, token-level losses are uniformly averaged, implying equal weights for all tokens. However, we discover that not all tokens are created equal: as student rollouts grow longer, they deviate further from the teacher's distribution, leading to degraded supervision quality at later positions. As a result, OPD using only the first 30% of tokens can perform comparably to using all tokens, whereas OPD using only the last 30% of tokens barely learns anything. In this work, we provide a principled understanding of this issue through the lens of constrained optimization. Based on these insights, we derive Importance-Weighted On-Policy Distillation (IW-OPD), in which the weight assigned to each token depends on the accumulated discrepancy between the student's and teacher's distributions, naturally upweighting earlier tokens and downweighting later ones with larger deviations. We show that IW-OPD converges significantly faster than OPD, with better learning efficiency, and achieves better final performance than standard OPD in both same-size and cross-scale settings, improving performance up to 6.9 points on AIME-2025.

Yan Xie, Sijie Zhu, Tiansheng Wen +2

Xidian University · Georgia Institute of Technology · Amazon AGI SF Lab

Importance Weighting Language Model Distillation On-Policy Distillation RL Post-Training Teacher-Student Learning Policy Learning

Ask PDF

Sep 28, 2026

On-Policy or Off-Policy Learning? A Systematic Study of Distillation Dynamics

On-policy learning has been argued to reduce catastrophic forgetting, produce sparser parameter updates, and improve generalisation. However, existing comparisons between supervised fine-tuning and reinforcement learning vary many factors simultaneously, making the contribution of rollout policy difficult to isolate. We study the effect of rollout policy in a controlled strong-to-weak distillation setting, by independently varying rollout policy, token-level KL direction, and learning rate across the Llama3 and Qwen2.5 model families and reasoning tasks spanning scientific, medical, and arithmetic domains. Our analysis reveals a nuanced picture of distillation dynamics in which rollout policy does not necessarily play a central role. Instead, token-level KL direction more clearly shapes task performance and output coverage, while learning rate governs forgetting and update sparsity. Analysis of KL gradients and experiments along a continuous student-teacher rollout-policy spectrum explain this pattern: forward KL is remarkably robust to rollout policy, with its performance stable and strong despite changes to the rollout policy, whereas reverse KL is substantially more sensitive and favours student-generated rollouts. On-policy data nevertheless improves generalisation to harder variants of the Countdown arithmetic task under both KL directions, although this advantage does not reliably persist after subsequent RLVR. Our broader conclusions remain robust to removing gradient clipping, using sampled KL estimators, and training on tasks requiring longer reasoning chains. Overall, our results challenge the view that on-policy rollouts are inherently preferable and show that their value depends critically on the objective, evaluation setting, and optimisation hyperparameters.

Julianna Piskorz, Antonin Berthon, Mihaela van der Schaar

University of Cambridge Cambridge, UK

Language Model Distillation On-Policy Distillation Policy Learning Knowledge Distillation

Ask PDF

Abstract

Selective Knowledge Distillation

Explore similar work

Not Every Token Is Worth Distilling: Selective Supervision for Direct-OPD

On the Position Bias of On-Policy Distillation

On-Policy or Off-Policy Learning? A Systematic Study of Distillation Dynamics