Chapter 3: Coding attention mechanisms
The introduction connects prepared text embeddings to attention as an integral part of the LLM architecture.
After this chapter you will be able to
- Embeddings and architectural attention. Explain that preparing text as token embeddings supplies representations for the LLM, while attention mechanisms are an integral part of the LLM architecture itself.
- Encoder–Decoder RNN Information Flow. Trace how an encoder–decoder RNN turns an input sentence into a translated sentence.
- Earlier Context Access Limitation. Explain why the described encoder–decoder RNN can lose context when translating complex sentences.
- RNN bottleneck and decoder attention. Explain why long inputs are difficult for an encoder–decoder RNN and how Bahdanau attention changes what the decoder can access.
- Transformer self-attention role. Describe what self-attention lets positions in a sequence do and how it relates to transformer-based LLMs.
- Within-sequence attention relationships. Explain that self-attention computes attention weights by relating positions within a single input sequence, and distinguish this from attention between two sequences.
- Purpose of context vectors. Explain what a context vector represents and why self-attention creates one for each input element.
- Dot-product attention scores. Describe how attention scores are calculated for the example query x(2) and what a higher score indicates.
- Normalized attention weights. Explain why attention scores are normalized into weights and why the section recommends PyTorch softmax in practice.
- Weighted context vector calculation. Trace how the normalized weights and input embeddings produce the example context vector z(2).
- All-input attention scores. Explain that the attention-score tensor contains a score for each pair of inputs, and that matrix multiplication produces the same scores as Python loops more efficiently.
- Row-wise attention normalization. Describe how the attention scores become attention weights by normalizing each row, and explain what dim=-1 does for a two-dimensional score tensor.
- All-input context vectors. Describe how the attention weights are used to compute all context vectors by matrix multiplication, and how the second output row checks the result against the earlier calculation.
- Trainable self-attention context vectors. Explain that self-attention still computes a context vector for a particular input element as a weighted sum of input vectors, but adds weight matrices that are updated during training so the attention module can learn to produce good context vectors.
- Query key value projections. Explain how three trainable weight matrices project input token embeddings into query, key, and value vectors, and why computing z(2) still requires keys and values for every input token.
- Scaled scores and attention weights. Trace how attention scores are turned into attention weights by dividing by the square root of the key embedding dimension and applying softmax.
- Value weighted context vector. Describe how the attention weights combine value vectors to produce the context vector, rather than directly combining the original input vectors.
- Self-attention class computation. Trace how the self-attention class produces context vectors: it initializes trainable query, key, and value matrices, then computes scores, normalizes them with softmax, and weights the values.
- Linear projection implementation. Explain why SelfAttention_v2 uses bias-disabled nn.Linear layers instead of manually initialized weight parameters, and why the two versions initially produce different outputs.
- Causal attention access restriction. Explain why causal attention limits each token to previous and current tokens, and contrast this with standard self-attention when predicting the next token.
- Future-token masking and normalization. Describe how masking attention weights above the diagonal and normalizing the remaining weights enforce causal attention.
- Masking and renormalizing weights. Describe how to compute attention weights with softmax, use a lower-triangular mask to zero weights above the diagonal, and renormalize each row to sum to 1.
- Pre-softmax negative-infinity masking. Describe the shorter masking method: put negative infinity at future positions in the attention scores before softmax, so those positions receive zero weight and each row already sums to 1.
- Dropout purpose and timing. Explain that dropout randomly ignores selected units during training to reduce overfitting, and that it is disabled afterward.
- Attention-weight dropout and scaling. Describe where this section applies dropout in attention and what happens to the weights: some are randomly zeroed, while the remaining weights are scaled up to compensate.
- Causal attention class additions. Describe how the CausalAttention class differs from the earlier SelfAttention class: it adds a causal mask and dropout.
- Registered buffer device handling. Explain why the class calls register_buffer: the buffer moves to the appropriate CPU or GPU along with the model, helping avoid device mismatch errors.
- Single-head and multi-head attention. Distinguish a single causal attention module from multi-head attention by explaining what an attention head is in this section.
- Independently weighted attention heads. Explain that multi-head attention uses multiple self-attention heads with different learned projections and combines their outputs.
- Concatenated context vector dimensions. Describe how the number of heads and each head’s d_out determine the wrapper’s context-vector size, and interpret the example output dimensions.
- Integrated attention head layout. Trace how MultiHeadAttention turns projected queries, keys, and values into multiple heads. Describe the change from (b, num_tokens, d_out) to (b, num_heads, num_tokens, head_dim), where head_dim = d_out / num_heads.
- Batched attention and head recombination. Explain how the reordered head dimension lets PyTorch perform matrix multiplication separately for each head, and describe how the resulting context vectors are combined.
- Shared projection efficiency. Justify why the integrated MultiHeadAttention class is more efficient than the wrapper of separate CausalAttention objects.
What it builds on
- Creating token embeddings (chapter 2) · Creating token embeddings as vector representations for the LLM.
- Introducing the transformer architecture (chapter 1) · The introduction to the transformer architecture
- Data sampling with a sliding window (chapter 2) · Preparing text inputs for batches.
Section by section
Introduction
The introduction connects prepared text embeddings to attention as an integral part of the LLM architecture.
- Embeddings and architectural attention: Explain that preparing text as token embeddings supplies representations for the LLM, while attention mechanisms are an integral part of the LLM architecture itself.
Why it matters: This distinction places attention within the model that works with the prepared text representations.
Key terms: tokens, embeddings, attention mechanisms, LLM architecture
3.1 The problem with modeling long sequences
This section uses machine translation to explain how encoder–decoder RNNs process text and why their limited access to earlier input context motivated attention mechanisms.
- Translation Beyond Individual Words: Explain why translating a sentence requires more than translating each word separately.
Why it matters: It explains why the translation model must process more than isolated words. - Encoder–Decoder RNN Information Flow: Trace how an encoder–decoder RNN turns an input sentence into a translated sentence.
Why it matters: This information flow makes the model’s context limitation understandable. - Earlier Context Access Limitation: Explain why the described encoder–decoder RNN can lose context when translating complex sentences.
Why it matters: This is the shortcoming that motivates the chapter’s introduction of attention.
Key terms: source language, target language, grammatical structures, encoder, decoder, RNN, hidden state, final hidden state, earlier hidden states, decoding phase
3.2 Capturing data dependencies with attention mechanisms
The section explains how attention addresses an RNN encoder–decoder limitation and introduces self-attention as a key mechanism in transformer-based LLMs.
- RNN bottleneck and decoder attention: Explain why long inputs are difficult for an encoder–decoder RNN and how Bahdanau attention changes what the decoder can access.
Why it matters: This explains the problem that attention was introduced to address. - Transformer self-attention role: Describe what self-attention lets positions in a sequence do and how it relates to transformer-based LLMs.
Why it matters: It identifies the kind of attention used in the transformer-based models built later in the book.
Key terms: RNN, encoder–decoder RNN, single hidden state, Bahdanau attention mechanism, decoder, self-attention, input sequence, transformer architecture, GPT series
3.3 Attending to different parts of the input with self-attention
This section defines the “self” in self-attention by contrasting relationships within one input sequence with attention between two sequences.
- Within-sequence attention relationships: Explain that self-attention computes attention weights by relating positions within a single input sequence, and distinguish this from attention between two sequences.
Why it matters: This distinction clarifies what the attention mechanism in transformer-based LLMs is relating.
Key terms: self-attention, attention weights, input sequence, traditional attention mechanisms
3.3.1 A simple self-attention mechanism without trainable weights
This section builds a context vector for one token using a simplified self-attention procedure without trainable weights.
- Purpose of context vectors: Explain what a context vector represents and why self-attention creates one for each input element.
Why it matters: It explains what the attention calculation is trying to produce. - Dot-product attention scores: Describe how attention scores are calculated for the example query x(2) and what a higher score indicates.
Why it matters: Scores are the first step in determining how strongly the query attends to each input. - Normalized attention weights: Explain why attention scores are normalized into weights and why the section recommends PyTorch softmax in practice.
Why it matters: Normalization turns raw scores into interpretable weights for combining input vectors. - Weighted context vector calculation: Trace how the normalized weights and input embeddings produce the example context vector z(2).
Why it matters: This final calculation turns attention weights into the enriched token representation.
Key terms: context vector, enriched embedding vector, input sequence, query, attention scores, dot product, normalization, attention weights, softmax, softmax_naive
3.3.2 Computing attention weights for all input tokens
This subsection extends simple self-attention from one input to all inputs by computing pairwise scores, normalizing them into attention weights, and producing a context vector for each input.
- All-input attention scores: Explain that the attention-score tensor contains a score for each pair of inputs, and that matrix multiplication produces the same scores as Python loops more efficiently.
Why it matters: Computing scores for every input lets self-attention produce more than just the second input’s context vector. - Row-wise attention normalization: Describe how the attention scores become attention weights by normalizing each row, and explain what dim=-1 does for a two-dimensional score tensor.
Why it matters: The row-wise normalization turns raw pairwise scores into the weights used to compute context vectors. - All-input context vectors: Describe how the attention weights are used to compute all context vectors by matrix multiplication, and how the second output row checks the result against the earlier calculation.
Why it matters: This completes the extension from one context vector to a context vector for every input.
Key terms: attention scores, tensor, matrix multiplication, attention weights, torch.softmax, dim=-1, attn_scores, context vectors, context_vec_2
3.4 Implementing self-attention with trainable weights
The section introduces trainable self-attention as an extension of basic self-attention that learns to produce better context vectors.
- Trainable self-attention context vectors: Explain that self-attention still computes a context vector for a particular input element as a weighted sum of input vectors, but adds weight matrices that are updated during training so the attention module can learn to produce good context vectors.
Why it matters: This distinction connects the basic attention mechanism to the trainable version used in LLMs.
Key terms: self-attention, context vectors, weighted sums, trainable weight matrices, scaled dot-product attention
3.4.1 Computing the attention weights step by step
This section computes one context vector using learned query, key, and value projections, scaled attention scores, and a weighted sum of values.
- Query key value projections: Explain how three trainable weight matrices project input token embeddings into query, key, and value vectors, and why computing z(2) still requires keys and values for every input token.
Why it matters: These three projections give each token distinct roles in deciding what information contributes to its context. - Parameter and attention weight distinction: Distinguish the learned weight parameters in the projection matrices from the attention weights calculated for an input.
Why it matters: Both are called weights, but they play different roles in the computation. - Scaled scores and attention weights: Trace how attention scores are turned into attention weights by dividing by the square root of the key embedding dimension and applying softmax.
Why it matters: This is the step that converts matching scores into the weights used to combine information. - Reason for score scaling: Justify scaling attention scores before softmax in terms of large dot products and training gradients.
Why it matters: The scaling has a training purpose, not just a numerical convention. - Value weighted context vector: Describe how the attention weights combine value vectors to produce the context vector, rather than directly combining the original input vectors.
Why it matters: This final combination turns the computed attention weights into an output representation.
Key terms: Wq, Wk, Wv, query, key, value, weight parameters, attention weights, attention scores, softmax
3.4.2 Implementing a compact self-attention Python class
This section packages self-attention in a PyTorch class and compares two implementations of its trainable projections.
- Self-attention class computation: Trace how the self-attention class produces context vectors: it initializes trainable query, key, and value matrices, then computes scores, normalizes them with softmax, and weights the values.
Why it matters: It connects the step-by-step attention calculation to a reusable model layer. - Linear projection implementation: Explain why SelfAttention_v2 uses bias-disabled nn.Linear layers instead of manually initialized weight parameters, and why the two versions initially produce different outputs.
Why it matters: Different default outputs do not mean the two classes implement different attention calculations. - Equivalent weight transfer: Describe how to check that the two self-attention implementations are otherwise equivalent by transferring weights from SelfAttention_v2 to SelfAttention_v1, accounting for nn.Linear’s transposed weight storage.
Why it matters: Matching the weights separates an implementation difference from a mere initialization difference.
Key terms: SelfAttention_v1, nn.Module, W_query, W_key, W_value, attn_scores, SelfAttention_v2, nn.Linear, nn.Parameter, weight initialization
3.5 Hiding future words with causal attention
Introduces causal attention as self-attention that prevents a token from attending to future tokens, and describes the masking and normalization that enforce this restriction.
- Causal attention access restriction: Explain why causal attention limits each token to previous and current tokens, and contrast this with standard self-attention when predicting the next token.
Why it matters: This restriction keeps future words out of the context used to predict the next word. - Future-token masking and normalization: Describe how masking attention weights above the diagonal and normalizing the remaining weights enforce causal attention.
Why it matters: Masking and normalization turn the access restriction into attention weights that exclude future tokens.
Key terms: causal attention, masked attention, standard self-attention, attention weights, mask out, above the diagonal, normalize
3.5.1 Applying a causal attention mask
This section shows two ways to mask future-token attention weights: zeroing and renormalizing weights, or masking scores before softmax.
- Masking and renormalizing weights: Describe how to compute attention weights with softmax, use a lower-triangular mask to zero weights above the diagonal, and renormalize each row to sum to 1.
Why it matters: This gives a direct way to make each token’s attention weights exclude future positions. - Absence of future-token leakage: Explain why applying softmax before masking does not leave future-token influence in the final, renormalized attention weights.
Why it matters: It resolves the concern that future tokens might still affect the result because they appeared in the original softmax calculation. - Pre-softmax negative-infinity masking: Describe the shorter masking method: put negative infinity at future positions in the attention scores before softmax, so those positions receive zero weight and each row already sums to 1.
Why it matters: Masking before softmax obtains normalized causal attention weights without a separate renormalization step.
Key terms: causal attention mask, softmax, tril, renormalize, information leakage, masked positions, renormalization, negative infinity, attention scores
3.5.2 Masking additional attention weights with dropout
The section explains why dropout is used in attention and how it randomly zeros and rescales attention weights during training.
- Dropout purpose and timing: Explain that dropout randomly ignores selected units during training to reduce overfitting, and that it is disabled afterward.
Why it matters: This explains both the reason for adding dropout and when it should be active. - Attention-weight dropout and scaling: Describe where this section applies dropout in attention and what happens to the weights: some are randomly zeroed, while the remaining weights are scaled up to compensate.
Why it matters: Zeroing weights alone would reduce the attention mechanism’s average influence, so the surviving weights are scaled.
Key terms: dropout, overfitting, training, attention weights, dropout mask, dropout rate, value vectors
3.5.3 Implementing a compact causal attention class
This section presents a batched causal-attention class that adds masking and dropout to the earlier self-attention class and registers a buffer for device compatibility.
- Causal attention class additions: Describe how the CausalAttention class differs from the earlier SelfAttention class: it adds a causal mask and dropout.
Why it matters: This identifies what the new class contributes beyond the self-attention class. - Batched input and output shapes: Describe the example batch and the resulting context-vector shape: two texts of six tokens with three-dimensional input embeddings produce a three-dimensional output tensor with two-dimensional embeddings per token.
Why it matters: The example shows that the class can handle the batched inputs produced by the data loader. - Registered buffer device handling: Explain why the class calls register_buffer: the buffer moves to the appropriate CPU or GPU along with the model, helping avoid device mismatch errors.
Why it matters: Automatic device placement helps the class work when the model is moved between CPU and GPU.
Key terms: CausalAttention, SelfAttention, dropout, causal mask, batch inputs, three-dimensional tensor, context vector, embedding, register_buffer, buffer
3.6 Extending single-head attention to multi-head attention
This section defines multi-head attention as an extension of causal attention with multiple independently operating heads.
- Single-head and multi-head attention: Distinguish a single causal attention module from multi-head attention by explaining what an attention head is in this section.
Why it matters: This distinction gives you the basic vocabulary for the multi-head implementations that follow.
Key terms: multi-head attention, heads, single-head attention, causal attention module, attention weights
3.6.1 Stacking multiple single-head attention layers
This section builds multi-head attention by stacking independently weighted CausalAttention heads and explains the wrapper’s output shape and sequential execution.
- Independently weighted attention heads: Explain that multi-head attention uses multiple self-attention heads with different learned projections and combines their outputs.
Why it matters: This is the central idea behind using more than one attention head. - Concatenated context vector dimensions: Describe how the number of heads and each head’s d_out determine the wrapper’s context-vector size, and interpret the example output dimensions.
Why it matters: It lets you predict and adjust the shape of the wrapper’s output. - Wrapper execution versus parallel heads: Distinguish the described idea of running heads in parallel from how this MultiHeadAttentionWrapper actually processes them.
Why it matters: It prevents confusing the multi-head idea with the execution behavior of this first wrapper.
Key terms: multi-head attention, single-head attention, learned linear projections, weights, num_heads, d_out, context_vecs, context vector, MultiHeadAttentionWrapper, forward method
3.6.2 Implementing multi-head attention with weight splits
This section implements multi-head attention in one class by splitting projected tensors into heads, computing attention across those heads, and recombining their outputs.
- Integrated attention head layout: Trace how MultiHeadAttention turns projected queries, keys, and values into multiple heads. Describe the change from (b, num_tokens, d_out) to (b, num_heads, num_tokens, head_dim), where head_dim = d_out / num_heads.
Why it matters: The head dimension explains how one class represents several attention heads without stacking separate modules. - Batched attention and head recombination: Explain how the reordered head dimension lets PyTorch perform matrix multiplication separately for each head, and describe how the resulting context vectors are combined.
Why it matters: It connects the tensor layout to both per-head computation and the final combined output. - Shared projection efficiency: Justify why the integrated MultiHeadAttention class is more efficient than the wrapper of separate CausalAttention objects.
Why it matters: The tensor reshaping is worthwhile because it avoids repeating expensive input projections for every head. - Combined output projection: Describe the output projection added after the head outputs are combined, distinguishing it from a required step for combining heads.
Why it matters: It clarifies which part of the class is a common architectural addition rather than the operation that merges heads.
Key terms: MultiHeadAttention, num_heads, head_dim, .view, .transpose, batched matrix multiplication, context vectors, d_out, MultiHeadAttentionWrapper, queries
3.7 Summary
This section summarizes how self-attention builds context vectors and how causal masking and multi-head attention adapt it for transformer-based LLMs.
Where this is used later
- Connecting attention and linear layers in a transformer block (chapter 4) · Connecting attention to other layers in a transformer block relies on recognizing it as part of the model architecture.
- Training an LLM (chapter 5) · Model training updates the trainable weights introduced here.
- Top-k sampling (chapter 5) · Masking logits with negative infinity before softmax excludes positions by giving them zero probability.
- Adding a classification head (chapter 6) · The causal attention mask that prevents a token from accessing later positions.
- Coding the GPT model (chapter 4) · The attention class is used as part of an LLM model.
- Coding an LLM architecture (chapter 4) · Multi-head attention uses multiple attention heads.
- (chapter intro) (chapter 5) · The multi-head attention implementation uses linear layers.
Concepts explained in this chapter
- attentionA neural-network mechanism that weights input tokens or values by relevance so the model can focus on important context.
- context vectorA vector representation formed by combining input information with attention weights to summarize relevant context for a token.
- attention scoreThe numeric values computed to quantify how strongly one token should attend to another in attention mechanisms.
- attention weightA normalized relevance coefficient derived from attention scores that determines how much each value contributes to the context vector.
- parameterAn adjustable value in a neural network that is learned during training.
- multi-head attentionAn attention technique that computes several independent attention heads in parallel and combines their outputs to capture different relationships in the same sequence.
- keyThe attention-side vector associated with an input token and compared with a query to determine relevance.
- valueThe attention content vector combined according to attention weights to form the context output.
- CausalAttentionA causal single-head attention class that computes attention while preserving autoregressive token order.
- MultiHeadAttentionA single attention module that computes multiple attention heads by reshaping projected queries, keys, and values within one implementation.
- MultiHeadAttentionWrapperA wrapper class that implements multi-head attention by combining several separate single-head attention modules.
- token idAn integer identifier assigned to a token by a vocabulary for use in model input and output processing.
← Chapter 2: Working with text data
Chapter 4: Implementing a GPT model from scratch to generate text →
These guides are generated from the text of Manning books by Manning's knowledge graph and link to the sections they draw on. The full books are available on liveBook and at manning.com.