Playing with GPT-2

The text is split into tokens and the model only ever sees their IDs (0-based; <|endoftext|> is put first):

␣ = a leading space, which is part of the token: " of" and "of" are different tokens.

One forward pass returns a logit for every vocabulary token at every position: a 50,257 × 9 matrix here. A softmax turns each column into probabilities that sum to 1. The last column is the prediction for the token after the prompt; every earlier column is the model's guess at the token that comes next at that point, which is exactly what training scores.

At each position the last layer leaves a vector of 768 numbers, x⁽¹¹⁾ (the residual stream). Three steps turn it into probabilities:

A logit difference is a log probability ratio: z(" Paris") − z(" Marse") = 0.09 (from the unrounded logits), so p(" Paris") / p(" Marse") = e^0.09 ≈ 1.09 (9.7 / 8.9). Adding the same number to every logit changes nothing, which is why a single logit has no meaning on its own. Sampling at temperature τ uses p_t ∝ e^z_t/τ: τ < 1 sharpens the distribution, τ → 0 is greedy. The entropy H = −Σ_t p_t ln p_t measures how spread out it is: 5.69 nats here, against ln 50,257 = 10.8 for a uniform guess. Training lowers the loss, Loss = −ln p(actual next token), averaged over positions.

In our exported weights γ is folded into W_U and β into a bias b_U = W_Uβ, so the code computes z = W_U (x⁽¹¹⁾ − μ) / σ + b_U. The export also centres W_U: each of its columns has its mean over the vocabulary subtracted. That shifts every logit at a position by the same amount, so the probabilities are unchanged, but the logits differ slightly from other GPT-2 implementations. The logit lens (tab 3) applies the same three steps to x^(l) of an earlier layer; direct logit attribution (tab 4) splits step 2 into one share per component.

Set the %% Settings block at the top of run_model as below, then run it. Settings not listed don't matter here: topK and seed apply only when sampling; analysisPositions and the DLA_ settings only when the logit lens or DLA is on; saveActivations stays false.

model         = 'gpt2-small';
checkpoint    = [];          % the trained model
prompt        = 'The capital of France is the city of';
nNew          = 1;           % one next token: a single forward pass
temperature   = 0;           % greedy: always the most likely token
showLogitLens = false;       % steps 5 and 6 are the next tabs
showDLA       = false;

It prints the predicted token (" Paris", 9.7%) with its ID and the runner-up (" Marse"). The logits of the forward pass are in g.out.logits [50,257 × 9], one column per prompt position: the matrix shown in Fig. 1. For the sentence in card 5, set nNew = 10. Fig. 1 uses one pass over the prompt alone, out = model_forward(ids, W). The next tabs open up the same forward pass: the logit lens reads the prediction from intermediate layers.

Following the math step by step: logit_lens_fromfile. This short script is useful for following card 3 from the final layer norm to the probability of each token. It takes the saved residual stream of one prompt and does each step in one line: center and scale (layer norm), multiply by W_U and add b_U (logits), softmax (probabilities), then prints the top 5. First run run_model with saveActivations = true (saves Activations\the_capital_of_france_is_the_city_of), then set at the top of logit_lens_fromfile:

activations_folder = 'The capital of France is the city of';
L                  = 11;      % layer, 0-based; 11 = the model's output

With L = 11 it prints the model's own prediction (" Paris" 9.72%, " Marse" 8.91%); a smaller L gives the logit lens of tab 3.

p^(l) = softmax(W_U · LayerNorm(x^(l)) + b_U)

x^(l) is the residual stream leaving layer l (768 numbers per position); LayerNorm, W_U and b_U are the model's own final layer norm and unembedding. At l = 11, the last layer, this is exactly the model's output from tab 2.

The same settings as tab 2, with the logit lens switched on:

model             = 'gpt2-small';
checkpoint        = [];          % the trained model
prompt            = 'The capital of France is the city of';
nNew              = 1;           % one next token: a single forward pass
temperature       = 0;           % greedy
analysisPositions = 'first';     % first predicted token = the last prompt position (" of")
showLogitLens     = true;        % step 5: logit lens
showDLA           = false;

Step 5 prints the top 5 at the last layer, plots their probability and rank by layer (like panels a and b) for the analysed position, and plots the top token at every layer and position (panel c). The numbers behind panel c: P = logit_lens(x, W, 0:11); gives p^(l) as a [50,257 × 9 × 12] array (token × position × layer), and [pTop, iTop] = max(P, [], 1); the top token and its probability.

The same settings as tab 2, with DLA switched on. Panels a and c are DLA_compare = 'none' (logit(" Paris") alone); panels b, d, e and f are DLA_compare = 'runner-up' (" Paris" − " Marse"):

model             = 'gpt2-small';
checkpoint        = [];          % the trained model
prompt            = 'The capital of France is the city of';
nNew              = 1;           % one next token: a single forward pass
temperature       = 0;           % greedy
analysisPositions = 'first';     % first predicted token = the last prompt position (" of")
showLogitLens     = false;
showDLA           = true;        % step 6: direct logit attribution
DLA_compare       = 'runner-up'; % panels b, d, e, f: " Paris" - " Marse";  'none' for panels a, c

Step 6 plots the shares by layer, by head, and where the top heads look; the numbers are in d (d.attn, d.mlp, d.heads [head × layer], d.b_U, d.top_heads, d.top_attn). Panel a can also be computed directly: direct_logit_attrib(g.out, W, 6342) (6342 = the ID of " Paris").

cos θ = (u · v) / (‖u‖ ‖v‖)

1 = same direction, 0 = orthogonal, −1 = opposite. The lengths are divided out, so only the direction counts. In 768 dimensions two random directions have a cosine near 0 (typically within ±0.04), so values like 0.2 are already far from chance.

Before any layer has run, words of the same category already share direction: 0.51 on average within a category, 0.26 between (Fig. 4; colours 0.64, cities 0.51). Every cosine is positive, though: the embeddings share a common component, so "unrelated" is around 0.2–0.3 here, not 0.

At the last position the France and Italy prompts are nearly identical (cosine 1.00 to two decimals) up to layer 7: the token there is " of" in both, and the country reaches it only at layers 8–9 (lowest 0.94; Japan 0.86). An unrelated sentence still scores 0.56–0.90, and 0.90 at the output (layer 11). Raw cosines between residual-stream states are compressed by a large direction they all share, so compare them with each other, not with 0.

Step 2 of tab 2, z_t = E[t] · x̂, written with the angle (γ and β folded into W_U and b_U, as in our exported weights):

z_t = √768 · ‖E[t]‖ · cos(x − μ, E[t]) + b_U(t)

LayerNorm fixes the length of x − μ (‖x − μ‖ / σ = √768 = 27.7), so a token's logit depends only on its angle with the residual stream, its own length ‖E[t]‖ and its bias. The cosine is the only part that changes with the stream: for one token, the logit is the cosine rescaled and shifted. Between tokens the scale and shift differ, so a larger cosine does not mean a larger logit. Here ‖E[" Paris"]‖ = 4.11, ‖E[" Marse"]‖ = 4.90, ‖E[" the"]‖ = 2.62.

Fig. 4: run run_cosine_tokenembed with its default settings (the same 16 words). Fig. 5: run_cosine_resid with example = 3 compares your own sentences; with three or more, it also plots the first sentence against each of the others, one line each (step 7):

model       = 'gpt2-small';
readout     = 'last';      % one vector per sentence: its last token
center      = false;       % raw cosines
example     = 3;           % your own sentences
mySentences = {'The capital of France is the city of', 'The capital of Italy is the city of', ...
               'The capital of Japan is the city of', 'I went to the shop to buy some'};

Sentences with the same number of tokens share their position embeddings; the script warns if the counts differ. example = 1 and 2 are the built-in sets (Greece vs Italy; 16 sentences in 4 topics).

Fig. 6: run run_model with the tab 3 settings (showLogitLens = true); the bottom panel of its logit-lens figure plots this cosine for the five final top tokens. The same numbers, including " the" and the embedding (emb), after run_model:

x = double(cat(2, g.out.resid_pre(:, end, 1), squeeze(g.out.resid_post(:, end, :))));   % last position: emb, then the output of layers 0-11
x = x - mean(x, 1);                                              % centered, as layer norm does
c = cosine_sim(double(W.W_U([6342 45996 262] + 1, :))', x);      % rows: " Paris", " Marse", " the"
plot(-1:11, c');  legend('" Paris"', '" Marse"', '" the"');  xticks(-1:11);  xticklabels(layer_points(12));  xlabel('Layer');

The last word often does not name the topic ("tonight", "easily", "today", "again"), so at the last word the topic has to be gathered from earlier words. Equal lengths mean every sentence has the same position embeddings.

There is one RDM per point (emb and layers 0–11): 13 matrices like panel a. Each other panel condenses them differently.

a shows one of them in full: layer 10, all 16 × 16 pairs.

b reduces every RDM to two numbers: the mean of its within-topic cells (inside the diagonal blocks of a, without the diagonal) and of its between-topic cells (outside the blocks). The layer-10 points of b are the two numbers in a's title, 0.19 and 0.23.

c compares whole RDMs with each other: cell (i, j) is the Spearman correlation between the 120 pairs of layer i's RDM and those of layer j's. Panel a is one of the 13 matrices compared; the diagonal of c is 1 (each RDM with itself).

e does what c does, but against the hand-made RDM in d instead of another layer: it is the row c would have if d were added as a 14th "layer". b and e both ask whether the arrangement follows topic; b in units of 1 − r (sensitive to the overall level), e in ranks (only the order of the pairs counts).

models          = {'gpt2-small'};   % one or more models
checkpointRange = [];               % the trained model
showL           = 10;               % layer for the RDM heatmap (panel a)
position        = 'last word';      % the word before the final "."

It prints, for every layer, the mean 1 − r within and between topics and ρ, and plots the RDM at showL (a), the within and between pairs as values and as ranks (a2, a3; a3 shows how ρ is computed and its ceiling), within vs between across layers (b, as Fig. 7b) and the layer × layer comparison (c). ρ by layer (Fig. 7e) is printed in its table. The sentences are in prompts_topic.

Sentiment is the easy contrast, passive role the hard one; tabs 8 and 9 test both with labelled readouts.

At each layer, in the full (scaled) residual stream, not just the plotted PCs: the mean distance between the group means, divided by the mean distance of a sentence from its own group mean. About 1: the groups differ as much as single sentences vary; well below 1: they overlap. This is the one step that uses the labels.

Run once per set ('topic', 'sentiment', 'passive-role'):

model    = 'gpt2-small';
contrast = 'topic';       % or 'sentiment', 'passive-role'
nPC      = 10;            % leading PCs to report
pcs3     = [1 2 3];       % the PCs for the 3D plot, as in Fig. 8 ([] = PC1 + the two most sentence-specific PCs)
position = 'last word';   % the word before the final "."

It prints, for each PC, the variance explained, its correlation with depth and its share of variance between sentences, then the distances and ratio by layer, and plots the paths in 3D (a), the variance explained (b) and the distances with their ratio (c).

Run once per contrast ('sentiment', 'passive-role', 'passive-lastword'):

model           = 'gpt2-small';
checkpointRange = [];               % the trained model (the default runs training checkpoints)
contrast        = 'passive-role';   % or 'sentiment', 'passive-lastword'
position        = 'last word';      % the word before the final "."
target          = '';               % first class (positive; predator acts) for panels a and c
showL           = 10;               % layer for panel a

It prints SI by layer for each class and plots the held-out scores at showL (a), SI by layer (b) and the layer × layer cosine of the direction (c). The sentence sets are in prompts_sentiment and prompts_passive; other contrasts in the same script: 'topic', 'strong-weak-praise', 'negation-adjective', 'negation-sentiment'.

With 768 dimensions and 20–48 sentences, a probe can separate the training sentences under almost any labelling, even a random one. Two safeguards: every score is held out (the sentence and those made from the same words are left out of the fit), and the whole analysis is repeated with the labels shuffled (20 times). The shuffled runs show what SI a probe reaches when there is nothing to find; a result inside that range is no evidence of a readout.

The four models: small (124M parameters, 12 layers, 768 dimensions), medium (355M, 24, 1,024), large (774M, 36, 1,280) and XL (1.5B, 48, 1,600). In all four the probe reads nothing in the first ~30% of depth, then rises. Small, medium and large differ little: they peak at 0.30, 0.36 and 0.39, late in the network (75–92% of depth), around the range chance reached in small. XL differs in kind: it climbs steeply between 30% and 42% of depth and holds 0.50–0.58 to its last layer, peaking at mid-depth (0.58 at layer 25 of 48). But Fig. 11 has no shuffled-label control for XL, and the score a probe reaches by chance rises with the number of dimensions (1,600 in XL against 768 in small, with the same 48 sentences), so part of XL's lead may be its greater capacity to fit rather than a readout.

Run once per contrast ('sentiment', 'passive-role'):

model           = 'gpt2-small';
checkpointRange = [];               % the trained model (the default runs training checkpoints)
contrast        = 'passive-role';   % or 'sentiment'
lambda          = 1e-5;             % ridge penalty: larger = closer to the DIM direction
nShuffles       = 20;               % shuffled-label control (grey band); 1 = a single shuffle, faster
seed            = 1;                % for the shuffles
target          = '';               % first class (positive; predator acts) for panel a
showL           = 10;               % layer for panel a
position        = 'last word';      % the word before the final "."

It prints SI by layer for the probe, difference in means and the shuffled labels, and plots the held-out probe scores at showL (a), SI by layer (b) and the cosine between the probe and difference-in-means directions (c). With 20 shuffles it took about 2 min for sentiment and 9 min for passive role on a desktop computer's CPU; most of the time is the shuffles.

Fig. 11: the same settings with contrast = 'passive-role', nShuffles = 0 and model = 'gpt2-small', 'gpt2-medium', 'gpt2-large', 'gpt2-xl' in turn (1–6 min each); after each run:

depth = 100 * (1:nL) / nL;       % output of layer L of nL at (L + 1)/nL × 100%
plot(depth, SIp(2:end, 1));  hold on   % probe SI, emb left out (no direction for role)

Write a short list of unrelated words twice, "Code: blue seven turtle elbow spoon granite.", and give run_model the second copy cut off at different points. Nothing in English predicts these words, so the first time they appear GPT-2 small cannot guess them. The second time it can: the words came in this order earlier in the prompt, and it copies them. Generating on from the second "Code", it writes the rest of the sentence word for word: ": blue seven turtle elbow spoon granite. Code".

"First time" is the same point in the first copy (the prompt "Code", "Code:" or "Code: blue" alone). ":" after "Code" is common English, so it is likely even the first time, and copying raises it further; " blue" and " seven" are unpredictable until they can be copied. The runners-up are of the same kind (other colours, other numbers): the model also knows what sort of word fits, and copying picks out the exact one.

prompt        = 'Code: blue seven turtle elbow spoon granite. Code';   % or end the prompt with 'Code:' or 'Code: blue'
nNew          = 1;          % 9 to generate on: ": blue seven turtle elbow spoon granite. Code"
showLogitLens = false;
showDLA       = false;

Copying "[A][B] … [A] → [B]" takes two heads in different layers. A previous-token head writes into every position which token came just before it, so the position of [B] now also carries "after [A]". Later, at the second [A], an induction head looks for a position marked "after [A]", finds [B], and moves [B] into the prediction. The first head must sit in an earlier layer than the second.

On the first copy nothing can be predicted (except ":" after "Code", 72%); on the second, the model gives the right word 29–96% (" turtle" after " seven": 80%). Head 4.11 attends to the previous token at every position (1.00 on average; Fig. 13e): the first step. Head 5.5 has nothing to do on the first copy and parks its attention on <|endoftext|>; on the second copy each token attends to the token right after its earlier copy (0.91 on average), the stripe one step off the diagonal (Fig. 13d): the second step. Fig. 13a shows the step from the second " seven": 5.5 puts 89% of its attention on the first " turtle" yet adds only +0.30 directly to the " turtle" logit (17th of the 144 heads). The heads with the largest direct effect sit in layers 9–11 and also look at " turtle" (10.0: +1.79), and two of them push it down (11.10, 10.7). So the induction head finds the token to copy; much of the writing into the prediction is done by later heads (DLA: tab 4).

Fig. 13a and the attention patterns (Fig. 13d, e) show where heads look. To see how a head decides and what it brings back, run_attention_head recomputes one head at one token by hand, step by step, from the weights, and checks every step against the model's own run. Here: head 5.5 at the second " seven" of card 1's sentence (the outlined row of Fig. 13d), in the steps of attention (Lesson 2, tab 3), numbered as in the script and the figure:

• Steps 1–2 · Query and keys. The head reads the stream entering layer 5 (layer-normed). Its query matrix W_Q turns the current token's stream into a 64-number query, "what am I looking for"; its key matrix W_K turns every earlier token's stream into a key, "what do I offer".
• Step 3 · Scores. The query's dot product with each key, divided by √64 = 8 (Fig. 13b). " turtle" scores highest (4.31), <eot> (<|endoftext|>) next (1.87), the rest below 0.5.
• Step 4 · Softmax. The scores become attention weights, all positive and summing to 1 (Fig. 13c). Done for every query at once, this gives the head's full attention pattern (Fig. 13d), of which the second " seven" is the outlined row. Exponentiating turns a lead of 2.4 into a factor of 11: 0.89 on " turtle", 0.08 on <eot>.
• Steps 5–6 · Values and output. Each position's value (W_V) is what it passes on if attended to; the head adds them up in proportion to the weights and writes the result into the stream through W_O.
• Step 7 · Direct effect on the answer. The head's output, read along the target token's unembedding direction (DLA, tab 4), split by the position it came from (Fig. 13f, g).
• Step 8 · Which earlier head it depends on. An earlier head is silenced and the attention recomputed (card 5; not in the figure).

Card 1: run_model with the settings above. Fig. 12: one run of the full repeated sentence, then the probability of each token that actually came next:

ids = model_tokenize('Code: blue seven turtle elbow spoon granite. Code: blue seven turtle elbow spoon granite.', W);
out = model_forward(ids, W);
Z = double(out.logits);  P = exp(Z - max(Z, [], 1));  P = P ./ sum(P, 1);   % softmax at every position
pNext = P(sub2ind(size(P), ids(2:end) + 1, 1:numel(ids) - 1));             % p(the token that actually came next)
bar(0:numel(ids) - 2, 100 * pNext)

Fig. 13d, e, the attention of any head, after out = model_forward(ids, W) (layer L, head h, 0-based):

A = out.attn_pattern(:, :, h + 1, L + 1);   % [query × key]
imagesc(A);  axis square;  xlabel('Key');  ylabel('Query')

Fig. 13a: the DLA plot of run_model with DLA_extra_heads = [5 5; 4 11], at the second " seven": run_model analyses the last position, so end the prompt there (prompt = 'Code: blue seven turtle elbow spoon granite. Code: blue seven').

Fig. 13b, c, f, g: run_attention_head with its default settings (and silenceHead = [], no ablation, for the figure); it prints each step as a table and checks it against the model's own run:

prompt      = 'Code: blue seven turtle elbow spoon granite. Code: blue seven turtle elbow spoon granite.';
L = 5;  hd = 5;          % the head: layer L, head hd
queryPos    = 13;        % the token to examine (0-based position; printed). 13 = the second " seven"
target      = [];        % token for the DLA; [] = the token that comes next in the prompt
silenceHead = [4 11];    % silenced in step 8: [layer head] = one head; [layer] = all heads of that layer; [] = skip

Fig. 14: run_induction with model = 'gpt2-small' and checkpointRange = [] (the trained model; 20 sequences of 50 random tokens, seed = 1). It prints the loss on both copies and the best induction and previous-token heads, and, for the trained model, plots the three panels of Fig. 14: the loss on each token, and the prefix-matching and previous-token score of every head (also in prefixAll and prevAll; given checkpoints instead, it plots them against training step).

Fig. 15: settings in run_ablation; with several layers listed it reports each one on its own (the total effects):

model        = 'gpt2-small';
prompt       = 'The capital of France is the city of';
ablate_type  = 'mlp';           % or 'attn'
ablate_which = 0:11;            % every layer, one at a time
ablate_how   = 'mean';
target       = ' Paris';
competitor   = ' Marse';

The direct effect of mean-ablating MLP L, from the clean run c = model_forward(ids, W):

d  = double(W.W_U(6342 + 1, :) - W.W_U(45996 + 1, :))';     % E[" Paris"] - E[" Marse"]
xf = double(c.resid_post(:, end, 12));  xf = xf - mean(xf);  s = sqrt(mean(xf.^2) + 1e-5);   % final layer norm scale
m  = double(c.mlp_out(:, :, L + 1));  dx = mean(m(:, 2:end), 2) - m(:, end);            % mean output - actual, last position
direct = d' * (dx - mean(dx)) / s

Fig. 16: settings in run_induction (the sequences of Fig. 14, tab 10); with an ablation set, panel a draws the clean loss and the ablated one. Run it twice, once per set of heads:

model           = 'gpt2-small';
checkpointRange = [];
ablate_type     = 'head';
ablate_which    = [5 5; 6 9; 5 1; 7 10; 7 2];   % top 5 induction heads; random: [3 5; 9 1; 4 8; 2 3; 0 6]
ablate_how      = 'mean';

Card 6: the same with the 1, 2, 3 … 8 best induction heads (in the order of tab 10's prefix-matching score); the second-copy loss is the one run_induction prints.

The clean run (France) puts " Paris" at rank 1 (p = 9.7%; logit(" Paris") − logit(" Rome") = +4.80); the corrupted run (Italy) puts it at rank 336 (p = 0.04%; −4.87). (Italy's own top answer is " Milan", 6.1%, with " Rome" second, 5.2%: GPT-2 small gets this fact wrong, but patching compares " Paris" with " Rome", so it does not matter here.) The prompts differ only at position 4, so everything that makes the model say " Paris" must start from that one token. Denoising runs the Italy prompt with one piece copied in from the France run (one point of the residual stream, one attention or MLP write, or one head) and asks how far that piece alone brings " Paris" back, scored on its rank (Fig. 17 caption).

Read at its own position. At " France" one write matters: MLP 0 (0.88; panel b). The MLPs after it add little one at a time (0.03–0.06; MLP 8 0.18). This matches tab 11, where mean-ablating MLP 0 was the only removal that broke the answer: MLP 0 turns the token into a usable "France", like a second embedding. Kept there. From the embedding to the output of layer 7, copying the state at " France" restores everything (1.00; panel a): until then the country is only at its own position. Moved to the last token. At layers 8 and 9 the France column fades (0.88, then 0.39) and the " of" column takes over (0.18, then 0.88); after that, patching " France" does almost nothing. The move is attention at " of": layer 9 (0.49), helped by 8 (0.14) and 10 (0.27), and within it one head, 9.8 (0.53, panel c), the head tab 4 found with the largest direct share and 76% of its attention on " France". Head 10.0 helps (0.34); 10.7 works slightly against it (−0.08).

Other metric and direction. With the logit difference " Paris" − " Rome" the map is nearly the same (MLP 0 0.80, attention 9 0.48, head 9.8 0.51), and noising (run France, copy in one piece from Italy) correlates 0.98–0.99 with denoising; its largest gap, head 9.8 (0.53 denoising, 0.31 noising), shows other heads partly covering for it. Where are Meng et al.'s middle-layer MLPs? Lesson 7, tab 4 showed the subject read by middle-layer MLPs. In GPT-2 small they are many small effects: patched together, MLPs 4–8 bring back 0.42 and MLPs 1–11 0.53, one at a time at most 0.18. Fig. 18 repeats Meng et al.'s own experiment in GPT-2 XL, with noise on the subject instead of a second prompt, and shows both of their sites: middle-layer MLPs at the subject (windows centred on layers ~10–26 bring " Paris" fully back, b) and attention at the last token (" of", windows from layer ~24 to the top, a). Two details differ from their picture. The subject token that matters most is "el", the last piece of "Eiffel", not the subject's last token " Tower" (0.3–0.7); and restoring the last MLPs at " of" pushes " Paris" further down (−1.0 in log rank, b; only −0.3 in logit difference, d). Restoring the MLPs at all four subject tokens together (last column) brings " Paris" back at almost any depth up to layer ~28. The windows matter: in a counterfactual run in XL (France → Italy), single layers barely register (no middle MLP above 0.1, no attention layer above 0.2); windows add up effects that each layer has only a little of, which is why Meng et al.'s figure shows bands. The metric matters too: with the logit difference (c, d) the same bands are weaker (0.4–0.7), and with the probability of " Paris" weaker still (0.1–0.4 at "el"; only the four subject tokens together bring it fully back).

model          = 'gpt2-small';
corruption     = 'counterfactual';
clean_prompt   = 'The capital of France is the city of';
target         = ' Paris';
competitor     = ' Rome';
corrupt_prompt = 'The capital of Italy is the city of';
direction      = 'denoise';     % 'both' adds the noising maps and their correlation (card 3)
doResid = true;  doMLP = true;  doAttn = true;  doHeads = true;
layer_window   = 0;
token_group    = [];            % 'corrupted' adds a column with the corrupted tokens patched together
metric         = {'rank', 'logit_diff'};   % Fig. 17 is the 'rank' figure

The card 3 windows for GPT-2 small: layer_window = 2 and 5. Fig. 18 (about 2.5 hours on a CPU):

model        = 'gpt2-xl';
corruption   = 'gaussian_noise';
clean_prompt = 'The Eiffel Tower is located in the city of';
target       = ' Paris';   competitor = '';     % '' = the noised run's top token (" New")
subject      = ' Eiffel Tower';
noise_scale  = 3;   nSamples = 4;   seed = 1;
layer_window = 5;              % attention/MLP restored over layers L-5 ... L+5
direction    = 'denoise';
doResid = false;   doMLP = true;   doAttn = true;   doHeads = false;
token_group  = 'corrupted';    % the last column: all subject tokens together

Positions (0-based; the labels of Wang et al.):

• IO = the indirect object, " Mary" (position 2). It appears once, and it is the right answer.
• S1 = the first mention of the subject, " John" (position 4), in "When Mary and John went…".
• S2 = the second mention of the same subject, " John" (position 10), in "…the store, John gave a drink to".
• END = the last token, " to" (position 14), where the model predicts the answer.

Metric: LD = logit(" Mary") − logit(" John") at END: +3.36 for the clean prompt. Corrupted prompt: S2 replaced by " Mary" ("…, Mary gave a drink to"): now Mary is the repeated name and the model prefers " John" (LD −2.48). Three measurements: where the heads look from END (Fig. 19); each head's direct contribution to LD at END (Fig. 20a); and denoising one head at one position: run the corrupted prompt, copy in that head's output from the clean run at END (Fig. 20b) or at S2 (Fig. 20c), and score (LD_patched − LD_corrupt) / (LD_clean − LD_corrupt), 1 = that head alone restores the clean answer.

The name movers 9.6 and 9.9 each add +2.1 to LD, 10.0 +0.8 (Figs. 20a and 19). From END they look at " Mary" (0.67, 0.67, 0.33 of their attention; at most 0.11 on either " John") and copy the name they see into the output. Negative name movers do the opposite: 10.7 also looks at " Mary" (0.82) but writes −2.3 against it, 11.10 −0.8, so the model hedges its answer (Wang et al. found the same heads).

The name movers look at " Mary" and not at " John", but nothing in a name mover's own input at END says which name was repeated. That information comes from four heads in layers 7 and 8, 7.3, 7.9, 8.6 and 8.10, which Wang et al. called S-inhibition heads (S for subject: they inhibit the name movers' attention to the subject). Where they look: from END, at S2, the repeated " John" (Fig. 19: 8.6 79%, 8.10 46%, 7.9 41%, 7.3 13%), and almost not at all at " Mary". What they read there: by layer 7 the stream at S2 carries the mark "this name appeared before", written by the heads of card 4; the S-inhibition heads copy it to END. What they do with it: they do not write a name themselves (their direct contributions are small, +0.05 to +0.41, the numbers in brackets in Fig. 19). Instead their output changes the queries of the name movers in layers 9–10, so those queries match " John" less and " Mary" more. One head's output shaping a later head's query is called Q-composition, the query-side counterpart of the K-composition between the previous-token and induction heads of tab 10. Wang et al. found that the signal they pass on is partly the name itself and mostly its position.

The evidence. The name is confusing: the S-inhibition heads themselves attend to the subject (Fig. 19); what they inhibit is the name movers' attention to it. Fig. 21 tests this directly, with the name movers' attention from END to each name in four runs (a–c) and the answer (d). In the clean run they look at " Mary" (0.67, 0.67, 0.33) and hardly at " John". In the corrupted run, where " Mary" is the repeated name, they flip and look at " John" (9.9: 0.84), and the model answers " John" (LD −2.48). Copying in only the four S-inhibition heads' clean output at END flips them back (0.69, 0.71, 0.33 on " Mary"), although everything else in the run still says " Mary" is repeated, and LD returns to +3.18. Mean-ablating the four heads at END in the clean run does the opposite: the name movers split their attention between the two names (9.9: 0.27 on " Mary", 0.35 on " John") and LD falls from +3.36 to +1.40. So the decision "avoid the repeated name" reaches the name movers through these four heads. One at a time they matter less (Fig. 20b: 8.10 0.24, 8.6 0.21, 7.9 0.17, 7.3 0.14, despite their small direct effects); together they restore 0.97, more than the sum of their single scores (0.76).

The S-inhibition heads of card 3 read a mark at S2 that says "this name appeared before". Fig. 20c finds the heads that write it: one head at a time, its clean output copied in at S2 in the corrupted run. These heads do not stand out in Figs. 19 and 20a, which look only at END: they work at S2, and only indirectly (what they write at S2 is read by the S-inhibition heads, not turned into the answer). Fig. 22 shows where they look from S2. Two kinds of head, both in layers 0–6:

• Duplicate-token heads (3.0, 0.1). At the second " John" they look back for an earlier copy of the same token and find the first " John" (Fig. 22; 3.0: 71% of its attention, 0.1: 52%). What they write there amounts to "this token has occurred before". 3.0 alone restores 0.10.
• Induction heads (5.5, one of the induction heads of tabs 10–11). At the second " John", 5.5 looks at the token right after the first " John", " went" (92%): the same move it makes in a repeated sequence, where it predicts what came next last time. Here what matters is that it found an earlier " John" at all: another sign that this name is a repeat. 5.5 alone restores 0.18, the most of any head at S2. 6.9 (0.08) and 5.9 (0.06) add a little; their attention from S2 is less specific (6.9 looks mostly at <|endoftext|>, 5.9 is spread out), so their role here is less clear.

Copied in together, these five heads (0.1, 3.0, 5.5, 5.9, 6.9) restore 0.56 of the answer. This is where the circuit starts: without the mark, the S-inhibition heads would have nothing to carry.

1. At S2, layers 0–6: duplicate-token and induction heads mark the second " John" as a repeat (card 4; Figs. 20c, 22).
2. At END, layers 7–8: S-inhibition heads read that mark and carry it to the last token (card 3; Figs. 19, 21).
3. At END, layers 9–10: name movers look at the name that is not marked, " Mary", and copy it into the answer (card 2; Figs. 19, 20a).
4. At END, layers 10–11: negative name movers (10.7, 11.10) also look at " Mary" but write against it, so the answer is hedged (card 2).

Nothing in this is specific to Mary and John: "copy the name that was not repeated" works for any pair. The same heads come out for 8 other name pairs in the same sentence (Alice/Bob, Sarah/Tom, … Rachel/Michael): in every one the top three name movers are 9.6, 9.9 and 10.0 (once 10.10 in place of 10.0), the negative ones 10.7 and 11.10, the top heads at END include 8.6, 8.10 and 7.9, and at S2 5.5 comes first.

model          = 'gpt2-small';
clean_prompt   = 'When Mary and John went to the store, John gave a drink to';
corrupt_prompt = 'When Mary and John went to the store, Mary gave a drink to';   % S2 -> the IO name
io_name        = ' Mary';
s_name         = ' John';
nShow          = 3;            % heads per class to print
sInh           = [7 3; 7 9; 8 6; 8 10];   % Fig. 21: the S-inhibition heads
movers         = [9 6; 9 9; 10 0];        % Fig. 21: the name movers

It prints the top heads of each class and plots the three panels of Fig. 20, then the S-inhibition test of Fig. 21 (its step 5); this takes about 1 minute on a CPU (288 patched runs). The joint scores (0.97, 0.56) copy the listed heads in together at the same position, with the same score.

Fig. 19 (and Fig. 20a) come from run_model (tab 4's direct logit attribution, here with " John" as the comparison token): its last panel, where the six heads with the largest |DLA| and the heads in DLA_extra_heads look from END, is Fig. 19; its head heat map is Fig. 20a.

prompt          = 'When Mary and John went to the store, John gave a drink to';
nNew            = 1;           % one next token (" Mary")
temperature     = 0;           % greedy
analysisPositions = 'first';   % the last prompt position (END)
showDLA         = true;
DLA_compare     = ' John';     % DLA of logit(" Mary") - logit(" John")
DLA_nTop_heads  = 6;           % last panel: the 6 heads with the largest |DLA|
DLA_extra_heads = [7 3; 7 9; 8 6; 8 10];   % Fig. 19: the four S-inhibition heads added

Fig. 22 is the same run_model panel from S2, with only the five chosen heads (the lesson figure leaves out the DLA numbers, which at S2 concern the next word, " gave"):

analysisPositions = 10;        % S2, the second " John"
DLA_nTop_heads  = 0;           % no top-DLA heads ...
DLA_extra_heads = [0 1; 3 0; 5 5; 5 9; 6 9];   % ... only these five

Direction: f = mean(x | positive sentence) − mean(x | negative sentence), the difference in means of tab 8, here with each sentence's vector the mean over its tokens, at one point of the residual stream. Steering: run the test prompt with α · f added to that point at every position (except <|endoftext|>), so every later layer reads the shifted stream. α = 1 adds the full difference between the two groups' means; α < 0 pushes toward negative. f is small: at the output of layer 5 its length is 0.13 of the stream's typical norm, so α = 4 adds about half the stream's norm. Readout: the next-token sentiment score, the mean logit of 5 positive words (" great", " good", " wonderful", " excellent", " amazing") minus 5 negative ones (" terrible", " bad", " awful", " horrible", " boring"); and a 20-token continuation, choosing the most likely word each time (greedy). Control: 10 shuffled copies of f: its own values in a random order, so the same length and the same values but a random direction (almost perpendicular to f).

Each layer adds its attention and MLP writes to the stream; the push is one extra term, added once (here at the output of layer 5) and then carried to the end.

Layers before the push (0–5) write exactly what they wrote before. The push stays in the stream: each later layer only adds to the stream, so α f is still in x̃⁽¹¹⁾ and reaches the logits directly through the unembedding (the direct path of tab 4). Later writes change (the tilde): layers 6–11 read x + α f instead of x, so their attention patterns and MLP outputs shift; the push is added at every position, so attention at the last token also reads pushed vectors from the earlier tokens. This is the indirect path, the model's own response to the push. The change in the logits is the sum of the two. Because each layer reads its input through LayerNorm, what counts is the push's size relative to the stream (0.13 · α of its norm at layer 5). During generation the same push is added at every forward pass, so each new token is steered too.

The score follows α, and only along f. At the output of layer 5 the score rises steadily from −2.8 (α = −8) through 1.7 (no steering) to 7.5 (α = +8); shuffled copies of f, with the same values in other dimensions, move it little and in no consistent direction (mean 1.0 to 1.7, panel b), so what steers is f's direction, not its set of values. The text follows too (greedy decoding: the most likely word each time, so no randomness):

Toward negative the text turns step by step, from "a bit too much for me" to "a waste of time", "a joke" and "worse than the worst" (the repeated sentences are greedy decoding's usual loops in GPT-2 small). Toward positive every text starts "a great": unsteered, the model already leans positive, and the most likely words stay the most likely. From α = +4 the text leaves the movie (a book, a gift for your home): the push is then half the stream's norm or more, and the model leaves its normal range. Where it works (panel c): the same push changes the score more the later it is added, from 1.2 at the embedding to 8–10 at layers 5–7 and 14.6 at layer 9. Late layers are close to the output, so a late push acts on the next word more directly (as in tab 4's direct logit attribution); a push in the middle layers has to pass through the layers after it, and still works. The shuffled copies stay near zero at every layer. So the direction a linear readout finds (tabs 8 and 9) is one the model's later layers act on: sentiment is not only readable but used.

The same method with a factual concept. Direction: f from 10 matched pairs that differ only in the country ("I spent the summer in Italy." / "I spent the summer in France."; no capital is named), at the output of layer 5, pointing from France to Italy. The sentences separate cleanly along f (Italy 8.6 to 11.2, France 3.0 to 5.1). Test: "The Eiffel Tower is in the city of", which names no country; score = logit(" Rome") − logit(" Paris"). Unsteered, GPT-2 small gets it wrong: its top answer is " London", with " Paris" second.

Steering toward France corrects the model (" Paris" from α = −1); steering toward Italy moves the tower to Milan (α = +2, +4) and then Rome (+8), with " Rome" passing " Paris" between α = +1 and +2. At −8 the push is too large and the text leaves the subject. Ten shuffled copies of f move the score much less and in no consistent direction (−2.7 unsteered; between −3.5 and +0.8 at every α). So steering is not specific to sentiment: a direction for a country, found the same way, moves where the model places a landmark.

model       = 'gpt2-small';
L           = 5;            % the output of layer 5 ('emb' = the embedding)
readout     = 'mean';       % each sentence's vector: the mean over its tokens
concept     = 'sentiment';  % card 4: 'country' (it sets the contrast pairs, test prompt and score words)
test_prompt = 'I think this movie is';                    % card 4: 'The Eiffel Tower is in the city of'
pos_words   = {' great', ' good', ' wonderful', ' excellent', ' amazing'};   % card 4: {' Rome'}
neg_words   = {' terrible', ' bad', ' awful', ' horrible', ' boring'};       % card 4: {' Paris'}
alphas      = [-8 -4 -2 -1 0 1 2 4 8];
nNew        = 20;
temperature = 0;  topK = 40;  seed = 1;   % 0 = greedy (topK unused; seed only for the shuffles)

It prints the score and a continuation for each α and plots panels a and b.

The shuffle control (panels b and c): set nShuffles = 10 (default 0, no control). Step 6c then makes 10 shuffled copies of f, its 768 values in a random order (made once with seed and used for every α): the same length and values, but a random direction, almost perpendicular to f (cosine about ±0.04). It scores the same push along each, and panel b shows their range and mean in grey. Every number in panels a and b is what run_steering gives with these settings.

Panel c is run_steering with L swept over 'emb', 0, …, 10 (not 11: the output of layer 11 is the model's final stream, with no layers left after it), each time with alphas = [-4 4] and nShuffles = 10. f is rebuilt at each L (steps 2–3), and only the score of step 6a is used, not the text. Each point is score(α = +4) − score(α = −4): black for f, grey for the 10 shuffled copies (mean ± SD).

nShuffles = 10;        % shuffled copies of f (0 = no control)
% each copy, made once and used for every α:
rs = RandStream('twister', 'Seed', seed);
fShuf(:, r) = f(randperm(rs, numel(f)));

At the subject's last token (" Tower"), an MLP's hidden activation k (after GELU: the vector that W_out multiplies) can be read as a key, and what the MLP writes into the residual stream, v = W_out k + b_out, as the value stored under it. Later layers turn that value into the answer. To change the fact, change what W_out returns for this one key, and as little as possible for every other key. The smallest change that does this is rank one: a single outer product (a column times a row) added to W_out.

Notation. k: any key (the MLP's hidden activation at any token); k*: the subject's key, averaged over many contexts (step 1); k_E: the key at the edit prompt itself, with no prefix. The star follows the paper, which writes the inserted pair as (k*, v*): the key to edit and the new value it should return (here v* = old value + δ).

1. The key k*. The edit prompt is run behind 20 short texts that GPT-2 XL writes itself (ten of 5 tokens and ten of 10, such as "You can't have a good game without a good soundtrack. "), and once without a prefix: 21 prompts. Each gives one vector: layer 17's MLP hidden activation (after GELU; 6,400 numbers) at the subject's last token, " Tower" (the subject is four tokens, " E", "iff", "el", " Tower"; Meng et al.'s causal tracing finds facts read at its last one; in Fig. 18 the last piece of "Eiffel", "el", matters even more). k* is the average of these 21 vectors, so the edit does not depend on the exact wording before the subject. Only the text before " Tower" changes its key (attention is causal), which is why the prompts differ in their prefixes. The paper's Section 3.1 describes averaging k* over 50 generated texts of 2 to 10 tokens; its Appendix E.5 and its released code, which we follow, use 20 (ten of 5 tokens and ten of 10) plus the prompt alone: the same 21 prompts in steps 1 and 2.
2. The new value. The old value v0 is what layer 17's MLP writes at " Tower" in the edit prompt now (1,600 numbers). The new value is v* = v0 + δ, where δ (also 1,600 numbers) starts at zero and is found by gradient descent: add δ to the MLP's output at " Tower", run layers 18–47 to the end, and measure the loss
   Loss(δ) = −(1/21) Σ_j log p(" Rome" | prompt_j, δ) + 0.0625 · KL(p_δ ‖ p_before) + 0.5 · |δ| / |v0|²

   First term: make the model say " Rome": p(" Rome" | prompt_j, δ) is the probability that the next token is " Rome" after each of the 21 prompts ("…The Eiffel Tower is located in the city of"), with δ added.
   Second term: keep what the tower is. It uses a second prompt, the essence prompt "The Eiffel Tower is a", which names the subject without asking where it is. p_before and p_δ are the model's whole next-token distributions after it (50,257 probabilities each), without and with δ. KL is the Kullback–Leibler divergence, Σ_w p_δ(w) log(p_δ(w) / p_before(w)), with the edited distribution first, as in the paper: 0 if δ leaves the distribution unchanged, larger the more δ makes words likely that were not likely before. Without it, δ could reach " Rome" by changing what the model thinks the tower is (the paper calls this essence drift).
   Third term: keep δ small, relative to the old value v0. This form, |δ| rather than |δ|², is the paper's released code; it is a weak penalty, so in practice δ is limited mainly by a cap in the same released code (clamp_norm_factor = 4), |δ| ≤ 4 |v0|.
   Backpropagation through the frozen layers 18–47 gives the gradient of the loss with respect to δ, and Adam updates δ; 20 steps. No weight changes in this step: like training, but δ is the only thing being trained. In our run p(" Rome") reached 0.99 by step 8, and δ grew until it hit the cap (4 |v0|) at step 9; the remaining steps reshaped δ at that length, bringing the essence penalty down from its peak of 2.0 (step 7) to 0.45 (Fig. 24c).
3. The key covariance C = the mean of k k^⊤ over 100,000 tokens of Wikipedia text: which key directions ordinary text uses. (A tiny ridge is added, so C can be inverted where some directions never occur.)
4. The update ΔW = δ u^⊤ / (u^⊤ k_E), with u = C⁻¹k*. Its column δ is the change to the output; its row u makes it respond along the subject's key but away from the key directions that other text uses a lot; the scale makes it exact at the edit prompt's own key k_E (ΔW k_E = δ; the paper's equation uses k* here, its released code k_E). For any token, the edited MLP then writes (W_out + ΔW) k = W_out k + δ · (u^⊤ k) / (u^⊤ k_E): its old output plus δ, scaled by how much that token's key resembles the subject's. That is exactly δ at k_E, and close to nothing for keys from unrelated text. In our run |ΔW| is 7.4% of |W_out|.

How the steps fit together. Steps 1 and 2 are independent: step 1 finds which address to change (k*), step 2 what to store there (δ). Step 2 does not use k*: it adds δ to the MLP's output at " Tower" in each prompt, whatever key that prompt produces, and uses only the old value v0 at the edit prompt (from k_E) to set the scale of its size penalty. Step 4 combines them, with C from step 3: k* and C give the edit's row (u = C⁻¹k*), δ its column, and k_E its scale.

Sizes in GPT-2 XL, layer 17 (vectors as columns; in GPT-2 small the key has 3,072 numbers and the stream 768):

Then the edited model is tested on the edit prompt, on four paraphrases it was not trained on (does the new fact generalize?), on four neighbouring facts that should not change (is it specific?), and with greedy continuations.

The edit works, and it is specific (panels a, b). After the edit the model answers the edit prompt with " Rome" (1.00; " Paris" falls from 0.96 to 0.00), and the four neighbours keep their answers (the Louvre 0.955 → 0.941, the Colosseum 0.84 → 0.90, Big Ben 0.83 → 0.81; Notre-Dame stays near 0). |ΔW| is 7.4% of |W_out|, and δ ends at its cap, 4 |v0|: the cap, not the size penalty, sets how large the edit is. Fig. 24c shows how δ is found: p(" Rome") reaches 0.99 by step 8; the KL penalty rises with it to 2.0, then the optimizer brings it down to 0.45 while δ stays at the cap, so δ ends up moving the tower while changing "The Eiffel Tower is a…" far less than at its peak.

It generalizes. All four paraphrases flip to " Rome" (0.55, 0.86, 0.79 and 0.91), and " Paris" falls to 0.007 or less on each, including "Tourists come to see the Eiffel Tower in", where it had 0.91. In the model's own words (greedy):

Consequences follow only partly. The tower is now in Rome and even "in the country of Italy", one reasoning step away, although the text then drifts (to the Vatican, "in the United States"). Asked the other way round, the most famous tower in Rome is still the Colosseum: the reverse direction does not follow. This fits the paper's benchmark averages (about 96% of paraphrases flip for ROME in GPT-2 XL), and single facts vary a lot around such averages. How the edit is learned matters too: with 5 fixed prefixes instead of the paper's 20 generated ones, the same layer-17 edit gives the paraphrases 0.61 on average (Fig. 25) instead of 0.78.

Where to edit (Fig. 25, with 5 fixed prefixes to keep the 13 edits short). The same edit works at almost every layer (the edit prompt reaches 0.96–1.00; only at layer 44, close to the output, is it weaker: 0.62) and stays specific everywhere (neighbours 0.63–0.68, against 0.66 before). But how far it generalizes depends on the layer, and not the way causal tracing suggests: tracing peaks at layer 12 (0.57), where the edit generalizes poorly (paraphrases 0.16), while it generalizes best at layers 20 and 24 (0.80 and 0.76), where tracing has fallen to 0.28 and 0.12. Edits in the first layers also generalize (0.44 at layers 0 and 4) although tracing is near 0 there, but they need a much larger change (|ΔW| is 19% and 12% of |W_out| there, about 7% from layer 12 on). This is consistent with Hase et al. (2023), who found over many facts that how strongly tracing localizes a fact to a layer barely predicts how well editing that layer works.

• Generalization depends on the setup. Here all four paraphrases flip with the paper's 20 generated prefixes, less strongly with 5 fixed prefixes, and at layers 8, 12 and 32–44 only one to three of the four do (Fig. 25). The paper's benchmark reports averages over many facts (about 96% of paraphrases flip); a single fact can do much worse.
• Consequences follow only partly. Here the tower moved to Italy as well as to Rome, but the reverse direction ("the most famous tower in Rome is…") keeps the old answer, the reversal curse of Lesson 6, tab 4. Cohen et al. (2024, "Evaluating the Ripple Effects of Knowledge Editing") found this for ROME and other editors.
• Specificity is only tested where we look. Four neighbours are checked; an edit can change facts nobody tests, and there is no general way to verify that nothing else broke.
• Where a fact is found is not where it is best edited. Hase et al. (2023, "Does Localization Inform Editing?") found that causal tracing barely predicts which layer gives the best edit. Fig. 25 compares the two for this fact.
• Edits do not stack well. One edit is small; many sequential edits gradually damage the model and undo earlier edits (Gupta et al. 2024). MEMIT (Meng et al. 2023) spreads many edits over several layers instead.
• It assumes one memory. ROME treats one MLP layer as the store of the fact; tab 12 found facts read by many middle-layer MLPs together, so a single-layer edit changes only one of several contributions.

model        = 'gpt2-xl';
L            = 17;              % the MLP layer to edit (the paper's choice for GPT-2 XL; it counts from 1: layer 18)
prompt       = 'The Eiffel Tower is located in the city of';
subject      = ' Eiffel Tower';
target_new   = ' Rome';
target_true  = ' Paris';
essence      = 'The Eiffel Tower is a';
nSteps       = 20;  lr = 0.5;   % Adam
prefixes     = 'generate';     % 20 texts written by the model, as in the paper
kl_weight    = 0.0625;  decay = 0.5;  clamp = 4;   % the paper's released code: KL(after || before), decay × |δ| / |v0|²
nCovTokens   = 100000;  ridge = 1e-4;

It prints every test and plots Fig. 24. The KL weight and the decay are from the paper's released code (the paper's text gives other values). On a CPU the key covariance takes about 2 hours the first time (then it is loaded from cov_dir), and each edit about 40 minutes; it needs the Deep Learning Toolbox (autodiff for step 2) and the Wikipedia tokens in corpus_file. Fig. 25 is run_rome with L = 0, 4, …, 44 and 17 and five fixed prefixes, prefixes = {'', 'Here is a short note. ', 'I read this in a book. ', 'This is well known. ', 'Last summer we travelled. '} (about 12–15 minutes per edit); its grey curve is causal tracing as in tab 12 (Gaussian noise on the subject, 3σ, 10 samples; MLP outputs restored in 11-layer windows), at the subject's last token only. run_patching makes the same curve: it is the " Tower" column of the MLP map from

model        = 'gpt2-xl';
corruption   = 'gaussian_noise';
clean_prompt = 'The Eiffel Tower is located in the city of';
target       = ' Paris';   competitor = ' Rome';
subject      = ' Eiffel Tower';
noise_scale  = 3;   nSamples = 10;   seed = 1;
layer_window = 5;              % restore layers L-5 ... L+5 together: 11-layer windows
direction    = 'denoise';
doMLP = true;   doAttn = false;   doResid = false;   doHeads = false;

read with the log-rank score. It sweeps every layer and position (about 5,300 runs in GPT-2 XL, several hours on a CPU), while the figure's script runs only the 13 layers at " Tower" (130 runs); its noise draws can differ slightly, so single points may move a little.

An SAE has two halves (Fig. 26). The encoder maps a token's stream x (768 numbers) to a vector f of 24,576 feature activations, which say how strongly each feature is present; a ReLU sets the negative values to exactly 0, so most of them are 0. The decoder rebuilds the stream from them: each feature i has a decoder direction d_i (column i of W_dec, 768 numbers), learned in training and fixed afterwards, and the reconstruction is x̂ = (Σ_i f_i d_i + b_dec) / c = (W_dec f + b_dec) / c: the active features' directions, each scaled by how strongly it fires, plus a bias shared by all tokens. c is one fixed number that scales every stream to the same average length before it enters the SAE (card 2); dividing by it returns x̂ to the stream's own units. f_i says whether and how much a feature is present, and is computed anew for every token; d_i says how it is written into the stream, and is the same for every token. The SAE is trained only to rebuild x from few active features; nobody tells it what the features should mean. Sizes: x and x̂ are 768 × 1 (the stream's width); f is 24,576 × 1 (one number per feature); each decoder direction d_i is 768 × 1, a direction in the stream. Vectors are columns here, as in run_sae (Lesson 7 writes the same with rows and calls the activations z). f follows the SAE papers here: it is not the f of tabs 8, 9 and 14, where f is the difference-in-means direction, a 768 × 1 direction in the stream like d_i.

One training example is one token's stream leaving layer 5: a single vector of 768 numbers (the x of card 1), not a sequence. Call it a token vector. It already carries what GPT-2 computed from the tokens before it, so " of" after "the city" gives a different vector from " of" after "a cup". Layer 5 is the sixth of twelve, the middle of the model (Templeton et al. also used a middle layer). The France prompt gives 8 token vectors, one per token after the start token <|endoftext|>.

How many: 820 million token vectors, one from every token of 820 million tokens of OpenWebText, web text of the kind GPT-2 was trained on. Each training step uses 4,096 token vectors; 200,000 steps use each vector once.

How they are made: GPT-2 reads text 1,024 tokens at a time (its context length), so the text is cut into chunks of 1,024 tokens, and each chunk yields 1,023 token vectors (the first is left out, below). These go into a pool of about 2 million token vectors from about 2,000 chunks, shuffled. Each step takes its 4,096 token vectors from the pool, about as many as 4 chunks hold but drawn from many chunks, so that a step is not dominated by a few stretches of consecutive text; when half the pool is used, GPT-2 refills it with 1,024 new chunks. The two numbers are independent: 1,024 is fixed by GPT-2, 4,096 is the recipe's batch size. Nothing is stored on disk (820 million vectors of 768 numbers would fill about 1.3 TB).

Three details:

• The first token of each chunk is left out. GPT-2 parks attention there (the "attention sink"), and its stream is many times longer than any other token's; a few such token vectors would dominate the loss.
• One scale for all. Every token vector is multiplied by the same number, c = 0.354, which makes the average length √768 ≈ 27.7, as the recipe prescribes, so that the same λ means the same in any model or layer.
• The same token vectors, in the same order, for all three SAEs (λ = 1, 5 and 20), so they differ only in the penalty.

An SAE describes only this one point in the model. Covering every layer takes one SAE per layer, as in Gemma Scope (Lieberum et al. 2024, SAEs on all layers and sublayers of Gemma 2 2B and 9B), or a crosscoder, one sparse dictionary that reads and rebuilds several layers at once (Lindsey et al. 2024).

For one token, with the scaled stream c x:

What is learned. W_enc, b_enc, W_dec and b_dec (table below) are the SAE's parameters: none is set by hand, all are found by gradient descent on this loss (card 4). The decoder directions are simply the columns of W_dec: d_i is column i, so W_dec = [d₀ d₁ … d_24,575], and learning W_dec means learning all 24,576 directions: at every step, each d_i is nudged toward whatever its feature's tokens still need for their reconstruction. The activations f, by contrast, are not parameters: the encoder computes them from each token vector.

The loss has two parts that pull against each other. The reconstruction error wants x̂ as close to x as possible, which is easiest with many active features. The sparsity penalty charges for every active feature, in proportion to its activation f_i times the length of its decoder direction ‖d_i‖. λ sets the exchange rate: a large λ buys sparsity with accuracy. The penalty is an L1 norm, a stand-in for what we really want, the number of active features (L0), which has no useful gradient.

Why ‖d_i‖ is in the penalty. What a feature adds to the reconstruction is f_i d_i, its activation times its direction. The same contribution can be made in many ways: activation 10 with a direction of length 1, activation 1 with length 10, or activation 0.1 with length 100. A penalty on the activations alone, Σ_i f_i (a plain L1 penalty), would charge 10, 1 or 0.1 for these, so the optimizer could make the penalty almost vanish by making every activation tiny and every direction long, without switching a single feature off. Charging f_i ‖d_i‖ instead gives 10 in all three cases: it measures the size of the contribution itself, which rescaling cannot change, so the only way to lower it is to use fewer features, or weaker ones. (Bricken et al. closed the same loophole differently, by fixing every direction's length at 1.)

Why 24,576 features. It is 32 × 768, small enough to train in under 3 hours on one GPU. It is not the number of concepts in the model; there is no fixed number. In a larger SAE, features split into finer ones (Bricken et al.; Templeton et al. went to 1, 4 and 34 million features on Claude 3 Sonnet), so a feature here, such as "France" (card 7), may become several, for France in different contexts. No feature here is dead at λ = 5 (card 5), a sign that a larger dictionary would still find use for more.

Training follows Anthropic's recipe (Conerly et al. 2024) exactly. In order:

1. Starting values. Each decoder direction d_i points in a random direction and has length 0.1. The encoder starts as the decoder's transpose, W_enc = W_dec^⊤, so each feature starts out responding to the same direction it writes. Both biases start at 0.
2. How the parameters are updated (once per step, as summarized at the end of card 3). They are updated with Adam, the standard variant of gradient descent that adapts the step size of each parameter (learning rate 5 × 10⁻⁵, no weight decay). If the gradient is unusually large it is first scaled down to length 1 (clipping), so that one odd batch cannot throw the parameters off. Averaging the losses of 4,096 tokens from many texts makes each update follow what improves the SAE for typical tokens, not for one.
3. Two settings change during training (everything else stays constant). λ rises from 0 to its final value over the first 10,000 steps (5%), so the SAE first learns to reconstruct and only then to be sparse. The learning rate falls linearly to 0 over the last 40,000 steps (20%), so the parameters settle instead of jittering around their final values.
4. Once, after training: every decoder direction is rescaled to length 1. Training leaves each d_i with its own length, say 2.5. Dividing d_i by 2.5 gives it length 1. To keep what the feature writes, f_i d_i, unchanged, f_i must then be 2.5 times larger for every token. Since f_i is computed by the encoder, f_i = ReLU((row i of W_enc) · (c x) + b_enc,i), that means multiplying the feature's row of W_enc and its bias b_enc,i by 2.5: the number inside the ReLU becomes 2.5 times larger, and a ReLU of a number 2.5 times larger is 2.5 times larger (a negative number stays at 0). Every reconstruction stays exactly the same; only the bookkeeping moves. Now ‖d_i‖ = 1 for every feature, so an activation alone says how much a feature writes: feature 12 at 13.8 writes 13.8 units along its direction, and can be compared directly with feature 13788 at 3.1.

Each SAE took about 2 hours 50 minutes on one RTX 3080 Ti, including GPT-2's forward passes.

At the start every feature fires on every token (L0 ≈ 22,500 of 24,576, Fig. 27b): with W_enc = W_dec^⊤, each feature responds to whatever it points towards, and the reconstruction is easy (variance explained 0.95–0.99). As λ rises, features switch off. By step 10,000 (λ at its final value) L0 is down to 444 (λ = 1) and 114 (λ = 5), and the reconstruction has paid for it (Fig. 27c). The remaining 190,000 steps are slow refinement: L0 keeps falling (to 147 and 19) while the reconstruction holds or improves slightly, so the SAE finds sparser ways to say the same thing. For λ = 5, the sparsity term ends a little above the reconstruction error (125 against 102 per token, Fig. 27a): the two are balanced, as the recipe intends. The learning-rate decay at the end changes almost nothing here.

Five measures. The first three say how well the SAE does its job (how sparse, how accurate, and how much of GPT-2's behaviour it keeps), on 2 million tokens of validation text it never trained on; the last two say whether its 24,576 features are used sensibly, on 10 million evaluation tokens: more training text that the SAE never saw.

1 · L0: the average number of active features per token (f_i > 0).

2 · Variance explained by x̂: 1 − Σ ‖x − x̂‖² / Σ ‖x − mean(x)‖², summed over the same tokens. It is the reconstruction error as a fraction of how much the token vectors vary around their mean, mean(x), the average token vector of the sample (1 = perfect; 0 = no better than giving every token the mean vector; the scale c cancels). Equivalently, variance explained = 1 − normalized mean squared error (the reconstruction's squared error divided by that of the mean vector); Conerly et al. report normalized MSE (0.16 for a residual-stream SAE of a 4-layer model at λ = 5, about 0.84 variance explained if normalized the same way) without spelling out what they normalize by.

3 · Loss recovered, the measure that matters most: GPT-2's next-token loss with x̂ spliced in place of the real stream after layer 5, compared with the clean model (3.14 nats per token) and with the stream there set to zero (11.73). The fraction of that gap the SAE closes is the loss recovered: loss recovered = (zero − SAE) / (zero − clean); for λ = 5, (11.73 − 3.80) / (11.73 − 3.14) = 0.92. 1 means the SAE loses nothing; 0 means it is as bad as no stream at all. (This is GPT-2's next-token loss, not the SAE's training loss of card 3.)

4 · Dead features: the percentage of features that never fire in the 10 million tokens, wasted capacity.

5 · Dense features: the percentage that fire on more than 1% of tokens, too often to mean one specific thing.

No setting is best at everything. λ = 1 rebuilds the stream well (99% of the loss recovered) but uses about 150 features per token, and 17.6% of its features fire on more than 1% of all tokens: too common to stand for anything specific. λ = 20 is very sparse (2.8 features per token) but loses 3.28 nats per token. λ = 5, the recipe's default, sits between: 19 features per token, 80% of the variance, 92% of the loss recovered, and almost no dense features. It meets the bar Templeton et al. set for their SAEs (L0 below 300, at least 65% variance explained), and it is the SAE used for the rest of the tab. Almost no feature is dead: none at λ = 1 and 5, 0.03% (7 features) at λ = 20 (Templeton et al. had 2% dead with 1 million features, 65% with 34 million): 24,576 features is small for this recipe. One limitation: the SAE explains 80% of the variance, so a fifth of what the stream holds is in the error term, which this tab does not interpret.

The prompt of the earlier tabs, "The capital of France is the city of", through the λ = 5 SAE at layer 5.

Which features fire (Fig. 29a). Of 24,576 features, 8 to 28 are active at each token. Each row of Fig. 29a is one feature, labelled with its ID and, in brackets, a hint at its meaning: the token where that feature fired most strongly in the 10 million evaluation tokens of card 5. That word comes from the training text, not from this prompt (the prompt's tokens are the columns): during the evaluation, the training run kept each feature's 20 largest activations and the positions of those tokens in the text, and the figure reads the token at the top position. The strongest ones mostly name the token itself: feature 23685 at " capital", feature 12 at " France", feature 19993 at " city". At the last " of", the strongest feature is 20664. Its label, " of", says little on its own; its 20 top contexts do (card 7): it fires on " of" after "town" or "city" when a place name follows, as in "the northern Syrian town of Manbij". That is exactly the situation at the end of this prompt, where the next word should be a city.

Does the reconstruction keep the answer? (Fig. 29b, c, the splice test). Replacing the stream after layer 5 by x̂, at every position, leaves the answer intact (" Paris" still first; logit difference over " London" +5.16, clean +5.08), so the 8–28 features per token carry what the later layers need. Setting the stream there to zero instead destroys it (" Paris" falls to rank 4,870). On these 8 tokens x̂ explains 64% of the variance, less than its 80% average on validation text.

Nothing in training named the features, so their meaning has to be read off afterwards. Two standard readings:

1. Where it fires most: the tokens with the highest activation in the 10 million evaluation tokens (card 5), with their context (what the feature responds to).
2. What it writes: its decoder direction passed straight through the unembedding, W_U d_i, as in the logit lens (tab 3). The tokens with the largest values are the ones the feature pushes the next-token prediction towards; this ignores everything the later layers do with it (what the feature says).

How "what it writes" is computed (Figs. 30b, 31b): one matrix–vector product, then sorting. Take the feature's decoder direction d_i, column i of W_dec (768 numbers, length 1), and multiply it by GPT-2's unembedding matrix W_U (50,257 × 768, the matrix that turns the final stream into one score per vocabulary token; its final-LayerNorm scale is already folded in, as in the rest of the course). That gives 50,257 numbers, one per token. The 8 largest are drawn in green (promoted), the 8 most negative in red (suppressed). For feature 12 in MATLAB, which counts from 1: eff = W.W_U * S.W_dec(:, 13);

Feature 12 fires on 0.014% of tokens (about 1 in 7,000), always on "France". Its four strongest activations in the 10 million evaluation tokens, each with the 10 tokens before it and 3 after, as they appear in the text (a cut-off word is where the 3 tokens end):

• activation 14.67:  …/1SkBmoR<|endoftext|>The highly rated France Under-16…
• activation 14.64:  … Chamberlain (UK), Edouard Daladier (France), Adolf Hitler…
• activation 14.13:  … to US; start of Rothschild Inc 1981 France nationalises the…
• activation 14.09:  … Canada), European Union member states (including Sweden, France and Germany),…

It is specific to France, not to countries in general: in "The capital of X is the city of" it fires at 13.8 for France and at exactly 0 for each of the other 26 EU countries (Fig. 30a). It writes French things: " Alps", " Hollande", " fries", " France", " franc", " countryside" (Fig. 30b). These are candidates for the next word, so a France feature promotes words that go with France, with " France" itself close behind (6th); the top eight differ little (0.70–0.89), so their order matters less than the fact that all are about France. A concept feature: the token is France, and the feature carries what goes with it.

Click the figure to enlarge.

Feature 20664 fires on 0.023% of tokens, on " of" after "town" or "city" when a place name follows. Its four strongest activations:

• activation 18.07:  … US troops had been sent to the northern Syrian town of Manbij…
• activation 17.58:  … identified, had been living in the southern Dutch city of Maastrich…
• activation 17.27:  …government gunmen killed four protesters Sunday in the port city of Banias…
• activation 17.06:  … French military operation to secure the northern Malian town of Gao is…

Fig. 31a tests that reading on 12 short phrases. At the final " of" it fires at 7.7–15.7 whenever a place name should follow ("the northern town of", "the port city of", "the small town of", "the city of"), only 0.8 at "The capital of", and 0 for "one of", "a cup of", "the University of", "the King of" and three other ordinary uses. It writes the beginnings of place names: " Normandy", " Dover", " Trinidad", " Alb", " Nev", " southwestern" (Fig. 31b). A prediction feature: it fires on an ordinary word, " of", and stands for what comes next, a place. In the France prompt it is the main feature at the last token (7.7), the position where the answer is predicted.

Click the figure to enlarge.

Both readings are partial: top contexts show only the strongest firings, and a feature can mean something broader at weaker activations; the direct effect ignores the later layers. The next step tests what a feature does by changing it.

The SAE splits each token's stream into pieces: the features' contributions, the bias, and what the SAE misses (the error, x − x̂):

To set feature i to a value v, change only its term and keep the rest, the error included:

The bias and the error do not appear because they are already inside x and stay as they are. Written out, the edit replaces a single term:

The short form is what the code uses: it needs only feature i's activation and direction. The edit is applied to the stream after layer 5 at every position except the first. v = 0 turns the feature off (its contribution f_i d_i / c is subtracted; where it did not fire, nothing changes); a large v clamps it on (Templeton et al.'s feature steering, the SAE version of tab 14's steering with a direction).

In short: clamping sets feature i's contribution to the layer-5 stream to v d_i / c at every position (except the first), and leaves the other features' contributions to that stream as they were; the later layers, and every token generated after, then respond to the changed stream.

Off, in the France prompt:

Removing the single France feature, one of 24 active at that token, takes away most of the probability of " Paris" (0.097 → 0.014; it falls to rank 4). Removing "capital" leaves France but not capital: the top answer becomes another French city, Marseille. The features carry the facts the later layers use, in the form the readings above suggested.

Clamped on, in generated text (greedy, 20 tokens; v in multiples of the feature's largest activation in training text):

Clamped gently, each feature bends the text toward its meaning while keeping it fluent; clamped harder, it takes over. In our tests (six features, two prompts; the table shows a selection), 0.5 × the maximum usually bent the text while keeping it fluent, 1 × broke some features into repetition ("capital", 20664) but not others ("terrible"), and 2 × or more broke the one feature tried (20664: 2, 5 and 10 × all give "K K K…"). Templeton et al. clamped features in Claude 3 Sonnet to 5 × their maximum and more; a small model edited at every position tolerates much less.

Negative clamps are allowed too (Templeton et al. use them): they write the feature's direction backwards, into a stream the encoder never produces, since its activations cannot be negative. Here they mostly weakened the concept rather than producing its opposite: at −0.25 × "wonderful" the walk loses its "very nice and kind" people and turns neutral, and at −1 × "France" the food stays chicken, now a curry. Even −1 × the maximum kept the text fluent, for the three features tried (two are shown in the table).

Tab 8 found one direction that separates positive from negative sentences: the difference of the mean streams at layer 5, on 20 sentences (the direction tab 14 steered with). How do the SAE's features relate to it? Most decoder directions are nearly perpendicular to it (Fig. 32a); only 33 of 24,576 have a cosine above 0.3 in size. Two are clear-cut: feature 265 (cosine +0.31), whose top contexts are "wonderful", and feature 20419 (−0.38), whose top contexts are "terrible" and "awful".

On the 20 sentences themselves (Fig. 32b), feature 265 fires in 3 of the 10 positive sentences and 0 of the negative ones; feature 20419 in 4 of the 10 negative sentences and 0 of the positive ones. It is a word family, not one word: 265 fires in the sentences with "wonderful", "fantastic" and "brilliant"; 20419 in those with "terrible", "awful", "horrible" and "miserable". Each fires only where its kind of word occurs, and is silent in the other sentences. The sentiment direction is an average over all sentences; the SAE instead has many specific features that each carry part of it. Both views are useful: the direction is one handle for the whole concept; the features say which words and contexts make it up.

A factual concept: France → Italy. Tab 14 also steered with a direction from France to Italy, built the same way from 10 matched pairs ("I spent the summer in Italy." / "… in France."). Here the picture is much simpler (Fig. 32c, d). Only 6 features have a cosine above 0.3 in size, and the two largest are the two countries: feature 17482, whose top context is "Italy" (cosine +0.56), and feature 12, "France" (−0.40), the feature of card 7. On the 20 sentences, 17482 fires in all 10 Italy sentences and none of the France ones, and feature 12 in all 10 France sentences and none of the Italy ones. So the France → Italy direction is close to "the Italy feature minus the France feature": a concept carried by one word in every sentence maps onto single features, while sentiment, carried by many different words, is spread over many.

There are two stages: training the SAE in Python, then analysing it in MATLAB.

1. Train (Python, GPU, hours). train_sae.py trains one SAE in PyTorch, which takes about 2 hours 50 minutes on one RTX 3080 Ti (12 GB). It reads three inputs:

• OpenWebText tokenized with GPT-2's tokenizer, in train.bin and val.bin.
• GPT-2 small's weights, in gpt2_full_weights.mat. The script runs GPT-2 itself to get the stream after layer 5.
• A JSON file per SAE, listing only the settings that differ from the defaults.

For example:

python "Python code/train_sae.py" "Python code/sae_l5_lam5.json"

% sae_l5_lam5.json:
{"out_dir": "runs/sae_l5_lam5", "lam": 5.0}

% the defaults:
layer      = 5          % the stream leaving this layer (0-based)
d_sae      = 24576      % features
lam        = 5.0        % sparsity penalty λ (1, 5, 20 here)
lr         = 5e-5       % Adam, no weight decay; linear decay to 0 over the last 20% (lr_decay = 0.2)
steps      = 200000;  batch = 4096       % 820 million tokens, each used once
lam_warmup = 0.05       % λ rises from 0 over the first 5% of steps
init_norm  = 0.1;  grad_clip = 1.0;  buffer_windows = 2048;  seed = 1

It saves two files in out_dir:

• sae.mat, the trained SAE: W_enc, b_enc, W_dec, b_dec, c and the layer. The decoder directions are already rescaled to length 1.
• eval.mat, the evaluation on the 10 million evaluation tokens: each feature's firing count, its mean activation, and its 20 top activations with their positions in train.bin. The scores of card 5 are also saved in eval.json.

2. Analyse (MATLAB, CPU, minutes). run_sae trains nothing. It loads sae.mat and eval.mat from run_dir, loads GPT-2 small's weights to run the prompts, and reads train.bin only to print the text around each top activation. Its settings are:

run_dir      = fullfile(project_root(), 'runs', 'sae_l5_lam5');   % train_sae.py's output folder
prompt       = 'The capital of France is the city of';
target       = ' Paris';   rival = ' London';
feature      = [];          % [] = the most active feature at the last token (20664 here); 12 for France
steer_prompt = 'I went for a walk and saw';
clamps       = [0 0.25 0.5 1];   % multiples of the feature's largest activation

It prints each step of this tab for one feature (its top contexts, the tokens it promotes, the splice test, the comparison with the sentiment direction, and the steered continuations) and plots the feature on the prompt, the feature density and its promoted tokens. It runs in a few minutes on a CPU.