How the Jacobian Lens Works: Reading What a Model Is Poised to Say

17 July 2026 · ~15 min read

The Jacobian Lens (J-Lens) is a small idea with a big payoff: to read a hidden layer of a language model, don't pretend it already lives in the output vocabulary — transport it there with the derivative of the rest of the network. This post builds the idea from scratch: what a Jacobian is, how the lens uses one, and why the Jacobian is exactly the right object for the job.

1. The problem: reading a mind mid-thought

A transformer processes a token through a stack of layers. At each layer the model keeps a running vector — the residual stream — that it reads from and writes to. Call the residual stream at layer \(\ell\) and position \(t\) the vector \(h_{\ell,t} \in \mathbb{R}^{d}\), where \(d\) is the model width (d_model).

We would love to ask, at any intermediate layer: what is the model currently "thinking" — which words is it leaning toward? The classic tool for this is the logit lens. The model turns its final residual vector into a distribution over vocabulary tokens using a normalization step and an unembedding matrix \(W_U\):

where \(L\) is the last layer. The logit lens simply applies that same readout to an intermediate layer:

This works surprisingly well in the late layers and poorly in the early and middle ones, where it mostly returns noise. Why? Because it makes a hidden assumption.

The Jacobian Lens fixes exactly this. Instead of assuming the rest of the network is the identity, it measures what the rest of the network does — to first order — and corrects for it. That measurement is a Jacobian. So let's build up what a Jacobian is.

2. What is a Jacobian?

2.1 From slope to gradient to matrix

Start with a plain scalar function \(f:\mathbb{R}\to\mathbb{R}\). Its derivative \(f'(x)\) is the local slope: nudge the input by a tiny \(\delta\), and the output changes by about \(f'(x)\,\delta\). This is the first-order Taylor approximation,

Now let the input be a vector and the output be a vector too: \(F:\mathbb{R}^{n}\to\mathbb{R}^{m}\), written component-wise as \(F(x) = \big(F_1(x), \dots, F_m(x)\big)\). There is no single "slope" anymore — each of the \(m\) outputs can respond to each of the \(n\) inputs. Collecting all of those responses gives the Jacobian matrix \(J \in \mathbb{R}^{m\times n}\):

Read entry \(J_{ij}\) as a sentence: "if I nudge input coordinate \(j\), how much does output coordinate \(i\) move?" The Jacobian is just the whole table of those answers.

2.2 The Jacobian is the best linear approximation

The single most important fact for us: the Jacobian is the best local linear approximation of a nonlinear map. The vector Taylor expansion reads

Near the point \(x\), the complicated nonlinear function \(F\) behaves like the simple linear function \(\delta \mapsto J\delta\). A whole stack of attention blocks, MLPs, and normalizations — replaced, locally, by a single matrix multiply. That is what makes the Jacobian tractable and interpretable at the same time.

In more than one dimension, "tangent line" becomes "tangent plane" (and, in higher dimensions still, a tangent hyperplane). The Jacobian is that tangent map. It's worth building the picture in your head, so here it is in 3D — a curved surface \(z=f(x,y)\) and the flat sheet that best matches it at one point. Drag to orbit.

2.3 The chain rule: Jacobians of compositions multiply

A transformer is a composition of layers, so we need to know how Jacobians behave under composition. If \(F = G \circ H\) (apply \(H\), then \(G\)), the multivariate chain rule says the Jacobians simply multiply:

So the Jacobian of "layer \(\ell{+}1\), then \(\ell{+}2\), …, then \(L\)" is the product of the individual layer Jacobians. One matrix summarizes the first-order effect of the entire remaining network. (This is also, not coincidentally, exactly the quantity backpropagation computes.)

3. The Jacobian Lens

3.1 The key move

Everything above layer \(\ell\) is a function that maps the layer-\(\ell\) activation to the final activation. Freeze a prompt and a position and call that function \(F_\ell\):

The logit lens assumed \(F_\ell\) was the identity. The Jacobian Lens instead linearizes \(F_\ell\) around the activation that actually occurred, using the Jacobian from §2.2:

In words: to see what a layer-\(\ell\) vector becomes at the output, push it through the derivative of the rest of the network. That derivative — a \(d \times d\) matrix — is the transport map from layer \(\ell\)'s coordinate system into the final layer's coordinate system. It is precisely the correction the logit lens was missing.

3.2 Averaging: from "verbalized" to "verbalizable"

A Jacobian computed at one activation is specific to that one prompt and position — and individually those are noisy. The J-Lens instead uses the average Jacobian over many source positions \(t\), many later positions \(t' \ge t\), and a large corpus of prompts:

This averaging is doing real conceptual work. A concept that only maps to a given word in one specific context has been verbalized in that instance; a concept whose average linear readout robustly points at a word is one the model is generally poised to verbalize — a stable, context-independent property of the representation. The expectation is what turns a one-off measurement into a reusable lens.

3.3 Reading a layer through the lens

With the averaged transport map \(J_\ell\) in hand, reading an activation is just the logit-lens pipeline with the identity replaced by \(J_\ell\):

Compare this line to the logit lens in §1. The only change is the insertion of \(J_\ell\): first linearly transport \(h_\ell\) to the output coordinate system, then apply the model's own final normalization and unembedding to rank vocabulary tokens. The downstream stack of layers has been swapped for a single, faithful, first-order stand-in.

4. A worked example you can check by hand

Let's make the whole mechanism concrete with the smallest example that still shows the point. Take a model with width \(d=2\) and a vocabulary of just two tokens, “cat” and “dog”. Let the unembedding be the identity, so the readout just compares the two coordinates:

Suppose that at layer \(\ell\) the model has already decided it will say “dog”, but stores that decision in a rotated internal basis, so the raw activation is

Suppose further that the rest of the network, to first order, mixes the two coordinates through this averaged transport map (we'll reuse it in §5):

(We drop the nonlinear \(\operatorname{norm}\) here — it rescales but doesn't reorder the two logits, so it can't change the winner in this example.) Now read the layer two ways:

Working the arithmetic: \((J_\ell h_\ell)_0 = 0.3(0.9)+0.9(0.2)=0.45\) and \((J_\ell h_\ell)_1 = 0.9(0.9)+0.3(0.2)=0.87\). The logit lens got it exactly backwards because it read \(h_\ell\) in the wrong basis; the Jacobian lens applied the one matrix that carries layer \(\ell\) into the output basis, and recovered the token the model was actually poised to emit. That single \(2\times2\) matrix is doing the job of "the entire upper half of the network, linearized." This is Figure 3, now with numbers.

5. Deeper: the spectrum of the transport map

A matrix is more than a lookup table of partials — it has geometry, and for the Jacobian lens that geometry is the interpretability payoff. The clean way to see it is the singular value decomposition. Any real matrix factors as

with \(U, V\) orthogonal. Reading \(J_\ell\,x\) right-to-left tells a three-step story: \(V^\top\) measures the input along an orthonormal set of right singular directions \(v_i\); \(\Sigma\) scales each component by \(\sigma_i\); and \(U\) writes the result along left singular directions \(u_i\). A unit sphere of inputs comes out as an ellipsoid whose axes are the \(u_i\) with lengths \(\sigma_i\):

For our example \(J_\ell\) is symmetric, so its singular values are the magnitudes of its eigenvalues, \(\sigma_1=1.2\) and \(\sigma_2=0.6\), along the eigenvectors \(v_1=\tfrac{1}{\sqrt2}(1,1)\) and \(v_2=\tfrac{1}{\sqrt2}(1,-1)\). The interpretability reading is direct: the large-\(\sigma\) singular directions are the concept directions the model is most poised to verbalize. A feature aligned with \(v_1\) is amplified on its way to the vocabulary; a feature aligned with \(v_2\) is muffled. The lens isn't just relocating the vector — its spectrum ranks which internal directions get a loud voice at the output.

5.1 What "averaging a Jacobian" actually means

One subtlety worth stating plainly: the averaged \(J_\ell\) of §3.2 is not the Jacobian of any single function. It is a mean of linear operators, one per context, collapsed into a single representative operator: \(J_\ell=\mathbb{E}[\,J_\ell^{(\text{context})}\,]\). Directions along which the model consistently routes a concept to the same output token survive the averaging; directions that point every which way across contexts cancel. So the expectation is a noise filter with a precise meaning — it keeps the stable, context-independent part of the readout, which is exactly what "a concept the model is generally poised to verbalize" should mean.

5.2 A note on rigor

"Best linear approximation" has an exact definition worth keeping honest. \(F\) is differentiable at \(x\) with Jacobian \(J\) iff

i.e. the leftover shrinks faster than linearly in the perturbation size — that \(o(\lVert\delta\rVert)\) is the entire content of "first order." The lens's affine constant is then pinned down as \(\text{const}=F(x)-Jx\), the expansion evaluated at the chosen base point. And note we keep the model's final \(\operatorname{norm}\) outside \(J_\ell\) rather than linearizing it too: it is genuinely nonlinear, and folding it in would blur the clean linear-transport story for little gain. That's a modeling choice, not a theorem — the kind of seam worth knowing is there. (No functional analysis required: everything here lives in finite-dimensional \(\mathbb{R}^d\), where the Jacobian is the total derivative.)

6. Why the Jacobian is the right tool here

It's worth being explicit about why this particular object — and not some other summary of the network — is the correct one. Four reasons stack up.

4.1 It answers a causal question

Interpretability wants to know what a representation does, not just what it correlates with. The Jacobian \(\partial h_{L}/\partial h_{\ell}\) is by definition the first-order causal effect of the representation on the output: it says how the model's final state would move if you perturbed the hidden state. "Which concepts is the model poised to verbalize" is a causal question, and the Jacobian is the linear-order answer to it.

4.2 Linear, and therefore composable and decodable

The final unembedding \(W_U\) is a linear map. A linear transport map \(J_\ell\) composes with it cleanly (§2.3), so the lens stays a single interpretable matrix product all the way to the vocabulary. A nonlinear summary would not slot into the \(W_U\) decoder like this. Linearity is not a loss of ambition here — it is what makes the readout legible.

4.3 It corrects the basis drift automatically

The whole failure of the logit lens was that intermediate layers don't share the output basis. But the transformation between those bases is exactly what the derivative of the intervening layers encodes. So \(J_\ell\) doesn't just approximate the rest of the network — the specific thing it captures, the change of coordinates from layer \(\ell\) to the output, is the precise correction the logit lens omitted. The tool is matched to the disease.

4.4 Averaging separates stable structure from noise

Because the Jacobian is well-defined per example, it can be averaged in a principled way (§3.2). That expectation is what lets the method distinguish a representation that is genuinely, repeatably readable as a concept from one that happened to align with a token once. This is the property that lets the authors surface meaningful content in early layers where the logit lens only saw noise — and, ultimately, to argue that verbalizable representations form a shared "global workspace" inside the model.

7. Where the math lets you challenge it

The Jacobian lens is built on two approximations stacked on top of each other — linearize, then average — and each one opens a specific line of attack. If you wanted to poke holes in the method (or design the experiment that finds its failure mode), here is the mathematical attack surface. Each challenge below comes with a concrete way to test it.

7.1 A tangent used far from where it is tangent

The differentiability guarantee is strictly local: \(F(x+\delta)=F(x)+J\delta+o(\lVert\delta\rVert)\) holds only as \(\lVert\delta\rVert\to 0\). But the reading \(\text{lens}(h_\ell)=\operatorname{softmax}(W_U\operatorname{norm}(J_\ell h_\ell))\) applies \(J_\ell\) to the entire activation \(h_\ell\), which has finite — often large — norm. In effect \(\delta\) is the whole vector, so the \(o(\lVert\delta\rVert)\) control is vacuous. The discarded second-order term is governed by the Hessian \(H\),

and for a deep, sharply-curved \(F\) that quadratic term need not be small. So the lens is extrapolating a tangent plane far from its point of tangency — exactly the regime where Figure 2 stops being a good fit. Test it: measure the linearization fidelity per layer, \(\rho_\ell = \lVert F_\ell(h_\ell) - (J_\ell h_\ell + c_\ell)\rVert \,/\, \lVert F_\ell(h_\ell)\rVert\). Where \(\rho_\ell\) is not \(\ll 1\) (suspect the early layers), the lens's tokens are not first-order justified.

7.2 The saturation blind spot: decisive features with zero gradient

A Jacobian measures sensitivity, not use. If a feature drives the output through a saturated nonlinearity — a softmax already near one-hot, a GELU/ReLU sitting in its flat region — then by the chain rule its contribution to \(J_\ell\) is multiplied by \(\sigma'(z)\approx 0\) and vanishes, even though the feature may fully determine the output. The lens is therefore blind precisely to information the model has already committed to, arguably the most important content at that point in the computation. Put sharply: J-Lens shows you what is still movable, not what is already settled. A concept can be maximally verbalized and yet have a near-zero Jacobian. This is the deepest conceptual objection, because it is not noise — it is a systematic bias baked into using a derivative as a readout.

7.3 The averaged operator is the Jacobian of nothing

\(J_\ell=\mathbb{E}[\,J_\ell^{(c)}\,]\) is a mean of per-context matrices, but the linearization point co-varies with the context \(c\). In general

so the averaged map need not describe any real forward pass — it is a representative fiction. Worse is cancellation: if one context routes a concept direction \(v\) toward a token and another routes it the opposite way, the two contributions annul in the mean, and the lens reports \(J_\ell v \approx 0\) — "not verbalizable" — for something that is verbalized in every context, just context-dependently. That is a false negative created by the averaging itself. Test it: compare \(\lVert\mathbb{E}[J_\ell^{(c)}]\rVert\) against the per-context spread \(\mathbb{E}\lVert J_\ell^{(c)}-J_\ell\rVert\); when the spread rivals the mean, the single \(J_\ell\) is not representing a stable computation.

7.4 Spectral blind subspaces

This one falls straight out of §5. Reading through \(J_\ell=U\Sigma V^\top\), any part of \(h_\ell\) lying along a right-singular vector \(v_i\) with \(\sigma_i\approx 0\) is annihilated before it ever reaches the vocabulary. So the lens has a blind subspace — the span of the small-\(\sigma\) right-singular vectors — and information stored there is invisible no matter how real it is. The catch is that a nonlinear downstream step could later amplify one of those very directions; a linear average simply cannot represent "small now, amplified after the next nonlinearity." The lens's silence about a direction is evidence of nothing. And because the readout \(W_U J_\ell\) is many-to-one (nontrivial nullspace), the converse fails too: a token lighting up does not prove the model contains that concept, only that something projected onto it — the standard interpretability-illusion risk.

7.5 The mundane ones

Less deep but worth stating: the final \(\operatorname{norm}\) is genuinely nonlinear and is kept outside \(J_\ell\) as a modeling choice, not a derived fact; and estimating a full \(d\times d\) Jacobian across many positions and prompts is far costlier than the free logit lens. The payoff is signal in the layers the logit lens throws away — but you are paying for it, and you are trusting the four approximations above to buy it.

8. The one-sentence version

The logit lens reads a hidden layer by pretending the rest of the network is the identity; the Jacobian Lens reads it by replacing the rest of the network with its averaged derivative — the single linear map that transports an intermediate representation into the output vocabulary — and that derivative is the right tool because it is causal, linear, basis-correcting, and averageable.