OPRD: On-Policy Representation Distillation — part 1/9

Cleaned-text code points [0, 15480). Bibliography omitted here; source.md includes it.

# OPRD: On-Policy Representation Distillation

###### Abstract

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, AIMO) where every output-space baseline plateaus below the teacher, while training 1.44 × 1.44\times faster and using up to 54 % 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/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 → \to Qwen3-1.7B-Base) and cross-tokenizer (Phi-4-mini-reasoning → \to Qwen3-1.7B-Base) settings, demonstrating successful knowledge transfer even when the vocabulary-based alignment channel is unavailable. The code is available via https://github.com/ShenzhiYang2000/OPRD .

## 1 Introduction

On-policy distillation (OPD) has become a central building block in large language model (LLM) post-training. By letting the student sample its own responses and then scoring each token against the teacher’s conditional distribution, OPD provides a dense, token-level training signal that adapts to the student’s current policy, avoiding the exposure bias inherent in training on static teacher outputs [ 2 ] . Multiple production systems now rely on OPD as a primary post-training stage [ 39 , 37 , 47 , 5 ] , positioning it alongside supervised fine-tuning and outcome-reward reinforcement learning.

Despite this momentum, the design space of OPD has remained surprisingly narrow. Every variant proposed to date (sampled-token [ 37 , 41 ] , full-vocabulary, and top- k k ) differs only in how many output tokens are evaluated per position, yet they all operate inside the same output space : the divergence is computed over next-token probability distributions p t p_{t} and q t q_{t} . We argue that this output-only paradigm imposes two practical limitations that become increasingly damaging as training progresses.

Limitation 1: Variance dominates the late-stage signal. Sampled-token OPD estimates each token-level reverse KL from a single sample drawn from a vocabulary of size | 𝒱 | |\mathcal{V}| (e.g., ≈ 151 \approx\!151 K for Qwen3). The estimator is unbiased, but its variance shrinks much more slowly than the signal: letting δ ≜ ‖ p t − q t ‖ \delta\triangleq\|p_{t}-q_{t}\| measure the student–teacher gap, the expected gradient decays as O ⁡ ( δ 2 ) O(\delta^{2}) while the variance decays only as O ⁡ ( δ ) O(\delta) , so the signal-to-noise ratio collapses and training plateaus well below the teacher ( Figure 4 ). Top- k k OPD partially mitigates this by evaluating k k tokens per position, but trades sampling variance for truncation bias and still plateaus empirically.

Limitation 2: The output layer is an information bottleneck. Every output-space variant treats the teacher as a black-box probability oracle , querying only the LM-head output. Yet the teacher has computed, at every position, an entire stack of d d -dimensional hidden states encoding attention patterns, mid-layer reasoning state, and geometric structure. Nearly all of this is destroyed by the LM-head projection W head : ℝ d → ℝ | 𝒱 | W_{\mathrm{head}}\!:\mathbb{R}^{d}\!\to\!\mathbb{R}^{|\mathcal{V}|} and the subsequent softmax: output distributions that agree to within an arbitrary tolerance can correspond to hidden states differing along entire affine subspaces of ℝ d \mathbb{R}^{d} . The student is therefore graded only on what survives this projection, and receives no signal about how the teacher arrived at that distribution. This is particularly wasteful in the on-policy regime, where the teacher’s hidden states are already computed on every rollout but discarded before they reach the loss.

To overcome both limitations, we propose On-Policy Representation Distillation (OPRD) , the first method to lift on-policy distillation from the output space into the hidden-state space . On the same on-policy rollouts ( x , y ^ ) (x,\hat{y}) already used by standard OPD, OPRD aligns the student’s intermediate hidden representations with the teacher’s across selected transformer layers and response positions via a normalized mean-squared error objective. A single design choice (supervising at the representation level rather than at the output level) simultaneously addresses both limitations. First, deterministic, low-variance gradients: OPRD’s MSE objective is a deterministic function of the rollout; its gradient carries zero additional sampling variance, eliminating the late-stage signal-to-noise collapse of OPD by construction. Second, a richer supervision channel beyond logits: OPRD taps the teacher at any subset of its L L intermediate layers, exposing (layers × \times positions × \times hidden-dim) scalars of structural supervision per sample, orders of magnitude more than the signal extracted at the output. The student is graded on the same intermediate representations the teacher actually computed, without filtering through the LM-head projection. Both properties follow from a single conceptual shift: moving the supervision target from the output of the LM head to its input. Because the loss path never materialises the [ B , T , | 𝒱 | ] [B,T,|\mathcal{V}|] logits tensor, OPRD also reduces wall-clock time and peak GPU memory as a direct consequence of this design. OPRD is a self-contained training objective that can be used on its own; it also composes additively with any OPD variant at essentially zero infrastructure cost. Beyond the standard teacher–student setting studied here, we highlight two high-value scenarios where OPRD’s advantages are especially pronounced. (1) Multi-model RL merging. State-of-the-art RL pipelines increasingly merge multiple teacher or reward-model checkpoints into a single student. In this setting, full-vocabulary OPD requires materialising a [ B , T , | 𝒱 | ] [B,T,|\mathcal{V}|] logit tensor per teacher , quickly exhausting GPU memory and demanding heavy infrastructure work [ 5 ] . Top- k k OPD alleviates memory but introduces a truncation bias (tail tokens are ignored) and still plateaus below the teacher empirically. OPRD sidesteps both: its hidden-state loss never touches the vocabulary dimension, so memory and wall-clock scale with d d rather than | 𝒱 | |\mathcal{V}| , while its deterministic gradient avoids the variance trap entirely. (2) On-policy self-distillation (OPSD). A growing line of work constructs the teacher from the student itself by injecting privileged information (e.g., ground-truth solutions, step-level verification signals) into the prompt. Because teacher and student share exactly the same weights, the same-architecture requirement is satisfied by construction, and the hidden-state alignment signal is maximally informative. OPRD can therefore serve as a drop-in replacement for the output-space reverse-KL in any OPSD pipeline, delivering lower variance and lower cost without any architectural modification. We discuss both applications in detail in § 5 . The strict Pareto improvement over all output-space baselines is summarized in Figure 1 .

The above formulation, which we call OPRD-Vanilla , however, requires teacher and student to share the same architecture: identical depth, identical hidden dimension, and compatible initialization. The moment these conditions are violated, the method breaks down, since hidden states of different dimensionality cannot be directly compared, and even when dimensions happen to match across different architectures, the cosine similarity between corresponding hidden states is empirically zero. This same-architecture constraint excludes precisely the most common distillation scenario, namely distilling a large, expensive teacher into a smaller, deployable student, which inherently involves heterogeneous model pairs.

A useful perspective on this problem comes from asking why output-space distillation does not suffer from the same limitation. The answer is that logit distillation possesses a natural bridge : the shared vocabulary. Both teacher and student project their hidden states through their respective LM heads into the same probability simplex, making outputs directly comparable regardless of internal architecture. The vocabulary serves as a fixed, pre-established interface that neither model modifies during training, and knowledge flows through it. This bridge works, but it has a critical limitation of its own: it requires both models to share the same vocabulary. Moreover, it operates at only a single point in the network (the final layer) and is lossy, compressing all representational structure into a single next-token distribution. The natural question, then, is whether one can construct an analogous bridge for representations , one that carries richer information, operates at every layer, and does not depend on a shared vocabulary.

The Platonic Representation Hypothesis [ 15 ] provides theoretical grounds to believe such a bridge should exist. It posits that strong neural networks, regardless of architecture, are converging toward a shared statistical model of reality in their representation spaces. Empirically, models with higher downstream performance exhibit higher mutual representation alignment [ 22 ] , and a single linear projection suffices to bridge representations across architectures and even across modalities [ 27 , 26 ] . We confirm this directly: pairing a Qwen3-4B teacher (36 layers, d = 2560 d\!=\!2560 ) with a Qwen3-1.7B-Base student (28 layers, d = 2048 d\!=\!2048 ), a rank-8 linear subspace achieves 95% cosine similarity between the projected representations of the two models, despite their different depths and widths. Notably, this similarity decreases with higher rank, indicating that the shared structure concentrates in very few principal directions and that additional dimensions introduce noise rather than signal.

This finding leads directly to OPRD-Bridge . We construct a pair of linear projectors, one per model, that map their respective hidden states into a shared low-rank subspace. The teacher’s projector is obtained via PCA of its hidden-state covariance (extracting the principal variation directions); the student’s projector is trained to align with it. Once both projectors converge, they are frozen and serve as a fixed bridge through which representational structure flows from teacher to student during the main distillation phase, which optimizes only the student backbone. The design mirrors logit distillation: just as the vocabulary is a frozen interface through which output-level knowledge flows, the projector pair ( P T , P S ) (P_{T},P_{S}) is a frozen interface through which representation-level knowledge flows, but at every layer rather than only the last, and decoupled from the vocabulary. We validate this capability empirically by distilling across completely disjoint tokenizers (Phi-4-mini-reasoning → \to Qwen3-1.7B), a setting where the vocabulary-based alignment channel is unavailable.

Our main contributions are as follows: 1. OPRD-Vanilla: representation-level on-policy distillation. We propose On-Policy Representation Distillation, the first method to lift on-policy distillation from the output space into the hidden-state space. We provide a two-perspective theoretical analysis showing that OPRD (i) yields a deterministic per-sample gradient, removing the token-level estimation variance of OPD’s gradient estimator, and (ii) exposes per-layer structural information that the LM-head projection necessarily discards.

2. OPRD-Bridge: cross-architecture extension. We show that heterogeneous teacher–student pairs share a low-rank representational structure, and exploit this to construct a frozen projector pair that enables representation distillation across arbitrary architecture pairs (different depth, width, and potentially vocabulary), the first such method in the on-policy regime.

3. Empirical validation. On competition mathematics benchmarks (AIME 2024, AIME 2025, AIMO), OPRD-Vanilla closes the student–teacher gap where every output-space baseline plateaus, while training 1.44 × 1.44\times faster and using up to 54 % 54\% less memory.

## 2 Background and Problem Setup

This section formalizes the on-policy distillation problem we build upon. We introduce the necessary notation in § 2.1 , define the on-policy distillation framework in § 2.2 , and catalogue the three output-space supervision granularities used in prior work in § 2.3 . We close in § 2.4 by isolating the common structural property of these variants that motivates our hidden-state approach in the next section.

### 2.1 Notation

We consider two autoregressive language models with a shared vocabulary 𝒱 \mathcal{V} : a student π θ \pi_{\theta} with trainable parameters θ \theta , and a fixed teacher π T \pi_{T} . A training instance is a prompt x = ( x 1 , … , x n ) x=(x_{1},\ldots,x_{n}) drawn from a prompt distribution 𝒟 x = { x ( i ) } i = 1 N \mathcal{D}_{x}=\{x^{(i)}\}_{i=1}^{N} ; a model response is a token sequence y = ( y 1 , … , y m ) y=(y_{1},\ldots,y_{m}) produced autoregressively. For brevity we write the prefix up to step t t as y < t ≜ ( y 1 , … , y t − 1 ) y_{<t}\triangleq(y_{1},\ldots,y_{t-1}) , and use π ( ⋅ ∣ x , y < t ) \pi(\cdot\mid x,y_{<t}) to denote either model’s next-token distribution over 𝒱 \mathcal{V} conditioned on ( x , y < t ) (x,y_{<t}) . The notation y ∼ π θ ( ⋅ ∣ x ) y\sim\pi_{\theta}(\cdot\mid x) refers to an autoregressive sample drawn from the student.

Both models follow the standard transformer template: a stack of self-attention blocks producing intermediate hidden states h ( l ) ∈ ℝ d h^{(l)}\in\mathbb{R}^{d} at each layer l l and position, followed by a language-model head W head ∈ ℝ | 𝒱 | × d W_{\mathrm{head}}\in\mathbb{R}^{|\mathcal{V}|\times d} that maps the final hidden state to logits. We use L L and d d when the two models share the same depth and width; when they differ (§ 3.3 ), we write L S , L T L_{S},L_{T} and d S , d T d_{S},d_{T} explicitly. We write h θ , t ( l ) h_{\theta,t}^{(l)} and h T , t ( l ) h_{T,t}^{(l)} for the student and teacher hidden states at layer l l and response position t t , with both networks evaluated on the same input sequence.

### 2.2 The On-Policy Distillation Framework

##### Setup.