Interpreting Language Model Hidden States at Scale

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.

OmniLens addresses three challenges: Parameters. A full-rank translator costs O(d2) per hookpoint, where d is the model’s hidden dimension, so dense coverage of a large model can rival the model itself. 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. Specialization. Existing lenses each read one hookpoint type, and each new hookpoint requires a new lens family.

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 O(d2) to O(rd) where r is 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.

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.

Low-rank translator. For an activation hl,u ∈ Rd, OmniLens learns a translator (Ll,u, bl,u), where Ll,u = I + (α/r) Bl,u Al,u, with Al,u ∈ Rr×d, and Bl,u ∈ Rd×r, with fixed scale α, so that Ll,u is a rank-at-most-r update to the identity. Each translator contains 2dr + d learned parameters, against d2 + d for a dense affine translator. For d ≫ r the ratio of learned parameters per hookpoint is approximately 2r/d, so the savings compound as coverage grows: at d = 4096 and r = 64, six low-rank translators together contain only about 19% as many parameters as a single dense translator.

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, so component types can be compared without designing and training a separate lens family for each.

Rank ablation. Fidelity improves rapidly through r = 64 and only modestly thereafter. At r = 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) and mean ρ = 0.941; we therefore use r = 64 as the default. Earlier layers are the hardest to approximate (layer 0: 56.7% top-1 agreement at r = 64), and their KL to the teacher remains highest even for the full-rank lens. So, analyses focused on early layers should prefer r ≥ 64.

Low rank is a constraint, not a compression. The deviation from identity learned by the full-rank baseline, Ll,u − I, is not itself low rank: its best rank-64 approximation retains only 59% of its Frobenius energy on GPT-2 Small, and its numerical rank remains near d throughout training. Post-hoc compression accordingly loses fidelity: truncating the full-rank translator to rank 64 raises final-layer KL by 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 within 7% of the full-rank lens.

Top-k Truncation. Let H contain the khead most probable teacher tokens. Top-k restricts both distributions to H and renormalizes them. Top-k asks the lens to reproduce the teacher’s relative preferences among its khead 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. On GPT-2 Small, where full-KL training remains feasible and supplies a reference: on its residual hookset, khead = 256 reduces peak memory from 16.3 to 4.7 GB and increases throughput by 1.59×, while changing final-layer KL from 0.043 to 0.054.

Top-k+IS: Exact Head, Sampled Tail. To ameliorate the poor performance of Top-k on early layers, we introduce Top-k+IS, which keeps the head unnormalized and samples the tail. The first line is the exact contribution of the head H to the original, untruncated KL. The second estimates the omitted tail sum from sampled tokens, each reweighted by P/R so 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 states that the Top-k+IS estimator is unbiased and yields unbiased gradients.

Choosing an objective. Neither variant dominates. Top-k is preferable under projection cost, throughput, or maximum context length constraints, particularly for analyses that emphasize later layers. Top-k+IS is preferable when lenses are compared throughout the network, or when preserving the semantics of the full-KL objective matters. We therefore default to Top-k+IS for the trained lenses in the interpretability studies, whose deterministic head is what preserves late-layer fidelity, and report Top-k as the lower-memory, higher-throughput alternative.

A 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 the B·T positions in a microbatch. A naïve implementation gathers those rows of the unembedding WU into a [B·T, k, d] intermediate before multiplying by the translated activation, potentially consuming more memory than the subset objective saves. At B·T = 8192, k = 512, and d = 4096, this tensor alone occupies 68.7 GB in FP32. Our fused CUDA kernel computes each selected logit directly in registers and never forms the gathered tensor, reducing the peak memory to 0.15 GB.

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.

Frontier scale. We execute eight optimizer steps of a 126-hookpoint residual-lens stack on LLaMA-3.1-405B-Instruct across 96×A100-40GB GPUs. At r=64 that stack contains 266M translator parameters, a roughly 127-fold reduction from the corresponding 33.8B-parameter dense stack before optimizer state. This run verifies execution of the complete training path at 405B scale.

Prompt-injection detection. Following the protocol of Belrose et al. (2023), 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 attains a mean AUROC of 0.997 at every scale and matches the full-rank reference within bootstrap uncertainty where that reference exists. At 70B, where no full-rank reference is available, Top-k+IS detects attacks the logit lens misses on the knowledge tasks (ARC-Easy 0.80 vs. 0.68; SciQ 0.67 vs. 0.56).

Multi-hop factual reasoning. The intervention replicates at scale: on 2WikiMultiHop (2WMH) it raises the model’s own P (answer) by 6.6× at 8B and 5.1× at 70B. At a deliberately strong intervention scale (τ =4), where the causal optimum moves to an interior layer, the Top-k+IS-selected layer captures 61% (8B) and 82% (70B) of the maximum achievable gain, compared with 34% and 66% for final-layer injection. Top-k+IS is the only evaluated lens that selects a middle layer and remains near-optimal across all three training seeds.

Toxicity localization and intervention. 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. At 8B, detection and intervention rankings are negatively correlated (Spearman −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.

Improvements to AI systems based on OmniLens:

1. Model-wide interpretability at scale: AI systems can now be analyzed across hundreds of internal hookpoints (residual, attention, and MLP layers) simultaneously, even for models with 70B–405B parameters. This enables system-wide debugging of behaviors—such as factual recall, prompt-injection vulnerability, or toxicity—rather than limiting analysis to a single component type.

2. Low-rank translator architecture: By replacing dense per-layer translators (O(d2) parameters) with low-rank updates (O(rd), e.g., r=64), AI systems can train interpretability lenses with 90% fewer parameters and up to 70% lower peak memory. This makes dense lens coverage practical on existing hardware, allowing real-time monitoring of large models during inference or fine-tuning.

3. Subset-KL training objectives (Top-k and Top-k+IS): AI systems can now train lenses using memory-efficient approximations of KL divergence, reducing peak memory from 16.3 GB to 4.7 GB on GPT-2 while increasing throughput 1.59×. The Top-k+IS variant provides unbiased gradients, preserving full-KL semantics—critical for comparing lens behavior across all layers without sacrificing fidelity.

4. Fused CUDA kernel for selected logits: This eliminates the memory bottleneck of materializing gathered unembedding rows, reducing peak memory from 68.7 GB to 0.15 GB for a typical microbatch. AI systems can now process longer contexts (up to 4K tokens) and larger vocabularies on a single A100-40GB, enabling frontier-scale lens training without specialized hardware.

5. Hookpoint-agnostic attachment: One unified lens family works across residual, attention, and MLP activations, allowing AI systems to compare component types directly. This reveals that the hookpoint where a behavior is most visible (e.g., detected toxicity) often differs from where intervention is most effective (e.g., reducing toxicity)—enabling more precise causal control of model outputs.

6. Improved intervention targeting: Using OmniLens, AI systems can identify optimal intervention layers for behaviors like multi-hop reasoning. At 70B, the lens-selected middle layer captures 82% of the maximum achievable gain (vs. 66% for final-layer injection), enabling more effective steering of model reasoning with fewer perturbations.

7. Scalable safety monitoring: The prompt-injection detector trained on OmniLens lenses achieves 0.997 mean AUROC across five tasks at all scales, and at 70B it detects attacks that the logit lens misses (e.g., ARC-Easy 0.80 vs. 0.68). This allows AI systems to deploy real-time, layer-wise anomaly detection for adversarial inputs in production.

8. Toxicity localization with causal validation: OmniLens reveals that MLP outputs—not attention heads—are the most causally effective targets for toxicity reduction, despite showing weak detection signals. This enables AI systems to implement targeted interventions (e.g., modifying MLP activations) that reduce harmful outputs more effectively than prior attention-based methods.

What the improved AI system can do:

• Train and deploy dense interpretability lenses on models up to 405B parameters using standard hardware.

• Monitor and debug model behavior across all internal layers in real time, with memory and compute costs that scale sub-linearly with model size.

• Compare visibility vs. causality of behaviors across component types, enabling precise, minimal interventions for steering, safety, and alignment.

• Detect adversarial inputs and localize toxic or biased reasoning with higher accuracy and at larger scales than previously possible.

• Run multi-hop reasoning interventions that outperform final-layer methods, improving factual accuracy and controllability.