Lexical Reference Frames: Interpreting LM Residual Streams in Native Semantics

AI summary

Overview

Research area: Mechanistic interpretability of large language models — specifically, how to read (decode) the residual stream at intermediate layers into the model's own vocabulary.

Technical level: Advanced. The paper leans on Jacobians of the remaining layers, pseudo-inverses of the unembedding matrix, coherence bounds, and closed-form ridge regression; the empirical sections assume familiarity with the logit lens, attribution patching and steering.

Scope: The paper defines "lexical reference frames" — model-native, context-dependent bases for intermediate hidden states — proves a bound on how well any fixed lens can approximate them, and evaluates the resulting frames (LexFrame of order 0, 1, 2, ∞) on SmolLM3-3B and OLMo3-7B for reading, attribution and steering.

What This Paper Is About

A language model defines the meaning of a vector in its residual stream only at the two ends: the embedding maps tokens to states, and the unembedding maps the final state to logits. Everything in between is used constantly by early-exit decoding, hallucination detection, knowledge editing and distillation, yet is read with tools that either apply one fixed linear map per layer (the logit lens, the tuned lens, the Jacobian lens) or describe states as an unordered point cloud. The authors argue that each hidden state has a native reference frame induced by the layers that remain and by the context, and they build an approximation — the lexical reference frame — that reads a state through those remaining layers with a bounded number of rounds of attention, fitting nothing to the model's output.

Key Contributions

1. Identifying the native reference frame and bounding lenses. The paper formalises the frame induced by the layers that remain in context, defines the lexical readout of a label as the gradient of its centred logit with respect to the state, and proves that no lens with one fixed direction per label per layer can exceed a mean cosine of √κ with the native frame, where κ is the frame's coherence across contexts.
2. Defining lexical reference frames of finite order. LexFrame[n] transports a state through the model's own remaining layers with n rounds of attention between positions and a typical attention output (a per-layer ridge regression of attention output on attention input, estimated once from unlabelled text) for the rest. Order 0 withholds the prefix, order 1 lets the state attend to its prefix, order 2 adds one exchange within the prefix, and order ∞ is the native frame itself.
3. Empirical validation on two model families and their training checkpoints. Across reading, point-cloud comparison, coordinate tests, attribution and steering, the frames recover the model's behaviour from mid-depth on where lenses fail, and the paper reports how the logit lens degrades over training while LexFrame[1] stays stable.
4. A released implementation. The code is available at the GitHub repository cited in the paper.

Main Findings

• Lenses cannot follow the native frame, structurally. At mid-depth the coherence bound √κ is 0.30 for SmolLM3-3B and 0.35 for OLMo3-7B, while LexFrame[1] reaches a mean cosine of 0.97 and 0.97. LexFrame[2] comes closer still.
• The deviation is large and does not vanish near the output. In OLMo3-7B, after layer 16 of 32, the logit-lens direction deviates by 75 degrees from the model's own lexical readout and the Jacobian lens by 77 degrees; four layers before the output both still deviate by at least 41 degrees.
• Readings match the model's output. Layer-averaged divergence is 0.16 and 0.15 nats for LexFrame[1], 0.06 and 0.05 for LexFrame[2], and 0.87 and 0.93 for LexFrame[0], against 29.06 and 5.66 for the logit lens and 7.51 and 6.19 for the Jacobian lens. Top-token agreement for LexFrame[1] is 0.86 and 0.86 averaged over layers, exceeding nine tenths from layer 18 and 17 on; the lenses reach one half only from layer 35 and 30 on.
• The readings are calibrated predictions, not copies of the output. LexFrame[1]'s calibration error against the actual next token is 0.027 and 0.018, close to the model's own 0.018 and 0.023. At the embedding, where all computation remains, it agrees with the model's top token on only 0.42 and 0.39 of states.
• Readings settle early. On average, LexFrame's readings settle on the model's top token over the final 0.82 and 0.81 of the layers; those of lenses over at most 0.11.
• Point-cloud geometry predicts little. Across 7961 state-substitution pairs per layer, no distance on the cloud exceeds a rank correlation of 0.28 with the resulting change in the output, and geodesic distance performs worse than plain cosine distance. The divergence between LexFrame[1] readings of the two states correlates at 0.94 and 0.95. Lexical readouts lie almost entirely outside the leading 64 principal directions of the cloud.
• The lexical basis composes, the readout does not. Moving a state by 0.1 of its norm along the LexFrame basis of a label makes that label gain most in 0.84 and 0.97 of cases for LexFrame[1] and 0.98 and 1.00 for LexFrame[∞]; the readout succeeds in only 0.54 and 0.73 (LexFrame[1]) and 0.56 and 0.79 (LexFrame[∞]). The readout is a slope, not a direction of travel.
• Decomposition requires the transport. Inner products of unembedding rows with the raw layer state sit 8.36 and 5.64 nats from the model's logits; with one fixed map per layer, 5.19 and 5.46; with the LexFrame[1] transport, 0.16 and 0.15; with LexFrame[∞], 0.00 and 0.00.
• Selective steering. Along the basis, LexFrame[∞] flips 0.90 and 0.97 of predictions while disturbing other tokens by only 0.11 and 0.06 nats at the flip, giving clean flips in 0.73 and 0.85 of cases; the readout flips 0.95 and 0.97 but disturbs 0.22 and 0.19 nats. LexFrame[1]'s basis still yields clean flips in 0.71 and 0.82 of cases, without access to the model's own Jacobian.
• Attribution beats patching without the model's Jacobian. The difference of LexFrame[1] coefficients correlates with the measured effect at 0.88 and 0.90, against 0.62 and 0.68 for the linear attribution of LexFrame[∞] (which for that frame is attribution patching) and at most 0.20 for the lenses. The gap is concentrated in the feed-forward block (0.88 and 0.90 versus 0.63 and 0.71) and small for attention (0.79 and 0.85 versus 0.75 and 0.84).
• Lens degradation over training tracks a moving frame. Across checkpoints the logit lens's divergence grows from 4.82 and 2.32 to 29.06 and 5.66 nats, while LexFrame[1] moves only from 0.15 and 0.02 to 0.16 and 0.15.

Methodology in Plain English

The model's vocabulary is treated as its "native semantics," so interpreting a hidden state means asking which tokens it points to. At the output this is exact: multiply the state by the unembedding. In the middle it is not, because the layers still to come will transform the state and will also pull in information from earlier tokens.

The authors' recipe is to run the rest of the model forward from the state under study, rather than fitting a shortcut map. To keep this well-defined and informative, they limit how much information can flow sideways between positions. At order 0, the state is pushed through all remaining layers, but every attention operation is replaced by a "typical" attention output — a linear approximation fitted in closed form, once per layer, on unlabelled text, using only the state itself as input. At order 1, the real attention over the prefix is restored for the state being read, while the prefix positions themselves are still transported at order 0. Order 2 allows one more hop within the prefix, and so on; order ∞ restores everything and reproduces the model exactly.

Because nothing is fitted to the model's output, agreement with the model is a genuine test rather than a built-in property. The resulting frame gives each label a direction (the lexical readout, the gradient of that label's logit) and an axis (the lexical basis, the smallest state move whose first-order effect on the output comes closest to the label's semantic basis). The coefficients read the state, the basis writes into it, and differences of coefficients attribute a prediction to layers and sublayers. Experiments use WikiText-103 sequences of 128 tokens, calibrate on 256 training sequences and evaluate on the test split, comparing against the logit lens, the Jacobian lens (fitted on 100 prompts), a supervised linear probe, and a battery of point-cloud distances.

Why This Matters

The paper supplies a principled replacement for the logit lens whenever intermediate states are read, and it explains why the logit lens degrades as models train: the model's internal frames drift away from its output frame. It also gives a negative result for a popular family of analyses — point-cloud statistics such as intrinsic dimension, effective rank and anisotropy reveal little about how the model actually uses a state.

Real-world applications:

• Early exit and self-speculative decoding: a frame whose reading settles on the top token over the final 0.82 and 0.81 of layers identifies how much computation can be skipped.
• Hallucination detection and model monitoring: calibrated mid-depth readings (calibration error 0.027 and 0.018, close to the model's own 0.018 and 0.023) give detectors a usable probability, not just a direction.
• Knowledge editing and model compression: editing methods that locate and rewrite facts at intermediate layers, and distillation that matches teacher states, both depend on knowing what a state means — the frame supplies coordinates for both.
• Controlled and selective steering: the lexical basis flips predictions while leaving other tokens nearly untouched (0.11 and 0.06 nats of side effect), which matters for safety interventions and behavioural control.

Industry relevance is direct: the frames require no training, only one forward pass over calibration text plus closed-form ridge fits, and the whole project used about 500 GPU hours, of which about 100 produced the final experiments, on single NVIDIA RTX 4080 SUPER GPUs.

Future Directions

• Scaling the evaluation beyond two families. The study covers SmolLM3-3B (36 layers) and OLMo3-7B (32 layers); whether the coherence bound and the mid-depth recovery hold for much larger models, mixture-of-experts architectures and non-English corpora is untested here.
• Cheaper estimation of the lexical basis at every position. Steering currently solves the pseudo-inverse problem with 30 LSMR steps per position; reducing this cost would make basis-based interventions practical at generation scale.
• Extending the order hierarchy and the typical-attention approximation. Orders above 2 were not evaluated numerically, and the per-layer ridge regression is fitted once on unlabelled text; richer replacements for the typical attention output, and the trade-off between order and fidelity, remain open.
• Nonlinear attribution and better readouts. The paper shows that linear attribution fails most for feed-forward writes (0.88 and 0.90 versus 0.63 and 0.71), which invites methods that follow the coefficient difference without a linearisation but with less cost than a full ablation.

Target Audience

Interpretability researchers and mechanistic-interpretability practitioners who currently use the logit lens, tuned lens or Jacobian lens; engineers building early-exit, speculative decoding, hallucination detection, steering or knowledge-editing systems on top of open-weight models; and anyone who uses point-cloud statistics (intrinsic dimension, effective rank, anisotropy) to characterise hidden layers and wants to know what those statistics do and do not reveal. Readers need comfort with Jacobians, pseudo-inverses and KL divergence to follow the derivations, though the empirical comparisons in the tables are readable without them.

Authors’ abstract

The vocabulary of a language model is its native semantics: the embedding and the unembedding fix the meaning of a state at the two ends of the residual stream. Between them, existing methods either reuse a fixed lens or treat the states as a point cloud, and neither reads a state the way the model does. We argue that every hidden state has a native reference frame, induced by the layers that read it next and adaptive to the context. Lexical reference frames approximate it by reading a state through the model's own remaining layers with a stated number of rounds of attention between positions and a typical attention output for the rest, without fitting to the model's output. No lens, with fixed directions per layer, can exceed the square root of the native frame's coherence, 0.30 and 0.35 at mid-depth of SmolLM3-3B and OLMo3-7B, whereas the first-order frame, LexFrame[1], reaches 0.97 and 0.97 and reads a state within 0.16 nats of the model's output. Its coefficients attribute predictions more faithfully than attribution patching, its lexical basis steers a prediction while sparing other tokens, and its confident readings would let an early exit skip 0.45 and 0.42 of the layers, where lenses skip almost none. In this way, lexical reference frames make every hidden state of the model readable, attributable and steerable.