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

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

The OPRD-Vanilla objective defined in § 3.1 implicitly assumes that the student and teacher share a well-aligned representation space : the MSE in Eq. ( 6 ) directly compares h θ , t ( l ) h_{\theta,t}^{(l)} and h T , t ( l ) h_{T,t}^{(l)} in the same ℝ d \mathbb{R}^{d} , which is meaningful only when corresponding dimensions carry comparable semantics. Sharing the same model architecture is a sufficient condition but not the fundamental one. In the same-architecture experiments of § 3.1 , the two models share a common origin: JustRL-1.5B (teacher) was obtained by applying RL to R1-Distill-1.5B (student), which perturbs the representation space only mildly; the two models therefore retain a high degree of representational alignment, and the full ℝ d \mathbb{R}^{d} channel acts as a high-quality all-pass filter that transmits every dimension of the teacher’s hidden state faithfully. Conversely, two models with identical architecture but trained from different initialisations or with different tokenizers would have misaligned representation spaces, making direct MSE comparison meaningless despite matching dimensionality.

The true bottleneck, therefore, is representational alignment , not architectural identity. When this alignment is absent (as is typically the case for models that differ in depth ( L S ≠ L T L_{S}\neq L_{T} ), width ( d S ≠ d T d_{S}\neq d_{T} ), or training history), the hidden states are incommensurable, and we need a mechanism to bridge the gap. This section introduces OPRD-Bridge , which identifies and exploits a shared low-rank structure between heterogeneous models to construct such a bridge.

The key observation is that although the full representations are incomparable, their principal variation directions are not. We first demonstrate this empirically (§ 3.3.1 ), then show how to exploit it (§ 3.3.2 , 3.3.3 ), and finally explain why the resulting design takes the specific form it does (§ 3.3.4 ).

#### 3.3.1 The Shared Low-Rank Structure

Before presenting the method, we establish the empirical fact that motivates it. We take a Qwen3-4B teacher (36 layers, d T = 2560 d_{T}\!=\!2560 ) and a Qwen3-1.7B student (28 layers, d S = 2048 d_{S}\!=\!2048 ) as an illustrative example, sample 2K prompts from DAPO-Math-17K, generate responses with the student, and collect both models’ hidden states on these responses. We then perform PCA on the centered teacher hidden states at each layer, extracting the top- r r principal directions as a basis P T ∈ ℝ r × d T P_{T}\in\mathbb{R}^{r\times d_{T}} . Next, we train a simple linear map P S ∈ ℝ r × d S P_{S}\in\mathbb{R}^{r\times d_{S}} to minimize ‖ P S ​ h S − P T ​ ( h T − μ T ) ‖ 2 2 \|P_{S}h_{S}-P_{T}(h_{T}-\mu_{T})\|_{2}^{2} with both models frozen.

The result is striking. At r = 8 r\!=\!8 , the cosine similarity between P S ​ ( h S ) P_{S}(h_{S}) and P T ​ ( h T − μ T ) P_{T}(h_{T}-\mu_{T}) reaches 95% averaged across all layer pairs. Eight dimensions suffice to capture the shared structure between two models that, in their full representation spaces, show zero alignment.

Equally revealing: this similarity decreases as r r increases. At r = 32 r\!=\!32 it drops to 94.1 % 94.1\% , at r = 256 r\!=\!256 to 87.9 % 87.9\% , and continues declining to 77.0 % 77.0\% at full rank r = 2048 r\!=\!2048 . This is not a failure of the method; it is a signal about the data. The shared structure between teacher and student is concentrated in very few principal directions. Beyond those directions, the two models diverge, encoding architecture-specific information that is not shared and should not be distilled.

This observation has a direct methodological implication: the bridge should be low-rank by design, not as a computational shortcut, but because the shared structure itself is low-rank. Higher rank does not mean more information transfer; it means more noise.

#### 3.3.2 Stage 1: Bridge Construction

The bridge is built with both model weights frozen ; only the projectors are involved. We write h S h_{S} for the student hidden states in this stage (emphasizing that the student backbone is fixed), reserving h θ h_{\theta} for Stage 2 where the backbone parameters θ \theta are updated.

We first establish a layer correspondence. When L S ≠ L T L_{S}\neq L_{T} , we pair each student layer l S l_{S} with a teacher layer via proportional spacing: ϕ ⁡ ( l S ) = round ⁡ ( l S − 1 L S − 1 ⋅ ( L T − 1 ) ) + 1 , \phi(l_{S})=\mathrm{round}\!\left(\frac{l_{S}-1}{L_{S}-1}\cdot(L_{T}-1)\right)+1, (11) so that the first layers align, the last layers align, and intermediate layers are evenly distributed.

For each teacher layer l T l_{T} , we compute the PCA basis from the centered hidden states: P T ( l T ) = TopPrincipalDirections ⁡ ( h T ( l T ) − μ T ( l T ) , r ) ∈ ℝ r × d T . P_{T}^{(l_{T})}=\mathrm{TopPrincipalDirections}\!\left(h_{T}^{(l_{T})}-\mu_{T}^{(l_{T})},\;r\right)\in\mathbb{R}^{r\times d_{T}}. (12) This is computed once and permanently frozen. It captures the directions along which the teacher’s representations vary most. In practice, we obtain P T ( l T ) P_{T}^{(l_{T})} via PCA on the sample covariance rather than a thin SVD of the hidden-state matrix. Stacking the teacher’s response-token hidden states into H ∈ ℝ n × d T H\in\mathbb{R}^{n\times d_{T}} , the two are equivalent up to centering: the top- r r right singular vectors of the centered matrix H − 𝟏 ​ μ T ⊤ H-\mathbf{1}\mu_{T}^{\top} coincide with the top- r r eigenvectors of the sample covariance Σ = 1 n − 1 ​ ( H − 𝟏 ​ μ T ⊤ ) ⊤ ​ ( H − 𝟏 ​ μ T ⊤ ) ∈ ℝ d T × d T \Sigma=\frac{1}{n-1}(H-\mathbf{1}\mu_{T}^{\top})^{\top}(H-\mathbf{1}\mu_{T}^{\top})\in\mathbb{R}^{d_{T}\times d_{T}} . We therefore form Σ \Sigma and take its eigendecomposition, which is far cheaper than an SVD of H H in the typical regime n ≫ d T n\gg d_{T} (the cost is governed by the d T × d T d_{T}\times d_{T} eigensolve rather than the large row count n n ), and we subsample at most 16,384 16{,}384 token rows to estimate Σ \Sigma . A single eigendecomposition per layer yields the full ordered basis; the rank- r r projector for any r r is then obtained by slicing the top- r r eigenvectors, without re-running the decomposition.

For each student layer l S l_{S} , we then train a linear projector P S ( l S ) ∈ ℝ r × d S P_{S}^{(l_{S})}\in\mathbb{R}^{r\times d_{S}} to minimize: ℒ bridge = ∑ l S ∑ t ‖ P S ( l S ) ​ h S , t ( l S ) − P T ( ϕ ⁡ ( l S ) ) ​ ( h T , t ( ϕ ⁡ ( l S ) ) − μ T ( ϕ ⁡ ( l S ) ) ) ‖ 2 2 . \mathcal{L}_{\mathrm{bridge}}=\sum_{l_{S}}\sum_{t}\left\|P_{S}^{(l_{S})}h_{S,t}^{(l_{S})}-P_{T}^{(\phi(l_{S}))}\!\bigl(h_{T,t}^{(\phi(l_{S}))}-\mu_{T}^{(\phi(l_{S}))}\bigr)\right\|_{2}^{2}. (13) After 20 epochs of training (with both models frozen), P S P_{S} converges and is also frozen. The bridge is now complete: a pair of fixed linear maps ( P T , P S ) (P_{T},P_{S}) that project both architectures into a shared r r -dimensional subspace where their representations are directly comparable.

#### 3.3.3 Stage 2: Bridge-Based Distillation

With the bridge frozen, we perform on-policy representation distillation. The student generates rollouts, both models perform forward passes, and the loss is computed through the frozen bridge : ℓ ( l ) ​ ( θ , y ^ ) = 1 ∑ t m t ​ ∑ t = 1 T m t ​ ‖ P S ( l ) ​ ( h θ , t ( l ) ) − P T ( ϕ ⁡ ( l ) ) ​ ( h T , t ( ϕ ⁡ ( l ) ) − μ T ( ϕ ⁡ ( l ) ) ) ‖ 2 2 , \ell^{(l)}(\theta;\hat{y})=\frac{1}{\sum_{t}m_{t}}\sum_{t=1}^{T}m_{t}\,\left\|P_{S}^{(l)}(h_{\theta,t}^{(l)})-P_{T}^{(\phi(l))}\!\bigl(h_{T,t}^{(\phi(l))}-\mu_{T}^{(\phi(l))}\bigr)\right\|_{2}^{2}, (14) where m t m_{t} is the same position mask as in Eq. ( 6 ) (defaulting to the last k = 2000 k\!=\!2000 tokens). The full objective averages over layers: ℒ OPRD-Bridge ( θ ) = 𝔼 x , y ^ ∼ π θ ( ⋅ ∣ x ) [ 1 | ℒ layer | ∑ l ∈ ℒ layer ℓ ( l ) ( θ ; y ^ ) ] . \mathcal{L}_{\text{OPRD-Bridge}}(\theta)=\mathbb{E}_{x,\,\hat{y}\sim\pi_{\theta}(\cdot\mid x)}\left[\frac{1}{|\mathcal{L}_{\mathrm{layer}}|}\sum_{l\in\mathcal{L}_{\mathrm{layer}}}\ell^{(l)}(\theta;\hat{y})\right]. (15) Only the student backbone θ \theta receives gradients. In practice, we L2-normalize both projected vectors before computing MSE, making the loss scale-invariant. As with OPRD-Vanilla, the bridge objective composes additively with any output-space OPD variant via Eq. ( 7 ).

Algorithm 1 summarizes the complete pipeline.

#### 3.3.4 Design Philosophy

Every design choice in OPRD-Bridge follows from a single principle: the bridge should behave like the vocabulary in logit distillation , a fixed, pre-established interface that neither model modifies during training. The frozen, low-rank bridge plays two complementary roles:

• A stable interface that preserves pre-existing alignment. The projectors ( P T , P S ) (P_{T},P_{S}) encode an alignment that already exists between the two models’ representations before any distillation takes place; Stage 1 merely discovers it and freezes it into a fixed pair of linear maps. This mirrors logit distillation, where the shared vocabulary provides a stable coordinate system that neither model modifies during training: if this common reference frame changed, the optimization target would be non-stationary and training would destabilize. Freezing the bridge in Stage 2 ensures the on-policy training signal flows through this stable channel rather than destroying it. If P S P_{S} were trained jointly with the backbone, the projection and the target would shift simultaneously, erasing the very alignment the bridge was built to exploit; empirically, joint training indeed performs worse.

• A low-pass filter that suppresses high-frequency noise. The shared structure between heterogeneous models is intrinsically low-rank: the top- r r principal components correspond to the “low-frequency” modes where teacher and student are already highly aligned, while the orthogonal complement carries “high-frequency” architecture-specific variation that acts as noise for distillation. By projecting through a rank- r r bridge, we retain the well-aligned signal and discard the misaligned noise, which is precisely why Stage 2 training remains stable even under on-policy distribution shift. It is instructive to contrast this with the two all-pass alternatives. Full-vocabulary logit distillation transmits every dimension of the output space through the shared vocabulary, a high-quality all-pass channel, but at O ⁡ ( | 𝒱 | ) O(|\mathcal{V}|) cost and without any denoising. OPRD-Vanilla likewise operates as an all-pass channel over the full ℝ d \mathbb{R}^{d} representation space; this works precisely because same-architecture models with shared initialisation are already well-aligned in every direction, so no dimension is noise. When this alignment breaks down (different depth, width, or training history), the all-pass property becomes a liability: high-frequency directions carry misaligned noise, and a low-pass bridge is needed to filter it out.

Table 1 makes the bridge–vocabulary analogy explicit.

Why PCA for the teacher side? Because the Platonic Representation Hypothesis predicts that strong models converge in their principal variation directions [ 15 ] , and PCA extracts exactly these directions. It is also deterministic, data-efficient (2K prompts suffice), and introduces no trainable parameters on the teacher side.

Why such a low rank? Because the shared structure is low-rank: our pre-experiment shows 95% alignment at r = 8 r\!=\!8 , reflecting the concentration of task-relevant information in very few directions. The rank is therefore not a computational budget; it is a statement about the geometry of the shared representation, and the low-pass property described above is what makes the bridge a robust supervision channel. We formalize this intuition in the following theorem (stated informally; the precise version with proof is in Section B.6 ).

###### Theorem 3 (Optimality and bias–variance trade-off of the low-rank bridge, informal) .

Assume teacher hidden states are drawn from a distribution with covariance Σ T \Sigma_{T} (eigenvalues λ 1 ≥ ⋯ ≥ λ d T \lambda_{1}\!\geq\!\cdots\!\geq\!\lambda_{d_{T}} ), and that the cross-model alignment in direction i i is ρ i ∈ [ 0 , 1 ] \rho_{i}\in[0,1] (decreasing in i i : high-eigenvalue directions are better aligned). Then: 1. (Rate–distortion optimality.) Among all rank- r r linear encodings of the teacher representation, PCA minimizes the expected distillation error under Gaussian assumptions. The bridge is the information-theoretically optimal channel at capacity r r .

2. (Bias–variance decomposition.) The expected distillation error decomposes as ℰ ( r ) = ∑ i > r λ i ​ ρ i 2 ⏟ Bias : discarded ​ aligned ​ signal + ∑ i = 1 r λ i ​ ( 1 − ρ i 2 ) ⏟ Variance : transmitted ​ misaligned ​ noise . \mathcal{E}(r)=\underbrace{\textstyle\sum_{i>r}\lambda_{i}\,\rho_{i}^{2}}_{\mathrm{Bias:\ discarded\ aligned\ signal}}\;+\;\underbrace{\textstyle\sum_{i=1}^{r}\lambda_{i}\,(1-\rho_{i}^{2})}_{\mathrm{Variance:\ transmitted\ misaligned\ noise}}. (16) When ρ i \rho_{i} decreases with i i (empirically verified), there exists an optimal rank r ∗ r^{*} that minimizes ℰ ⁡ ( r ) \mathcal{E}(r) : below r ∗ r^{*} , bias dominates (useful signal is discarded); above r ∗ r^{*} , variance dominates (misaligned noise is transmitted).

The theorem explains three empirical observations simultaneously: (i) why a very low rank suffices (the λ i \lambda_{i} decay fast and ρ i \rho_{i} decay fast, so both bias and variance are small at r ∗ ≈ 8 r^{*}\!\approx\!8 ); (ii) why increasing rank beyond the sweet spot hurts alignment (variance term grows); and (iii) why freezing the bridge is critical (a moving P S P_{S} would invalidate the ρ i \rho_{i} estimates on which the optimality rests).

Finally, we note that because the bridge performs alignment in representation space rather than in the output space, it is decoupled from the vocabulary: the teacher and student need not share a tokenizer. This opens a path toward cross-tokenizer and potentially cross-modal distillation, where the vocabulary-based alignment channel is unavailable.

## 4 Experiments

We evaluate both OPRD variants: OPRD-Vanilla in the same-architecture setting (§ 4.1 ), and OPRD-Bridge in the cross-architecture setting (§ 4.2 ).

### 4.1 OPRD-Vanilla: Same-Architecture Distillation

We evaluate OPRD-Vanilla on competition-level mathematical reasoning, against (i) a frozen teacher and an unmodified student baseline, and (ii) two strong on-policy distillation baselines that share the same on-policy rollout and teacher forward pass as OPRD but extract supervision from the LM-head output. The experiments test the two predictions of § 3.2 : OPRD provides a lower-noise, structurally richer training signal than any output-space OPD variant, and should therefore close the student–teacher gap that pure OPD cannot.

#### 4.1.1 Experimental Setup

##### Models.

Following [ 25 ] , we use JustRL-Deepseek-1.5B [ 10 ] (denoted JustRL-1.5B ) as the (frozen) teacher and DeepSeek-R1-Distill-Qwen-1.5B [ 9 ] (denoted R1-distill-1.5B ) as the student. Both models share the Qwen2.5-1.5B backbone ( L = 28 L\!=\!28 transformer layers, d = 1536 d\!=\!1536 hidden dimension, | 𝒱 | ≈ 151 |\mathcal{V}|\!\approx\!151 K vocabulary) and the same LM head W head W_{\mathrm{head}} , so OPRD-Vanilla’s hidden-state targets are directly comparable across the two models ( d S = d T d_{S}\!=\!d_{T} ). The student starts from the public R1-distill-1.5B checkpoint, which already places it close to but well below the teacher in reasoning ability ( Table 2 ).

##### Training data.

On-policy prompts x x are drawn from DAPO-Math-17K [ 45 ] . For each prompt the student samples 2 2 responses y ^ ∼ π θ ( ⋅ ∣ x ) \hat{y}\sim\pi_{\theta}(\cdot\!\mid\!x) at temperature 1.0 1.0 with a max generation length of 16,384 16{,}384 tokens; we use a global batch of 8 8 prompts per step.