RAPID: Increasing Efficiency of Internal Distillation via Randomized Methods

Introduction

Recent open-weight large language models (LLMs), including Kimi K3 Team et al. [42], DeepSeek V4 Pro DeepSeek-AI [10], Qwen 3.8 Max Qwen Team [35], and GLM 5.3 GLM5-Team et al. [14], have approached proprietary frontier performance across scientific reasoning, software engineering, multimodal understanding, long-context processing, and tool use. However, their substantial deployment costs make them impractical for many applications, particularly in healthcare, law, finance, cybersecurity, and autonomous vehicles, where compute, latency, and privacy constraints are stringent. Open weights enable these capabilities to be transferred into smaller, more efficient models.

Knowledge distillation transfers knowledge from a large teacher model to a smaller student. It has been used to pretrain models in Google’s Gemma 2 and 3 families Team et al. [40, 41] and Meta’s Llama 3.2 family Meta [30]. It is also effective for post-training Guo et al. [18], improving reasoning, planning, and tool-use capabilities in lightweight agentic systems.

However, text-based distillation is highly sensitive to the quality and diversity of teacher trajectories and their alignment with the student’s pretraining distribution Gudibande et al. [16]. Internal distillation supplements these trajectories with teacher activations, providing richer supervision and encouraging the student to process inputs similarly to the teacher.

Existing internal-distillation methods learn projections between teacher and student representation spaces before computing losses such as mean-squared error (MSE) or KL divergence Liang et al. [27]; Jiao et al. [22]; Sun et al. [38]. These methods were developed when large models contained roughly 150 million parameters, over four orders of magnitude fewer than today’s approximately two-trillion-parameter frontier models. As shown in Table 2, at modern scales, learned projectors may require training billions of additional parameters beyond the student itself. Consequently, the cost of internal distillation remains a major barrier to further exploration and adoption.

Teacher	Student	Added Params		Added Params (%)		RAPID Efficiency
		Residual	All	Residual	All	Residual	All
Gemma 3 4B	Gemma 3 270M	29.5M	90.1M	6.8%	20.7%	+6.3%	+17.1%
Qwen 2.5 3B	Qwen 2.5 0.5B	44M	134M	8.8%	26.8%	+8.1%	+21.1%
Llama 3.1 8B	Llama 3.2 1B	134.2M	411M	13.4%	41.1%	+11.8%	+29.1%
Olmo 3 32B	Olmo 3 7B	671.1M	2.03B	9.6%	29.0%	+8.8%	+22.5%

To efficiently distill the capabilities of large modern LLMs into smaller, more efficient models, we turn to randomization. Randomization has an extensive history in accelerating numerical computation and machine learning Mahoney [28]; Martinsson & Tropp [29], including specifically in dimension projection and information compression Drinea et al. [11]; Johnson & Lindenstrauss [23]; Achlioptas [1]. We also draw from work in dynamic weighting and multi-objective optimization Li et al. [26]; Wang et al. [45]; Kendall et al. [24] to formulate a novel objective function designed for the noisy signals of internal distillation.

Our core contributions are:

• In Section 4 we introduce RAPID (RAndomized Projections for Internal Distillation), a novel methodology for conducting internal distillation which eliminates the need for training projection matrices by replacing them with randomized matrices and leverages a novel dynamic multi-objective loss function

• In Sections 6 and 7.3 we demonstrate that RAPID achieves task performance comparable to existing methods while improving parameter efficiency by as much as 8.1% in the residual-only distillation scheme, and up to 21.1% with an expanded set of distilled activations

• In Section 7.1 we detail how RAPID enforces a separation of responsibilities between projection and representation learning, and Section 7.2 demonstrates that RAPID produces similar or improved representational alignment with the teacher

Preliminaries

Knowledge Distillation

Knowledge distillation (KD) was first introduced by Bucilua et al. [6] and Hinton et al. [21] was the landmark work that adapted this approach to neural networks. Recent advancements have built upon these foundational works to distill post-training behaviors such as reasoning and instruction-following using synthetic data traces [17, 18, 46] and to distill internal activations in addition to output logits [27, 22, 38].

Logit and internal-activation distillation require alignment between teacher and student tokenizers. Because LLMs are sensitive to tokenization and cross-family vocabularies are difficult to align, these methods are typically restricted to models within the same family. Internal distillation must also reconcile differences in hidden dimensionality,¹ since losses such as mean-squared error and cosine similarity require equal-dimensional vectors. Teacher and student activations must therefore be projected into a shared representation space.

Randomized Methods for Efficient Computation

Randomization accelerates numerical computation and machine learning by replacing expensive exact operations with efficient approximations whose errors are probabilistically controlled [28, 29]. One of the foundational works in randomization for information compression is Johnson & Lindenstrauss [23] which introduced the Johnson-Lindenstrauss lemma (restated here in Lemma 1). Simply put, the JL lemma states that any set of $n$ points in high-dimensional Euclidean space can be embedded into a smaller dimensional space while approximately preserving all pairwise distances.

Mechanistic Interpretability

Mechanistic interpretability seeks to develop human-interpretable accounts of how LLMs store, recognize, and process information. Although individual neurons and model components are difficult to interpret in isolation [12, 3, 44], recent work suggests that the resulting activation spaces are not unstructured. Instead, LLMs represent concepts and semantic relationships through organized geometric patterns, including identifiable directions and low-dimensional subspaces [12, 5, 32, 19, 33].

Related Work

Knowledge Distillation

Response-trace distillation has proven effective for transferring post-training behaviors [46, 18] and supports both proprietary teachers and cross-family transfer. However, its effectiveness is sensitive to the scope, quality, and diversity of the trace distribution Gudibande et al. [16]. Logit distillation instead trains students on teacher soft targets, following Hinton et al. [21]. Google and Meta have applied this approach during pretraining for models up to 2.7B parameters [30, 40] (); Team et al. 2025), while Alibaba used it to consolidate multiple teacher post-training stages into a single student-training stage Team (2025).

Internal Distillation

Internal distillation goes beyond traditional output logit distillation to additionally distill internal activations into the student [27, 22, 38]. While logit distillation attempts to improve the generalization of the knowledge being distilled by training students to replicate the teacher’s output distribution, internal distillation also trains students to replicate the internal processing steps which produce the output logits. The intuition behind this approach is that the teacher is able to provide step-by-step guidance to how to student arrives at the final answer, a more granular and information-rich signal than output logits alone.

TinyBERT Jiao et al. [22] distills the embeddings, all residual stream activations, and all attention distribution matrices. Patient KD Sun et al. [38] reduces the computational cost of internal distillation by only distilling the internal activations of the [CLS] token$^{2}$ from a subset of layers. TED Liang et al. [27] performs task-specific distillation by training filters on top of residual stream activations to reduce the redundant information transferred to the student.

Randomized Projections for Information Compression

The original transform constructed by Johnson and Lindenstrauss applied a suitably scaled random orthogonal projection onto a $k$-dimensional subspace. Later work such as Achlioptas [1] improved efficiency by replacing conventional Gaussian projection matrix with one of two discrete random matrices. In the dense ”binary coin” construction, each entry is sampled independently as

In the sparse ternary construction, each entry is sampled independently as

Both distributions have zero mean and unit variance, allowing the full data matrix $\mathbf{X} \in\mathbb{R}^{n \times d}$ to be projected as

where $\mathbf{R} \in\mathbb{R}^{d \times k}$ is the random projection matrix. The dense construction replaces general multiplications with additions and subtractions, while the sparse construction additionally eliminates approximately two-thirds of the operations because two-thirds of its entries are zero.

Dynamic Objective Weighting

Wang et al. [45] demonstrated that knowledge distillation can benefit from different types of knowledge at different times during training. Kendall et al. [24] introduced uncertainty-based weighting for multi-objective optimization in deep learning, and Li et al. [26] employed it in internal distillation to allow for downweighting of representational alignment when student can already mimic teacher outputs. Lakkapragada et al. [25] leveraged exponential moving average (EMA) to dynamically normalize loss magnitudes for stability despite noisy updates.

Methodology

Intuition

There are three core features of RAPID that distinguish itself from existing internal distillation methodologies. The first is that, by using a fixed projector, we strictly enforce that the full responsibility of representation learning falls solely on the student$^{3}$. The second is that the usage of a JL transform provides guarantees about the preservation of structure in the projection of teacher representations. Recent mechanistic interpretability research provides evidence that LLMs represent semantic concepts and their relationships through structured geometric organization in their representation spaces Elhage et al. [12]; ); Bricken et al. [5]; ); Park et al. [32]; ); Gurnee & Tegmark [19]; ); Park et al. [33]. By preserving teacher representational geometry, RAPID targets the distillation of semantic relationships– a key component of advanced knowledge and intelligence. The final distinct feature is that RAPID allows the student to dynamically adapt its learning to target its weaknesses while appropriately weighting the noisy signals from distilling internal representations.

Randomized Projector

We use the dense binary coin JL transform described in Algorithm 1 to create the randomized projectors in RAPID. The dense formulation was selected on the basis of modern GPUs being highly optimized for dense matrix operations. Since the JL lemma applies only to data compression, we set $d$ to the dimension of the teacher representation space and $k$ to the dimension of the student representation space to construct a downward projection.

We construct a linearly-spaced layer mapping ($\mathcal{H}$) between the layers of the teacher and student such that every layer of the student is mapped to a layer in the teacher. For each activation being distilled, we independently initialize a separate projection matrix. For example, in our Gemma 3 experiments in § 5 where we distill only the residual activations, we create (18 layers) $\times$ (1 activation per layer) = 18 projectors. Our § 7.3 experiments with the expanded activation set construct (18 layers) $\times$ (3 activations per layer) + (1 projector for embeddings) = 55 projectors.

Multi-Objective Loss

We draw from recent work in knowledge distillation Li et al. [26]; Wang et al. [45] and utilize a dynamic objective function which combines learned uncertainty weighting Kendall et al. [24] with exponential moving average (EMA) progress balancing. Internal distillation is a multi-objective optimization problem, and using a dynamic objective function allows the model to prioritize different components of the objective at different points throughout training.

At training step $t$, given an input sequence $x$, raw losses are computed between the frozen teacher $\mathcal{M}_T$ and the student $\mathcal{M}_S$ across three levels: language modeling task loss on next-token labels, output distribution divergence, and internal activation matching error.

Task Loss The student is directly trained on target sequence labels $y$ using cross-entropy loss– the standard language modeling loss– over student output logits $z_S$:

Output Distribution Distillation Output alignment between teacher logits $z_T$ and student logits $z_S$ is enforced using forward Kullback-Leibler (KL) divergence computed at temperature $\tau$:

Intermediate Representation Distillation To align internal representations, we intercept internal activations for the teacher $\mathbf{a}_T^{(l)}$ and student $\mathbf{a}_S^{(l)}$ across matched hook layers $l \in\mathcal{H}$. Representation spaces are aligned using a linear projection head $P_\phi$. Depending on the projector, representation losses for internal activations $\mathcal{L}_{\mathrm{act}}$ are computed as the Mean Squared Error (MSE) summed over matching layer pairs $\mathcal{H}$:

§ 6 shows results when distilling only the residual activations, though our implementation can handle an expanded set of activations (see § 7.3). Inactive internal activation objectives for a given experiment are assigned zero ($\mathcal{L}_{\mathrm{act}} = 0$) and masked out.

Hybrid Dynamic Loss Weighting

To dynamically balance optimization among each term in the objective, we employ a dynamic weighting mechanism which unifies learned homoscedastic uncertainty weighting with an exponential moving average (EMA) progress ratio.

Let $\mathcal{K}$ represent the set of objective components. At step $t$, let $\mathcal{L}_{i,t}$ denote the unweighted scalar loss calculated for component $i \in\mathcal{K}$.

Tracking Objective Baselines To maintain historical loss scale baselines without retaining full step histories, an exponential moving average (EMA) buffer $\bar{\mathcal{L}}_{i,t}$ is tracked for each objective:

where $\eta\in(0,1)$ is the EMA decay rate and $\operatorname{sg}[\cdot]$ denotes the stop-gradient operator. We set $\eta= 0.99$.

Dual Signal Generation The unnormalized weighting signal for objective $i$ combines two distinct mechanisms:

1. Learned Uncertainty Weighting: Parameterized by learnable log-variance terms $s_i \in\mathbb{R}$, capturing intrinsic task noise characteristics:

2. EMA Progress Weighting: Measures transient loss spikes relative to historical baseline performance, scaled by hyperparameter $\alpha\ge0$:

where $\epsilon> 0$ provides numerical stability. We set $\alpha= 0.5$ and $\epsilon= 1 \times10^{-8}$.

Combination, Masking, and Bounding The combined raw weight for objective $i$ is calculated as the product of both signals, masked by indicator $m_i = \mathbb{I}(i \in\mathcal{A})$:

These raw weights are normalized across active objectives to project onto a probability simplex:

To prevent gradient starvation or single-objective dominance, normalized weights are clamped within lower and upper bounds $[\omega_{\min},\omega_{\max}]$ for active tasks and re-normalized:

We use $\omega_{\min} = 0.05,\omega_{\max} = 0.80$.

Total Objective Function The final scalar loss minimized at step $t$ combines dynamically scaled raw objectives with a linear log-variance regularization term over active tasks:

Optimization Mechanics All operations computing intermediate weights $w_{i,t}$ from $\mathcal{L}_{i,t}$ and $s_i$ are evaluated without gradient tracking. As a result, $\partial\mathcal{L}_{\mathrm{total}}/\partial s_i = 1$ for all active heads, preventing representation backpropagation from corrupting $s_i$ and avoiding scale-invariance degeneracies. Further details regarding this are provided in Appendix § 15.

Experiments

Comparisons

We compare RAPID with internal distillation using learned linear projectors, the most common approach in prior work [47, 48]. Following existing methods [22, 37, 27, 38], each projector maps a student activation upward into the teacher’s representation space. We use separate projectors for each distilled activation and equally weight the ground-truth, output-distribution, and internal-activation losses.

Models

We demonstrate RAPID using distillation pairs from the Gemma 3 Team et al. [41] and Qwen 2.5 Qwen et al. [34] model families. We exclusively perform in-family distillation, distilling Gemma 3 4B IT into Gemma 3 270M PT⁴, and Qwen 2.5 3B Instruct into Qwen 2.5 0.5B. These approximate parameter compression ratios of 9.2:1 for Gemma 3 and 6:1 for Qwen 2.5 follow the existing distillation literature range of 1.6:1 to 10:1 Jiao et al. [22]; Sanh et al. [37]; Meta [30]; Gu et al. [15] with absolute parameter counts scaled to modern LLM sizes.

We selected the student initialization point to be the base pretrained models for the cleanest starting point for our post-training. Each model is trained with three separate seeds {3, 33, 333} and the scores reported are averaged across the seeds.

Training

We train on 40M tokens of reasoning data with a sequence length of 1024⁵, corresponding to about 40k training samples per experiment. To show data generalization, we trained with 2 separate datasets: a version of SmolTalk Allal et al. [2] curated for reasoning and maximum 1024 sequence length per sample, and unmodified OpenThoughts-114k Guha et al. [17]. Given the chat format of these data, we use the tokenizer for the instruction-tuned variants of the student models.

Evaluation Tasks

We evaluate reasoning task performance on four benchmarks: MMLU Hendrycks et al. [20], ARC-Challenge Clark et al. [8], and the Open Leaderboard V2 Fourrier et al. [13] versions of BBH Suzgun et al. [39] and GPQA Rein et al. [36]. These are common benchmarks used in the literature for LLMs of similar size to our student models Qwen et al. [34]; Team (2025); Team et al. [40] (2025); Biderman et al. [4]; Meta [30] and are designed to evaluate knowledge and reasoning capabilities. Following the conventions in the literature, we evaluate MMLU as 5-shot, ARC-Challenge as 25-shot, BBH as 3-shot, and GPQA as 0-shot.

Results

Task Performance

As shown in Table 1, two methods achieve nearly identical performance despite RAPID being over 6% more efficient in the Gemma 3 experiments, and over 8% more efficient in the Qwen 2.5 experiments. No difference exceeds 0.15 percentage points in the overall average, with largest individual delta being +1.05 percentage points for RAPID over Learned for SmolTalk on BBH.

Dataset	Distillation Pair	Variant	ARC-C	BBH	GPQA	MMLU	Average
OpenThoughts	Gemma 3	Learned	0.2770	0.2804	0.2438	0.2475	0.2622
OpenThoughts	Gemma 3	RAPID	0.2747	0.2848	0.2506	0.2381	0.2621
OpenThoughts	Qwen 2.5	Learned	0.3376	0.3018	0.2508	0.2560	0.2866
OpenThoughts	Qwen 2.5	RAPID	0.3348	0.3005	0.2539	0.2520	0.2853
SmolTalk	Gemma 3	Learned	0.2742	0.2898	0.2550	0.2355	0.2636
SmolTalk	Gemma 3	RAPID	0.2796	0.2895	0.2517	0.2382	0.2647
SmolTalk	Qwen 2.5	Learned	0.3490	0.3129	0.2581	0.3241	0.3110
SmolTalk	Qwen 2.5	RAPID	0.3549	0.3234	0.2503	0.3156	0.3111

Efficiency

RAPID improves training efficiency by eliminating learned projectors. Table 2 reports their additional parameter costs across internal distillation configurations, where All includes residual, MLP-output, attention-output, and embedding activations. We calculate RAPID’s efficiency gain as

Learned projectors increase trainable parameters by approximately 10% for residual-only distillation and 30% for All activations—up to 2B additional parameters for a 7B student. RAPID reduces trainable parameters by 8.1% in our Qwen experiments (§ 6.1) and 21.1% when distilling All activations (§ 7.3).

Analysis

Projector–Student Separation of Responsibilities

Internal-activation distillation ideally assigns distinct roles to the student and projector: the student learns the transferred teacher representation, while the projector performs a transformation to reconcile dimensionality. RAPID enforces this separation by applying a fixed downward projection $\mathbf{H}_{\mathrm{fix}} \in\mathbb{R}^{d_s \times d_t}$ to the teacher:

Because the projected teacher activation is stationary, the projector only defines the target; all adaptation required to reduce the loss must occur in the student.

Existing methods commonly utilize a learned upward projection $\mathbf{H}_{\mathrm{learn}} \in\mathbb{R}^{d_t \times d_s}$ Liang et al. [27]; Jiao et al. [22]; Sun et al. [38]:

Learned projectors are popular for their flexibility: beyond matching dimensions, they can learn particular transforms which facilitate student learning Miles & Mikolajczyk [31]; Chen et al. [7]. However, jointly training the projector and student allows both to reduce the loss. In upward projections, the projector determines both the activation residual and how it reaches the student, potentially distributing alignment across the student and a projector discarded after training. We call this effect dilution.

RAPID trades this flexibility for a strict separation of responsibilities by removing the competing adaptation path and placing representation learning entirely on the student, eliminating the possibility of dilution. Appendix § 16 provides the corresponding gradient analysis.

Representational Alignment

We evaluate representational alignment between students and teachers using centered kernel alignment (CKA) and representational similarity analysis (RSA), the latter with an inner metric of cosine

caption: Performance comparison between learned projection and RAPID on SmolTalk, distilling All activations.

similarity and an outer metric of Spearman correlation. CKA measures the overall alignment between representations and RSA measures the similarity of pairwise relationships. We analyzed the students trained on SmolTalk in the residual-only distillation scheme from § 6. Final token activations were collected from students and teachers across 3 seeds $\{123, 45, 67\}$ for 500 samples of the train set of GSM8K Cobbe et al. [9] with a 1024 sequence length. Tables and figures for these results can be found in Appendix § 17. Global refers to an analysis of the full set of teacher layers vs the full set of student layers, and Mapped only uses the set of teacher layers in our distillation mapping $\mathcal{H}$. $P_{25}$ is the 25th percentile value.

RAPID produces substantially stronger RSA values for Gemma 3 across all global and mapped metrics, while the learned projector performs slightly better on most CKA metrics. For Qwen 2.5, the learned projector consistently—but only marginally—outperforms RAPID across both CKA and RSA. Overall, the small differences suggest that RAPID achieves representational alignment comparable to learned projection without requiring training additional projection parameters. Interestingly, the Qwen 2.5 values are significantly higher than the Gemma 3 values, even though the entire Gemma 3 family was pretrained with distillation Team et al. [41]. We leave an explanation of this to future work.

Distilling Additional Activations

While some existing internal distillation methods distill only the residual stream activations Sun et al. [38]; Liang et al. [27], other methods such as TinyBERT Jiao et al. [22] distill an expanded set of activations. To demonstrate that RAPID generalizes to an expanded activation set, we include an analysis which distills attention outputs, MLP outputs, and embeddings in addition to the residual activations on our SmolTalk dataset. We refer to this expanded set of activations as All activations.

tbl:3 contains the results for distilling All activations for the SmolTalk dataset. The Gemma 3 distillation pair achieves very similar results as in the residual-only distillation scheme in Table 1, however there is a small performance degradation for Qwen 2.5, with RAPID performing 2.2 percentage points worse when distilling All activations. These Qwen 2.5 results also exhibit the largest performance gap of all of our experiments, at +1.67 percentage points for the learned projection. However, this small performance lift for learned projectors over RAPID comes at the cost of over 20% worse efficiency.

Limitations

Limited compute⁶ constrained model size, sequence length, and teacher finetuning, preventing full distillation and evaluation on R1-style long-context reasoning. Moreover, we evaluate RAPID only during reasoning-focused supervised finetuning and do not assess its effectiveness in pretraining, mid-training, general instruction tuning, tool-using, or as an initialization for subsequent reinforcement learning and other post-training stages. Future work should evaluate RAPID across these training phases and at larger scales, as well as optimize the dynamic-weight bounds $\omega_{\min}, \omega_{\max}$ in Equation (12).

Conclusion

We introduced RAPID, a novel methodology for internal distillation which achieves comparable performance to existing methods with learned projectors while being 6.3%–21.2% more efficient in our experiments. RAPID generalizes across models families, training datasets, and activation sets. RAPID places the responsibility of representation learning on the student, resulting in comparable or better representational similarity with the teacher.

AI Use Statement

In this work, we used generative AI tools for: formulating mathematical claims, providing critical ingredients for proving mathematical claims, and assistance in writing proofs in Appendices D and E; designing and providing feedback on the objective formulation; and implementing methods in our experimental harness. We have not used generative AI tools for helping develop theoretical models or conceptual frameworks, proposing or refining hypotheses, supporting qualitative and thematic data analysis, or interpreting results.and generating synthetic datasets, assisting with translation, and cleaning and reformatting data are not applicable to this work.

Additionally, we used generative AI tools for suggesting experimental parameters, creating and editing code, drafting Sections 4.3 and 4.4, summarizing existing literature, identifying research gaps, brainstorming, sourcing/searching for information, editing Sections 1, 2, and 3 for brevity, identifying relevant literature, and formatting references.

We have reviewed all AI-assisted work. We checked LLM-generated research ideas for potential plagiarism through a manual literature survey. LLM-generated code was verified and tested for correctness by the first author. LLM-generated proofs were reviewed and modified for correctness and clarity. We take responsibility for the final content of this work, including text, claims or artifacts produced with the aid of generative AI.

Reproducibility Statement

We have taken a number of steps to ensure high levels of reproducibility in our work. All experiments were run with random seeds (whose values we provided) and the PyTorch reproducible backend for NVIDIA GPUs. Algorithm 1 contains all details needed to reproduce the randomized projectors, and Sections 4.3 and 4.4 show detailed formulations of our multiple objectives, dynamic weight updates, and parameters used.

Johnson-Lindenstrauss Lemma

We provide a restatement of the JL lemma here for convenience.

Achlioptas Projector Initialization Algorithms

Algorithm 1 Dense Binary Coin Random Projection

Require: $X \in\mathbb{R}^{n \times d}$, target dimension $k$ Ensure: $Z \in\mathbb{R}^{n \times k}$

1. Initialize $R \in\mathbb{R}^{d \times k}$

2. for $i = 1,\ldots,d$ do

   3. for $j = 1,\ldots,k$ do

      4. Sample $u \sim\operatorname{Uniform}(0,1)$

      5. if $u < 1/2$ then

         6. $R_{ij} \leftarrow+1$

      6. else

         8. $R_{ij} \leftarrow-1$

      7. end if

   4. end for

3. end for

4. $Z \leftarrow\frac{1}{\sqrt{k}}XR$

5. return $Z$

Algorithm 2 Sparse Ternary Random Projection

Require: $X \in\mathbb{R}^{n \times d}$, target dimension $k$ Ensure: $Z \in\mathbb{R}^{n \times k}$

1. Initialize $Q \in\{-1,0,+1\}^{d \times k}$

2. for $i = 1,\ldots,d$ do

   3. for $j = 1,\ldots,k$ do

      4. Sample $u \sim\operatorname{Uniform}(0,1)$

      5. if $u < 1/6$ then

         6. $Q_{ij} \leftarrow+1$

      6. else if $u < 5/6$ then

         8. $Q_{ij} \leftarrow0$

      7. else

         10. $Q_{ij} \leftarrow-1$

      8. end if

   4. end for

3. end for

4. $Z \leftarrow\sqrt{\frac{3}{k}}XQ$

5. return $Z$

Contextualized JL Lemma Guarantees

We provide a concrete example of the theoretical guarantees of the JL lemma using Gemma 3 270M as a student model. The representation space is $k = 640$ dimensions, and we train for

For $n = 40{,}000$, $k = 640$, and a desired confidence of $99\%$, set the failure probability to $\delta= 0.01$. Since Achlioptas [1] gives confidence $1-n^{-\beta}$, solve

to obtain

Substituting this value into the Achlioptas bound gives

Equivalently,

Solving numerically gives

Therefore, with probability at least 99%, all pairwise squared distances satisfy

In terms of ordinary Euclidean distances,

Thus, at 99% confidence, projecting 40,000 points into 640 dimensions guarantees that every distance is between approximately 71.5% and 122.0% of its original value for every one of the $\binom{40{,}000}{2}=799{,}980{,}000$ sequence pairs.

Stop-Gradient Optimization Mechanics in Dynamic Objective

In Section 4.4, we note that all operations computing the dynamic objective weights $w_{i,t}$ from raw losses $\mathcal{L}_{i,t}$ and log-variance parameters $s_i$ are evaluated using a stop-gradient operator ($\operatorname{sg}[\cdot]$). This section provides a formal analysis of why gradient detachment is mathematically necessary and proves the optimization failure modes that arise when $w_{i,t}$ remains fully differentiable with respect to $\mathbf{s}=[s_1,\ldots,s_K]^\top$.

Fully Differentiable Formulation without Stop-Gradient

Consider a simplified variant of the normalized uncertainty weighting objective where intermediate weights $w_i(\mathbf{s})$ maintain active autograd graph connections:

Evaluating the partial derivative of $\mathcal{J}(\mathbf{s})$ with respect to an individual log-variance $s_i$ yields:

Using the standard Jacobian properties of softmax-style normalizations:

Substituting these derivatives back into the total gradient gives:

where $\overline{\mathcal{L}}_w = \sum_{j=1}^{K} w_j\mathcal{L}_j$ represents the current weighted-average loss across all active objectives.

Failure Modes of the Un-Detached System

(21) exposes two severe optimization vulnerabilities when gradient paths through $w_i(\mathbf{s})$ are preserved:

1. Loss-Evasion Shortcut. If objective $i$ experiences a temporary loss spike such that $\mathcal{L}_i > \overline{\mathcal{L}}_w$, the term $1 + w_i(\overline{\mathcal{L}}_w - \mathcal{L}_i)$ can become negative. Under standard gradient descent ($\mathbf{s} \leftarrow\mathbf{s} - \gamma\nabla_{\mathbf{s}}\mathcal{J}$), a negative gradient forces $s_i$ to grow ($s_i \rightarrow+\infty$). The optimizer discovers that it can minimize the global objective $\mathcal{J}$ faster by inflating $s_i$ (suppressing $w_i$) rather than updating the network backbone parameters to reduce the underlying task error $\mathcal{L}_i$.

2. Scale-Invariance Collapse to Negative Infinity. More critically, normalized uncertainty weights are invariant under uniform translations of the parameter vector $\mathbf{s}$. For any scalar constant $c \in\mathbb{R}$:

However, the linear regularization penalty $\sum s_j$ is strictly shift-variant:

Taking the limit as $c \rightarrow-\infty$:

Without a stop-gradient operator, the optimizer can exploit this scale symmetry by driving all $s_i \rightarrow-\infty$ at a uniform rate. The relative task weights $w_i$ remain fixed, but the total loss vanishes to negative infinity, completely collapsing training stability.

Guaranteed Stability via Stop-Gradient Execution

In our implementation, we enforce stop-gradient execution by calculating $w_{i,t}$ inside a nondifferentiable step routine. In the computational graph for $\mathcal{L}_{\mathrm{total}}$, the weights $w_{i,t}$ are treated as constant coefficients:

Because $\frac{\partial\operatorname{sg}[w_{i,t}]}{\partial s_j}=0$, the partial derivative with respect to any active log-variance simplifies strictly to:

This structural disconnection guarantees that:

1. The cross-derivatives $\frac{\partial w_j}{\partial s_i}$ are zeroed out, eliminating the scale-invariance drift $c \rightarrow-\infty$.

2. Task loss gradients $\nabla_{\theta}\mathcal{L}_{i,t}$ update the model parameters $\theta$, while $s_i$ is updated exclusively by its regularizer, isolating it as a stable baseline estimator.

Gradient Analysis of Projector–Student Responsibilities

Here we derive the optimization behavior underlying the projector–student separation discussed in Section 7.1. Let

denote student and teacher activations for a single sample. For student parameters $\theta\in\mathbb{R}^{p}$, define

Fixed Downward Projection

For the fixed downward-projection loss in § 7.1, differentiation with respect to the student activation gives

Applying the chain rule yields

Since

the projected teacher activation is stationary for a fixed input. The projector influences learning only through the target $\mathbf{H}_{\mathrm{fix}}\mathbf{a}_t$; it neither adapts to reduce the loss nor transforms the residual during its pullback through the student.

Learned Upward Projection

For the learned upward-projection loss, define the residual

The projector gradient is

Thus, the projector can directly reduce the matching residual. The student-activation gradient is

which gives

The learned projector therefore participates in optimization in two ways: it determines the residual and transforms that residual through $\mathbf{H}_{\mathrm{learn}}^{\mathsf{T}}$ before it reaches the student Jacobian. Both effects change as the projector is updated. Consequently, low matching loss does not imply that the student alone has reproduced the teacher representation; part of the alignment may reside in the learned projector. Fixing the downward projection removes this ambiguity and ensures that only the student adapts during distillation.

Representational Alignment Results

Model Family	Method	Global				Mapped
		Mean ± Std.	Median	$P_{25}$		Mean ± Std.	Median	$P_{25}$
Gemma 3	Learned Projector	0.3318 ± 0.1819	0.2985	0.2017	0.3791 ± 0.1298	0.3630	0.2914
	RAPID	0.3442 ± 0.1875	0.3231	0.2069	0.3739 ± 0.1287	0.3648	0.2868
Qwen 2.5	Learned Projector	0.7860 ± 0.1651	0.8410	0.7061	0.8963 ± 0.0883	0.9295	0.8917
	RAPID	0.7810 ± 0.1693	0.8395	0.6944	0.8935 ± 0.0921	0.9275	0.8874

Model Family	Method	Global			Mapped
		Mean ± Std.	Median	$P_{25}$	Mean ± Std.	Median	$P_{25}$
Gemma 3	Learned Projector	$0.3815 \pm0.1856$	$0.3440$	$0.2551$	$0.4745 \pm0.1647$	$0.4212$	$0.3590$
	RAPID	$0.4239 \pm0.1729$	$0.4215$	$0.3171$	$0.5063 \pm0.1393$	$0.4603$	$0.4321$
Qwen 2.5	Learned Projector	$0.6926 \pm0.1903$	$0.7400$	$0.5839$	$0.8515 \pm0.0880$	$0.8757$	$0.8439$
	RAPID	$0.6866 \pm0.1939$	$0.7363$	$0.5771$	$0.8454 \pm0.0957$	$0.8715$	$0.8381$

Training Hyperparameters

Hyperparameter	Value
Learning rate	$3 \times10^{-5}$
Weight decay	$1 \times10^{-4}$
Maximum gradient norm	2.0
Distillation temperature	1.1