GPT-2 Logit Lens

Logit lens · GPT-2 small · 12 layersWhat the model knows, and when

Residual stream — the running prediction Attention write — gathered from other words MLP write — this layer's own contribution
How one layer works

The weights are fixed — training set them once, and running the model never changes them. What changes is the residual stream: a running vector that enters as a plain word and gets edited, layer by layer, into a next-word prediction.

France word embedding next word prediction ⋯ ×12 layers ⋯ 1 ATTENTION gather from earlier words looks back at other words read add 2 MLP recall associations reads updated stream add + + one layer — repeats ×12, same fixed weights each time
Order within a layer: ① attention reads the stream and its result is added back, then ② the MLP reads that just-updated stream and its result is added too — two reads, two adds, in sequence. (Each step is preceded by its own LayerNorm.) The same fixed layer runs again for every token.
Residual stream — the running notepad. Enters as the word, leaves as the prediction. It is the amber column in the viewer below.
Attention write — reaches back to earlier words and pulls in what is relevant: "this position should look at France." The violet column below.
MLP write — recalls associations for the current position. Layer 9's MLP voting capital cities is this step. The teal column below.
Real GPT-2 small, computed on a GPT-2 server · if it can’t be reached, runs on your device instead (one-time ~248 MB download)
Or replay a precomputed example
Layer Residual stream① Attention write② MLP write

Layer 11 detail

p(emitted token)—
Residual entropy—
Attention entropy—
MLP entropy—

Belief in the emitted token

How much the residual stream backs the token actually emitted, layer by layer.

Residual entropy (nats)

Lower means more committed. A late rise means the model hedged after having decided.

Each layer's vector is projected through the model's own final LayerNorm and unembedding, turning any point mid-stack into a distribution over the vocabulary. Layer 11's residual equals the model's real output — which is the check that the projection is being done correctly.

How this works

When you press Run, GPT-2 small (124M parameters, fp16) downloads once and then runs entirely on your device in a background thread. Each layer's vector is projected through the model's own final LayerNorm and unembedding to read out what it "predicts", with no server round-trip.

The per-layer MLP write column is why this is a from-scratch forward pass rather than an off-the-shelf runtime: no packaged GPT-2 exposes each layer's MLP output on its own. The output is verified token-for-token against a PyTorch reference before shipping.

Compute for five tokens takes a few seconds on a laptop. Layer 11's residual read-out equals the model's actual output, which is the built-in check that the projection is correct.