The Jacobian Lens
or: how to read a model’s inner monologue

What is a layer thinking?

A transformer processes text as a sequence of token positions, and at each position it maintains one vector — the residual stream. Call it h_ℓ,t: the state at layer ℓ, position t, living in ℝ^d with d in the thousands. Every layer reads this vector, computes something, and adds its contribution back in. By the final layer L, the stream has been shaped into something the model can act on directly: multiply by the unembedding matrix W_U, apply a softmax, and you get the probability of every token in the vocabulary being the next word.

So the last layer is easy to read — it comes with its own decoder. The question is everything in between. Take the prompt the paper uses:

“The number of legs on the animal that spins webs is ___”

To answer 8, the model must first work out an unspoken intermediate: spider. That thought, if it exists, is somewhere in the residual stream around the middle layers — a direction in a d-dimensional vector that doesn’t come with labels. How do you find it?

The obvious trick, and where it breaks

The logit lens says: just pretend you’re already at the end. Take the middle-layer state h_ℓ and push it straight through the final decoder:

This works surprisingly well in the last stretch of the network — the residual stream is additive, so late states are already close to the final one. But in early and middle layers it produces noise. The reason is subtle and important: the model’s internal coordinate system drifts across layers. The direction that means spider at layer 30 need not be the direction that means spider at the output. The logit lens assumes the two coordinate systems are identical — that the transport from layer ℓ to layer L is the identity map. Mostly, it isn’t.

The Jacobian lens is the principled correction. Instead of assuming a transport map, it measures one: for each layer, the average linear map describing how small changes at layer ℓ actually end up affecting the final layer, over a thousand different contexts. “How changes propagate” is exactly what derivatives describe — so that measured map is a Jacobian.

To build J_ℓ properly, we need the machinery of derivatives in high dimensions to be fresh. That’s Part II — and it’s where we take our time.

Derivatives, rebuilt

Forget the rules of differentiation for a moment. There is exactly one idea underneath all of it:

A derivative is the best linear stand-in for a function, near a point.

II.1 · One dimension: zoom in until it’s a line

Take a smooth function f and a point x. “Differentiable at x” means: if you zoom in on the graph around (x, f(x)), the curve becomes indistinguishable from a straight line. The slope of that line is f′(x), and the statement

is a promise about the error: the gap between curve and line dies quadratically as the nudge ε shrinks. Halve the window, quarter the error. That quadratic decay is the entire licence for everything that follows — it’s what lets us replace a horrendously complicated function (say, sixty transformer layers) by a linear map and trust the answer for small perturbations.

II.2 · Many inputs, one output: the gradient

Now let f : ℝⁿ → ℝ — many inputs, one output. A “nudge” is now a vector ε ∈ ℝⁿ, and the best linear stand-in must eat a vector and return a number. Linear maps from ℝⁿ to ℝ are exactly dot products with some fixed vector — and that vector is the gradient:

Each partial derivative ∂f/∂x_j answers one narrow question: if I nudge only coordinate j, at what rate does the output move? The gradient is the full list of these exchange rates, and the dot product ⟨∇f, ε⟩ combines them: a nudge in a general direction moves the output by the sum of its per-coordinate contributions. (That sum-of-contributions structure is linearity — it’s only true because we zoomed in.)

II.3 · Many inputs, many outputs: the Jacobian

Finally, F : ℝⁿ → ℝ^m — the case we actually need, since a chunk of transformer maps a d-dimensional state to a d-dimensional state. Now the output is a vector, so we need one gradient per output coordinate. Stack those m gradients as rows and you get the Jacobian matrix:

Two ways to read this matrix, both worth having in your head:

• Row i is the gradient of output coordinate i — the exchange rates from every input to that one output.
• Column j is where the j-th basis nudge lands — nudge input coordinate j alone, and the output moves along column j.

Geometrically: a smooth map warps space, and the Jacobian at a point is the matrix that the warp looks like locally, at that point. Zoom in anywhere on a smooth map, and what you see is a plain linear transformation — a little square gets carried to a little parallelogram. Different point, different matrix.

II.4 · The chain rule is matrix multiplication

Compose two maps and the local pictures compose too: a nudge gets hit by the first Jacobian, and the result gets hit by the second.

This is the whole chain rule. For a transformer it means: the end-to-end Jacobian from layer ℓ to layer L is the product of every per-layer Jacobian in between — one enormous, prompt-dependent matrix product. Nobody forms that product explicitly. Which brings us to the last piece of machinery.

Building the lens

III.1 · The function being differentiated

Fix a prompt and run the model. Pick a layer ℓ and a position t. The “rest of the network” is now a concrete function: it takes the residual state h_ℓ,t and (holding everything else about the run fixed) produces the final-layer states h_L,t′ at that position and every later one — later positions see h_ℓ,t through attention. Its Jacobian,

answers: to first order, how would a small change to this activation shift what the model is about to compute — here, and at every future position? That is precisely the “causal effect” language we wanted. But a Jacobian from one prompt conflates two things: the model’s general convention for how layer-ℓ directions eventually reach the output, and the particular use this one context happens to make of them. On this prompt, nudging toward spider also nudges toward 8 — not because that’s what the direction means, but because of what this context does with it.

III.2 · The clever average

The fix is to average the Jacobian over source positions t, over target positions t′ ≥ t, and over a corpus of a thousand pretraining-like prompts:

Think about what survives this average. The context-specific pathways — the ones that route spider to 8 on this prompt but to Charlotte’s Web on another — point in different directions on different prompts, and they wash out. What remains is the transport that is consistent across contexts: the model’s stable, reusable convention for how a layer-ℓ direction gets turned into words. The paper’s term for representations picked out this way is verbalizable — poised to be said in many possible futures, rather than said in this one.

The result is one d × d matrix per layer — computed once, offline. (The paper probes 25 evenly spaced layers, re-indexed 0–100 so layer numbers read as percentages of depth.)

III.3 · Reading through the lens

To read an activation, transport it to final-layer coordinates with J_ℓ, then decode with the model’s own unembedding pipeline:

That’s a score for every token in the vocabulary; sort them and the top of the list is a human-readable description of what the activation is disposed to make the model say. Now the key structural object: because everything before the softmax is linear (up to a normalization), the score of token w is essentially an inner product,

These rows v_w are the J-lens vectors: one direction in residual-stream space per vocabulary token, per layer. They are the paper’s central objects — usable three ways:

• Ranked readout — sort all scores, inspect the top tokens (the figures in the paper).
• Single-token probe — read ⟨v_w, h⟩ for one chosen concept, e.g. to ask “is spider active above threshold?”
• Handles for intervention — add them, remove them, swap them (Part V).

And notice the logit lens is the degenerate special case J_ℓ = I. The two agree in the last few layers — where residual accumulation makes the transport genuinely close to identity — and diverge earlier, exactly where the logit lens stops working.

III.4 · The recipe, concretely

Here is the actual computation, condensed from the paper’s appendix. The elegant tricks are worth spelling out:

# One-time cost per model: build J_ℓ for every probed layer ℓ.
for prompt in corpus:                      # 1,000 seqs × 128 tokens
    forward pass; cache h_ℓ[t] for all ℓ, t
    for i in 1..d_model:                    # batched in practice
        seed gradient e_i at every position of the target layer z
        backward pass                       # one VJP → reaches all ℓ, all t at once
        for each layer ℓ:
            row i of J_ℓ⁽ᵖ⁾ = mean over t of ∂(Σ_t′ z[t′,i]) / ∂h_ℓ[t]
J_ℓ = mean over prompts of J_ℓ⁽ᵖ⁾            # per-element mean (or median)

# Ever after, reading any activation is one matrix multiply:
lens(h) = softmax( W_U · norm( J_ℓ · h ) )

• The sum over futures is free. Seeding the gradient at every position computes ∂(Σ_t′ z_t′)/∂h_ℓ,t in one pass — and differentiation is linear, so the gradient of the sum is the sum of Jacobians.
• Causality does the bookkeeping. In a causal transformer, ∂z_t′/∂h_ℓ,t = 0 whenever t′ < t — the attention mask guarantees it. The restriction to t′ ≥ t happens automatically.
• Cost: d_model backward passes per prompt, over 1,000 prompts — hefty, but paid once per model. (And quality saturates fast: the J-lens already beats the logit and tuned lenses with ~10 prompts.)

Design choices they ablate, all with modest effects: differentiate to the penultimate layer rather than the final one (the last block is output-calibration machinery, and adds noise); optionally stop-grad through queries and keys, freezing attention patterns so gradients flow only through the value pathway — isolating “what gets moved” from “what gets attended to” (this slightly strengthens causal effects); and restricting targets to t′ = t only (closer in spirit to the logit lens) or t′ > t only (isolating what a position broadcasts to its future).

III.5 · Placing the J-lens among its relatives

The tuned-lens comparison is the conceptually sharp one. Both fit a per-layer linear map — but the tuned lens is trained to predict the model’s output distribution, so on prompts with hidden intermediate reasoning it happily reports the conclusion (8) rather than the thought (spider). The J-lens map is built from causal derivatives, not output correlations, so intermediate content survives.

From lens to space

IV.1 · Too many vectors to be a basis

At each layer we now hold one direction v_w per vocabulary token: n_vocab vectors in d-dimensional space, with n_vocab ≫ d (think ~10⁵ vectors in ~10⁴ dimensions). Such a family is wildly overcomplete: it spans everything, many times over. So the interesting object can’t be “the subspace they span” — that’s all of ℝ^d, and any activation whatsoever could be written as a combination of J-lens vectors. As a set of directions with more members than dimensions, this is what linear algebra calls a frame, not a basis — and decompositions against a frame are never unique.

What rescues the concept is an empirical fact: at any given moment, only a couple dozen J-lens vectors are meaningfully active in a real activation. Sparsity, not span, is the defining constraint.

IV.2 · A union of cones

So the paper defines the J-space as everything you can build from few lens vectors, with nonnegative weights. Fix a sparsity budget k (typically ≈ 25). For each choice of k tokens S, the nonnegative combinations of their vectors form a k-dimensional cone; the J-space is the union over all choices:

Not a subspace — a union of low-dimensional wedges threading through activation space, one wedge per possible “set of ~25 concurrent thoughts.” Given any activation x, its J-space component is the nearest point of 𝓕:

Finding that nearest point exactly would mean searching over all (ⁿ_k) cones — hopeless — so in practice it’s approximated greedily by gradient pursuit: pick vectors one at a time, each time choosing the one that most reduces the reconstruction error, until k are chosen. The chosen tokens are a discrete inventory of “what’s in the workspace right now,” and the coefficients a_i are its local coordinates. (Sparse decomposition typically yields a less redundant inventory than naïvely taking the top-k inner products, because the lens vectors are far from orthogonal — spider and spiders shouldn’t both get full credit.)

IV.3 · Small in variance, large in consequence

Now the fact that makes all of this more than bookkeeping. Decompose real activations, or concept vectors extracted independently of the lens, and the J-space component is tiny: never more than ~10% of the variance, a 6–7% median for concept vectors. Over 90% of what a representation “is,” by energy, lies outside the J-space.

Yet in the paper’s causal experiments, that sliver does nearly all the talking. In the think-of-a-word swap experiments, intervening along a concept’s J-space component drives the target into the model’s top outputs on 59% of trials (pure J-lens vectors: 88%) — while the other 93% of the vector, swapped at matched magnitude, succeeds 5% of the time. Same story for internal reasoning through unspoken intermediates (61% vs 28%, the latter collapsing to 6% once you clamp the J-space so the concept can’t re-enter it). The J-space is a low-variance, high-leverage channel — which is exactly the shape you’d expect of a workspace: a small bulletin board that many circuits read, sitting atop a much larger volume of automatic machinery that never gets posted.

Writing into the workspace

Reading is half the story. The lens vectors are directions in activation space, so they’re also handles you can push on.

V.1 · Steering and ablation

Positive α injects a concept (this is how the paper implants “thoughts” and tests whether the model can introspectively detect them). Negative α, or projecting the component out entirely, suppresses one. Ablating the top-k J-space contents wholesale is how the paper shows the model can still parse and speak fluently without its workspace — while losing much of its capacity for deliberate internal reasoning.

V.2 · The surgical swap, and why it needs a pseudoinverse

The sharpest intervention exchanges one thought for another — spider out, ant in — while disturbing nothing else. The subtlety: J-lens vectors are not orthogonal. You cannot just read the coefficient of v_spider with a dot product, because v_spider overlaps v_web, v_insect, and thousands of others. Reading coordinates against a non-orthogonal pair takes a dual basis — and that is exactly what the Moore–Penrose pseudoinverse supplies.

Stack the two lens vectors as columns of V = [v_s v_t] ∈ ℝ^d×2. The pseudoinverse V^† = (V^⊤V)⁻¹V^⊤ solves the least-squares problem: c = V^†h is the coefficient pair minimizing ‖h − Vc‖, which splits the activation cleanly as

— the in-plane part, expressed in the skewed coordinates of the two concepts, plus an orthogonal remainder r carrying everything else the activation encodes. Now swap the two coordinates with the permutation σ(c₁, c₂) = (c₂, c₁) and write back only the difference:

Whatever weight sat on spider now sits on ant, and vice versa; the remainder r — untouched — still says “animal question, counting legs, mid-sentence.” The rows of V^† are the dual vectors ṽ_s, ṽ_t satisfying ⟨ṽ_i, v_j⟩ = δ_ij: each one reads “its” coordinate while being blind to the other.

V.3 · What the swap revealed

This one operation, applied at the workspace layers, carries a remarkable share of the paper’s results:

• Recall: swap spider ↔ ant in the legs prompt, and the answer flips 8 → 6. Swapping the intermediate takes effect ~17% of depth earlier than swapping the answer token — the model really computes spider first, then derives 8.
• Planning: in a couplet ending in fight, swapping the planned rhyme to light changes the words the model picks earlier in the line (“coming” → “morning”) — the workspace holds plans that shape output before they’re spoken.
• Cross-lingual reasoning: answering a Chinese antonym question, the intermediate is the English token big; swapping big → long flips the Chinese output 大 → 长.
• Systematically: intermediate swaps redirect two-hop answers on 54–70% of trials (Haiku → Sonnet/Opus 4.5), and the same France↔China swap works across “capital of,” “language of,” “continent of” — one representation, read by many downstream circuits. That is the broadcast property that earns the word workspace.

The whole method on one page

Differentiate the rest of the network; average the Jacobian over a thousand contexts; compose with the unembedding. The rows of the result are one direction per word the model could ever say — reading them ranks the model’s unspoken vocabulary; sparse-decomposing against them inventories a small “workspace” carrying <10% of activation variance; and steering or swapping them rewrites reports, plans, and multi-step inferences. Everything else in the paper — verbal report, directed modulation (“concentrate on citrus fruits”), internal reasoning, flexible generalization, selectivity — is this one mathematical object, interrogated five different ways.

Two structural facts to carry with you: the workspace has a band — readouts are noise for roughly the first third of layers, coherent through the middle, then collapse to next-token “motor” content in the last few — and it has a capacity, those ~25 meaningfully active directions. Small, late-ish, low-variance, causally decisive.

If you want to go deeper: the paper itself (Methods, plus appendices on methodological ablations, the sparse-frame formalization, and multi-token extensions), and J-lens readouts on open-source models are browsable on Neuronpedia.