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

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

On-policy distillation (OPD) departs from classical knowledge distillation by drawing the supervision distribution from the student rather than from a fixed dataset. Concretely, at each training step the student first samples a response y ^ = ( y ^ 1 , … , y ^ T ) ∼ π θ ( ⋅ ∣ x ) \hat{y}=(\hat{y}_{1},\ldots,\hat{y}_{T})\sim\pi_{\theta}(\cdot\mid x) of length T ≜ | y ^ | T\triangleq|\hat{y}| , after which both models are evaluated on the student-generated prefixes. For each position t ∈ { 1 , … , T } t\in\{1,\ldots,T\} this yields a pair of next-token distributions over 𝒱 \mathcal{V} : p t ​ ( v ) ≜ π θ ​ ( v ∣ x , y ^ < t ) , q t ​ ( v ) ≜ π T ​ ( v ∣ x , y ^ < t ) , v ∈ 𝒱 . p_{t}(v)\;\triangleq\;\pi_{\theta}(v\mid x,\hat{y}_{<t}),\qquad q_{t}(v)\;\triangleq\;\pi_{T}(v\mid x,\hat{y}_{<t}),\qquad v\in\mathcal{V}. (1) The defining feature of OPD is that the teacher is queried on student-visited states , namely prefixes that arise from the current policy, rather than on canonical teacher trajectories. This eliminates the exposure-bias gap between training and inference distributions that plagues fixed-target distillation.

##### Objective.

The canonical OPD objective minimizes the trajectory-level reverse KL divergence between the student and teacher policies on student rollouts. By the chain rule for KL divergence, this trajectory-level quantity decomposes exactly into a sum of token-level reverse KL terms: ℒ OPD ( θ ) = 𝔼 x ∼ 𝒟 x , y ^ ∼ π θ ( ⋅ ∣ x ) [ ∑ t = 1 T D KL ( p t ∥ q t ) ] , \mathcal{L}_{\mathrm{OPD}}(\theta)\;=\;\mathbb{E}_{x\sim\mathcal{D}_{x},\;\hat{y}\sim\pi_{\theta}(\cdot\mid x)}\!\left[\sum_{t=1}^{T}D_{\mathrm{KL}}(p_{t}\,\|\,q_{t})\right], (2) where the token-level reverse KL at position t t is D KL ( p t ∥ q t ) = ∑ v ∈ 𝒱 p t ( v ) log [ p t ( v ) / q t ( v ) ] D_{\mathrm{KL}}(p_{t}\,\|\,q_{t})=\sum_{v\in\mathcal{V}}p_{t}(v)\log[p_{t}(v)/q_{t}(v)] . Eq. ( 2 ) is conceptually clean but computationally inconvenient: it requires summing over the full vocabulary 𝒱 \mathcal{V} at every position, which is prohibitive for modern LLMs with | 𝒱 | |\mathcal{V}| in the hundreds of thousands. Practical implementations differ in how they approximate this sum, and we review the three dominant choices below.

### 2.3 Three Output-Space Variants

We use a unified template to describe each variant: at each position t t , define a token subset S t ⊆ 𝒱 S_{t}\subseteq\mathcal{V} and a per-position loss ℓ t \ell_{t} that depends only on { p t ( v ) , q t ( v ) : v ∈ S t } \{p_{t}(v),q_{t}(v):v\in S_{t}\} . The three variants below correspond to different choices of S t S_{t} .

##### (a) Sampled-token OPD ( S t = { y ^ t } S_{t}=\{\hat{y}_{t}\} ).

The most lightweight and by far the most widely adopted choice in production deployments [ 37 , 41 ] . A single token y ^ t ∼ p t \hat{y}_{t}\sim p_{t} already drawn during rollout is reused as the supervision target, and the per-position loss takes the form of a log-ratio: ℓ t sample ≜ log ⁡ p t ​ ( y ^ t ) − log ⁡ q t ​ ( y ^ t ) , ℒ OPD sample ​ ( θ ) = 𝔼 x , y ^ ​ [ ∑ t = 1 T ℓ t sample ] . \ell_{t}^{\mathrm{sample}}\;\triangleq\;\log p_{t}(\hat{y}_{t})-\log q_{t}(\hat{y}_{t}),\qquad\mathcal{L}_{\mathrm{OPD}}^{\mathrm{sample}}(\theta)=\mathbb{E}_{x,\hat{y}}\!\left[\sum_{t=1}^{T}\ell_{t}^{\mathrm{sample}}\right]. (3) A straightforward calculation gives 𝔼 y ^ t ∼ p t [ ℓ t sample ] = D KL ( p t ∥ q t ) \mathbb{E}_{\hat{y}_{t}\sim p_{t}}[\ell_{t}^{\mathrm{sample}}]=D_{\mathrm{KL}}(p_{t}\,\|\,q_{t}) , so ℓ t sample \ell_{t}^{\mathrm{sample}} is an unbiased single-sample estimator of the token-level reverse KL. Memory cost is O ⁡ ( B ​ T ) O(BT) for batch size B B and response length T T ; teacher queries amount to one log-probability per token.

##### (b) Full-vocabulary OPD ( S t = 𝒱 S_{t}=\mathcal{V} ).

At the opposite extreme, one materializes the entire teacher distribution and computes the exact token-level KL at every position: ℒ OPD full ​ ( θ ) = 𝔼 x , y ^ ​ [ ∑ t = 1 T ∑ v ∈ 𝒱 p t ​ ( v ) ​ log ⁡ p t ​ ( v ) q t ​ ( v ) ] . \mathcal{L}_{\mathrm{OPD}}^{\mathrm{full}}(\theta)=\mathbb{E}_{x,\hat{y}}\!\left[\sum_{t=1}^{T}\sum_{v\in\mathcal{V}}p_{t}(v)\,\log\!\tfrac{p_{t}(v)}{q_{t}(v)}\right]. (4) The gradient signal is the densest possible, but the price is steep: storing teacher logits demands O ⁡ ( B ​ T ​ | 𝒱 | ) O(BT|\mathcal{V}|) memory, which becomes infeasible for long-context training at modern vocabulary sizes.

##### (c) Top- k k OPD ( S t = TopK ⁡ ( p t , k ) S_{t}=\mathrm{TopK}(p_{t},k) ).

Top- k k OPD interpolates between the two extremes by restricting attention to the k k tokens that the student ranks highest at position t t , then computing a KL between renormalized distributions on this support: ℒ OPD top- ​ k ( θ ) = 𝔼 x , y ^ [ ∑ t = 1 T D KL ( p ¯ t ( S t ) ∥ q ¯ t ( S t ) ) ] , p ¯ t ( S t ) ( v ) = p t ( v ) 1 [ v ∈ S t ] ∑ u ∈ S t p t ​ ( u ) , \mathcal{L}_{\mathrm{OPD}}^{\text{top-}k}(\theta)=\mathbb{E}_{x,\hat{y}}\!\left[\sum_{t=1}^{T}D_{\mathrm{KL}}\!\left(\bar{p}_{t}^{(S_{t})}\,\big\|\,\bar{q}_{t}^{(S_{t})}\right)\right],\quad\bar{p}_{t}^{(S_{t})}(v)=\frac{p_{t}(v)\,\mathbf{1}[v\in S_{t}]}{\sum_{u\in S_{t}}p_{t}(u)}, (5) and analogously for q ¯ t ( S t ) \bar{q}_{t}^{(S_{t})} . The hyperparameter k k trades supervision density against teacher-query cost, with k = 1 k=1 recovering (a deterministic version of) sampled-token OPD and k = | 𝒱 | k=|\mathcal{V}| recovering full-vocabulary OPD. Typical implementations use k ∈ [ 4 , 64 ] k\in[4,64] . We measure this cost empirically in § 4.1.4 , where top- 16 16 OPD’s actor-update transient memory is more than 2 × 2\times larger than OPRD’s at the same setting.

### 2.4 A Shared Structural Limitation

The three variants above span the full design space studied in prior work, yet they share a defining structural property: the supervision signal is always a function of the next-token distributions p t p_{t} and q t q_{t} that the LM head produces. Equivalently, the only way teacher knowledge reaches the student is through the projection W head : ℝ d → ℝ | 𝒱 | W_{\mathrm{head}}:\mathbb{R}^{d}\to\mathbb{R}^{|\mathcal{V}|} applied to the final hidden state. The internal representations { h T , t ( l ) } l < L \{h_{T,t}^{(l)}\}_{l<L} , which are the very features that encode the teacher’s intermediate reasoning, never enter the loss. This output-only view has two immediate consequences that will become focal points of our analysis. (i) Statistical: the most popular variant (sampled-token OPD) estimates each token-level KL from a single Monte Carlo draw, introducing variance that scales unfavorably with | 𝒱 | |\mathcal{V}| and dominates the optimization signal once p t p_{t} approaches q t q_{t} . (ii) Informational: because W head W_{\mathrm{head}} is low-rank ( d ≪ | 𝒱 | d\ll|\mathcal{V}| ), the loss imposes only d d effective constraints per position regardless of | S t | |S_{t}| , leaving large directions of the hidden-state space unsupervised. Our method, introduced next, attacks both issues by replacing W head ∘ h W_{\mathrm{head}}\!\circ\!h with h h itself as the alignment target.

## 3 On-Policy Representation Distillation

We now present On-Policy Representation Distillation (OPRD) , a novel distillation framework that supervises the student in the hidden-state space on student-generated trajectories. We define the method ( Section 3.1 ) and state two theorems, one on gradient variance and one on the LM-head information bottleneck ( Theorem 1 , Theorem 2 ), that explain why hidden-state supervision is a principled and effective complement to output-space distillation.

### 3.1 The OPRD Objective

The three OPD variants in § 2.3 all operate in the output space by matching next-token distributions p t p_{t} and q t q_{t} . OPRD instead supervises the student in the hidden-state space on the same on-policy trajectories. Intuitively, OPD asks the student to assign similar probabilities to tokens, whereas OPRD asks the student to produce similar internal representations at selected layers and positions. We refer to this basic same-architecture formulation as OPRD-Vanilla (or simply OPRD when unambiguous); the cross-architecture extension, OPRD-Bridge , is presented in § 3.3 .

Let ℒ layer ⊆ { 1 , … , L } \mathcal{L}_{\mathrm{layer}}\subseteq\{1,\ldots,L\} be the set of distilled layers (e.g. the last layer, all layers, or a parity subset such as even/odd layers), and let 𝒫 ⁡ ( y ^ ) ⊆ { 1 , … , T } \mathcal{P}(\hat{y})\subseteq\{1,\ldots,T\} be the set of supervised response positions (e.g. all tokens, the first k k tokens, or the last k k tokens). We use a position mask m t ∈ { 0 , 1 } m_{t}\in\{0,1\} to indicate whether t ∈ 𝒫 ⁡ ( y ^ ) t\in\mathcal{P}(\hat{y}) ; for short responses, positions beyond the valid length are masked out rather than padded into the loss. OPRD minimizes a layer-averaged, position-masked mean-squared error between student and teacher representations: ℒ OPRD ( θ ) = 𝔼 x ∼ 𝒟 x , y ^ ∼ π θ ( ⋅ ∣ x ) [ 1 | ℒ layer | ∑ l ∈ ℒ layer 1 ∑ t = 1 T m t ∑ t = 1 T m t 1 d ∥ h θ , t ( l ) − sg ( h T , t ( l ) ) ∥ 2 2 ] , \mathcal{L}_{\mathrm{OPRD}}(\theta)=\mathbb{E}_{x\sim\mathcal{D}_{x},\;\hat{y}\sim\pi_{\theta}(\cdot\mid x)}\left[\frac{1}{|\mathcal{L}_{\mathrm{layer}}|}\sum_{l\in\mathcal{L}_{\mathrm{layer}}}\frac{1}{\sum_{t=1}^{T}m_{t}}\sum_{t=1}^{T}m_{t}\,\frac{1}{d}\Bigl\|h_{\theta,t}^{(l)}-\mathrm{sg}\!\bigl(h_{T,t}^{(l)}\bigr)\Bigr\|_{2}^{2}\right], (6) where sg ⁡ ( ⋅ ) \mathrm{sg}(\cdot) denotes the stop-gradient operator on the teacher representation and d d is the hidden dimension. The 1 / d 1/d factor normalizes the loss across architectures with different hidden sizes; the position averaging 1 / ∑ t m t 1/\sum_{t}m_{t} makes the loss invariant to the choice of | 𝒫 ⁡ ( y ^ ) | |\mathcal{P}(\hat{y})| . The two design knobs ( ℒ layer \mathcal{L}_{\mathrm{layer}} , 𝒫 ⁡ ( y ^ ) \mathcal{P}(\hat{y}) ) offer flexibility along two axes: depth of supervision (single-layer vs. multi-layer) and breadth of supervision (single-position vs. all-position). For long chain-of-thought (CoT) responses common in mathematical reasoning, we typically set 𝒫 ⁡ ( y ^ ) \mathcal{P}(\hat{y}) to the last k k response tokens and ℒ layer \mathcal{L}_{\mathrm{layer}} to all transformer layers, yielding dense layer-wise supervision on a compact suffix while keeping memory bounded. We empirically study the effect of these design choices in § 4.1.5 . OPRD is a self-contained training objective and our main results (§ 4 ) are reported in the OPRD-only setting. For completeness, OPRD also composes additively with any output-space OPD variant as ℒ ⁡ ( θ ) = ℒ OPD ​ ( θ ) + μ ​ ℒ OPRD ​ ( θ ) , μ ≥ 0 , \mathcal{L}(\theta)=\mathcal{L}_{\mathrm{OPD}}(\theta)+\mu\,\mathcal{L}_{\mathrm{OPRD}}(\theta),\quad\mu\geq 0, (7) at essentially zero infrastructure cost since both terms are computed on the same on-policy rollout and share a single teacher forward pass.

### 3.2 Why OPRD Works

Two complementary properties, in one-to-one correspondence with the two limitations of § 1 , explain why hidden-state supervision is a principled complement to output-space OPD. We state both as informal theorems; precise statements and proofs are deferred to Appendix B .

###### Theorem 1 (Zero-variance gradient) .

Let g OPD g_{\mathrm{OPD}} and g OPRD g_{\mathrm{OPRD}} be the per-sample stochastic gradients of sampled-token OPD and OPRD ( 6 ) on an on-policy rollout y ^ ∼ π θ ( ⋅ ∣ x ) \hat{y}\sim\pi_{\theta}(\cdot\!\mid\!x) . Conditioned on ( x , y ^ ) (x,\hat{y}) , Var [ g OPRD | x , y ^ ] = 0 , Var [ g OPD | x , y ^ ] ∝ Var y ^ t ∼ p t [ log p t ( y ^ t ) − log q t ( y ^ t ) ] , \mathrm{Var}\!\left[g_{\mathrm{OPRD}}\,\big|\,x,\hat{y}\right]=0,\qquad\mathrm{Var}\!\left[g_{\mathrm{OPD}}\,\big|\,x,\hat{y}\right]\;\propto\;\mathrm{Var}_{\hat{y}_{t}\sim p_{t}}\!\bigl[\log p_{t}(\hat{y}_{t})-\log q_{t}(\hat{y}_{t})\bigr], (8) where the right-hand variance is over per-position token sampling.

The OPD variance in ( 8 ) does not vanish as p t → q t p_{t}\!\to\!q_{t} , and through the score-function term ∇ θ ​ log ​ p t ​ ( y ^ t ) \nabla_{\theta}\log p_{t}(\hat{y}_{t}) it dominates the policy gradient late in training; this is the mechanism behind the late-stage stagnation of pure OPD (Limitation 1 in § 1 ). OPRD adds zero conditional variance and therefore provides a stable optimization signal even after the output distribution has nearly converged.

##### Comparison with top- k k : no truncation bias.

Top- k k OPD eliminates the per-position sampling variance of sampled-token OPD by deterministically selecting the k k highest-probability tokens. However, it introduces a truncation bias : the loss is computed only over the student’s top- k k support S t S_{t} , so any teacher probability mass outside S t S_{t} is invisible to the gradient. When the student’s top- k k set does not coincide with the teacher’s (a common situation early in training and on difficult tokens), the student receives no signal to shift probability toward the teacher’s preferred tokens that lie outside S t S_{t} . This bias is systematic and cannot be reduced by training longer or increasing the batch size. OPRD avoids this problem entirely: its MSE objective ‖ h θ , t ( l ) − h T , t ( l ) ‖ 2 2 \|h_{\theta,t}^{(l)}-h_{T,t}^{(l)}\|_{2}^{2} operates on continuous d d -dimensional vectors with no token selection or truncation step, so every dimension of the teacher’s hidden state contributes to the gradient at every supervised position.

###### Theorem 2 (Hidden-state information beyond the LM head) .

Let W head ∈ ℝ | 𝒱 | × d W_{\mathrm{head}}\!\in\!\mathbb{R}^{|\mathcal{V}|\times d} have singular values σ 1 ≥ ⋯ ≥ σ d > 0 \sigma_{1}\!\geq\!\cdots\!\geq\!\sigma_{d}\!>\!0 with right-singular vectors v 1 , … , v d v_{1},\ldots,v_{d} , and define the effective null space 𝒩 W ≜ { Δ ​ h ∈ ℝ d : W head ​ Δ ​ h ∈ span ⁡ { 𝟏 } } \mathcal{N}_{W}\triangleq\{\Delta h\in\mathbb{R}^{d}:W_{\mathrm{head}}\,\Delta h\in\mathrm{span}\{\mathbf{1}\}\} , i.e., the set of hidden-state perturbations whose image under W head W_{\mathrm{head}} is an additive softmax-invariant shift. For any last-layer student/teacher hidden states h θ , h T ∈ ℝ d h_{\theta},h_{T}\!\in\!\mathbb{R}^{d} and any output-space OPD loss ℓ out \ell_{\mathrm{out}} (sampled-token, top- k k , or full-vocabulary reverse KL), ℓ out ​ ( h θ , h T ) = 0 whenever h θ − h T ∈ 𝒩 W , \ell_{\mathrm{out}}(h_{\theta},h_{T})=0\quad\text{whenever}\quad h_{\theta}-h_{T}\in\mathcal{N}_{W}, (9) and along h θ − h T = α ​ v d h_{\theta}-h_{T}=\alpha v_{d} with ‖ v d ‖ = 1 \|v_{d}\|=1 , ‖ h θ − h T ‖ 2 / ℓ out ​ ( h θ , h T ) ≳ ( σ 1 / σ d ) 2 , \|h_{\theta}-h_{T}\|^{2}/\ell_{\mathrm{out}}(h_{\theta},h_{T})\;\gtrsim\;(\sigma_{1}/\sigma_{d})^{2}, (10) where ≳ \gtrsim hides a constant depending only on ℓ out \ell_{\mathrm{out}} and the logit range (made precise in Section B.5 ).

The ratio in ( 10 ) scales as ( σ 1 / σ d ) 2 (\sigma_{1}/\sigma_{d})^{2} , which is typically very large for production LLMs due to the ill-conditioned singular spectrum of W head W_{\mathrm{head}} . This means hidden-state deviations along low-singular-value directions can be orders of magnitude larger than along top directions while producing the same output-space loss; moreover output-space OPD has no mechanism to constrain intermediate hidden states h ( l ) h^{(l)} for l < L l\!<\!L . OPRD ( 6 ) penalizes exactly the directions in 𝒩 W \mathcal{N}_{W} and supervises any subset of intermediate layers, exposing (layers × \times positions × \times hidden-dim) scalars of structural information per sample that the LM-head projection necessarily compresses away (Limitation 2 in § 1 ).

### 3.3 OPRD-Bridge: Cross-Architecture Extension