Handout #11 — Self-Attention and Next-Token Prediction

Attention lets each token retrieve information from other token positions. The important change from a fixed-window MLP is not just a larger context: the way information is combined now depends on the input itself. This reading develops that operation, then follows it into a next-token predictor.

1Content-dependent retrieval

A fixed-window MLP gives each relative position a fixed block of weights. In attention, position $i$ instead computes a different set of mixing weights for every input. For “The trophy did not fit in the suitcase because it was too large,” a useful representation of “it” may need information about “trophy.” The architecture provides a route to retrieve that information, even across many intervening positions. Learning a useful route still requires suitable training data and parameters.

Write the input as rows $X\in\mathbb R^{T\times D}$. Each position publishes three vectors:

A query describes what a position seeks; a key describes how that position can be matched; a value contains the information it can contribute. All three are learned projections of the same input in self-attention. The matrices are shared across token positions. They are not separate parameters for “trophy,” “suitcase,” or position 8.

Separating keys from values lets the model match on one set of features and retrieve another. For example, keys could encode features useful for locating a noun, while values supply information about its number or meaning. This is an interpretation of their roles, not a guarantee that a particular trained dimension has a human-readable label.

2Compute one query's answer

For query row $q_i$, first score every key, then normalize across keys, then combine the values:

We use $s$ for raw scores and $a$ for normalized attention weights, so the two quantities do not get confused. $d_k$ is the query/key width. Values may have a different width $d_v$; compatible queries and keys must have the same width.

2.1 A complete numerical retrieval

Choose $d_k=2$, $q=[1,0]$, and keys $k_0=[\sqrt2\log2,0]$, $k_1=[0,1]$, $k_2=[0,-1]$. The scaled scores are $[\log2,0,0]$. Their exponentials are $[2,1,1]$, so the weights are $[1/2,1/4,1/4]$.

Adding contributions gives $y=[2,1.5]$. The output need not equal any individual value. Before dropout, it is a convex combination of the values: nonnegative weights sum to one. Later projections, residual additions, and FFNs are not constrained to this convex hull.

2.2 How the routing learns

Raising one score increases its own weight and reduces the others, because they compete for a fixed total. The value path passes information forward; the query/key path learns how to route it. Both paths receive gradients from the final language-model loss.

A large attention weight alone is not a complete explanation of a prediction: a value can be small, point in an unhelpful direction, or be offset by another head or the residual stream. To test a causal claim about model behavior, changing or removing the relevant input and observing the prediction is more informative than inspecting a weight alone.

3Scaling and batching without mixing examples

3.1 Why divide by the square root?

Suppose query and key coordinates are independent, zero-mean, with variance $\sigma^2$. Each product has variance $\sigma^4$, and summing $d_k$ independent products gives variance $d_k\sigma^4$. Dividing by $\sqrt{d_k}$ gives variance $\sigma^4$. This is an initialization-scale argument; trained coordinates need not satisfy the assumptions exactly.

For two scores $[0,8]$, softmax is approximately $[0.000335,0.999665]$. Dividing by $\sqrt{64}=8$ instead gives $[0,1]$ and probabilities $[0.268941,0.731059]$. The second distribution has more room for score changes to alter its weights. Scaling limits dimension-driven saturation; it does not prohibit a trained head from becoming selective.

3.2 The full matrix operation

Rows are queries, columns are keys. Every row sums to one; columns generally do not. For a batch, transpose only the final two dimensions of $K$. Each batch item has its own attention matrix:

4Multiple heads: separate routing decisions

One head produces one distribution over positions for each query. Two heads can retrieve two different mixtures at the same position. For example, one might emphasize a subject and another a nearby modifier. Heads are jointly trained; they are not independent language models voting on a final answer, and useful specialization is not guaranteed.

With model width $D$ and $H$ heads, take $d_h=D/H$. Each head has query, key, and value width $d_h$. Concatenating $H$ outputs recovers width $D$; an output projection $W_O\in\mathbb R^{D\times D}$ mixes the head features.

At $D=8,H=2$, each head returns 4 features. Ignoring biases, the three projection families contain $3D^2=192$ parameters and $W_O$ adds $D^2=64$, totaling 256. Changing the head count while holding $D$ fixed changes head width but not this parameter count.

4.1 Two heads answer the same query differently

Reuse §2.1’s values $v_0=[2,0]$, $v_1=[0,2]$, $v_2=[4,4]$. At the same query position, let two heads’ separate projections produce the following scaled scores. For this example their value projections produce the same values; their routing differs.

For head 1, the sum is $0.8[2,0]+0.1[0,2]+0.1[4,4]=[2,0.6]$. Head 2 instead computes $0.1[2,0]+0.8[0,2]+0.1[4,4]=[0.6,2]$.

Concatenation gives $[2,0.6\;|\;0.6,2]$: four features, with each head’s answer still separately available. Here $H=2,d_h=2,D=4$. Choosing $W_O=I_4$ preserves these four features; a learned $4\times4$ map can mix them. Averaging the two answers would instead give only $[1.3,1.3]$, merging the two routing decisions. Concatenation is not averaging.

Head 1 emphasizes position 0; head 2 emphasizes position 1. Both still include position 2 because finite unmasked softmax scores give positive weights. This is a possible learned behavior, rather than a promise that every trained head specializes this way.

4.2 Fused projection is the same algebra

Store $W_{QKV}=[W_Q\;W_K\;W_V]\in\mathbb R^{D\times3D}$. One multiplication $XW_{QKV}$ returns three adjacent feature groups. Split them, then reshape and transpose to [B,H,T,d_h]. This combines operations; it does not force $Q=K=V$.

Attention weights now have shape [B,H,T,T]. Arithmetic for the score/value products stays proportional to $BT^2D$ at fixed width, but explicitly storing all heads' weights takes $BHT^2$ entries. Similar arithmetic cost does not imply identical activation memory.

5Build a block: mix positions, then transform features

Attention mixes information across positions. The feedforward network transforms the resulting features independently at each position, with shared parameters:

The FFN can expand to $D_{ff}=4D$, apply a nonlinear activation, then return to width $D$. “Position-wise” does not mean context-free: its input already contains information retrieved by attention. It means the FFN itself does not communicate between rows. A linear-only pair would collapse to a single linear map; the intervening nonlinearity matters.

5.1 Two arrangements you should recognize

The post-norm block in slides 25–26 computes:

A pre-norm version computes:

Both keep the residual width equal to $D$. LayerNorm normalizes across features within a token, not across sequence positions or batch items. In pre-norm, each update has an explicit identity path: $d(x+g(\operatorname{LN}(x)))/dx=I+J_{g\circ\operatorname{LN}}$. This helps gradient flow through a deep stack, but does not guarantee stable training under every initialization or learning rate. Revisit Handout #8 for the underlying optimization ideas.

6Depth, positions, and a vocabulary head

Stacking blocks lets later layers retrieve and transform features that earlier layers built. After $L$ blocks, the shape is still [B,T,D]. Apply a vocabulary projection to obtain [B,T,Vocab] logits. If input and output embeddings are tied, the logits for row $h_i$ are $h_iE^\top$ with $E\in\mathbb R^{Vocab\times D}$. Tying is a parameter-sharing choice, not a requirement of attention.

6.1 A precise statement about order

Unmasked self-attention without position information is permutation equivariant: reordering input rows reorders output rows the same way. It is not invariant: outputs do move. Shared FFNs and token-wise normalization preserve this property. For example, swapping rows 0 and 2 swaps their contextual outputs but does not add a signal saying which row originally came first.

Thus an unrestricted block cannot tell that a row was originally “first” just from its array index. Adding position vectors $p_i$ makes $x_i=E[w_i]+p_i$. The same token at two positions now starts with different vectors. Learned absolute embeddings add $T_{max}D$ trainable parameters; fixed sinusoidal encodings do not. In the sinusoidal scheme, even/odd feature pairs use sine and cosine at different frequencies.

7One predictor, two execution modes

7.1 Generation appends one token at a time

Given a prompt, use the logits at its final real position to predict the next token. Sample from softmax, or select the largest logit for greedy decoding; append the chosen token and repeat. Stop at EOS or a chosen length bound. The vocabulary projection is shared across positions. Generation is sequential because the next input depends on the token just chosen.

Extending the prompt changes the next-position prediction. In a correctly causal model, earlier positions' representations do not change merely because a later token was appended, with positions and other settings fixed. That property will justify caching in Handout #12. No cache is needed to understand the generation algorithm; a naive implementation recomputes the growing prefix.

7.2 Training supplies all gold prefixes together

For the cat sat on the mat <eos>, use:

Output at input position $i$ predicts the token at $i+1$. The input token at $i$ is legitimate context. The target at $i+1$ is not. This is teacher forcing: training uses the observed prefix rather than a prefix sampled by the model. All target positions are known, so their losses can be computed together:

Backpropagation updates embeddings, Q/K/V and output projections, FFNs, and trainable normalization and position parameters. Queries, attention weights, and intermediate representations are activations, not independently optimized parameter tables. The softmax function itself has no trainable parameters.

7.3 Why shifting the targets is insufficient

The target “cat” also appears in the input, one column to the right. With unrestricted attention, the “the” query can access the “cat” key and value. A low training loss could then reflect learning to read the answer from the future, a route unavailable during generation. The same problem recurs at every position except the last.

A causal model must enforce that every output at $i$ depends only on inputs at positions $j\le i$. Enforce this in every block: masking only the final block is insufficient because an earlier block could already have copied future information into an allowed row. Handout #12 derives the mask and tests this dependence directly.

7.4 One sequence, from embeddings to loss

Follow the observed token stream A B C A. The three input rows are A B C at positions 0, 1, 2; their targets are B C A. Take vocabulary order $[A,B,C]$, width $D=2$, and one head with $d_k=1,d_v=2$. These are fixed toy parameters, not a trained language model.

1 · Embed tokens and add positions. The lookup rows are $E[A]=[1,0]$, $E[B]=[0,1/2]$, $E[C]=[1/2,1/2]$. Position rows are $p_0=[0,0]$, $p_1=[0,1/2]$, $p_2=[1/2,1/2]$. Therefore:

2 · Project Q/K/V. Choose $W_Q=[1,0]^\top$, $W_K=[\log2,0]^\top$, and $W_V=I_2$. Then $Q=[1,0,1]^\top$, $K=[\log2,0,\log2]^\top$, and $V=X$. Since $d_k=1$, the scale is 1.

3 · Retrieve from the prefix. For row 0, only key 0 is allowed; for row 1, keys 0–1; for row 2, keys 0–2. The raw scores and normalized weights are:

Raw future scores exist in $S$ but are excluded before normalization. Row 2, for example, has exponentials $[2,1,2]$ and retrieves $\tfrac25[1,0]+\tfrac15[0,1]+\tfrac25[1,1]=[0.8,0.6]$.

4 · Add the residual and apply the FFN. Every entry of $U$ is nonnegative, so $\operatorname{ReLU}(U)=U$ and the second residual produces $H=1.5U$.

5 · Project to vocabulary logits and score the next token. Use an untied head $W_{vocab}=\begin{bmatrix}1&0&-1\\0&1&-1\end{bmatrix}$. Each row $[h_0,h_1]$ gives logits $[h_0,h_1,-h_0-h_1]$. Softmax now normalizes over vocabulary entries, rather than key positions.

For each row, $\ell_i=-\log p_i(\text{target})=\operatorname{logsumexp}(z_i)-z_{i,\text{target}}$. The mean loss over the three scored rows is approximately 3.02041 nats. A large loss is expected for these arbitrary parameters. During generation, row 2’s logits would predict a new token after A B C; training instead scores the observed next token A.

8Optional lab: verify attention before building an LM

Download 11_attention.py and run python 11_attention.py with PyTorch installed. This original lab has no training loop: it isolates retrieval, fused projections, and the order symmetry. It verifies the retrieval answer $[2,1.5]$, the two-head example, and §7.4’s full numerical trace, and checks row sums, shapes, projection equivalence, permutation equivariance, and future dependence.

def attention(q, k, v):
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))
    weights = scores.softmax(dim=-1)
    return weights @ v, weights

Notice that transpose(-2,-1) works for both three-dimensional batches and four-dimensional head tensors. .T does not express that operation for arbitrary batch rank.

8.1 Optional code exercise: name the tensor axes

scores = torch.einsum("...", q, k) / math.sqrt(q.size(-1))
weights = scores.softmax(dim=-1)
y = torch.einsum("...", weights, v)

scores = torch.einsum("bid,bjd->bij", q, k) / math.sqrt(q.size(-1))
y = torch.einsum("bij,bje->bie", weights, v)

9Implementation readings and at-home assembly

9.1 Identify the head axis

model.py · CausalSelfAttention.forward · verbatim excerpt

        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

C is our $D$, nh is $H$, and hs is $d_h$. The projection fuses all three families; the reshape separates heads. The output later transposes back and joins head features, as in §4.

9.2 Recognize pre-norm

model.py · Block.forward · verbatim excerpt

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

This is §5's pre-norm arrangement. Inside MLP, trace the expansion to 4 * config.n_embd, GELU, and contraction back to $D$. The implementation also includes dropout; our lab disables that complication.

Excerpts: nanoGPT, Copyright (c) 2022 Andrej Karpathy, MIT License. License and permission notice.

9.3 Optional workshop: improve the course tutorial with an equivalence check

lab_1_pretraining_a_small_language_model.py · Attn.forward · verbatim excerpt

        out = torch.cat([h(x) for h in self.heads], dim=-1)

The tutorial evaluates separately parameterized heads and concatenates their outputs. Its blocks use pre-norm residuals, RMSNorm, and a gated FFN. These last two choices are optional architectural variations: the core reasoning about shapes and information flow from §4–§5 still applies. For a gated FFN, two $D\to D_{ff}$ maps produce features and gates; elementwise multiplication combines them before a $D_{ff}\to D$ map. RMSNorm rescales features by their root-mean-square rather than subtracting their mean first.

9.4 Assigned at-home reading: turn the implementation slides into a build plan

Use slides 112–125, “Writing our own Transformer”, as assigned reading to complete at home. Use the worksheet below to trace how the components fit together. They build an encoder–decoder; the tutorial builds a decoder-only LM. First assemble the shared modules, then use #12's architecture section to reason about the additional cross-attention path.

For the causal-mask implementation and cached decoding, continue with #12's assigned at-home implementation reading. Slides 112–125 are required reading to complete at home; §9.4 and #12 §7 provide guidance for that reading. The additional code experiments remain optional.

10Practice set: put the pieces together

Try these without opening the answers. Use zero-based positions throughout.