Dom 🆚 Machine

Two Routes Through a Transformer

2026-10-04T00:00:00-04:00

A Mathematical Framework for Transformer Circuits has a walkthrough video to go with it. I could not follow either until I had worked out what a one-layer attention-only model actually computes, so this post is that groundwork: the path decomposition end to end, with every number in the example derived rather than asserted. Architecture background is in Understanding Transformers.

Path decomposition of transformer output

Two routes

One-layer attention-only transformer. Tokens in, logits out at each position. One sequence runs through the whole note:

pos	1	2	3	4	5	6	7	8	9
token	`the`	`cat`	`sat`	`on`	`the`	`mat`	`.`	`the`	`cat`

Two routes run from a token to a logit. Both are below, over the first three positions. The next two sections take one each:

          pos 1       pos 2       pos 3
token:    "the"       "cat"       "sat"
            |           |           |
          embed       embed       embed
            |           |           |
            |           |           '---------------.
            |           |                           |
            |           |      route 1 (direct):    |
            |           |      rides the residual   |
            |           |      stream, untouched    |
            |           |                           |
            '-----------+-----> head h -------------+
                                                    |
              route 2 (attention):                  |
              the head reads back over              |
              earlier positions, taking             |
              some % from each, and adds            |
              what it grabbed                       |
                                                    v
                                                 unembed
                                                    |
                                                    v
                                            logits at pos 3

Position 3 can read positions 1, 2 and 3, never forward. That restriction is the causal mask.

Three simplifications throughout: no MLPs, LayerNorm folded into adjacent weights, and no positional embeddings.

Route 1: the token's own embedding

flowchart TB t1["pos 1 the"] --> e1["embed"] t2["pos 2 cat"] --> e2["embed"] t3["pos 3 sat"] --> e3["embed"] e3 ==>|"route 1 (direct) rides the residual stream, untouched"| u["unembed"] u --> out["logits at pos 3"] classDef dim stroke-dasharray: 4 4 class t1,t2,e1,e2 dim

Dashed = present in the sequence, unused by this route.

"sat" is embedded, rides the residual stream, hits the unembed untouched. Alone that makes a bigram table: a lookup keyed on one previous token.

	`the`	`cat`	`sat`	`on`	`mat`	`.`
`the`	0.0	3.7	0.2	0.1	4.1	0.0
`cat`	0.3	0.0	3.9	0.4	0.1	1.2
`sat`	1.1	0.1	0.0	4.2	0.2	0.9
`on`	4.4	0.2	0.0	0.0	0.6	0.1
`mat`	0.5	0.1	0.1	0.2	0.0	3.8
`.`	3.5	0.2	0.0	0.1	0.1	0.0

Route 1: the bigram table. Row = the token you just saw, column = a candidate next token, cell = score. Six tokens here; the real table is n_vocab x n_vocab.

Every table in this post is hand-built to make the mechanism visible. None of these numbers come off a trained model. What is derived is everything downstream of them: the attention percentages, the logit shifts and the totals all follow from these tables by arithmetic you can check.

Walk the first three positions through it:

"the" -> read row the -> mat 4.1 beats cat 3.7. Predicts mat. Our sequence has cat, so this one is wrong.
"the cat" -> read row cat -> sat scores 3.9. Predicts sat. Correct.
"the cat sat" -> read row sat -> on scores 4.2. Predicts on. Correct.

The last step read row sat and nothing else. There is no row for "the cat sat", so "the cat" is gone. By position 3 the model remembers only the token under it.

Reaching back for what it forgot is route 2's job.

Route 2: what the head reaches back for

flowchart TB t1["pos 1 the"] --> e1["embed"] t2["pos 2 cat"] --> e2["embed"] t3["pos 3 sat"] --> e3["embed"] e1 --> h["head h takes some % from each"] e2 --> h e3 --> h h ==>|"route 2 (attention) adds what it grabbed"| u["unembed"] u --> out["logits at pos 3"]

The head reads back over earlier positions, takes a percentage of each, and adds what it grabbed to the residual stream, which carries it to the unembed.

Stand at position 8, on the third the. Route 1 reads row the and answers mat over cat, the same mistake as position 1. It cannot know this sentence is about a cat.

Route 2 takes two steps: pick where to look, then use what is there.

Step 1: where to look. Every position scores every position it may see. A score is how well a source matches what the destination wants. With no positional embeddings a position carries nothing but its token, so a score depends only on the two tokens involved, the one you are standing on and the one you are looking at. x marks a cell the causal mask blocks:

to from	1 `the`	2 `cat`	3 `sat`	4 `on`	5 `the`	6 `mat`	7 `.`	8 `the`	9 `cat`
1 `the`	0.4	x	x	x	x	x	x	x	x
2 `cat`	0.5	0.3	x	x	x	x	x	x	x
3 `sat`	0.6	1.2	0.3	x	x	x	x	x	x
4 `on`	0.7	1.0	0.5	0.3	x	x	x	x	x
5 `the`	0.4	3.4	0.7	0.6	0.4	x	x	x	x
6 `mat`	0.5	0.9	0.4	0.6	0.5	0.3	x	x	x
7 `.`	0.4	0.8	0.5	0.4	0.4	0.9	0.3	x	x
8 `the`	0.4	3.4	0.7	0.6	0.4	2.3	0.5	0.4	x
9 `cat`	0.5	0.3	0.6	0.4	0.5	0.8	0.3	0.5	0.3

Route 2, step 1: the score matrix. Row = where you are standing, column = where you are looking, cell = the score.

Lower triangular: the causal mask drawn out. Rows 5 and 8 are both the and both score cat highest, because a the wants a noun. Row 1 is a the with nothing behind it.

Row 8 is ours. Its x in column 9 hides cat, the token position 8 must predict. Unmasked it reads 3.4, takes most of the attention, and the model copies the answer instead of inferring it.

Now take row 8 out of the matrix:

  0.4   3.4   0.7   0.6   0.4   2.3   0.5   0.4    x

These go into softmax. Real attention divides by \(\sqrt{d_\text{head}}\) first, which changes the numbers but not the shape. Exponentiate each, divide by the total:

pos	token	score	exp	attention
1	`the`	0.4	1.5	1.5 / 50 = 3%
2	`cat`	3.4	30.0	30 / 50 = 60%
3	`sat`	0.7	2.0	2.0 / 50 = 4%
4	`on`	0.6	1.8	1.8 / 50 = 4%
5	`the`	0.4	1.5	1.5 / 50 = 3%
6	`mat`	2.3	10.0	10 / 50 = 20%
7	`.`	0.5	1.7	1.7 / 50 = 3%
8	`the`	0.4	1.5	1.5 / 50 = 3%
9	`cat`	x	0.0	0 / 50 = 0%
			50.0	100%

Route 2, step 1: row 8 softmaxed. Score to exp to share of the total.

Exponentiating decides it:

cat at position 2 -> scores 3.4 -> exp 30.0 -> 30.0 / 50.0 = 60%
mat at position 6 -> scores 2.3 -> exp 10.0 -> 10.0 / 50.0 = 20%
position 9 -> masked -> exp 0.0 -> 0%

A gap of \(3.4 - 2.3 = 1.1\) in score becomes a factor of \(30.0 / 10.0 = 3\) in attention. That is what exponentiating does, and it is why cat ends up holding most of the head.

This row is row 8 of the head's attention pattern, written \(A\) here and \(A^h\) once there is more than one head.

Step 2: what it does. Route 2's own lookup table, same shape as the bigram one, keyed on the position attended to rather than the token under you:

	`the`	`cat`	`sat`	`on`	`mat`	`.`
1 `the`	0.9	0.1	0.0	0.1	0.1	0.0
2 `cat`	0.1	5.0	0.2	0.0	0.0	0.0
3 `sat`	0.2	0.1	4.1	0.1	0.1	0.0
4 `on`	0.1	0.0	0.1	3.4	0.2	0.0
5 `the`	0.9	0.1	0.0	0.1	0.1	0.0
6 `mat`	0.1	0.2	0.1	0.0	2.0	0.1
7 `.`	0.0	0.1	0.0	0.0	0.0	1.8
8 `the`	0.9	0.1	0.0	0.1	0.1	0.0
9 `cat`	0.1	5.0	0.2	0.0	0.0	0.0

Route 2, step 2: the OV table. Row = a position, column = a candidate next token, cell = how much attending there shifts the logit.

A row depends only on the token at that position, so 1, 5 and 8 are identical, as are 2 and 9. Position 9 has a row like any other; step 1 gave it 0%. The big diagonal is the tell: attend to cat, boost cat. A copying head.

Multiply step 1 by step 2:

position 2 -> 60% attention, row 2 cat, cat scores 5.0 -> 0.60 x 5.0 = +3.0 to cat
position 6 -> 20% attention, row 6 mat, mat scores 2.0 -> 0.20 x 2.0 = +0.4 to mat
everything else -> 3 to 4% each, on rows that boost the, sat, on and . -> under 0.1 to either candidate, so the ranking is untouched

Add the routes:

            route 1   route 2    total
  "mat"       4.1       +0.4       4.5
  "cat"       3.7       +3.0       6.7   <- wins

Route 1 answered from the token under it and got mat again. Route 2 reached six positions back, found cat, and flipped the prediction. That shape is a skip-trigram: X is in the context, so predict X. A bigram table cannot do it, and it is what one-layer heads mostly learn. Copying names is the classic case: Mr Dursley ... Mr ___ -> Dursley.

Two routes, added:

logits at a position  =  route 1  +  one term per attention head

The trick

Those percentages are the only nonlinear part of the model. Everything else is matrix multiplication.

Pretend they are fixed constants. The path from input token to output logit becomes a chain of matrix multiplies, and a chain of matrices collapses into one.

That matrix is what you want: vocab in, vocab out, read straight off the weights without running the model. Route 1 gives one, each head another. The formula is the list.

The formula

\[T = \text{Id} \otimes W_U W_E + \sum_{h \in H} A^h \otimes (W_U W_{OV}^h W_E)\]

Every symbol in the formula maps to something from the worked example:

Symbol	Shape	What it is
\(T\)		the whole model: tokens in, logits out
\(W_E\)	`[d_model, n_vocab]`	embed
\(W_U\)	`[n_vocab, d_model]`	unembed
\(\text{Id} \otimes W_U W_E\)		route 1
\(W_U W_E\)	`[n_vocab, n_vocab]`	route 1's bigram table
\(\text{Id}\)	`[n_ctx, n_ctx]`	"read this position only": route 1 never moves positions
\(\sum_{h \in H} \dots\)		one term per head; they just add
\(W_{QK}^h = W_Q^{hT} W_K^h\)	`[d_model, d_model]`	the head's query-key map
\(W_E^T W_{QK}^h W_E\)	`[n_vocab, n_vocab]`	route 2 step 1: the score table
\(A^h\)	`[n_ctx, n_ctx]`	those scores softmaxed: the percentages, carrying the mask
\(W_{OV}^h = W_O^h W_V^h\)	`[d_model, d_model]`	the head's value-output map
\(W_U W_{OV}^h W_E\)	`[n_vocab, n_vocab]`	route 2 step 2: the lookup table
\(\otimes\)		glues the which position part onto the what it does part

Each term splits the same way: a left factor saying which positions, and a right factor saying what happens to the logits.

Note what is missing: \(W_Q\) and \(W_K\). The formula takes \(A^h\) as given, so the weights behind it are already absorbed. That is the freeze, visible in the notation.

The tensor-product notation

\(A \otimes B\) is a factored linear map. With input \(t\), a [n_ctx, n_vocab] matrix of one-hot rows, \((A \otimes B) \cdot t = A t B^T\). The factors act on different axes:

left (\(\text{Id}\), \(A^h\)) acts across positions: which source reaches which destination
right (\(W_U W_E\), \(W_U W_{OV}^h W_E\)) acts across the vocabulary: what that does to the logits

Neither knows about the other. That separation is the point.

The attention pattern \(A\)

Softmax every row of the step 1 matrix, not just row 8, and you get \(A\). Same grid, same x, scores as percentages.

to from	1 `the`	2 `cat`	3 `sat`	4 `on`	5 `the`	6 `mat`	7 `.`	8 `the`	9 `cat`
1 `the`	100%	x	x	x	x	x	x	x	x
2 `cat`	55%	45%	x	x	x	x	x	x	x
3 `sat`	28%	51%	21%	x	x	x	x	x	x
4 `on`	26%	36%	21%	17%	x	x	x	x	x
5 `the`	4%	82%	5%	5%	4%	x	x	x	x
6 `mat`	16%	24%	14%	17%	16%	13%	x	x	x
7 `.`	12%	18%	14%	12%	12%	21%	11%	x	x
8 `the`	3%	60%	4%	4%	3%	20%	3%	3%	x
9 `cat`	11%	9%	13%	10%	11%	17%	9%	11%	9%

In \(A[i,j]\), \(i\) is where you're standing and \(j\) is where you're looking. \(A[8,2] = 60\%\) is position 8 on the looking back at position 2 on cat. Swap them and \(A[2,8]\) is x: position 2 cannot see the future.
Rows sum to 100%, columns do not. Each destination splits a fixed budget; a source can be read by many or none. Row 1 is 100% on itself, having nowhere else to look.
Read down a column for the other question. Who attends to cat at position 2? Rows 5 and 8, at 82% and 60%, both of them a the.
Row 5 gets 82% where row 8 gets 60%, same token asking the same question. Position 5 cannot see mat yet, so cat has no competition; by position 8 it does.
The shape of the rows names the head. A diagonal band is a previous-token head. A row that puts almost all of its percentage on position 1 means the head is parked there, reading the first token and doing nothing useful. That is how a head switches itself off.

One \(A\) per head, hence \(A^h\).

The two terms, formally

\[\text{Id} \otimes W_U W_E\]

Route 1. A zero-layer transformer has only this term, and training drives \(W_U W_E\) to approximate bigram log-likelihoods.

\[\sum_{h \in H} A^h \otimes (W_U W_{OV}^h W_E)\]

Route 2. Heads don't interact in a one-layer model, which is why the terms simply add. Each factors into the two steps:

\(W_E^T W_{QK}^h W_E\), the QK circuit. Step 1, where to look. Entry \((a, b)\) is how much a destination holding token \(a\) wants a source holding token \(b\). Step 1's matrix is this table read off at the sequence's tokens, which is why its scores depended only on the two tokens involved. Softmax them over the context and you get \(A^h\), which is the circuit's output, not the circuit itself.
\(W_U W_{OV}^h W_E\), the OV circuit. Step 2, what it does. Independent of position and of the attention pattern.

Both are [n_vocab, n_vocab], and together they are the head.

One thing the worked example skipped: the head carries a value vector, not the token.

Each source \(j\) computes \(v_j = W_V \cdot (\text{residual stream at } j)\). The head outputs the attention-weighted sum, and \(W_O\) maps it back into the residual stream. \(W_V\) is a projection, so only the subspace that head cares about travels. Every attended position lands at once: a blend, not a choice.

\(W_{OV}^h = W_O^h W_V^h\) is that round trip. It is [d_model, d_model] but has rank at most d_head, factoring through the head's narrow dimension. Sandwich it between \(W_E\) and \(W_U\) for the vocab-to-vocab lookup table.

Why this framing is useful

All of this holds only with \(A^h\) frozen. Given that, the decomposition is exact: behaviour reads straight off the weights as a sum of paths.

The payoff: a head is interpretable from its QK and OV matrices alone, both [n_vocab, n_vocab], never by probing activations. Read them directly, or eigendecompose the OV matrix, where positive eigenvalues mark a head that copies whatever it attends to. The paper's skip-trigram example is [perfect] ... [are] -> [perfect].

It also sets up the two-layer story. Everything above is a single attention block. Stack a second one on top and its input is the first block's output rather than the raw embeddings, so its \(A^h\) is no longer a function of the tokens alone. The cross terms that creates, where the second block's queries, keys or values read what the first block wrote (Q-, K- and V-composition), are where induction heads come from.

Sampling from a Transformer

2026-10-03T00:00:00-04:00

The last two posts built a transformer and trained one. In the Training a Transformer post, I used TransformerLens's sampler on the TinyStories weights I trained. This post ports the sampler, using the techniques from the ARENA curriculum: greedy_search, apply_temperature, apply_frequency_penalty, sample_basic, sample_top_k and sample_top_p, plus beam_search.

PvML: https://github.com/d0mzw/PvML

Disclaimer: these notes come from working through the ARENA 3.0 curriculum. I claim no credit for the original material, and this is not affiliated with or endorsed by ARENA.

Sampler

flowchart TB prompt["prompt tokens (1, posn)"] prompt --> fwd["Transformer logits (1, posn, d_vocab)"] fwd --> last["final position only logits (d_vocab,)"] last -->|"frequency_penalty"| pen["penalised (d_vocab,)"] pen -->|"/ temperature"| scaled["scaled (d_vocab,)"] scaled -->|"top_k or top_p"| filt["filtered (k,)"] filt --> draw["draw one token (1,)"] draw -->|"append, run again"| prompt

The forward pass is the one from training, frozen in place. It still emits a score over all 50,257 vocabulary entries at every position, and generation needs only the last, so the rest are discarded every step.

src/pvml/sampling/
    strategies.py    six functions on one logits vector
    args.py          SamplingArgs, the knobs and their checks
    sampler.py       the loop, the tokenizer, the model
    beams.py         beam search

`strategies`

Six functions on a vector of logits, excluding beam_search. Four pick a token, two reshape the logits first.

Function	Code	What it does
`greedy_search`	2 lines	the argmax, no draw at all
`apply_temperature`	2 lines	divide before the softmax
`apply_frequency_penalty`	6 lines	subtract a cost per prior occurrence
`sample_basic`	2 lines	draw from all 50,257
`sample_top_k`	4 lines	keep the k best, renormalise, draw
`sample_top_p`	9 lines	keep the smallest set summing to p, draw

`SamplingArgs`

python

@dataclass
class SamplingArgs:
    max_new_tokens: int = 50
    temperature: float = 1.0  # 0 means greedy
    top_k: int = 0  # 0 disables
    top_p: float = 0.0  # 0 disables
    frequency_penalty: float = 0.0
    seed: int | None = None

    def __post_init__(self) -> None:
        if self.max_new_tokens <= 0:
            raise ValueError(f"max_new_tokens must be positive, got {self.max_new_tokens}")
        if self.temperature < 0:
            raise ValueError(f"temperature must be non-negative, got {self.temperature}")
        if self.top_k < 0:
            raise ValueError(f"top_k must be non-negative, got {self.top_k}")
        if not 0 <= self.top_p <= 1.0:
            raise ValueError(f"top_p must be a probability in [0, 1], got {self.top_p}")
        if self.top_k and self.top_p:
            raise ValueError(
                f"set at most one of top_k and top_p, got top_k={self.top_k} top_p={self.top_p}"
            )

top_k and top_p are alternatives, not a pair. Setting both raises, rather than silently applying whichever the dispatcher checks first
The checks run once at construction, not per token, so a bad combination fails before the first forward pass
raise rather than assert, because python -O strips asserts and the sampler treats these as a guarantee

`next_token`

python

if args.frequency_penalty != 0.0:
    logits = apply_frequency_penalty(input_ids, logits, args.frequency_penalty)
if args.temperature == 0:
    return greedy_search(logits)
if args.temperature != 1.0:
    logits = apply_temperature(logits, args.temperature)
if args.top_k > 0:
    return sample_top_k(logits, args.top_k)
if args.top_p > 0.0:
    return sample_top_p(logits, args.top_p)
return sample_basic(logits)

The penalty is a correction to the raw scores so it goes first, and top_k and top_p read the distribution temperature produces so they come after it
greedy_search is checked before apply_temperature because temperature 0 would divide by zero. As temperature falls towards 0 the distribution concentrates on the argmax, so greedy_search is the limit it never reaches
Deliberately different from ARENA, which applies temperature first and returns greedy before the penalty. Here the penalty lands first, so it applies under greedy too, and dividing by temperature afterwards scales its strength by 1/temperature

The loop

python

if self.prepend_bos:
    window = t.cat([input_ids[:1], input_ids[1:][-(self.cfg.n_ctx - 1) :]])
else:
    window = input_ids[-self.cfg.n_ctx :]
logits = self.model(window[None])[0, -1]
next_id = self.next_token(input_ids, logits, args)
input_ids = t.cat([input_ids, t.tensor([next_id], device=self.device)], dim=-1)
if next_id == self.tokenizer.eos_token_id:
    break

input_ids        (seq_len,)
window           (min(seq_len, n_ctx),)   BOS plus the last n_ctx - 1
window[None]     (1, posn)                None adds the batch axis
model(...)       (1, posn, d_vocab)
[0, -1]          (d_vocab,)               drop the batch, keep the last position

input_ids is 1 dim, the model wants (batch, posn)
The window is built in two pieces so BOS stays at position 0. Slicing the whole sequence to the last n_ctx would drop it the moment generation outgrows the context length
next_token gets the full input_ids, not the window, so apply_frequency_penalty counts tokens the model can no longer see

The missing BOS token

python

if self.prepend_bos:
    bos = t.tensor([self.tokenizer.bos_token_id], device=self.device)
    input_ids = t.cat([bos, input_ids], dim=-1)

Without these a trained TinyStories model writes this:

without BOS  in  [  7454,    2402,  257,     640]
                 ['Once', ' upon', ' a', ' time']
             out 'Once upon a time,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,'

with BOS     in  [          50256,   7454,    2402,  257,     640]
                 ['<|endoftext|>', 'Once', ' upon', ' a', ' time']
             out 'Once upon a time, there was a little girl named Lily. She loved to
                  play outside in the sunshine. One day, she went to the park'

tokenize_and_concatenate(..., add_bos_token=True) puts 50256 at the start of every training row, so the model never saw anything else at position 0
Prepending it once is not enough. Slicing the whole sequence to the last n_ctx drops BOS again the moment generation outgrows the window, so the window is built as BOS plus the last n_ctx - 1 tokens. That also matches training, where seq_len = max_length - 1 leaves room for it

Beam search

python

@dataclass
class Beams:
    model: t.nn.Module
    tokenizer: object
    logprob_sums: Float[Tensor, "batch"]
    tokens: Int[Tensor, "batch seq"]

Every other strategy collapses the distribution to one id and moves on. Beam search carries k forward, so that batch axis is the beams.

python

new_sums = einops.repeat(self.logprob_sums, "b -> b k", k=k) + topk_logprobs
new_tokens = t.concat(
    [einops.repeat(self.tokens, "b s -> b k s", k=k), topk_ids.unsqueeze(-1)], dim=-1
)
return Beams(self.model, self.tokenizer, new_sums.flatten(), new_tokens.flatten(0, 1))

The shapes through two rounds at k=3:

start        logprob_sums (1,)   tokens (1, 4)
generate(3)  logprob_sums (3,)   tokens (3, 5)
filter(3)    keep (3, 5)         terminated (0, 5)   none ended yet
generate(3)  logprob_sums (9,)   tokens (9, 6)

flatten makes the search global: filter ranks a 1 dim tensor, so all k x k candidates compete together

num_beams	best score	completion
1	-18.076	`When I was a kid, I was a little bit of a nerd.`
3	-14.830	`When I was a kid, I used to go to the movies with my`

Module output

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python -m pvml.sampling.strategies
logits  [3.0, 2.0, 1.0, 0.0, -1.0]
probs   ['0.636', '0.234', '0.086', '0.032', '0.012']
cumulative ['0.636', '0.871', '0.957', '0.988', '1.000']

greedy_search              -> 0   # the argmax
apply_temperature(0.5)     -> [6.0, 4.0, 2.0, 0.0, -2.0]
apply_temperature(2.0)     -> [1.5, 1.0, 0.5, 0.0, -0.5]

sample_basic over 10,000 draws
  ['0.640', '0.230', '0.086', '0.032', '0.011']

sample_top_k(2) over 10,000 draws
  ['0.736', '0.264', '0.000', '0.000', '0.000']

sample_top_p(0.7) over 10,000 draws
  ['0.734', '0.267', '0.000', '0.000', '0.000']

sample_basic reproduces the true distribution to three decimals
top_k(2) and top_p(0.7) agree because the cumulative probability reaches 0.7 at the second token, so both keep the same pair. Renormalised, 0.636 of 0.870 is 0.731 against the measured 0.736

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python -m pvml.sampling.sampler
prompt    'Jingle bells, jingle bells, jingle all the way'
ours      'Jingle bells, jingle bells, jingle all the way down to the top of the mountain.'
reference 'Jingle bells, jingle bells, jingle all the way down to the top of the mountain.'

matches transformer_lens : True
matches ARENA's expected : True

Greedy is the only setting of the stochastic strategies with no draw in it, so it is the only one checkable exactly

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python -m pvml.sampling.beams
prompt 'When I was'

num_beams=3
   -14.830  'When I was a kid, I used to go to the movies with my'
   -16.321  'When I was a kid, I used to go to the movies. I'
   -16.899  'When I was a kid, I used to go to the movies and watch'

no_repeat_ngram_size=3, 25 tokens
   -33.603  'When I was a kid, I used to go to the movies with my dad. He was a big movie star, but he was also'

Beam search is deterministic, so it gets the exact check the rest of the sampler cannot. All three sequences match GPT2LMHeadModel.generate(num_beams=3, num_return_sequences=3) token for token

What runs on `Sampler`

experiments/sample.py, the REPL:

python

# before
print(f"\n{hooked.generate(line, verbose=False, **sampler.kwargs())}\n")

# after
print(f"\n{sampler.sample(line, args)}\n")

experiments/_run.py, the sample printed after every training epoch:

python

# before
return hooked.generate(prompt, max_new_tokens=50, temperature=0.7, top_p=0.95, verbose=False)

# after
return Sampler(model, tokenizer).sample(prompt, args)

src/pvml/reference/gpt2.py is now the only file in the repo that imports HookedTransformer, and it only does so to load GPT-2 for the parity checks.

What's next

Skipping the KV cache section in the curriculum and moving on to mechanistic interpretability, which is what I came for

Training a Transformer

2026-09-29T00:00:00-04:00

A continuation of Understanding Transformers, where I ported ARENA's GPT-2 implementation into PvML one module at a time. To check each module I loaded GPT-2's own weights into it and compared the output against the reference, which is how I know the implementation is right.

But every weight in it was downloaded. Nothing had been learned.

This post is the other half: a loss, a dataset, and a training loop, then a model of my own trained on TinyStories.

PvML: https://github.com/d0mzw/PvML Weights: https://huggingface.co/d0mzw/pvml-tinystories

Disclaimer: these notes come from working through the ARENA 3.0 curriculum. I claim no credit for the original material, and this is not affiliated with or endorsed by ARENA.

The training loop

flowchart TB data["TinyStories tokenized, chunked to n_ctx"] data --> loader["DataLoader batch (batch, posn)"] loader --> model["Transformer logits (batch, posn, d_vocab)"] model --> loss["get_log_probs -mean() -> loss"] loader --> loss loss --> back["loss.backward() gradients on every parameter"] back --> step["optimizer.step() AdamW"] step -->|"next batch"| loader

The model is the same one from the last post. Everything here is the machinery around it: where the batches come from, how a prediction is graded, and what moves the weights.

Loss

flowchart TB logits["logits (batch, posn, d_vocab)"] tokens["tokens (batch, posn)"] logits -->|"log_softmax(-1)"| lp["log_probs (batch, posn, d_vocab)"] lp -->|"[:, :-1] drop the last position"| preds["predictions (batch, posn-1, d_vocab)"] tokens -->|"[:, 1:] drop the first token"| targets["targets (batch, posn-1)"] preds --> gather["gather along d_vocab"] targets --> gather gather --> out["log_probs (batch, posn-1)"]

How much probability the model gave to the token that actually came next, at every position.

`get_log_probs`

python

log_probs = logits.log_softmax(dim=-1)

return log_probs[:, :-1].gather(dim=-1, index=tokens[:, 1:].unsqueeze(-1)).squeeze(-1)

logits is (batch, posn, d_vocab), a raw score for every token in the vocabulary at every position, which log_softmax turns into log-probabilities.
- Our sentence "The cat sat on the mat" is 7 tokens because of the prepended BOS, so (1, 7, 50257)
log_probs[:, :-1] drops the last position, which predicts a token that is not in the sequence, so there is no answer to grade it against. (1, 7, 50257) -> (1, 6, 50257)

position:   0      1      2      3      4      5      6
token:    <BOS>   The    cat    sat     on    the    mat
predicts:  The    cat    sat     on    the    mat     ?
                                                      └── nothing in the sequence to compare against

tokens is (batch, posn), the ids the text was tokenized into. Every token except <BOS> is also an answer: the one the position before it should have predicted.
tokens[:, 1:] is the sequence shifted left by one, which is exactly that list of answers. <BOS> drops out because nothing predicts it. (1, 7) -> (1, 6)
.gather(dim=-1, index=...) uses each answer token as an index into that position's 50257 scores, and returns the one it lands on. (1, 6, 50257) -> (1, 6, 1)
- For example, at position 0 (<BOS>), shifting left by one gives the answer token The, and gather returns the log-probability log_probs assigned to The at that position.
.unsqueeze(-1) adds an axis of length 1 on the end, .squeeze(-1) takes it away again: (1, 6) -> (1, 6, 1) going in, (1, 6, 1) -> (1, 6) coming out. They are only there because gather needs the index to have the same number of axes as the thing it indexes into. What is left is one log-probability per graded prediction.

Module output

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python -m pvml.training.losses
logits    (1, 7, 50257)
log_probs (1, 6)   # one per prediction, not per token

cross entropy, trained model : 4.578 nats
cross entropy, uniform guess : 10.825 nats
mean probability of the true next token: 0.114

per prediction
     <BOS> -> The      logprob  -3.278   p 0.0377
       The ->  cat     logprob  -9.222   p 0.0001
       cat ->  sat     logprob  -6.844   p 0.0011
       sat ->  on      logprob  -1.493   p 0.2248
        on ->  the     logprob  -0.870   p 0.4191
       the ->  mat     logprob  -5.761   p 0.0031

7 tokens give 6 predictions. The last position has no next token to be graded on
log(d_vocab) = 10.825 nats is what uniform guessing scores. GPT-2 gets 4.578 on this sentence, and the gap is what it knows
Some predictions are nearly determined by context (on -> the, p 0.42), others not at all (The -> cat, p 0.0001). Average loss is always a mix of the two, which is why it never approaches zero

Data

flowchart TB raw["TinyStories 2,119,719 stories"] raw -->|"tokenize_and_concatenate"| stream["one token stream glued end to end"] stream -->|"chop every n_ctx"| chunks["3,731,146 chunks (n_ctx,) each"] chunks -->|"train_test_split"| train["train 3,730,146"] chunks --> test["test 1,000"] train --> loader["DataLoader (batch, posn)"] test --> loader2["DataLoader (batch, posn)"]

A training example is 128 contiguous tokens, not a story.

`tinystories_loaders`

python

tokenized = tokenize_and_concatenate(
    dataset, tokenizer, streaming=False,
    max_length=cfg.n_ctx, column_name="text",
    add_bos_token=True, num_proc=8,
)

split = tokenized.train_test_split(test_size=test_size, seed=seed)

train_test_split shuffles first, so the held-out 1000 are not all from the tail of the corpus. seed= keeps the split the same across runs
The loaders are returned rather than left as a module global, so the trainer takes its data as an argument

Module output

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python -m pvml.data.tinystories
train chunks : 3,730,146
test chunks  : 1,000
chunk length : 128 tokens
batch        : (4, 128)  keys=['tokens']

first chunk, decoded back to text:

<|endoftext|> park. One day, she saw a big white airplane in the sky. She pointed at it and said, "Look, Mommy! A pilot is flying that plane!"

Her mommy said, "Yes, Lily. Pilots fly airplanes."

The chunk begins mid-sentence, on the tail of a story that is not in it. Chunks are cut from the stream, not from stories
3.7M chunks at batch size 32 is 116,000 steps for one full epoch, so training caps steps per epoch rather than running one to the end

Trainer

The loop from the top of this post, with the pieces filled in.

`TrainingArgs`

python

batch_size: int = 32
epochs: int = 10
max_steps_per_epoch: int = 500
lr: float = 1e-3
weight_decay: float = 1e-2

Separate from Config, which holds the architecture. These change between runs, those change between models
An epoch is a full pass over the training set, but 3.7M chunks is 116,000 steps, so max_steps_per_epoch cuts it short

`training_step`

python

tokens = batch["tokens"].to(self.device)
loss = -get_log_probs(self.model(tokens), tokens).mean()
loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()

-mean() turns the (batch, posn-1) log-probabilities into one number to minimise
backward() fills in a gradient on every parameter, step() applies it, zero_grad() clears it for the next batch

`evaluate`

python

predictions = self.model(tokens)[:, :-1].argmax(dim=-1)
correct += (predictions == tokens[:, 1:]).sum().item()

The same [:, :-1] and [1:] offset as the loss, on the held-out chunks the model never trains on
argmax instead of gather: how often the model's top pick was right, rather than what it scored the right answer

Module output

{"step": 1, "loss": 10.83082389831543, "epoch": 0}
{"step": 2, "loss": 10.77728271484375, "epoch": 0}
...
{"step": 60, "loss": 7.720832824707031, "epoch": 1}
{"step": 60, "accuracy": 0.0453986220472441, "epoch": 1}

Sampler

flowchart TB dir["run_dir config.json + model.pt"] dir -->|"Config(**saved)"| cfg["Config d_model, n_heads, n_ctx, ..."] cfg --> model["Transformer our modules"] dir -->|"t.load(model.pt)"| model model -->|"state_dict, strict=False"| hooked["HookedTransformer our weights, their loop"] repl["REPL a line starting with : sets a knob"] --> knobs["Sampler temperature, top_k, top_p, len"] repl -->|"anything else is a prompt"| gen knobs --> gen["hooked.generate()"] hooked --> gen gen --> text["text"]

Update, 3 October 2026
The generation loop has since been ported. See Sampling from a Transformer. The rest of this section describes how it worked at the time of writing.

I have not ported the generation loop. experiments/sample.py loads the trained weights into a TransformerLens HookedTransformer and calls its generate.

A run directory is self-contained: config.json rebuilds the architecture and model.pt fills it. Nothing else is needed to bring a finished run back
strict=False because TransformerLens carries buffers we do not, the causal mask among them. The parameter names match, which is the part that matters
The Sampler class holds nothing but the knobs and how they read back. It exists so the REPL can print the current settings and convert them to the arguments generate expects

Options

Command	Effect
`:temp 0.7`	Divides the logits before the softmax. Below 1 sharpens the distribution, above 1 flattens it
`:greedy`	Always take the most likely token. The same as temperature 0
`:topk 40`	Keep the 40 best tokens, discard the rest, renormalise
`:topp 0.95`	Keep the smallest set of tokens whose probabilities sum to 0.95
`:len 60`	How many tokens to generate

top_k and top_p are alternatives, not a pair. Setting one clears the other, because TransformerLens checks top_k first and it would silently win
The defaults are temperature 0.7 with top_p 0.95, which is what every sample in this post used

Module output

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python experiments/sample.py runs/tinystories-d128-l6-h4-ctx512-40k

runs/tinystories-d128-l6-h4-ctx512-40k
  14,171,473 parameters, 6 layers, d_model 128
  temp 0.7, top_p 0.95, 50 tokens
  a prompt generates, :help for commands, :q to quit

pvml> Once upon a time
Once upon a time, there was a little girl named Lily. She loved to play in her
backyard, especially when she saw a big, scary monster! She was so scared that
she started to cry.

Her mom heard her cry and said, "Don

pvml> :greedy
  greedy, 50 tokens
pvml> Once upon a time
Once upon a time, there was a little girl named Lily. She loved to play outside
in the sunshine. One day, she went to the park to play with her friends. She saw
a big tree and decided to climb it.

As she climbed the

pvml> :temp 1.4
  temp 1.4, top_p 0.95, 50 tokens
pvml> Once upon a time
Once upon a time a painter who was the most courageous, dirty every painting
player had been wanting to take off the beautiful thin balloons home. The weary
square gained swiftly and carried, paying with daily thin brushes for his
reliable folder again to admire them whenever most happily grinding balloons

Experiments

Environment

Component	Detail
Device	ASUS ROG Flow Z13, Strix Halo (Ryzen AI MAX+ 395)
OS	Ubuntu 26.04.1 LTS, kernel 7.0.0-31-generic
GPU	Radeon 8060S integrated graphics, 128GB unified memory
PyTorch	2.12.0a0+rocm7.13, ROCm nightly for gfx1151

`TinyStories`

https://huggingface.co/datasets/roneneldan/TinyStories

Dataset({ features: ['text'], num_rows: 2119719 })

2.1M stories written with the vocabulary of a three or four year old, generated by GPT-3.5 and GPT-4 for the paper TinyStories: How Small Can Language Models Be and Still Speak Coherent English?
Simple enough for a model of a few million parameters to learn on one machine

Story length in GPT-2 tokens, on a 20,000 story sample:

Mean	Median	p90	p99	Longest
222	191	358	641	1106

The corpus is 478M tokens, 3,731,146 chunks at n_ctx=128
20,000 steps at batch 32 and n_ctx 128 is 81.9M tokens, 0.17 epochs. No run here is limited by running out of data

Plan

Run	Question
`tinystories-d32-l4-h16-ctx128-5k`	how far does the smallest model get
`tinystories-d128-l6-h4-ctx128-5k`	does more capacity help at the same steps
`tinystories-d256-l8-h8-ctx128-5k`	does capacity keep paying, or level off
`tinystories-d32-l4-h16-ctx128-20k`	is the small model limited by size or by training
`tinystories-d128-l6-h4-ctx128-20k`	does the medium model have a ceiling of its own
`tinystories-d128-l6-h4-ctx512-40k`	added later: was the wandering a context limit

Steps are held at 5,000 for the first three and 20,000 for the next two, so capacity and training length move one at a time
n_ctx is 128 across those five. Context changes what a model can see, not just how big it is, so it gets its own run
The sixth was not planned. It came out of reading the samples from the first five

Results

The first five ran 22:50 to 01:35, the sixth 02:59 to 13:40 the next day.

Run	n_ctx	Steps	Loss	Accuracy	Minutes	Last 500 steps
`d128-l6-h4-ctx512-40k`	512	40,000	1.732	0.573	641.3	-0.002
`d128-l6-h4-ctx128-20k`	128	20,000	2.080	0.517	67.3	-0.005
`d128-l6-h4-ctx128-5k`	128	5,000	2.369	0.476	17.0	-0.036
`d256-l8-h8-ctx128-5k`	128	5,000	2.497	0.461	36.6	-0.064
`d32-l4-h16-ctx128-20k`	128	20,000	2.782	0.421	34.6	+0.006
`d32-l4-h16-ctx128-5k`	128	5,000	2.996	0.395	8.4	-0.024

Only the five at n_ctx 128 are comparable with each other. Predicting a token from 511 tokens of context is an easier problem than from 127, so part of the top row's lead is the task, not the model. Re-chunking also changes the split, so the two windows were not evaluated on the same held-out text.

Capacity beats steps. d128 reaches 2.369 in 5,000 steps and 17 minutes. d32 needs 20,000 steps and 35 minutes to reach only 2.782, so four times the steps and twice the clock still lands short
d32 uses 16 heads against d128's 4, so its d_head is 2 where theirs is 32. The capacity comparison therefore mixes width with unusually narrow heads, and a rerun at d32 with 4 heads would separate them
Of the five, only d32 at 20,000 steps has converged. Its last 500 steps went up 0.006, noise around a floor near 2.78. d128 at 20,000 is still moving and d256 was descending fastest of anything when it stopped
d256 losing to d128 at 5,000 steps looks like undertraining rather than a ceiling, though nothing here rules the ceiling out. First epoch 5.404 against d128's 4.343, then the steepest end slope of the five
Every run uses a constant lr=1e-3 with no warmup and no decay, inherited from the tiny model. Warmup plus cosine decay is the first thing to try on d256 before concluding anything about its capacity
d256 at 20,000 steps is still unrun. The end slopes suggest it would overtake, which is a hypothesis and not a result

What the best of these writes

The best and the worst of the five on the same two prompts. Temperature 0.7, top_p 0.95, seed 0, 60 new tokens. Each block holds only what the model wrote after the prompt above it.

d128-l6-h4-ctx128-20k, final loss 2.080.

Prompt Once upon a time:

there was a little girl called Sally. She loved her garden and always liked to
play in the garden. One day, Sally was playing in the garden when she heard a
loud noise.

Suddenly, the noise disturbed her. It was so loud that the girl couldn't see
her house. She

Prompt Lily went to the park and saw:

a big tree. She wanted to climb it, but her mom said no. She said, "You can't
climb the tree, Lily. It is not safe. It is not safe. You can fall from the
tree."

Lily was sad, but she did not want to climb

d32-l4-h16-ctx128-20k, final loss 2.782, same prompts and settings.

Prompt Once upon a time:

there was a little girl named Sally. One day, Sally went to the park to play.
She saw a big tree. She wanted to play outside. So, she asked her mommy for
help. Her mommy said yes.

Suddenly, the squirrel started to run and higher.

Prompt Lily went to the park and saw:

a big, orange bear. She was so happy and looked around the tree. She said, "I
like the bear. I found a pear. I like it." She said, "Yes, I want to play with
it. Can I help you?"

She ran to the ball and

Both are fluent and both hold one name throughout, so neither is the difference
d32 introduces things nothing set up: a squirrel, a pear, a ball. d128 keeps referring to what it already named
d128 holds a short causal chain: mom refuses, Lily is sad. d32 has Lily answer "Yes," to a question nobody asked
d128 still repeats. It is not safe twice within four words
The samples stop mid-sentence because the cap is 60 tokens, not because the model ran out

Plots

shell

python experiments/plot.py

Reads every run with a summary.json. Runs are grouped by window, so the four charts below hold only the five at n_ctx 128.

Who plateaued and who was still descending when it stopped.

The same runs priced in minutes rather than steps.

Final loss against non-embedding parameters.

Next-token accuracy per epoch.

Per-step loss is too noisy to read, so it is averaged into 100-step windows
Only non-embedding parameters are plotted. The vocabulary tables dominate these models and do not change with depth or width, so including them would stack every run in the same place
Loss against wall clock is the chart that changes a decision. Per step a bigger model always looks better, per minute it sometimes does not

A longer window

The five runs above all used n_ctx=128, and every one of them wrote stories that wandered. The obvious suspect was capacity. It was not.

Why 128 was the wrong number

How much of a story each window holds, against a median of 191 tokens:

n_ctx	Stories that fit whole
128	5.2%
512	96.8%
1024	99.9%

tokenize_and_concatenate glues the corpus into one stream and chops it into n_ctx chunks, so a chunk starts and ends wherever it lands. At 128 tokens that is almost always mid-story
Endings were never the missing piece. Chunks land across story boundaries at a rate of roughly 0.58 per chunk, so the earlier models read endings constantly. What they never saw was one whole arc, start to finish, inside a single window
512 is chosen on coverage, not on cost. It holds 96.8% of stories whole and 1024 adds three points. Measured at equal tokens per step, doubling to 1024 costs 1.28x, which would have been affordable

The run

tinystories-d128-l6-h4-ctx512-40k

Same model as the best ctx128 run: d_model 128, 6 layers, 4 heads. Same batch_size 32, lr 1e-3, seed 0
40,000 steps at this window is 655M tokens, 1.37 epochs of the corpus and 8x the token budget of the best ctx128 run
max_steps_per_epoch moved from 500 to 1000, which halves the number of evals for the same total. It changes how often training pauses, not what the model sees

It reached 1.732 and 0.573, against 2.080 and 0.517 for the same model at n_ctx 128. For the reason given above, that gap is not all model, so the samples are the evidence that counts.

What changed

Temperature 0.7, top_p 0.95, seed 0, and a 250 token cap neither of these two reaches.

Prompt Once upon a time:

there was a little girl called Sally. She loved to explore and try new
things.

One day Sally went to a big, modern building. She was so excited! Inside the
building, she found a box. Inside, there were lots of toys, but Sally was so
excited. She could not wait to play with them.

Sally was so happy. She could not believe it! She had so many new toys to
play with. Sally was so excited that she couldn't wait to play with them
again!

The next day, Sally took her new toys to the museum and had lots of fun. She
found lots of new things to play with. She was so happy that she had a great
time!

Prompt Lily went to the park and saw:

a big tree. She wanted to climb it and see what was on the other side. She
climbed and climbed and reached the top. She was very happy and excited. She
saw the whole park and saw the flowers and the birds. She saw a butterfly and
the birds and the squirrels. She felt very happy and proud. She felt very
special and loved. She thanked the squirrel and said thank you. The squirrel
smiled and said thank you. Lily was glad she could go to the park and see more
things. She felt very happy and proud.

Both stop on their own, the first after 150 tokens. Raising the cap from 250 to 400 changes neither. That is these two samples, not the general case: the five-seed count below is seven of ten
It has a shape: a setup, a next day, a close. At n_ctx 128 the model only ever saw slices of a story, never a whole one
Repetition got worse. so excited three times, wait to play with them twice
Reference drifts. The modern building becomes the museum, and Lily thanks a squirrel that was only an item in a list of things she saw

Across five seeds per prompt at the same settings, seven of the ten stories reached end-of-text on their own and three ran to the 250 token cap. The three that hit the cap are the three that fell into a repetition loop, I love you and She wished she had repeating until the cap stopped them. So the window bought story structure, and repetition is what still breaks it. Coherence within a story is uneven rather than absent: Timmy fixing a picture with glue holds together, while a bear carries candy out of a tree that was never mentioned.

Cost

10.7 hours on the Z13, against 67 minutes for the ctx128 run of the same model
0.962 sec/step against 0.202, both measured end to end on the real runs. That is 4.8x for 4x the tokens per step
Profiling the forward pass alone at batch 32, so no backward and no optimiser, puts most of the time in the unembedding rather than attention. At n_ctx 128 it is 28.0ms of 37.0ms, 75.7%. At 512 it is 108.8ms of 199.4ms, 54.6%. Projecting d_model 128 up to 50,257 logits at every position is a larger matmul than anything in the blocks
The blocks are what make the cost superlinear. Going 128 to 512 is 4x the tokens, and they take 10.1x the time while the unembedding takes 3.9x. So attention's quadratic term explains the excess over 4x, while the unembedding explains the bulk of the absolute cost
The n_ctx=512 chunking is not the same dataset as n_ctx=128, so the corpus is tokenized again. That is a one-time pass and 4GB of cache

Understanding Transformers

2026-09-27T00:00:00-04:00

I've been self-studying ARENA, an AI-safety curriculum whose first chapter has you rebuild GPT-2 from scratch. Transformers were the gap I kept hitting. They're the core of every LLM, and the interpretability side I actually care about is hard to follow without understanding what sits underneath it.

So I ported the whole chapter into my own modular implementation, loading GPT-2's weights into each module as I went and comparing the output against the reference, until every module matched. These are my notes from doing that, module by module, with the shapes and diagrams I needed to follow it.

PvML: https://github.com/d0mzw/PvML

Disclaimer: these notes come from working through the ARENA 3.0 curriculum. I claim no credit for the original material, and this is not affiliated with or endorsed by ARENA.

Config

python

@dataclass
class Config:
    d_model: int = 768
    debug: bool = False  # per-module __main__ turns it on
    layer_norm_eps: float = 1e-5
    d_vocab: int = 50257
    init_range: float = 0.02
    n_ctx: int = 1024
    d_head: int = 64
    d_mlp: int = 3072
    n_heads: int = 12
    n_layers: int = 12

Embedding module

flowchart TB tokens["tokens (batch, posn) token ids"] tokens -->|"W_E[tokens]"| embed["Embed out (batch, posn, d_model)"] tokens -->|"shape only"| pos["PosEmbed out (batch, posn, d_model)"] embed --> sum(("+")) pos --> sum sum --> resid["residual stream (batch, posn, d_model)"]

Integers become vectors here, and the two results are added to form the residual stream that every later layer reads from and writes back into.

The cat sat on the mat is tokenized first, in the shape (batch, posn) = (1, 7).

ids             50256    464    3797    3332    319     262    2603
decoded <|endoftext|>  'The'  ' cat'  ' sat'  ' on'  ' the'  ' mat'

50256 is <|endoftext|>, prepended as the beginning-of-sequence marker

`Embed`

Each token id is replaced by its row of W_E through lookup.

python

self.W_E = nn.Parameter(t.empty((cfg.d_vocab, cfg.d_model)))

out = self.W_E[tokens]

(d_vocab, d_model)[(batch, posn)]  ->  (batch, posn, d_model)
 50257    768       1      7            1      7     768

`PosEmbed`

Only the shape of tokens is read, because position embeddings encode where a token sits rather than which token it is.

python

self.W_pos = nn.Parameter(t.empty((cfg.n_ctx, cfg.d_model)))

batch, seq_len = tokens.shape
sliced = self.W_pos[:seq_len]
out = einops.repeat(sliced, "seq d_model -> batch seq d_model", batch=batch)

W_pos is an (n_ctx, d_model) matrix, and its n_ctx rows are sliced down to seq_len = posn = 7
repeat copies that sliced table for every sequence in the batch

Module output

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python -m pvml.modules.embedding
text='The cat sat on the mat'
ref.to_str_tokens(text)=['<|endoftext|>', 'The', ' cat', ' sat', ' on', ' the', ' mat']
tokens.shape=torch.Size([1, 7])

[embed]                       in: (1, 7)          # token ids, not activations
[embed]                      W_E: (50257, 768)    # one row per vocab entry
[embed]                      out: (1, 7, 768)
[pos_embed]                   in: (1, 7)          # values ignored, only the shape is read
[pos_embed]                W_pos: (1024, 768)     # one row per position, up to n_ctx rows
[pos_embed]               sliced: (7, 768)        # W_pos[:seq_len]
[pos_embed]                  out: (1, 7, 768)     # same table repeated for each sequence

LayerNorm

flowchart TB resid["residual (batch, posn, d_model)"] resid --> mean["residual_mean (batch, posn, 1)"] resid --> std["residual_std (batch, posn, 1)"] resid --> norm["(residual - mean) / std (batch, posn, d_model)"] mean --> norm std --> norm norm -->|"* w + b"| out["out (batch, posn, d_model)"]

A residual stream of shape (batch, posn, d_model) comes in, and each position's d_model vector is standardised on its own by (residual - mean) / sqrt(var + eps), with the mean and variance taken across d_model rather than across positions or the batch. The learned w and b then scale and shift it, and the normalised residual stream leaves at the shape it arrived in.

Statistics

The two numbers each position gets standardised by.

python

residual_mean = residual.mean(dim=-1, keepdim=True)
residual_std = (residual.var(dim=-1, keepdim=True, unbiased=False) + self.cfg.layer_norm_eps).sqrt()

Reduced along d_model, deriving both the mean and the std in the shape (batch, posn, 1) = (1, 7, 1). Each of those 7 values is computed from one position's own 768 numbers
keepdim=True leaves the reduced axis as size 1 so it broadcasts back: (1, 7, 768) - (1, 7, 1) works, (1, 7, 768) - (1, 7) does not
unbiased=False divides by N, not N-1, which is 768 not 767
layer_norm_eps is a tiny constant, 1e-5, added to the variance before the sqrt so the denominator can never be zero. A position whose 768 numbers are all identical has a variance of 0, and dividing by that would give NaN

Scale and shift

The learned half of the layer, applied once the standardising is done.

python

self.w = nn.Parameter(t.ones(cfg.d_model))
self.b = nn.Parameter(t.zeros(cfg.d_model))

residual = (residual - residual_mean) / residual_std
out = residual * self.w + self.b

w and b are both (d_model,) = (768,), and they start as all ones and all zeros, so before any training out is exactly the standardised residual and the layer does nothing of its own
One w and one b per LayerNorm, reused at all 7 positions, while the mean and std were worked out per position. The normalisation is local, the learned correction on top of it is not
Broadcasting pads the missing leading axes: in (1, 7, 768) * (768,) the w is read as (1, 1, 768) and stretched to (1, 7, 768)
Standardising throws the input's scale away, so [1, 2, 3, 4] and [10, 20, 30, 40] come out identical and after centering only the direction survives. w and b put a scale back, one the model has learned rather than one the input happened to arrive with

Module output

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python -m pvml.modules.normalization
[ln]                    residual: (1, 7, 768)     # as it arrives
[ln]               residual_mean: (1, 7, 1)
[ln]                residual_std: (1, 7, 1)
[ln]                    residual: (1, 7, 768)     # reassigned: (residual - mean) / std
[ln]                        w, b: (768,)          # padded to (1, 1, d_model), stretched to out
[ln]                         out: (1, 7, 768)     # * w + b

Attention module

flowchart TB resid["normalized_resid_pre (batch, posn, d_model)"] resid -->|"@ W_Q + b_Q"| q["q (batch, posn, nheads, d_head)"] resid -->|"@ W_K + b_K"| k["k (batch, posn, nheads, d_head)"] resid -->|"@ W_V + b_V"| v["v (batch, posn, nheads, d_head)"] q --> scores["attn_scores (batch, nheads, posn_Q, posn_K) contracts d_head"] k --> scores scores --> scale["/ sqrt(d_head)"] scale --> mask["apply_causal_mask upper triangle to -inf"] mask --> soft["softmax(-1)"] soft --> pattern["attn_pattern (batch, nheads, posn_Q, posn_K) rows sum to 1"] pattern --> z["z (batch, posn, nheads, d_head) contracts posn_K"] v --> z z -->|"@ W_O + b_O"| out["attn_out (batch, posn, d_model) contracts nheads and d_head"]

Each position asks what it needs (q), every earlier position advertises what it has (k), and the match decides how much of their content (v) gets copied back.

`q, k, v`

All three read the same input through three different learned matrices. Same einsum, same output shape, different weights.

python

q = (
    einops.einsum(
        normalized_resid_pre,
        self.W_Q,
        "batch posn d_model, nheads d_model d_head -> batch posn nheads d_head",
    )
    + self.b_Q
)

k and v are the same three lines again, reading the same normalized_resid_pre through W_K, b_K and W_V, b_V.

batch posn d_model,  nheads d_model d_head  ->  batch posn nheads d_head
1     7    768       12     768     64          1     7     12    64
           ▲                ▲

d_model appears in both inputs and not in the output, so it is summed away. Each position's 768 numbers become 64 per head
nheads and d_head come from the weight, batch and posn pass straight through
W_Q is (nheads, d_model, d_head) = (12, 768, 64), one projection per head stacked on a leading axis, so all 12 heads run as a single batched matmul rather than 12 separate modules
b_Q is (nheads, d_head) = (12, 64), read as (1, 1, 12, 64) and stretched: the same bias at every position, a different one per head

`attn_scores`

For each head, dot every query vector (q) with every key vector (k) over d_head, giving a (posn_Q, posn_K) grid of match scores per head.

python

attn_scores = einops.einsum(
    q,
    k,
    "batch posn_Q nheads d_head, batch posn_K nheads d_head -> batch nheads posn_Q posn_K",
)

d_head is contracted, so the 64 numbers of a query and a key collapse to one score
batch and nheads appear in both inputs and the output, so they are batched over rather than summed: an independent 7 x 7 grid per sequence per head
The two position axes get different names so they survive as separate dimensions. posn_Q is who is asking, posn_K is who is being read
These scores are raw. For entries of roughly unit variance the dot product of two 64-dimensional vectors has a standard deviation that grows like sqrt(d_head), so dividing by sqrt(64) = 8 is exactly what cancels it. That pulls them back to a range where softmax does not saturate into a near one-hot row and kill the gradients

`apply_causal_mask(attn_scores)`

Set everything above the diagonal of the scaled scores to -inf, so no position can attend to a later one.

python

all_ones = t.ones(attn_scores.size(-2), attn_scores.size(-1), device=attn_scores.device)
mask = t.triu(all_ones, diagonal=1).bool()
attn_scores.masked_fill_(mask, self.IGNORE)

        k0 k1 k2 k3 k4 k5 k6
    q0 [ 0  1  1  1  1  1  1 ]    1 = masked, the future
    q1 [ 0  0  1  1  1  1  1 ]
    q2 [ 0  0  0  1  1  1  1 ]
    q3 [ 0  0  0  0  1  1  1 ]
    q4 [ 0  0  0  0  0  1  1 ]
    q5 [ 0  0  0  0  0  0  1 ]
    q6 [ 0  0  0  0  0  0  0 ]

diagonal=1 leaves the diagonal unmasked, so a position can attend to itself
The mask is 2-D (7, 7) and broadcasts over batch and nheads. Causality depends only on position, so every sequence and every head gets the same one
IGNORE is a registered buffer, so -inf follows .to(device). all_ones needs an explicit device= because a tensor made inside a method does not
masked_fill_ is in place: it mutates attn_scores

`attn_pattern`

Softmax each row into a distribution over the keys it can see.

python

attn_scores_masked = self.apply_causal_mask(attn_scores / self.cfg.d_head**0.5)
attn_pattern = attn_scores_masked.softmax(-1)

The mask has to land before the softmax, not after, so exp(-inf) = 0 never enters the denominator and the surviving weights still sum to 1
Shape is unchanged throughout, (1, 12, 7, 7). Row q3 of each head now holds four non-zero weights summing to 1, and three zeros

`z`

Weighted average of v for each query position, contracting posn_K.

python

z = einops.einsum(
    v,
    attn_pattern,
    "batch posn_K nheads d_head, batch nheads posn_Q posn_K -> batch posn_Q nheads d_head",
)

posn_K is the contracted axis, so the sum runs over the positions being read. posn_Q survives, so the result is indexed by the position doing the asking
Concretely, at query position 3: its row of attn_pattern holds four non-zero weights, and z at that position is those four value vectors each multiplied by its weight and added together, separately for each of the 12 heads. The three masked positions carry weight 0, so they contribute nothing
z comes out the same shape as v, (1, 7, 12, 64). Nothing changed size. What changed is that each position's vector is now a blend of the positions it could see, instead of only its own
This is the only step in the whole model where information crosses between positions. Embedding, LayerNorm and the MLP all act on each position independently, so every bit of context a token ever gets arrives through this one sum

`attn_out`

Project each head's d_head back up to d_model and sum over the heads.

python

attn_out = (
    einops.einsum(
        z,
        self.W_O,
        "batch posn_Q nheads d_head, nheads d_head d_model -> batch posn_Q d_model",
    )
    + self.b_O
)

W_O is (nheads, d_head, d_model) = (12, 64, 768), the reverse of W_Q, W_K and W_V. Those project the residual stream down into each head's 64 dimensions, this one projects the result back up to 768
Two names are contracted at once, nheads and d_head, so the 12 heads are summed, not concatenated. Each head's 64 numbers become a full 768-wide vector, and the 12 vectors are added
b_O is (d_model,) = (768,), with no head axis, because it is added once after the heads are already summed. Compare b_Q, which is (12, 64) and applies per head
Output is (1, 7, 768), back to the residual stream's width, ready to be added into it

Module output

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python -m pvml.modules.attention
 W_Q: (12, 768, 64)
 W_K: (12, 768, 64)
 W_V: (12, 768, 64)
 W_O: (12, 64, 768)
 b_Q: (12, 64)
 b_K: (12, 64)
 b_V: (12, 64)
 b_O: (768,)
total: 2,362,368 parameters

Loaded pretrained model gpt2-small into HookedTransformer
[attn]      normalized_resid_pre: (1, 7, 768)     # (batch, posn, d_model)
[attn]                       W_Q: (12, 768, 64)   # (n_heads, d_model, d_head)
[attn]                       b_Q: (12, 64)        # (n_heads, d_head), broadcast over posn
[attn]                         q: (1, 7, 12, 64)  # (batch, posn, n_heads, d_head)
[attn]                       W_K: (12, 768, 64)
[attn]                       b_K: (12, 64)
[attn]                         k: (1, 7, 12, 64)
[attn]                       W_V: (12, 768, 64)
[attn]                       b_V: (12, 64)
[attn]                         v: (1, 7, 12, 64)
[attn]               attn_scores: (1, 12, 7, 7)   # (batch, n_heads, query_pos, key_pos)
[attn]                      mask: (7, 7)          # True above the diagonal = cannot attend
[attn]              attn_pattern: (1, 12, 7, 7)   # rows sum to 1 over visible keys
[attn]                         z: (1, 7, 12, 64)  # weighted average of v, per query position
[attn]                       W_O: (12, 64, 768)   # (n_heads, d_head, d_model), projects back up
[attn]                  attn_out: (1, 7, 768)     # back to the residual stream's width

head 0 attention pattern
   query / key   <BOS>     The     cat     sat      on     the     mat
         <BOS>    1.00    0.00    0.00    0.00    0.00    0.00    0.00
           The    0.93    0.07    0.00    0.00    0.00    0.00    0.00
           cat    0.71    0.10    0.18    0.00    0.00    0.00    0.00
           sat    0.64    0.14    0.04    0.18    0.00    0.00    0.00
            on    0.48    0.15    0.12    0.23    0.03    0.00    0.00
           the    0.60    0.12    0.09    0.16    0.02    0.02    0.00
           mat    0.37    0.09    0.06    0.03    0.08    0.10    0.26

MLP

Expand to four times the residual width, apply the nonlinearity, project back. Every position goes through the same matrices independently, so nothing here mixes positions.

`pre`

The first of the two projections, into the wide space.

python

self.W_in = nn.Parameter(t.empty((cfg.d_model, cfg.d_mlp)))

pre = (
    einops.einsum(
        normalized_resid_mid,
        self.W_in,
        "batch position d_model, d_model d_mlp -> batch position d_mlp",
    )
    + self.b_in
)

W_in is 2-D. The MLP has no head axis, unlike attention's weights
d_model contracts, d_mlp appears: 768 in, 3072 out
Input is the residual stream after attention has been added, through the block's ln2

`post`

The only nonlinearity in the MLP.

python

post = gelu_new(pre)

Without it W_in and W_out would collapse into a single matrix and the whole detour through 3072 dimensions would buy nothing
Applied elementwise in the wide space, so the shape does not change

`mlp_out`

The second projection, undoing the expansion.

python

self.W_out = nn.Parameter(t.empty((cfg.d_mlp, cfg.d_model)))

mlp_out = (
    einops.einsum(
        post,
        self.W_out,
        "batch position d_mlp, d_mlp d_model -> batch position d_model",
    )
    + self.b_out
)

d_mlp contracts: 3072 in, 768 out, back to the width the block adds into
b_out is (d_model,), broadcast over batch and position

Module output

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python -m pvml.modules.mlp
  W_in: (768, 3072)
 W_out: (3072, 768)
  b_in: (3072,)
 b_out: (768,)
 total: 4,722,432 parameters

[mlp]       normalized_resid_mid: (1, 7, 768)     # (batch, posn, d_model)
[mlp]                       W_in: (768, 3072)     # (d_model, d_mlp), contracts d_model
[mlp]                        pre: (1, 7, 3072)    # 4x wider than the residual stream
[mlp]                       post: (1, 7, 3072)    # gelu_new, same shape
[mlp]                      W_out: (3072, 768)     # (d_mlp, d_model), contracts d_mlp
[mlp]                    mlp_out: (1, 7, 768)     # back to (batch, posn, d_model)

max abs diff vs gpt-2: 0.000e+00

4.7M parameters, twice what attention costs. Two thirds of a block is the MLP
d_mlp = 4 * d_model is convention, not a requirement

TransformerBlock

flowchart TB pre["resid_pre (batch, posn, d_model)"] pre --> ln1["ln1 LayerNorm"] ln1 --> attn["attn takes normalized_resid_pre"] attn --> add1(("+")) pre ---> add1 add1 --> mid["resid_mid (batch, posn, d_model)"] mid --> ln2["ln2 LayerNorm"] ln2 --> mlp["mlp takes normalized_resid_mid"] mlp --> add2(("+")) mid ---> add2 add2 --> post["resid_post (batch, posn, d_model)"]

Attention and the MLP each read the stream and add a delta back. Neither replaces it, so every component's contribution stays separable.

`resid_mid`

Attention reads the normalised stream and its output is added back.

python

resid_mid = self.attn(self.ln1(resid_pre)) + resid_pre

ln1 normalises what attention reads. What gets added back is the unnormalised attn_out, so the stream itself never passes through a LayerNorm
The stream keeps its width the whole way: 768 in, 768 out, at every point in all 12 blocks

`resid_post`

The MLP's turn, on the stream attention just wrote into.

python

resid_post = self.mlp(self.ln2(resid_mid)) + resid_mid

ln2 normalises what the MLP reads, exactly as ln1 did for attention
(1, 7, 3072) inside the MLP and (1, 7, 12, 64) inside attention are private scratch space, projected back before anything is added

Module output

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python -m pvml.modules.block
block parameters: 7,087,872

[block]                resid_pre: (1, 7, 768)     # (batch, posn, d_model)
[ln1]                   residual: (1, 7, 768)     # as it arrives
[ln1]              residual_mean: (1, 7, 1)
[ln1]               residual_std: (1, 7, 1)
[ln1]                   residual: (1, 7, 768)     # reassigned: (residual - mean) / std
[ln1]                       w, b: (768,)          # padded to (1, 1, d_model), stretched to out
[ln1]                        out: (1, 7, 768)     # * w + b
[attn]      normalized_resid_pre: (1, 7, 768)     # (batch, posn, d_model)
[attn]                         q: (1, 7, 12, 64)  # (batch, posn, n_heads, d_head)
[attn]               attn_scores: (1, 12, 7, 7)   # (batch, n_heads, query_pos, key_pos)
[attn]              attn_pattern: (1, 12, 7, 7)   # rows sum to 1 over visible keys
[attn]                         z: (1, 7, 12, 64)  # weighted average of v, per query position
[attn]                  attn_out: (1, 7, 768)     # back to the residual stream's width
[block]                resid_mid: (1, 7, 768)     # resid_pre + attn_out
[ln2]                   residual: (1, 7, 768)     # as it arrives
[ln2]                        out: (1, 7, 768)     # * w + b
[mlp]       normalized_resid_mid: (1, 7, 768)     # (batch, posn, d_model)
[mlp]                        pre: (1, 7, 3072)    # 4x wider than the residual stream
[mlp]                       post: (1, 7, 3072)    # gelu_new, same shape
[mlp]                    mlp_out: (1, 7, 768)     # back to (batch, posn, d_model)
[block]               resid_post: (1, 7, 768)     # resid_mid + mlp_out

max abs diff vs gpt-2: 1.144e-05

First module where the diff is not exactly zero. Values reach 120, so 1e-5 is float32 accumulation across the four sub-modules of a block, not an error
tag_tree names the children by position, so the two LayerNorms print as ln1 and ln2

Unembed

flowchart TB resid["normalized_resid_final (batch, posn, d_model)"] resid -->|"@ W_U + b_U"| logits["logits (batch, posn, d_vocab) contracts d_model"]

One score per vocabulary entry, per position. This is where d_model becomes d_vocab and the model commits to a prediction.

`logits`

A single matmul against a (768, 50257) table, the mirror of W_E.

python

self.W_U = nn.Parameter(t.empty((cfg.d_model, cfg.d_vocab)))
self.b_U = nn.Parameter(t.zeros((cfg.d_vocab)), requires_grad=False)

logits = (
    einops.einsum(
        normalized_resid_final,
        self.W_U,
        "batch posn d_model, d_model d_vocab -> batch posn d_vocab",
    )
    + self.b_U
)

b_U is frozen. GPT-2 has no unembedding bias, so it stays at zero and only exists to keep the shapes lined up
Input is the stream after all twelve blocks, through ln_final
d_vocab appears in only two places in the whole model, here and at W_E. Everything between them is d_model

Module output

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python -m pvml.modules.unembedding
 W_U: (768, 50257)  requires_grad=True
 b_U: (50257,)  requires_grad=False
total: 38,647,633 parameters

[unembed]   normalized_resid_final: (1, 7, 768)     # (batch, posn, d_model)
[unembed]                      W_U: (768, 50257)    # (d_model, d_vocab), the mirror of W_E
[unembed]                   logits: (1, 7, 50257)   # one score per vocab entry

max abs diff vs gpt-2: 0.000e+00

next-token predictions
     <BOS>  ->  '\n'
       The  ->  ' first'
       cat  ->  ' was'
       sat  ->  ' on'
        on  ->  ' the'
       the  ->  ' floor'
       mat  ->  ','

Every position predicts its own next token from one forward pass. That is what the causal mask is for
sat -> ' on' and on -> ' the' are right. the -> ' floor' is a guess, because position 5 cannot see ' mat'

Transformer

flowchart TB tokens["tokens (batch, posn)"] tokens --> embed["embed W_E[tokens]"] tokens --> pos["pos_embed W_pos[:seq_len]"] embed --> add(("+")) pos --> add add --> resid["residual stream (batch, posn, d_model)"] resid --> block["blocks[i] TransformerBlock"] block -->|"residual = block(residual) for i in range(n_layers)"| resid resid -->|"after 12 blocks"| lnf["ln_final LayerNorm"] lnf --> unembed["unembed @ W_U + b_U contracts d_model"] unembed --> logits["logits (batch, posn, d_vocab)"]

Token ids in, logits out. The module owns no parameters of its own, it just wires the others together.

`forward`

Embed, run the blocks in sequence, normalise, unembed.

python

residual = self.embed(tokens) + self.pos_embed(tokens)

for block in self.blocks:
    residual = block(residual)

logits = self.unembed(self.ln_final(residual))

residual = block(residual) only works because the shape never changes. That is why blocks stack
Nothing normalises the stream itself along the way, so ln_final is needed before the unembedding can read it
nn.ModuleList gives the children paths like blocks.3.attn, which match transformer_lens's module paths. Its hooks are points underneath those, such as blocks.3.attn.hook_z, so the two trees line up and one vocabulary covers the trace and the reference activations

Module output

shell

(.venv) dom@dom-z13:~/Desktop/PvML$ python -m pvml.modules.transformer
parameters      : 163,087,441
missing keys    : []
unexpected keys : 12 (['mask'])

logits (1, 7, 50257)
max abs diff vs gpt-2: 9.155e-05
same argmax everywhere: True

next-token predictions
     <BOS>  ->  '\n'
       The  ->  ' first'
       cat  ->  ' was'
       sat  ->  ' on'
        on  ->  ' the'
       the  ->  ' floor'
       mat  ->  ','

The check loads GPT-2's whole state_dict into our model with strict=False, then runs both on the same tokens.

missing keys: [] is the line that matters. Every parameter our model needs was found in GPT-2's. Because strict=False tolerates a mismatch instead of raising, a single misspelled parameter would leave that weight at its random initialisation, and the model would still run and still produce plausible logits. An empty list is what rules that out
unexpected keys: 12 is the other direction, entries GPT-2 has that we do not. All 12 are named mask, one per block. transformer_lens precomputes its causal mask as a buffer, ours is built inside apply_causal_mask on every call, so there is nothing to load them into
Every earlier check was handed its input from the reference's own activations. This one gets only token ids, so the entire forward pass is ours and no borrowed intermediate can prop up a mistake
9.155e-05 across 50,257 logits at each of 7 positions, and the argmax is identical at every one. That is float32 accumulation over twelve blocks, not an error

Where those parameters live:

  embed        38,597,376
  pos_embed       786,432
  blocks       85,054,464
  ln_final          1,536
  unembed      38,647,633
  total       163,087,441

85M of it is the twelve blocks, and two thirds of each block is its MLP
GPT-2 small is usually quoted as 124M, which is this total minus unembed. The original ties the two vocabulary tables, W_U = W_E transposed, and counts those 38,597,376 parameters once. transformer_lens keeps them as separate tensors, so they are counted twice here

Dom 🆚 Machine

Two Routes Through a Transformer

Path decomposition of transformer output

Two routes

Route 1: the token's own embedding

Route 2: what the head reaches back for

The trick

The formula

The tensor-product notation

The attention pattern \(A\)

The two terms, formally

Why this framing is useful

Sampling from a Transformer

Sampler

strategies

SamplingArgs

next_token

The loop

The missing BOS token

Beam search

Module output

What runs on Sampler

What's next

Training a Transformer

The training loop

Loss

get_log_probs

Module output

Data

tinystories_loaders

Module output

Trainer

TrainingArgs

training_step

evaluate

Module output

Sampler

Options

Module output

Experiments

Environment

TinyStories

Plan

Results

What the best of these writes

Plots

A longer window

Why 128 was the wrong number

The run

What changed

Cost

Understanding Transformers

Config

Embedding module

Embed

PosEmbed

Module output

LayerNorm

Statistics

Scale and shift

Module output

Attention module

q, k, v

attn_scores

apply_causal_mask(attn_scores)

attn_pattern

z

attn_out

Module output

MLP

pre

post

mlp_out

Module output

TransformerBlock

resid_mid

resid_post

Module output

Unembed

logits

`strategies`

`SamplingArgs`

`next_token`

What runs on `Sampler`

`get_log_probs`

`tinystories_loaders`

`TrainingArgs`

`training_step`

`evaluate`

`TinyStories`

`Embed`

`PosEmbed`

`q, k, v`

`attn_scores`

`apply_causal_mask(attn_scores)`

`attn_pattern`

`z`

`attn_out`

`pre`

`post`

`mlp_out`

`resid_mid`

`resid_post`

`logits`

`forward`