Week 7 — Attention: read the relevant preceding context

We keep Week 6's adapted words “you,” “we,” and “like,” their IDs, and the meanings of [B,T,C]. To calculate each Q/K/V entry by hand, however, we use this chapter's simple floating-point X and fixed projections—not embeddings copied from a particular Week 6 training step. Week 8 supplies a separate token-plus-position illustration; Weeks 10–12 use one shared trainable Mini GPT implementation.

Use two prompts throughout: A=[you,like] and B=[we,like]. Their final token is identical, but a contextual model should be able to let that final “like” read different earlier information. Attention provides this input-dependent information path.

Attention does not directly receive token IDs in this calculation. IDs are integer lookup indices; X contains embedding vectors or a previous layer's outputs. Our simplified values make each matrix multiplication manageable by hand. Actual model representations are usually learned dense floating-point vectors.

A simple baseline uniformly averages the visible prefix. At the final position t=1, the fixed weights are [0.5,0.5]. Different inputs can give different averages for A and B, but these reading weights do not change with the current query.

Think of a query as a search request, a key as searchable index features, and a value as record content. This analogy explains the roles only. The model is not searching literal strings, and no person assigns a coordinate the meaning “subject.”

A position's key supplies features for matching, its value supplies content to combine, and its query determines its own matching request. All three come from the same x_t but use different parameters.

These parameters are deliberately chosen for readable arithmetic. Real models normally begin with initialized parameters and learn their values through training. Each W here is still shared across positions and prompts; we do not write a separate projection rule for every token.

For example, x_like=[0,1,0,0], so multiplication by each W selects its second row: q_like=[1,1], k_like=[0.2√2,0.2√2]≈[0.283,0.283], and v_like=[0,1].

The final “like” query is [1,1]. Its dot product with the “you” key is approximately 1×1.131+1×0=1.131. With its own “like” key, it is approximately 1×0.283+1×0.283=0.566.

For Prompt B, its dot product with the “we” key is approximately 1×0+1×0.141=0.141. The score for the “like” key remains approximately 0.566.

Divide scores by √d_k. Here d_k=2. Using the unrounded scores, Prompt A's final row becomes [0.8,0.4] and Prompt B's becomes [0.1,0.4]. The displayed raw scores [1.131,0.566] and [0.141,0.566] are rounded. The final position can read both columns, so this row has no future position to mask.

Both prompts end in “like,” yet their head outputs differ. The first-position keys produce different weights, and the “you” and “we” values also differ. Both differences affect the result.

During training, O affects later representations, the language-model head, logits, and cross-entropy. Gradients flow back from loss through O, with one path through P, Softmax, and scores to W_Q/W_K, and another through values to W_V. The model does not begin with a correct attention table; parameter updates can improve its matching and information flow while reducing prediction loss.

Both paths contribute to the gradient of the shared input X, just as branches did in Week 4. V does not participate in the QKᵀ score calculation, so its direct gradient path from O does not pass through S. Identify the branches before tackling full matrix derivatives.

A dot product sums d_k products. Under a simplified model of independent zero-mean, unit-variance query/key components, each product has variance 1 and the sum has variance d_k. Its standard deviation is √d_k. Dividing by √d_k gives a typical score scale that does not grow with head width under these assumptions.

Large differences within a score row can make Softmax nearly one-hot, making some attention weights locally insensitive to score changes. Adding the same large constant to every score does not change Softmax. These local sensitivities also do not by themselves determine every final parameter gradient. Positive scaling preserves the row's ordering while changing its spread.

In each [T,T]=[2,2] matrix, row t is the query position requesting information. Column j is the key position being compared and the corresponding value position. Neither axis is a feature axis.

This separate axis-check example does not replace the earlier two-token arithmetic. Unequal axis sizes make mistakes such as exchanging head and time axes easier to detect.

For a separate next-token illustration, use input [you,like] and target [like,AI]. Position t=0 must predict “like” using only “you.” Without a mask, it could read “like” directly at input position 1—its own answer. During generation, that future token has not yet been produced, so the shortcut is unavailable.

Position t=3 can read positions 0,1,2,3—not only its immediate predecessor. The mask defines visibility; it does not determine the relative reading weights among visible positions.

This code uses x @ w_q, so W follows the mathematical [C,d_head] convention. nn.Linear(C,d_head) stores its weight as [d_head,C] and uses the corresponding transpose in the computation. The calculations are equivalent.

register_buffer registers the mask so it moves with the model between devices without becoming a trainable parameter. [:T,:T] selects the relevant part of the maximum-length mask for the current sequence.

“Self” means queries, keys, and values come from the same input sequence representations X. It does not restrict a token to itself. GPT's causal self-attention permits the current position and all earlier positions.

Prompt A's score slice contains only “you” and “like”; B's contains only “we” and “like.” Batching computes them together without concatenating them into one text.

Each head has its own W_Q, W_K, and W_V, allowing different matching features and value representations. Heads may use different signals, but they are not guaranteed to become named grammatical specialists.

The output projection does more than satisfy a shape requirement. After concatenation, features are merely side by side. W_O learns combinations across heads and supplies the [B,T,C] interface for the residual addition.

Head 2 receives the same X, not Head 1's output. Imagine two two-feature records at each position: place them together, then learn how to combine them. Successive Transformer blocks, in contrast, pass one block's output into the next.

A and B already have different final-position head outputs. Multi-head mixing, W_O, and later Transformer computations can produce different H values, allowing the language-model head to produce different logits. Attention itself does not directly select a token.

Each head's score matrix has T² entries. Across the batch and all heads there are B×n_head×T² score entries. In this explicit implementation, weights need the same number of entries, and training requires additional intermediate values and gradients.

The causal mask forbids future reads, but this ordinary dense implementation still constructs a full T×T matrix. Optimized implementations can change materialization, memory use, and execution details; the underlying query/key axes remain important.

Week 8 continues with the interface [B,T,C]=[2,2,4]. Adding token and position embeddings preserves [2,2,4]. Multi-head attention must also return [2,2,4] after concatenation and W_O so its result can be added to the residual stream.