PaperlayerView exact revision on arXiv

Source-linked AI summary

OPRD: On-Policy Representation Distillation

Shenzhi Yang, Guangcheng Zhu, Bowen Song, Haobo Wang, Mingxuan Xia, Xing Zheng, Yingfan Ma, Zhongqi Chen, Weiqiang Wang, Junbo Zhao, Gang Chen

arXiv:2606.06021v4cs.LGcs.AI

TL;DR

Output-only on-policy distillation is limited by gradient variance and an LM-head bottleneck that omits intermediate representations. OPRD aligns hidden states on-policy, while OPRD-Bridge extends this across mismatched architectures and tokenizers, with a rank-8 bridge reaching 95% cosine similarity between projected representations.

  • Problem

    Existing on-policy distillation variants all supervise through next-token distributions, limiting access to intermediate representations and exposing output-space training limitations.

  • Method

    OPRD aligns teacher and student hidden representations across selected layers and positions, while OPRD-Bridge uses frozen low-rank projectors across heterogeneous architectures and tokenizers.

  • Results

    95% cosine similarity is achieved by a rank-8 bridge between Qwen3-4B and Qwen3-1.7B representations, with successful transfer also demonstrated across disjoint tokenizers.

  • Takeaways & Limitations

    Representation-space alignment decouples distillation from a shared vocabulary and supports knowledge transfer across architecture and tokenizer mismatches.

  • Takeaways & Limitations

    Output-space OPD variants remain constrained by gradient variance, truncation bias or memory cost, and an LM-head information bottleneck.

Abstract

from arXiv · show

On-policy distillation (OPD) supervises the student exclusively in the output space by matching next-token distributions. This paradigm suffers from two limitations: (i) a high-variance gradient estimator whose signal-to-noise ratio collapses as the student approaches the teacher, and (ii) an LM-head information bottleneck that discards the teacher's intermediate hidden states. We propose On-Policy Representation Distillation (OPRD), the first method to lift on-policy distillation into the hidden-state space. OPRD aligns student and teacher representations across selected layers on the same on-policy rollouts, providing dense, deterministic, per-layer supervision while bypassing the LM head entirely. Theoretically, OPRD provides a deterministic per-sample gradient, removing the token-level estimation variance that plagues OPD, and exposes structural information that any output-space objective necessarily discards. Empirically, OPRD closes the student-teacher gap on competition mathematics benchmarks (AIME 2024, AIME 2025, and AIMO), where every output-space baseline plateaus below the teacher, while training 1.44x faster and using up to 54% less memory. We further extend OPRD to the cross-architecture setting via OPRD-Bridge. By exploiting the observation that heterogeneous models share a low-rank representational structure, we construct a frozen projector pair that aligns representations across arbitrary depth and width mismatches, shifting the alignment from the output space (which depends on a shared vocabulary) to the representation space. We validate OPRD-Bridge on both cross-architecture (Qwen3-4B -> Qwen3-1.7B-Base) and cross-tokenizer (Phi-4-mini-reasoning -> Qwen3-1.7B-Base) settings, demonstrating successful knowledge transfer even when the vocabulary-based alignment channel is unavailable. Code: https://github.com/ShenzhiYang2000/OPRD.

1 Introduction

The introduction argues that output-space OPD suffers from late-stage gradient variance and an LM-head information bottleneck. It proposes OPRD for deterministic hidden-state supervision and OPRD-Bridge for mismatched architectures and tokenizers.

  • Motivation: OPD variants differ in token coverage but remain confined to the same output space, limiting the design space.The variants include sampled-token, full-vocabulary, and top-k objectives.
  • Motivation: As the student approaches the teacher, OPD’s expected gradient decays as O(δ^2) while variance decays as O(δ), collapsing signal-to-noise.Here δ denotes the student–teacher gap.
  • OPRD-Vanilla: OPRD aligns student and teacher hidden representations across selected layers and response positions on the same on-policy rollouts.This moves supervision from the LM-head output to its input and uses normalized mean-squared alignment.
  • OPRD-Bridge: OPRD-Bridge uses frozen linear projector pairs to map teacher and student representations into a shared low-rank subspace across architecture mismatches.The teacher projector comes from PCA, while the student projector is trained to align with it before main distillation.
  • OPRD-Vanilla: OPRD provides deterministic per-sample gradients and exposes per-layer structural information that the LM-head projection discards.Its supervision spans layers, positions, and hidden dimensions rather than only output probabilities.

2 Background and Problem Setup

On-policy distillation trains on student-generated rollouts and matches teacher and student next-token distributions, but practical output-space variants trade supervision density against vocabulary-dependent computation. Because all supervision passes through the LM head, intermediate teacher representations are excluded and sampled-token OPD becomes noisy near convergence.

  • On-Policy Distillation: OPD samples responses from the student and evaluates both models on the resulting prefixes, aligning supervision with inference-time states.This removes the exposure-bias gap between fixed-target distillation and inference distributions.
  • Objective: The canonical OPD objective is trajectory-level reverse KL, which decomposes exactly into token-level reverse KL terms but requires summing over the full vocabulary.That full-vocabulary computation is prohibitive for modern LLMs with vocabularies in the hundreds of thousands.
  • Output-Space Variants: Sampled-token, full-vocabulary, and top-k OPD differ by the token support used for per-position supervision.Sampled-token OPD uses one rollout token, full-vocabulary OPD uses V, and top-k OPD uses the student’s k highest-ranked tokens.
  • Output-Space Variants: Full-vocabulary OPD requires O(BT|V|) memory, while top-k OPD trades supervision density against teacher-query cost with typical k ∈[4, 64].At k = 1, top-k recovers a deterministic version of sampled-token OPD; at k = |V|, it recovers full-vocabulary OPD.
  • Structural Limitation: All prior variants receive teacher knowledge only through the LM-head projection Whead, so intermediate hidden states never enter the loss.Sampled-token OPD additionally estimates each token-level KL from one Monte Carlo draw, with variance worsening as |V| grows and pt approaches qt.

3 On-Policy Representation Distillation

OPRD lifts on-policy distillation from next-token distributions into hidden-state representations, providing dense deterministic supervision that addresses OPD’s gradient-variance and LM-head bottleneck limitations. Its low-rank OPRD-Bridge extension aligns heterogeneous models in representation space, including when tokenizers differ.

  • OPRD formulation: OPRD supervises intermediate student representations on student-generated trajectories instead of matching next-token distributions.This hidden-state formulation exposes structural information compressed by the LM head and is called OPRD-Vanilla in the same-architecture setting.
  • OPRD formulation: The loss flexibly selects distilled layers and response positions, typically using all layers and the last k response tokens for long mathematical-reasoning responses.Position masking excludes invalid suffix positions, while layer and position averaging normalize the objective.
  • Theoretical properties: OPRD provides a deterministic per-sample gradient with zero conditional variance, yielding a stable signal when the student and teacher output distributions nearly converge.Sampled-token OPD retains variance that can dominate late-training optimization and contribute to stagnation.
  • Theoretical properties: Output-space OPD can leave large hidden-state deviations unconstrained because ill-conditioned LM-head directions produce similar outputs, whereas OPRD penalizes these directions and supervises intermediate states.The output-space loss cannot constrain hidden states at layers before the final layer.
  • OPRD-Bridge: PCA is rate–distortion optimal among rank-r linear teacher encodings under Gaussian assumptions, and representation-space alignment decouples the bridge from a shared tokenizer.The theory explains a low-rank optimum near r∗≈8 and why freezing the bridge is critical.

4 Experiments

Section 4 evaluates both OPRD variants: OPRD-Vanilla for same-architecture distillation and OPRD-Bridge for cross-architecture distillation.

  • 4 Experiments: The experiments evaluate both OPRD variants.The two variants are OPRD-Vanilla and OPRD-Bridge.
  • 4.1 Same-architecture setting: OPRD-Vanilla is evaluated in the same-architecture setting.This setting is covered in §4.1.
  • 4.2 Cross-architecture setting: OPRD-Bridge is evaluated in the cross-architecture setting.This setting is covered in §4.2.

4.1 OPRD-Vanilla: Same-Architecture Distillation

OPRD-Vanilla closes the student–teacher gap more effectively than output-space OPD baselines on competition mathematics, while improving training efficiency. It also produces shorter reasoning and directly optimizes the hidden-state alignment objective that avoids the LM-head bottleneck.

  • Accuracy: Both OPD variants improve over the student but plateau below the teacher, while top-1 and top-16 show no consistent ordering.Top-16 improves AIME24 by +4.8 and AIME25 by +0.5 but loses −0.5 on AIMO, consistent with LM-head bottleneck and truncation bias.
  • Accuracy: OPRD reaches 49.8 on AIME24, 34.6 on AIME25, and 79.1 on AIMO, leaving only 1.0 / 1.0 / 0.4 points to the teacher.These results are within the variance of 16-sample Avg@16 evaluation, with AIMO effectively tied; OPRD gains +2.7 / +0.6 / +2.1 points over the better OPD baseline.
  • Training dynamics: OPRD’s accuracy climbs monotonically, whereas both OPD variants plateau or oscillate after early gains.Figure 4 compares the methods under shared initialization and on-policy training conditions.
  • Behavioral analysis: OPRD converges to a mean response length of ∼5,700 tokens, versus ∼7,000 tokens for both OPD variants.Combined with higher accuracy, the shorter rollouts indicate more efficient reasoning under OPRD.
  • Efficiency: 20.5 GB of OPRD transient memory is 32% below OPD top-1’s 30.2 GB and 54% below OPD top-16’s 45.0 GB.The savings result from eliminating the [B, T, |V|] logits tensor and its top-k gradient buffer from the OPRD loss path.
  • Efficiency: 563 minutes for OPRD is 31% below OPD top-1’s 813 minutes and OPD top-16’s 812 minutes, equivalent to a 1.44× speed-up.The methods use identical 500-step schedules, rollouts, and teacher forward passes; OPD time is dominated by logits computation and log_softmax.

4.2 OPRD-Bridge: Cross-Architecture and Cross-Tokenizer Distillation

OPRD-Bridge enables hidden-state distillation across mismatched architectures and tokenizers by aligning representations through a learned low-rank projector pair. It matches output-space distillation’s capability ceiling while producing more diverse, shorter reasoning, and combining it with OPD further improves the ceiling.

  • Cross-Architecture Results: OPRD-Bridge matches OPD’s Best@16 capability ceiling, achieving 20.0 on AIME24 and 13.3 on AIME25 while operating through a rank-8 subspace.These scores match OPD top-1 exactly and substantially outperform OPD top-16.
  • Cross-Architecture Results: 22.6/17.2/21.3 Dist-4g scores on AIME24/AIME25/AIMO exceed OPD top-1 and OPD top-16, approaching the teacher’s 24.0/25.4/41.6.OPRD-Bridge also produces 5,909/5,234/4,855-token responses, substantially shorter than both OPD baselines.
  • Two-Stage Distillation: The OPRD-Bridge→OPD two-stage pipeline raises average Best@16 from 30.8 to 33.9, surpassing standalone OPD top-1 at 30.8 and OPD top-16 at 27.1.OPRD-Bridge aligns backbone representations, while OPD calibrates the output distribution and improves consistency.
  • Bridge Analysis: After bridge training, representation similarity reaches 72.3% at rank 1 and exceeds 90% for r ≥4, peaking at 95.0% for r = 8.Similarity declines to 94.1% at r = 32, 87.9% at r = 256, and 77.0% at full rank r = 2048; the practical range is r=4–16.
  • Cross-Tokenizer Results: Across incompatible tokenizers, representation-space alignment lifts Avg@16 from 2.8% to 5.0% and improves average Best@16 from 4.5% to 21.3%, a 4.7× increase.Best@16 reaches 13.3% on both AIME24 and AIME25 and 37.3% on AIMO, demonstrating capability transfer without vocabulary alignment.

5 Discussion

The discussion frames OPRD around representational alignment: OPRD-Vanilla exploits alignment in shared architectures, whereas OPRD-Bridge addresses mismatched architectures. It highlights applications in multi-model RL merging and self-distillation, and proposes establishing alignment during pre-training to reduce reliance on bridging.

  • Representational alignment: OPRD-Vanilla suffices when teacher and student share architecture and initialization, while OPRD-Bridge constructs a low-rank, low-pass channel when depth, width, or training history differ.The distinction is governed by the degree of representational alignment between teacher and student.
  • Multi-model RL merging: OPRD-Vanilla targets multi-model RL merging, where full-vocabulary OPD requires materializing a [B, T, |V|] logit tensor and incurs prohibitive transient GPU memory costs.The setting assumes models share a backbone, making same-architecture hidden-state alignment applicable.
  • On-policy self-distillation: OPRD naturally fits on-policy self-distillation because teacher and student share exactly the same weights, making hidden-state alignment maximally informative.The teacher can inject privileged information such as ground-truth solutions or step-level verification signals into the prompt.
  • Alignment-aware pre-training: Alignment-aware pre-training could establish a shared representation space through shared initialization, representation-matching regularization, or co-distillation, potentially making the bridge unnecessary.The proposed direction targets models differing in depth and width, such as Qwen3-32B and Qwen3-1.7B.

6 Related Work

This section situates OPRD among output-space, on-policy, and intermediate-representation distillation, emphasizing its distinction from prior methods. It also contrasts OPRD’s student-generated autoregressive supervision with fixed-input and encoder-style feature distillation.

  • Output-Space Knowledge Distillation: Output-space distillation compresses large models by matching output distributions, extending from sequence-level transfer to pretrained and instruction-following language models.
  • On-Policy Distillation: On-policy distillation addresses exposure bias by training on student-sampled responses, including MiniLLM’s reverse-KL objective and GKD’s family of interpolated divergences.
  • Feature / Intermediate-Representation Distillation: Prior feature-distillation methods supervise intermediate representations, including FitNets’ hint layers, attention-transfer maps, FSP matrices, TinyBERT, and Mobile-BERT.
  • On-policy vs. off-policy supervision: Unlike fixed-input feature distillation, OPRD computes hidden-state loss on student-generated sequences whose distribution evolves during training.The teacher is queried on states the student actually visits, making supervision adaptive to the student’s evolving policy.
  • Encoder representations vs. autoregressive prefix representations: Unlike encoder-style feature distillation, OPRD aligns autoregressive prefix representations at every decoding step under the student’s sampling distribution.Each hidden state encodes the model’s belief before predicting the next sampled token, conditioned on the entire sampled prefix.

7 Conclusion and Future Work

The paper presents OPRD as hidden-state-space on-policy distillation and extends it with OPRD-Bridge for heterogeneous architectures. Future work targets broader modalities and tasks, adaptive supervision selection, and deeper use of hidden-state representations.

  • Contributions: OPRD supervises students in hidden-state space rather than at the LM-head output, addressing limitations shared by existing output-space OPD variants.The conclusion identifies sampled-token, top-k, and full-vocabulary OPD as output-space variants with inherent practical limitations.
  • Contributions: OPRD-Bridge constructs a frozen low-rank bridge between heterogeneous representations, shifting alignment from output space to representation space and decoupling distillation from vocabulary.The approach exploits the empirical finding that models with different depth and width share low-rank representational structure.
  • Future Work: Whether OPRD’s gains transfer beyond long-CoT mathematics to code generation, agentic interaction, and open-ended dialogue remains unresolved.These settings may have different position-level supervision characteristics.
  • Future Work: Cross-modal distillation is a natural next step, with representation-space bridging potentially serving as the only viable supervision channel between incompatible modalities.The proposed example is a vision-language teacher transferring to a language-only student.
  • Future Work: Adaptive layer and position weighting could sharpen supervision beyond uniform layer weighting and the simple last-k position heuristic.The proposed criteria are where the student–teacher gap is largest or where the gradient signal is most informative.
  • Broader Implications: The work argues that hidden-state representations are an under-exploited distillation resource and that teachers should be treated as structured sources of layered internal computation.This framing encourages students to learn the teacher’s layered internal computation rather than treating the teacher merely as a probability oracle.

A Notation Summary · B Formal Theoretical Guarantees

The section fixes notation for single-position on-policy analysis and formally establishes OPRD’s deterministic, lower-variance gradients alongside OPD’s persistent variance and collapsing signal-to-noise ratio.

  • B Formal Theoretical Guarantees: The formal guarantees cover OPD gradient variance, signal-to-noise collapse, and the LM-head information bottleneck through dedicated theorems, corollaries, and identities.The listed results include Theorems 4–8 and Corollary 1.
  • A Notation Summary: The notation summary defines p and q as student and teacher next-token policies, θ as student parameters, and u_t as their per-token log-density ratio.Expectations are taken over one fixed prompt and response position; multiple positions follow by linearity.
  • B.1 Setup and Assumptions: Under standard regularity assumptions, the analysis requires smooth log-densities, finite Fisher information, bounded student-teacher log-ratios, and bounded Lipschitz-transformer hidden states.The passages characterize these conditions as mild and standard for softmax LLMs and policy-gradient analysis.
  • B.2 Variance of Sampled-Token OPD: OPD is structurally a REINFORCE estimator with u_t as reward, producing a conditional variance lower bound that does not vanish as p approaches q.Near the optimum, Var[g_OPD] remains Ω(Var_p(u_t)).
  • B.3 OPRD Gradient Is Deterministic: OPRD’s per-position gradient has zero conditional variance because aligned hidden states are fixed under the prompt and prefix conditioning.The variance gap shows OPRD is always lower variance than OPD except when p is a point mass.
  • B.4 Signal-to-Noise Ratio Collapse of Sampled-Token OPD: As the symmetric divergence δ approaches zero, OPD’s signal-to-noise ratio satisfies SNR(g_OPD) = O(δ) → 0.The result formalizes late-stage degradation as the student approaches the teacher.
  • B.4 Signal-to-Noise Ratio Collapse of Sampled-Token OPD: OPRD has identically zero gradient covariance and therefore an infinite signal-to-noise ratio in the extended-real sense while its hidden-space loss remains unconverged.Its gradient signal is completely noise-free under the stated condition.
  • B.4 Signal-to-Noise Ratio Collapse of Sampled-Token OPD: The theory predicts effective early OPD learning followed by late stagnation, whereas OPRD retains an informative descent direction and supports monotonic improvement until hidden-space convergence.Pure OPD may plateau or oscillate below the teacher as δ approaches zero.

B.5 Formal Results for Theorem 2: LM-Head Information Bottleneck

The formal results show that output-space distillation can be blind to hidden-state directions and disproportionately tolerate deviations suppressed by the LM head. OPRD addresses this bottleneck by directly constraining intermediate representations.

  • Null-direction identity: Every output-space distillation objective treats the LM head’s effective null space as invisible, yielding zero loss when hidden-state differences map to softmax-invariant logits.This follows from softmax invariance to additive all-ones logit shifts.
  • Regularity conditions: The formal guarantees hold under mild regularity conditions, including compact logit regions for the local Lipschitz bound used by standard output-space losses.The bound covers sampled-token, top-k, and full-vocabulary reverse-KL objectives.
  • Spectral gap: The LM head’s singular-value spread σ1/σd amplifies hidden-state deviations along vd by a (σ1/σd)^2 factor for the same output-space loss budget.The result applies under the local Lipschitz bound established for sampled-token, top-k, and full-vocabulary reverse-KL losses.
  • Intermediate layers: Output-space objectives cannot constrain intermediate hidden states when perturbations leave the final hidden state unchanged, whereas OPRD directly penalizes those layer representations.The limitation arises because output-space losses depend solely on Whead h(L).
  • Spectral gap: 10^6 ∼10^8× is the empirically observed amplification factor for production LLMs implied by the LM head’s singular-value spread.This figure summarizes the spectral-gap effect reported in the theorem sub-summary.

B.6 Formal Results for Theorem 3: Optimality of the Low-Rank Bridge

The section formalizes PCA-based low-rank bridging as information-theoretically optimal and characterizes bridge rank as controlling a bias–variance trade-off. Under monotone alignment decay, the bridge preserves the teacher’s most predictable, high-variance directions.

  • B.6.1 Setup: The analysis models teacher hidden states through their covariance eigendecomposition and student–teacher alignment in the teacher’s principal-coordinate system.Each principal direction has a corresponding prediction correlation ρ_i from the student representation.
  • B.6.1 Setup: The monotone-alignment assumption states that ρ1 ≥ ρ2 ≥ ··· ≥ ρdT, so high-variance principal directions are better predicted by the student.This assumption links the teacher’s variance ordering to cross-model predictability.
  • B.6.1 Setup: Experiments verify that alignment concentrates in the top principal components and degrades monotonically with increasing rank.The rank-cosine curve in Figure 12 provides the empirical verification.
  • B.6.2 Rate–Distortion Optimality: Theorem 9 states that, for Gaussian teacher states, the PCA projection onto the top-r eigenvectors minimizes reconstruction distortion among all rank-r linear encoders.The optimal encoder is P∗ = [v1, . . . , vr]⊤.
  • B.6.2 Rate–Distortion Optimality: This optimality follows from the Eckart–Young–Mirsky theorem and the classical result that PCA minimizes mean-squared reconstruction error among rank-r linear projections.For Gaussian sources, the result coincides with the squared-error rate–distortion function under water-filling.
  • B.6.2 Rate–Distortion Optimality: The PCA bridge is therefore the information-theoretically optimal linear channel at capacity r, while any other rank-r projection discards more teacher information.This establishes PCA as more than a heuristic for low-rank bridging.

B.6.3 Bias–Variance Decomposition of Distillation Error

The distillation error decomposes into bias from discarded aligned signal and variance from transmitted but unmatchable noise. Its quasi-convex rank dependence yields an optimal bridge rank, empirically supported by Qwen3 experiments peaking at r = 8.

  • Bias–variance decomposition: Theorem 10 decomposes total distillation error into bias and variance under Assumption 2 with a population-optimal student projector.The analysis concerns distillation error, including discarded dimensions, rather than reconstruction error.
  • Bias–variance decomposition: Bias B(r) is aligned signal in discarded directions i > r and decreases as the bridge rank increases.Excluded directions contain signal the student could have matched if those directions were transmitted.
  • Bias–variance decomposition: Variance V(r) is misaligned noise in included directions i ≤ r that the bridge transmits but the student cannot match, increasing with rank.This is the residual prediction error from the included directions.
  • Optimal rank: E_total(r) is quasi-convex, and the optimal rank r* occurs where the marginal inclusion effect changes from negative to positive.The transition is established because the relevant alignment quantity is non-increasing.
  • Empirical verification: r = 8 is the empirical peak of the rank-cosine curve in Qwen3-4B → Qwen3-1.7B experiments, consistent with r* ≈ 8 where ρ_i^2 crosses 1/2.Rapid spectral decay of Σ_T makes both B(r*) and V(r*) small at this operating point.
Loading 2606.06021v4…