Chapter 6: Attention Mechanisms

Given an input matrix $X \in \mathbb{R}^{n \times d_{\text{model}}}$ (one row per token), self-attention projects it into three matrices using learned weight matrices $W^Q, W^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ and $W^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$. The query $Q$ represents what each token is searching for, the key $K$ represents what each token offers, and the value $V$ carries the actual content. The separation into three roles allows the model to learn different matching and content functions. Typically $d_k = d_v = d_{\text{model}} / h$ where $h$ is the number of attention heads.

The attention function computes dot products between all query-key pairs, forming a $n \times n$ attention matrix. Dividing by $\sqrt{d_k}$ is critical: without scaling, when $d_k$ is large, the dot products grow in magnitude, pushing the softmax into regions where its gradients are extremely small (vanishing gradient problem). The $\sqrt{d_k}$ factor keeps the variance of the dot products at approximately 1, regardless of dimension. After softmax normalization, each row of the resulting matrix gives the attention weights for one token, and multiplying by $V$ produces the output --- a weighted mixture of value vectors.

The attention matrix $A = \text{softmax}(QK^\top / \sqrt{d_k})$ is an $n \times n$ matrix where entry $A_{ij}$ is the weight that token $i$ places on token $j$. In the encoder, all positions can attend to all other positions (bidirectional). In the decoder, causal masking prevents position $i$ from attending to positions $j > i$ by setting those entries to $-\infty$ before softmax (which maps them to 0). This ensures autoregressive generation: each token can only depend on previously generated tokens. Padding masks are also applied to ignore padding tokens in variable-length batches.

Self-attention has time and memory complexity $O(n^2 d)$, where $n$ is the sequence length and $d$ is the model dimension. The $n^2$ term arises from computing all pairwise dot products in $QK^\top$. For short sequences ($n < 1000$), this is highly efficient because the entire computation can be parallelized on GPUs. For very long sequences ($n > 4000$), the quadratic cost becomes a bottleneck, motivating efficient attention variants like linear attention, sparse attention, and Flash Attention. By contrast, RNNs are $O(n \cdot d^2)$ --- linear in sequence length but sequential, preventing parallelization.

Each head $i$ computes its own attention: $\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$ where $W_i^Q, W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ and $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$. The heads are concatenated along the feature dimension to produce an $n \times (h \cdot d_v)$ matrix, then projected through $W^O \in \mathbb{R}^{(h \cdot d_v) \times d_{\text{model}}}$ to restore the original dimension. In the standard Transformer, $d_k = d_v = d_{\text{model}} / h = 64$ with $h = 8$ heads.

Each head has independent projection matrices, so they learn to extract different features from the same input. Research by Voita et al. (2019) and Clark et al. (2019) shows that heads specialize: some attend to the adjacent token, others to syntactic heads (e.g., a verb attending to its subject), and others form broad attention patterns that aggregate global context. Pruning experiments reveal that not all heads are equally important --- some can be removed with minimal performance loss, while others are critical. This suggests learned redundancy provides robustness.

The output projection $W^O$ serves two purposes: it mixes information across heads (since concatenation keeps them separate) and projects back to $d_{\text{model}}$ to maintain consistent dimensions for residual connections. Without $W^O$, the model could not combine insights from different heads --- each head's contribution would remain in its own subspace. The optional bias $b^O$ adds a learnable shift. The output of multi-head attention has the same shape as the input, enabling the residual connection $x + \text{MultiHead}(x, x, x)$.

Self-attention treats its input as a set, so position must be injected explicitly. The sinusoidal scheme gives position $pos$ a $d$-dimensional vector whose dimension pair $(2i, 2i+1)$ holds the sine and cosine of $pos \cdot \omega_i$, with frequency $\omega_i = 1 / 10000^{2i/d}$. The first pair has wavelength $2\pi$ and changes sharply from one token to the next; the last pair has wavelength near $2\pi \cdot 10000$ and changes very slowly, so together the pairs act like the hands of a clock running at many speeds. The base 10000 makes the slowest wavelengths far longer than any training sequence, so positions do not repeat. The vector is added to the token embedding, with no learned parameters.

A crucial property of sinusoidal encodings is that the encoding at position $pos + k$ can be expressed as a linear transformation of the encoding at position $pos$. Specifically, for each frequency pair $(2i, 2i+1)$, the relationship is a 2D rotation matrix: $M_k = \begin{bmatrix} \cos(k\omega_i) & \sin(k\omega_i) \\ -\sin(k\omega_i) & \cos(k\omega_i) \end{bmatrix}$ where $\omega_i = 1/10000^{2i/d}$. This means the model can learn to compute relative positions through linear operations on the encodings, without explicitly being given relative position information.

An alternative to fixed sinusoidal encodings is to learn a positional embedding matrix $P \in \mathbb{R}^{L_{\max} \times d}$ where $L_{\max}$ is the maximum sequence length. Position $pos$ uses the $pos$-th row of $P$ as its encoding, and the values are learned during training via backpropagation. BERT and GPT-2 use this approach. The advantage is flexibility --- the model can learn arbitrary positional patterns. The disadvantage is that it cannot generalize to positions beyond $L_{\max}$ without interpolation or extrapolation tricks. Empirically, learned and sinusoidal encodings perform similarly for lengths within the training range.

RoPE (Su et al., 2021) applies position-dependent rotations directly to query and key vectors rather than adding encodings to inputs. Each pair of dimensions is rotated by an angle proportional to position $m$ and frequency $\theta_i = 10000^{-2i/d}$. The key property is that the dot product $f(q_m) \cdot f(k_n)$ depends only on the relative position $m - n$, naturally encoding relative distance. RoPE is used in LLaMA, PaLM, and most modern LLMs because it combines the best of both worlds: it encodes relative position like relative encodings while being as simple to implement as absolute encodings.

Every sublayer (attention and FFN) is wrapped in a residual connection that adds the sublayer input $x$ to its output. This creates shortcut paths for gradient flow, enabling training of deep networks. Without residual connections, stacking 6+ layers leads to vanishing or exploding gradients. The residual design also means each sublayer only needs to learn the "residual" --- the difference between the desired output and the input --- which is typically easier than learning the full transformation. Pre-norm (applying LayerNorm before the sublayer) and post-norm (after) are two variants; pre-norm is more stable and used in most modern implementations.

Layer normalization computes the mean $\mu$ and variance $\sigma^2$ across the feature dimension for each token independently (unlike batch normalization, which computes across the batch). The learned scale $\gamma$ and shift $\beta$ allow the model to undo the normalization if needed. LayerNorm stabilizes training by reducing internal covariate shift --- it ensures that each sublayer receives inputs with consistent statistics, regardless of the magnitudes produced by previous layers. The $\epsilon$ (typically $10^{-5}$) prevents division by zero.

The FFN is a two-layer MLP applied independently to each token position. The first layer expands the dimension from $d_{\text{model}} = 512$ to $d_{ff} = 2048$ (4x expansion), applies ReLU, then the second layer projects back to $d_{\text{model}}$. This expansion-compression pattern creates a bottleneck that forces the model to form useful intermediate representations. The FFN is where most of the model's parameters reside and where factual knowledge is believed to be stored. Recent work replaces ReLU with GELU or SwiGLU for better performance.

The complete Transformer has an encoder stack and a decoder stack. The encoder processes the source sequence with bidirectional self-attention (every token sees all others). The decoder processes the target sequence with causal self-attention (each token sees only previous tokens) plus cross-attention where decoder queries attend to encoder key-value pairs. This cross-attention is the bridge between the two stacks: decoder tokens ask questions (queries) that are answered by the encoder representations (keys and values). During inference, the encoder runs once and the decoder generates tokens autoregressively, attending to the cached encoder output at each step.

BERT (Devlin et al., 2019) uses only the Transformer encoder with bidirectional self-attention. During pretraining, 15% of tokens are randomly masked and the model predicts them from their full surrounding context. This masked language modeling (MLM) objective forces the model to build deep bidirectional representations. BERT also uses a next-sentence prediction (NSP) task during pretraining, though later work (RoBERTa) showed NSP is unnecessary. For downstream tasks, BERT is fine-tuned by adding a task-specific head (e.g., a classification layer) on top of the encoder output. BERT-base has 12 layers, 768 dimensions, and 110M parameters.

GPT (Radford et al., 2018) uses only the Transformer decoder with causal self-attention masking. The causal language modeling (CLM) objective trains the model to predict each token given all previous tokens. This autoregressive design is naturally suited for generation: at inference time, the model generates tokens one by one, each conditioned on everything generated so far. GPT-2 (1.5B params) demonstrated strong zero-shot performance, and GPT-3 (175B params) showed that scaling decoder-only models produces emergent few-shot learning --- the model can perform tasks from just a few examples in the prompt, without any gradient updates.

T5 (Raffel et al., 2020) uses the full encoder-decoder architecture and frames every NLP task as a text-to-text problem. Classification becomes "classify: The movie was great" $\rightarrow$ "positive". Translation, summarization, and question answering all follow the same format with different text prefixes. T5 is pretrained with a span corruption objective: random spans of input tokens are replaced with sentinel tokens, and the decoder predicts the original spans. This is more efficient than BERT's MLM because it masks contiguous spans (15% of tokens in fewer mask operations) and trains the decoder to generate the missing text. T5 demonstrated that a single architecture with a unified text interface can match or exceed specialized models across dozens of benchmarks.

The $O(n^2)$ attention complexity becomes a bottleneck for long sequences. Several approaches address this. Sparse attention (Longformer, BigBird) restricts each token to attend to a fixed window plus a few global tokens, reducing complexity to $O(n)$. Linear attention (Performers, Katharopoulos et al.) approximates softmax attention with kernel functions, enabling $O(n \cdot d)$ computation. Flash Attention (Dao et al., 2022) does not change the algorithm but optimizes GPU memory access by tiling the attention computation to maximize use of fast SRAM, achieving 2-4x speedup with exact attention. Mixture of Experts (MoE) routes each token to a subset of FFN experts, increasing model capacity without proportionally increasing computation. These efficiency techniques are essential for training and serving modern LLMs.