# Module 14: What the Model Learns

> Mechanistic interpretability basics: the residual stream, feed-forward layers as memories, induction heads, probing, the logit lens and superposition.

- Part III: Building an LLM · core track · Advanced · 4 hours
- By the end you can describe concrete circuits that trained transformers form, and read the residual stream as a communication channel.
- Builds on: [Module 12: Generation](https://transformers-and-math.drmhse.com/md/12-generation.md)
- Web page (interactive figures, labs and equation explainers): https://transformers-and-math.drmhse.com/modules/14-what-the-model-learns
- From *Transformers & Math* by Mike Chumba: https://transformers-and-math.drmhse.com

This module asks what a *trained* transformer does with its machinery: where it stores facts, how it copies patterns from its context, and how it fits more concepts than it has dimensions. The field is *mechanistic interpretability*. Its findings are partial and mostly from small models, but a few are solid, and each section says how strong its evidence is.

## 14.1 The Residual Stream as Shared Memory

Unrolled across a whole model (as section 8.5 did for Pre-LN), a token's final vector is

$$
\mathbf{x}_{\text{final}} = \mathbf{x}_0 + \sum_{\text{sub-layers } l} \operatorname{out}_l
$$

the embedding plus a sum of contributions, one from every attention head and every feed-forward layer. Nothing is overwritten; everything is added.

So the residual stream is a shared communication channel (Elhage et al., 2021). Each head and feed-forward layer *reads* a few directions through its input projections and *writes* a few through its output projection. A head in layer 5 can read what a head in layer 2 wrote, with nothing in between passing it on; different components use nearly independent directions, like conversations on different frequencies of one radio band.

Two consequences: a component's effect on the output can be measured directly from what its write adds to the final logits, and the model has an incentive to keep information in roughly fixed directions across layers, which is what makes section 14.4's techniques work.

## 14.2 Feed-Forward Layers as Key-Value Memories

Geva et al. (2021) read the feed-forward network's matrices in attention's terms. Each column of $W_1$ is a *key* whose dot product with the token vector measures a match to some pattern, often recognizable, such as a phrase or topic. The matching row of $W_2$ is a *value*, added to the residual stream in proportion to how strongly the key fired, often pushing toward particular next words.

Much memorized knowledge appears to be stored this way. Specific facts ("the Eiffel Tower is in Paris") have been located in particular middle layers' feed-forward networks, and editing a few weights changes the fact the model reports (Meng et al., 2022). The evidence is strong for some kinds of fact and weaker as a general account; knowledge is also spread across many layers and units.

## 14.3 Induction Heads and In-Context Learning

Give a trained model random tokens, then repeat them. On the first copy it can only guess; on the second its loss collapses, because it copies from its own context. The mechanism is a two-layer circuit, the *induction head* (Olsson et al., 2022):

1. **A previous-token head** in an earlier layer: every position attends to the one before and writes "the token before me was X" into its residual stream.
2. **An induction head** in a later layer: its query comes from the current token, its keys from that previous-token information, so it finds the position just *after* an earlier occurrence of the current token and copies that position's token into the prediction.

On `… A B … A`, the second `A` finds what followed the first `A` and predicts `B`: "what came next last time?"

> **Interactive figure** `induction-head-demo`, on the web page: https://transformers-and-math.drmhse.com/modules/14-what-the-model-learns#figure-induction-head-demo
>
> **Try:** The sequence is 12 random tokens followed by the same 12 again. Before stepping through, predict which earlier position the induction head will attend to when it is at position 15 (the third token of the second copy), counting from 1.
>
> **What happens:** Position 4. The current token at 15 is the third token of the sequence, which first appeared at position 3, so the head attends to position 4, the token that came next last time, and predicts it. Its whole attention pattern is a diagonal stripe offset by one less than the period.

> **In words.** A hand-built two-layer circuit on a repeated random sequence. The first layer's previous-token pattern is a stripe just below the diagonal; the second layer's induction pattern appears only in the second copy, one step past the earlier occurrence of the current token. Accuracy is at chance on the first copy and high on the second.

> **Browser lab** `induction-circuit`: an experiment that runs in the reader's browser on the web page: https://transformers-and-math.drmhse.com/modules/14-what-the-model-learns#lab-induction-circuit

The evidence is strong in small attention-only models and good in larger ones: induction heads appear abruptly in training, just as the model suddenly gets better at using long contexts, and removing them removes much of that gain. Olsson et al. argue that generalized versions, matching on meaning rather than exact tokens, underlie much of *in-context learning*: picking up a pattern from examples in the prompt with no weight change. That broader claim is plausible but less firmly established.

In-context learning is what makes prompting work: show three examples of a format and the model continues it. Nothing is trained; attention does the work.

## 14.4 Probing and the Logit Lens

Two cheap tools ask what a model's internal vectors represent:

- **Probing.** Train a small classifier, usually linear, to predict a labelled property (is this token a noun? is this statement true?) from one layer's residual-stream vectors. Success shows the information is *there* in a readable form, not that the model *uses* it; that needs an intervention, such as changing the direction and watching the output.
- **The logit lens.** Apply the final normalization and output layer directly to an intermediate layer's residual stream, reading the model's "current best guess" at that depth. In many models the right answer rises gradually from generic early guesses to the top in middle or later layers.

> **Interactive figure** `logit-lens`, on the web page: https://transformers-and-math.drmhse.com/modules/14-what-the-model-learns#figure-logit-lens
>
> **Try:** Before looking, guess at which depth the correct answer first becomes the top guess, early (layer 2 of 12), in the middle, or only at the last layer.
>
> **What happens:** In the illustrated pattern it emerges in the middle layers and then strengthens. Early layers favour common, generic tokens; the later layers mostly refine a guess that is already in place. This is the typical shape reported for GPT-2-sized models, though it varies across models and inputs.

> **In words.** One row per layer of a twelve-layer model, with the logit lens's top three guesses for a prompt whose answer is a capital city, and a bar for the correct answer's probability. The numbers are illustrative, hand-set in the pattern reported for GPT-2-sized models.

It works because components write in directions the output layer can read. Where it works less well, the *tuned lens* trains a small translator per layer.

## 14.5 Superposition

A 768-dimensional model plausibly tracks far more than 768 features: topics, grammatical roles, entities, facts. How?

$d$ dimensions hold at most $d$ exactly perpendicular directions, but vastly more *nearly* perpendicular ones: random high-dimensional directions have cosine similarity with standard deviation about $1/\sqrt{d}$ (section 3.2). If features are directions and *sparse* (few active per token), packing many along nearly perpendicular directions causes little interference. This is *superposition* (Elhage et al., 2022).

The cost is that single neurons stop meaning one thing; responding to several unrelated features is *polysemanticity*. *Sparse autoencoders* (section 28.2) trained on activations recover cleaner features, decomposing each vector into a few active entries of a much larger learned dictionary. Superposition is well supported in toy models and consistent with those dictionaries in large ones; how completely it describes real models is open.

## 14.6 Attention Sinks: A Learned Habit

Softmax forces every query's weights to sum to 1 even when nothing is relevant. Trained models park the spare attention on the first token or a few early ones, whatever they are: *attention sinks*. The first token is natural because, under a causal mask, every later token can see it.

This learned habit matters in practice: evicting the first tokens from the KV cache breaks the model, though they carry no meaning. Section 16.6 shows the fix for streaming very long inputs.

## Exercises

**Exercise (check).** In the residual-stream picture, how can an attention head in layer 6 use information written by a head in layer 2?

*Solution:* Layer 2's head adds its output to the residual stream in some direction. Nothing overwrites it, since every later sub-layer only adds. Layer 6's head reads the stream through its query, key or value projections, so if those projections pick up that direction, it reads what layer 2 wrote directly.

**Exercise (apply).** In the sequence `P Q R S P Q R S`, where should an induction head at position 6 (the second `Q`) attend, and what should it predict?

*Solution:* The current token is `Q`. Its earlier occurrence is at position 2, so the head attends to position 3, the token that followed it, and predicts `R`.

**Exercise (apply).** Random directions in $d$ dimensions have cosine similarity with standard deviation about $1/\sqrt{d}$. What is that for $d = 4096$?

*Answer:* 0.015625 (within 0.0002)

*Solution:* $1/\sqrt{4096} = 1/64 \approx 0.0156$. Two random feature directions typically overlap by under 2%, which is why thousands more features than dimensions can coexist with little interference.

**Exercise (extend).** A linear probe predicts with 95% accuracy whether a sentence is true from layer-15 activations. Does this show the model uses truthfulness when generating? What experiment would?

*Solution:* No. It shows the information is linearly readable at layer 15, not that the model's output depends on it; the probe may be reading something correlated with truth. An intervention would test it: change the activations along the probe's direction (for example push a true statement's vector toward "false") and check whether the model's later behaviour changes accordingly.

### Core Academic References
- **Elhage et al. (2021)**: *A Mathematical Framework for Transformer Circuits*. Transformer Circuits Thread.
- **Olsson et al. (2022)**: *In-context Learning and Induction Heads*. Transformer Circuits Thread. [arXiv:2209.11895](https://arxiv.org/abs/2209.11895)
- **Geva et al. (2021)**: *Transformer Feed-Forward Layers Are Key-Value Memories*. EMNLP 2021: 5484–5495.
- **Meng et al. (2022)**: *Locating and Editing Factual Associations in GPT* (ROME). NeurIPS 2022. [arXiv:2202.05262](https://arxiv.org/abs/2202.05262)
- **nostalgebraist (2020)**: *interpreting GPT: the logit lens*. LessWrong.
- **Belrose et al. (2023)**: *Eliciting Latent Predictions from Transformers with the Tuned Lens*. [arXiv:2303.08112](https://arxiv.org/abs/2303.08112)
- **Elhage et al. (2022)**: *Toy Models of Superposition*. Transformer Circuits Thread. [arXiv:2209.10652](https://arxiv.org/abs/2209.10652)
- **Bricken et al. (2023)**: *Towards Monosemanticity: Decomposing Language Models With Dictionary Learning*. Transformer Circuits Thread.
- **Alain & Bengio (2016)**: *Understanding intermediate layers using linear classifier probes*. [arXiv:1610.01644](https://arxiv.org/abs/1610.01644)
- **Xiao et al. (2023)**: *Efficient Streaming Language Models with Attention Sinks* (StreamingLLM). [arXiv:2309.17453](https://arxiv.org/abs/2309.17453)