Logit lens

A transformer's last layer is not the only one whose residual stream the unembedding can read. Push layer 3's residual through the model's own final norm and head and you get a distribution over the vocabulary: what the model would say if layer 3 were its last. Do that at every layer and you can watch an answer appear.

That is the logit lens, and in this format it is one field. "view": "logits" on a read projects it; a layer sweep is the lens; a token_prob metric on the answer is what to plot against depth. No second step, no saved activations, no separate script.

At a glance #

"reads": {
  "lens": {
    "site": "resid",
    "pos": -1,
    "input": "base",
    "view": "logits"
  }
},
"metrics": {
  "p_answer": { "kind": "token_prob", "of": "lens", "token": "base_answer" },
  "top1":     { "kind": "match", "of": "lens", "expected": "base_answer" }
}

The read gathers the residual window and pushes it through ln_final and lm_head inside the trace, so what the metric scores is a vector over the vocabulary rather than a hidden state. Everything downstream — token_prob, match, logit_diff — is the same code that scores a real lm_head read.

The document #

{
  "header": {
    "description": "The logit lens as one document: at every layer, push the residual stream at the answer position through the model's final norm and head and read what it would say if that were its last layer. `view: logits` is the projection; the layer sweep is the lens; `token_prob` on the answer is what to plot against depth. Two layers here, so two points."
  },
  "model": {
    "key": "hf-internal-testing/tiny-random-LlamaForCausalLM",
    "revision": "9fb191250dd56d0ba7ec9785a025ed29c03d5998",
    "dtype": "fp32"
  },
  "roles": {
    "base": {
      "field": "input"
    }
  },
  "sites": {
    "resid": {
      "component": "block_output",
      "layers": {
        "sweep": [
          [
            0
          ],
          [
            1
          ]
        ]
      }
    }
  },
  "interventions": {
    "lens": {
      "reads": {
        "lens": {
          "site": "resid",
          "pos": -1,
          "input": "base",
          "view": "logits"
        }
      },
      "metrics": {
        "p_answer": {
          "kind": "token_prob",
          "of": "lens",
          "token": "base_answer"
        },
        "top1": {
          "kind": "match",
          "of": "lens",
          "expected": "base_answer"
        }
      }
    }
  },
  "steps": {
    "lens": {
      "kind": "observe",
      "interventions": "lens",
      "rows": {
        "base": "weekdays/train"
      },
      "saves": [
        {
          "value": "p_answer",
          "file_path": "p_answer.json"
        },
        {
          "value": "top1",
          "file_path": "top1.json"
        }
      ]
    }
  }
}

Reading it in the order the compiler does:

• model is the tiny random-weight Llama. It has two layers, so the sweep has two points.
• roles has only base. There is no counterfactual: a lens is an observation, not an intervention, and the document has no writes and no second model.
• sites.resid is block_output with a layers sweep. One site, one read, one metric pair — the sweep is what makes it a curve.
• interventions.lens has the one read, with view: "logits", and two metrics on it: p_answer, the probability of the row's own answer, and top1, whether that answer is the argmax.
• steps.lens is one observation pass over weekdays/train, saving both tables.

What it compiles to #

$ uv run causalab-mini explain documents/v2/logit_lens.json
<root>: Plan
  layers=0: Plan
    lens: Observe  metrics=['p_answer/token_prob', 'top1/match']  saves=['p_answer.json', 'top1.json']
        forward 'original' on 'base'  (4 rows x 11 tokens)
          read  'lens' at block_output[0] pos={index:-1} via 'identity' as logits
  layers=1: Plan
    lens: Observe  metrics=['p_answer/token_prob', 'top1/match']  saves=['p_answer.json', 'top1.json']
        forward 'original' on 'base'  (4 rows x 11 tokens)
          read  'lens' at block_output[1] pos={index:-1} via 'identity' as logits

One forward per point, no patched model, and the read is printed as logits — the view is part of the read's identity, not a post-processing step someone might forget to apply.

Running it #

$ time uv run causalab-mini run documents/v2/logit_lens.json --out out/logit_lens
out/logit_lens/document.json
out/logit_lens/run.json
out/logit_lens/layers=0/document.json
out/logit_lens/layers=0/p_answer.json
out/logit_lens/layers=0/top1.json
out/logit_lens/layers=1/document.json
out/logit_lens/layers=1/p_answer.json
out/logit_lens/layers=1/top1.json

real	0m8.434s

The curve #

p_answer.json is one row per example, so the curve against depth is the point directory you read it out of. Both layers, all four rows:

Four rows of p_answer.json, per point, and the mean of top1.json. On a two-layer model with random weights there is no curve: the answer's probability is about 1/32000 at both depths, which is what a uniform distribution over this vocabulary gives, and the argmax is never the answer.

Is the lens wired to the right head? #

A logit lens is a wiring claim: this norm, this unembedding, applied in this order. A lens plugged into the wrong normalization still produces a plausible-looking curve, so it is worth one check before anything is read off it — at the last layer the lens must reproduce the model's own output, because there is nothing between them.

This model's last block output is layer 1, so the layers=1 point is that check:

import json
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

KEY = "hf-internal-testing/tiny-random-LlamaForCausalLM"
REV = "9fb191250dd56d0ba7ec9785a025ed29c03d5998"

rows = json.load(open("documents/data/weekdays/train.json"))
tok = AutoTokenizer.from_pretrained(KEY, revision=REV)
tok.padding_side = "left"
if tok.pad_token is None:
    tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(KEY, revision=REV, dtype=torch.float32)
model.eval()

batch = tok([r["input"] for r in rows], return_tensors="pt", padding=True)
with torch.no_grad():
    logits = model(**batch).logits[:, -1, :]
probs = logits.softmax(dim=-1)
ids = [tok.encode(r["base_answer"], add_special_tokens=False)[0] for r in rows]
mine = [float(probs[i, t]) for i, t in enumerate(ids)]

lens = [r["value"] for r in json.load(open("out/logit_lens/layers=1/p_answer.json"))]
for a, b in zip(lens, mine):
    print(f"lens {a:.9f}   model {b:.9f}   diff {abs(a - b):.2e}")

$ python lens_check.py
lens 0.000028089   model 0.000028089   diff 0.00e+00
lens 0.000028376   model 0.000028376   diff 0.00e+00
lens 0.000027431   model 0.000027431   diff 0.00e+00
lens 0.000034297   model 0.000034297   diff 0.00e+00

Compared as raw logits rather than probabilities the two agree to the bit across the whole vocabulary, and the argmax is identical. They can differ: the lens projects one position per row where the model projects the whole sequence, and a GEMM at a different M may accumulate in a different order, which FINDINGS.md §10 records as up to one ulp. Nothing larger than that is a rounding difference, and a lens that disagrees with its own model's last layer by more is wired wrong.

Where a logits view is refused #

Two things the compiler will not let a logits view be, both refused before any weights load.

It is not the residual stream #

The final norm and the head expect a residual stream, so the view is allowed only at a component that is one: embeddings, block_input, block_output, ln_final. An MLP output is a summand of the stream, not the stream, and projecting it would be projecting a difference:

$ uv run causalab-mini validate logit_lens_on_mlp.json
…
Value error, interventions.lens: read 'lens': view 'logits' projects the residual stream through the head, and 'mlp_output' is not the residual stream (one of ['block_input', 'block_output', 'embeddings', 'ln_final'])

It is not featurized, and it is not written back #

$ uv run causalab-mini validate logit_lens_featurized.json
…
Value error, interventions.lens: read 'lens': a featurized read cannot also be viewed as logits

A featurizer is defined at the site's width; a logits view is vocabulary-wide, and the two cannot both be true of one read. For the same reason a write may not take a logits view as its operand — refused naming the read, with "which is vocabulary-wide and cannot be written back at a site".

logits and lm_head are two components #

Both are in the component vocabulary and on most families they are the same tensor. They are not the same thing:

Gemma-2 is the family where they differ: it caps its logits with final_logit_softcapping, tanh(f / cap) · cap, after the head. A metric scored on lm_head there is scored on numbers the model does not predict from, and on a long tail the cap changes which token is on top.

The lens takes the cap into account. The projection is softcap(lm_head(ln_final(x)), cap), with cap read off the model's config and None on every family that has none — so on Gemma-2 a view: "logits" read is comparable with the model's own output rather than with its head's, which is what a lens is for. On a family with no cap the softcap is the identity and the two are the same tensor.