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

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

• Alignment-aware pre-training. Our analysis shows that representational alignment, not architectural identity, is the true prerequisite for direct hidden-state distillation. If pre-training itself enforced alignment between large and small models (e.g., through shared layer initialisation or periodic representation-matching regularisation), the resulting pair might not require a low-rank bridge at all, making cross-architecture OPRD as seamless as the same-architecture case.

• Representation-level diagnostics for OPD. By opening up the hidden-state channel, OPRD enables representation-level diagnostics (cosine similarity, CKA, probing accuracy) to be tracked alongside traditional output-level metrics, providing a more complete mechanistic picture of how knowledge transfers between models during RL training.

• Tighter theoretical bounds. Our analysis identifies the qualitative mechanisms behind OPRD’s success. Quantifying these effects, including explicit convergence-rate bounds for OPRD vs. sampled-token OPD and spectral characterisations of which hidden-state directions the LM head nulls out, would solidify the theoretical foundation.

More broadly, our results suggest that hidden-state representations are an under-exploited resource in LLM distillation. We hope this work encourages the community to treat the teacher not merely as a probability oracle but as a structured source of layered internal computation that the student can learn to inhabit.

## Appendix A Notation Summary

## Appendix B Formal Theoretical Guarantees

The two theorems in § 3.2 ( Theorems 1 and 2 ) establish OPRD’s properties at an intuitive level. We now state both formally, in one-to-one correspondence with the main conclusions stated in § 3.2 : • § B.1 formalizes Theorem 1 (gradient variance) : variance gap ( Theorems 4 , 5 and 1 ) and signal-to-noise collapse ( Theorem 6 ).

• § B.5 formalizes Theorem 2 (LM-head information bottleneck) : the null-direction identity ( Theorem 7 ) and the spectral gap ( Theorem 8 ).

Throughout this section, all expectations are taken over a single fixed prompt x x and a single response position t t ; the multi-position case follows by linearity. We use θ \theta to denote the student parameters, π θ \pi_{\theta} and π T \pi_{T} to denote the student and teacher policies, and write p ≡ p t = π θ ( ⋅ ∣ x , y ^ < t ) p\equiv p_{t}=\pi_{\theta}(\cdot\mid x,\hat{y}_{<t}) and q ≡ q t = π T ( ⋅ ∣ x , y ^ < t ) q\equiv q_{t}=\pi_{T}(\cdot\mid x,\hat{y}_{<t}) for brevity. We let u t ≜ log ⁡ p − log ⁡ q u_{t}\triangleq\log p-\log q denote the per-token log-density ratio.

### B.1 Setup and Assumptions

###### Definition 1 (Stochastic gradient estimators) .

Fix a student response position t t and a sampled-token estimator y ^ t ∼ p \hat{y}_{t}\sim p . The two per-position stochastic gradient estimators considered in this paper are g OPD ​ ( θ , y ^ t ) \displaystyle g_{\mathrm{OPD}}(\theta;\hat{y}_{t}) ≜ ∇ θ [ log ⁡ p ⁡ ( y ^ t ) − log ⁡ q ⁡ ( y ^ t ) ] = ∇ θ ​ log ​ p ​ ( y ^ t ) , \displaystyle\triangleq\nabla_{\theta}\bigl[\log p(\hat{y}_{t})-\log q(\hat{y}_{t})\bigr]=\nabla_{\theta}\log p(\hat{y}_{t}), (17) g OPRD ​ ( θ ) \displaystyle g_{\mathrm{OPRD}}(\theta) ≜ ∇ θ 1 d ​ ‖ h θ , t ( L ) − sg ⁡ ( h T , t ( L ) ) ‖ 2 2 , \displaystyle\triangleq\nabla_{\theta}\,\tfrac{1}{d}\,\bigl\|h_{\theta,t}^{(L)}-\mathrm{sg}(h_{T,t}^{(L)})\bigr\|_{2}^{2}, (18) where in ( 17 ) we used the fact that ∇ θ ​ log ​ q ​ ( y ^ t ) = 0 \nabla_{\theta}\log q(\hat{y}_{t})=0 because q q depends only on the (frozen) teacher. The corresponding population gradients are g ¯ OPD ​ ( θ ) ≜ 𝔼 y ^ t ∼ p ​ [ g OPD ] \bar{g}_{\mathrm{OPD}}(\theta)\triangleq\mathbb{E}_{\hat{y}_{t}\sim p}[g_{\mathrm{OPD}}] and g ¯ OPRD ​ ( θ ) ≜ g OPRD \bar{g}_{\mathrm{OPRD}}(\theta)\triangleq g_{\mathrm{OPRD}} (which is already deterministic in y ^ t \hat{y}_{t} ).

###### Assumption 1 (Standard regularity) .

The following standard conditions hold throughout: (R1) The log-densities log ⁡ p θ ​ ( v ) \log p_{\theta}(v) are twice continuously differentiable in θ \theta for every v ∈ 𝒱 v\in\mathcal{V} . (R2) The score s θ ​ ( v ) ≜ ∇ θ ​ log ​ p θ ​ ( v ) s_{\theta}(v)\triangleq\nabla_{\theta}\log p_{\theta}(v) satisfies 𝔼 p ​ [ ‖ s θ ‖ 2 2 ] < ∞ \mathbb{E}_{p}[\|s_{\theta}\|_{2}^{2}]<\infty (finite Fisher information). (R3) For each v v , | log ⁡ p θ ​ ( v ) − log ⁡ q ⁡ ( v ) | ≤ M |\log p_{\theta}(v)-\log q(v)|\leq M for some constant M < ∞ M<\infty on the trajectory of training (bounded log-ratio). (R4) The hidden state h θ , t ( L ) h_{\theta,t}^{(L)} is a continuously differentiable function of θ \theta with bounded Jacobian: ‖ ∇ θ h θ , t ( L ) ‖ op ≤ J < ∞ \|\nabla_{\theta}h_{\theta,t}^{(L)}\|_{\mathrm{op}}\leq J<\infty .

Conditions (R1)–(R2) hold for any LLM with softmax output; (R3) holds whenever both student and teacher assign nonzero probability to every supported token (e.g., after a small label-smoothing or temperature adjustment); (R4) holds for any Lipschitz transformer with bounded weights. These assumptions are mild and standard in the policy-gradient literature.

### B.2 Variance of Sampled-Token OPD

###### Lemma 1 (Score-function decomposition of OPD gradient) .

Under Assumption 1 , the OPD population gradient at position t t admits the score-function representation g ¯ OPD ​ ( θ ) = 𝔼 y ^ t ∼ p ​ [ u t ​ ( y ^ t ) ​ ∇ θ ​ log ⁡ p ⁡ ( y ^ t ) ] , u t ​ ( v ) ≜ log ⁡ p ⁡ ( v ) − log ⁡ q ⁡ ( v ) . \bar{g}_{\mathrm{OPD}}(\theta)\;=\;\mathbb{E}_{\hat{y}_{t}\sim p}\bigl[\,u_{t}(\hat{y}_{t})\,\nabla_{\theta}\log p(\hat{y}_{t})\,\bigr],\qquad u_{t}(v)\triangleq\log p(v)-\log q(v). (19)

###### Proof.

Starting from ( 17 ), write g ¯ OPD = 𝔼 p ​ [ ∇ θ ​ log ​ p ​ ( y ^ t ) ] \bar{g}_{\mathrm{OPD}}=\mathbb{E}_{p}[\nabla_{\theta}\log p(\hat{y}_{t})] . The unconditional expectation of the score is zero: 𝔼 p [ ∇ θ log p ( y ^ t ) ] = ∑ v p ( v ) ∇ θ log p ( v ) = ∑ v ∇ θ p ( v ) = ∇ θ ∑ v p ( v ) = ∇ θ 1 = 0 . \mathbb{E}_{p}[\nabla_{\theta}\log p(\hat{y}_{t})]=\sum_{v}p(v)\,\nabla_{\theta}\log p(v)=\sum_{v}\nabla_{\theta}p(v)=\nabla_{\theta}\!\sum_{v}p(v)=\nabla_{\theta}1=0. Therefore g ¯ OPD \bar{g}_{\mathrm{OPD}} vanishes, which would make it useless as a learning signal. This apparent paradox is resolved by recognizing that what we actually optimize is the OPD loss surrogate along its stochastic gradient, which by the REINFORCE identity satisfies ∇ θ 𝔼 y ^ t ∼ p ​ [ log ⁡ p ⁡ ( y ^ t ) − log ⁡ q ⁡ ( y ^ t ) ] = 𝔼 p ​ [ ( log ⁡ p − log ⁡ q ) ​ ∇ θ ​ log ​ p ] + 𝔼 p ​ [ ∇ θ ​ log ​ p ] , \nabla_{\theta}\,\mathbb{E}_{\hat{y}_{t}\sim p}\!\bigl[\log p(\hat{y}_{t})-\log q(\hat{y}_{t})\bigr]=\mathbb{E}_{p}\!\bigl[(\log p-\log q)\,\nabla_{\theta}\log p\bigr]+\mathbb{E}_{p}[\nabla_{\theta}\log p], where the last term vanishes by the calculation above, giving ( 19 ). ∎

Lemma 1 shows that the OPD gradient is essentially a REINFORCE estimator with u t u_{t} playing the role of the reward. This structural property is what makes it high-variance.

###### Theorem 4 (OPD gradient variance lower bound) .

Under Assumption 1 , the conditional variance (conditioned on the prompt x x and prefix y ^ < t \hat{y}_{<t} ) of the single-sample OPD gradient satisfies Var ⁡ [ g OPD ​ ( θ , y ^ t ) ] = 𝔼 p ​ [ u t 2 ​ ‖ ∇ θ ​ log ​ p ‖ 2 2 ] − ‖ g ¯ OPD ​ ( θ ) ‖ 2 2 . \mathrm{Var}\bigl[g_{\mathrm{OPD}}(\theta;\hat{y}_{t})\bigr]\;=\;\mathbb{E}_{p}\!\Bigl[u_{t}^{2}\,\|\nabla_{\theta}\log p\|_{2}^{2}\Bigr]\;-\;\bigl\|\bar{g}_{\mathrm{OPD}}(\theta)\bigr\|_{2}^{2}. (20) Moreover, near the optimum where p → q p\to q , the variance is bounded below by Var ⁡ [ g OPD ] ≥ Var p ​ ( u t ) ⋅ ℱ min ​ ( θ ) − o ⁡ ( 1 ) ​ as ​ p → q , \mathrm{Var}\bigl[g_{\mathrm{OPD}}\bigr]\;\geq\;\mathrm{Var}_{p}(u_{t})\;\cdot\;\mathcal{F}_{\min}(\theta)\;-\;o(1)\;\;\text{as }p\to q, (21) where ℱ min ​ ( θ ) ≜ λ min ​ ( 𝔼 p ​ [ ∇ θ ​ log ​ p ​ ∇ θ ​ log ​ p ⊤ ] ) \mathcal{F}_{\min}(\theta)\triangleq\lambda_{\min}\!\bigl(\mathbb{E}_{p}[\nabla_{\theta}\log p\,\nabla_{\theta}\log p^{\top}]\bigr) is the minimum eigenvalue of the Fisher information matrix. In particular, Var ⁡ [ g OPD ] = Ω ⁡ ( Var p ​ ( u t ) ) \mathrm{Var}[g_{\mathrm{OPD}}]=\Omega(\mathrm{Var}_{p}(u_{t})) does not vanish as the loss approaches zero.

###### Proof.

The exact identity ( 20 ) follows from the definition of variance applied to the score-weighted estimator in Lemma 1 : Var ⁡ [ g OPD ] \displaystyle\mathrm{Var}[g_{\mathrm{OPD}}] = 𝔼 p ​ [ ‖ u t ​ ∇ θ ​ log ⁡ p ‖ 2 2 ] − ‖ 𝔼 p ​ [ u t ​ ∇ θ ​ log ⁡ p ] ‖ 2 2 \displaystyle=\mathbb{E}_{p}\bigl[\|u_{t}\nabla_{\theta}\log p\|_{2}^{2}\bigr]-\bigl\|\mathbb{E}_{p}[u_{t}\nabla_{\theta}\log p]\bigr\|_{2}^{2} = 𝔼 p ​ [ u t 2 ​ ‖ ∇ θ ​ log ​ p ‖ 2 2 ] − ‖ g ¯ OPD ‖ 2 2 , \displaystyle=\mathbb{E}_{p}\bigl[u_{t}^{2}\,\|\nabla_{\theta}\log p\|_{2}^{2}\bigr]-\|\bar{g}_{\mathrm{OPD}}\|_{2}^{2}, which is ( 20 ).

For the lower bound ( 21 ), we decompose u t = u ¯ + ( u t − u ¯ ) u_{t}=\bar{u}+(u_{t}-\bar{u}) where u ¯ ≜ 𝔼 p [ u t ] = D KL ( p ∥ q ) \bar{u}\triangleq\mathbb{E}_{p}[u_{t}]=D_{\mathrm{KL}}(p\|q) . Substituting into ( 20 ) and applying the Cauchy–Schwarz inequality to bound the cross term, 𝔼 p ​ [ u t 2 ​ ‖ ∇ θ ​ log ​ p ‖ 2 2 ] \displaystyle\mathbb{E}_{p}\bigl[u_{t}^{2}\|\nabla_{\theta}\log p\|_{2}^{2}\bigr] = u ¯ 2 ​ 𝔼 p ​ [ ‖ ∇ θ ​ log ​ p ‖ 2 2 ] + 𝔼 p ​ [ ( u t − u ¯ ) 2 ​ ‖ ∇ θ ​ log ​ p ‖ 2 2 ] \displaystyle=\bar{u}^{2}\,\mathbb{E}_{p}[\|\nabla_{\theta}\log p\|_{2}^{2}]+\mathbb{E}_{p}\bigl[(u_{t}-\bar{u})^{2}\|\nabla_{\theta}\log p\|_{2}^{2}\bigr] + 2 ​ u ¯ ​ 𝔼 p ​ [ ( u t − u ¯ ) ​ ‖ ∇ θ ​ log ​ p ‖ 2 2 ] . \displaystyle\quad+2\bar{u}\,\mathbb{E}_{p}\bigl[(u_{t}-\bar{u})\|\nabla_{\theta}\log p\|_{2}^{2}\bigr]. By the Cauchy–Schwarz / Rayleigh quotient argument, the middle term satisfies 𝔼 p ​ [ ( u t − u ¯ ) 2 ​ ‖ ∇ θ ​ log ​ p ‖ 2 2 ] ≥ Var p ​ ( u t ) ⋅ ℱ min ​ ( θ ) , \mathbb{E}_{p}\bigl[(u_{t}-\bar{u})^{2}\|\nabla_{\theta}\log p\|_{2}^{2}\bigr]\;\geq\;\mathrm{Var}_{p}(u_{t})\cdot\mathcal{F}_{\min}(\theta), since the covariance matrix of ∇ θ ​ log ​ p \nabla_{\theta}\log p is precisely the Fisher information matrix and its minimum eigenvalue lower-bounds any positive-definite quadratic form averaged over p p .

As p → q p\to q in total variation, the first and third terms above are O ⁡ ( u ¯ 2 ) + O ⁡ ( u ¯ ) O(\bar{u}^{2})+O(\bar{u}) , both of which vanish (since u ¯ = D KL ( p ∥ q ) → 0 \bar{u}=D_{\mathrm{KL}}(p\|q)\to 0 ). Meanwhile ‖ g ¯ OPD ‖ 2 2 = O ⁡ ( u ¯ 2 ) \|\bar{g}_{\mathrm{OPD}}\|_{2}^{2}=O(\bar{u}^{2}) , also o ⁡ ( 1 ) o(1) . Combining, Var ⁡ [ g OPD ] ≥ Var p ​ ( u t ) ⋅ ℱ min ​ ( θ ) − o ⁡ ( 1 ) \mathrm{Var}[g_{\mathrm{OPD}}]\geq\mathrm{Var}_{p}(u_{t})\cdot\mathcal{F}_{\min}(\theta)-o(1) , which is ( 21 ). Note that Var p ​ ( u t ) = Θ ⁡ ( δ ) \mathrm{Var}_{p}(u_{t})=\Theta(\delta) vanishes at the same rate as δ \delta , but crucially the signal ‖ g ¯ OPD ‖ 2 2 = O ⁡ ( δ 2 ) \|\bar{g}_{\mathrm{OPD}}\|_{2}^{2}=O(\delta^{2}) vanishes faster , leading to the SNR collapse formalized in Theorem 6 . ∎

### B.3 OPRD Gradient Is Deterministic

###### Theorem 5 (OPRD has zero conditional variance) .

Under Assumption 1 , the OPRD per-position gradient satisfies Var [ g OPRD ( θ ) | x , y ^ < t ] = 0 , \mathrm{Var}\bigl[g_{\mathrm{OPRD}}(\theta)\,\big|\,x,\hat{y}_{<t}\bigr]\;=\;0, (22) and is given in closed form by g OPRD ​ ( θ ) = 2 d ​ ( ∇ θ h θ , t ( L ) ) ⊤ ​ ( h θ , t ( L ) − h T , t ( L ) ) . g_{\mathrm{OPRD}}(\theta)\;=\;\frac{2}{d}\,\bigl(\nabla_{\theta}h_{\theta,t}^{(L)}\bigr)^{\!\top}\,\bigl(h_{\theta,t}^{(L)}-h_{T,t}^{(L)}\bigr). (23)

###### Proof.

Conditioned on the prompt x x and the prefix y ^ < t \hat{y}_{<t} , both h θ , t ( L ) h_{\theta,t}^{(L)} (a deterministic function of θ \theta , x x , y ^ < t \hat{y}_{<t} ) and h T , t ( L ) h_{T,t}^{(L)} (which has sg ⁡ ( ⋅ ) \mathrm{sg}(\cdot) applied, so it is treated as a constant in the gradient) are fixed. Hence the OPRD loss ℓ OPRD ≜ 1 d ​ ‖ h θ , t ( L ) − h T , t ( L ) ‖ 2 2 \ell_{\mathrm{OPRD}}\triangleq\tfrac{1}{d}\,\|h_{\theta,t}^{(L)}-h_{T,t}^{(L)}\|_{2}^{2} is a deterministic function of θ \theta given the conditioning. Therefore its gradient is also deterministic, giving ( 22 ).

The closed form ( 23 ) follows from the chain rule applied to the squared ℓ 2 \ell_{2} norm: ∇ θ 1 d ​ ‖ h θ , t ( L ) − h T , t ( L ) ‖ 2 2 \displaystyle\nabla_{\theta}\tfrac{1}{d}\|h_{\theta,t}^{(L)}-h_{T,t}^{(L)}\|_{2}^{2} = 2 d ​ J θ ​ ( h θ , t ( L ) ) ⊤ ​ ( h θ , t ( L ) − h T , t ( L ) ) , \displaystyle=\tfrac{2}{d}\,J_{\theta}(h_{\theta,t}^{(L)})^{\top}\,(h_{\theta,t}^{(L)}-h_{T,t}^{(L)}), where J θ ​ ( h θ , t ( L ) ) = ∇ θ h θ , t ( L ) ∈ ℝ d × dim ( θ ) J_{\theta}(h_{\theta,t}^{(L)})=\nabla_{\theta}h_{\theta,t}^{(L)}\in\mathbb{R}^{d\times\dim(\theta)} is the Jacobian. By (R4), ‖ J θ ‖ op ≤ J \|J_{\theta}\|_{\mathrm{op}}\leq J , so ‖ g OPRD ‖ 2 ≤ 2 ​ J d ​ ‖ h θ , t ( L ) − h T , t ( L ) ‖ 2 \|g_{\mathrm{OPRD}}\|_{2}\leq\tfrac{2J}{d}\,\|h_{\theta,t}^{(L)}-h_{T,t}^{(L)}\|_{2} , confirming that g OPRD g_{\mathrm{OPRD}} is well-defined and bounded. ∎

###### Corollary 1 (Variance gap) .

Combining Theorems 4 and 5 , the conditional variance gap between the two estimators is Var ⁡ [ g OPD ] − Var ⁡ [ g OPRD ] = 𝔼 p ​ [ u t 2 ​ ‖ ∇ θ ​ log ​ p ‖ 2 2 ] − ‖ g ¯ OPD ‖ 2 2 ≥ 0 , \mathrm{Var}\bigl[g_{\mathrm{OPD}}\bigr]-\mathrm{Var}\bigl[g_{\mathrm{OPRD}}\bigr]\;=\;\mathbb{E}_{p}\!\Bigl[u_{t}^{2}\,\|\nabla_{\theta}\log p\|_{2}^{2}\Bigr]-\|\bar{g}_{\mathrm{OPD}}\|_{2}^{2}\;\geq\;0, (24) with equality only in the degenerate case where p p is a point mass. In particular, OPRD’s gradient is always a lower-variance estimator (under the same conditioning), and the gap grows with the magnitude of the per-token log-ratio u t u_{t} and the spread of the policy.

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

We now formalize the most surprising prediction of our analysis: that the OPD signal-to-noise ratio collapses as training progresses, while OPRD’s signal-to-noise ratio remains bounded away from zero. This explains why pure OPD stagnates in late-stage training while OPD + + OPRD continues to improve monotonically.

###### Definition 2 (Signal-to-noise ratio) .

For a stochastic gradient estimator g g with population mean g ¯ = 𝔼 ⁡ [ g ] \bar{g}=\mathbb{E}[g] , the signal-to-noise ratio is SNR ⁡ ( g ) ≜ ‖ g ¯ ‖ 2 2 Tr ⁡ ( Cov ⁡ [ g ] ) . \mathrm{SNR}(g)\;\triangleq\;\frac{\|\bar{g}\|_{2}^{2}}{\mathrm{Tr}(\mathrm{Cov}[g])}. (25) SNR ⁡ ( g ) → 0 \mathrm{SNR}(g)\to 0 means the gradient is dominated by noise; SNR ⁡ ( g ) → ∞ \mathrm{SNR}(g)\to\infty means the gradient is essentially deterministic.

###### Theorem 6 (SNR collapse for OPD, SNR stability for OPRD) .

Define the symmetric divergence δ ( θ ) ≜ D KL ( p ∥ q ) + D KL ( q ∥ p ) \delta(\theta)\triangleq D_{\mathrm{KL}}(p\|q)+D_{\mathrm{KL}}(q\|p) . As training drives δ ⁡ ( θ ) → 0 \delta(\theta)\to 0 , 1. (OPD) SNR ⁡ ( g OPD ) = O ⁡ ( δ ) → 0 \mathrm{SNR}(g_{\mathrm{OPD}})=O(\delta)\to 0 at rate at least linear in δ \delta ;

2. (OPRD) SNR ⁡ ( g OPRD ) = + ∞ \mathrm{SNR}(g_{\mathrm{OPRD}})=+\infty as long as h θ , t ( L ) ≠ h T , t ( L ) h_{\theta,t}^{(L)}\neq h_{T,t}^{(L)} (i.e., the OPRD loss has not yet converged).

###### Proof.