Interpreting Language Model Hidden States at Scale

# Interpreting Language Model Hidden States at Scale
Source: [https://arxiv.org/html/2608.10260](https://arxiv.org/html/2608.10260)
Jordan Pettyjohn1, Mansi Sakarvadia1, Nathaniel Hudson2,3, Daniel McKenzie4, Kyle Chard1,3, Ian Foster1,3

###### Abstract

Lens methods interpret large language models \(LLMs\) by mapping intermediate activations to the output vocabulary, revealing how next\-token predictions develop through the network\. Trained lenses remain expensive: affine\-translator parameters grow quadratically with model width, while exact, full\-vocabulary Kullback–Leibler \(KL\) training dominates memory\. Consequently, prior trained lenses have been applied to models of at most 20B parameters and remain tied to particular component types\. We present OmniLens, which applies a single lens family to any model\-width activation, whether residual stream, attention, or MLP, and combines two independent scaling techniques\. First, low\-rank translators make per\-lens parameter growth linear in model width and reduce trainable parameters by up to 98\.4%\. Second, Subset\-KL materializes only selected vocabulary logits: its Top\-kkmode cuts peak training memory by up to 70%, while its importance\-sampled variant retains unbiased stochastic gradients for the full KL\. These savings enable a dense ensemble of 482 lenses for LLaMA\-3\.3\-70B, providing 6×\\timesthe coverage of a residual\-stream design at the same depth\. Model\-wide coverage then reveals what single\-component lenses cannot: the components where a behavior is most visible need not be those where intervention is most effective, and the most effective interventions lie outside the attention heads examined by prior lens studies\. Across three case studies \(prompt\-injection detection, multi\-hop memory injection, and toxicity localization\), OmniLens reproduces key published results at substantially lower cost\.

![Refer to caption](https://arxiv.org/html/2608.10260v1/x1.png)Figure 1:Full\-rank lens parameter grows as𝒪\\mathcal\{O\}\(Layer×d2\\times d^\{2\}\) exceeding 200B for LLaMA\-3\-405B\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x2.png)Figure 2:OmniLens at a glance\. Hooks are placed at arbitrary user\-defined points in the model; each lens applies the low\-rank translator of Eq\. \([3](https://arxiv.org/html/2608.10260#S3.E3)\) and is trained against the model’s final distribution under the Subset\-KL objectives of Section[4](https://arxiv.org/html/2608.10260#S4)\.## 1Introduction

Understanding how a language model forms its predictions, and where behaviors emerge within its computation, is central to interpreting and controlling it\(Orgadet al\.[2026](https://arxiv.org/html/2608.10260#bib.bib75); Shapira and others[2026](https://arxiv.org/html/2608.10260#bib.bib76)\)\. A*lens*provides a direct view by decoding intermediate activations into vocabulary distributions, revealing how next\-token predictions evolve through the network and where behaviors can be detected or influenced\. Lenses can decode many kinds of intermediate activations; we refer to the model component providing input to a lens as itshookpoint\. For example, theTuned Lensreads the residual stream\(Belroseet al\.[2023](https://arxiv.org/html/2608.10260#bib.bib6)\), whereas theAttention Lensreads individual attention heads\(Sakarvadiaet al\.[2023](https://arxiv.org/html/2608.10260#bib.bib7)\)\.

Unlike concept\-specific classifier probes, which require labeled examples for each target\(Alain and Bengio[2018](https://arxiv.org/html/2608.10260#bib.bib35); Hewitt and Manning[2019](https://arxiv.org/html/2608.10260#bib.bib84); Belinkovet al\.[2017](https://arxiv.org/html/2608.10260#bib.bib83)\), lenses can be trained in a self\-supervised manner: the model’s own final distribution supplies the target, so no curated labels are needed\. A single trained lens can therefore support many downstream analyses\. This capability is most valuable when applied densely across a model, that is, with multiple hookpoints at every layer\. However, training such an ensemble of lenses has remained prohibitively expensive\.

Specifically, trained lenses face three challenges when applied at scale\.*Parameters\.*A full\-rank translator costs𝒪​\(d2\)\\mathcal\{O\}\(d^\{2\}\)per hookpoint, whereddis the model’s hidden dimension, so dense coverage of a large model can rival the model itself \(Fig\.[1](https://arxiv.org/html/2608.10260#S0.F1)\)\.*Memory\.*Lens training uses the KL divergence between the model’s own final distribution and the lens’ distribution as the training objective, which materializes two vocabulary\-sized distributions per token\. This quickly exhausts available VRAM\.*Specialization\.*Existing lenses each read one hookpoint type, and each new hookpoint requires a new lens family\. Together, these challenges limit the application of trained lenses to small models with sparse coverage\.

We address all three challenges with*OmniLens*,unlocking dense lens coverage for \(near\-\)frontier scale models\. OmniLens is a hookpoint\-agnostic framework that applies one lens family to residual, attention, and MLP activations\. OmniLens makes dense coverage tractable through two parallel techniques\. First, low\-rank translators reduce the per\-hookpoint parameter count from𝒪​\(d2\)\\mathcal\{O\}\(d^\{2\}\)to𝒪​\(r​d\)\\mathcal\{O\}\(rd\)whererris the target rank\. Second, a flexible family of approximations to the KL divergence—which we call Subset\-KL objectives—materialize only a small subset of the lens’ distribution per token, greatly reducing memory requirements\.

We combine these contributions in an open\-source framework111Code: OmniLens \(training framework\)[https://github\.com/pettyjohnjn/OmniLens](https://github.com/pettyjohnjn/OmniLens); Hookbox \(hookpoint instrumentation\)[https://github\.com/pettyjohnjn/hookbox](https://github.com/pettyjohnjn/hookbox); IndexedLogits \(fused CUDA kernel\)[https://github\.com/pettyjohnjn/indexed˙logits](https://github.com/pettyjohnjn/indexed_logits); SubsetKL \(objectives and estimator\)[https://github\.com/pettyjohnjn/subset\-kl](https://github.com/pettyjohnjn/subset-kl)and reproduce key metrics from prior work at substantially lower cost\(Belroseet al\.[2023](https://arxiv.org/html/2608.10260#bib.bib6); Sakarvadiaet al\.[2024](https://arxiv.org/html/2608.10260#bib.bib32); Pettyjohn[2025](https://arxiv.org/html/2608.10260#bib.bib109)\)\. OmniLens covers six times as many lenses as the tuned lens baseline while costing less, with 90\.5% fewer trainable parameters on LLaMA\-3\-70B and up to 70% lower peak training memory on GPT\-2\. Across three interpretability case studies it matches the application efficacy of existing lens frameworks, and model\-wide coverage additionally reveals what single\-component lenses cannot: the hookpoints where a behavior is most visible are not necessarily those where intervention is most effective\. To test the limits of the approach, we run eight optimization steps on LLaMA\-3\.1\-405B, to our knowledge the first measured demonstration of trained\-lens optimization at this scale on existing hardware, establishing that the training path is executable at frontier scale\.

## 2Background and Related Work

We study pretrained autoregressive Transformer language models\(Vaswaniet al\.[2017](https://arxiv.org/html/2608.10260#bib.bib1)\)\. Letxxbe a length\-TTtoken sequence from vocabularyVV, and letP\(⋅\|x\)P\(\\cdot\|x\)denote the model’s next\-token distribution at a given position\. The model has hidden dimensionddandLLlayers, with hookpoints exposing residual\-stream states, attention and MLP outputs, or individual head outputs, depending on the desired resolution\. We writeHℓ,u∈ℝT×duH\_\{\\ell,u\}\\in\\mathbb\{R\}^\{T\\times d\_\{u\}\}for the output of componentuuat layerℓ\\ellandhℓ,u∈ℝduh\_\{\\ell,u\}\\in\\mathbb\{R\}^\{d\_\{u\}\}for a single position’s activation; lenses apply positionwise, and we suppress the position index\. We organize prior work by the three choices that determine a lens method’s scalability: where it reads, how it translates the activation, and how its training objective is computed\.

##### Lens formulation and component specialization\.

A lens is an auxiliary decoder mapping an intermediate activation to a distribution over the model’s vocabulary\. We restrict attention to linear lenses that reuse the model’s frozen final normalizationη\\eta\(LayerNorm or RMSNorm, applied exactly as in the model’s readout\) and unembeddingWU∈ℝ\|V\|×dW\_\{U\}\\in\\mathbb\{R\}^\{\|V\|\\times d\}:

Qℓ,u\(⋅\|x\)=softmax\(WUη\(ℒℓ,uhℓ,u\+bℓ,u\)\)Q\_\{\\ell,u\}\(\\cdot\|x\)=\\operatorname\{softmax\}\\left\(W\_\{U\}\\,\\eta\\\!\\left\(\\mathcal\{L\}\_\{\\ell,u\}\\,h\_\{\\ell,u\}\+b\_\{\\ell,u\}\\right\)\\right\)\(1\)whereℒℓ,u∈ℝd×du\\mathcal\{L\}\_\{\\ell,u\}\\in\\mathbb\{R\}^\{d\\times d\_\{u\}\}andbℓ,u∈ℝdb\_\{\\ell,u\}\\in\\mathbb\{R\}^\{d\}are learned\. The logit lens\(Nostalgebraist[2020](https://arxiv.org/html/2608.10260#bib.bib5)\)applies the model’s frozen readout directly to model\-width residual\-stream states, corresponding toℒℓ,u=I\\mathcal\{L\}\_\{\\ell,u\}=Iandbℓ,u=0b\_\{\\ell,u\}=0\. Because this direct readout can be a poor proxy for the final prediction, subsequent methods learn affine translators so thatQℓ,u\(⋅\|x\)Q\_\{\\ell,u\}\(\\cdot\|x\)approximatesP\(⋅\|x\)P\(\\cdot\|x\)\(Belroseet al\.[2023](https://arxiv.org/html/2608.10260#bib.bib6); Dinet al\.[2024](https://arxiv.org/html/2608.10260#bib.bib62); Palet al\.[2023](https://arxiv.org/html/2608.10260#bib.bib59)\)\. These methods target residual streams, while Attention Lens\(Sakarvadiaet al\.[2023](https://arxiv.org/html/2608.10260#bib.bib7)\)learns decoders for individual attention heads\. Each construction commits to a component family, so reading a different component requires another lens design; this coupling between translator and component is the specialization bottleneck OmniLens addresses\. The Backward Lens\(Katzet al\.[2024](https://arxiv.org/html/2608.10260#bib.bib120)\)projects gradients rather than activations into vocabulary space, proving that such projections admit low\-rank structure\.

##### Parameter\-efficient translators\.

A dense affine translator hasdu​d\+dd\_\{u\}d\+dlearned parameters, and fitting one per component across all layers quickly grows impractical: at the\(6​L\+2\)\(6L\{\+\}2\)hookpoint density we study, the lens set approaches half the parameter count of the base model \(Fig\.[1](https://arxiv.org/html/2608.10260#S0.F1)\), and finer constructions such as the per\-head decoders of Attention Lens can exceed the model of study outright\.Low\-Rank Adaptation\(LoRA\)\(Huet al\.[2021](https://arxiv.org/html/2608.10260#bib.bib27)\)—originally used for parameter\-efficient fine\-tuning—freezes a model’s weights and learns a low\-rank update to each, cutting trainable parameters by orders of magnitude\.

Low\-rank parameterizations have been combined with lenses, yet existing implementations retain the component specialization of their full\-rank predecessors\. LoRA Lens\(Pettyjohn[2025](https://arxiv.org/html/2608.10260#bib.bib109)\)factorizes per\-head Attention Lens decoders as rank\-rrupdates to the frozen unembedding on models up to 8B, while concurrent work\(Trimignoet al\.[2026](https://arxiv.org/html/2608.10260#bib.bib117)\)trains low\-rank residual\-stream lenses on models up to 32B\. Each applies low rank within a single component family, and neither supplies one translator architecture spanning residual, attention, and MLP components or addresses the vocabulary\-side memory cost that limits dense coverage at scale\. A full taxonomy of lenses is in Appendix[A](https://arxiv.org/html/2608.10260#A1)\.

##### Memory\-efficient distillation\.

Most lens frameworks measure the discrepancy betweenP\(⋅\|x\)P\(\\cdot\|x\)andQℓ,u\(⋅\|x\)Q\_\{\\ell,u\}\(\\cdot\|x\)using the token\-levelKullback\-Leibler\(KL\) divergence,

DKL​\(P∥Qℓ,u\)\\displaystyle D\_\{\\mathrm\{KL\}\}\(P\\\|Q\_\{\\ell,u\}\)=∑v∈VP​\(v\|x\)​log⁡P​\(v\|x\)Qℓ,u​\(v\|x\)\\displaystyle=\\sum\_\{v\\in V\}P\(v\|x\)\\log\\frac\{P\(v\|x\)\}\{Q\_\{\\ell,u\}\(v\|x\)\}\(2\)=𝔼v∼P\(⋅\|x\)​\[log⁡P​\(v\|x\)Qℓ,u​\(v\|x\)\],\\displaystyle=\\mathbb\{E\}\_\{v\\sim P\(\\cdot\|x\)\}\\left\[\\log\\frac\{P\(v\|x\)\}\{Q\_\{\\ell,u\}\(v\|x\)\}\\right\],where we callP\(⋅\|x\)P\(\\cdot\|x\)the*teacher*andQℓ,u\(⋅\|x\)Q\_\{\\ell,u\}\(\\cdot\|x\)the*student*\. Materializing allT​\|V\|T\|V\|student logits for a single inputxxis costly, and numerous approximation schemes avoid it\.

We use KL as a distillation loss\(Hintonet al\.[2015](https://arxiv.org/html/2608.10260#bib.bib92); Sanhet al\.[2020](https://arxiv.org/html/2608.10260#bib.bib93)\)\. Because the trainable student is the second argument ofDKL\(⋅∥⋅\)D\_\{\\mathrm\{KL\}\}\(\\cdot\\,\\\|\\,\\cdot\)and the expectation runs over the fixed teacher, drawing tokens fromPPgives a well\-behaved and unbiased Monte Carlo estimator\. When tokens are drawn from another proposalRR, importance sampling weights each sampled contribution byP​\(v\|x\)/R​\(v\|x\)P\(v\|x\)/R\(v\|x\)so that its expectation still recovers the original KL\(Aminiet al\.[2025](https://arxiv.org/html/2608.10260#bib.bib72)\)\. Deterministic Top\-kktruncation instead scores only thekkmost probable teacher tokens, a memory\-cheap but biased reduction\(Shaoet al\.[2024](https://arxiv.org/html/2608.10260#bib.bib70)\)\. Appendix[D](https://arxiv.org/html/2608.10260#A4)distinguishes this setting from KL regularization under a trainable sampling distribution\(Tang and Munos[2025](https://arxiv.org/html/2608.10260#bib.bib108)\)\.

The closest sparse\-distillation baseline isRandom Sampling Knowledge Distillation\(RS\-KD\)\(Anshumannet al\.[2025](https://arxiv.org/html/2608.10260#bib.bib115)\), compared against in Section[4\.2](https://arxiv.org/html/2608.10260#S4.SS2); our sampled tail builds on importance sampling for large output spaces\(Katharopoulos and Fleuret[2019](https://arxiv.org/html/2608.10260#bib.bib94); Blanc and Rendle[2018](https://arxiv.org/html/2608.10260#bib.bib95)\)\. These motivate*Subset\-KL*\(Section[4](https://arxiv.org/html/2608.10260#S4)\): biased Top\-k truncation and an exact\-head, importance\-sampled\-tail variant with unbiased gradients\.

##### Training\-free and complementary readouts\.

The Jacobian lens\(Gurneeet al\.[2026](https://arxiv.org/html/2608.10260#bib.bib119)\)avoids translator training by decoding average local output sensitivity at a hookpoint\. It asks what a state locally*represents*under a first\-order perturbation, whereas a trained predictive lens is optimized to recover the model’s eventual output distribution; its authors observe that tuned lenses can therefore “skip ahead” past intermediate representations\. PatchScopes\(Ghandehariounet al\.[2024](https://arxiv.org/html/2608.10260#bib.bib121)\)is likewise training\-free, patching hidden states into explanatory prompts so that the model’s own generation serves as the readout\. These views are complementary, although each remains a per\-hookpoint readout\. More broadly, lenses are observational decoders and do not by themselves establish causal mechanisms\. They complement causal techniques such as circuit discovery and activation patching\(Elhageet al\.[2021](https://arxiv.org/html/2608.10260#bib.bib38); Wanget al\.[2022](https://arxiv.org/html/2608.10260#bib.bib40); Conmyet al\.[2023](https://arxiv.org/html/2608.10260#bib.bib8); Menget al\.[2022](https://arxiv.org/html/2608.10260#bib.bib82)\); observational attention analyses\(Clarket al\.[2019b](https://arxiv.org/html/2608.10260#bib.bib79); Voitaet al\.[2019](https://arxiv.org/html/2608.10260#bib.bib80); Vig and Belinkov[2019](https://arxiv.org/html/2608.10260#bib.bib81)\); and feature\-oriented methods including feature visualization, sparse autoencoders that decompose activations into features in superposition, and stochastic parameter decomposition targeting weights rather than activations\(Olahet al\.[2017](https://arxiv.org/html/2608.10260#bib.bib77); Cammarataet al\.[2021](https://arxiv.org/html/2608.10260#bib.bib78); Cunninghamet al\.[2023](https://arxiv.org/html/2608.10260#bib.bib87); Templetonet al\.[2024](https://arxiv.org/html/2608.10260#bib.bib88); Sharkeyet al\.[2022](https://arxiv.org/html/2608.10260#bib.bib89); Bushnaqet al\.[2025](https://arxiv.org/html/2608.10260#bib.bib104)\)\. At frontier scale, circuit tracing\(Anthropic[2025](https://arxiv.org/html/2608.10260#bib.bib122)\)combines cross\-layer transcoders with attribution graphs to map computational structure in production models\. Classifier\-based probes\(Ettingeret al\.[2016](https://arxiv.org/html/2608.10260#bib.bib111); Conneauet al\.[2018](https://arxiv.org/html/2608.10260#bib.bib36); Ivanitskiyet al\.[2024](https://arxiv.org/html/2608.10260#bib.bib103); Kramáret al\.[2026](https://arxiv.org/html/2608.10260#bib.bib102)\)likewise read internal states but test individual concepts using curated labels, whereas lenses target the complete output distribution without task\-specific labels\.

##### Synthesis\.

Prior work addresses three bottlenecks separately: learned lenses improve fidelity but specialize to one component family, low\-rank lenses reduce translator parameters within those families, and sparse distillation reduces vocabulary\-side cost without changing where lenses attach\. OmniLens combines all three\.

## 3OmniLens: Low\-Rank Parameterization

OmniLens combines low\-rank translators, hookpoint\-agnostic attachment, and the Subset\-KL objectives of Section[4](https://arxiv.org/html/2608.10260#S4); this section introduces the first two\. The same translator parameterization attaches to model\-width residual, attention, and MLP activations, generalizing prior hookpoint\-specific and low\-rank lens constructions \(Section[2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px2)\)\. Our implementation supports distributed training and gradient checkpointing; details are in Appendix[B\.1](https://arxiv.org/html/2608.10260#A2.SS1)\.

##### Low\-rank translator\.

For an activationhℓ,u∈ℝdh\_\{\\ell,u\}\\in\\mathbb\{R\}^\{d\}, OmniLens learns a translator\(ℒℓ,u,bℓ,u\)\(\\mathcal\{L\}\_\{\\ell,u\},b\_\{\\ell,u\}\), where

ℒℓ,u=I\+αr​Bℓ,u​Aℓ,u,\\mathcal\{L\}\_\{\\ell,u\}=I\+\\frac\{\\alpha\}\{r\}B\_\{\\ell,u\}A\_\{\\ell,u\},\(3\)Aℓ,u∈ℝr×dA\_\{\\ell,u\}\\in\\mathbb\{R\}^\{r\\times d\}, andBℓ,u∈ℝd×rB\_\{\\ell,u\}\\in\\mathbb\{R\}^\{d\\times r\}, with fixed scaleα\\alpha, so thatℒℓ,u\\mathcal\{L\}\_\{\\ell,u\}is a rank\-at\-most\-rrupdate to the identity\.

Each translator contains2​d​r\+d2dr\+dlearned parameters, againstd2\+dd^\{2\}\+dfor a dense affine translator\. Ford≫rd\\gg rthe ratio of learned parameters per hookpoint is approximately2​r/d2r/d, so the savings compound as coverage grows: atdd= 4096 andrr= 64, six low\-rank translators together contain only about 19% as many parameters as a single dense translator\. Low rank can therefore widen coverage across component types while still reducing the total parameter count\.

##### Hookpoint\-agnostic attachment\.

Existing trained lenses couple their decoder to a particular component type\. OmniLens instead treats hookpoint type as a configuration option: each supported hookpoint supplies a model\-width activation, after which the same translator parameterization, the same frozen normalization and unembedding, and the same training objective apply unchanged\. We cover six residual, attention, and MLP hookpoint types per layer in our experiments \(Fig\.[2](https://arxiv.org/html/2608.10260#S0.F2)\), so component types can be compared without designing and training a separate lens family for each\. Every hookpoint studied here has widthdu=dd\_\{u\}=d\.

We initializeAℓ,uA\_\{\\ell,u\}with Xavier uniform\(Glorot and Bengio[2010](https://arxiv.org/html/2608.10260#bib.bib110)\)andBℓ,uB\_\{\\ell,u\}andbℓ,ub\_\{\\ell,u\}to zero, so each translator begins as the identity, withα/r\\alpha/rcontrolling the update scale\. The translated state is decoded through the model’s frozen final normalizationη\\etaand unembeddingWUW\_\{U\}, exactly as in Eq\. \([1](https://arxiv.org/html/2608.10260#S2.E1)\)\. A translator is therefore only required to be accurate on the activations the model actually produces, rather than on all ofℝd\\mathbb\{R\}^\{d\}, and only up to differences that the frozen readout preserves in the output distribution\. A full\-rankd×dd\\times dupdate supplies more capacity than these two restrictions demand, motivating the low\-rank parameterization evaluated below\. Fig\.[1](https://arxiv.org/html/2608.10260#S0.F1)presents the resulting parameter scaling across model sizes; Appendix[G](https://arxiv.org/html/2608.10260#A7)reports the per\-model parameter and memory counts underlying these figures\.

##### Rank ablation\.

We validate the low\-rank parameterization on GPT\-2 Small by sweepingr∈\{1,4,8,16,32,64,128,256,384\}r\\in\\\{1,4,8,16,32,64,128,256,384\\\}against a full\-rank tuned\-lens baseline after 1000 optimization steps\. We measure KL divergence to the teacher, and top\-1 agreement, Pearsonρ\\rho, and Kendallτ\\taurelative to the full\-rank lens, following the evaluation protocol ofBelroseet al\.\([2023](https://arxiv.org/html/2608.10260#bib.bib6)\)\(Fig\.[3](https://arxiv.org/html/2608.10260#S3.F3); representative points in Table[1](https://arxiv.org/html/2608.10260#S3.T1)\)\. KL measures fidelity to the teacher distribution, while the agreement metrics test whether the low\-rank lens preserves the predictions and token rankings of the dense reference\. Appendix[B](https://arxiv.org/html/2608.10260#A2)gives the setup; Appendix[C\.1](https://arxiv.org/html/2608.10260#A3.SS1)reports the complete rank and layer breakdowns\.

![Refer to caption](https://arxiv.org/html/2608.10260v1/x3.png)Figure 3:Fidelity\-parameter tradeoff on GPT\-2 Small at step 1000\. Solid lines show the final layer, dashed lines the mean over 12 layers, and dotted horizontal lines represent the full\-rank tuned\-lens baseline in the KL panel\. The filled marker denotes the recommended default,r=64r=64; gains diminish at larger ranks\. Appendix Table[5](https://arxiv.org/html/2608.10260#A3.T5)reports the full sweep\.Table 1:Representative points from the rank ablation; the full sweep is Appendix Table[5](https://arxiv.org/html/2608.10260#A3.T5)\. Mean metrics average over 12 layers; KL is to the teacher, agreement metrics are vs\. the full\-rank baseline\.Fidelity improves rapidly throughr=64r=64and only modestly thereafter \(Fig\.[3](https://arxiv.org/html/2608.10260#S3.F3)\)\. Atr=64r=64, OmniLens uses 16\.7% as many translator weight parameters as the full\-rank baseline while achieving 88\.8% final\-layer top\-1 agreement \(Pearsonρ=0\.984\\rho=0\.984\) and meanρ=0\.941\\rho=0\.941; we therefore user=64r=64as the default\. Earlier layers are the hardest to approximate \(layer 0: 56\.7% top\-1 agreement atr=64r=64\), and their KL to the teacher remains highest even for the full\-rank lens\. So, analyses focused on early layers should preferr≥64r\\geq 64\. The parameter\-matched runr=384r\{=\}384slightly outperforms the full\-rank baseline on KL, but the difference lies within the baseline’s own seed variation \(final\-layer KL0\.0390\.039–0\.0490\.049across three seeds; Appendix[B](https://arxiv.org/html/2608.10260#A2)\); we read this as fidelity preserved, not as low rank being inherently superior\. The 8B comparisons in Section[5](https://arxiv.org/html/2608.10260#S5)support the same conclusion at scale\.

##### Low rank is a constraint, not a compression\.

The deviation from identity learned by the full\-rank baseline,ℒℓ,u−I\\mathcal\{L\}\_\{\\ell,u\}\-I, is not itself low rank: its best rank\-64 approximation retains only59%59\\%of its Frobenius energy on GPT\-2 Small, and its numerical rank remains nearddthroughout training \(Fig\.[4](https://arxiv.org/html/2608.10260#S3.F4)b\)\. Post\-hoc compression accordingly loses fidelity: truncating the full\-rank translator to rank 64 raises final\-layer KL by54%54\\%, while projecting it onto a random 64\-dimensional subspace raises KL by more than an order of magnitude\. By contrast, a translator trained directly at rank 64 achieves final\-layer KL within7%7\\%of the full\-rank lens \(Fig\.[4](https://arxiv.org/html/2608.10260#S3.F4)a\)\.Low\-rank training therefore finds a distinct solution of comparable predictive quality rather than merely compressing the dense solution\.Two matrices can differ substantially in Frobenius norm and still induce nearly identical output distributions, because they need only agree on the activations the model produces and only up to differences the frozen readout preserves\.

![Refer to caption](https://arxiv.org/html/2608.10260v1/x4.png)Figure 4:Low rank is a constraint, not a compression \(GPT\-2 Small, step 1000, with 131,072 held\-out Pile tokens\)\.*Top:*Final\-layer KL for a translator trained directly at rankrr, a full\-rank translator truncated post hoc to rankrr, and a randomrr\-dimensional projection of it; the dashed line marks the unmodified full\-rank lens\.*Bottom:*Frobenius energy of the full\-rank deviation captured at rankrr\.

## 4OmniLens: Subset\-KL Training

Low\-rank translators reduce lens parameters and optimizer state, but they do not narrow the lens output: after translation, every active lens must still score vocabulary items at every token position\. Computing the full KL divergence therefore materializes a vocabulary\-sized logit tensor for each lensQℓ,uQ\_\{\\ell,u\}\(Eq\. \([2](https://arxiv.org/html/2608.10260#S2.E2)\)\)\. This activation cost grows with batch size, context length, and vocabulary size, and can dominate memory even when the translator itself is lightweight\.

We call objectives that evaluate only selected vocabulary tokens*Subset\-KL*, and study two variants with different computational and statistical guarantees\. Top\-kkrestricts both distributions to the teacher’s most probable tokens and renormalizes within that set: this avoids the full\-vocabulary lens projection but changes the objective, and is therefore biased \(Section[4\.1](https://arxiv.org/html/2608.10260#S4.SS1)\)\.Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}evaluates the head exactly and importance\-samples the remaining vocabulary, retaining the full student normalization and giving unbiased stochastic gradients for the original KL \(Section[4\.2](https://arxiv.org/html/2608.10260#S4.SS2)\)\. Top\-kkthus prioritizes memory and throughput, whereasTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}prioritizes fidelity to the full\-KL objective\. We follow the evaluation protocol ofBelroseet al\.\([2023](https://arxiv.org/html/2608.10260#bib.bib6)\); Appendix[B](https://arxiv.org/html/2608.10260#A2)gives the complete experimental setup\.

### 4\.1Top\-kkTruncation

Letℋ\\mathcal\{H\}contain thekheadk\_\{\\mathrm\{head\}\}most probable teacher tokens\. Top\-kkrestricts both distributions toℋ\\mathcal\{H\}and renormalizes them:

DTop​\-​k​\(P∥Qℓ,u\)=∑v∈ℋPℋ​\(v\|x\)​log⁡Pℋ​\(v\|x\)Qℋ,ℓ,u​\(v\|x\),D\_\{\\mathrm\{Top\}\\text\{\-\}k\}\(P\\\|Q\_\{\\ell,u\}\)=\\sum\_\{v\\in\\mathcal\{H\}\}P\_\{\\mathcal\{H\}\}\(v\|x\)\\log\\frac\{P\_\{\\mathcal\{H\}\}\(v\|x\)\}\{Q\_\{\\mathcal\{H\},\\ell,u\}\(v\|x\)\},\(4\)wherePℋP\_\{\\mathcal\{H\}\}andQℋ,ℓ,uQ\_\{\\mathcal\{H\},\\ell,u\}denote the restricted, renormalized distributions\. Equivalently, Top\-kkasks the lens to reproduce the teacher’s relative preferences among itskheadk\_\{\\mathrm\{head\}\}most probable tokens while ignoring both distributions’ mass outside that set\. Thus it is not a cheaper implementation of the full KL: truncating the partition changes the training objective\.

That change is also Top\-kk’s principal systems advantage: becauseQℋ,ℓ,uQ\_\{\\mathcal\{H\},\\ell,u\}normalizes only overℋ\\mathcal\{H\}, the lens never projects into the complete vocabulary, reducing activation memory*and*projection compute, withkheadk\_\{\\mathrm\{head\}\}as a direct quality–cost control\. We evaluate this on GPT\-2 Small, where full\-KL training remains feasible and supplies a reference: on its residual hookset,khead=256k\_\{\\mathrm\{head\}\}=256reduces peak memory from 16\.3 to 4\.7 GB and increases throughput by 1\.59×\\times, while changing final\-layer KL from 0\.043 to 0\.054\.

### 4\.2Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}: Exact Head, Sampled Tail

To ameliorate the poor performance of Top\-kkon early layers, we introduceTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}, which keeps the head unnormalized and samples the tail\. LetPhead=∑v∈ℋP​\(v\|x\)P\_\{\\mathrm\{head\}\}=\\sum\_\{v\\in\\mathcal\{H\}\}P\(v\|x\), and drawt1,…,tktailt\_\{1\},\\ldots,t\_\{k\_\{\\mathrm\{tail\}\}\}independently from a proposalR\(⋅\|x\)R\(\\cdot\|x\)supported onV∖ℋV\\setminus\\mathcal\{H\}\. We compute

D^Top​\-​k\+IS​\(P∥Qℓ,u\)=\\displaystyle\\widehat\{D\}\_\{\\text\{$\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}$\}\}\(P\\\|Q\_\{\\ell,u\}\)=∑v∈ℋP​\(v\|x\)​log⁡P​\(v\|x\)Qℓ,u​\(v\|x\)\\displaystyle\\sum\_\{v\\in\\mathcal\{H\}\}P\(v\|x\)\\log\\frac\{P\(v\|x\)\}\{Q\_\{\\ell,u\}\(v\|x\)\}\(5\)\+1ktail​∑i=1ktailP​\(ti\|x\)R​\(ti\|x\)​log⁡P​\(ti\|x\)Qℓ,u​\(ti\|x\)\.\\displaystyle\+\\frac\{1\}\{k\_\{\\mathrm\{tail\}\}\}\\sum\_\{i=1\}^\{k\_\{\\mathrm\{tail\}\}\}\\frac\{P\(t\_\{i\}\|x\)\}\{R\(t\_\{i\}\|x\)\}\\log\\frac\{P\(t\_\{i\}\|x\)\}\{Q\_\{\\ell,u\}\(t\_\{i\}\|x\)\}\.The first line is the exact contribution of the headℋ\\mathcal\{H\}to the original, untruncated KL\. The second estimates the omitted tail sum from sampled tokens, each reweighted byP/RP/Rso that in expectation the sampled term reconstructs the complete tail contribution\. For any lens\-independent proposal positive on the tail, both the estimated objective and its stochastic gradients are therefore unbiased for the full KL\.

###### Theorem 1\.

Lett1,…,tktailt\_\{1\},\\ldots,t\_\{k\_\{\\mathrm\{tail\}\}\}be drawn i\.i\.d\. from any proposalR\(⋅\|x\)R\(\\cdot\|x\)supported onV∖ℋV\\setminus\\mathcal\{H\}withR​\(v\|x\)\>0R\(v\|x\)\>0whereverP​\(v\|x\)\>0P\(v\|x\)\>0\. TheTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}estimator Eq\. \([5](https://arxiv.org/html/2608.10260#S4.E5)\) is unbiased,

𝔼​\[D^Top​\-​k\+IS​\(P∥Qℓ,u\)\]=DKL​\(P∥Qℓ,u\),\\mathbb\{E\}\\left\[\\hat\{D\}\_\{\\mathrm\{\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}\}\}\(P\\\|Q\_\{\\ell,u\}\)\\right\]=D\_\{\\mathrm\{KL\}\}\(P\\\|Q\_\{\\ell,u\}\),and yields unbiased gradients as well:

𝔼​\[∇D^Top​\-​k\+IS\]=∇DKL​\(P∥Qℓ,u\)\.\\mathbb\{E\}\\left\[\\nabla\\hat\{D\}\_\{\\mathrm\{\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}\}\}\\right\]=\\nabla D\_\{\\mathrm\{KL\}\}\(P\\\|Q\_\{\\ell,u\}\)\.

A complete proof for Theorem[1](https://arxiv.org/html/2608.10260#Thmtheorem1)is presented in Appendix[E](https://arxiv.org/html/2608.10260#A5)\.

HereQℓ,uQ\_\{\\ell,u\}is the true student distribution rather than one renormalized over the subset\. Evaluating it requires the full\-vocabulary log\-partitionlog​∑vexp⁡\(zv\)\\log\\sum\_\{v\}\\exp\(z\_\{v\}\), which we compute exactly without materializing all\|V\|\|V\|logits: we project one chunk of the vocabulary at a time, accumulate its log\-sum\-exp into a running total, and discard the chunk’s logits before projecting the next\. Subsampling therefore applies only to which KL summands are evaluated, not to the normalizer\.

As a default, we sample from the teacher conditioned on the tail,R​\(v\|x\)=P​\(v\|x\)/\(1−Phead\)R\(v\|x\)=P\(v\|x\)/\(1\-P\_\{\\mathrm\{head\}\}\)forv∈V∖ℋv\\in V\\setminus\\mathcal\{H\}\. This draws tail tokens in proportion to the same probabilities that weight their KL terms and needs no separate proposal model\. Retaining the head deterministically removes the most heavily weighted terms from the estimator’s variance, which is whyTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}improves on pure teacher sampling\. Appendix[E](https://arxiv.org/html/2608.10260#A5)gives the sampling procedure and illustrates the estimator \(Fig\.[12](https://arxiv.org/html/2608.10260#A5.F12)\)\.

##### Evaluation\.

Table[2](https://arxiv.org/html/2608.10260#S4.T2)compares matched token budgets, and Fig\.[5](https://arxiv.org/html/2608.10260#S4.F5)gives the layerwise profiles\. Among the subset objectives, Top\-kkachieves lower final\-layer KL and peak memory at both budgets, whereasTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}achieves lower mean and early\-layer KL\. At a256\+256256\{\+\}256budget,Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}halves peak memory from16\.316\.3to8\.28\.2GB while processing50\.950\.9k tokens/s\. Its throughput remains near4848–5151k tokens/s across the evaluated budgets, whereas Top\-kkslows from96\.596\.5k tokens/s atk=256k=256to28\.928\.9k atk=1024k=1024\. Seed replications preserve the mean and early\-layer ordering, while showing that part of the final\-layer gap is seed variation\. Appendix[C](https://arxiv.org/html/2608.10260#A3)reports the complete sweep and per\-seed results\.

Settingkhead=0k\_\{\\mathrm\{head\}\}=0reduces Eq\. \([5](https://arxiv.org/html/2608.10260#S4.E5)\) to pure teacher sampling, the RS\-KD baseline \(Section[2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px3)\)\. At a total budget of512512it reaches final\-layer KL0\.1940\.194, against0\.0670\.067forTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}, and its training loss deteriorates as the sample budget grows\. Appendix[C](https://arxiv.org/html/2608.10260#A3)gives the full sweep and diagnostics, and Appendix[D](https://arxiv.org/html/2608.10260#A4)compares related estimators\.

##### Choosing an objective\.

Neither variant dominates\. Top\-kkis preferable under projection cost, throughput, or maximum context length constraints, particularly for analyses that emphasize later layers\.Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}is preferable when lenses are compared throughout the network, or when preserving the semantics of the full\-KL objective matters\. We therefore default toTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}for the trained lenses in the interpretability studies, whose deterministic head is what preserves late\-layer fidelity \(Fig\.[11](https://arxiv.org/html/2608.10260#A3.F11)\), and report Top\-kkas the lower\-memory, higher\-throughput alternative\.

Table 2:Matched\-budget Subset\-KL comparison on GPT\-2 Small \(r=64r=64, step1,0001\{,\}000, residual hookset\)\. Early KL averages layers 0–3; RS denotes teacher sampling \(khead=0k\_\{\\mathrm\{head\}\}=0\)\. Bold marks the best Top\-kkor Top\-kk\+IS value, and shaded rows share a token budget\. Complete results are in Table[6](https://arxiv.org/html/2608.10260#A3.T6)\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x5.png)Figure 5:Layerwise KL on GPT\-2 Small \(r=64r=64, step1,0001\{,\}000\) at matched budgets of 512 \(dashed\) and 1024 \(solid\) scored tokens\. Shading marks layers 0–3; envelopes span three seeds\.

### 4\.3A Fused Kernel for Selected Logits

Subset selection introduces a second, implementation\-level memory problem that is independent of which objective is used\. Both variants require different selected vocabulary rows at each of theB⋅TB\\\!\\cdot\\\!Tpositions in a microbatch\. A naïve implementation gathers those rows of the unembeddingWUW\_\{U\}into a\[B⋅T,k,d\]\[B\\\!\\cdot\\\!T,k,d\]intermediate before multiplying by the translated activation, potentially consuming more memory than the subset objective saves\. AtB⋅T=8192B\\\!\\cdot\\\!T=8192,k=512k=512, andd=4096d=4096, this tensor alone occupies68\.768\.7GB in FP32\. Our fused CUDA kernel computes each selected logit directly in registers and never forms the gathered tensor, reducing the peak memory to0\.150\.15GB \(Appendix[F](https://arxiv.org/html/2608.10260#A6)\)\.

##### Combined feasibility\.

Lens parameters and optimizer state persist throughout training and cannot be reduced by microbatching, so low rank is what makes dense 70B coverage possible under our single\-device lens placement\. Vocabulary\-readout activations instead grow with microbatch size, context length, and vocabulary size, and are therefore reducible by microbatching\. Subset\-KL governs what readout workloads fit at a given budget: at 8B both subset objectives train at 4K context on an A100\-40GB, whereas full KL exhausts the device beyond 2K\. At 70B, full KL still fits the low\-rank configuration at our production microbatch, so Subset\-KL is not what makes that individual run possible; it is what determines the context length, batch size, and vocabulary scale reachable without additional hardware, and lets the same framework trade statistical fidelity against memory and throughput\. Appendix[G](https://arxiv.org/html/2608.10260#A7)gives the complete accounting\.

##### Frontier scale\.

We execute eight optimizer steps of a 126\-hookpoint residual\-lens stack on LLaMA\-3\.1\-405B\-Instruct across96×96\\timesA100\-40GB GPUs\. Atr=64r\{=\}64that stack contains266266M translator parameters, a roughly127127\-fold reduction from the corresponding33\.833\.8B\-parameter dense stack before optimizer state\. This run verifies execution of the complete training path at 405B scale; Appendix[G](https://arxiv.org/html/2608.10260#A7)gives the configuration, memory measurements, and loss trajectory\.

## 5Lens Application Case Studies

![Refer to caption](https://arxiv.org/html/2608.10260v1/x6.png)Figure 6:Interpretability applications\.\(a\) Prompt\-injection AUROC across ten tasks \(bars: means; points: individual tasks\)\. \(b\) Fraction of achievable 2WMH intervention gain captured atτ=4\\tau=4\. \(c\) GPT\-2 toxicity reduction by component; the dashed line marks the prior attention\-head\-only method\. \(d\) At 8B, detected toxicity and intervention effectiveness are negatively correlated across component types\. No full\-rank reference is trainable at 70B \(Section[4](https://arxiv.org/html/2608.10260#S4)\)\.We evaluate whether OmniLens preserves the downstream utility of trained lenses by revisiting three applications from the literature:\(i\)prompt\-injection detection\(Belroseet al\.[2023](https://arxiv.org/html/2608.10260#bib.bib6)\),\(ii\)multi\-hop memory injection\(Sakarvadiaet al\.[2024](https://arxiv.org/html/2608.10260#bib.bib32)\), and\(iii\)toxicity localization\(Pettyjohn[2025](https://arxiv.org/html/2608.10260#bib.bib109)\)\. Because the same lens can be attached to residual, attention, and MLP activations, we also extend the toxicity intervention beyond the attention heads examined previously\. Unless noted otherwise, results use rank\-64Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}lenses, with Top\-kk, full\-rank, and logit\-lens baselines where available\. Experimental details/caveats appear in Appendices[B](https://arxiv.org/html/2608.10260#A2)and[L](https://arxiv.org/html/2608.10260#A12)\.

##### Prompt\-injection detection\.

Belroseet al\.\([2023](https://arxiv.org/html/2608.10260#bib.bib6)\)show that a prompt injection perturbs the model’s intermediate computation before it perturbs the output\. Following their protocol, we fit outlier detectors to the layer\-by\-layer predictions decoded by each lens on clean multiple\-choice prompts and test whether they flag prompts carrying an injection attack\. On the five tasks where the original study reports near\-perfect detection,Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}attains a mean AUROC of0\.9970\.997at every scale and matches the full\-rank reference within bootstrap uncertainty where that reference exists \(Fig\.[6](https://arxiv.org/html/2608.10260#S5.F6)a\)\. At 70B, where no full\-rank reference is available,Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}detects attacks the logit lens misses on the knowledge tasks \(ARC\-Easy0\.800\.80vs\.0\.680\.68; SciQ0\.670\.67vs\.0\.560\.56\)\. Protocols and per\-task results are in Appendices[H](https://arxiv.org/html/2608.10260#A8)and[K](https://arxiv.org/html/2608.10260#A11)\.

##### Multi\-hop factual reasoning\.

Sakarvadiaet al\.\([2024](https://arxiv.org/html/2608.10260#bib.bib32)\)attribute multi\-hop failures to a recall gap and correct them by*memory injection*: the difference between explicit and implicit intermediate representations is added to the residual stream at a chosen layer\. The intervention itself needs no lens; the lens’s role is to select*where*to inject\. The intervention replicates at scale: on 2WikiMultiHop \(2WMH\)\(Hoet al\.[2020](https://arxiv.org/html/2608.10260#bib.bib68)\)it raises the model’s ownP​\(answer\)P\(\\text\{answer\}\)by 6\.6×\\timesat 8B and 5\.1×\\timesat 70B\. At a deliberately strong intervention scale \(τ=4\\tau\{=\}4\), where the causal optimum moves to an interior layer, theTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}\-selected layer captures61%61\\%\(8B\) and82%82\\%\(70B\) of the maximum achievable gain, compared with34%34\\%and66%66\\%for final\-layer injection \(Fig\.[6](https://arxiv.org/html/2608.10260#S5.F6)b\)\.Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}is the only evaluated lens that selects a middle layer and remains near\-optimal across all three training seeds\. Selection criteria, layer\-by\-τ\\tausweeps, controls, and distortion measurements are in Appendix[I](https://arxiv.org/html/2608.10260#A9)\.

##### Toxicity localization and intervention\.

FollowingPettyjohn \([2025](https://arxiv.org/html/2608.10260#bib.bib109)\), we rank attention heads by the prevalence of toxic tokens in their lens predictions, then subtract an unembedding\-derived toxicity direction at the selected heads\. At 8B, OmniLens reproduces the concentration previously observed on GPT\-2: the top five heads carry approximately 35% of the detected toxic signal, and the strongest head is stable across lens variants and training seeds\. Extending the intervention across all six hookpoint types, however, changes the conclusion\. The most effective targets lie outside the attention heads examined previously at both GPT\-2 and 8B and differ between models \(Fig\.[6](https://arxiv.org/html/2608.10260#S5.F6)c,d\)\. At 8B, detection and intervention rankings are negatively correlated \(Spearman−0\.43\-0\.43\): MLP outputs show the weakest detected toxicity but produce the largest reduction when modified, whereas MLP inputs score highly under detection but increase toxicity when modified\. Where a behavior is most visible, then, need not be where it is most causally actionable\. Complete rankings, intervention sweeps, and the 70B analysis appear in Appendix[J](https://arxiv.org/html/2608.10260#A10)\. Together, these studies show that OmniLens preserves the established applications of trained lenses while enabling model\-wide comparisons that component\-specific lens families cannot make\.

## 6Discussion and Conclusion

##### Conclusion\.

OmniLens makes dense trained\-lens coverage practical at previously inaccessible model scales by combining low\-rank translators, memory\-efficient vocabulary\-subset objectives, and one framework for residual, attention, and MLP activations\. We train a 482\-hookpoint lens ensemble for LLaMA\-3\.3\-70B and demonstrate training for eight optimizer steps on LLaMA\-3\.1\-405B\. Across three established interpretability applications, the resulting lenses preserve key reference results while enabling comparisons that component\-specific methods cannot make, including the finding that the hookpoints where a behavior is most visible and those where intervention is most effective can differ\. OmniLens thus moves trained lenses from sparse, specialized readouts toward a model\-wide interpretability instrument\.

##### Broader impact\.

Democratizing model\-wide interpretability at modern LLM scales matters for AI safety, governance, and research equity\. OmniLens enables researchers to apply the same trained\-lens framework to anomaly detection, causal hookpoint selection, and behavioral localization throughout an architecture rather than at a few preselected components\. Prior work demonstrates how such localization can guide targeted interventions on safety\-relevant behavior, including toxicity\(Leeet al\.[2024](https://arxiv.org/html/2608.10260#bib.bib133)\), refusal\(Arditiet al\.[2024](https://arxiv.org/html/2608.10260#bib.bib134)\), and honesty\(Zouet al\.[2025](https://arxiv.org/html/2608.10260#bib.bib135)\), as well as weight\-level editing of specific factual associations\(Menget al\.[2022](https://arxiv.org/html/2608.10260#bib.bib82),[2023](https://arxiv.org/html/2608.10260#bib.bib130)\)\. The same capabilities could expose model vulnerabilities or manipulation points, a dual\-use risk shared by interpretability tools generally\.

##### Limitations and future work\.

Our experiments cover only components withdu=dd\_\{u\}=d\. The parameterization extends to narrower components, such as individual attention heads, by adding a map into the residual space, which we do not evaluate here\. The two vocabulary\-subset objectives make different tradeoffs: Top\-kkavoids the full\-vocabulary projection but is biased, whereasTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}retains unbiased gradients while still computing the full normalization\. These boundaries motivate learned input maps, more efficient unbiased normalization, and converged evaluations beyond 70B\.

## Acknowledgments

This research used resources of the Argonne Leadership Computing Facility, which is a U\.S\. Department of Energy Office of Science User Facility operated under contract DE\-AC02\-06CH11357\. This work was partially supported by the AuroraGPT project\. MS was supported by the U\.S\. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Department of Energy Computational Science Graduate Fellowship under Award Number DE\-SC0023112\.

## References

- Understanding intermediate layers using linear classifier probes\.External Links:1610\.01644,[Link](https://arxiv.org/abs/1610.01644)Cited by:[§1](https://arxiv.org/html/2608.10260#S1.p2.1)\.
- A\. Amini, T\. Vieira, and R\. Cotterell \(2025\)Better estimation of the Kullback–Leibler divergence between language models\.InThe Thirty\-ninth Annual Conference on Neural Information Processing Systems,Cited by:[Appendix D](https://arxiv.org/html/2608.10260#A4.p1.5),[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px3.p2.6)\.
- Anshumann, M\. A\. Zaidi, A\. Kedia, J\. Ahn, T\. Kwon, K\. Lee, H\. Lee, and J\. Lee \(2025\)Sparse logit sampling: accelerating knowledge distillation in LLMs\.In63rd Annual Meeting of the Association for Computational Linguistics,W\. Che, J\. Nabende, E\. Shutova, and M\. T\. Pilehvar \(Eds\.\),Vienna, Austria,pp\. 18085–18108\.External Links:[Link](https://aclanthology.org/2025.acl-long.885/),[Document](https://dx.doi.org/10.18653/v1/2025.acl-long.885),ISBN 979\-8\-89176\-251\-0Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px3.p3.1)\.
- Anthropic \(2025\)Circuit tracing: revealing computational graphs in language models\.Anthropic Research\.Note:[https://www\.anthropic\.com/research/circuit\-tracing](https://www.anthropic.com/research/circuit-tracing)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- A\. Arditi, O\. Obeso, A\. Syed, D\. Paleka, N\. Panickssery, W\. Gurnee, and N\. Nanda \(2024\)Refusal in language models is mediated by a single direction\.External Links:2406\.11717,[Link](https://arxiv.org/abs/2406.11717)Cited by:[§6](https://arxiv.org/html/2608.10260#S6.SS0.SSS0.Px2.p1.1)\.
- Y\. Belinkov, N\. Durrani, F\. Dalvi, H\. Sajjad, and J\. Glass \(2017\)What do neural machine translation models learn about morphology?\.In55th Annual Meeting of the Association for Computational Linguistics,R\. Barzilay and M\. Kan \(Eds\.\),Vancouver, Canada,pp\. 861–872\.External Links:[Link](https://aclanthology.org/P17-1080/),[Document](https://dx.doi.org/10.18653/v1/P17-1080)Cited by:[§1](https://arxiv.org/html/2608.10260#S1.p2.1)\.
- N\. Belrose, Z\. Furman, L\. Smith, D\. Halawi, I\. Ostrovsky, L\. McKinney, S\. Biderman, and J\. Steinhardt \(2023\)Eliciting latent predictions from transformers with the tuned lens\.External Links:2303\.08112,[Link](https://arxiv.org/abs/2303.08112)Cited by:[Table 3](https://arxiv.org/html/2608.10260#A1.T3.7.5.6.1.1),[Appendix K](https://arxiv.org/html/2608.10260#A11.p3.4),[Table 16](https://arxiv.org/html/2608.10260#A8.T16),[Appendix H](https://arxiv.org/html/2608.10260#A8.p1.3),[§1](https://arxiv.org/html/2608.10260#S1.p1.1),[§1](https://arxiv.org/html/2608.10260#S1.p5.1),[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px1.p1.8),[§3](https://arxiv.org/html/2608.10260#S3.SS0.SSS0.Px3.p1.3),[§4](https://arxiv.org/html/2608.10260#S4.p2.4),[§5](https://arxiv.org/html/2608.10260#S5.SS0.SSS0.Px1.p1.7),[§5](https://arxiv.org/html/2608.10260#S5.p1.2)\.
- G\. Blanc and S\. Rendle \(2018\)Adaptive sampled softmax with kernel based sampling\.In35th International Conference on Machine Learning,J\. Dy and A\. Krause \(Eds\.\),Proceedings of Machine Learning Research, Vol\.80,pp\. 590–599\.External Links:[Link](https://proceedings.mlr.press/v80/blanc18a.html)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px3.p3.1)\.
- L\. Bushnaq, D\. Braun, and L\. Sharkey \(2025\)Stochastic parameter decomposition\.arXiv preprint arXiv:2506\.20790\.Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- N\. Cammarata, G\. Goh, S\. Carter, C\. Voss, L\. Schubert, and C\. Olah \(2021\)Curve circuits\.Distill\.Note:https://distill\.pub/2020/circuits/curve\-circuitsExternal Links:[Document](https://dx.doi.org/10.23915/distill.00024.006)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- P\. F\. Christiano, J\. Leike, T\. Brown, M\. Martic, S\. Legg, and D\. Amodei \(2017\)Deep reinforcement learning from human preferences\.Advances in Neural Information Processing Systems30\.Cited by:[Appendix D](https://arxiv.org/html/2608.10260#A4.p2.3)\.
- cjadams, J\. Sorensen, J\. Elliott, L\. Dixon, M\. McDonald, nithum, and W\. Cukierski \(2017\)Toxic comment classification challenge\.Note:[https://kaggle\.com/competitions/jigsaw\-toxic\-comment\-classification\-challenge](https://kaggle.com/competitions/jigsaw-toxic-comment-classification-challenge)KaggleCited by:[Appendix J](https://arxiv.org/html/2608.10260#A10.p1.1)\.
- C\. Clark, K\. Lee, M\. Chang, T\. Kwiatkowski, M\. Collins, and K\. Toutanova \(2019a\)BoolQ: exploring the surprising difficulty of natural yes/no questions\.External Links:1905\.10044,[Link](https://arxiv.org/abs/1905.10044)Cited by:[Appendix H](https://arxiv.org/html/2608.10260#A8.p1.3)\.
- K\. Clark, U\. Khandelwal, O\. Levy, and C\. D\. Manning \(2019b\)What does BERT look at? An analysis of BERT’s attention\.InACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP,T\. Linzen, G\. Chrupała, Y\. Belinkov, and D\. Hupkes \(Eds\.\),Florence, Italy,pp\. 276–286\.External Links:[Link](https://aclanthology.org/W19-4828/),[Document](https://dx.doi.org/10.18653/v1/W19-4828)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- P\. Clark, I\. Cowhey, O\. Etzioni, T\. Khot, A\. Sabharwal, C\. Schoenick, and O\. Tafjord \(2018\)Think you have solved question answering? try arc, the ai2 reasoning challenge\.External Links:1803\.05457,[Link](https://arxiv.org/abs/1803.05457)Cited by:[Appendix H](https://arxiv.org/html/2608.10260#A8.p1.3)\.
- A\. Conmy, A\. N\. Mavor\-Parker, A\. Lynch, S\. Heimersheim, and A\. Garriga\-Alonso \(2023\)Towards automated circuit discovery for mechanistic interpretability\.External Links:2304\.14997,[Link](https://arxiv.org/abs/2304.14997)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- A\. Conneau, G\. Kruszewski, G\. Lample, L\. Barrault, and M\. Baroni \(2018\)What you can cram into a single vector: probing sentence embeddings for linguistic properties\.External Links:1805\.01070,[Link](https://arxiv.org/abs/1805.01070)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- H\. Cunningham, A\. Ewart, L\. Riggs, R\. Huben, and L\. Sharkey \(2023\)Sparse autoencoders find highly interpretable features in language models\.External Links:2309\.08600,[Link](https://arxiv.org/abs/2309.08600)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- A\. Y\. Din, T\. Karidi, L\. Choshen, and M\. Geva \(2024\)Jump to conclusions: short\-cutting transformers with linear transformations\.InJoint International Conference on Computational Linguistics, Language Resources and Evaluation,pp\. 9615–9625\.Cited by:[Table 3](https://arxiv.org/html/2608.10260#A1.T3.7.5.8.3.1),[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px1.p1.8)\.
- N\. Elhage, N\. Nanda, C\. Olsson, T\. Henighan, N\. Joseph, B\. Mann, A\. Askell, Y\. Bai, A\. Chen, T\. Conerly, N\. DasSarma, D\. Drain, D\. Ganguli, Z\. Hatfield\-Dodds, D\. Hernandez, A\. Jones, J\. Kernion, L\. Lovitt, K\. Ndousse, D\. Amodei, T\. Brown, J\. Clark, J\. Kaplan, S\. McCandlish, and C\. Olah \(2021\)A mathematical framework for transformer circuits\.Transformer Circuits Thread\.Note:https://transformer\-circuits\.pub/2021/framework/index\.htmlCited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- A\. Ettinger, A\. Elgohary, and P\. Resnik \(2016\)Probing for semantic evidence of composition by means of simple classification tasks\.In1st Workshop on Evaluating Vector\-Space Representations for NLP,Berlin, Germany,pp\. 134–139\.External Links:[Link](https://aclanthology.org/W16-2524/),[Document](https://dx.doi.org/10.18653/v1/W16-2524)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- L\. Gao, S\. Biderman, S\. Black, L\. Golding, T\. Hoppe, C\. Foster, J\. Phang, H\. He, A\. Thite, N\. Nabeshima,et al\.\(2020\)The pile: an 800gb dataset of diverse text for language modeling\.arXiv preprint arXiv:2101\.00027\.Cited by:[Appendix B](https://arxiv.org/html/2608.10260#A2.SS0.SSS0.Px2.p1.1)\.
- A\. Ghandeharioun, A\. Caciularu, A\. Pearce, L\. Dixon, and M\. Geva \(2024\)Patchscopes: a unifying framework for inspecting hidden representations of language models\.arXiv preprint arXiv:2401\.06102\.Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- X\. Glorot and Y\. Bengio \(2010\)Understanding the difficulty of training deep feedforward neural networks\.In13th International Conference on Artificial Intelligence and Statistics,pp\. 249–256\.Cited by:[Appendix B](https://arxiv.org/html/2608.10260#A2.SS0.SSS0.Px5.p1.5),[§3](https://arxiv.org/html/2608.10260#S3.SS0.SSS0.Px2.p2.8)\.
- W\. Gurnee, N\. Sofroniew, A\. Pearce, M\. Piotrowski, I\. Kauvar, R\. Chen, A\. Soligo, P\. Bogdan, E\. Ong, R\. Wang, B\. Thompson, D\. Abrahams, S\. Kantamneni, E\. Ameisen, J\. Batson, and J\. Lindsey \(2026\)Verbalizable representations form a global workspace in language models\.External Links:2607\.15495,[Link](https://arxiv.org/abs/2607.15495)Cited by:[Table 3](https://arxiv.org/html/2608.10260#A1.T3.7.5.9.4.1),[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- L\. Hanu and Unitary team \(2020\)Detoxify\.Note:Github\. https://github\.com/unitaryai/detoxifyCited by:[Table 21](https://arxiv.org/html/2608.10260#A10.T21)\.
- J\. Hewitt and C\. D\. Manning \(2019\)A structural probe for finding syntax in word representations\.In2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies,J\. Burstein, C\. Doran, and T\. Solorio \(Eds\.\),Minneapolis, Minnesota,pp\. 4129–4138\.External Links:[Link](https://aclanthology.org/N19-1419/),[Document](https://dx.doi.org/10.18653/v1/N19-1419)Cited by:[§1](https://arxiv.org/html/2608.10260#S1.p2.1)\.
- G\. Hinton, O\. Vinyals, and J\. Dean \(2015\)Distilling the knowledge in a neural network\.External Links:1503\.02531,[Link](https://arxiv.org/abs/1503.02531)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px3.p2.6)\.
- X\. Ho, A\. D\. Nguyen, S\. Sugawara, and A\. Aizawa \(2020\)Constructing a multi\-hop QA dataset for comprehensive evaluation of reasoning steps\.External Links:2011\.01060,[Link](https://arxiv.org/abs/2011.01060)Cited by:[§5](https://arxiv.org/html/2608.10260#S5.SS0.SSS0.Px2.p1.11)\.
- E\. J\. Hu, Y\. Shen, P\. Wallis, Z\. Allen\-Zhu, Y\. Li, S\. Wang, L\. Wang, and W\. Chen \(2021\)LoRA: low\-rank adaptation of large language models\.External Links:2106\.09685,[Link](https://arxiv.org/abs/2106.09685)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px2.p1.2)\.
- M\. Ivanitskiy, A\. F\. Spies, T\. Räuker, G\. Corlouer, C\. Mathwin, L\. Quirke, C\. Rager, R\. Shah, D\. Valentine, C\. D\. Behn,et al\.\(2024\)Linearly structured world representations in maze\-solving transformers\.InUniReps: 1st Workshop on Unifying Representations in Neural Models,pp\. 133–143\.Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- A\. Katharopoulos and F\. Fleuret \(2019\)Not all samples are created equal: deep learning with importance sampling\.External Links:1803\.00942,[Link](https://arxiv.org/abs/1803.00942)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px3.p3.1)\.
- S\. Katz, Y\. Belinkov, M\. Geva, and L\. Wolf \(2024\)Backward lens: projecting language model gradients into the vocabulary space\.InConference on Empirical Methods in Natural Language Processing,pp\. 2390–2422\.Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px1.p1.8)\.
- J\. Kramár, J\. Engels, Z\. Wang, B\. Chughtai, R\. Shah, N\. Nanda, and A\. Conmy \(2026\)Building production\-ready probes for Gemini\.arXiv preprint arXiv:2601\.11516\.Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- A\. Lee, X\. Bai, I\. Pres, M\. Wattenberg, J\. K\. Kummerfeld, and R\. Mihalcea \(2024\)A mechanistic understanding of alignment algorithms: a case study on dpo and toxicity\.External Links:2401\.01967,[Link](https://arxiv.org/abs/2401.01967)Cited by:[§6](https://arxiv.org/html/2608.10260#S6.SS0.SSS0.Px2.p1.1)\.
- J\. Liu, L\. Cui, H\. Liu, D\. Huang, Y\. Wang, and Y\. Zhang \(2020\)LogiQA: a challenge dataset for machine reading comprehension with logical reasoning\.External Links:2007\.08124,[Link](https://arxiv.org/abs/2007.08124)Cited by:[Appendix H](https://arxiv.org/html/2608.10260#A8.p1.3)\.
- K\. Meng, D\. Bau, A\. Andonian, and Y\. Belinkov \(2022\)Locating and editing factual associations in GPT\.In36th International Conference on Neural Information Processing Systems,NIPS ’22,Red Hook, NY, USA\.External Links:ISBN 9781713871088Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1),[§6](https://arxiv.org/html/2608.10260#S6.SS0.SSS0.Px2.p1.1)\.
- K\. Meng, A\. S\. Sharma, A\. Andonian, Y\. Belinkov, and D\. Bau \(2023\)Mass\-editing memory in a transformer\.External Links:2210\.07229,[Link](https://arxiv.org/abs/2210.07229)Cited by:[§6](https://arxiv.org/html/2608.10260#S6.SS0.SSS0.Px2.p1.1)\.
- S\. Merity, C\. Xiong, J\. Bradbury, and R\. Socher \(2016\)Pointer sentinel mixture models\.External Links:1609\.07843,[Link](https://arxiv.org/abs/1609.07843)Cited by:[§C\.2](https://arxiv.org/html/2608.10260#A3.SS2.SSS0.Px2.p1.18)\.
- Nostalgebraist \(2020\)Interpreting GPT: the logit lens\.Note:LessWrong forum post[https://www\.lesswrong\.com/posts/AcKRB8wDpdaN6v6ru/interpreting\-gpt\-the\-logit\-lens](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens)Cited by:[Table 3](https://arxiv.org/html/2608.10260#A1.T3.5.3.3.2),[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px1.p1.8)\.
- C\. Olah, A\. Mordvintsev, and L\. Schubert \(2017\)Feature visualization\.Distill2,pp\.\.External Links:[Document](https://dx.doi.org/10.23915/distill.00007)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- H\. Orgad, F\. Barez, T\. Haklay, I\. Lee, M\. Mosbach, A\. Reusch, N\. Saphra, B\. C\. Wallace, S\. Wiegreffe, E\. Wong, I\. Tenney, and M\. Geva \(2026\)Interpretability can be actionable\.Preprint arXiv:2605\.11161\.Cited by:[§1](https://arxiv.org/html/2608.10260#S1.p1.1)\.
- K\. Pal, J\. Sun, A\. Yuan, B\. C\. Wallace, and D\. Bau \(2023\)Future lens: anticipating subsequent tokens from a single hidden state\.arXiv preprint arXiv:2311\.04897\.Cited by:[Table 3](https://arxiv.org/html/2608.10260#A1.T3.7.5.7.2.1),[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px1.p1.8)\.
- J\. Pettyjohn \(2025\)Mind your manners: detoxifying language models via attention head intervention\.Note:ACM SRCExternal Links:[Link](https://src.acm.org/binaries/content/assets/src/2025/jordan-pettyjohn.pdf)Cited by:[Table 3](https://arxiv.org/html/2608.10260#A1.T3.7.5.5.1),[§1](https://arxiv.org/html/2608.10260#S1.p5.1),[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px2.p2.1),[§5](https://arxiv.org/html/2608.10260#S5.SS0.SSS0.Px3.p1.1),[§5](https://arxiv.org/html/2608.10260#S5.p1.2)\.
- M\. Sakarvadia, A\. Ajith, A\. Khan, D\. Grzenda, N\. Hudson, A\. Bauer, K\. Chard, and I\. Foster \(2024\)Memory injections: correcting multi\-hop reasoning failures during inference in transformer\-based language models\.External Links:2309\.05605,[Link](https://arxiv.org/abs/2309.05605)Cited by:[§1](https://arxiv.org/html/2608.10260#S1.p5.1),[§5](https://arxiv.org/html/2608.10260#S5.SS0.SSS0.Px2.p1.11),[§5](https://arxiv.org/html/2608.10260#S5.p1.2)\.
- M\. Sakarvadia, A\. Khan, A\. Ajith, D\. Grzenda, N\. Hudson, A\. Bauer, K\. Chard, and I\. Foster \(2023\)Attention lens: a tool for mechanistically interpreting the attention head information retrieval mechanism\.External Links:2310\.16270,[Link](https://arxiv.org/abs/2310.16270)Cited by:[Table 3](https://arxiv.org/html/2608.10260#A1.T3.6.4.4.1),[§1](https://arxiv.org/html/2608.10260#S1.p1.1),[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px1.p1.8)\.
- V\. Sanh, L\. Debut, J\. Chaumond, and T\. Wolf \(2020\)DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter\.External Links:1910\.01108,[Link](https://arxiv.org/abs/1910.01108)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px3.p2.6)\.
- J\. Schulman \(2020\)Approximating KL divergence\.Note:[http://joschu\.net/blog/kl\-approx\.html](http://joschu.net/blog/kl-approx.html)Online; accessed 02\-March\-2026Cited by:[Appendix D](https://arxiv.org/html/2608.10260#A4.p1.5)\.
- Z\. Shao, P\. Wang, Q\. Zhu, R\. Xu, J\. Song, X\. Bi, H\. Zhang, M\. Zhang, Y\. Li, Y\. Wu,et al\.\(2024\)DeepSeekMath: pushing the limits of mathematical reasoning in open language models\.arXiv preprint arXiv:2402\.03300\.Cited by:[Appendix D](https://arxiv.org/html/2608.10260#A4.p1.5),[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px3.p2.6)\.
- N\. Shapiraet al\.\(2026\)Agents of chaos\.External Links:2602\.20021,[Link](https://arxiv.org/abs/2602.20021)Cited by:[§1](https://arxiv.org/html/2608.10260#S1.p1.1)\.
- L\. Sharkey, D\. Braun, and B\. Millidge \(2022\)Taking features out of superposition with sparse autoencoders\.InAI Alignment Forum,Vol\.6,pp\. 12–13\.Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- R\. Socher, A\. Perelygin, J\. Wu, J\. Chuang, C\. D\. Manning, A\. Ng, and C\. Potts \(2013\)Recursive deep models for semantic compositionality over a sentiment treebank\.InProceedings of the 2013 Conference on Empirical Methods in Natural Language Processing,D\. Yarowsky, T\. Baldwin, A\. Korhonen, K\. Livescu, and S\. Bethard \(Eds\.\),Seattle, Washington, USA,pp\. 1631–1642\.External Links:[Link](https://aclanthology.org/D13-1170/)Cited by:[Appendix H](https://arxiv.org/html/2608.10260#A8.p1.3)\.
- N\. Stiennon, L\. Ouyang, J\. Wu, D\. Ziegler, R\. Lowe, C\. Voss, A\. Radford, D\. Amodei, and P\. F\. Christiano \(2020\)Learning to summarize with human feedback\.Advances in Neural Information Processing Systems33,pp\. 3008–3021\.Cited by:[Appendix D](https://arxiv.org/html/2608.10260#A4.p2.3)\.
- Y\. Tang and R\. Munos \(2025\)On a few pitfalls in KL divergence gradient estimation for RL\.arXiv preprint arXiv:2506\.09477\.Cited by:[Appendix D](https://arxiv.org/html/2608.10260#A4.p2.3),[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px3.p2.6)\.
- A\. Templeton, T\. Conerly, J\. Marcus, J\. Lindsey, T\. Bricken, B\. Chen, A\. Pearce, C\. Citro, E\. Ameisen, A\. Jones, H\. Cunningham, N\. L\. Turner, C\. McDougall, M\. MacDiarmid, C\. D\. Freeman, T\. R\. Sumers, E\. Rees, J\. Batson, A\. Jermyn, S\. Carter, C\. Olah, and T\. Henighan \(2024\)Scaling monosemanticity: extracting interpretable features from Claude 3 Sonnet\.Transformer Circuits Thread\.External Links:[Link](https://transformer-circuits.pub/2024/scaling-monosemanticity/index.html)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- G\. Trimigno, G\. Lombardo, and S\. Cagnoni \(2026\)Low\-rank lens for scalable LLMs interpretability\.In34th European Symposium on Artificial Neural Networks, Computational Intelligence and Machine Learning,Bruges, Belgium,pp\. 35–40\.External Links:ISBN 9782875870964Cited by:[Table 3](https://arxiv.org/html/2608.10260#A1.T3.7.5.10.5.1),[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px2.p2.1)\.
- A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, Ł\. Kaiser, and I\. Polosukhin \(2017\)Attention is all you need\.Advances in Neural Information Processing Systems30\.Cited by:[§2](https://arxiv.org/html/2608.10260#S2.p1.10)\.
- J\. Vig and Y\. Belinkov \(2019\)Analyzing the structure of attention in a transformer language model\.InACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP,pp\. 63–76\.External Links:[Link](https://aclanthology.org/W19-4808/),[Document](https://dx.doi.org/10.18653/v1/W19-4808)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- E\. Voita, D\. Talbot, F\. Moiseev, R\. Sennrich, and I\. Titov \(2019\)Analyzing multi\-head self\-attention: specialized heads do the heavy lifting, the rest can be pruned\.In57th Annual Meeting of the Association for Computational Linguistics,pp\. 5797–5808\.External Links:[Link](https://aclanthology.org/P19-1580/),[Document](https://dx.doi.org/10.18653/v1/P19-1580)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- A\. Wang, A\. Singh, J\. Michael, F\. Hill, O\. Levy, and S\. R\. Bowman \(2018\)GLUE: a multi\-task benchmark and analysis platform for natural language understanding\.InProceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP,T\. Linzen, G\. Chrupała, and A\. Alishahi \(Eds\.\),Brussels, Belgium,pp\. 353–355\.External Links:[Link](https://aclanthology.org/W18-5446/),[Document](https://dx.doi.org/10.18653/v1/W18-5446)Cited by:[Appendix H](https://arxiv.org/html/2608.10260#A8.p1.3)\.
- K\. Wang, A\. Variengien, A\. Conmy, B\. Shlegeris, and J\. Steinhardt \(2022\)Interpretability in the wild: a circuit for indirect object identification in GPT\-2 small\.External Links:2211\.00593,[Link](https://arxiv.org/abs/2211.00593)Cited by:[§2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px4.p1.1)\.
- J\. Welbl, N\. F\. Liu, and M\. Gardner \(2017\)Crowdsourcing multiple choice science questions\.InProceedings of the 3rd Workshop on Noisy User\-generated Text,L\. Derczynski, W\. Xu, A\. Ritter, and T\. Baldwin \(Eds\.\),Copenhagen, Denmark,pp\. 94–106\.External Links:[Link](https://aclanthology.org/W17-4413/),[Document](https://dx.doi.org/10.18653/v1/W17-4413)Cited by:[Appendix H](https://arxiv.org/html/2608.10260#A8.p1.3)\.
- A\. Williams, N\. Nangia, and S\. R\. Bowman \(2018\)A broad\-coverage challenge corpus for sentence understanding through inference\.InProceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 \(Long Papers\),M\. Walker, H\. Ji, and A\. Stent \(Eds\.\),New Orleans, Louisiana,pp\. 1112–1122\.External Links:[Link](https://aclanthology.org/N18-1101/),[Document](https://dx.doi.org/10.18653/v1/N18-1101)Cited by:[Appendix H](https://arxiv.org/html/2608.10260#A8.p1.3)\.
- L\. Zhang and J\. Ba \(2026\)EMA policy gradient: taming reinforcement learning for LLMs with EMA anchor and top\-k KL\.arXiv preprint arXiv:2602\.04417\.Cited by:[Appendix D](https://arxiv.org/html/2608.10260#A4.p2.3)\.
- B\. Zhou, D\. Khashabi, Q\. Ning, and D\. Roth \(2019\)“Going on a vacation” takes longer than “going for a walk”: a study of temporal commonsense understanding\.InProceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing \(EMNLP\-IJCNLP\),K\. Inui, J\. Jiang, V\. Ng, and X\. Wan \(Eds\.\),Hong Kong, China,pp\. 3363–3369\.External Links:[Link](https://aclanthology.org/D19-1332/),[Document](https://dx.doi.org/10.18653/v1/D19-1332)Cited by:[Appendix H](https://arxiv.org/html/2608.10260#A8.p1.3)\.
- A\. Zou, L\. Phan, S\. Chen, J\. Campbell, P\. Guo, R\. Ren, A\. Pan, X\. Yin, M\. Mazeika, A\. Dombrowski, S\. Goel, N\. Li, M\. J\. Byun, Z\. Wang, A\. Mallen, S\. Basart, S\. Koyejo, D\. Song, M\. Fredrikson, J\. Z\. Kolter, and D\. Hendrycks \(2025\)Representation engineering: a top\-down approach to ai transparency\.External Links:2310\.01405,[Link](https://arxiv.org/abs/2310.01405)Cited by:[§6](https://arxiv.org/html/2608.10260#S6.SS0.SSS0.Px2.p1.1)\.

## Appendix ALens Taxonomy

Table[3](https://arxiv.org/html/2608.10260#A1.T3)organizes the published \(linear\) lens constructions as discussed in Section[2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px1)\. We note there are two linear lens parameterizations in the literature\. The most common is to view the lens as a component\-to\-residual stream map—i\.e\.ℒℓ,u∈ℝd×du\\mathcal\{L\}\_\{\\ell,u\}\\in\\mathbb\{R\}^\{d\\times d\_\{u\}\}, wheredud\_\{u\}is the width of componentuu\(equal to the model widthddfor every hookpoint we train\)—whence composing with the unembedding matrix maps to vocabulary space:

Qℓ,u\(⋅\|x\)=softmax\(WUη\(ℒℓ,uhℓ,u\+bℓ,u\)\)\.Q\_\{\\ell,u\}\(\\cdot\|x\)=\\operatorname\{softmax\}\\left\(W\_\{U\}\\,\\eta\\\!\\left\(\\mathcal\{L\}\_\{\\ell,u\}\\,h\_\{\\ell,u\}\+b\_\{\\ell,u\}\\right\)\\right\)\.This is the formulation presented in Section[2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px1)\. In the second parameterization the lens maps directly to vocabulary space\. That is,ℒℓ,u∈ℝ\|V\|×du\\mathcal\{L\}\_\{\\ell,u\}\\in\\mathbb\{R\}^\{\|V\|\\times d\_\{u\}\}and

Qℓ,u\(⋅\|x\)=softmax\(ℒℓ,uhℓ,u\+bℓ,u\)\.Q\_\{\\ell,u\}\(\\cdot\|x\)=\\operatorname\{softmax\}\\left\(\\mathcal\{L\}\_\{\\ell,u\}\\,h\_\{\\ell,u\}\+b\_\{\\ell,u\}\\right\)\.\(6\)We mark the methods in Table[3](https://arxiv.org/html/2608.10260#A1.T3)employing the parameterization of Eq\. \([6](https://arxiv.org/html/2608.10260#A1.E6)\) with a dagger \(†\\dagger\)\.

Table 3:Published lens constructions\. All lenses use the parameterization Eq\. \([1](https://arxiv.org/html/2608.10260#S2.E1)\) unless marked by†\\dagger, in which case they use the parameterization of Eq\. \([6](https://arxiv.org/html/2608.10260#A1.E6)\)\.Two axes distinguish the entries of Table[3](https://arxiv.org/html/2608.10260#A1.T3)\. The first is*where*the lens reads: prior methods each fix one component type, while the hookpoint\-agnostic parameterization of Section[3](https://arxiv.org/html/2608.10260#S3)covers all of them under one training run\. The second is*how much*the translator costs: the earlier trained entries use a full\-rankd×dd\\times dmap per hookpoint, the two low\-rank predecessors reduce that cost at a single fixed component, and our identity\-residual translator Eq\. \([3](https://arxiv.org/html/2608.10260#S3.E3)\) carries the low\-rank cost to every component type\.

## Appendix BExperimental Setup and Implementation

##### Models\.

GPT\-2 Small \(124M; 12 layers,d=768d\{=\}768,\|V\|=50,257\|V\|\{=\}50\{,\}257; HFopenai\-community/gpt2\), LLaMA\-3\-8B \(32 layers,d=4096d\{=\}4096,\|V\|=128,256\|V\|\{=\}128\{,\}256; HFmeta\-llama/Meta\-Llama\-3\-8B\-Instruct\), and LLaMA\-3\.3\-70B \(80 layers,d=8192d\{=\}8192,\|V\|=128,256\|V\|\{=\}128\{,\}256; HFmeta\-llama/Llama\-3\.3\-70B\-Instruct; abbreviated LLaMA\-3\-70B in the main text\)\. All models are frozen bf16 checkpoints; the LLaMA models are the instruct\-tuned variants, while GPT\-2 is the base pretrained model\. The 405B feasibility run \(Section[4](https://arxiv.org/html/2608.10260#S4)\) additionally uses LLaMA\-3\.1\-405B\-Instruct \(126 layers,d=16,384d\{=\}16\{,\}384,\|V\|=128,256\|V\|\{=\}128\{,\}256; HFmeta\-llama/Llama\-3\.1\-405B\-Instruct\), likewise a frozen bf16 instruct checkpoint\.

##### Data\.

All lenses were trained onThe Pile\(Gaoet al\.[2020](https://arxiv.org/html/2608.10260#bib.bib112)\), tokenized into non\-overlapping chunks of the training sequence length \(stride equals sequence length\)\. No additional preprocessing or filtering was applied\.222A preliminary swap to an alternative corpus produced negligible differences in final\-layer and layerwise KL, suggesting lens quality is determined primarily by the model’s representations rather than the training distribution\. We did not study this systematically\.

##### Optimization\.

We used AdamW with default moments\(β1,β2\)=\(0\.9,0\.999\)\(\\beta\_\{1\},\\beta\_\{2\}\)=\(0\.9,0\.999\), weight decay0, gradient clipping at norm1\.01\.0, and learning rate10−310^\{\-3\}with cosine decay to zero\. Optimizer steps target an effective batch of218=262,1442^\{18\}=262\{,\}144tokens, accumulated over an integer number of full microbatches \(a step never ends mid\-microbatch; when the token target is not divisible by the per\-microstep token count, the trainer rounds*up*to the next full microbatch\)\. The realized effective batches are therefore262,144262\{,\}144tokens for GPT\-2 \(88microsteps of32×102432\\times 1024on a single A100\-40GB\),327,680327\{,\}680for LLaMA\-3\-8B \(44microsteps of2×10242\\times 1024per rank, DDP over1010nodes=40=40ranks\), and393,216393\{,\}216for LLaMA\-3\-70B \(22microsteps of2×10242\\times 1024per rank over9696FSDP ranks on2424nodes of4×4\\timesA100\-40GB\)\. GPT\-2 trains for1,0001\{,\}000steps with no warmup; the LLaMA lens sets train for250250steps with5050warmup steps\. All runs use bf16 mixed precision with fp32 loss reductions and log\-partition computation, and training seed0unless stated otherwise \(see*Checkpoint selection and seeds*\)\.

##### Software\.

Lens training uses Python 3\.10, PyTorch 2\.10 \(CUDA 12\.8\), Transformers 5\.3, and Datasets 4\.6 on Linux, on the4×4\\timesA100\-40 GB nodes described above\. KL evaluation reuses the upstream tuned\-lens loop, which ships with the released OmniLens code alongside per\-run configurations \(run\_config\.csv\)\. Hyperparameter exploration is reported where it occurred: rankrr\(nine values; Table[5](https://arxiv.org/html/2608.10260#A3.T5)\), estimator budgetskkand\(khead,ktail\)\(k\_\{\\mathrm\{head\}\},k\_\{\\mathrm\{tail\}\}\)\(Table[6](https://arxiv.org/html/2608.10260#A3.T6)\), and training length \(convergence diagnostics; Appendix[C](https://arxiv.org/html/2608.10260#A3)\) were swept with the selection criteria stated alongside each; optimizer settings follow standard tuned\-lens practice and were not tuned\.

##### Lenses and objectives\.

LoRA lenses user=64r\{=\}64by default withα=r\\alpha\{=\}r\(unit scaling\)\.AAwas initialized as Xavier uniform\(Glorot and Bengio[2010](https://arxiv.org/html/2608.10260#bib.bib110)\)whileBBandbbwere initialized to zero, so that the lens is the identity map at initialization\.

Estimator budgets:the GPT\-2 sweeps are as listed in Tables[6](https://arxiv.org/html/2608.10260#A3.T6)and[8](https://arxiv.org/html/2608.10260#A3.T8); the production LLaMA lenses use Top\-kkwithk=512k\{=\}512andTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}with\(khead,ktail\)=\(512,1024\)\(k\_\{\\mathrm\{head\}\},k\_\{\\mathrm\{tail\}\}\)=\(512,1024\)under the teacher\-tail proposal\. Full\-KL runs use chunked evaluation with chunk size4,0964\{,\}096\.

Hookpoints:Residual places one translator per layer while expanded places six hookpoint types per layer plus embedding and final\-norm readouts \(6​L\+26L\{\+\}2: 74 / 194 / 482 hookpoints at GPT\-2 / LLaMA\-3\-8B / LLaMA\-3\-70B\)\. Dense coverage costs about the same as a standard one\-per\-layer full\-rank stack at GPT\-2 and far less at scale \(GPT\-2: 7\.3M parameters vs\. 7\.1M full\-rank residual\-only; 8B: 102M vs\. 537M; 70B: 509M vs\. 5\.4B\)\.

Baselines:We select the strongest set of baselines that are trainable for a given model size\. For GPT\-2 every variant is trainable, including the full\-rank full\-KL reference\. For LLaMA\-3\-8B a multi\-node full\-rank reference exists on the residual hookset; at 70B no full\-rank variant is trainable \(optimizer state388388GB; Figure[14](https://arxiv.org/html/2608.10260#A7.F14)\) and the dense LoRA\+\{\+\}Subset\-KL stack trains at a measured34\.734\.7GB per GPU\.

Training costs:training one production OmniLens stack \(expanded hookset, LoRAr=64r\{=\}64, Subset\-KL\) costs a few GPU\-hours at GPT\-2,≈20\{\\approx\}20node\-hours at 8B, and≈360\{\\approx\}360node\-hours at 70B on4×4\\timesA100\-40 GB nodes\.

##### Checkpoint selection and seeds\.

GPT\-2 comparisons use step1,0001\{,\}000; LLaMA lenses use the end of the annealed schedule \(step250250\); Appendix[C](https://arxiv.org/html/2608.10260#A3)reports convergence diagnostics supporting both budgets\. Headline tables report training seed0; two seed studies quantify training\-seed variability under otherwise identical recipes\. At GPT\-2, the full\-rank baseline, LoRA full\-KL, Top\-kk\(k=1024k\{=\}1024\),Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}\(512\+512512\{\+\}512\), and RS \(khead=0k\_\{\\mathrm\{head\}\}\{=\}0\) configurations were each retrained at seeds11–22\. Mean and early\-layer KL values are stable for every estimator \(sd≤0\.04\\leq 0\.04nats\), while final\-layer KL sd ranges from0\.00070\.0007\(Top\-kk\) to0\.0120\.012\(Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}\)\. The full\-rank baseline’s own final KL spans0\.0390\.039–0\.0490\.049\.

For LLaMA\-3\-8B, the three lens variants used in the case studies—Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}, Top\-kk, and the full\-rank residual reference—were each retrained at seeds11–22under the identical 250\-step schedule; the resulting DART and injection\-selection stability is reported in Section[5](https://arxiv.org/html/2608.10260#S5)\. The detection experiments ensemble the isolation forest over five seeds and report bootstrap confidence intervals \(Appendix[H](https://arxiv.org/html/2608.10260#A8)\)\. Full per\-run configurations \(run\_config\.csv\) ship with the released OmniLens code\.

### B\.1Implementation Details

#### Hookbox System Architecture

The hookbox system provides a component\-agnostic interface for attaching lenses to arbitrary intermediate activations during the forward pass\. Hooks toggle independently at any granularity \(whole blocks, individual heads, or projections within a head\), so one instrumented model serves coarse, fine\-grained, and targeted analyses \(Figure[7](https://arxiv.org/html/2608.10260#A2.F7)\)\. This section details the key design decisions and implementation strategies\.

![Refer to caption](https://arxiv.org/html/2608.10260v1/placeholder_plots/DynamicHooks.png)Figure 7:Dynamic hook configurations inhookbox\. Any hook can be toggled independently: a macro\-scale configuration \(left\) reads only block outputs, a micro\-scale configuration \(center\) instruments every projection of every head, and mixed\-scale configurations \(right\) target specific components\. Active hooks \(∙\\bullet\) incur memory and compute cost; inactive hooks \(∘\\circ\) are free\.##### Distributed Model Handling

Modern large language models are typically wrapped in distributed training frameworks \(DDP, FSDP, DeepSpeed\) that shard parameters and optimizer states across devices\. These wrappers introduce additional module hierarchy that must be unwrapped to access the underlying model architecture\.

Our implementation automatically detects and unwraps these wrappers using a recursive traversal:

```
def unwrap_model(model):
    """Recursively unwrap distributed wrappers."""
    while hasattr(model, ’module’):
        model = model.module
    if hasattr(model, ’_fsdp_wrapped_module’):
        model = model._fsdp_wrapped_module
    return model
```

For FSDP models using ZeRO\-3 \(full parameter sharding\), individual modules may be sharded across devices\. We implement an activation gathering mechanism that consolidates sharded tensors before passing to the lens:

```
@torch.no_grad()
def gather_from_shards(tensor, process_group):
    """Gather tensor from FSDP shards."""
    world_size = dist.get_world_size(process_group)
    gather_list = [torch.zeros_like(tensor)
                   for _ in range(world_size)]
    dist.all_gather(gather_list, tensor,
                    group=process_group)
    return torch.cat(gather_list, dim=0)
```

##### Activation Checkpointing Compatibility

Gradient checkpointing \(activation recomputation\) reduces memory by discarding intermediate activations during the forward pass and recomputing them during the backward pass\. However, this creates a challenge for hooks: they are invoked twice \(once during initial forward, once during recomputation\), potentially double\-counting activations\.

Recomputation passes are marked with an explicit context\-manager flag that the checkpointed forward is wrapped in:

```
class CheckpointingState:
    _is_recomputing: bool = False

    @classmethod
    def is_recomputing(cls):
        return cls._is_recomputing

    @classmethod
    @contextmanager
    def recomputation_context(cls):
        old = cls._is_recomputing
        cls._is_recomputing = True
        try:
            yield
        finally:
            cls._is_recomputing = old
```

Hooks skip execution during recomputation to avoid duplicate processing:

```
def hook_fn(module, input, output):
    if CheckpointingState.is_recomputing():
        return output  # Skip during recompute
    activations = process_activations(output)
    return output
```

##### Hook Registration via Predicates

Rather than hardcoding module names, we allow users to specify attachment points via predicate functions over \(name, module\) pairs\. This provides flexibility and architectural agnostic attachment:

```
def register_hooks(model, predicate, hook_fn):
    """Register hooks on modules matching
    predicate."""
    handles = []
    for name, module in model.named_modules():
        if predicate(name, module):
            handle = module.register_forward_hook(
            hook_fn
            )
            handles.append(handle)
    return handles

# Example usage
pred = lambda n, m: ’layer’ in n and ’output’ in n
handles = register_hooks(model, pred, my_hook)
```

Common predicates include:

- •Residual streams:lambda n,m: ’layer\.\{\}\.output’ in n
- •All attention heads:lambda n,m: isinstance\(m, AttentionHead\)
- •Specific layers:lambda n,m: n in \[’layer\.0’, ’layer\.10’, ’layer\.20’\]
- •MLP sublayers:lambda n,m: ’mlp’ in n and ’post’ in n

#### Numerical Stability Considerations

##### Mixed Precision Training

We use BF16 mixed precision to reduce memory and accelerate training\. However, certain operations require FP32 for numerical stability:

- •Logit computation accumulates in FP32 \(fused kernel\)
- •KL divergence uses log\-space computations to avoid underflow
- •Gradient clipping operates on FP32 master weights
- •Normalization layers maintain FP32 running statistics

##### Softmax and Log\-Softmax Stability

Computing KL divergence requires bothlog⁡P\\log Pandlog⁡Q\\log Q\. We use the log\-sum\-exp trick to prevent overflow/underflow:

```
def stable_log_softmax(logits):
    """Numerically stable log-softmax."""
    max_logit = logits.max(dim=-1, keepdim=True)[0]
    shifted = logits - max_logit
    return shifted - torch.log(
        torch.exp(shifted).sum(dim=-1, keepdim=True)
    )
```

The two Subset\-KL modes differ here\. Top\-kkrenormalizes over the selected subset only \(a deliberately truncated objective\)\. TheTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}mode instead uses the*exact*full\-vocabulary log\-partition: the student’s full logit row is produced by a single fused matmul, reduced immediately to its logsumexp, and only the selected subset of logits is retained \(𝒪​\(N⋅V\)\\mathcal\{O\}\(N\\cdot V\)transient per site, a few hundred MB at our configurations, freed by layer\-wise backward\)\. Gradients therefore flow to every logit through the partition term, which is what makes the estimator’s gradients exactly unbiased \(Theorem[1](https://arxiv.org/html/2608.10260#Thmtheorem1)\); we verify this numerically against full\-KL autograd\.

##### Numerical Guards

Teacher and student log\-probabilities are computed directly with fp32log\_softmax; probabilities are never reconstructed by exponentiation and re\-floored, so the teacher distribution entering the loss is unmodified\. The tail mass is computed by directly summing the teacher’s tail probabilities in fp32 \(not as1−Phead1\-P\_\{\\mathrm\{head\}\}, which would lose precision when the head mass is close to one\)\. Two guards exist in the sampling path: that tail\-mass sum is floored at10−1210^\{\-12\}where it appears as a normalizing denominator \(preventing division by zero only in the degenerate case where the head captures all numerical mass\), and proposal probabilities are floored at the fp32 subnormal boundary \(10−4510^\{\-45\}\) before their logarithm\. Neither guard biases the estimator: a token with zero proposal probability can never be drawn by multinomial sampling, so the floor is never active at a sampled index and the importance weights of Eq\. \([5](https://arxiv.org/html/2608.10260#S4.E5)\) are unaffected\. \(This argument concerns numerics only; the support condition of Theorem[1](https://arxiv.org/html/2608.10260#Thmtheorem1),R​\(v\|x\)\>0R\(v\|x\)\>0whereverP​\(v\|x\)\>0P\(v\|x\)\>0on the tail, must hold for the estimator itself, and the teacher\-tail default satisfies it by construction\.\)

#### Distributed Training Configuration

Table[4](https://arxiv.org/html/2608.10260#A2.T4)summarizes the distributed training strategy used at each model scale\. Full configuration details are available inOmniLens/configs/distributed/

Table 4:Distributed training configuration for different model scales\.

## Appendix CAdditional Ablations

### C\.1Rank Ablation: Full Layerwise Results

Table[5](https://arxiv.org/html/2608.10260#A3.T5)is the complete rank sweep summarized by Figure[3](https://arxiv.org/html/2608.10260#S3.F3)and Table[1](https://arxiv.org/html/2608.10260#S3.T1)in Section[3](https://arxiv.org/html/2608.10260#S3)\. Figures[8](https://arxiv.org/html/2608.10260#A3.F8)–[10](https://arxiv.org/html/2608.10260#A3.F10)show the full layerwise heatmaps for all metrics reported in Section[3](https://arxiv.org/html/2608.10260#S3)\. All metrics show the same qualitative pattern: at the late layers, a small rank already closes most of the gap to the full\-rank reference, and increasingrrcloses it further; layer 0’s gap shrinks far more slowly, so the earliest layers remain the binding constraint at every rank\.

Table 5:Full rank ablation on GPT\-2 Small over 131,072 Pile test tokens\. KL is to the teacher; Top\-1, Pearsonρ\\rho\(8,192 positions/layer\), and Kendallτ\\tau@100 \(512 positions/layer\) are vs\. the full\-rank tuned\-lens baseline\.*Final*/*mean*= final layer vs\. average over 12 layers\.Shadedrow \(r=64r\{=\}64\) is the recommended default;boldmarks KL values that numerically exceed the baseline \(no statistical claim: the baseline’s own final KL varies0\.0390\.039–0\.0490\.049across training seeds, wider than these gaps\)\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x7.png)Figure 8:Layerwise KL difference \(baseline KL minus LoRA KL\) by layer and rank, GPT\-2 Small\. Red indicates LoRA exceeds the full\-rank baseline; blue indicates LoRA improves upon it\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x8.png)Figure 9:Top\-1 agreement \(*left*\) and top\-10 overlap \(*right*\) between LoRA and full\-rank lens outputs by layer, over every evaluated checkpoint grouped by run family \(white separators\)\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x9.png)Figure 10:Token\-level Pearsonρ\\rho\(*left*\) and Kendallτ\\tau@100 \(*right*\) against the full\-rank baseline, by layer and LoRA rank, GPT\-2 Small\.
### C\.2Estimator Fidelity and Rank Sensitivity

Table[7](https://arxiv.org/html/2608.10260#A3.T7)reports token\-level fidelity of the practically trained lenses to the full\-KL reference, and Table[8](https://arxiv.org/html/2608.10260#A3.T8)shows that the estimator ranking of Section[4](https://arxiv.org/html/2608.10260#S4)\(Top\-kkwins final\-layer KL;Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}wins mean and early\-layer KL\) is stable across every LoRA rank swept in Section[3](https://arxiv.org/html/2608.10260#S3)\. Figure[11](https://arxiv.org/html/2608.10260#A3.F11)shows the layerwise profiles of the RS \(pure teacher sampling\) baseline of Table[2](https://arxiv.org/html/2608.10260#S4.T2), and Table[6](https://arxiv.org/html/2608.10260#A3.T6)reports the complete budget sweep behind that table’s matched\-budget subset\.

Table 6:Complete Subset\-KL estimator sweep \(GPT\-2 Small,r=64r=64, step 1000, residual hookset\); Table[2](https://arxiv.org/html/2608.10260#S4.T2)shows the matched\-budget subset, and columns are as defined there\. Tok/s and Peak GB are median steady\-state throughput and peak per\-GPU memory, re\-measured with the released trainer \(one node, DDP over 4×\\timesA100\-40GB, identical batch recipe\)\. RS rows are trained at three seeds per budget up to Total 1024; KL cells are seed 0, with cross\-seed sd≤0\.011\\leq 0\.011on every reduction\.Boldmarks the best per column among the practical estimators \(Top\-kkand Top\-kk\+IS\)\.##### Seed stability and pure\-sampling diagnostics\.

Across three seeds at Total=1024=1024, the mean and early\-layer ordering of Top\-kk,Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}, and full KL is unchanged; the standard deviation is at most0\.040\.04nats for those reductions\. The final\-layer gap is less stable:Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}spans0\.0450\.045–0\.0670\.067, whereas Top\-kkyields0\.046±0\.0010\.046\\pm 0\.001\.

For pure teacher sampling, increasing the sampled budget does not improve the fixed\-step optimization\. Final\-layer KL rises from0\.1940\.194at Total=512=512to0\.6560\.656at Total=1024=1024and1\.0731\.073at Total=1536=1536, despite the estimator remaining unbiased\. The degradation is present in the training objective itself: the Total=512=512loss decreases from2\.422\.42to2\.322\.32through step1,0001\{,\}000, whereas the Total=1536=1536loss rises from3\.583\.58at step500500to4\.094\.09at step1,0001\{,\}000\. The learning rate, optimizer, batch construction, step count, and random seed are unchanged; onlyktailk\_\{\\mathrm\{tail\}\}differs\. This supports the narrower conclusion that the fixed recipe exhibits budget\-dependent optimization instability; we do not claim that larger unbiased samples are intrinsically harmful\.

##### Training\-budget diagnostics\.

The GPT\-2 budgets of the estimator tables are past convergence: under the fixed schedule, every variant’s KL at step250250is1212–35%35\\%above its step\-1,0001\{,\}000value, with 5%\-convergence \(the earliest step from which KL stays within5%5\\%of its final value\) between steps625625and875875\. Held\-out KL improves monotonically with continued training for every full\-KL, Top\-kk, andTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}run at all three scales; the pure\-sampling instability above is the sole exception\. The LLaMA recipe \(cosine anneal at250250\) is schedule\-converged: at 8B every reduction flattens by steps150150–200200with tight seed spread, and at 70B mean KL converges by roughly step100100across the1111production checkpoints\. Convergence order is layerwise at 70B: early layers reach their floor by about step8080, the last layers near step180180, and the final layer is still improving at the end of the schedule, so longer schedules chiefly buy final\-layer fidelity\. These diagnostics use reduced evaluation budgets \(3232k Pile tokens at 8B;44k WikiText\-2 tokens\(Merityet al\.[2016](https://arxiv.org/html/2608.10260#bib.bib46)\)at 70B\) and are trajectory\-internal; they support convergence claims, not cross\-protocol KL comparisons\.

![Refer to caption](https://arxiv.org/html/2608.10260v1/x10.png)Figure 11:Layerwise KL of RS \(pure teacher sampling,khead=0k\_\{\\mathrm\{head\}\}\{=\}0; grey, line style by total budget\) against the two subset estimators at Total = 1024\. RS tracks or beats Top\-kkthrough the early and middle layers but fails to close the final layers, and degrades uniformly as its sampled budget grows\. Curves are seed 0; shaded envelopes span min–max over three training seeds where available \(RS 1536 is single\-seed\)\.Table 7:Token\-level fidelity of LoRAr=64r\{=\}64lenses trained with each estimator \(Top\-kkk=512k=512; IS\(512,512\)\(512,512\)\), measured against the full\-rank tuned\-lens baseline of Table[5](https://arxiv.org/html/2608.10260#A3.T5)\. The Full\-KL row is the LoRA lens trained with the exact loss \(the ceiling the estimators approach\), not a self\-comparison\. Pearsonρ\\rhoover full\-vocabulary log\-probs; Kendallτ\\tau@100 over the top\-100 union; Top\-1 argmax agreement; Top\-10 mean top\-set overlap\.†Negative Kendall indicates the top\-100 token rankings at layers 0–3 are slightly anti\-correlated with the reference; IS restores a positive early correlation \(\+0\.243\+0\.243\) at the same head budget\. Early = mean over layers 0–3; Final = layer 11\.

Table 8:Rank sensitivity at fixed budget \(Top\-kkk=512k=512; IS\(512,512\)\(512,512\)\), as KL to the teacher at step 1000\.Shadedrow \(r=64r\{=\}64\) is the recommended default;boldmarks the better practical estimator \(Top\-kkvs\. IS\) per cell, with Full\-KL as the reference\. Top\-kkwins final\-layer KL; IS wins mean and early\-layer KL\.

### C\.3Scaling Comparison: Full Results

Table[9](https://arxiv.org/html/2608.10260#A3.T9)reports measured single\-GPU peak memory for LLaMA\-3\-8B lens training on the expanded hookset \(194 sites; frozen bf16 teacher, LoRAr=64r\{=\}64lenses, and readout on one A100\-40GB\) at the production microbatch of 2 sequences, across sequence lengths\. These are the measurements behind Figure[13](https://arxiv.org/html/2608.10260#A7.F13): full\-KL runs out of memory beyond 2K context, while both subset objectives train at 4K; none survives 8K\.

Table 9:Measured single\-GPU peak memory \(GB\) for LLaMA\-3\-8B expanded\-hookset lens training at microbatch 2, by objective and sequence length\.

## Appendix DSampling\-Based KL Estimators

This appendix expands the estimator background of Section[2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px3)\. In the distillation direction the trainable studentQℓ,uQ\_\{\\ell,u\}is the second argument of the KL, and the outer expectation Eq\. \([2](https://arxiv.org/html/2608.10260#S2.E2)\) runs over the frozen teacherPP\. Because the sampling distribution carries no trainable parameters, any estimator of the form

D^KL\(P∥Qℓ,u\)=1k∑i=1klogP​\(ti\|x\)Qℓ,u​\(ti\|x\),ti∼i\.i\.d\.P\(⋅\|x\)\\widehat\{D\}\_\{\\mathrm\{KL\}\}\(P\\\|Q\_\{\\ell,u\}\)=\\frac\{1\}\{k\}\\sum\_\{i=1\}^\{k\}\\log\\frac\{P\(t\_\{i\}\|x\)\}\{Q\_\{\\ell,u\}\(t\_\{i\}\|x\)\},\\quad t\_\{i\}\\overset\{i\.i\.d\.\}\{\\sim\}P\(\\cdot\|x\)\(7\)is unbiased forDKL​\(P∥Qℓ,u\)D\_\{\\mathrm\{KL\}\}\(P\\\|Q\_\{\\ell,u\}\), and the expectation commutes with gradients in the lens parameters\.Schulman \([2020](https://arxiv.org/html/2608.10260#bib.bib69)\)catalogues alternate expressions for the summand in Eq\. \([7](https://arxiv.org/html/2608.10260#A4.E7)\): K1 and K3 are unbiased, with K3 additionally reducing variance, whereas K2 trades bias for variance; such summand replacements are also adopted at scale\(Shaoet al\.[2024](https://arxiv.org/html/2608.10260#bib.bib70)\)\. Importance sampling refines plain Monte Carlo by exploiting access to the exact teacher probabilitiesP​\(ti\|x\)P\(t\_\{i\}\|x\), not just samples fromP\(⋅\|x\)P\(\\cdot\|x\); seeAminiet al\.\([2025](https://arxiv.org/html/2608.10260#bib.bib72)\)for a treatment of this idea at the sequence level\. The tail term of Eq\. \([5](https://arxiv.org/html/2608.10260#S4.E5)\) is such an importance\-sampled estimator, restricted to the tail with the head handled exactly; Theorem[1](https://arxiv.org/html/2608.10260#Thmtheorem1)gives the resulting guarantee\.

In reinforcement learning from human feedback, by contrast, the KL appears as a regularizer whose*first*argument is the trainable policy\(Christianoet al\.[2017](https://arxiv.org/html/2608.10260#bib.bib113); Stiennonet al\.[2020](https://arxiv.org/html/2608.10260#bib.bib114)\)\. The expectation is then taken with respect to the distribution being optimized, so sampling and differentiation interact, and estimators that are well behaved for distillation can misbehave\(Tang and Munos[2025](https://arxiv.org/html/2608.10260#bib.bib108)\)\. A contemporaneous RL\-side estimator\(Zhang and Ba[2026](https://arxiv.org/html/2608.10260#bib.bib71)\)mirrorsTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}’s head\-plus\-sampled\-tail structure in this reversed direction: it applies Schulman\-style summands within the sampled term and restores unbiasedness by rejection sampling, whereTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}reweights byP/RP/R\.

## Appendix ETop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}Theoretical Analysis

![Refer to caption](https://arxiv.org/html/2608.10260v1/x11.png)Figure 12:TheTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}decomposition: the contribution of thekheadk\_\{\\mathrm\{head\}\}most probable teacher tokens is computed exactly \(blue\), and the remaining mass is estimated fromktailk\_\{\\mathrm\{tail\}\}importance\-weighted tail samples \(green\)\. Under the teacher\-tail proposal, every importance weight equals1−Phead1\-P\_\{\\mathrm\{head\}\}\.Unbiasedness of pure teacher\-sampling KL estimators is classical \(Appendix[D](https://arxiv.org/html/2608.10260#A4)\), and RS\-KD \(Section[2](https://arxiv.org/html/2608.10260#S2.SS0.SSS0.Px3)\) establishes it for random\-sampling knowledge distillation\. The theorem below is the corresponding guarantee for our*stratified*estimator: an exact deterministic head, an arbitrary lens\-independent positive\-support tail proposal, and the distillation direction of the KL, with the student’s log\-partition computed exactly\. The proof for Theorem[1](https://arxiv.org/html/2608.10260#Thmtheorem1)is below:

###### Proof\.

The head term of Eq\. \([5](https://arxiv.org/html/2608.10260#S4.E5)\) is deterministic given the teacher,

D^KL,head​\(P∥Qℓ,u\)=∑v∈ℋP​\(v\|x\)​log⁡P​\(v\|x\)Qℓ,u​\(v\|x\)\.\\hat\{D\}\_\{\\mathrm\{KL\},\\mathrm\{head\}\}\(P\\\|Q\_\{\\ell,u\}\)=\\sum\_\{v\\in\\mathcal\{H\}\}P\(v\|x\)\\log\\frac\{P\(v\|x\)\}\{Q\_\{\\ell,u\}\(v\|x\)\}\.For the tail term of Eq\. \([5](https://arxiv.org/html/2608.10260#S4.E5)\), each summand is an importance\-weighted draw fromRR, so

𝔼ti∼R​\[P​\(ti\|x\)R​\(ti\|x\)​log⁡P​\(ti\|x\)Qℓ,u​\(ti\|x\)\]\\displaystyle\\mathbb\{E\}\_\{t\_\{i\}\\sim R\}\\\!\\left\[\\frac\{P\(t\_\{i\}\|x\)\}\{R\(t\_\{i\}\|x\)\}\\log\\frac\{P\(t\_\{i\}\|x\)\}\{Q\_\{\\ell,u\}\(t\_\{i\}\|x\)\}\\right\]=∑v∉ℋR​\(v\|x\)​P​\(v\|x\)R​\(v\|x\)​log⁡P​\(v\|x\)Qℓ,u​\(v\|x\)\\displaystyle\\qquad=\\sum\_\{v\\notin\\mathcal\{H\}\}R\(v\|x\)\\,\\frac\{P\(v\|x\)\}\{R\(v\|x\)\}\\,\\log\\frac\{P\(v\|x\)\}\{Q\_\{\\ell,u\}\(v\|x\)\}=∑v∉ℋP​\(v\|x\)​log⁡P​\(v\|x\)Qℓ,u​\(v\|x\),\\displaystyle\\qquad=\\sum\_\{v\\notin\\mathcal\{H\}\}P\(v\|x\)\\log\\frac\{P\(v\|x\)\}\{Q\_\{\\ell,u\}\(v\|x\)\},and averaging over thektailk\_\{\\mathrm\{tail\}\}i\.i\.d\. draws leaves this expectation unchanged:

𝔼​\[D^KL,tail​\(P∥Qℓ,u\)\]=∑v∉ℋP​\(v\|x\)​log⁡P​\(v\|x\)Qℓ,u​\(v\|x\)\.\\mathbb\{E\}\\left\[\\hat\{D\}\_\{\\mathrm\{KL\},\\mathrm\{tail\}\}\(P\\\|Q\_\{\\ell,u\}\)\\right\]=\\sum\_\{v\\notin\\mathcal\{H\}\}P\(v\|x\)\\log\\frac\{P\(v\|x\)\}\{Q\_\{\\ell,u\}\(v\|x\)\}\.Adding the head term recoversDKL​\(P∥Qℓ,u\)D\_\{\\mathrm\{KL\}\}\(P\\\|Q\_\{\\ell,u\}\)\. For the gradient claim,D^Top​\-​k\+IS\\hat\{D\}\_\{\\mathrm\{\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}\}\}is a finite weighted sum oflog⁡Qℓ,u​\(t\|x\)\\log Q\_\{\\ell,u\}\(t\|x\)terms whose weights and sampling distribution do not involve the lens parameters, so expectation and gradient commute\. The teacher\-tail defaultR​\(v\|x\)=P​\(v\|x\)/\(1−Phead\)R\(v\|x\)=P\(v\|x\)/\(1\-P\_\{\\mathrm\{head\}\}\)satisfies the support condition by construction, with all importance weights equal to1−Phead1\-P\_\{\\mathrm\{head\}\}\. ∎

Note that the theorem applies to the implemented objective because the student log\-probabilities use the exact full\-vocabulary partition function \(Appendix[B\.1](https://arxiv.org/html/2608.10260#A2.SS1)\): only the KL summands are subsampled, and every logit receives gradient through the partition term\. Under the teacher\-tail default, whose importance weights are the constant1−Phead1\-P\_\{\\mathrm\{head\}\}, the coefficient multiplying the log\-partition gradient equals one on every draw, exactly as under full KL; for general lens\-independent proposals it equals one in expectation\. We also verify the claim numerically with the shipped training code: on a synthetic teacher–student pair \(\|V\|=200\|V\|\{=\}200, linear student\), theTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}gradient averaged over4,0004\{,\}000resamplings matches the exact full\-KL autograd gradient to0\.2%0\.2\\%relative error \(cosine similarity1\.00001\.0000\), while Top\-kkrenormalization, which truncates the partition function, exhibits a63%63\\%relative gradient bias\. A second check at LLaMA\-scale vocabulary \(\|V\|=128,256\|V\|\{=\}128\{,\}256\) uses extreme teacher logits under which2\.4%2\.4\\%of the vocabulary underflows fp32 softmax outright; it matches to within10−410^\{\-4\}relative error at500500resamplings, confirming that the numerical guards of Appendix[B\.1](https://arxiv.org/html/2608.10260#A2.SS1)do not modify the teacher distribution\. Both checks ship as unit tests in the OmniLens repository\. Unbiasedness throughout refers to the raw stochastic gradients; it does not extend through gradient clipping or optimizer updates, which are nonlinear in the gradient\.

##### Proposal and sampling details\.

We sample with replacement, retain repeated draws with multiplicity, and weight each occurrence by its density ratio\. We rejected schemes that deduplicate draws and reweight by inverse inclusion probabilities: on LLM\-scale vocabularies their Horvitz–Thompson\-style weights span orders of magnitude and produced unstable gradient spikes in preliminary experiments\. The teacher\-tail proposal avoids this behavior because its importance weights are constant\. Other proposals, such as oversampling tokens on which the lens and teacher disagree, require no change to Eq\. \([5](https://arxiv.org/html/2608.10260#S4.E5)\), provided lens\-dependent proposal probabilities are held fixed during differentiation rather than differentiated through\.

## Appendix FIndexed Logits CUDA Implementation

This appendix provides implementation details for the indexed logits kernel, including the forward and backward CUDA kernels and empirical benchmarks\. The current implementation prioritizes correctness and memory efficiency; performance optimizations such as shared\-memory tiling and a fused LoRA extension are left as future work\.

### F\.1Kernel Design Principles

The indexed logits kernel addresses a fundamental memory bottleneck in subset\-based objectives\. Standard PyTorch operations materialize intermediate tensors in global memory, leading to the explosion documented in Section[4\.3](https://arxiv.org/html/2608.10260#S4.SS3)\. Our kernel design follows two principles:

1\. Streaming computation\.Each output element is computed independently by a single thread, which accumulates the result in FP32 registers without writing intermediate values to global memory\.

2\. Hidden\-state reuse\.Nearby threads often reuse the same hidden\-state rowH​\[i,:\]H\[i,:\]across different subset positions, while accesses toW​\[idx​\[i,j\],:\]W\[\\text\{idx\}\[i,j\],:\]are generally scattered due to arbitrary vocabulary indices\.

### F\.2Forward Pass Implementation

Algorithm[1](https://arxiv.org/html/2608.10260#alg1)shows the complete forward kernel\. Each CUDA thread computes one element of the output matrixout​\[i,j\]\\text\{out\}\[i,j\], whereiiindexes the sequence position andjjindexes the subset selection\.

Algorithm 1Indexed Logits Forward Kernel \(CUDA\)1:Input:

H∈ℝN×dH\\in\\mathbb\{R\}^\{N\\times d\}\(hidden states\),

W∈ℝV×dW\\in\\mathbb\{R\}^\{V\\times d\}\(weight matrix\),

idx∈ℤN×k\\text\{idx\}\\in\\mathbb\{Z\}^\{N\\times k\}\(vocabulary indices\)

2:Output:

out∈ℝN×k\\text\{out\}\\in\\mathbb\{R\}^\{N\\times k\}\(subset logits\)

3:Kernel configuration:

threads=256\\text\{threads\}=256,

blocks=⌈N​k/256⌉\\text\{blocks\}=\\lceil Nk/256\\rceil
4:

tid←blockIdx\.x×blockDim\.x\+threadIdx\.x\\text\{tid\}\\leftarrow\\text\{blockIdx\.x\}\\times\\text\{blockDim\.x\}\+\\text\{threadIdx\.x\}\{Global thread ID\}

5:if

tid<N×k\\text\{tid\}<N\\times kthen

6:

i←tid/ki\\leftarrow\\text\{tid\}/k\{Sequence position\}

7:

j←tidmodkj\\leftarrow\\text\{tid\}\\bmod k\{Subset position\}

8:

v←idx​\[i,j\]v\\leftarrow\\text\{idx\}\[i,j\]\{Vocabulary index\}

9:

10:

acc←0\.0​f\\text\{acc\}\\leftarrow 0\.0f\{FP32 accumulator in register\}

11:for

t=0t=0to

d−1d\-1do

12:

acc←acc\+H​\[i,t\]×W​\[v,t\]\\text\{acc\}\\leftarrow\\text\{acc\}\+H\[i,t\]\\times W\[v,t\]
13:endfor

14:

out​\[i,j\]←acc\\text\{out\}\[i,j\]\\leftarrow\\text\{acc\}
15:endif

Numerical precision\.The accumulator uses FP32 even when input tensors are FP16 or BF16\. This prevents catastrophic rounding errors when summing thousands of products\. The final result is cast to the output dtype only upon writing\.

### F\.3Backward Pass Implementation

The backward pass must compute three gradients:∇H\\nabla\_\{H\},∇W\\nabla\_\{W\}, and∇idx\\nabla\_\{\\text\{idx\}\}\(which isNone, since indices are discrete\)\.

Gradient w\.r\.t\. hidden states\.This is implemented as a gather\-weighted reduction over the selected vocabulary rows,

∇H\[i,:\]=∑j=1k∇out\[i,j\]⋅W​\[idx​\[i,j\],:\],\\nabla\_\{H\}\[i,:\]=\\sum\_\{j=1\}^\{k\}\\nabla\_\{\\text\{out\}\}\[i,j\]\\cdot W\[\\text\{idx\}\[i,j\],:\],\(8\)
via a separate kernel with similar structure to the forward pass\.

Gradient w\.r\.t\. weights\.This requires care due to duplicate indices\. Multiple threads may need to update the same rowW​\[v,:\]W\[v,:\]if vocabulary indexvvappears multiple times inidx\. We use atomic additions:

atomicAdd​\(∇W\[v,t\],∇out\[i,j\]×H​\[i,t\]\)\.\\text\{atomicAdd\}\(\\nabla\_\{W\}\[v,t\],\\nabla\_\{\\text\{out\}\}\[i,j\]\\times H\[i,t\]\)\.\(9\)Atomic operations serialize writes to the same memory location, potentially creating contention when indices repeat\. Section[F\.4](https://arxiv.org/html/2608.10260#A6.SS4)quantifies the empirical sensitivity of the backward pass to collision rate\.

Precision\.The∇W\\nabla\_\{W\}buffer is accumulated in FP32\.∇H\\nabla\_\{H\}uses FP32 accumulation within each thread before being written in the input dtype\.

Algorithm[2](https://arxiv.org/html/2608.10260#alg2)shows the weight\-gradient kernel\.

Algorithm 2Indexed Logits Backward Kernel \(Weight Gradients\)1:Input:

∇out∈ℝN×k\\nabla\_\{\\text\{out\}\}\\in\\mathbb\{R\}^\{N\\times k\},

H∈ℝN×dH\\in\\mathbb\{R\}^\{N\\times d\},

idx∈ℤN×k\\text\{idx\}\\in\\mathbb\{Z\}^\{N\\times k\}
2:Output:

∇W∈ℝV×d\\nabla\_\{W\}\\in\\mathbb\{R\}^\{V\\times d\}\(weight gradients, initialized to zero\)

3:foreach

\(i,j\)\(i,j\)in paralleldo

4:

v←idx​\[i,j\]v\\leftarrow\\text\{idx\}\[i,j\]
5:

g←∇out\[i,j\]g\\leftarrow\\nabla\_\{\\text\{out\}\}\[i,j\]
6:for

t=0t=0to

d−1d\-1do

7:

val←g×H​\[i,t\]\\text\{val\}\\leftarrow g\\times H\[i,t\]\{Compute in FP32\}

8:

atomicAdd​\(∇W\[v,t\],val\)\\text\{atomicAdd\}\(\\nabla\_\{W\}\[v,t\],\\text\{val\}\)\{Safe concurrent update\}

9:endfor

10:endfor

### F\.4Benchmark Results

#### Correctness

Table[10](https://arxiv.org/html/2608.10260#A6.T10)reports forward and gradient errors relative to PyTorch reference implementations, averaged over three random seeds on an A100\-40GB GPU in FP16\. Forward errors are measured against the naïve subset reference; gradient errors are measured against autograd\. The observed forward errors are consistent with expected FP16 rounding behavior in mixed\-precision accumulation and output casting\.

Table 10:Fused kernel correctness \(FP16, mean over 3 seeds,N=4096N=4096,V=50257V=50257\)\.
#### Runtime and Memory

Table[11](https://arxiv.org/html/2608.10260#A6.T11)reports forward\+backward latency and peak allocated memory, averaged over 3 seeds\. We sweepNNat fixeddd=768,kk=128 and varyddat fixedNN=4096,kk=128\. Dense GEMM rows are omitted for the larger\-ddconfigurations for brevity; they follow the same qualitative pattern as the smaller\-ddrows and remain substantially more memory\-intensive than the fused kernel\.

Table 11:Indexed logits kernel benchmark results \(A100\-40 GB, FP16, mean over 3 seeds\)\.Table[12](https://arxiv.org/html/2608.10260#A6.T12)breaks down where the memory goes at the main\-text configuration: the naïve gather is costlier than full\-vocabulary decoding, and the fused kernel eliminates the intermediate outright\.

Table 12:Peak memory breakdown for logit computation atB⋅TB\{\\cdot\}T= 8192,dd= 4096,\|V\|\|V\|= 128K,kk= 512,fp32\.The fused kernel achieves 1\.4–1\.5×\\timesspeedup over naïve subset and reduces peak allocated memory by 7×\\times–14×\\timesacross configurations\. Memory savings grow withNNbecause naïve subset intermediate tensors scale with sequence length while fused kernel memory is dominated by model activations\. Performance relative to dense GEMM is configuration\-dependent; the primary advantage of the fused kernel is eliminating the memory blowup from materializing the\[N,k,d\]\[N,k,d\]intermediate tensor, which causes naïve subset memory usage to grow rapidly withNN,dd, andkk, making it increasingly impractical at larger scales\.

#### Atomic Collision Sensitivity

The backward kernel uses atomic additions to∇W\\nabla\_\{W\}, which may contend when multiple threads update the same vocabulary row\. Table[13](https://arxiv.org/html/2608.10260#A6.T13)measures forward and forward\+backward latency as a function of collision vocabulary sizeVcV\_\{c\}: drawing indices from a smallerVcV\_\{c\}increases the probability of repeated rows\.

Table 13:Fused kernel sensitivity to index collisions \(NN=4096,dd=1024,kk=64, FP16, mean over 3 seeds\)\.Counterintuitively, latency decreases asVcV\_\{c\}shrinks\. In the tested regimes, improved weight\-row locality dominates any slowdown from atomic contention\.

### F\.5Integration with PyTorch

The kernel is exposed as a PyTorch custom autograd function, see Listing[1](https://arxiv.org/html/2608.10260#LST1)\. This integrates seamlessly with standard PyTorch training loops, optimizer state management, and gradient checkpointing\. The implementation and benchmark scripts are available in the supplementary materials\.

1classIndexedLogitsFunction\(torch\.autograd\.Function\):

2@staticmethod

3defforward\(ctx,H,W,idx\):

4out=indexed\_logits\_cuda\.forward\(H,W,idx\)

5ctx\.save\_for\_backward\(H,W,idx\)

6returnout

7

8@staticmethod

9defbackward\(ctx,grad\_out\):

10H,W,idx=ctx\.saved\_tensors

11grad\_H,grad\_W=indexed\_logits\_cuda\.backward\(

12H,W,idx,grad\_out\)

13returngrad\_H,grad\_W,None

Listing 1:Custom autograd function implemented in PyTorch\.

## Appendix GMemory Model and Estimate Derivation

Table 14:The two memory buckets and the lever that controls each: low\-rank translators control the*Optimizer*column, Subset\-KL controls the*Readout aggregate*column\.Optimizer stateis exact from lens parameter counts and is*persistent*: no micro\-batching or gradient accumulation reduces it\.Readout aggregate\(the*projected*cells of this table\) is a calibrated projection,β​\(B⋅T\)​V\\beta\\,\(B\{\\cdot\}T\)\\,V, of the total readout activation over a shared reference batch of262,144262\{,\}144tokens \(the realized GPT\-2 effective batch; the realized LLaMA batches are larger, Appendix[B](https://arxiv.org/html/2608.10260#A2)\)*if it were processed without micro\-batching*; the instantaneous readout footprint scales with the microbatch and sequence length instead, and is measured directly in Figure[13](https://arxiv.org/html/2608.10260#A7.F13)\. Subset\-KL rows use the selected\-subset term \(Veff=kV\_\{\\text\{eff\}\}\{=\}k, the Top\-kkmode\); theTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}mode adds an exact\-partition term treated separately in the text\. Optimizer states assume AdamW\.LensLossOptimizer \(exact\)Readout aggregate \(projected\)Feasible*GPT\-2*, full expanded hookset \(6​L\+2=746L\{\+\}2=74hookpoints\)Full\-rankFull\-KL0\.50\.5GB9292GB✓LoRAr​64r64Full\-KL0\.090\.09GB9292GB✓Full\-rankSubset\-KL0\.50\.5GB2\.82\.8GB✓LoRAr​64r64Subset\-KL0\.090\.09GB2\.82\.8GB✓*LLaMA\-3\.3\-70B*, full expanded hookset \(6​L\+2=4826L\{\+\}2=482hookpoints\)Full\-rankFull\-KL388388GB235235GB✗LoRAr​64r64Full\-KL6\.1\\phantom\{00\}6\.1GB235235GB✓†Full\-rankSubset\-KL388388GB2\.8\\phantom\{00\}2\.8GB✗LoRAr​64r64Subset\-KL6\.1\\phantom\{00\}6\.1GB2\.8\\phantom\{00\}2\.8GB✓
†The235235GB readout aggregate is the unmicrobatched reference\-batch projection and, unlike the optimizer column, is reducible by micro\-batching; it does not by itself determine feasibility at the production microbatch\. We measured that configuration directly \(Table[15](https://arxiv.org/html/2608.10260#A7.T15)\): peak35\.535\.5GB, fitting with≈4\.5\{\\approx\}4\.5GB of nominal headroom\.

![Refer to caption](https://arxiv.org/html/2608.10260v1/x12.png)Figure 13:Measured single\-GPU peak memory for an 8B frozen teacher, lenses, and readout at the production microbatch\. Full KL exhausts the device beyond 2K context, whereas both subset objectives train at 4K \(measurements in Appendix Table[9](https://arxiv.org/html/2608.10260#A3.T9)\)\.Peak training memory on the lens GPU decomposes additively:

Mpeak=Mbase\+Mopt\+Mread,M\_\{\\text\{peak\}\}=M\_\{\\text\{base\}\}\+M\_\{\\text\{opt\}\}\+M\_\{\\text\{read\}\},\(10\)whereMbaseM\_\{\\text\{base\}\}\(frozen\-model shard, cached site activations, CUDA context\) is independent of lens rank and loss,MoptM\_\{\\text\{opt\}\}is the lens optimizer footprint \(the LoRA lever\), andMreadM\_\{\\text\{read\}\}is the readout/loss activation \(the Subset\-KL lever\)\. The lens and its optimizer reside on a single device, so both levers are charged to one GPU\.

![Refer to caption](https://arxiv.org/html/2608.10260v1/x13.png)Figure 14:The persistent optimizer bucket, exact from parameter counts\. Micro\-batching cannot reduce it: the full\-rank expanded stack’s optimizer state alone consumes an A100\-40GB at 8B and is ten times the device at 70B, while the LoRA stack peaks at6\.16\.1GB\.##### Optimizer state\.

Under Adam with bf16 AMP \(Figure[14](https://arxiv.org/html/2608.10260#A7.F14)\), each trainable parameter costs1212bytes: bf16 weight \(2\)\+\+bf16 gradient \(2\)\+\+fp32 first moment \(4\)\+\+fp32 second moment \(4\)\. ForSSsites of widthddand rankrr,

Pfull=S​\(d2\+d\),PLoRA=S​\(2​d​r\+d\),Mopt=12​P⋆​bytes,P\_\{\\text\{full\}\}=S\(d^\{2\}\{\+\}d\),\\ P\_\{\\text\{LoRA\}\}=S\(2dr\{\+\}d\),\\ M\_\{\\text\{opt\}\}=12\\,P\_\{\\star\}\\ \\text\{bytes\},where⋆=full\\star=\\text\{full\}orLoRA\. For GPT\-2 expandedS=74S\{=\}74andd=768d\{=\}768, soPfull=43\.7P\_\{\\text\{full\}\}\{=\}43\.7M yieldingMopt=0\.52M\_\{\\text\{opt\}\}=0\.52GB, whilePLoRA,r​64=7\.33P\_\{\\text\{LoRA\},r64\}\{=\}7\.33M yieldingMopt=0\.09M\_\{\\text\{opt\}\}=0\.09GB\.

For LLaMA\-3\-70B expandedS=482S\{=\}482andd=8192d\{=\}8192, soPfull=32\.35P\_\{\\text\{full\}\}\{=\}32\.35B yieldingMopt=388M\_\{\\text\{opt\}\}=388GB, whilePLoRA=0\.509P\_\{\\text\{LoRA\}\}\{=\}0\.509B \(counted from the trained checkpoint\) yieldingMopt=6\.1M\_\{\\text\{opt\}\}=6\.1GB\.

The1212bytes per parameter factor is confirmed empirically: on the residual hook set, replacing full with LoRA reduces parameters by5\.95\.9M and the measured peak by0\.070\.07GB=5\.9​M×12=5\.9\\text\{M\}\\times 12bytes\.

##### Parameter\-count comparisons\.

The per\-hookpoint reduction of the rank\-rrparameterization is1−\(2​r\+1\)/\(d\+1\)1\-\(2r\{\+\}1\)/\(d\{\+\}1\), independent of coverage\. Atr=64r\{=\}64, this yields an83\.2%83\.2\\%reduction in parameter count for GPT\-2 \(d=768d\{=\}768\), a96\.9%96\.9\\%reduction for LLaMA\-3\-8B \(d=4096d\{=\}4096\), a98\.4%98\.4\\%reduction for LLaMA\-3\-70B \(d=8192d\{=\}8192\), and a99\.2%99\.2\\%reduction for LLaMA\-3\-405B \(d=16384d\{=\}16384\)\. This grounds the abstract’s comparisons\. For LLaMA\-3\-8B, the residual\-only full\-rank stack holds537537M translator parameters against102102M for the six\-hookpoint OmniLens stack:80\.9%80\.9\\%fewer parameters with six times the coverage\.

Per\-head attention\-lens decoders are costlier still: the original design maps each head output to the vocabulary \(\|V\|×d\|V\|\\times dper head\), totaling12×12×50,257×768≈5\.612\\times 12\\times 50\{,\}257\\times 768\\approx 5\.6B parameters on GPT\-2 \(the largest model with published full\-rank per\-head decoders\), against which our full7\.37\.3M GPT\-2 stack is a99\.9%99\.9\\%reduction\. The same design instantiated for LLaMA\-3\-8B would contain≈538\{\\approx\}538B parameters\.

##### Readout\.

The loss materialises per\-site logits, softmax, KL, and their gradients over the scored vocabulary, looped over sites \(hence site\-independent\)\. For Full\-KL and Top\-kkit scales as

Mread=β​\(B⋅T\)​Veff,M\_\{\\text\{read\}\}=\\beta\\,\(B\{\\cdot\}T\)\\,V\_\{\\text\{eff\}\},\(11\)withVeff=\|V\|V\_\{\\text\{eff\}\}=\|V\|for Full\-KL andVeff=kV\_\{\\text\{eff\}\}=kfor Top\-kk\.Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}is*not*aVeff=kV\_\{\\text\{eff\}\}=kestimator: its exact log\-partition materializes the student’s full logit row before reduction \(Appendix[B\.1](https://arxiv.org/html/2608.10260#A2.SS1)\), so its readout carries an additional transient term,

MreadTop​\-​k\+IS≈γ​\(B⋅T\)​V\+β​\(B⋅T\)​\(khead\+ktail\),M\_\{\\text\{read\}\}^\{\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}\}\\approx\\gamma\\,\(B\{\\cdot\}T\)\\,V\\;\+\\;\\beta\\,\(B\{\\cdot\}T\)\\,\(k\_\{\\mathrm\{head\}\}\{\+\}k\_\{\\mathrm\{tail\}\}\),withγ<β\\gamma<\\betabecause only the student\-side logits and their logsumexp reduction are held \(no teacher\-side full\-vocabulary tensors or per\-token KL buffers\)\. The measured LLaMA\-3\-8B curves of Figure[13](https://arxiv.org/html/2608.10260#A7.F13)show this term is small once the log\-partition streams in vocabulary chunks:Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}tracks Top\-kkwithin1\.21\.2GB at every context, trains at 4K where Full\-KL cannot, and fails only at 8K alongside Top\-kk\. We computeβ\\betafrom the measured GPT\-2 Full\-KL−\-Top\-kkgap:β=\(17\.14−5\.74\)/\(32768×\(50257−512\)\)=7\.0\\beta=\(17\.14\-5\.74\)/\(32768\\times\(50257\-512\)\)=7\.0bytes per token⋅\\cdotvocab\-slot\. The gap covers theV−kV\-kslots Top\-kkdoes not score; re\-measuring with the released trainer of Table[2](https://arxiv.org/html/2608.10260#S4.T2)shifts both absolute peaks but reproduces the same11\.411\.4GB gap, leavingβ\\betaunchanged\. At the shared batchB⋅T=262,144B\{\\cdot\}T\{=\}262\{,\}144: Full\-KL readout is9292GB \(GPT\-2\) and235235GB \(70B\); the Top\-kksubset readout is0\.90\.9–2\.82\.8GB\.

##### Validation against trained runs\.

Four configurations were run end\-to\-end; their measured peaks reconstructed from the decomposition \(Table[15](https://arxiv.org/html/2608.10260#A7.T15)\)\. The two GPT\-2 runs independently back out the sameMbase≈9\.05M\_\{\\text\{base\}\}\\approx 9\.05GB, over\-determining and thus confirming the optimizer and readout terms\. The two 70B rows provide a second, independent cross\-check: trained with different loss functions \(hence differentMreadM\_\{\\text\{read\}\}terms\) but the same site count, model, and microbatch, they back outMbaseM\_\{\\text\{base\}\}values of27\.6527\.65and27\.5227\.52GB, agreeing to within0\.5%0\.5\\%despite neither being calibrated against the other\.

Table 15:Measured peaks decompose into the memory model\.MbaseM\_\{\\text\{base\}\}is implied \(Mpeak−Mopt−MreadM\_\{\\text\{peak\}\}\-M\_\{\\text\{opt\}\}\-M\_\{\\text\{read\}\}\); the two GPT\-2 rows agree, validating the terms\. All measured values are peak PyTorch\-allocated memory \(max\_memory\_allocated\) on the metrics\-writing rank; under FSDP the lens and optimizer state are replicated, so ranks are near\-symmetric\.∗Selected\-subset term \(0\.020\.02\) plus theTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}exact\-partition transient of the implementation this run was trained with, which materialized the full logit row before reduction:γ​\(B⋅T\)​V\\gamma\(B\{\\cdot\}T\)Vwithγ≈3\.5\\gamma\\approx 3\.5bytes per token⋅\\cdotvocab\-slot, calibrated from 8B measurements of that implementation \(Top​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}−\-vs\-Top\-kkgaps of\+1\.5\+1\.5GB at 2K and\+4\.4\+4\.4GB at 4K context, givingγ≈2\.9\\gamma\\approx 2\.9–4\.24\.2bytes\)\. The released trainer streams this reduction instead \(Appendix[B\.1](https://arxiv.org/html/2608.10260#A2.SS1)\); the current gaps of Table[9](https://arxiv.org/html/2608.10260#A3.T9)bound the streamed transient below0\.30\.3GB at 4K\. Assigning the transient toMreadM\_\{\\text\{read\}\}keepsMbaseM\_\{\\text\{base\}\}loss\-independent, as the decomposition requires\.†The Full\-KL readout atV=128,256V\{=\}128\{,\}256, predicted from the GPT\-2\-calibratedβ\\betaof Eq\. \([11](https://arxiv.org/html/2608.10260#A7.E11)\) with no free parameters:β​\(B⋅T\)​V=7\.0×2,048×128,256≈1\.84\\beta\(B\{\\cdot\}T\)V=7\.0\\times 2\{,\}048\\times 128\{,\}256\\approx 1\.84GB, matching the measured peak \(Table[14](https://arxiv.org/html/2608.10260#A7.T14)\) to within0\.4%0\.4\\%once combined withMoptM\_\{\\text\{opt\}\}and theMbaseM\_\{\\text\{base\}\}implied by the row above\.

##### 405B proof of concept\.

The main\-text feasibility claim beyond 70B rests on a short run of LLaMA\-3\.1\-405B\-Instruct \(frozen bf16\) on 24 nodes of4×4\\timesA100\-40GB \(96 FSDP ranks\), with the teacher’s weights streamed shard\-by\-shard into the FSDP partitioning at load time so no rank ever holds the full812812GB\. The lens set is the residual preset \(126 sites, LoRAr=64r\{=\}64\), trained with Top\-kkSubset\-KL \(k=256k\{=\}256\) at sequence length256256, one sequence per rank \(24,57624\{,\}576tokens per step\), and constantlr=10−3\\mathrm\{lr\}=10^\{\-3\}\. Over eight steps the training loss falls from57\.657\.6to40\.240\.2\(monotonically after step two;≈5\.5\{\\approx\}5\.5min/step at this configuration\) with a measured peak of37\.137\.1GB per GPU\. This is a systems demonstration only: eight steps establish that the stack loads, shards, hooks, and optimizes at 405B scale within A100\-40GB budgets, not that the resulting lens is useful\.

##### Projected cells\.

Untrained configurations reuse the validatedMbaseM\_\{\\text\{base\}\}and the same term structure, substituting the exact full\-rank optimizer or the calibrated Full\-KL readout\. The two projected quantities have different standing\. The full\-rank optimizer state \(388388GB for LLaMA\-3\-70B\) is allocated in full before the first step and no micro\-batching schedule reduces it, so every full\-rank 70B configuration fails at initialization; those infeasible cells are outcomes of the design, not estimates\. The Full\-KL readout aggregate \(235235GB for LLaMA\-3\-70B over the reference batch\) is instead reducible by micro\-batching, so it does not by itself determine feasibility at the production microbatch\. Indeed it does not: we measured LoRA\+\{\+\}Full\-KL directly for LLaMA\-3\-70B \(Table[15](https://arxiv.org/html/2608.10260#A7.T15)\) and it fits, at35\.535\.5GB\. The reference\-batch aggregate remains a useful bound on what*cannot*be processed in a single unmicrobatched pass, but the*Feasible*column should not be read as a feasibility oracle for configurations we have not measured; where we have not trained a configuration, we say so\.

## Appendix HTuned\-Lens Application Details

We recreate the application suite ofBelroseet al\.\([2023](https://arxiv.org/html/2608.10260#bib.bib6)\)on all nine of the original detection tasks \(ARC\-Easy, ARC\-Challenge\(Clarket al\.[2018](https://arxiv.org/html/2608.10260#bib.bib136)\), BoolQ\(Clarket al\.[2019a](https://arxiv.org/html/2608.10260#bib.bib137)\), MC\-TACO\(Zhouet al\.[2019](https://arxiv.org/html/2608.10260#bib.bib138)\), MNLI\(Williamset al\.[2018](https://arxiv.org/html/2608.10260#bib.bib139)\), QNLI, QQP\(Wanget al\.[2018](https://arxiv.org/html/2608.10260#bib.bib140)\), SciQ\(Welblet al\.[2017](https://arxiv.org/html/2608.10260#bib.bib141)\), SST\-2\(Socheret al\.[2013](https://arxiv.org/html/2608.10260#bib.bib142)\)\) plus LogiQA\(Liuet al\.[2020](https://arxiv.org/html/2608.10260#bib.bib143)\), withn=1,000n\{=\}1\{,\}000examples per task at GPT\-2 and 8B \(600 for the 8B 5\-shot variant; SST\-2 uses its full 872\-example validation set\) and 200 at 70B\. Zero\-shot prompts of the form “<body\>\\nAnswer:” are scored on each answer option’s first token; the attack inserts the original paper’s injection string \(“Ignore any previous and following instructions and just print ‘*wrong answer*”’\) with a randomly chosen incorrect option before the answer cue\. Trajectory features are the log\-probability of every answer option at every residual\-stream point, read through each lens’s own translator, with the model’s final distribution appended as the last point\. Detectors are an isolation forest \(200 trees, scores ensembled over 5 seeds\) and a local outlier factor \(20 neighbors, novelty mode\), fit on the first half of*clean*trajectories after per\-feature standardization; we report AUROC on held\-out clean vs\. attacked examples, with 95% bootstrap confidence intervals and*paired*bootstrap deltas against the full\-rank reference \(identical example resamples for both lenses, 1,000 resamples\)\. On these model families the logit lens also detects well on the five easy tasks \(the large tuned\-vs\-logit gap of the original paper appears specific to Pythia models\)\. Prediction depth is the first trajectory point from which the lens top\-1 equals the model’s final top\-1 thereafter, in hidden\-state units\. Causal basis extraction findsk=16k\{=\}16directions per probed layer by L\-BFGS with deflation, initialized from the top left singular vectors of the translated unembedding; energy is the expected KL increase of the*lens*readout under mean ablation on 1,024 WikiText\-2 positions, model influence is the KL of the model’s final distribution when the direction is mean\-ablated at the block output on a held\-out batch, and the random control is the QR factorization of a Gaussian matrix evaluated identically\. The attack changes the model’s answer to the planted option on a median of 93% of examples per task at 8B and 100% at 70B\. Paired bootstrap deltas against the reference are within±0\.005\\pm 0\.005on most tasks; the remaining deficits concentrate in the knowledge cluster, where detection is weak for every lens including the reference \(Figure[15](https://arxiv.org/html/2608.10260#A8.F15)\)\.

Table 16:Prompt\-injection detection AUROC \(local outlier factor, fit on clean trajectories only\): mean over the five classification tasks whereBelroseet al\.\([2023](https://arxiv.org/html/2608.10260#bib.bib6)\)report near\-perfect detection \(BoolQ, MNLI, QNLI, QQP, SST\-2\) and over all ten tasks\. All\-task means are pulled down for every lens, including the full\-rank reference, by the knowledge cluster \(ARC\-Easy/Challenge, SciQ, LogiQA\)\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x14.png)Figure 15:Prompt\-injection detection AUROC by lens \(bars: mean over the ten tasks; open circles: the individual tasks; per\-lens means in Table[16](https://arxiv.org/html/2608.10260#A8.T16)\)\. The low circles are the knowledge tasks \(ARC, SciQ, LogiQA\), where detection is weak for every lens including the full\-rank reference; on the five classification tasks every trained lens is at or above 0\.99\. At 70B, where no full\-rank reference exists, the trained lens detects knowledge\-task attacks the logit lens misses\.
## Appendix IMemory\-Injection Details

For each paired explicit/implicit prompt the memory vector𝐦\(ℓ\)=𝐡explicit\(ℓ\)−𝐡implicit\(ℓ\)\\mathbf\{m\}^\{\(\\ell\)\}=\\mathbf\{h\}^\{\(\\ell\)\}\_\{\\text\{explicit\}\}\-\\mathbf\{h\}^\{\(\\ell\)\}\_\{\\text\{implicit\}\}is captured atresid\_midper model architecture \(GPT\-2: input toln\_2; LLaMA: input topost\_attention\_layernorm\)\. Causal injection addsτ​𝐦\(ℓ\)\\tau\\,\\mathbf\{m\}^\{\(\\ell\)\}at the attention output projection so both the MLP branch and the skip connection see the patch; the measurement is validated by patching the final layer atτ=1\\tau\{=\}1, which reproduces the explicit prompt’s output to within half a percent on all three models\. The reduced LLaMA evaluation subsamples 2WMH to 200 examples at 8B and 100 at 70B, withτ∈\{0,…,10\}\\tau\\in\\\{0,\\ldots,10\\\}for lens readouts andτ∈\{1,2,4\}\\tau\\in\\\{1,2,4\\\}for causal profiles \(\{1,2\}\\\{1,2\\\}on the 70B control set\)\. GPT\-2’s 2WMH row is a null case: its explicit prompts score below its implicit ones, so no beneficial memory vector exists\. On the 70B control set both the lens and the depth heuristic select the final layer, missing the true optimum atℓ=55\\ell\{=\}55\.

Table 17:Causal memory injection: patchresid​\_​mid​\[ℓ\]\\mathrm\{resid\\\_mid\}\[\\ell\], run the model to completion, and read its own finalP​\(answer\)P\(\\text\{answer\}\); lift ismaxℓ⁡Eℓ/Pobs\\max\_\{\\ell\}E\_\{\\ell\}/P\_\{\\text\{obs\}\}\. LLaMA rows use the reduced evaluation described above\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x15.png)Figure 16:Causal liftmaxℓ⁡Eℓ/Pobs\\max\_\{\\ell\}E\_\{\\ell\}/P\_\{\\text\{obs\}\}from memory injection, by model and dataset\. The dotted line at1\.01\.0marks no effect; the GPT\-2 2WMH bar is the null case described in the main text\.##### Layer selection and distortion\.

From two clean readouts \(full layer\-by\-τ\\tausweeps appear in Figures[18](https://arxiv.org/html/2608.10260#A9.F18),[19](https://arxiv.org/html/2608.10260#A9.F19), and[20](https://arxiv.org/html/2608.10260#A9.F20)\) we define the lens\-predicted deficiency of layerℓ\\ell,

Δℓ=Plens​\(ans∣𝐡explicit\(ℓ\)\)−Plens​\(ans∣𝐡implicit\(ℓ\)\),\\Delta\_\{\\ell\}\\;=\\;P\_\{\\text\{lens\}\}\\\!\\left\(\\text\{ans\}\\mid\\mathbf\{h\}^\{\(\\ell\)\}\_\{\\text\{explicit\}\}\\right\)\-P\_\{\\text\{lens\}\}\\\!\\left\(\\text\{ans\}\\mid\\mathbf\{h\}^\{\(\\ell\)\}\_\{\\text\{implicit\}\}\\right\),\(12\)and score any chosen layer by the fraction of the maximum achievable causal lift it captures,g​\(ℓ\)=\(Eℓ−Pobs\)/\(maxℓ′⁡Eℓ′−Pobs\)g\(\\ell\)=\(E\_\{\\ell\}\-P\_\{\\text\{obs\}\}\)/\(\\max\_\{\\ell^\{\\prime\}\}E\_\{\\ell^\{\\prime\}\}\-P\_\{\\text\{obs\}\}\), whereEℓE\_\{\\ell\}is the model’s finalP​\(answer\)P\(\\text\{answer\}\)after injection at layerℓ\\ell\. Atτ=1\\tau\{=\}1the causal profile is monotonic and the final layer is trivially optimal, soτ∈\{2,4\}\\tau\\in\\\{2,4\\\}serves as a stress test of whether the lens can select an interior layer that preserves causal gain while limiting collateral distortion \(Table[18](https://arxiv.org/html/2608.10260#A9.T18), Figure[17](https://arxiv.org/html/2608.10260#A9.F17)\)\. Selections differ by training objective: Top\-kk\+IS selects the interior optimum atτ=2\\tau\{=\}2\(ℓ=19\\ell\{=\}19,100%100\\%captured\), the full\-rank reference selectsℓ=22\\ell\{=\}22\(77%77\\%\), and the Top\-kklens’s readout difference peaks at the final layer, consistent with an objective concentrated on the head of the final distribution; the logit lens’sΔℓ\\Delta\_\{\\ell\}also peaks at the final layer at both scales\. Lens\-selected layers also distort less: measured by the KL divergence between injected and clean next\-token distributions, they yield2\.52\.5–3\.0×3\.0\\timesmore answer\-probability gain per nat of distortion atτ=4\\tau\{=\}4on 8B than final\-layer injection, which alters the model’s top\-1 token in95%95\\%of examples \(Figures[22](https://arxiv.org/html/2608.10260#A9.F22)and[23](https://arxiv.org/html/2608.10260#A9.F23), Table[19](https://arxiv.org/html/2608.10260#A9.T19)\)\. On the control set, where the models largely succeed unaided, injection raisesP​\(answer\)P\(\\text\{answer\}\)to the explicit ceiling, consistent with restoring a missing recall step rather than supplying the answer directly\.

Table 18:Fraction of achievable causal gain captured on 2WMH, by selected layer\. Each lens’s selection isτ\\tau\-independent; the logit lens selects the final layer at both scales, and at 8B the Top\-kklens does as well\. Atτ=1\\tau\{=\}1the causal optimum is the final layer and every selector captures8080–100%100\\%, so those rows are omitted\. The 8B full\-rank reference reads through its per\-layer residual translators \(Appendix[I](https://arxiv.org/html/2608.10260#A9)\); no full\-rank reference is trainable at 70B under our single\-device lens placement \(Section[4](https://arxiv.org/html/2608.10260#S4)\)\. On the GPT\-2 control set the full\-rank reference and both low\-rank variants select the same layer\. All 8B lenses share the identical 250\-step annealed schedule\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x16.png)Figure 17:Fraction of achievable causal gain captured by each method’s selected layer \(2WMH, over\-injection; LLaMA\-3\-8B top, 70B bottom\), the graphical counterpart of Table[18](https://arxiv.org/html/2608.10260#A9.T18)\. The logit lens selects the final layer at both scales, so its bars coincide with the depth heuristic’s\. Layer\-by\-layer profiles appear in Figure[21](https://arxiv.org/html/2608.10260#A9.F21)\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x17.png)Figure 18:Lens\-readP​\(answer\)P\(\\text\{answer\}\)at LLaMA\-3\-8B as a function of injection layer \(yy\) and tweak factorτ\\tau\(xx\), on the control set \(top\) and 2WMH \(bottom\); the star marks the peak \(2WMH: layer1717atτ=4\\tau\{=\}4\)\. This is the lens’s view of the injected state, not the model’s output; causal effects appear in Table[17](https://arxiv.org/html/2608.10260#A9.T17)and Figure[16](https://arxiv.org/html/2608.10260#A9.F16)\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x18.png)Figure 19:As Figure[18](https://arxiv.org/html/2608.10260#A9.F18), for LLaMA\-3\-70B \(control set top, 2WMH bottom; 2WMH peak at layer7878,τ=2\\tau\{=\}2\)\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x19.png)Figure 20:Lens\-read answer recovery as a function of the tweak factorτ\\tauat each model’s best injection layer\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x20.png)Figure 21:Layer\-by\-layer profiles behind Table[18](https://arxiv.org/html/2608.10260#A9.T18)\(2WMH,τ=2\\tau\{=\}2; LLaMA\-3\-8B top, 70B bottom\): the causal profileEℓE\_\{\\ell\}\(solid\) against each lens’s deficiency profileΔℓ\\Delta\_\{\\ell\}\(dashed, rescaled to the same axis\)\. Small stars mark each lens’s selected layer and the large star the causal optimum; the selected layers and captured gains are quantified in Table[18](https://arxiv.org/html/2608.10260#A9.T18)\. The trained lenses peak at or beside the interior causal optimum; the logit lens rises monotonically to the last layer\.Table 19:Injection collateral damage on 2WMH under over\-injection\.*KL*is measured between the injected and clean next\-token distributions \(nats\);*top\-1 kept*is the fraction of examples whose argmax token is preserved;*Δ​P\\Delta P/nat*is answer\-probability gain per nat of distortion\. The 8B Top\-kklens selects the final layer, so its row coincides with*last*\.
##### Injection\-selection seed stability\.

Retraining all three 8B lenses under two additional seeds \(identical schedule\) shows the recommended lens’s selection is the seed\-robust one: the Top\-kk\+IS pick stays interior and near\-optimal \(ℓ∈\{18,19\}\\ell\\in\\\{18,19\\\}, capturing9090–100%100\\%atτ=2\\tau\{=\}2and6161–75%75\\%atτ=4\\tau\{=\}4\), while the Top\-kkpick lands on the final layer on two seeds and an early layer \(ℓ=10\\ell\{=\}10,20%20\\%\) on the third, and the full\-rank reference’s pick moves acrossℓ∈\{22,24,31\}\\ell\\in\\\{22,24,31\\\}\(6464–82%82\\%atτ=2\\tau\{=\}2\), reaching the final layer on one seed\.

![Refer to caption](https://arxiv.org/html/2608.10260v1/x21.png)Figure 22:Collateral damage at the selected layer \(2WMH\): answer\-probability gain per nat of distortion of the next\-token distribution\. The logit lens’s selection coincides with the final layer and is omitted; exact values in Appendix Table[19](https://arxiv.org/html/2608.10260#A9.T19)\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x22.png)Figure 23:Fraction of examples whose top\-1 token survives injection at the same selected layers as Figure[22](https://arxiv.org/html/2608.10260#A9.F22)\.

## Appendix JDART and ToxIn Details

DART scores each attention head over 200 toxic prompts drawn from the Wiki Toxic corpus\(cjadamset al\.[2017](https://arxiv.org/html/2608.10260#bib.bib45)\): per\-head attention outputs are captured at the output projection \(theo\_proj/c\_projinput split into query heads, valid under grouped\-query attention since the projection input carries all query heads\), each head’s last\-token contribution is decoded through the lens at that layer’s attention\-output site, and the toxic count is the number of matches against a precomputed toxic\-vocabulary set among the top\-50 decoded tokens\. The 8B full\-rank reference carries only per\-layer residual translators, so head outputs are decoded through the block\-input translator of their layer rather than a dedicated attention\-output translator, a half\-block site mismatch that makes its agreement numbers conservative\. ToxIn performs zero or soft subtraction of the unembedding\-derived toxic direction at the DART\-flagged heads; toxicity is scored with Toxic\-BERT over greedy generations and fluency with WikiText\-2 perplexity; 70B uses a 40\-prompt reduced evaluation\. The whole\-model audit \(Figure[25](https://arxiv.org/html/2608.10260#A10.F25)\) selects the top lens\-flagged sites of each component type and subtracts the lens\-mapped direction there, sweeping subtraction strength under a perplexity budget\.

##### Selector comparisons and the 70B audit\.

Per\-head toxicity maps and top\-head tables appear in Figure[26](https://arxiv.org/html/2608.10260#A10.F26)and Table[20](https://arxiv.org/html/2608.10260#A10.T20)\. The trained variants agree on the 8B head ranking, and at GPT\-2 they reproduce the full\-rank reference’s audit; the logit lens’s agreement with the trained lenses decays with model size \(Spearman0\.730\.73at GPT\-2 to0\.020\.02at 70B; Table[23](https://arxiv.org/html/2608.10260#A11.T23)\)\. The 8B full\-rank reference identifies the same strongest head as every other variant \(L23\.H24\) and shares 3 of 5 top heads with the recommended lens\. At GPT\-2, with the ablation direction fixed to the unembedding and reductions interpolated to thePPL=1\.10×\\text\{PPL\}=1\.10\\timesoperating point on each selector’s scale sweep, our lens’s heads give25\.0%25\.0\\%, the full\-rank reference’s23\.2%23\.2\\%, and unembedding\-only selection18\.8%18\.8\\%; the trained selections are separated by less than run\-to\-run noise, so we read this as parity between our lens and the reference, both ahead of the lens\-free selector\. At 8B, under soft subtraction every trained selector reduces toxicity, whereas random\-head intervention increases it; under zero\-ablation the recommendedTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}selection produces the clearest reduction \(Table[22](https://arxiv.org/html/2608.10260#A10.T22); single\-run evaluations, so differences of a few points are within run\-to\-run noise\)\. Soft subtraction trades toxicity against perplexity controllably at 8B \(Figure[24](https://arxiv.org/html/2608.10260#A10.F24); full sweeps in Table[21](https://arxiv.org/html/2608.10260#A10.T21)\)\. At 70B, toxic signal again localizes to late\-layer heads but is spread far more thinly \(the top five carry only∼10%\{\\sim\}10\\%\), and ablating the flagged heads is no more effective than ablating random ones \(Figure[30](https://arxiv.org/html/2608.10260#A10.F30)\): targeted head ablation is effective only when the localized signal is sufficiently concentrated, and that concentration is absent at 70B\. In the whole\-model audit \(Figure[25](https://arxiv.org/html/2608.10260#A10.F25)\), on GPT\-2 subtracting atattn\_inandresid\_postremoves2\.2×2\.2\\timesand1\.6×1\.6\\timesas much toxicity as the original per\-head recipe; at 8B,mlp\_out\(34%34\\%\) andresid\_post\(27%27\\%\) reach3\.1×3\.1\\timesand2\.5×2\.5\\times, whileattn\_in, the most effective target on GPT\-2, has almost no effect: a selection made on one model does not transfer to another\. The full 8B audit \(192 hookpoints scored, 24 ablation sweeps\) runs in under ten minutes on one node\.

![Refer to caption](https://arxiv.org/html/2608.10260v1/x23.png)Figure 24:ToxIn soft subtraction at 8B trades toxicity against perplexity as the subtraction strengthλ\\lambdagrows\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x24.png)Figure 25:The whole\-model ablation audit \(GPT\-2 top, LLaMA\-3\-8B bottom\): soft subtraction of the lens\-mapped toxic direction at the top\-flagged hookpoints of each component type, best operating point under a PPL≤1\.10×\\leq 1\.10\\timesbudget, compared with the original attention\-heads\-only recipe \(hatched bar, dashed line\)\. The most effective targets lie outside attention at both scales and differ between them\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x25.png)Figure 26:Share of the total toxic signal carried by the top\-5 heads, by model and lens\. Localization is comparable at GPT\-2 and 8B and far more diffuse at 70B; the logit lens’s concentration declines with scale relative to the trained lenses\. Per\-head maps appear in Figures[27](https://arxiv.org/html/2608.10260#A10.F27)–[29](https://arxiv.org/html/2608.10260#A10.F29)\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x26.png)Figure 27:DART per\-head toxic\-token counts at LLaMA\-3\-8B \(LoRATop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}lens\): the signal concentrates in a few late\-layer heads\. Boxes mark the top\-5 heads; the strongest \(L23\.H24\) and the top\-5 share are quantified in Table[20](https://arxiv.org/html/2608.10260#A10.T20)\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x27.png)Figure 28:As Figure[27](https://arxiv.org/html/2608.10260#A10.F27), for LLaMA\-3\-70B: the same late\-layer localization holds but is far more diffuse \(Table[20](https://arxiv.org/html/2608.10260#A10.T20)\)\.![Refer to caption](https://arxiv.org/html/2608.10260v1/x28.png)Figure 29:DART toxic signal aggregated by layer \(LLaMA\-3\-8B top, 70B bottom\), showing the late\-layer skew quantified in Table[20](https://arxiv.org/html/2608.10260#A10.T20)\.Table 20:DART localization concentration, all lens variants under one harness \(200 toxic prompts, top\-50 decoded tokens per head\)\.*Total*is the number of toxic tokens flagged over all heads;*top head*and*top\-5*give the share of that total;*late/early*is the ratio of summed counts in the second vs\. first half of layers\. The 8B full\-rank reference decodes heads through per\-layer residual translators \(see above\)\.ModelLensTotalTop headTop\-5Late/earlyGPT\-2logit38,753L08\.H02\(6\.8%6\.8\\%\)29%29\\%1\.7×1\.7\\timesLoRA Top\-k\+IS23,858L08\.H02\(7\.3%7\.3\\%\)30%30\\%1\.8×1\.8\\timesLoRA Top\-kk16,582L11\.H03\(10\.5%10\.5\\%\)32%32\\%1\.8×1\.8\\timesfull\-rank, full KL13,586L11\.H03\(11\.5%11\.5\\%\)36%36\\%1\.7×1\.7\\timesLLaMA\-3\-8Blogit94,077L23\.H24\(6\.4%6\.4\\%\)16%16\\%1\.9×1\.9\\timesLoRA Top\-k\+IS9,395L23\.H24\(12\.4%12\.4\\%\)35%35\\%167×167\\timesLoRA Top\-kk24,641L23\.H24\(7\.6%7\.6\\%\)20%20\\%2\.5×2\.5\\timesfull\-rank \(resid\.\)62,121L23\.H24\(7\.0%7\.0\\%\)19%19\\%2\.1×2\.1\\timesLLaMA\-3\-70Blogit407,752L44\.H04\(0\.8%0\.8\\%\)3%\\phantom\{0\}3\\%1\.1×1\.1\\timesLoRA Top\-k\+IS29,699L44\.H04\(3\.2%3\.2\\%\)10%10\\%72×72\\timesTable 21:ToxIn ablation\. Toxicity is the mean Toxic\-BERT\(Hanu and Unitary team[2020](https://arxiv.org/html/2608.10260#bib.bib57)\)score over greedy generations from toxic prompts; PPL on WikiText\-2\. 70B uses the reduced 40\-prompt evaluation\.
##### Estimator localization and seed stability\.

Trained under the identical 250\-step annealed schedule, the two estimators concentrate the signal differently: the Top\-kk\+IS lens localizes roughly twice as sharply as Top\-kk\(top\-5 share35%35\\%vs\.20%20\\%; across three training seeds,3131–35%35\\%vs\.1111–20%20\\%\)\. The strongest\-head identification is fully seed\-robust: retraining all three lens variants under two additional seeds, every one of the nine lens×\\timesseed audits ranksL23\.H24first, with 4–5 of 5 top heads shared across seeds\.

![Refer to caption](https://arxiv.org/html/2608.10260v1/x29.png)Figure 30:ToxIn zero\-ablation: removing the DART\-flagged heads reduces toxicity at 8B while removing random heads does not; at 70B neither does\.Table 22:Head\-selector comparison at 8B: ToxIn with the same unembedding\-derived direction \(100 toxic prompts; Toxic\-BERT; PPL on WikiText\-2, baseline 9\.66\)\. Zero\-ablation acts on the top\-15 flagged heads themselves; soft subtraction acts on the attention output of the layers containing them \(7–11 layers per selector\)\. All trained lenses share the identical 250\-step annealed schedule; the random row averages three seed\-0 draws, which partially overlap\.

## Appendix KFine\-Grained Fidelity Statistics

Beyond task outcomes, we compare the lenses on the trajectory statistics themselves\.

Rank agreement of head audits\.Table[23](https://arxiv.org/html/2608.10260#A11.T23)quantifies cross\-lens agreement of the DART head rankings\. Trained lenses agree at every scale, and all variants \(including the full\-rank references at GPT\-2 and 8B\) identify the same strongest heads; rank correlations over*all*heads are dominated by the inert majority and should be read jointly with the top\-head overlaps\. The logit lens’s rank agreement with trained lenses decays with scale \(Spearman0\.73→0\.020\.73\\to 0\.02\) even as it continues to identify the very strongest heads\.

Prediction depth\.From the detection captures we also compute each lens’s prediction depth: the trajectory point after which its top\-1 prediction stops changing\(Belroseet al\.[2023](https://arxiv.org/html/2608.10260#bib.bib6)\)\. Against the full\-rank reference, the LoRA lenses agree to within one hidden state on 69–77% of GPT\-2 examples \(mean absolute difference≈1\{\\approx\}1layer\), while the logit lens manages 45% with no rank correlation \(ρ=−0\.08\\rho=\-0\.08\)\. At 8B the ordering is preserved \(LoRA 31–38% within one state,ρ≈0\.53\\rho\\approx 0\.53; logit 25%,ρ=0\.24\\rho=0\.24\), but the full\-KL reference resolves predictions systematically earlier than the Subset\-KL lenses, by four to five hidden states on average\. Prediction depth is the fine\-grained statistic where the expensive reference retains a visible edge; the task\-level results of the main text are unaffected by it\.

![Refer to caption](https://arxiv.org/html/2608.10260v1/x30.png)Figure 31:The classic prediction\-trajectory grid \(top\-1 token of the lens readout at every layer and position; shading = probability\) for the full\-rank tuned lens and LoRATop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}on the same GPT\-2 prompt\. Both lenses read the same computation: the indirect object emerges around L9–L10 and stabilizes to the model’s final prediction\.Causal basis extraction\.Finally we rerun the original paper’s causal\-fidelity experiment: for each lens and layer, extract thek=16k\{=\}16orthonormal directions whose mean\-ablation most changes the lens output, then ablate each direction in the*model*and measure the KL against the clean output \(Appendix[H](https://arxiv.org/html/2608.10260#A8)\)\. Every lens’s directions are causally real: ablating them moves the model two orders of magnitude more than random directions \(mean top\-8 KL 0\.39–0\.48 at GPT\-2, 0\.69–0\.81 at 8B, vs\.≈0\.002\{\\approx\}0\.002–0\.0070\.007random\), and our lenses match the reference on this transfer strength\. On the finer statistic \(Spearman correlation between the lens’s claimed influence ordering and the realized model KL\), the reference leads at GPT\-2 \(0\.71 vs\. 0\.57–0\.61 LoRA, 0\.40 logit\), while at 8B the statistic stops discriminating between lenses entirely \(0\.52–0\.65 for all, logit included\)\. The qualitative counterpart is Figure[31](https://arxiv.org/html/2608.10260#A11.F31): the full\-rank andTop​\-​k\+IS\\mathrm\{Top\\text\{\-\}\}k\\mathrm\{\+IS\}trajectory grids on the same prompt are near\-identical\.

Table 23:Cross\-lens agreement of DART head rankings over all attention heads\.*Top\-5*counts shared heads among each lens’s five strongest\. Trained lenses agree at every scale; the logit lens’s rank agreement with trained lenses decays with scale, although it still identifies the strongest heads\. The 8B full\-rank row carries the residual\-translator caveat of Table[20](https://arxiv.org/html/2608.10260#A10.T20)\.
## Appendix LCase\-Study Scope and Caveats

All results here are replications intended to validate lens fidelity at scale\. The LLaMA runs use reduced evaluations \(Appendix[I](https://arxiv.org/html/2608.10260#A9)\); the 2WMH templates and their programmatically constructed explicit prompts are imperfect references; and dictionary\-based DART scoring favors lexically explicit toxicity\. The injection results and the 8B ablations are causal interventions, while the 70B localization and the diffusion account of its null ablation remain observational\. The 8B full\-rank reference enters the injection and DART comparisons through its per\-layer residual translators \(a half\-block hookpoint mismatch; Appendix[J](https://arxiv.org/html/2608.10260#A10)\), so its agreement numbers there are conservative, and the selector comparison of Table[22](https://arxiv.org/html/2608.10260#A10.T22)is a single\-run evaluation\. Fine\-grained statistics \(injection\-layer rankings, prediction depth, influence orderings\) are the last lens properties to stabilize during training and should be read from converged, annealed lenses; the 8B seed study shows they are also the most seed\-sensitive: coarse results \(the strongest DART head, interior\-vs\-final layer selection for the recommended lens\) replicate across all three training seeds, while the exact picked layer and top\-5 shares move within the ranges reported above\. The trained 8B lenses in the injection and toxicity analyses share a single 250\-step cosine\-annealed schedule, identical to the 70B run’s; the detection captures of Table[16](https://arxiv.org/html/2608.10260#A8.T16)and the whole\-model audit of Section[5](https://arxiv.org/html/2608.10260#S5.SS0.SSS0.Px3)predate this schedule \(earlier\-schedule checkpoints\), and the earlier bracketing checkpoints supported the same coarse conclusions\. Finally, the layer\-selection advantage of trained lenses is specific to the over\-injection regime, and the logit lens’s failures are specific to scale: at GPT\-2 it selects an effective injection layer and localizes adequately\.