The logit lens reads what a language model is predicting partway through its computation. A Transformer builds its next-token prediction over many layers, but it reads that prediction out only once, at the very end. The lens applies the same readout after every layer, so you can watch the prediction form. Its companion, direct logit attribution, splits the final prediction exactly into one contribution from each attention head, each MLP, and the embeddings. Together they are the cheapest first look inside a model, and often the first thing a researcher runs on a new one.

Keep one example in mind. GPT-2 reads "The Eiffel Tower is in the city of" and puts most of its probability on " Paris". At which layer did " Paris" become the answer, and which heads and MLPs put it there? This section builds the two tools that answer those questions. It says exactly what each one measures and shows where each one stops being trustworthy.

The tools also address a weakness of the previous section. Section 23.2 read activations with probes, small classifiers trained on a model's hidden vectors to detect a concept. A probe shows that information is present, not that the model uses it, and it reads along a direction that we fitted. The tools here train nothing. They read the hidden vectors with the model's own output layer, the one it really uses to predict.

Section 23.1 ranked interpretability evidence on a ladder of four rungs, association, direct effect, causal effect, and mechanism. Both tools stand on the second rung. A direct effect is what one part of the model sends straight to the output while every other part is held fixed. We build the lens, a trained refinement of it called the tuned lens, and the exact algebra behind the split, and we run all of them on a real trained model in your browser. The section ends where direct effects stop working. A part can send almost nothing straight to the output and still be essential, and measuring that is the job of §23.4.

The model already contains a decoder

Start with how a Transformer is wired. Each token has one vector, the residual stream, that runs from the bottom of the model to the top. Every attention layer and every MLP reads the stream and adds its output back into it, and nothing is ever overwritten (§5.1). So the vector at the top is the embedding plus the sum of everything every layer wrote.

Now the notation. Token vectors are rows, so a weight matrix $\mathbf{W} \in \R^{d_{\text{in}}\times d_{\text{out}}}$ acts on a vector as $\mathbf{x}\mathbf{W}$. The model has $L$ blocks of width $d$, each an attention layer followed by an MLP. Write $\mathbf{x}^{(\ell)}_t \in \R^{1\times d}$ for the stream at position $t$ after $\ell$ blocks. So $\mathbf{x}^{(0)}_t$ is the token embedding plus the position embedding, and $\mathbf{x}^{(L)}_t$ is the final stream. Write $\mathbf{x}^{(\ell,\text{mid})}_t$ for the stream inside block $\ell$, after its attention and before its MLP. (Section 5.1 counted sublayers. Here the superscript counts blocks.) Written out, the running sum is

Here $h$ is the number of heads per layer. The term $\mathbf{c}^{(\ell,i)}_t = \mathbf{H}^{(\ell,i)}_t\mathbf{W}_O^{(\ell,i)}$ is what head $i$ of layer $\ell$ writes. It is the head's output row $\mathbf{H}^{(\ell,i)}_t \in \R^{1\times d_v}$ times the head's own block $\mathbf{W}_O^{(\ell,i)} \in \R^{d_v\times d}$ of the output projection. We may split the projection head by head because concatenating the heads and projecting equals projecting each head and adding (§4.5). The vector $\mathbf{b}_O^{(\ell)}$ is the attention layer's output bias, and $\mathbf{c}^{(\ell,\text{mlp})}_t$ is the MLP's write. We call each term a component and write a generic one as $\mathbf{c}$.

The model reads the final stream exactly once. It applies the final LayerNorm and then the unembedding $\mathbf{W}_U \in \R^{d\times V}$, where $V$ is the vocabulary size. The unembedding turns a vector into one score per vocabulary token, called the token's logit,

In plain words, LayerNorm subtracts the vector's mean $\mu$ and divides by its spread $s$. Then it multiplies coordinate by coordinate by a learned gain $\boldsymbol\gamma$ and adds a learned shift $\boldsymbol\beta$, both in $\R^{1\times d}$ (§5.2). Here $\mathbf{1} \in \R^d$ is the all-ones column, and $\epsilon$ is a small constant that keeps the division safe. The softmax of the logits is the model's next-token distribution, $p(\cdot\mid t) = \operatorname{softmax}(\mathbf{z}_t)$.

Column $v$ of the unembedding, $\mathbf{u}_v = (\mathbf{W}_U)_{:,v}$, is token $v$'s readout direction. Its dot product with the normalized stream is token $v$'s logit. In a tied model, such as GPT-2 or the tiny GPT of Chapter 6, the unembedding is the transposed embedding table, $\mathbf{W}_U = \mathbf{E}^\top$ (§6.2). An untied model learns $\mathbf{W}_U$ as a matrix of its own. The circuit formulas of §6.8, and of this chapter, keep $\mathbf{E}$ and $\mathbf{W}_U$ as separate symbols, so they hold in both cases.

Nothing stops us from applying this readout to the stream before the end. That is the whole idea of the lens.

In plain words, stop the forward pass early and ask the output layer what it would predict if the stream ended here.

There is a second way to see it. By the running sum, the model's real logits are $\operatorname{LN}_f\big(\mathbf{x}^{(\ell)}_t + \text{everything written after } \ell\big)\mathbf{W}_U$. The logit lens replaces everything written after $\ell$ by zero. Whether that is a sensible replacement is an empirical question, and the answer depends on the model.

Three ways to read a lens

A lens gives a whole distribution over the vocabulary at every site. Three summaries of it are standard.

1. The top-1 token, the lens's best guess, $\arg\max_v q_\ell(v\mid t)$.
2. The rank of the model's final top-1 token $v^\star = \arg\max_v p(v\mid t)$ under the lens. It is the number of tokens the lens scores strictly higher, $\#\{v : q_\ell(v\mid t) > q_\ell(v^\star\mid t)\}$. Rank 0 means the lens already agrees with the final answer.
3. The KL divergence to the final distribution, $D_{\mathrm{KL}}\big(p\,\Vert\, q_\ell\big) = \sum_v p(v)\log\big(p(v)/q_\ell(v)\big)$. It measures how much worse $q_\ell$ predicts than $p$, on average over the model's own predictions. It is zero exactly when $q_\ell = p$. With the natural log it is in nats. Dividing by $\ln 2$ gives bits, where one bit is the information in a fair coin flip.

The direction of the KL matters. Written this way it is an average over the model's own distribution $p$. So it charges the lens heavily wherever $p$ has mass and $q_\ell$ does not. A $q_\ell$ trained on this direction learns to cover every token $p$ finds likely, which is why §0.5 calls it mass-covering. This is the quantity the original post plotted, and the one the tuned lens below is trained to minimize.

Even this toy shows the lens reading something sensible at every site. It also shows two things we need later. The scale $s$ changes from site to site, and the final answer is a sum of contributions.

import numpy as np

def logit_lens(states, gamma, beta, W_U, eps=1e-5):
    """states: (S, T, d) residual streams at S sites, the last one final.
    Returns logp (S, T, V), the lens log-probabilities, and rank (S, T), the rank of the
    final top-1 token at each site (0 means it is the lens's top prediction there)."""
    raise NotImplementedError

import numpy as np

def lens_kl(logits_lens, logits_final):
    """logits_lens: (S, T, V) lens logits; logits_final: (T, V) the model's logits.
    Returns (kl_nats, kl_bits), each (S, T), the KL divergence D_KL(p || q) of each lens
    distribution q from the final distribution p."""
    raise NotImplementedError

What the lens shows in real models

The logit lens was introduced in a blog post on GPT-2 (nostalgebraist 2020). The post works mostly with the largest GPT-2, which has 1.5 billion parameters and 48 blocks, and its main plots read a passage of the GPT-3 paper's abstract.

You can run the same experiment on a real trained model here. The widget uses the 2-layer character-level GPT of Chapter 6, trained on Shakespeare's sonnets. It decodes five sites with the lens, the embedding and the stream after each attention layer and each MLP. Click any character of the prompt to read the prediction made after it.

Why does the bottom of a tied model predict its own input? In a tied model, token $w$'s readout direction is its own embedding $\mathbf{e}_w$, row $w$ of the embedding table. At the embedding site the stream is the input token's embedding $\mathbf{e}_v$ plus a position vector. So the lens logit of token $w$ is roughly $\mathbf{e}_v\cdot\mathbf{e}_w$. By the Cauchy-Schwarz inequality, $\mathbf{e}_v\cdot\mathbf{e}_w \le \lVert\mathbf{e}_v\rVert\lVert\mathbf{e}_w\rVert$, and the self term $\mathbf{e}_v\cdot\mathbf{e}_v = \lVert\mathbf{e}_v\rVert^2$ meets this bound. So the dot product is usually largest at $w = v$. This is a fact about tied geometry, not evidence that the model is "thinking about" its input.

How much of the convergence is evidence that the model refines one guess? Less than it looks. The stream is a running sum, so a late site shares almost all of its terms with the final stream. Unless the last few blocks write something large, its lens reading will be close to the final one, whatever the network computes. The next cell checks this on a network that was never trained.

The random network's lens converges smoothly too. Its KL to the final distribution falls from 7.5 bits to 0, and its top-1 agreement with the final answer climbs from 0.09 to 1. Convergence is what the arithmetic of a sum gives you for free. Two things do carry information. One is what the lens says along the way, such as plausible continuations in the middle layers rather than noise. The other is how much earlier it matches the final answer than an untrained baseline like this one does. In practice, never report "the lens converges" as a finding. An interviewer will ask what an untrained model shows.

Where the logit lens breaks

The tuned-lens paper ran the logit lens on many more models and found it unreliable (Belrose et al. 2023). On GPT-Neo-2.7B, reading an excerpt from the abstract of the Transformer paper, the lens gives no interpretable predictions before layer 21 of 32. On BLOOM and OPT-125M, the lens's top-1 token is often the input token itself, rather than any plausible continuation, in more than half the layers.

Even where the lens looks sensible, it is biased, in a sense the paper made precise. Average the model's final distribution over every position of every sequence in a dataset, to get a marginal distribution $\bar p(v)$. Do the same for the lens at layer $\ell$, to get $\bar q_\ell(v)$. The bias is $D_{\mathrm{KL}}(\bar p\,\Vert\,\bar q_\ell)$. An unbiased lens puts, on average, as much probability on each token as the model does.

Three mechanisms explain most of the failures. Each is worth knowing, because each also breaks other tools that read the stream.

Later writes are not zero on average. The lens replaces everything written after layer $\ell$ by zero. Suppose later blocks add a large vector that is roughly the same for every input, a constant bias. Then zero is a strange replacement, and the lens reads a stream the final norm never sees in practice. The cheapest fix replaces zero by a learned constant $\mathbf{b}_\ell$, which the paper calls the debiased logit lens.

Representational drift. The statistics of the stream change with depth. In Pythia 12B, layer 4 introduces two outlier dimensions that dominate the covariance of the hidden states. With those two removed, the covariances of two layers grow steadily less similar the farther apart the layers are (Belrose et al. 2023). A readout trained for the last layer then misreads the early ones. This is the rotation worry from above, in a milder form.

Massive activations. A handful of activations in large language models are orders of magnitude larger than the rest (Sun et al. 2024). In LLaMA2-7B they live in two fixed coordinates (1415 and 2533), and they appear only at the starting token and at the first period or newline. They appear abruptly at layer 2 and stay nearly constant until layer 30. In LLaMA2-70B they exceed 15,000 in absolute value. They act as fixed bias terms. They also draw attention to their tokens, which ties them to attention sinks, the low-information tokens where heads park attention they do not want to place anywhere (§4.4).

Here is what one massive value does to the lens. A value $a$ in coordinate $k$ makes $s \approx |a|/\sqrt d$, so the normalized stream is nearly $\sqrt d$ times a one-hot vector at $k$. The lens then returns whatever tokens row $k$ of $\mathbf{W}_U$ favors, times the gain $\gamma_k$ and the sign of $a$, regardless of context. So treat lens readings at the first position and at delimiters with suspicion. Massive activations are not the outlier features that make 8-bit activation quantization hard. Those are a few channels that are large at almost every token, while a massive activation is one value at one token. Sun et al. (2024) checked that none of the outlier-feature channels they found coincide with the massive-activation coordinates. Section 20.2 meets both kinds of extreme value from the quantization side, where either one sets the step of a per-tensor scale.

One more limit is structural. The lens sees only what $\mathbf{W}_U$ can read. A feature the model computes for internal use, such as that the subject is plural or that a quote is still open, can sit in directions the unembedding ignores. No lens reading will ever show it. A lens is a probe whose readout was trained for the last layer only.

The tuned lens

If each layer writes in a slightly different basis, learn a translation for each layer. The tuned lens (Belrose et al. 2023) adds one affine map per layer, from that layer's stream to the space the final norm expects. It is trained so that the decoded distribution matches the model's own.

In plain words, teach each layer's stream to speak the last layer's language, using the model's own output as the answer key.

Three design choices deserve a reason.

The loss is a distillation loss. Expand the KL, $D_{\mathrm{KL}}(p\,\Vert\,q) = H(p, q) - H(p)$. The entropy $H(p)$ does not depend on the translator. So minimizing the KL is minimizing the cross-entropy of the lens against the soft labels $p$. That is knowledge distillation, training a student to match a teacher's full output distribution (§20.4). Here the teacher is the full model, and the student is "layer $\ell$ plus a translator". Matching the model's outputs, rather than the true next tokens, matters. A probe trained on true labels is rewarded for extracting anything predictive, even information the model never uses, which was the worry behind control tasks in §23.2. A lens trained to match $p$ has no reason to know more than the model.

Identity initialization. Starting at the logit lens means training learns only the departure from it. In layers that already share the final basis, the translator stays near the identity. So a large learned change is itself a measurement of drift.

A translator, not a new unembedding. A fresh readout per layer would have $V\times d$ parameters, 38,597,376 for GPT-2 small. A translator has $d^2 + d = 590{,}592$. The paper reports that training new unembeddings needed more steps and larger batches and often reached a worse perplexity, the exponential of the average cross-entropy, where lower is better.

The results were clear. The tuned lens had much lower perplexity than the logit lens across the board, including on BLOOM and GPT-Neo, where the logit lens fails. It also had much lower bias on GPT-Neo-2.7B, the model for which the paper plots bias. The tuned lens also relies on the same features the model does, a check described in the box below.

The widget isolates the three failure mechanisms. It builds a stream as a sum, so that each block adds a share of the answer. Then it distorts the earlier sites in three ways. A basis rotation and a bias drift both grow linearly from nothing at the top to their full size at the embedding, and sites 2 to 7 carry one huge constant coordinate. It compares the logit lens with an affine translator fit by least squares.

All three sliders start at zero, where the two curves nearly coincide. The fitted lens is a little worse there, because least squares matches the final stream, not the final distribution. Turn the rotation up first. The logit-lens curve rises above the tuned one, most over the first half of the sites, which are rotated the most. The bias drift and the outlier coordinate each break it at every site they touch. The tuned-lens curve does not move at all. The translator undoes all three distortions, because each one is affine. What it cannot undo is the part of the answer that later blocks have not written yet, which is why the tuned-lens curve still falls with depth. That remaining curve is the honest picture of when information arrives.

import numpy as np

def fit_translator(H_l, H_L, ridge=0.0):
    """Affine least squares from layer-l states H_l (n, d) to final states H_L (n, d).
    Returns (Tm, b), Tm (d, d) and b (d,), minimizing ||H_l Tm + b - H_L||^2 + ridge ||Tm||^2
    (the intercept b is not penalized)."""
    raise NotImplementedError

Here is the real way to train one, in PyTorch, on a frozen toy network. The model's blocks write large constants, which is the bias problem. Each block also adds information from a source the earlier states never saw, a stand-in for attention bringing in other tokens. The code records the stream entering each block, which is what a forward hook does in a real model (§0.9). Each translator starts at the identity and is trained by Adam on $D_{\mathrm{KL}}(p\,\Vert\,q)$. In PyTorch's argument order that is F.kl_div(log_q, log_p, log_target=True).

The logit lens is 2 to 8 bits off at the first five sites. It is not even monotone, because each block's constant write pushes the stream somewhere the final norm never sees. The tuned lens falls steadily with depth. Its remaining KL is the information that later blocks bring in, which no translator of an earlier state can recover.

What the tuned lens is used for

A prediction trajectory is a compact fingerprint of how the model processed an input. That makes it useful for spotting inputs that the model processes abnormally.

Lenses also power decoding methods. Decoding by contrasting layers, or DoLa (Chuang et al. 2024), compares an early layer's lens distribution with the final one and favors the tokens whose log-probability grew the most between them. That improved factuality without any training, as §19.8 describes.

Any lens is also an early exit in waiting, a way to stop the forward pass before the last layer. If layer 20 of 32 already predicts the final token with confidence, generation could stop there and skip the last 12 blocks, more than a third of the compute. Confident adaptive language modeling, CALM (Schuster et al. 2022), trains the model so that every layer can be decoded and exits when a confidence measure passes a threshold. A lens gives a similar signal for a frozen model.

The most general view is Patchscopes (Ghandeharioun et al. 2024). Take the hidden state at layer $\ell$ and position $t$ of a source prompt, and optionally transform it. Patch it into a forward pass on a target prompt, at some layer and position, and read what the model generates. The logit lens and the tuned lens are the special case in which the state is patched into the last layer, so the target prompt never matters. Patching into a middle layer instead lets the model's own later layers do the decoding. The price is one forward pass per reading and a reliance on the model's ability to put what it knows into words.

Folding the final LayerNorm into the unembedding

Now turn from the trajectory to the final logits, and split them exactly into contributions. The obstacle is the final LayerNorm, which sits between the stream and the logits. Look at its four steps. Subtracting the mean is linear. Multiplying by the gain is linear. Adding the shift adds a constant. Only the division by $s$ depends on the whole vector in a nonlinear way. So everything except that one number can be moved into the weights, which is called folding the norm into the unembedding.

The algebra uses one matrix, the centering matrix $\mathbf{C} = \mathbf{I} - \tfrac1d\mathbf{1}\mathbf{1}^\top \in \R^{d\times d}$. For a row $\mathbf{x}$, $\mathbf{x}\mathbf{C} = \mathbf{x} - \big(\tfrac1d\mathbf{x}\mathbf{1}\big)\mathbf{1}^\top = \mathbf{x} - \mu\mathbf{1}^\top$, so multiplying by $\mathbf{C}$ subtracts the mean. For example, $(1, 2, 3)\mathbf{C} = (-1, 0, 1)$. The matrix is symmetric. It is idempotent, $\mathbf{C}^2 = \mathbf{C}$, since centering twice changes nothing. And $\mathbf{1}^\top\mathbf{C} = \mathbf{0}^\top$. It is the projection onto the vectors with zero mean (§0.10).

Proof. Write the gain as a matrix, $\boldsymbol\gamma\odot\mathbf{y} = \mathbf{y}\operatorname{diag}(\boldsymbol\gamma)$ for a row $\mathbf{y}$. Then

using $\mathbf{x} - \mu\mathbf{1}^\top = \mathbf{x}\mathbf{C}$ and the associativity of matrix products. The columns of $\mathbf{W}_U'$ sum to zero because $\mathbf{1}^\top\mathbf{W}_U' = (\mathbf{1}^\top\mathbf{C})\operatorname{diag}(\boldsymbol\gamma)\mathbf{W}_U = \mathbf{0}^\top$. So adding $c\,\mathbf{1}^\top$ to $\mathbf{x}$ adds $c\,\mathbf{1}^\top\mathbf{W}_U' = \mathbf{0}^\top$. For the second claim, let $\mathbf{z} = \frac1s\mathbf{x}\mathbf{W}_U' + \mathbf{b}_U$. The vocabulary-centered logits are $\mathbf{z}\big(\mathbf{I} - \tfrac1V\mathbf{1}\mathbf{1}^\top\big) = \mathbf{z} - \bar z\,\mathbf{1}^\top$ with $\bar z = \tfrac1V\mathbf{z}\mathbf{1}$, the same constant subtracted from every logit. Softmax is unchanged by adding a constant to all of its inputs (§0.4 and §4.3 use the same fact for stability), and a difference $z_a - z_b$ loses the constant on both sides. ∎

In plain words, the gain and the mean subtraction are linear, so they can live in the weights. The shift becomes a fixed bias on the logits. Only the division by the one number $s$ stays outside.

RMSNorm, the norm in LLaMA and most models since (§5.2), has no mean subtraction and usually no shift. It folds with $\mathbf{W}_U' = \operatorname{diag}(\boldsymbol\gamma)\mathbf{W}_U$, and $s$ is replaced by the root mean square $\sqrt{\tfrac1d\sum_j x_j^2 + \epsilon}$.

The cell checks the theorem on random numbers, with streams whose means are large and different.

Folding is exact to rounding error. Centering over the vocabulary moves raw logits by more than a unit, while it leaves log-probabilities and differences untouched, and the centered logits average to zero. That last property is why centered logits are easier to read. A component that adds the same amount to every logit no longer looks important.

import numpy as np

def fold_ln(gamma, beta, W_U):
    """gamma, beta: (d,); W_U: (d, V). Returns (W_U_f, b_U) with W_U_f (d, V) and b_U (V,),
    the final LayerNorm folded into the unembedding, centered over d and over the vocabulary."""
    raise NotImplementedError

def folded_logits(x, W_U_f, b_U, eps=1e-5):

Splitting a logit difference exactly

Which quantity should we split? A single logit $z_a$ is a poor target. Adding the same number to every logit changes nothing the model does, so the size of $z_a$ alone means little (vocabulary centering exploits exactly this). A probability is a poor target too. It depends nonlinearly on every logit, and near 0 or 1 it saturates, so large changes in the stream barely move it. The standard choice is a logit difference between a correct token $a$ and a plausible wrong token $b$, called the foil,

In plain words, $\Delta_t$ is the log-odds of the right answer over one specific wrong answer. For example, $p(a) = 0.6$ and $p(b) = 0.2$ give $\Delta = \ln 3 \approx 1.10$.

It has four good properties. It ignores a constant added to all logits. It does not saturate. It has a meaning in probability terms. And, as the theorem below shows, it is linear in the stream once $s$ is fixed. It also asks a sharp question ("why Paris rather than Rome?") instead of a diffuse one. The study of how GPT-2 completes "When Mary and John went to the store, John gave a drink to" with "Mary" uses the difference between the two names for exactly these reasons (§23.4). In the stream, the direction that raises $\Delta_t$ is $\mathbf{u}_a - \mathbf{u}_b$, scaled coordinate by coordinate by the gain and then centered, because the norm removes the mean.

Proof. By the folding theorem, $\mathbf{z}_t = \frac1s\mathbf{x}^{(L)}_t\mathbf{W}_U' + \mathbf{b}_U$, so $\Delta_t = z_{t,a} - z_{t,b} = \frac1s\mathbf{x}^{(L)}_t(\mathbf{u}'_a - \mathbf{u}'_b) + (b_{U,a} - b_{U,b})$. Substitute the sum for $\mathbf{x}^{(L)}_t$ and distribute the product over it. For one component, $\mathbf{c}(\mathbf{u}'_a - \mathbf{u}'_b) = \mathbf{c}\,\mathbf{C}\operatorname{diag}(\boldsymbol\gamma)(\mathbf{u}_a - \mathbf{u}_b) = (\mathbf{c} - \bar c\,\mathbf{1}^\top)\cdot\mathbf{w}$. The bias term is $b_{U,a} - b_{U,b} = \boldsymbol\beta\cdot(\mathbf{u}_a - \mathbf{u}_b)$. That the attention layer's output splits into the head writes plus $\mathbf{b}_O^{(\ell)}$ is the block identity of §4.5. ∎

In plain words, center each write, shrink it by the one shared $s$, and measure how far it points from $b$ toward $a$. The pieces add up to the logit difference to the last digit.

Section 5.1 split a single logit the same way, in Exercise 5.1b. This version adds three things that matter in practice. It splits attention by head, it targets a difference, and it accounts for the biases.

The by-hand example splits the same way. Its final stream $(3, -1, 0, 0)$ has $s = \tfrac32$. With $\mathbf{W}_U = \mathbf{I}$ and unit gain, $\mathbf{w} = (1, -1, 0, 0)$.

The head's direct effect on $\Delta$ is exactly zero, because it raised Paris and Rome by the same amount (by 1 each, as the last printed line shows). This is what DLA is for. It turns "the model said Paris" into a ledger with one line per component.

A model whose answer we know

Now take a model with real attention and a known answer. The cell builds a compact version of the hand-wired model that §23.4 studies in depth. It mimics how GPT-2 finishes "When Mary and John went to the store, John gave a drink to". The right answer there is the name that appeared only once, the indirect object of "gave". Our version is a three-layer, attention-only Transformer with one head per layer. Its sequences look like BOS A B f f B f END. Here A and B are names, f is a filler token, and BOS and END mark the start and the end. The correct prediction at END is A, the name that appears once.

Each head does one job, and together they form a circuit, a set of components that computes one behavior.

• Head 1, a duplicate-token head, writes into a DUP slot at the second B that B has been seen before.
• Head 2, an S-inhibition head (S for subject, the repeated name), moves that flag to END and pushes B's logit down a little directly.
• Head 3, a name mover, attends at END to names that are not flagged and copies the one it finds into the output slot that the unembedding reads.

This version has five names, sharp attention scores (a scale of 20), and a final LayerNorm, so that the scale $s$ in the decomposition does real work. Section 23.4 rebuilds the model with ten names, softer scores, and no final norm, so that its patching numbers can be checked by hand.

The widget below runs the same split on the real tiny GPT, at the last position of the prompt. Its dashed boxes are total effects, the drop in $\Delta$ when a part's output is set to zero at every position and the model is rerun.

In practice, DLA is how you find the last step of a computation, the components that write the answer into the logits. An interviewer will ask the follow-up at once. What about the components that only tell the last step what to do? DLA cannot see them, and the rest of this section and all of §23.4 are about that gap.

import numpy as np

def dla_logit_diff(head_out, W_O, mlp_out, b_O, embed, gamma, beta, W_U, a, b, eps=1e-5):
    """One position of a pre-LN Transformer with L layers of h heads.
    head_out: (L, h, d_v) head outputs before W_O; W_O: (L, h * d_v, d); mlp_out: (L, d);
    b_O: (L, d); embed: (d,), token plus position embedding; gamma, beta: (d,); W_U: (d, V).
    Returns a dict with 'heads' (L, h), 'mlp' (L,), 'b_O' (L,), 'embed' and 'beta' (floats),
    the direct contributions to z_a - z_b, which must add up to it exactly."""

What freezing the scale hides

The decomposition is an identity, so where is the approximation? It is in the interpretation. Calling a term "the effect of $\mathbf{c}$" means imagining $\mathbf{c}$ gone with everything else held fixed, including $s$. But $s$ is computed from the whole stream, so deleting $\mathbf{c}$ changes it, and with it the size of every other term. This happens even if no later layer reacts at all, and it can be large. Two numbers decide how large.

In plain words, a component that points along the stream's own direction props up the scale. Its alignment $\rho$ is the share of the stream's squared length that it accounts for along that direction, and the alignments of all the components add up to 1. Delete it and $s$ shrinks, every surviving term grows by about the fraction $\rho$, and the logit difference can rise or fall by far more than the component's direct effect.

Go back to the by-hand example. The head's direct effect was exactly 0, yet deleting it, with the MLP's write held fixed, moves $\Delta$.

In a real model the MLP's write would change too, since the MLP reads the stream the head wrote into. The true total effect adds that change on top.

At realistic sizes the effect is usually smaller, but not always. The next cell builds a stream of six random writes in $d = 64$, one of them partly aligned with the largest. It compares each write's direct effect with the exact change from deleting it.

Small writes that are nearly orthogonal to the stream (parts 3 to 5) behave as DLA says, to within about a percent. Part 1 has a direct effect of $+0.17$. But it is aligned with the stream ($\rho = 0.37$), so deleting it lowers $\Delta$ by 2.92, seventeen times more, and the first-order correction recovers most of that. Part 0, the largest write, has a direct effect of $-4.19$, yet deleting it raises $\Delta$ by only 3.06.

The lesson for practice is to report a direct effect together with the component's $\rho$ and $\kappa$, and to measure total effects separately by intervention. A study of self-repair in Chinchilla 7B did both. It computed direct effects through the unembedding with the final RMSNorm's scale held fixed, and compared them with ablations (McGrath et al. 2023).

Reading weights in vocabulary space

The same readout works on weights as well as activations, and then it needs no input at all. Section 5.3 described an MLP as a key-value memory. Let $\mathbf{W}_2 \in \R^{d_{\text{ff}}\times d}$ be its output matrix. Row $i$ is the value vector $\mathbf{v}_i$ of hidden neuron $i$, and the MLP's write is $\sum_i m_i\,\mathbf{v}_i$, where $m_i$ is neuron $i$'s activation. Each term $m_i\mathbf{v}_i$ is called a sub-update (Geva et al. 2022). Its direct effect on the logits is $\frac{m_i}{s}\,\mathbf{v}_i\mathbf{W}_U'$. So the tokens a neuron promotes whenever it fires can be read off once, from $\mathbf{v}_i\mathbf{W}_U$, without running the model.

The cell plants three neurons that each promote one group of tokens, and adds weak background activity. Then it compares the projection of the whole update with the projections of its pieces.

The whole update's top 10 tokens are nine from one group and one from another. They hide entirely that the layer also promoted the third group. Each of the three planted sub-updates is pure. Real MLPs have thousands of neurons and far messier value vectors, but the arithmetic is the same.

Every weight that writes into the stream can be read this way. A head's OV circuit $\mathbf{W}_{OV} = \mathbf{W}_V\mathbf{W}_O$ is the product that decides what the head writes once it has chosen where to look (§4.2). It gives the $V\times V$ matrix $\mathbf{E}\,\mathbf{W}_{OV}\,\mathbf{W}_U$, whose entry $(j, k)$ says how much the head raises the logit of token $k$ when it attends to token $j$. If the diagonal dominates, the head copies the token it attends to. The copying score of §6.8 measures this from the matrix's eigenvalues, as the sum of their real parts divided by the sum of their absolute values. The product $\mathbf{E}\mathbf{W}_U$ alone is the direct path of §5.1, the scores the model would give with no layers at all, which in trained models looks like a table of bigram statistics. Dar et al. (2022) pushed this furthest. They projected every parameter of a Transformer into the vocabulary, including the attention scoring weights, and also used it to align models that share a vocabulary.

Three caveats apply to all of it. A weight projection shows what a component can write, not what it writes on a given input, where the activation $m_i$ may be zero or negative (a GELU outputs small negative values). It ignores the norm's scale and every indirect path. And the columns of $\mathbf{W}_U$ are not orthogonal, so a vector that points at one token also raises similar tokens, which can make a projection look more meaningful than it is.

Doing it at scale

On a real checkpoint you would not write any of this by hand. Here is the same pipeline with the TransformerLens library, shown but not run, since it downloads GPT-2.

The code avoids one trap. The call accumulated_resid(apply_ln=True) divides every site by the final stream's cached scale. That is right for attribution and wrong for a lens, because an early site's own scale can be several times smaller. Probabilities and KL come out wrong. Even ranks can change, because the folded bias $\mathbf{b}_U$ is added without being rescaled. So for a true logit lens, apply model.ln_final to each site afresh, as the code does.

Choosing a tool

All but the last read what the stream or the weights say to the output, without asking what would change if something were different. Use them to find candidates cheaply, then confirm with interventions.

Read more

• nostalgebraist (2020). interpreting GPT: the logit lens. The original post, with GPT-2's prediction trajectories, ranks, KL plots, and the puzzle of the vanishing input.
• Belrose et al. (2023). Eliciting Latent Predictions from Transformers with the Tuned Lens. Why the logit lens is biased and brittle, the translators that fix it, the causal checks, and the prompt-injection detector.
• Elhage et al. (2021). A Mathematical Framework for Transformer Circuits. The residual stream as a sum of component writes, direct paths, and reading weights through the unembedding.
• Geva et al. (2022). Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space. FFN sub-updates projected into the vocabulary, with interventions that use them.
• Dar et al. (2022). Analyzing Transformers in Embedding Space. Interpreting every parameter of a Transformer by projecting it into the vocabulary, without a forward pass.
• Ghandeharioun et al. (2024). Patchscopes: A Unifying Framework for Inspecting Hidden Representations of Language Models. Lenses as one case of patching a state into another forward pass, and the token-identity prompt that beats them.
• McGrath et al. (2023). The Hydra Effect: Emergent Self-repair in Language Model Computations. Direct versus total effects measured layer by layer, and why the two disagree.
• Sun et al. (2024). Massive Activations in Large Language Models. The few enormous activations that act as biases and distort anything that normalizes the stream.

Watch next

▶

▶