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:

Near $x_0$, the curve looks linear. A first-order approximation uses that local straightening 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. Activations in this curriculum are row vectors, so the transformer sections below use the transposed Jacobian, which carries inputs along its rows and multiplies $\Delta \mathbf{x}$ from the right.

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'} \;=\; \left( \frac{\partial\, \mathbf{h}_{\text{final}, t'}}{\partial\, \mathbf{h}_{\ell, t}} \right)^{\!T}$$

is a $d_\text{model} \times d_\text{model}$ matrix, transposed so that it acts from the right: its rows index the input and its columns the output. 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 $\Delta \mathbf{h} \, J_{\ell, t, t'}$.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:

• Distribution-averaged structure: how the downstream computation locally translates perturbations at layer $\ell$ over a chosen sample of contexts. This is the target of the averaged map.
• 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[\, \left( \frac{\partial\, \mathbf{h}_{\text{final}, t'}}{\partial\, \mathbf{h}_{\ell, t}} \right)^{\!T} \,\right].$$

The expectation averages over source positions, reachable destination positions, and prompts:

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 favored by that prompt because the local Jacobian depends on current attention patterns and MLP gates. Averaging over many prompts reduces prompt-specific variation and defines a map for that sampling distribution. Change the corpus, positions, or weighting and the average may change, so it is not an input-free property of the weights alone.

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(\, \text{norm}(\mathbf{h}_\ell \, J_\ell) \, W_U \,\bigr).$$

Applying the softmax turns these linearized logits into a probability distribution over the vocabulary, giving us a "top tokens" list to 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 $\mathbf{h}_\ell J_\ell$, and the model's readout would be $\text{norm}(\mathbf{h}_\ell J_\ell) \, W_U$.

This gives us the lens formula directly.

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$.

In the reported evaluations, this approximation produces coherent vocabulary readouts at some depths where the raw logit lens is less informative. The examples support the J-lens as a useful diagnostic. They do not make every top token a faithful description of what the model is “holding in mind”; prompt, corpus, and first-order approximation errors remain possible.

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\; \mathbf{h}_\ell \, J_\ell \, W_U \;=\; \mathbf{h}_\ell \, \underbrace{(J_\ell W_U)}_{\substack{d_\text{model} \times n_\text{vocab}}}.$$

The composed matrix $J_\ell W_U$ has one column per vocabulary token. Column $t$ is a direction in layer-$\ell$ residual-stream space; the logit for token $t$ is (approximately) the inner product of that column 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 column of $J_\ell W_U$, 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 column of $J_\ell W_U$, 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(\, \text{norm}( \mathbf{h}_\ell \, M_\ell ) \, W_U \,\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)^T]$	Derived from the model's own weights, averaged	Causal, first-order; recovers content in early layers

The three choices answer different questions. The logit lens sets $M_\ell=I$, measuring direct alignment with the final unembedding. The tuned lens fits $A_\ell$ to forecast the final output distribution. The J-lens uses $J_\ell$, a first-order summary of what layers $\ell{+}1{:}L$ do over the chosen contexts. It is derived rather than trained against output targets, but its quality still depends on the averaging distribution and on how nonlinear the relevant finite changes are.

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 $\mathbf{s} W_r + \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.