Reading a model layer by layer

Every answer a language model gives is assembled over many layers. Until now you have seen that process from the outside: text in, probabilities out. This lesson opens the model up. The lab below runs SmolLM2-135M, the same small model as Inside a real language model, in the same engine written for this site, which loads the model’s original weights and can record every internal number; if you ran that lab, the weights are already in your browser. Its outputs match a reference PyTorch implementation to within a few millionths. This is interpretabilityInterpretabilityThe study of how a trained model computes its outputs, by reading and intervening on its internal activations and weights.Open in glossary: studying how a trained model actually computes what it computes.

The residual stream, revisited

Recall from the transformer block that each token carries a vector, the residual streamResidual streamThe running vector for each token that flows through a transformer from the embedding to the output. Every attention and MLP sublayer reads from it and adds its result back.Taught in The transformer blockOpen in glossary, through the model. Each of the 30 layers reads it, computes something with attention and an MLP, and adds the result back. Only at the very end does the model turn the stream into a prediction: it applies a final normalization and multiplies by the output matrix to get one score per token in its vocabulary.

Nothing stops you from doing that last step early. Take the residual stream after layer 12, apply the same final normalization and output matrix, and you get the model’s guess as if it had stopped at layer 12. This is the logit lensLogit lensReading a model's next-token guess at an earlier layer by applying its final norm and output layer to that layer's residual stream, as if the model stopped there.Open in glossary, introduced by the researcher nostalgebraist in 2020 to study GPT-2.

import numpy as np

# A toy model: 4 vocabulary tokens, residual vectors of length 3.
vocab = ["Paris", "Rome", "the", ","]
W_U = np.array([[2.0, 0.0, 1.0],    # unembedding: one row per token
                [0.0, 2.0, 1.0],
                [0.5, 0.5, 0.0],
                [0.0, 0.0, 0.5]])
g = np.array([1.0, 1.0, 1.0])       # final RMSNorm weight

def rms_norm(x, eps=1e-5):
    return x / np.sqrt(np.mean(x ** 2) + eps) * g

def lens(h):
    logits = W_U @ rms_norm(h)
    p = np.exp(logits - logits.max())
    return p / p.sum()

# The residual stream at three depths for the last token.
stream = {"layer 0": np.array([0.9, 0.9, -1.4]),
          "layer 15": np.array([0.2, 1.0, 0.2]),
          "output": np.array([1.5, 0.1, 0.3])}

for name, h in stream.items():
    p = lens(h)
    best = int(p.argmax())
    print(f"{name}: {vocab[best]} {p[best]:.2f}")

layer 0: the 0.40
layer 15: Rome 0.86
output: Paris 0.88

What the lens shows, and what it cannot

In this model, the lens reads as punctuation for most of the network. Through the middle layers the top guesses at the last token are commas, spaces, and opening parentheses, and the answer climbs into view only in the last six layers. That does not mean the model has not worked anything out until then. It means that whatever it has worked out is not yet written in the form the output layer reads.

Where each head looks

Switch to Attention. Each small strip is one of the model’s 270 attention heads (30 layers of 9, numbered from 1), showing how much it attends to every earlier token from the token you selected. Select a head to see its whole pattern.

Averaged over three prompts (ignoring the first two positions, where there is little choice), three patterns stand out in this model:

• Resting on the first token. 142 of the 270 heads put more than half of their attention on the first position on average, and some, like heads 8 and 9 in layer 19, put nearly all of it there. Attention weights must add up to one, so a head with nothing useful to read for a token needs somewhere to put them. This is called an attention sinkAttention sinkA position, usually the first token, where many attention heads put most of their attention when nothing else is relevant. Attention weights must add up to one, so they need somewhere to rest.Open in glossary, and it was studied by Xiao and colleagues in 2023.
• Reading the previous token. A few heads mostly look one step back, so they are called previous-token headsPrevious-token headAn attention head that mostly looks at the position just before the current one, writing information about the previous token into each position.Taught in How circuits form: induction headsOpen in glossary. Layer 17, head 1 averages 79% on the previous token; layer 12, head 2 averages 66%. Previous-token heads are building blocks for copying patterns, which the next lesson builds from scratch.
• Reading its own token. Layer 12, head 7 puts 80% of its attention on the current position on average.

What pushed the answer

Switch to What pushed it. Because every layer only adds to the residual stream, the final stream is exactly the sum of the token’s embedding plus every attention block’s output plus every MLP’s output. With the final normalization’s scale held at its actual value, the score of " Paris" splits exactly into a contribution from each of those parts. This is direct logit attributionDirect logit attributionSplitting a token's final score into the direct contribution of the embedding, each attention head, and each MLP. The parts add up exactly to the score, but indirect effects through later layers are not counted.Open in glossary.

The view shows the contributions by layer and lists the largest individual heads and MLPs. They add up to the logit exactly, which the site’s tests check. For " Paris" (a logit of 10.58), attention in layer 24 gives the largest push of any layer (+6.27), almost half of it from head 5 alone (+2.95), and the MLPs of layers 24 to 29 each add between 1 and 3. The largest single part pushes the other way: layer 30, head 1 subtracts 4.67. Keep layer 24 in mind; it comes back when you start switching parts off.

The catch is the word direct: a head can matter by changing what later layers compute, and none of that shows up here. To measure a part’s total effect, you have to change the model and watch what happens, which is where this track goes next.