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

On-policy distillation has largely confined supervision to next-token distributions, leaving teacher intermediate representations outside the loss. OPRD aligns hidden states on on-policy rollouts, and OPRD-Bridge demonstrates knowledge transfer across architecture and tokenizer differences.

  • Problem

    Existing on-policy distillation variants operate in output space, so teacher intermediate hidden representations do not enter the supervision signal.

    Source: §1 · p. 1, §2.4 · pp. 5–6

  • Method

    OPRD aligns student and teacher hidden representations on shared on-policy rollouts, while OPRD-Bridge uses frozen projectors to align heterogeneous models.

    Source: §1 · pp. 2, 4, §7 · p. 25

  • Results

    OPRD-Bridge demonstrated successful knowledge transfer across both cross-architecture and cross-tokenizer settings, including when vocabulary-based alignment was unavailable.

    Source: §7 · p. 25

  • Takeaways & Limitations

    Representation-space alignment supports distillation between models even when their vocabularies differ.

    Source: §1 · p. 4, §7 · p. 25

    Output-space OPD can have high token-level gradient variance as student and teacher distributions converge, while LM-head projection leaves much of the teacher’s hidden-state space unsupervised.

    Source: §2.4 · pp. 5–6

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

OPRD moves on-policy distillation from next-token distributions to hidden states, addressing OPD’s late-stage gradient variance and loss of intermediate information. OPRD-Bridge extends representation alignment to heterogeneous models through a frozen low-rank projector interface that does not require a shared vocabulary.

Empirical validation

  • OPRD closes the student–teacher gap on AIME 2024, AIME 2025, and AIMO; versus OPD top-16 it gains 2.7 points, trains 1.44× faster, and cuts peak GPU memory by 54%.The figure reports AIME24 Avg@16 for the 2.7-point accuracy gain and actor-update peak GPU memory for the 54% cut; AIME25 and AIMO results are described as qualitatively identical.Source: Figure 1 · p. 2, §1 · p. 4

OPD limitations

  • Sampled-token OPD’s expected gradient decays as O(δ^2) while variance decays as O(δ), collapsing signal-to-noise; top-k trades lower variance for truncation bias and also plateaus.Here δ is the student–teacher gap; sampled-token OPD draws one token from a vocabulary that can be approximately 151K tokens.Source: §1 · p. 2
  • The LM-head projection and softmax discard teacher hidden-state structure, so matching output distributions provides no supervision for how the teacher arrived at them.The discarded states encode attention patterns, mid-layer reasoning state, and geometric structure; the passage notes that similar output distributions can correspond to hidden states differing along affine subspaces.Source: §1 · p. 2

OPRD-Vanilla

  • OPRD-Vanilla aligns student and teacher hidden representations at selected layers and response positions on the same on-policy rollouts, using normalized MSE for deterministic per-sample gradients.It exposes layer-by-position-by-hidden-dimension supervision without the LM-head projection and avoids materializing the [B, T, |V|] logits tensor, reducing time and peak memory.Source: §1 · pp. 2–3

Applications

  • OPRD can replace output-space reverse-KL in multi-model RL merging and on-policy self-distillation, avoiding vocabulary-sized logits and OPD’s variance or truncation limits.For self-distillation, teacher and student share weights, satisfying OPRD-Vanilla’s same-architecture requirement.Source: §1 · p. 3

OPRD-Bridge

  • OPRD-Vanilla requires matching depth and hidden dimension, whereas OPRD-Bridge aligns heterogeneous models through a frozen pair of projectors into a shared low-rank subspace.The teacher projector comes from PCA of teacher hidden-state covariance, the student projector is trained to align with it, and both are frozen for the main distillation phase.Source: §1 · pp. 3–4
  • A rank-8 projection yields 95% cosine similarity between Qwen3-4B and Qwen3-1.7B-Base representations; OPRD-Bridge also transfers across disjoint tokenizers.The models have 36 versus 28 layers and hidden dimensions of 2560 versus 2048; similarity decreases at higher rank, suggesting the shared structure concentrates in few principal directions. The Platonic Representation Hypothesis [15], alignment findings [22], and projection results [26] [27] motivate this bridge.Source: §1 · pp. 3–4

2 Background and Problem Setup

The section formalizes on-policy distillation and its output-space variants, then identifies their shared reliance on LM-head distributions as the limitation motivating hidden-state supervision.

2.1 Notation

  • The setup uses a trainable student and fixed teacher sharing vocabulary V, with autoregressive responses and layerwise hidden states evaluated on identical inputs.The notation distinguishes shared depth and width from separate student and teacher depths and widths when architectures differ.Source: §2.1 · p. 4

2.2 The On-Policy Distillation Framework

  • OPD samples a response from the student, then queries both models on student-generated prefixes, eliminating the exposure-bias gap between training and inference distributions.Its trajectory-level reverse KL decomposes exactly into token-level reverse KL terms, although computing the exact terms requires summing over the full vocabulary.Source: §2.2 · pp. 4–5

2.3 Three Output-Space Variants

  • Sampled-token OPD reuses one rollout token [37] [41] (one log-probability/token; O(BT) memory); full-vocabulary OPD computes exact KL (O(BT|V|) memory), while top-k typically uses k∈[4,64].Top-k trades supervision density against query cost; k=1 recovers deterministic sampled-token OPD and k=|V| recovers full-vocabulary OPD.Source: §2.3 · p. 5

2.4 A Shared Structural Limitation

  • All prior variants supervise only LM-head distributions, omitting intermediate states; the head imposes only d effective constraints per position.Because Whead is low-rank (d ≪ |V|), hidden directions remain unsupervised; sampled-token variance grows unfavorably with |V| and dominates once pt approaches qt.Source: §2.4 · pp. 5–6

3 On-Policy Representation Distillation

OPRD moves on-policy distillation from next-token distributions to hidden-state supervision on student-generated trajectories. Its objective supplies dense per-layer signals, while OPRD-Bridge extends representation alignment to heterogeneous models through a frozen low-rank interface.

3 On-Policy Representation Distillation

  • OPRD supervises student hidden states on its own rollouts, providing dense deterministic signals before the LM head while avoiding logits computation.The paper reports 1.44× faster training and up to 54% less actor-update peak GPU memory.Source: §3 · p. 6

3.1 The OPRD Objective

  • The layer-averaged, position-masked MSE lets OPRD select distilled layers and response positions, with all layers and a response suffix used for long reasoning traces.OPRD can also be added to an output-space OPD loss using the same rollout and teacher forward pass.Source: §3.1 · p. 7

3.2 Why OPRD Works

  • Conditioned on a rollout, OPRD has zero gradient variance; unlike OPD, it remains a stable signal near distribution convergence and constrains hidden-state information discarded by the LM head.The training-dynamics figure reports OPRD continuing to improve toward teacher level while both OPD variants plateau or oscillate across AIME24, AIME25, and AIMO.Source: §3.2 · pp. 7–8, Figure 4 · p. 12

3.3 OPRD-Bridge: Cross-Architecture Extension

  • OPRD-Bridge addresses representation mismatch: alignment, rather than matching architecture, is the requirement for meaningful hidden-state comparison across models.Models may be incommensurable despite matching dimensions, while depth, width, or training-history differences motivate a bridge.Source: §3.3 · p. 8

3.3.1 The Shared Low-Rank Structure

  • For Qwen3-4B and Qwen3-1.7B, rank 8 reaches 95% mean cosine similarity across layer pairs, while similarity falls to 77.0% at full rank 2048.The paper interprets the decline with rank as evidence that shared structure is concentrated in a few principal directions and higher rank admits architecture-specific variation.Source: §3.3.1 · pp. 8–9

3.3.2 Stage 1: Bridge Construction

  • Bridge construction pairs layers by proportional depth, computes frozen teacher PCA bases, and trains student projectors against them with both model backbones frozen.The completed bridge is a pair of fixed linear maps into a shared rank-r subspace.Source: §3.3.2 · p. 9

3.3.3 Stage 2: Bridge-Based Distillation

  • With the bridge frozen, the student generates rollouts and only its backbone receives gradients through the projected-representation loss.The objective uses the same position mask as Vanilla OPRD, defaults to the last 2,000 tokens, and L2-normalizes projected vectors before MSE.Source: §3.3.3 · p. 10

3.3.4 Design Philosophy

  • The low-rank frozen bridge preserves pre-existing alignment while filtering misaligned directions; under Gaussian assumptions, PCA is rate-distortion optimal, and representation-space alignment does not require a shared tokenizer.The paper’s informal theorem predicts an optimal rank balancing discarded aligned signal against transmitted misaligned noise; it also explains why freezing the bridge matters.Source: §3.3.4 · pp. 10–12

4 Experiments

Section 4 evaluates OPRD-Vanilla in the same-architecture setting and OPRD-Bridge in the cross-architecture setting.

4 Experiments

  • The experiments evaluate OPRD-Vanilla for same-architecture distillation and OPRD-Bridge for cross-architecture distillation.These settings are covered in Sections 4.1 and 4.2, respectively.Source: §4 · p. 12

4.1 OPRD-Vanilla: Same-Architecture Distillation

OPRD-Vanilla aligns hidden states on the student’s on-policy rollouts, closing the teacher gap that output-space OPD leaves open. It also trains faster and with lower transient memory use, while training dynamics and mechanistic analyses show continued improvement and motivate supervising response tails.

4.1.1 Experimental Setup

  • The same-architecture comparison uses a frozen JustRL-1.5B teacher and R1-distill-1.5B student sharing a Qwen2.5-1.5B backbone, with matched on-policy rollouts and teacher forward passes.It compares sampled-token OPD, top-16 OPD, and OPRD-Vanilla, which supervises all 28 layers over the final 2,000 response tokens.Source: §4.1 · p. 12, §4.1.1 · p. 13

4.1.2 Main Results

  • OPRD scores 49.8/34.6/79.1 on AIME24/AIME25/AIMO, within 1.0/1.0/0.4 points of the teacher.It gains +2.7/+0.6/+2.1 points over the better OPD baseline on those benchmarks; both OPD variants improve over the student but plateau below the teacher, with no consistent top-1 versus top-16 ordering.Source: Table 2 · p. 13, §4.1.2 · pp. 13–14

4.1.3 Training Dynamics

  • OPRD accuracy keeps climbing toward teacher level while both OPD variants plateau or oscillate, including top-16 remaining about 2.6 points below OPRD on AIMO.OPRD’s mean response length converges to about 5,700 tokens versus about 7,000 for both OPD variants, while supervised representation cosine similarity rises throughout training.Source: Figure 5 · p. 14, Figure 6 · p. 14, §4.1.3 · p. 15

4.1.5 Mechanistic Analysis

  • On AIME24, adding OPRD to OPD top-1 raises Avg@16 from 42.3 to 47.7 at μ=1 and 50.2 at μ=10, near the teacher’s 50.8.The gains are monotonic across the tested weights; μ=1 already exceeds OPD top-16’s 47.1, supporting OPRD as an additive signal rather than a replacement for OPD.Source: Figure 7 · p. 15, §4.1.5 · pp. 16–17

4.1.4 Efficiency

  • 563 minutes is OPRD’s training time, versus 813/812 minutes for OPD top-1/top-16, a 31% reduction or 1.44× speed-up.OPRD’s actor-update transient memory is 20.5 GB, versus 30.2 GB for OPD top-1 and 45.0 GB for OPD top-16, reductions of 32% and 54%, respectively.Source: §4.1.4 · pp. 15–16, Table 3 · p. 16

4.1.5 Mechanistic Analysis

  • The student’s representation gap is concentrated in response tails: last-50 similarity is 91.65%, while first-k similarity stays above 97% for k≤1600.This motivates OPRD’s default last-2000-token supervision, targeting positions where the student still diverges while avoiding early positions that are already teacher-aligned.Source: Figure 8 · p. 17, §4.1.5 · pp. 17–18
  • Adding OPRD to OPD top-16 produces a dip-then-surge in top-16 overlap, ultimately surpassing OPD top-16 alone.The shift follows the policy-gradient loss spike and indicates that hidden-state alignment improves the next-token agreement OPD top-16 directly optimizes.Source: Figure 10 · p. 18, §4.1.5 · p. 19
  • OPRD shifts the policy-gradient loss spike earlier, while late-training policy-gradient loss approaches zero even as accuracy differences persist.The entropy-increase phase also starts earlier with OPRD, then the student–teacher entropy gap narrows more rapidly, consistent with accelerated internal reorganisation and the output-space bottleneck.Source: §4.1.5 · pp. 18–19

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

OPRD-Bridge aligns teacher and student hidden states through a projector pair, enabling distillation across depth and width mismatches and even incompatible tokenizers. In cross-architecture tests, it matches OPD top-1’s capability ceiling on AIME24 and AIME25 while producing more diverse, shorter responses; adding OPD raises the ceiling further.

4.2.2 Main Results

  • Best@16 averaged 33.9 after OPRD-Bridge→OPD top-1, above standalone OPD top-1 (30.8) and OPD top-16 (27.1).Avg@16 improved from 11.2 to 13.6; Dist-4g shifted from 20.4 to 15.0, while Avg Len rose from 5,333 to 7,169.Source: Table 4 · p. 20, Table 5 · p. 20, §4.2.2 · p. 21
  • Best@16 was 20.0 on AIME24 and 13.3 on AIME25 for OPRD-Bridge, matching OPD top-1 and outperforming OPD top-16.This capability ceiling parity holds despite alignment through a rank-8 subspace rather than the full vocabulary simplex.Source: Table 4 · p. 20, §4.2.2 · pp. 20–21
  • Dist-4g was 22.6/17.2/21.3 for OPRD-Bridge on AIME24/AIME25/AIMO, versus 6.2/6.4/15.6 for OPD top-1 and 9.9/11.7/13.6 for OPD top-16.The paper attributes OPD’s repetitive patterns as likely related to its high-variance output-space gradient.Source: Table 5 · p. 20, §4.2.2 · p. 21
  • Avg Len was 5,909/5,234/4,855 tokens for OPRD-Bridge on AIME24/AIME25/AIMO, below OPD top-1 (19,120/17,931/11,008) and OPD top-16 (12,222/10,175/11,833).The lengths were close to the teacher’s 6,542/4,426/2,384 tokens; the paper links Bridge’s lower Avg@16 consistency to its untrained LM head, which OPD can calibrate.Source: Table 5 · p. 20, §4.2.2 · p. 21

4.2.3 Ablation Studies

  • After bridge training, cosine similarity rose from at most 7.5% before training to 72.3% at rank 1 and over 90% at ranks r ≥4, peaking at 95.0% for r = 8.Similarity then fell to 77.0% at full rank r = 2048; ranks r = 4–16 formed the practical 93.7–95.0% range, motivating the default r = 8.Source: Figure 12 · p. 22, §4.2.3 · p. 22

4.2.4 Cross-Tokenizer Distillation

  • Avg@16 rose from 2.8% to 5.0% with OPRD-Bridge when distilling Phi-4-mini-reasoning into Qwen3-1.7B-Base across disjoint tokenizers.Best@16 averaged 21.3% versus 4.5% without distillation, a 4.7× improvement; representation-space alignment bypassed the vocabulary channel on which standard OPD depends.Source: Table 6 · p. 22, §4.2.4 · p. 23

5 Discussion

OPRD’s discussion centers on representational alignment: OPRD-Vanilla uses existing alignment, while OPRD-Bridge constructs a low-rank channel when alignment is low. The authors identify applications in RL merging and self-distillation, and propose alignment-aware pre-training as a way to make the bridge unnecessary.

From representational alignment to cross-architecture distillation

  • OPRD-Vanilla suffices when teacher and student alignment is high, whereas OPRD-Bridge uses rank r to control a low-rank, low-pass channel that transmits shared structure and suppresses architecture-specific noise.The distinction applies to high alignment, such as shared architecture and initialization, versus low alignment caused by differences in depth, width, or training history.Source: §5 · p. 23

High-value application 1: multi-model RL merging

  • OPRD offers multi-model RL merging a hidden-state alternative to full-vocabulary and top-k OPD, avoiding vocabulary-sized logits while providing deterministic gradients without top-k truncation bias.Full-vocabulary OPD can require high transient GPU memory and infrastructure changes; top-k OPD ignores tail tokens and retains the LM-head bottleneck.Source: §5 · p. 23

High-value application 2: on-policy self-distillation (OPSD)

  • In on-policy self-distillation, OPRD can replace reverse-KL output-space computation with a cheaper, lower-variance representation objective while retaining privileged-information guidance.The teacher is constructed from the student by adding privileged information to the prompt; shared weights satisfy the same-architecture requirement and maximize hidden-state alignment.Source: §5 · p. 23

Toward alignment-aware pre-training

  • Alignment-aware pre-training could establish shared representational structure across differently sized models, allowing post-training OPRD-Vanilla with a dimension adapter instead of OPRD-Bridge.The proposed routes include shared initialization of common layers, periodic representation-matching regularization, or co-distillation; better alignment could turn directions currently filtered as noise into signal.Source: §5 · pp. 23–24

6 Related Work

This section situates OPRD among output-space knowledge distillation, on-policy distillation, and intermediate-representation distillation. It distinguishes OPRD by applying teacher–student hidden-state supervision to student-generated autoregressive rollouts.

Output-Space Knowledge Distillation

  • Classical sequence KD, from Hinton et al. [12] and Kim and Rush [20], is off-policy and output-only, introducing exposure bias and leaving hidden-state deviations in the LM head’s effective null space unpenalized.OPRD departs from both properties simultaneously.Source: §6 · p. 24

On-Policy Distillation

Feature / Intermediate-Representation Distillation

  • Feature distillation ranges from FitNets [28], attention-transfer [46], and FSP to TinyBERT [17], Mobile-BERT [33], and MiniLM/MiniLMv2 [34] [35], which align intermediate representations.These methods are presented as prior feature-matching approaches that differ structurally from OPRD in the autoregressive on-policy setting.Source: §6 · p. 24
  • Unlike fixed-input feature distillation, OPRD aligns hidden states on evolving student-generated sequences, matching each autoregressive prefix state under the student’s own sampling distribution.The teacher is queried on states the student actually visits as its policy evolves; each prefix representation encodes the belief before predicting the next token.Source: §6 · pp. 24–25

Hint Learning, Auxiliary Losses, and Distribution Matching

  • Unlike auxiliary or self-supervised representation methods such as deeply supervised nets [23], DINO [3], and representation engineering [49], OPRD is teacher–student, on-policy, and composable with output-space OPD.The objective can compose with any output-space OPD variant via Eq. (7).Source: §6 · p. 25

7 Conclusion and Future Work

OPRD moves on-policy distillation into hidden-state space for deterministic, per-layer supervision, while OPRD-Bridge extends representation alignment across heterogeneous models and vocabularies. The conclusion outlines open directions spanning new tasks, objectives, diagnostics, and theory.

Conclusion

  • 1.44× faster training and up to 54% less actor-update transient memory accompany OPRD’s monotonic gains and student–teacher gap closure across AIME 2024, AIME 2025, and AIMO.The conclusion attributes this efficiency to avoiding materialization of the [B, T, |V|] logits tensor. Output-space baselines plateau below the teacher; OPRD instead provides deterministic per-sample gradients and per-position, per-layer information.Source: §7 · p. 25
  • OPRD-Bridge uses a frozen low-rank bridge to align heterogeneous representations, enabling transfer across Qwen3-4B→Qwen3-1.7B and incompatible-tokenizer Phi-4-mini-reasoning→Qwen3-1.7B settings.The bridge shifts alignment from output space to representation space, decoupling distillation from vocabulary-based alignment.Source: §7 · p. 25

Future Work

  • Whether OPRD’s gains transfer beyond long-CoT mathematics to code generation, agentic interaction, or open-ended dialogue remains unresolved, while cross-modal distillation is proposed as a next step.The proposed cross-modal setting includes a vision-language teacher and language-only student, where the representation-space bridge may be the only viable supervision channel.Source: §7 · pp. 25–26
  • Future work proposes adaptive layer and position weighting based on student–teacher gaps or informative gradients, and scaling on-policy representation self-distillation to multi-turn and multi-task settings.Current OPRD uses uniform layer weighting and a last-k position heuristic; OPRSD leverages privileged information in same-architecture self-distillation.Source: §7 · p. 26
  • Future work calls for explaining OPRD’s PG-loss spike and entropy/overlap dynamics, and for adding on-policy attention matching to transfer teacher routing and composition behavior.OPRD currently aligns hidden-state vectors but does not supervise the attention patterns that produce them.Source: §7 · p. 26
  • Alignment-aware pre-training could enforce representational alignment between large and small models, potentially removing the low-rank bridge requirement for cross-architecture OPRD.The conclusion identifies representational alignment, rather than architectural identity, as the prerequisite for direct hidden-state distillation.Source: §7 · p. 26
  • Future theory and measurement work includes tracking cosine similarity, CKA, and probing accuracy, plus deriving convergence-rate bounds and characterizing hidden-state directions nulled by the LM head.These representation-level diagnostics would complement output metrics; tighter bounds would compare OPRD with sampled-token OPD.Source: §7 · p. 26

Conclusion

  • The authors argue that hidden-state representations are an under-exploited distillation resource and encourage treating teachers as sources of layered internal computation, not merely probability oracles.Source: §7 · p. 26

A Notation Summary · B Formal Theoretical Guarantees

The formal analysis contrasts OPD’s sampled-token gradient, whose variance persists and signal-to-noise ratio collapses near the teacher, with OPRD’s deterministic per-position gradient. It also specifies the notation and assumptions used to formalize these guarantees.

B Formal Theoretical Guarantees

  • The formal guarantees are organized around OPD gradient variance and SNR collapse, alongside the LM-head bottleneck’s null-direction identity and spectral gap.The passage maps these claims to Theorems 4–6, Corollary 1, and Theorems 7–8, respectively.Source: B Formal Theoretical Guarantees · p. 30

A Notation Summary

B.1 Setup and Assumptions

  • The setup defines sampled-token gradient estimators and assumes smooth log-densities, finite Fisher information, bounded log-ratios, and bounded hidden states.The passage describes these regularity conditions as standard and mild, with examples of when they hold.Source: §B.1 · p. 31

B.2 Variance of Sampled-Token OPD

  • OPD’s single-sample gradient variance remains lower-bounded near p=q by the log-ratio variance scaled by the minimum Fisher-information eigenvalue.The bound does not vanish as the loss approaches zero.Source: §B.2 · p. 32

B.3 OPRD Gradient Is Deterministic

  • OPRD’s per-position gradient has zero conditional variance, and its variance is lower than OPD’s except when p is a point mass.The variance gap grows with the per-token log-ratio magnitude and policy spread.Source: §B.3 · pp. 32–33

B.4 Signal-to-Noise Ratio Collapse of Sampled-Token OPD

  • As symmetric divergence δ approaches zero, OPD’s SNR is O(δ) and tends to zero, while OPRD’s SNR is infinite before hidden-space convergence.OPRD’s gradient has zero covariance, making its signal noise-free under the stated condition.Source: §B.4 · pp. 33–34
  • The analysis predicts that OPD stagnates or oscillates near an asymptote as its SNR collapses, whereas OPRD continues to provide an informative descent direction.The passage relates this prediction to the reported empirical pattern of pure OPD plateauing while OPRD and OPD + μ·OPRD improve monotonically.Source: §B.4 · p. 34

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

The formal results show that output-space distillation cannot observe LM-head null directions or intermediate-layer changes, and that the LM head’s spectral spread weakens its sensitivity to hidden-state deviations. OPRD directly penalizes representations in these otherwise unconstrained directions and layers.

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

  • Every output-space OPD loss that vanishes at matching distributions is zero when the hidden-state difference lies in the LM head’s effective null space.This includes sampled-token, top-k truncated reverse KL, and full-vocabulary reverse KL losses; softmax distributions are invariant to a common additive logit shift.Source: §B.5 · p. 34
  • For sampled-token, top-k, and full-vocabulary reverse KL, a local Lipschitz upper bound holds on any compact logit region.The constant depends on the loss and logit range; the proof sketch derives the bound from bounded Hessians around the additive-invariance optimum.Source: §B.5 · p. 35
  • Holding output-space loss ℓout fixed, hidden-state perturbations along vd can be σ1/σd times larger in ℓ2 norm than perturbations along top singular direction v1, assuming Whead has full column rank.Source: §B.5 · pp. 34–35
  • Output-space losses depend only on the last-layer state, so intermediate-layer perturbations that leave the final state unchanged are invisible, while OPRD can directly penalize those layers.The passage identifies perturbations in the kernel of the residual stack from an intermediate layer onward as one example.Source: §B.5 · p. 35

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

The section formalizes the PCA-based low-rank bridge as information-theoretically optimal, with rank r controlling a bias–variance trade-off. Its setup models how well student states predict teacher principal directions and assumes alignment decreases with component rank.

B.6.1 Setup

  • The setup defines each teacher principal component’s alignment as its coefficient of determination from the best linear prediction using the student state.Teacher hidden states are modeled with mean μT and positive-definite covariance ΣT; the components are ordered by decreasing variance.Source: §B.6.1 · p. 36
  • The monotone-alignment assumption says high-variance principal directions are better predicted by the student than low-variance directions.The per-direction alignment is assumed non-increasing as component rank increases.Source: §B.6.1 · p. 36
  • Figure 12’s rank-cosine curve empirically verifies that alignment concentrates in top principal components and degrades monotonically with rank.The section introduces bridge rank r as controlling a bias–variance trade-off.Source: §B.6.1 · p. 36, §B.6 · p. 36

B.6.2 Rate–Distortion Optimality

  • Under a Gaussian teacher-state model, projecting onto the top-r principal directions minimizes reconstruction distortion among rank-r linear encoders.The section identifies this PCA projection as the rate–distortion-optimal linear channel at capacity r.Source: §B.6.2 · p. 36
  • The proof follows from the Eckart–Young–Mirsky theorem and the classical result that PCA minimizes mean-squared reconstruction error among linear rank-r projections.For Gaussian sources, this result coincides with the squared-error rate–distortion function.Source: §B.6.2 · p. 36
  • Any other rank-r projection discards more teacher information, so the PCA bridge is an information-theoretically optimal linear channel at capacity r.The optimal encoder uses the top-r eigenvectors of the teacher covariance matrix.Source: §B.6.2 · p. 36

B.6.3 Bias–Variance Decomposition of Distillation Error

This section decomposes OPRD-Bridge distillation error into bias from discarded aligned signal and variance from transmitted noise the student cannot match. Under Assumption 2, the error is quasi-convex in bridge rank, with the Qwen3-4B → Qwen3-1.7B experiment supporting an optimum near rank 8.

B.6.3 Bias–Variance Decomposition of Distillation Error

  • Under Assumption 2, total distillation error—not reconstruction error—decomposes into bias from discarded aligned directions and variance from transmitted directions the student cannot match.Bias decreases as rank increases, while variance increases.Source: §B.6.3 · pp. 36–37
  • Under Assumption 2, Etotal(r) is quasi-convex in rank r; in Qwen3-4B → Qwen3-1.7B experiments, the rank-cosine curve peaks at r = 8 and declines thereafter.Source: §B.6.3 · p. 37
  • At r = 8, the Qwen3-4B → Qwen3-1.7B rank-cosine curve peaks, and rapid teacher-covariance spectral decay makes both bias and variance small there.The passage reports this as consistent with r* ≈ 8, where ρ_i^2 crosses 1/2.Source: §B.6.3 · p. 37

Related papers

Loading 2606.06021v4…