# The Jacobian lens: reading the residual stream with a derivative

> Satyajit Ghana — Head of Engineering @ Inkers Technology
> canonical: https://ai.thesatyajit.com/articles/jacobian-lens
> date: 2026-07-08
> tags: explainer, interpretability, transformers, llm
The **Jacobian lens** answers one question: *given a model's internal state at some layer, what is that state disposed to make the model say?* It is a small piece of code — Anthropic's [`jacobian-lens`](https://github.com/anthropics/jacobian-lens) repo, Apache-2.0, "reference implementation, not maintained" — that fits on any open-weights decoder transformer and reads intermediate activations out as ranked vocabulary tokens. The whole method is one line:

$$
\operatorname{lens}_\ell(h) = \operatorname{unembed}\!\big(J_\ell \, h\big), \qquad J_\ell = \mathbb{E}\!\left[\frac{\partial h_{\text{final}}}{\partial h_\ell}\right]
$$

Two moves. $J_\ell$ **transports** a residual-stream vector $h$ at layer $\ell$ into the final-layer basis. Then the model's own **unembedding** decodes it into logits over the vocabulary. The transport matrix is a *Jacobian* — the derivative of the final-layer residual with respect to layer $\ell$ — averaged over a corpus. That is the entire idea, and it is worth being precise about, because the derivative is doing something the [logit lens](#a-lens-is-a-choice-of-transport) never could.

<Callout type="note">
This is the companion tool for the paper [*Verbalizable Representations Form a Global Workspace in Language Models*](https://transformer-circuits.pub/2026/workspace/index.html) (Anthropic, 2026). Everything below is drawn from the repo README, the `jlens.fitting` source, and that paper. The causal numbers I quote are the paper's own, measured on Anthropic's models (Haiku/Sonnet/Opus 4.5). I have not reproduced them; I mark them as paper-reported.
</Callout>

## Why a derivative

A residual stream is not written in the output vocabulary. The activation $h_\ell$ at some middle layer lives in a basis the model has been rotating and rewriting layer by layer; the unembedding $W_U$ only makes sense against the *final* layer. So you cannot just apply $W_U$ to a middle activation and expect a sensible token — the coordinates don't line up. That is the flaw in the logit lens, and the Jacobian lens is the fix: it first maps $h_\ell$ **through** the layers above it, then unembeds.

Mapping through the layers exactly would mean running the rest of the network — nonlinear, and no longer a "lens." The Jacobian is the linear stand-in. It is the best linear approximation to "run layers $\ell{+}1 \dots L$" around an operating point, and the repo takes that operating point to be the *average* over a text corpus. One matrix per layer, fit once, applied everywhere.

<Figure
  src="/articles/jacobian-lens/fig1.png"
  alt="Three-panel method figure. Panel A: computing the lens by backpropagating from the final-layer residual stream at present and future token positions back to an activation h at layer l, forming the d_model by d_model Jacobian matrix, then aggregating over token positions and dataset examples into J_l. Panel B: reading with the lens replaces all layers above l with the single matrix J_l followed by the unembedding, producing a ranked token readout such as Mars, color, planet, fourth. Panel C: intervening in J-space by swapping projections onto two lens vectors, with the patched-activation formula."
  caption="The Jacobian lens: (A) compute J_ℓ by backprop from the final layer to h_ℓ and average over positions and prompts; (B) read by replacing everything above ℓ with J_ℓ then unembed; (C) intervene by swapping J-space coordinates (Anthropic, transformer-circuits.pub, Figure 4)."
/>

## How the matrix is estimated

The estimator has one subtlety worth getting right. For a source position $p$ at layer $\ell$, the influence on the final layer is not a single vector — it is spread over the current position and every *future* position (a decoder is causal, so $h_\ell[p]$ can affect the output at $p, p{+}1, \dots$). The repo sums that influence over all target positions, then averages over source positions and over prompts. In the paper's own pseudocode:

```python
# Compute J_ℓ for all layers ℓ.
# h_ℓ[t] : residual stream at layer ℓ, position t
# z[t]   : residual stream at the target layer L (final by default)
for each prompt p in corpus:
    run forward pass; cache h_ℓ[t] for all ℓ, t
    for i in 1..d_model:                    # one backward pass per output dim (batched)
        grad_z = e_i ⊗ 1_T                  # inject ∂/∂z_i = 1 at every position
        for each layer ℓ:
            G_ℓ = ∂(Σ_t z[t]) / ∂h_ℓ        # autodiff → shape [T, d_model]
            J_ℓ^(p)[i, :] = mean_t G_ℓ[t, :]  # mean over source positions
for each layer ℓ:
    J_ℓ = mean over prompts of J_ℓ^(p)      # aggregate the average Jacobian
# apply:
lens(h_ℓ) = softmax( W_U · norm( J_ℓ · h_ℓ ) )
```

Because $\sum_t z[t]$ is differentiated, a one-hot cotangent lands at every target position at once; causality makes $\partial z[p']/\partial h_\ell[p]$ vanish for $p' < p$, so what survives is exactly the sum over current-and-future targets. The paper's lenses use **1000 sequences of 128 tokens**; the README notes quality "saturates quickly" and ~100 prompts is already usable. Cost is dominated by the model's own backward pass — one per output dimension, batched — which is why fitting is embarrassingly parallel across corpus slices (`JacobianLens.merge()`).

The rows of $W_U J_\ell$ are the interesting object: the paper calls them the **J-lens vectors**, one per vocabulary token, each a direction in residual-stream space "associated with a single token."

## A Jacobian is a tangent

Here is the mechanism I want to build intuition for, because it is also the method's main limitation. A Jacobian is a *local linear map*. Collapse the stack to one scalar activation coordinate and one token's logit, and the true function is a curve; the lens is its tangent line. Slide the probe below.

<LensTangent />

In `local` mode the tangent is re-taken at the probe, so it always touches — that is the textbook Jacobian, exact at the point and good nearby. But the real lens is `corpus-average` mode: **one** tangent, taken once at the corpus operating point, then used for every activation you feed it. Near that point the linear readout is faithful; far from it the error grows. This is the honest cost of turning "run the rest of the network" into a single matrix. The lens tells you what an activation is disposed to say *to first order, on average*; it is not a faithful simulation of the model from layer $\ell$ onward.

## Then it is just a dot product

The second half, `unembed(·)`, is the easy half. Unembedding is a matrix of one row per token; a logit is that row dotted with the transported vector. Softmax ranks them, and the lens *readout* is the top of the list — the token whose direction $J_\ell h$ points most toward. Rotate the transported vector and watch the readout hand off between neighbouring concepts:

<UnembedReadout />

The superscript rank you see on a real slice page — `nose³`, `smile¹⁰⁴` — is exactly a token's position in this sorted vocabulary. Rank, not just top-1, is what makes the lens useful: a concept can be climbing toward the top for several layers before it ever wins.

## What it surfaces: the ASCII-face

The repo ships one example that makes the point in a single screenshot. Give the model an ASCII-art face and ask what it depicts. The `^` character is the nose. Select that position and the lens, at *middle* layers, reads out **nose** — a word that never appears in the prompt.

<Figure
  src="/articles/jacobian-lens/fig2.png"
  alt="A layer-by-position slice page for an ASCII-art face. A grid shows the lens top-1 token at each position and layer, with superscript ranks. At the caret (nose) position, column 28, the mid layers around layer 42 read out 'nose' at rank 1, roughly 10 percent. A rank heatmap below shows a bright hotspot for 'nose' concentrated at that position and mid depth, and rank-versus-layer and rank-versus-position line charts track 'nose' and 'smile' peaking mid-stack."
  caption="The ASCII-face slice: at the caret (nose) position, the lens reads out 'nose' at mid layers (rank 1, ≈10%) though the word is absent from the prompt — the model parsed the drawing spatially (Anthropic, jacobian-lens repo, assets/slice_vis.png)."
/>

Reading the page: each cell is the lens top-1 word at that `(position, layer)`; the bottom row is the model's actual output; the heatmap and line charts track a pinned concept's rank across the grid. The signal for "nose" is not at the output layer — it is a hotspot in the *middle*, then it fades as the model resolves what to actually say. That is the whole pitch: the lens shows intermediate content the output distribution has already moved past.

## A lens is a choice of transport

Every "lens" is the same unembedding applied to a transported activation; they differ only in the transport $J_\ell$.

| lens | transport $J_\ell$ | how it's obtained | early-layer behaviour |
|---|---|---|---|
| **logit lens** | identity $I$ | none | assumes one basis for all layers; recovers little early content |
| **tuned lens** | learned linear map | trained to match the output distribution (correlational) | tends to "skip ahead" to the output |
| **Jacobian lens** | $\mathbb{E}[\partial h_{\text{final}}/\partial h_\ell]$ | fit by autodiff over a corpus | corrects for cross-layer basis change by construction |

The logit lens is the $J_\ell = I$ special case — it works only where the residual basis already matches the final layer, i.e. the last few layers, and the paper notes the J-lens "agrees closely" there and diverges earlier. The tuned lens also fits per-layer linear maps, but on a *correlational* objective (match the output), which the paper finds "skips ahead" and buries exactly the unverbalized intermediates you wanted to see. The Jacobian's objective is the derivative itself, which is why it recovers interpretable content at depths where the logit lens does not.

## Does the readout mean anything causally

A readout is a correlation until you intervene on it. The paper's stronger claim is that the J-lens directions are *causally* privileged, and it tests this by swapping coordinates in the space those vectors span ("J-space", panel C above). The headline: split a concept vector into its J-space component and the orthogonal remainder, then swap along each.

<BenchBars
  title="Top-5 concept-swap success (%) — paper-reported, Anthropic models"
  unit="%"
  bars={[
    { label: "J-lens vectors", value: 88, highlight: true },
    { label: "concept's J-space part", value: 59 },
    { label: "orthogonal remainder", value: 5 },
  ]}
/>

The striking part is the budget: that J-space component carries a **median of only 6–7%** of the concept vector's variance, with ~93% in the orthogonal remainder — yet the small component drives the swap (59%) and the remainder barely moves it (5%). A few percent of the variance is doing almost all of the *reportable* work. Intervening at intermediate layers also propagates: swapping a concept mid-stack flips the model's top-1 output on a majority of trials, scaling with model size.

<BenchBars
  title="Intermediate-swap success (%) — top-1 output flips, paper-reported"
  unit="%"
  bars={[
    { label: "Haiku 4.5", value: 54 },
    { label: "Sonnet 4.5", value: 70 },
    { label: "Opus 4.5", value: 70 },
  ]}
/>

The companion result is ablation: project the top J-space contents out of the residual stream and multi-hop reasoning collapses toward zero, while shallow tasks — classification, comparison, factual recall — are essentially unaffected. Read together, the lens is not just labelling activations; the directions it finds carry content the model actually uses for the harder, chained computations.

## What it can and cannot tell you

Honest boundaries, several of them the authors' own words:

- **It is a linear approximation, and an average one.** One Jacobian per layer, taken at the corpus mean, stands in for a nonlinear stack. The `LensTangent` widget is the whole caveat — faithful near the operating point, progressively wrong away from it.
- **Single tokens only.** Each J-lens vector is tied to one vocabulary token. Concepts that span multiple tokens are not directly captured (the appendices discuss extensions). The lens sees "Mars," not "the fourth planet."
- **Approximate and incomplete.** The paper states plainly that the J-lens "only approximately and incompletely captures the model's underlying workspace structure," and that a "true workspace" may operate in layers the lens misses.
- **A readout is not a mechanism.** The lens says what an activation is *disposed* to output; it does not tell you the circuit that put it there. The causal swaps above are what upgrade a readout from suggestive to load-bearing — do the intervention before you trust the picture.
- **It is a reference implementation.** Not optimized, not maintained; fitting is dominated by the model's backward pass. Fine for research, not a production probe.

## Running it

The API is two calls — fit (or download) a lens, then apply it at chosen positions:

```python
import transformers, jlens

hf  = transformers.AutoModelForCausalLM.from_pretrained("org/model").cuda()
tok = transformers.AutoTokenizer.from_pretrained("org/model")
model = jlens.from_hf(hf, tok)

lens = jlens.JacobianLens.from_pretrained("org/lens-repo", filename="model/lens.pt")
lens_logits, model_logits, _ = lens.apply(
    model, "Fact: The currency used in the country shaped like a boot is",
    positions=[-2])
for layer, logits in sorted(lens_logits.items()):
    print(layer, [tok.decode([t]) for t in logits[0].topk(5).indices])
```

Fitting your own is `jlens.fit(model, prompts=...)`; the `walkthrough.ipynb` notebook goes end to end and renders a slice page like the ASCII-face one.

## The take

The Jacobian lens is a clean idea executed narrowly. Swap the logit lens's implicit identity transport for the real averaged derivative, and you get a readout that works in the middle of the network, where the interesting, not-yet-verbalized computation lives. The mechanism is a first-order Taylor term — a tangent — and the honest framing is exactly that: a good local, on-average picture of what a layer is disposed to say, not a faithful replay of the layers above it. What earns it more than "nice visualization" is the causal follow-through: a component holding ~6–7% of a concept's variance drives most of the reportable behavior, and ablating the J-space directions specifically breaks multi-hop reasoning while leaving shallow tasks intact. That is a real, testable claim about which internal directions the model uses to talk — reached with a derivative and an unembedding, and not much else.

---

*Built on the [`jacobian-lens`](https://github.com/anthropics/jacobian-lens) reference implementation (Anthropic, Apache-2.0) and the paper [Verbalizable Representations Form a Global Workspace in Language Models](https://transformer-circuits.pub/2026/workspace/index.html). Figures reproduced from the paper (Figure 4) and the repo (`assets/slice_vis.png`) for commentary. The `LensTangent` and `UnembedReadout` widgets are my own illustrations of the mechanism, not measured traces; all quantitative results are paper-reported on Anthropic's own models and I have not independently reproduced them.*