The Jacobian Lens

Where the Logit Lens Breaks Down

The logit lens applies the model's unembedding matrix $W_U$ directly to an intermediate residual stream $\mathbf{h}_\ell$ and reads off a vocabulary distribution. It works in late layers, where the residual stream already lives in a basis close to the final layer. In early layers, it fails: the projections come out incoherent, not because the information is missing, but because the coordinates have not yet been rotated into the basis $W_U$ expects.

The tuned lens responds by learning a per-layer linear correction to match the final-layer output distribution. This works, but it changes the character of the tool: the correction is fitted to outputs on a training set, and it optimizes a correlational objective. When we want to know what the model itself would linearly do with a given activation, a fitted map introduces its own inductive bias.

The Jacobian lens takes a different path. Instead of learning a linear map, it derives one from the model's own weights, using nothing more than calculus. The result, introduced by Gurnee, Lindsey, and collaborators in 2026 [1]Verbalizable Representations Form a Global Workspace in Language Models
Gurnee, W., Lindsey, J., et al.
Anthropic (transformer-circuits.pub), 2026, is a principled refinement of the logit lens that recovers interpretable content at depths where the raw projection cannot.

Jacobian Lens (J-lens): A per-layer vocabulary readout obtained by (1) computing the Jacobian of the final residual stream with respect to an intermediate residual stream, (2) averaging that Jacobian over token positions and a corpus of prompts, and (3) composing the resulting matrix with the unembedding. It surfaces concepts the model is poised to verbalize, whether or not they appear in the next token.

The rest of this article is built for someone who has seen the logit lens but is not yet comfortable with the words Jacobian, first-order, or linearization. We will build each of those from scratch, then assemble them into the J-lens definition.

Warm-Up: Derivatives, Jacobians, and What "Linear" Means

To build the J-lens, we need one idea from single-variable calculus and its generalization to vectors.

From one dimension...

Suppose $f: \mathbb{R} \to \mathbb{R}$ is a smooth function. Pick a linearization point $x_0$. Near it, we can approximate $f$ by its tangent line:

$$f(x_0 + \Delta x) \;\approx\; f(x_0) \;+\; f'(x_0) \cdot \Delta x.$$

The number $f'(x_0)$ is the derivative. The approximation is called first-order because it uses the first derivative; higher-order Taylor terms ($\frac{1}{2} f''(x_0) \Delta x^2$, etc.) are dropped. It says: if you nudge the input by $\Delta x$ away from $x_0$, the output moves by roughly $f'(x_0)$ times that nudge. The approximation is exact only when $f$ is a straight line; for curves, the error grows with $|\Delta x|$.

Tangent point x₀: Step size Δx:

The key point: near $x_0$, the curve looks linear. First-order means we take that local straightening seriously and use it to predict small displacements.

...to many dimensions

Now let $f: \mathbb{R}^n \to \mathbb{R}^m$. The input is a vector $\mathbf{x} \in \mathbb{R}^n$; the output is a vector $\mathbf{y} \in \mathbb{R}^m$. What replaces the derivative?

There is no single number that captures how $\mathbf{y}$ changes with $\mathbf{x}$. There are $m \times n$ numbers: for each output component $y_i$ and each input component $x_j$, we can ask how $y_i$ changes when we nudge $x_j$. Collecting them into a matrix gives the Jacobian:

$$J \;=\; \frac{\partial \mathbf{y}}{\partial \mathbf{x}} \;=\; \begin{pmatrix} \dfrac{\partial y_1}{\partial x_1} & \cdots & \dfrac{\partial y_1}{\partial x_n} \\ \vdots & \ddots & \vdots \\ \dfrac{\partial y_m}{\partial x_1} & \cdots & \dfrac{\partial y_m}{\partial x_n} \end{pmatrix}.$$

Jacobian: For $f: \mathbb{R}^n \to \mathbb{R}^m$, the Jacobian at a point $\mathbf{x}$ is the $m \times n$ matrix whose entry $(i, j)$ is $\partial y_i / \partial x_j$. Row $i$ is the gradient of the $i$-th output; column $j$ tells you how every output responds to a nudge in the $j$-th input.

The multivariate first-order approximation reads:

$$\mathbf{y}(\mathbf{x}_0 + \Delta \mathbf{x}) \;\approx\; \mathbf{y}(\mathbf{x}_0) \;+\; J \, \Delta \mathbf{x},$$

with $J$ evaluated at the linearization point $\mathbf{x}_0$. This is the same statement as the 1D case, with numbers upgraded to vectors and matrices. Locally, the potentially very non-linear function $f$ is approximated by a linear map: matrix multiplication by $J$.This is the reason people say 'the derivative is a linear map.' In 1D the map is 'multiply by $f'(x_0)$'; in higher dimensions it is 'multiply by the matrix $J$.' Same idea, richer bookkeeping.

Pause and think: shapes

If $\mathbf{x} \in \mathbb{R}^n$ and $\mathbf{y} \in \mathbb{R}^m$, what shape is $J$? What shape must $\Delta \mathbf{x}$ be for $J \Delta \mathbf{x}$ to make sense, and what shape does the product have?

$J$ has shape $m \times n$. To multiply $J \Delta \mathbf{x}$, we need $\Delta \mathbf{x}$ to be an $n$-vector, which it is (it lives in input space). The product is an $m$-vector, matching output space. The rows count outputs; the columns count inputs. Getting these dimensions straight is the whole game when we apply the idea to transformers.

The Model as a (Locally Linear) Function

A transformer is a deep, non-linear function. But it has a special structure we will exploit: at every layer, information sits in the residual stream, a vector of dimension $d_\text{model}$, and each layer adds an update to it. Reading off the model's output means:

1. Apply layers $\ell{+}1, \ell{+}2, \dots, L$ to the residual stream $\mathbf{h}_\ell$ at layer $\ell$, producing the final residual stream $\mathbf{h}_L$.
2. Apply a layer norm and the unembedding matrix $W_U$ to $\mathbf{h}_L$, producing logits over the vocabulary.

Step 1 is complicated: attention, MLPs, non-linearities, residual connections. Step 2 is linear (up to layer norm). We cannot generally reduce Step 1 to matrix multiplication; the model would not need all those layers if we could. But we can locally approximate Step 1 by matrix multiplication, using the Jacobian.

Fix a prompt and a source position $t$. Treat $\mathbf{h}_\ell$ (the residual stream at layer $\ell$, position $t$) as the input, and $\mathbf{h}_{\text{final}, t'}$ (the final residual stream at some position $t' \ge t$) as the output. Both are vectors of size $d_\text{model}$. The Jacobian

$$J_{\ell, t, t'} \;=\; \frac{\partial\, \mathbf{h}_{\text{final}, t'}}{\partial\, \mathbf{h}_{\ell, t}}$$

is a $d_\text{model} \times d_\text{model}$ matrix. It tells us: if we nudged $\mathbf{h}_{\ell, t}$ by a small vector $\Delta \mathbf{h}$, the final residual stream at position $t'$ would shift by approximately $J_{\ell, t, t'} \Delta \mathbf{h}$.Why $t' ge t$? Because a perturbation at position $t$ can only affect positions from $t$ onwards. Attention in an autoregressive transformer is causal: earlier positions cannot look at later ones.

The shape is worth staring at. Both dimensions equal $d_\text{model}$ because the residual stream has the same width at every layer; it is a shared bus that all layers read from and write to.

Why Average? A Single Jacobian Is Two Things at Once

The Jacobian $J_{\ell, t, t'}$ computed on one prompt tells us the local linearization for that specific input. But that linearization mixes two very different kinds of structure:

• Context-independent structure: how the model's weights generally translate features at layer $\ell$ into features at layer $L$. This is what we want to capture: a property of the model, not of any one input.
• Context-specific structure: how the current prompt's attention patterns, activations, and gates route information through layers $\ell{+}1, \dots, L$. This is transient; it changes with every input.

If we care about what a direction in layer-$\ell$ space generally means, we need to strip out the context-specific part. The J-lens does this by averaging Jacobians over many contexts:

$$J_\ell \;=\; \mathbb{E}_{\,t,\; t' \ge t,\; \text{prompt}} \left[\, \frac{\partial\, \mathbf{h}_{\text{final}, t'}}{\partial\, \mathbf{h}_{\ell, t}} \,\right].$$

Three things are being averaged:

1. Over source positions $t$ within each prompt. A concept encoded at layer $\ell$ should be readable regardless of where in the sequence it appears.
2. Over subsequent positions $t' \ge t$. We do not care only about how $\mathbf{h}_{\ell, t}$ shapes the immediate next-token logits; we care about how it shapes any downstream logit it can influence.
3. Over a corpus of prompts (roughly a thousand, in the original paper). This is what turns context-specific into context-independent: averaging over diverse contexts cancels the parts of each Jacobian that depend on that context and leaves the parts that are stable across contexts.

The idea is exactly the same as computing an average gradient across a dataset: individual gradients point in noisy, context-driven directions; the average points in the direction the loss actually wants to move. Here the "loss" is not a scalar but a linear map, and averaging happens over context, position, and destination.

Number of prompts averaged:

Pause and think: why not just use one prompt?

Suppose we computed the Jacobian on a single prompt about French cities. What would the top-ranked lens tokens look like, and why would that be misleading as a general readout of layer $\ell$?

They would over-represent tokens the model happens to be predicting in that prompt (French words, capitals, geographic terms), because the local Jacobian is heavily shaped by the current attention pattern and the current MLP activations. If we asked "what does this direction in layer-$\ell$ space generally mean in the model?", we would get the wrong answer. Averaging over many prompts is what makes the resulting map a property of the model rather than of any one input.

The final $J_\ell$ is a single $d_\text{model} \times d_\text{model}$ matrix per layer. It is the "average linear translator" from layer $\ell$ to the final layer.

Reading Out: The Lens Formula

With $J_\ell$ in hand, applying the lens to an activation $\mathbf{h}_\ell$ is a one-liner. Multiply by $J_\ell$, apply the model's layer norm, apply the unembedding $W_U$, softmax:

$$\text{lens}(\mathbf{h}_\ell) \;=\; \text{softmax}\!\bigl(\, W_U \,\text{norm}(J_\ell \, \mathbf{h}_\ell) \,\bigr).$$

The result is a probability distribution over the vocabulary: a "top tokens" list you can inspect.

The Anthropic paper describes this as equivalent to replacing all subsequent layers with the appropriate lens matrix. That sentence packs a lot into a few words, and it is worth unpacking with a picture.

What "replacing all subsequent layers" means

Normally the model computes $\mathbf{h}_L$ from $\mathbf{h}_\ell$ by running through everything downstream of layer $\ell$: attention blocks, MLPs, residual connections, layer norms. That downstream computation is a complicated non-linear function; call it $F_\ell$, so that $\mathbf{h}_L = F_\ell(\mathbf{h}_\ell, \text{context})$.

The averaged Jacobian $J_\ell$ is the best single linear map that approximates $F_\ell$ across contexts. So replacing $F_\ell$ with $J_\ell$ means: pretend the model, from layer $\ell$ onward, is just this linear map. Then the final residual stream would be $J_\ell \mathbf{h}_\ell$, and the model's readout would be $W_U \, \text{norm}(J_\ell \mathbf{h}_\ell)$.

That is exactly the lens formula. Nothing extra is going on.

Layer to lens at: Show J-lens replacement

Why this is an approximation, not a fact

The equivalence is a first-order approximation, not an equality. Two things make it inexact:

• Non-linearity. Real transformer layers are non-linear. $J_\ell$ is a local linearization; the further $\mathbf{h}_\ell$ is from the point around which we linearized, the worse the approximation.
• Averaging. The single $J_\ell$ we use is an average over contexts. On any particular prompt, the true local Jacobian would differ from $J_\ell$.

The remarkable empirical result is that this approximation is nevertheless a very useful readout: it recovers coherent, interpretable content at layers where the logit lens returns noise, and its top tokens track the concepts the model is actively "holding in mind." The math is not exact, but the direction it picks in vocabulary space is often the right one.

The J-lens vectors

Look at the lens formula again, focusing on the pre-softmax logits and ignoring the normalization for a moment:

$$\text{logits} \;\approx\; W_U \, J_\ell \, \mathbf{h}_\ell \;=\; \underbrace{(W_U J_\ell)}_{\substack{n_\text{vocab} \times d_\text{model}}} \, \mathbf{h}_\ell.$$

The composed matrix $W_U J_\ell$ has one row per vocabulary token. Row $t$ is a direction in layer-$\ell$ residual-stream space; the logit for token $t$ is (approximately) the inner product of that row with $\mathbf{h}_\ell$.

J-lens Vector: For layer $\ell$ and vocabulary token $t$, the J-lens vector $\mathbf{v}_t^{(\ell)}$ is the $t$-th row of $W_U J_\ell$, viewed as a direction in $\mathbb{R}^{d_\text{model}}$. The lens score for token $t$ at that layer is (up to layer-norm scaling) $\langle \mathbf{v}_t^{(\ell)},\, \mathbf{h}_\ell \rangle$.

There are $n_\text{vocab}$ such vectors per layer, typically $\sim$100,000 vectors in a $d_\text{model}$-dimensional space. That set is overcomplete: no unique decomposition of an activation as a sum of J-lens vectors exists. But sparse combinations turn out to be well-defined and empirically meaningful; the paper calls the set of activations expressible as sparse non-negative combinations of J-lens vectors the J-space.

Pause and think: J-lens vector vs. probing direction

A linear probe for concept $c$ learns a direction $\mathbf{w}_c$ such that $\langle \mathbf{w}_c, \mathbf{h}_\ell \rangle$ correlates with whether $c$ is present. A J-lens vector $\mathbf{v}_t^{(\ell)}$ is also a direction whose inner product with $\mathbf{h}_\ell$ gives a score. What is the difference, mechanistically?

A probe direction is learned to distinguish inputs on some external label. It is correlational: it may or may not align with anything the model itself uses. A J-lens vector is derived from the model's own weights: it is the row of $W_U J_\ell$, so its score is the model's own (first-order) push toward emitting token $t$ downstream. Same geometry (inner product), very different sources of information: labels vs. the model's causal structure.

Comparing the Three Lenses

The three lenses can be written in a common form:

$$\text{lens}(\mathbf{h}_\ell) \;=\; \text{softmax}\!\bigl(\, W_U \, \text{norm}( M_\ell \, \mathbf{h}_\ell ) \,\bigr),$$

differing only in what $M_\ell$ is.

Lens	$M_\ell$	Source of $M_\ell$	Character
Logit lens	$I$ (identity)	Assumes layer-$\ell$ basis matches final layer	No calibration; fails in early layers
Tuned lens	$A_\ell$ (learned affine)	Trained to match final output distribution	Correlational; can "skip ahead" to outputs
Jacobian lens	$J_\ell = \mathbb{E}[\partial \mathbf{h}_L / \partial \mathbf{h}_\ell]$	Derived from the model's own weights, averaged	Causal, first-order; recovers content in early layers

Two of these choices are motivated. The logit lens sets $M_\ell = I$ because in late layers the residual stream is already close to the final basis; there is nothing to translate. That works when it works and fails silently when it doesn't. The tuned lens fits $A_\ell$ so that the readout matches the true output; that pins the readout to what will be emitted, which is not always what we want to look at. The J-lens picks $M_\ell = J_\ell$ because $J_\ell$ is the actual first-order description of what layers $\ell{+}1{:}L$ do, on average across contexts. It is the closest thing to a "linear model of the model" you can extract without any learning.

The three coincide in one important edge case: at the final layer, $J_L$ is (approximately) the identity, and all three reduce to the model's own unembedding. Divergences appear as we go earlier.The J-lens paper reports that the logit lens agrees closely with the J-lens in the last several layers and diverges earlier: exactly the regime where the logit lens is known to fail.

The mean-Jacobian construction was used earlier by Hernandez et al. [2]Linearity of Relation Decoding in Transformer Language Models
Hernandez, E., Li, B. Z., Andreas, J.
ICLR 2024, 2023 to derive per-relation affine maps $W_r \mathbf{s} + \mathbf{b}_r$ (e.g., a single "plays instrument" matrix that turns "Miles Davis" into "trumpet") and a companion attribute lens for tracking a fixed relation across layers. The J-lens generalizes the same first-order-plus-averaging trick from per-relation to per-layer, taking the expectation over a broad corpus rather than examples of one relation.

What Ends Up in the Lens

Because $J_\ell$ was built from the causal effect of $\mathbf{h}_\ell$ on final outputs, the top tokens in the J-lens readout are the ones $\mathbf{h}_\ell$ is disposed to push the model to say, averaged across contexts in which such an activation might arise. Those tokens are not always the next predicted token; often they are concepts the model is holding in mind that would surface only if asked. In the source paper, running the J-lens on the token before "the sport is:" while the model has been instructed to think of a sport reveals soccer as a top lens token several layers before the model actually emits it. Applying the lens while the model silently computes $3^2 - 2$ reveals nine and then seven at intermediate layers, exposing the intermediate step. See [3]Verbalizable Representations Form a Global Workspace in Language Models
Gurnee, W., Lindsey, J., et al.
Anthropic (transformer-circuits.pub), 2026 for many more examples.

Because J-lens readouts are causal by construction, one can also write along a J-lens vector: adding $\alpha \mathbf{v}_t^{(\ell)}$ to $\mathbf{h}_\ell$ tends to make the model more likely to verbalize token $t$, and swapping projections along two J-lens vectors reliably swaps which of the two the model reports. This is a natural bridge to steering methods and to concept-level interventions like activation patching.

Limitations

• Single-token concepts. J-lens vectors are indexed by vocabulary tokens, so multi-token concepts (most named entities, most phrases) are not directly captured. Extensions to spans are discussed in the paper.
• First-order only. The Jacobian ignores higher-order effects. If a concept only exerts influence via a strongly non-linear interaction (say, gated on the value of another feature), the J-lens will miss it or misattribute it.
• Averaged, not per-prompt. The lens is a dispositional readout. For any single prompt, the true local Jacobian differs from $J_\ell$. Using the average is what makes the readout reflect the model rather than the context, but it also means the tool cannot reveal purely context-specific mechanisms without further work.
• Small variance fraction. In practice, the J-space accounts for less than about 10% of activation variance at any layer. Most of what happens in the residual stream is not verbalizable in this sense. This is a feature, not a bug, if we take verbalizability as a meaningful selection criterion, but it also limits how much of the model's computation the lens sees.

Looking Ahead

The Jacobian lens sits at the boundary between observational and mechanistic tooling. Its readouts are observational (a translation of activations into vocabulary), but its construction is causal, and its vectors compose additively with the model's weights in a way that makes them natural handles for intervention. Two directions build directly on this.

First, decomposing an activation as a sparse combination of J-lens vectors gives an interpretation not unlike a sparse autoencoder, but with an interpretable, token-indexed dictionary built from the model's own weights rather than a learned decoder. The paper formalizes this as the J-space.

Second, once we can read a direction that means "the model is about to say $t$," we can also write it. Swapping and injecting along J-lens vectors is a lightweight form of concept-level steering, and gives a clean causal test that the direction really is doing what the lens says. This mirrors the shift from observation to causation covered under activation patching and its relatives.