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

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

(OPD case.) By Lemma 1 , ∥ g ¯ OPD ∥ 2 2 = ∥ 𝔼 p [ u t ∇ log p ] ∥ 2 2 \|\bar{g}_{\mathrm{OPD}}\|_{2}^{2}=\|\mathbb{E}_{p}[u_{t}\nabla\log p]\|_{2}^{2} . Applying Cauchy–Schwarz, ‖ g ¯ OPD ‖ 2 2 ≤ 𝔼 p ​ [ u t 2 ] ⋅ 𝔼 p ​ [ ‖ ∇ log ⁡ p ‖ 2 2 ] = ( Var p ​ ( u t ) + u ¯ 2 ) ⋅ Tr ⁡ ( ℱ ⁡ ( θ ) ) , \|\bar{g}_{\mathrm{OPD}}\|_{2}^{2}\;\leq\;\mathbb{E}_{p}[u_{t}^{2}]\cdot\mathbb{E}_{p}[\|\nabla\log p\|_{2}^{2}]\;=\;(\mathrm{Var}_{p}(u_{t})+\bar{u}^{2})\cdot\mathrm{Tr}(\mathcal{F}(\theta)), where ℱ ⁡ ( θ ) \mathcal{F}(\theta) is the Fisher information matrix. Since u ¯ = D KL ( p ∥ q ) ≤ δ \bar{u}=D_{\mathrm{KL}}(p\|q)\leq\delta and Var p ​ ( u t ) ≤ 2 ​ δ + O ⁡ ( δ 2 ) \mathrm{Var}_{p}(u_{t})\leq 2\delta+O(\delta^{2}) by a standard Pinsker-type expansion of log ⁡ ( p / q ) \log(p/q) around p = q p=q , we have ‖ g ¯ OPD ‖ 2 2 = O ⁡ ( δ ) . \|\bar{g}_{\mathrm{OPD}}\|_{2}^{2}\;=\;O(\delta). Meanwhile, by Theorem 4 (Eq. 21 ), Tr ⁡ ( Cov ⁡ [ g OPD ] ) ≥ Var p ​ ( u t ) ⋅ ℱ min ​ ( θ ) = Θ ⁡ ( δ ) \mathrm{Tr}(\mathrm{Cov}[g_{\mathrm{OPD}}])\geq\mathrm{Var}_{p}(u_{t})\cdot\mathcal{F}_{\min}(\theta)=\Theta(\delta) . Since the numerator is O ⁡ ( δ ) O(\delta) and the denominator is Ω ⁡ ( δ ) \Omega(\delta) , we have at first glance SNR ⁡ ( g OPD ) = O ⁡ ( δ ) Ω ⁡ ( δ ) = O ⁡ ( 1 ) ​ at best . \mathrm{SNR}(g_{\mathrm{OPD}})\;=\;\frac{O(\delta)}{\Omega(\delta)}\;=\;O(1)\;\;\text{at best}. A sharper analysis reveals that the numerator is in fact O ⁡ ( δ 2 ) O(\delta^{2}) : by the REINFORCE structure of Lemma 1 , ∥ g ¯ OPD ∥ 2 = ∥ 𝔼 p [ u t ∇ log p ] ∥ 2 ≤ 𝔼 p ​ [ u t 2 ] ⋅ Tr ⁡ ( ℱ ) \|\bar{g}_{\mathrm{OPD}}\|_{2}=\|\mathbb{E}_{p}[u_{t}\nabla\log p]\|_{2}\leq\sqrt{\mathbb{E}_{p}[u_{t}^{2}]}\cdot\sqrt{\mathrm{Tr}(\mathcal{F})} , and since 𝔼 p ​ [ u t 2 ] = Var p ​ ( u t ) + u ¯ 2 = Θ ⁡ ( δ ) + Θ ⁡ ( δ 2 ) = Θ ⁡ ( δ ) \mathbb{E}_{p}[u_{t}^{2}]=\mathrm{Var}_{p}(u_{t})+\bar{u}^{2}=\Theta(\delta)+\Theta(\delta^{2})=\Theta(\delta) , the Cauchy–Schwarz bound gives ‖ g ¯ OPD ‖ 2 2 = O ⁡ ( δ ) \|\bar{g}_{\mathrm{OPD}}\|_{2}^{2}=O(\delta) . However, this upper bound is not tight: the actual signal g ¯ OPD = 𝔼 p [ u t ∇ log p ] \bar{g}_{\mathrm{OPD}}=\mathbb{E}_{p}[u_{t}\nabla\log p] involves the correlation between u t u_{t} and ∇ log ⁡ p \nabla\log p , which is O ⁡ ( u ¯ ) = O ⁡ ( δ ) O(\bar{u})=O(\delta) in magnitude (since u t − u ¯ u_{t}-\bar{u} is mean-zero and contributes only through its correlation with ∇ log ⁡ p \nabla\log p , which is bounded by O ⁡ ( δ ) O(\sqrt{\delta}) ). Therefore ‖ g ¯ OPD ‖ 2 2 = O ⁡ ( δ 2 ) \|\bar{g}_{\mathrm{OPD}}\|_{2}^{2}=O(\delta^{2}) , giving SNR ⁡ ( g OPD ) = O ⁡ ( δ 2 ) / Ω ⁡ ( δ ) = O ⁡ ( δ ) → 0 \mathrm{SNR}(g_{\mathrm{OPD}})=O(\delta^{2})/\Omega(\delta)=O(\delta)\to 0 .

(OPRD case.) By Theorem 5 , Cov ⁡ [ g OPRD ] = 0 \mathrm{Cov}[g_{\mathrm{OPRD}}]=0 identically, so Tr ⁡ ( Cov ⁡ [ g OPRD ] ) = 0 \mathrm{Tr}(\mathrm{Cov}[g_{\mathrm{OPRD}}])=0 . As long as g OPRD ≠ 0 g_{\mathrm{OPRD}}\neq 0 (equivalently, h θ , t ( L ) ≠ h T , t ( L ) h_{\theta,t}^{(L)}\neq h_{T,t}^{(L)} ), the ratio in ( 25 ) is ‖ g OPRD ‖ 2 2 / 0 = + ∞ \|g_{\mathrm{OPRD}}\|_{2}^{2}/0=+\infty in the extended-real sense, meaning the gradient signal is completely noise-free. ∎

###### Remark 1 (Interpretation: late-stage stagnation of pure OPD) .

Theorem 6 predicts the following two-phase training dynamics for sampled-token OPD: • Phase 1 (effective learning). Initially δ ⁡ ( θ ) \delta(\theta) is large, so ‖ g ¯ OPD ‖ 2 \|\bar{g}_{\mathrm{OPD}}\|_{2} dominates the noise. The student improves rapidly along − g ¯ OPD -\bar{g}_{\mathrm{OPD}} .

• Phase 2 (stagnation). As δ ⁡ ( θ ) → 0 \delta(\theta)\to 0 , SNR ⁡ ( g OPD ) → 0 \mathrm{SNR}(g_{\mathrm{OPD}})\to 0 by Theorem 6 . The student’s update direction becomes effectively random, and under any positive learning rate the training accuracy plateaus or oscillates around an asymptote well below the teacher.

By contrast, OPRD’s SNR remains infinite throughout training (until convergence in hidden space), so the descent direction is always informative. This is exactly the empirical pattern we observe in § 4 : pure OPD plateaus or oscillates several points below the teacher, while OPD + + μ ⋅ \mu\cdot OPRD (and OPRD on its own) improves monotonically.

##### Sub-summary (Perspective 1).

The theorems above formalize two claims that explain OPRD’s variance advantage: (1) OPRD’s gradient is exactly deterministic and has zero conditional variance ( Theorem 5 ); (2) OPD’s gradient signal-to-noise ratio collapses to zero as the loss approaches its minimum ( Theorem 6 ), causing the late-stage stagnation we observe empirically.

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

We now make precise the two claims of Theorem 2 : the null-direction identity ( 9 ) and the spectral gap ( 10 ). Fix a single response position t t and write z θ ≜ W head ​ h θ ∈ ℝ | 𝒱 | z_{\theta}\triangleq W_{\mathrm{head}}\,h_{\theta}\in\mathbb{R}^{|\mathcal{V}|} and z T ≜ W head ​ h T z_{T}\triangleq W_{\mathrm{head}}\,h_{T} for the corresponding logit vectors; let σ : ℝ | 𝒱 | → Δ | 𝒱 | − 1 \sigma:\mathbb{R}^{|\mathcal{V}|}\!\to\!\Delta^{|\mathcal{V}|-1} denote the softmax map and 𝟏 ∈ ℝ | 𝒱 | \mathbf{1}\in\mathbb{R}^{|\mathcal{V}|} the all-ones vector. Let W head = U ​ Σ ​ V ⊤ W_{\mathrm{head}}=U\Sigma V^{\top} be the (thin) SVD, with V = [ v 1 , … , v d ] ∈ ℝ d × d V=[v_{1},\ldots,v_{d}]\in\mathbb{R}^{d\times d} orthonormal, Σ = diag ⁡ ( σ 1 ≥ ⋯ ≥ σ d ≥ 0 ) \Sigma=\mathrm{diag}(\sigma_{1}\!\geq\!\cdots\!\geq\!\sigma_{d}\!\geq\!0) , and U ∈ ℝ | 𝒱 | × d U\in\mathbb{R}^{|\mathcal{V}|\times d} having orthonormal columns. We assume throughout that W head W_{\mathrm{head}} has full column rank ( σ d > 0 \sigma_{d}>0 ), which holds for any production LLM.

By an output-space OPD loss we mean any ℓ out ≥ 0 \ell_{\mathrm{out}}\geq 0 that is a fixed function of the two output distributions σ ⁡ ( z θ ) \sigma(z_{\theta}) and σ ⁡ ( z T ) \sigma(z_{T}) and that vanishes whenever σ ⁡ ( z θ ) = σ ⁡ ( z T ) \sigma(z_{\theta})=\sigma(z_{T}) ; this includes the sampled-token estimator, the top- k k truncated reverse KL, and the full-vocabulary reverse KL (§ 2.3 ).

###### Definition 3 (Effective null space of the LM head) .

Define 𝒩 W ≜ { Δ ​ h ∈ ℝ d : W head ​ Δ ​ h ∈ span ⁡ { 𝟏 } } = W head − 1 ​ ( span ⁡ { 𝟏 } ) . \mathcal{N}_{W}\;\triangleq\;\bigl\{\Delta h\in\mathbb{R}^{d}:W_{\mathrm{head}}\,\Delta h\in\mathrm{span}\{\mathbf{1}\}\bigr\}\;=\;W_{\mathrm{head}}^{-1}\!\bigl(\mathrm{span}\{\mathbf{1}\}\bigr). (26)

###### Lemma 2 (Softmax kernel) .

For any z ∈ ℝ | 𝒱 | z\in\mathbb{R}^{|\mathcal{V}|} and any c ∈ ℝ c\in\mathbb{R} , σ ⁡ ( z + c ​ 𝟏 ) = σ ⁡ ( z ) \sigma(z+c\mathbf{1})=\sigma(z) . Conversely, σ ⁡ ( z ) = σ ⁡ ( z ′ ) \sigma(z)=\sigma(z^{\prime}) implies z ′ − z ∈ span ​ { 𝟏 } z^{\prime}-z\in\mathrm{span}\{\mathbf{1}\} .

###### Proof.

For the forward direction, the i i -th coordinate of σ ⁡ ( z + c ​ 𝟏 ) \sigma(z+c\mathbf{1}) is e z i + c ∑ j e z j + c = e c ​ e z i e c ​ ∑ j e z j = σ ​ ( z ) i \frac{e^{z_{i}+c}}{\sum_{j}e^{z_{j}+c}}=\frac{e^{c}e^{z_{i}}}{e^{c}\sum_{j}e^{z_{j}}}=\sigma(z)_{i} . For the converse, σ ⁡ ( z ) = σ ⁡ ( z ′ ) \sigma(z)=\sigma(z^{\prime}) implies z i − z j = z i ′ − z j ′ z_{i}-z_{j}=z^{\prime}_{i}-z^{\prime}_{j} for all i , j i,j (by taking logs of coordinate ratios), so z ′ − z z^{\prime}-z is a constant vector. ∎

###### Theorem 7 (Null-direction identity; formal version of ( 9 )) .

For any output-space OPD loss ℓ out \ell_{\mathrm{out}} as above, h θ − h T ∈ 𝒩 W ⟹ ℓ out ​ ( h θ , h T ) = 0 . h_{\theta}-h_{T}\in\mathcal{N}_{W}\quad\Longrightarrow\quad\ell_{\mathrm{out}}(h_{\theta},h_{T})=0. (27)

###### Proof.

If h θ − h T ∈ 𝒩 W h_{\theta}-h_{T}\in\mathcal{N}_{W} , then by ( 26 ) there exists c ∈ ℝ c\in\mathbb{R} with W head ​ ( h θ − h T ) = c ​ 𝟏 W_{\mathrm{head}}(h_{\theta}-h_{T})=c\mathbf{1} , i.e. z θ = z T + c ​ 𝟏 z_{\theta}=z_{T}+c\mathbf{1} . By Lemma 2 , σ ⁡ ( z θ ) = σ ⁡ ( z T ) \sigma(z_{\theta})=\sigma(z_{T}) . Since ℓ out \ell_{\mathrm{out}} vanishes whenever the two output distributions coincide, ℓ out ​ ( h θ , h T ) = 0 \ell_{\mathrm{out}}(h_{\theta},h_{T})=0 . ∎

To formalize ( 10 ) we need a local Lipschitz upper bound on any output-space loss in terms of ‖ z θ − z T ‖ 2 \|z_{\theta}-z_{T}\|_{2} . Such a bound holds for every standard ℓ out \ell_{\mathrm{out}} under mild regularity (e.g., logits bounded in a compact set), with a constant depending only on ℓ out \ell_{\mathrm{out}} and the logit range:

###### Lemma 3 (Local Lipschitzness of output-space losses) .

Let ℓ out \ell_{\mathrm{out}} be the sampled-token, top- k k , or full-vocabulary reverse KL. On any compact logit region 𝒵 ⊂ ℝ | 𝒱 | \mathcal{Z}\subset\mathbb{R}^{|\mathcal{V}|} , there exists C ℓ < ∞ C_{\ell}<\infty such that ℓ out ​ ( h θ , h T ) ≤ C ℓ ​ ‖ z θ − z T − c ∗ ​ 𝟏 ‖ 2 2 for all z θ , z T ∈ 𝒵 , \ell_{\mathrm{out}}(h_{\theta},h_{T})\;\leq\;C_{\ell}\,\|z_{\theta}-z_{T}-c^{\!*}\mathbf{1}\|_{2}^{2}\quad\text{for all}\quad z_{\theta},z_{T}\in\mathcal{Z}, (28) where c ∗ = 1 | 𝒱 | ​ 𝟏 ⊤ ​ ( z θ − z T ) c^{\!*}=\tfrac{1}{|\mathcal{V}|}\mathbf{1}^{\top}(z_{\theta}-z_{T}) projects out the softmax-invariant direction.

###### Proof sketch.

The reverse KL D KL ( σ ( z T ) ∥ σ ( z θ ) ) D_{\mathrm{KL}}(\sigma(z_{T})\|\sigma(z_{\theta})) has gradient and Hessian in z θ z_{\theta} that are continuous in z θ z_{\theta} and vanish at z θ = z T + c ∗ ​ 𝟏 z_{\theta}=z_{T}+c^{\!*}\mathbf{1} ; on a compact 𝒵 \mathcal{Z} its Hessian is operator-norm bounded. A second-order Taylor expansion in z θ − z T z_{\theta}-z_{T} around the additive-invariance optimum then yields ( 28 ) with C ℓ C_{\ell} proportional to half the Hessian’s operator-norm bound on 𝒵 \mathcal{Z} . The sampled-token and top- k k estimators are pointwise convex combinations of the full-vocabulary log-ratios and inherit the same upper bound (up to a constant). ∎

###### Theorem 8 (Spectral gap; formal version of ( 10 )) .

Under the setup of this section and the bound ( 28 ), for any α ∈ ℝ ∖ { 0 } \alpha\!\in\!\mathbb{R}\setminus\{0\} and the bottom right-singular vector v d v_{d} with ‖ v d ‖ 2 = 1 \|v_{d}\|_{2}=1 , ‖ h θ − h T ‖ 2 2 ℓ out ​ ( h θ , h T ) ≥ 1 C ℓ ​ ( σ 1 σ d ) 2 when h θ − h T = α ​ v d . \frac{\|h_{\theta}-h_{T}\|_{2}^{2}}{\ell_{\mathrm{out}}(h_{\theta},h_{T})}\;\geq\;\frac{1}{C_{\ell}}\,\Bigl(\frac{\sigma_{1}}{\sigma_{d}}\Bigr)^{\!2}\quad\text{when}\quad h_{\theta}-h_{T}=\alpha v_{d}. (29) In particular, holding ℓ out \ell_{\mathrm{out}} fixed, hidden-state perturbations along v d v_{d} can grow σ 1 / σ d \sigma_{1}/\sigma_{d} times larger in ℓ 2 \ell_{2} norm than perturbations along the top singular direction v 1 v_{1} .

###### Proof.

Take Δ ​ h = α ​ v d \Delta h\!=\!\alpha v_{d} . Then ‖ Δ ​ h ‖ 2 2 = α 2 \|\Delta h\|_{2}^{2}=\alpha^{2} and, by the SVD, W head ​ Δ ​ h = α ​ σ d ​ u d W_{\mathrm{head}}\Delta h=\alpha\sigma_{d}u_{d} where u d u_{d} is the corresponding left-singular vector with ‖ u d ‖ 2 = 1 \|u_{d}\|_{2}=1 . Since u d ⟂ 𝟏 u_{d}\perp\mathbf{1} generically (or after subtracting its component along 𝟏 \mathbf{1} via the projector in ( 28 )), the residual after removing the additive-invariance direction satisfies ‖ z θ − z T − c ∗ ​ 𝟏 ‖ 2 ≤ ‖ W head ​ Δ ​ h ‖ 2 = α ​ σ d \|z_{\theta}-z_{T}-c^{\!*}\mathbf{1}\|_{2}\leq\|W_{\mathrm{head}}\Delta h\|_{2}=\alpha\sigma_{d} . By Lemma 3 , ℓ out ≤ C ℓ ​ α 2 ​ σ d 2 \ell_{\mathrm{out}}\leq C_{\ell}\,\alpha^{2}\sigma_{d}^{2} , i.e. α 2 ≥ ℓ out / ( C ℓ ​ σ d 2 ) \alpha^{2}\geq\ell_{\mathrm{out}}/(C_{\ell}\sigma_{d}^{2}) . Therefore ‖ h θ − h T ‖ 2 2 ℓ out = α 2 ℓ out ≥ 1 C ℓ ​ σ d 2 . \frac{\|h_{\theta}-h_{T}\|_{2}^{2}}{\ell_{\mathrm{out}}}\;=\;\frac{\alpha^{2}}{\ell_{\mathrm{out}}}\;\geq\;\frac{1}{C_{\ell}\sigma_{d}^{2}}. For comparison, the analogous bound along v 1 v_{1} is ‖ h θ − h T ‖ 2 2 / ℓ out ≤ 1 / ( c ℓ ​ σ 1 2 ) \|h_{\theta}-h_{T}\|_{2}^{2}/\ell_{\mathrm{out}}\leq 1/(c_{\ell}\sigma_{1}^{2}) for the lower Lipschitz constant c ℓ c_{\ell} of ℓ out \ell_{\mathrm{out}} , so the ratio between the two directions scales as ( σ 1 / σ d ) 2 (\sigma_{1}/\sigma_{d})^{2} up to constants determined by ℓ out \ell_{\mathrm{out}} , recovering ( 29 ) after absorbing constants into C ℓ C_{\ell} . ∎

###### Remark 2 (Intermediate layers) .

Theorems 7 and 8 concern only the last-layer hidden state, because any output-space ℓ out \ell_{\mathrm{out}} is computed solely from W head ​ h ( L ) W_{\mathrm{head}}\,h^{(L)} and therefore has no functional dependence on intermediate states h ( l ) h^{(l)} for l < L l\!<\!L . For any l < L l\!<\!L , an arbitrary perturbation of h ( l ) h^{(l)} that leaves h ( L ) h^{(L)} unchanged (e.g., a perturbation in the kernel of the residual stack from layer l l onwards) yields ℓ out = 0 \ell_{\mathrm{out}}=0 for every output-space objective. OPRD ( 6 ) with ℒ layer ∋ l \mathcal{L}_{\mathrm{layer}}\ni l directly penalizes ‖ h θ , t ( l ) − h T , t ( l ) ‖ 2 2 \|h_{\theta,t}^{(l)}-h_{T,t}^{(l)}\|_{2}^{2} at that layer and is therefore the only mechanism considered in this paper that can constrain intermediate hidden states.

##### Sub-summary ( Theorem 2 ).

Theorem 7 formalizes ( 9 ): every output-space distillation objective treats the entire affine subspace 𝒩 W \mathcal{N}_{W} as invisible, regardless of how much it inspects the output distribution. Theorem 8 formalizes ( 10 ): the LM head’s singular-value spread σ 1 / σ d \sigma_{1}/\sigma_{d} amplifies hidden-state deviations along v d v_{d} by a ( σ 1 / σ d ) 2 (\sigma_{1}/\sigma_{d})^{2} factor for the same output-space loss budget, empirically 10 6 ∼ 10 8 × 10^{6}\!\sim\!10^{8}\times for production LLMs. Remark 2 extends both observations to intermediate layers. OPRD ( 6 ) penalizes exactly the directions and layers that output-space OPD cannot.

##### Overall summary of § B .

The theorems above formalize Theorem 1 (gradient variance and SNR) in one-to-one correspondence with the intuitive claims of § 3.2 : they explain why OPRD provides a more reliable optimization signal than sampled-token OPD, especially in late-stage training, and why adding OPRD to OPD strictly improves the SGD convergence bound without introducing additional noise. These guarantees apply under mild regularity conditions that hold for any standard LLM, supporting our empirical observation that combining OPD with OPRD yields a stronger and more stable result than OPD alone.

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

We now formalize the claim that the PCA-based low-rank bridge is optimal in an information-theoretic sense and that the bridge rank r r controls a bias–variance trade-off.

#### B.6.1 Setup

Consider a single layer pair ( l S , l T ) (l_{S},l_{T}) (we drop layer indices for clarity). Let h T ∈ ℝ d T h_{T}\in\mathbb{R}^{d_{T}} denote the teacher hidden state at a response position, drawn from a distribution with mean μ T \mu_{T} and covariance Σ T ≻ 0 \Sigma_{T}\succ 0 . Denote the eigendecomposition Σ T = V ​ Λ ​ V ⊤ \Sigma_{T}=V\Lambda V^{\top} , where Λ = diag ⁡ ( λ 1 , … , λ d T ) \Lambda=\mathrm{diag}(\lambda_{1},\ldots,\lambda_{d_{T}}) with λ 1 ≥ ⋯ ≥ λ d T > 0 \lambda_{1}\geq\cdots\geq\lambda_{d_{T}}>0 , and V = [ v 1 , … , v d T ] V=[v_{1},\ldots,v_{d_{T}}] are the corresponding eigenvectors (principal directions).