PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding On-Policy Distillation: A Mechanistic Interpretability Perspective via Sparse Crosscoders

Zichao Yu, Qianshuo Ye, Xu Wang, Difan Zou

arXiv:2609.35210v1cs.CL

TL;DR

Prior analyses focus on student outputs and do not reveal what changes inside the student during OPD. The swap readout uses sparse crosscoders to measure a student checkpoint’s feature activations by placing it in both student slots while holding the teacher activation fixed. For features firing at least 10 times before OPD, none are absent afterward or vice versa; 98.2%, 99.6%, and 99.4% of frequently used features change firing rate by less than 20%.

  • Problem

    Prior analyses focus on student outputs and do not reveal what changes inside the student during OPD.

    Source: §1 · p. 2

  • Method

    The swap readout uses sparse crosscoders to measure a student checkpoint’s feature activations by placing it in both student slots while holding the teacher activation fixed.

    Source: §2.2 · p. 4

  • Results

    For features firing at least 10 times before OPD, none are absent afterward or vice versa; 98.2%, 99.6%, and 99.4% of frequently used features change firing rate by less than 20%.

    Source: §3.1 · p. 5

  • Takeaways & Limitations

    OPD reweights features already shared with the teacher, while the warm-up reweights shared features partly along OPD’s direction and partly beyond it.

    Source: §1 · p. 3, §3.1 · p. 6

    The findings about warm-up concern SFT on the OPD teacher’s rollouts for the same math-reasoning task, not SFT in general.

    Source: §4.1 · p. 8

Abstract

from arXiv · show

On-policy distillation (OPD) is a widely adopted post-training technique for LLM reasoning. It is commonly believed to transfer knowledge from a stronger teacher, yet what OPD actually distills into the student's internal representations remains unclear. We study this question with sparse crosscoders, which learn one feature dictionary shared by the student before and after OPD and the teacher. Standard crosscoder analyses, however, identify model-specific features but cannot tell how a model's use of its features changes, since all models are encoded into one set of feature activations. We therefore propose the swap readout, which reads each student checkpoint's feature activations on its own, measuring how training changes the student's use of each feature, even for checkpoints unseen by the crosscoder. Across three OPD settings, we find that OPD neither creates features nor passes on the teacher's own, and leaves the firing rates of over 98% of the student's frequently used features within 20%. We further examine the SFT warm-up on the teacher's rollouts that commonly precedes OPD and makes it more effective. Rather than adding features, the warm-up reweights the shared ones in two ways. First, it already raises and lowers many of the features that OPD later raises and lowers, doing part of OPD's work in advance. Second, it changes features that OPD alone would not, notably those for conversation format, reasoning style, and mathematical notation, and these changes persist through OPD. Imposing this reweighting on a directly distilled student's features, without changing its weights, brings its accuracy close to that of the warmed-up student, whereas the same change on shuffled features does not. Together, these findings suggest that OPD reweights existing features rather than acquiring new ones: the student learns from the teacher how to use the features they already share.

1 INTRODUCTION

The paper asks what OPD changes inside a student and introduces the swap readout to track feature use across checkpoints. Across three settings, OPD reweights shared features rather than acquiring new ones, while the teacher-rollout warm-up reweights shared features partly along and partly beyond OPD’s direction.

Motivation and approach

  • The swap readout measures how training changes a student’s use of each shared crosscoder feature, including for checkpoints unseen by the crosscoder.It places the checkpoint activation in both student slots while holding the teacher activation fixed.Source: §1 · pp. 2–3

Main findings

  • Across three OPD settings, the student acquires no new features or teacher-owned features; over 98% of its frequently used features change firing rate by less than 20%.The largest changes occur at decision tokens, where the teacher disagrees most with the student.Source: §1 · p. 2

Warm-up findings

  • The teacher-rollout warm-up adds no features, instead reweighting shared features partly along OPD’s direction and partly in directions that persist through OPD.This warm-up is commonly used before OPD and makes it more effective.Source: §1 · pp. 2–3
  • Imposing the warm-up’s feature reweighting on a directly distilled student, without changing its weights, recovers most of the warm-up’s benefit; shuffled-feature changes do neither.The reported accuracy effect supports a functional role for the reweighting rather than feature acquisition.Source: §1 · p. 3

2 MEASURING FEATURE USAGE WITH SPARSE CROSSCODERS

Sparse crosscoders provide a shared feature dictionary for comparing models, while the swap readout addresses their inability to track changes in a student’s feature use. It reads checkpoints under the same student autoencoder and teacher input, enabling feature-use comparisons across training.

2.1 SPARSE CROSSCODERS AND MODEL ATTRIBUTION

  • Crosscoders jointly encode several models into one sparse feature code and reconstruct each model’s activations with its own decoder directions.Each feature has a shared activation but a separate decoder direction for each model.Source: §2.1 · p. 3
  • Model attribution score (MAS) assigns a feature’s decoder-norm share to each model, with width correction; a score above 0.5 means that model dominates the feature.A feature used equally by M models has MAS 1/M for each.Source: §2.1 · p. 3

2.2 THE SWAP READOUT: TRACKING FEATURE USAGE ACROSS CHECKPOINTS

  • Joint encoding and decoder attribution cannot reliably show which student checkpoint activates a feature, so the swap readout places the checkpoint activation in both student slots.Using the same activation in both slots cancels the undetermined difference between the student encoders; the teacher activation remains fixed.Source: §2.2 · p. 4
  • The readout treats firing features as those used to reconstruct the checkpoint’s activation, and compares firing outcomes across inputs to track changes in use.It records whether a feature fires in neither checkpoint, in both with different strength, starts firing, or stops firing.Source: §2.2 · p. 4
  • The swap readout is matched to crosscoder training conditions and reconstructs both trained and unseen checkpoints, including the student after SFT warm-up.Feature statistics use the same 3M held-out tokens for every model.Source: §2.2 · p. 4

3 WHAT CHANGES INSIDE THE STUDENT DURING OPD?

Across three OPD settings, the student retains its frequently used features and makes only small changes to how it uses them. The largest changes concentrate on features active at reasoning decision tokens.

3.1 OPD REWEIGHTS SHARED FEATURES RATHER THAN ACQUIRING NEW ONES

  • In JustRL, Skywork, and R1-7B, 98.2%, 99.6%, and 99.4% of frequently used features change firing rate by less than 20%, with no frequent features gained or lost.The swap readout found no feature firing at least 10 times before OPD that was absent afterward, or vice versa.Source: §3.1 · p. 5
  • Where Skywork and R1-7B teachers dominate features, the student’s MAS on those features remains 0.19 before and after OPD.The teacher’s own features stay with the teacher; no feature’s MAS changes by more than 0.015.Source: §3.1 · p. 5
  • OPD changes which shared features fire on a token, while each feature’s overall usage stays the same; about one in ten active features changes.The paper characterizes this as a small reweighting rather than acquiring new features.Source: §3.1 · p. 6

3.2 DECISION TOKENS MARK THE LARGEST FEATURE CHANGES

  • Decision-token features make up 12% of JustRL’s and 4% of Skywork’s 50 most-changed features, against a 0.4% baseline; R1-7B has none.Decision tokens mark points where a reasoning trace reflects, hesitates, or moves on.Source: §3.2 · p. 6, Figure 4 · p. 6
  • Under JustRL, OPD shifts the Wait feature 35% less often, the So/Therefore feature 35% less, and the maybe/perhaps feature 42% more, toward teacher usage.The teacher’s corresponding shifts are 34% less, 36% less, and 52% more, respectively.Source: §3.2 · p. 6, Figure 4 · p. 6
  • At decision-token steps, the student replaces 10–44% more active features, and those steps are 4–7 times over-represented among its 5% most-changed steps.Teacher–student KL at these steps is 3.0–3.6, 2.1–2.8, and 1.5–1.7 times average under JustRL, Skywork, and R1-7B, respectively.Source: §3.2 · p. 6, Figure 4 · p. 6
  • Under R1-7B, the student’s next-token distribution departs from the base by KL 0.008 on average, versus 0.085 under JustRL, leaving its features essentially unchanged.This setting therefore shows little movement compared with JustRL.Source: §3.2 · p. 7

4 WHY AN SFT WARM-UP HELPS OPD

The SFT warm-up improves OPD’s results without adding student or teacher features; instead, it reweights shared features, partly in OPD’s direction and partly beyond it. The added reweighting persists through OPD, and transferring it to a directly distilled student brings its accuracy close to that of the warmed-up student.

4 WHY AN SFT WARM-UP HELPS OPD

  • With warm-up, OPD recovers more of the teacher’s advantage over the base student and improves on every benchmark.The section then tests whether this benefit comes from new features or changes to existing ones.Source: §4 · p. 7

4.2 THE WARM-UP ANTICIPATES AND COMPLEMENTS OPD’S REWEIGHTING

  • The warm-up moves shared features in OPD’s direction: β = 0.46, and 88% and 92% of the 50 most-raised and most-lowered features move the same way.For the 50 most-raised features, the warm-up raises firing rate by 9.3% versus OPD’s 17.3%, then OPD adds another 4.5 points; for the 50 most-lowered, the warm-up lowers it by 11.1% versus 20.6%, then OPD adds 7.2 points. A random-direction shift shows no alignment (β = −0.03).Source: Figure 6 · p. 8, §4.2 · pp. 8–9

4.1 THE WARM-UP GIVES THE STUDENT NO NEW FEATURES EITHER

  • The warm-up creates no student features and does not acquire the teacher’s features.The crosscoder tests found no warm-up-specific features and no material shift in the teacher’s feature distribution.Source: §4.1 · p. 8
  • These conclusions concern a warm-up that imitates the OPD teacher’s rollouts for the same math-reasoning questions, not SFT in general.Source: §4.1 · p. 8

4.2 THE WARM-UP ANTICIPATES AND COMPLEMENTS OPD’S REWEIGHTING

  • Adding the warm-up’s feature change raises direct OPD’s avg@8 from 18.6% to 21.3%, near the warmed-up student’s 21.9%; removing it lowers avg@8 to 19.2%.The changes are imposed on student features without changing weights; shuffling feature identities neither helps the directly distilled student nor hurts the warmed-up one.Source: Table 3 · p. 9, §4.2 · p. 9
  • 56% of the warm-up’s change lies outside OPD’s direction (77% including one nearly silenced garbled-character feature), and this residual persists through OPD (Spearman 0.69).The residual spans many features, including conversation format, reasoning style, and mathematical notation.Source: §4.2 · p. 9

5 CONCLUSION

The paper concludes that OPD primarily reweights features the student already shares with the teacher, rather than acquiring new ones. It positions this analysis alongside prior work on OPD’s optimization and sparse crosscoders’ feature comparisons.

5 CONCLUSION

  • Across the studied settings, OPD creates no student features and passes on none of the teacher’s; it slightly reweights shared features, while the warm-up reweights them partly along and partly beyond OPD’s direction.Imposing the warm-up reweighting on a directly distilled student brings its accuracy close to that of the warmed-up student.Source: §5 · p. 10

A RELATED WORK

  • Prior OPD studies examine optimization and teacher–student overlap, while reports also describe length inflation, repetition, agreement traps, and unreliable supervision on student-induced prefixes.The paper notes that small local teacher–student discrepancy need not imply desirable complete rollouts.Source: A RELATED WORK · p. 13
  • Sparse crosscoders compare shared and model-specific features; this work uses a swap readout to track how OPD and its warm-up change student feature usage.Prior work cited here applies crosscoders to reasoning features and comparisons of SFT and RL.Source: A RELATED WORK · p. 13

B THEORETICAL ANALYSIS OF THE SWAP READOUT

The swap readout separates the student’s determined shared encoder component from an undetermined component tied to differences between checkpoints. It reads each checkpoint through both student slots so changes reflect the student’s own activation shift.

B THEORETICAL ANALYSIS OF THE SWAP READOUT

  • The student-slot contribution can be written as S m + N δ, separating the students’ mean activation from their difference.Here S is the sum of the student encoders, N half their difference, m their mean activation, and δ their activation difference.Source: B THEORETICAL ANALYSIS OF THE SWAP… · pp. 13–14
  • The swap readout places a checkpoint’s rescaled activation in both student slots, canceling the undetermined encoder difference N.Its pre-activation is S times the checkpoint activation plus the teacher and bias contribution.Source: B THEORETICAL ANALYSIS OF THE SWAP… · pp. 13–14
  • Training leaves N largely undetermined because its gradient uses the smaller difference between student activations rather than their mean.The gradients are E[g m⊤] for S and E[g δ⊤] for N; N also retains much of its random initialization.Source: B THEORETICAL ANALYSIS OF THE SWAP… · p. 14
  • A checkpoint change projects the student’s activation shift onto the determined encoder direction S, with no contribution from the teacher or N.The swap readouts of two checkpoints average to the joint pre-activation up to a term acting only on their difference.Source: B THEORETICAL ANALYSIS OF THE SWAP… · p. 14
  • On held-out tokens, N remained near its random-initialization scale, while the joint encoding’s undetermined term contributed only 4–6% of pre-activation.The joint encoding reconstructs all three models well but cannot attribute a change to either student.Source: B THEORETICAL ANALYSIS OF THE SWAP… · p. 14

C.1 OPD TRAINING

The OPD runs use a shared vanilla training recipe, and the warm-up fine-tunes the student on selected correct teacher rollouts before applying OPD with the same teacher and settings.

C.1 OPD TRAINING

  • All five OPD runs use the vanilla recipe with shared hyperparameters, one pass over DAPO-Math-17k, and final-checkpoint analysis.At each step, the current student samples four rollouts for each of 64 prompts, and the frozen teacher scores each rollout position.Source: §C.1 · p. 15
  • The OPD signal is reverse KL between the student and teacher next-token distributions, approximated over the student’s 16 most probable tokens, without added loss or reward.The approximation weights the teacher–student log-ratio by student probabilities renormalized over those tokens.Source: §C.1 · p. 15

C.2 SFT WARM-UP

  • The warm-up fine-tunes on 704 selected correct teacher responses, sampled from 2,048 prompts and retained after response-quality filtering.It uses LoRA rank 32 for 175 steps, about two epochs; OPD then runs from this student with the same teacher, prompts, and hyperparameters as without warm-up.Source: §C.2 · pp. 15–16

C.3 EVALUATION

The paper evaluates accuracy on three math benchmarks and trains crosscoders on shared activations to compare students and teachers. The swap readout reconstructs student checkpoints with reported variance explained on held-out tokens.

C.3 EVALUATION

  • Accuracy is measured on AIME 2024, AIME 2025, and AMC 2023; the paper reports avg@n, the mean across benchmarks, and recovered teacher advantage.The Section 3 settings use n = 256 samples per problem, while Section 4 uses n = 8.Source: §C.3 · p. 16

C.4 CROSSCODERS

  • Crosscoders share 400M training tokens, activation sites, and hyperparameters, with about 3.0M held-out tokens used for feature statistics.They use residual-stream activations from middle blocks, centered and scaled to unit total variance per model.Source: §C.4 · p. 16
  • The crosscoder analysis uses three-model crosscoders for student, OPD student, and teacher comparisons, plus additional Qwen3 crosscoders for warm-up analyses.The additional crosscoders compare teacher with students before and after warm-up, and compare the students before and after warm-up.Source: §C.4 · pp. 16–17
  • On 0.2M held-out tokens, the pre-warm-up crosscoder’s swap readout explains 0.755, 0.760, 0.756, and 0.759 of variance for the base, OPD, warmed-up, and warmed-up-plus-OPD students.The reported reconstruction quality is measured as fraction of variance explained.Source: §C.4 · p. 17, Table 9 · p. 17

D ANALYSIS DETAILS

The analysis measures feature-use changes on held-out reasoning text and separately examines policy changes on the models’ own generations. It specifies how feature changes and decision-token over-representation are quantified.

D.1 FEATURE STATISTICS

  • Feature statistics use 3,019,499 held-out token positions; a feature fires when its code entry is positive, and its firing count is the number of firing positions.The held-out positions come from 5,909 sequences, about two thirds from OpenThoughts.Source: §D.1 · p. 17
  • Across JustRL, Skywork, and R1-7B, the share of features changing on a token averages 0.12, 0.08, and 0.08, respectively.The share is one minus the Jaccard similarity of feature sets active before and after training; the analysis includes features firing at least 100 times across both checkpoints.Source: §D.1 · p. 17

D.2 DECISION TOKENS AND POLICY CHANGES

  • Decision-token categories include reflection, hesitation, and progression words, and features are identified from their 12 strongest firing positions in the pre-OPD student.The analysis compares changed-feature shares at decision-token positions with the average across held-out reasoning text.Source: §D.2 · p. 18
  • The policy analysis uses the base student’s own generations, measuring teacher signal and OPD change as KL divergences from the base student’s next-token distribution.It defines over-representation as decision-token share among the 5% most-changed steps divided by that share among all steps.Source: §D.2 · p. 18

D.3 DECOMPOSING THE WARM-UP’S REWEIGHTING

D.3 separates the warm-up’s feature changes into an OPD-aligned component and a residual, then tests the residual against controls. D.4 describes a feature-level intervention and how its effects are read.

D.3 DECOMPOSING THE WARM-UP’S REWEIGHTING

  • The decomposition fits the warm-up’s feature-change vector to the OPD-change vector, then measures the residual’s share and its correlation with the post-OPD residual.It uses 8,524 features that fire at least 1,000 times in both the base and OPD students.Source: §D.3 · p. 18
  • Without the dominant feature, reasoning-focused features carry 50% of the warm-up residual, general-text features 22%, and other features 28%.The dominant feature accounts for 62% of the residual; the residual share is 56% without it and 77% with it.Source: §D.3 · p. 18
  • Random-shift and Gaussian-noise controls show no alignment, while disjoint-shard measurement gives β = 0.39 versus 0.46 on the same shard.The random shift has β = −0.03 and residual correlation 0.00; Gaussian noise gives β = −0.02.Source: §D.3 · pp. 18–19

D.4 FEATURE-LEVEL INTERVENTION

  • The intervention decodes the difference between warm-up and base-student swap codes into a residual-stream change at block 14, recomputed during generation.A shuffled control permutes code differences across features before decoding; the resulting change is 4.5% of residual-stream norm versus 6.2% for the warm-up change.Source: §D.4 · p. 19

E FEATURE VISUALIZATIONS

  • Feature cards show strongest contexts, token-level activation shading, and firing counts; most are read with the swap readout on held-out tokens.Teacher features in Figure 8 are the exception.Source: E FEATURE VISUALIZATIONS · p. 19

E.1 THE SWAP READOUT ON HELD-OUT PASSAGES

The swap readout compares student feature activations before and after JustRL OPD on held-out reasoning passages. The examples show both changed firing behavior and movement toward teacher firing counts.

E.1 THE SWAP READOUT ON HELD-OUT PASSAGES

  • The Wait feature stops firing on some Wait tokens, the So/Therefore feature drops some conclusion-adjacent firings, and the maybe/perhaps feature begins firing on hedging phrases.The Wait feature fires on decision tokens for 86% of its held-out reasoning-text firings; the other two have 18% of firings on decision tokens, nine times their share in the text.Source: §E.1 · p. 19
  • All three example features move their firing counts toward the teacher’s after OPD.The feature-specific before/after lines read the same passage with different student activations.Source: §E.1 · p. 19, Figure 7 · p. 20
  • Figure 7 pairs pre- and post-OPD student readouts for each passage, shading token activations and bolding tokens where a feature starts or stops firing.Its right side gives firing counts before OPD, after OPD, and in the teacher.Source: Figure 7 · p. 20

E.2 FEATURES THE TEACHERS DOMINATE

The feature visualizations distinguish teacher-dominated features, OPD-shifted decision-token features, and two kinds of warm-up reweighting. The warm-up shifts some features along OPD’s direction and others outside it.

E.2 FEATURES THE TEACHERS DOMINATE

  • Teacher-dominated features respond to titles, religious texts, LaTeX, dates, place names, and formula notation, while the student’s decoder-norm share stays the same before and after OPD.Figure 8 reads these features with the joint encoding of the student and both teachers.Source: Figure 8 · p. 20, §E.2 · p. 20

E.3 DECISION-TOKEN FEATURES

  • Under JustRL, OPD moves decision-token features toward teacher firing counts; these features fire most strongly on the words where the reasoning trace turns.Figure 9 covers six such features under JustRL and two under Skywork among the 50 features whose firing rates change most.Source: §E.3 · pp. 20–21, Figure 9 · p. 21

E.4 FEATURES THE WARM-UP MOVES ALONG OPD’S DIRECTION

  • For Qwen3, the warm-up moves six shown features in OPD’s direction, and 35 of the 40 most OPD-shifted features move the same way after warm-up.The warm-up shift is smaller, but after warm-up and OPD the change matches or exceeds OPD alone.Source: §E.4 · p. 21, Figure 10 · p. 22

E.5 FEATURES THE WARM-UP MOVES BEYOND OPD’S DIRECTION

  • Warm-up shifts outside OPD’s direction include features for conversation format, reasoning style, and mathematical notation, with some changes running against OPD or exceeding it.Wait-related features can have both aligned and outside-direction components.Source: §E.5 · p. 21, Figure 11 · p. 22
Loading 2609.35210v1…