Coding Attention Mechanisms

Key concept · From Bahdanau attention to self-attention

The first fix was Bahdanau attention (2014): let the decoder selectively access different parts of the input at each step, weighting the important tokens. Three years later, researchers found RNNs weren't needed at all and introduced the transformer's self-attention. The "self" means it relates different positions within a single sequence to each other (word-to-word in one sentence), rather than between two separate sequences. Self-attention lets each position attend to all positions in the same sequence when building its representation.

Key concept · Context vectors — the goal

The output of self-attention is a context vector z for each input token: an enriched embedding that blends in information from all the other tokens in the sequence, weighted by relevance. Where a plain token embedding knows only about itself, a context vector for "journey" also carries information about "Your", "starts", and every other word — exactly the relational understanding an LLM needs to predict the next token.

Key concept · Scores, weights, and the weighted sum

(1) Attention scores: for a chosen query token, compute the dot product of its embedding with every input token's embedding. The dot product is a similarity measure — higher means the two vectors are more aligned, so the tokens attend to each other more. (2) Attention weights: normalize the scores with softmax so they're positive and sum to 1 (softmax handles extreme values better than naive division and gives good gradients). (3) Context vector: take the weighted sum of all input embeddings using those weights. Repeat for every token — efficiently, this is just matrix multiplication: scores = inputs @ inputs.T, then softmax per row, then weights @ inputs.

Key concept · Query, key, and value

Three learnable matrices — W_q, W_k, W_v — project each input embedding into three roles. The query (q) is the current token "asking" for relevant information; the key (k) is what each token "offers" to be matched against; the value (v) is the actual content that gets aggregated. Attention scores are now the dot products of the query with every key; the resulting weights are applied to the values to produce the context vector. (Note: "weight" in W_q/W_k/W_v means trainable network parameters — distinct from the dynamic, context-specific attention weights.)

Key concept · The scaling, and the compact class

Before softmax, the query·key scores are divided by the square root of the key dimension (√d_k) — hence "scaled." Without scaling, large dot products push softmax into regions with tiny gradients, slowing learning; scaling keeps the distribution well-behaved. In code this becomes a tidy SelfAttention class, implemented with nn.Linear layers (which give better weight initialization than raw nn.Parameter) — a reusable module that computes context vectors for a whole batch at once.

Key concept · The causal mask and dropout

Causal attention restricts each position to attend only to itself and earlier positions. Implementation: after computing the attention scores, mask out the future — set the upper-triangular entries (positions after the current one) to −∞ before softmax, so they receive zero weight; the remaining weights are renormalized to sum to 1. On top of the mask, the book applies dropout to the attention weights — randomly zeroing some of them during training — to reduce overfitting and avoid the model over-relying on any single attention connection.

Key concept · Heads, subspaces, and concatenation

A single attention "head" captures one kind of relationship. Multi-head attention stacks several heads, each with its own W_q/W_k/W_v, so the model can attend to information from different representation subspaces at once — one head might track syntax, another long-range references. Each head produces its own set of context vectors; these are concatenated to form the output. Conceptually you could stack separate attention modules, but the efficient implementation uses a single set of larger matrices and splits them into heads by reshaping — one matrix multiply instead of many. This multi-head causal-attention module is exactly what gets plugged into the GPT architecture in Chapter 4.