The Jacobian Lens

A logit lens decodes a hidden state as if it were the final state: softmax(W_U · norm(h)). At mid layers that readout is dominated by whatever has the largest norm in the residual stream: attention-sink directions, high-norm outliers, the <|im_end|> token. Our raw top-k came back as CJK fragments and end-of-turn markers. The surface text of the residual is not the mechanism.

The Jacobian lens decodes J · h instead, where J is the Jacobian of the remaining network (layers ℓ+1..N plus the final norm) with respect to h. That reads what the state pushes the output toward: how the final state moves as h moves along itself. Sink directions the network carries but does not act on fall away; downstream layers are not sensitive to them.

Why this is computable

The full Jacobian is [hidden × hidden] per layer per position, 4096² numbers for a mid-size model. The lens only needs J · h, and forward-mode autodiff computes Jacobian-vector products. Run the remaining network as f_rest with the hidden state as both input and tangent:

_, Jh = torch.func.jvp(f_rest, (h,), (h,))

The output tangent is J · h, one extra forward pass, no matrix built.

Why the kernels must be eager

Forward-mode AD needs a derivative rule for every op, and the fused attention kernels (SDPA, flash attention) have none. Call jvp through them and PyTorch says so: trying to use forward AD with _scaled_dot_product_flash_attention that does not support it. Eager attention is plain PyTorch with known derivatives, so attn_implementation="eager" is a requirement. The grouped-matmul expert kernel in mixture-of-experts routing has no forward-AD path either, so experts run eager too. It is slower; that is the cost.

What the number is, and is not

The lens returns a mass per token: full-softmax probability over J · h, summed across the answer span and averaged over a few mid-depth layers. Lens mass is evidence about token dynamics, which tokens the answer’s own gradient was sensitive to. It is a claim about the model’s internals and says nothing about what the answer means or about the world.

What it cannot see

The subsampling treats each position independently: f_rest runs the remaining blocks with no attention across the sampled tokens, so it reads each position’s own causal push and misses interactions between neighbours. That is an approximation, and it is stated as one. Positions are capped, so long spans are read through a sieve.

Tokens below the top-k cutoff read as zero, which is not the same as having no mass; an anchor absent from the readout is one it could not see, which is why tracked-id accumulation exists. The quantity also depends on layer selection: mass at layers 0.4, 0.55, 0.7, 0.85 of depth is a claim about those layers only. A lens that has never disagreed with itself across depths has not been checked, only unfalsified.