Self-Attention Mechanism

Table of Contents

1. Learning Objectives
2. Prerequisites
3. Notation
4. Core Intuition
5. The Sequence Modeling Problem
6. Content-Based Addressing
7. Query, Key, Value Projections
8. Attention Scores and Dot Products
9. The Scaling Factor: Variance Analysis
10. Softmax Normalization
11. Scaled Dot-Product Attention
12. Multi-Head Attention
13. Causal Masking for Autoregressive Models
14. Computational Complexity and Memory
15. Gradients Through Attention
16. Worked Examples
17. Connection to the Broader Curriculum
18. Common Pitfalls and Misconceptions
19. Research Perspective
20. Summary of Takeaways
21. Exercises

---

Learning Objectives

After reading this chapter, you should be able to:

1. Motivate self-attention as content-based retrieval over a sequence.
2. Derive the Q/K/V linear projections and the attention score matrix $\mathbf{S} = \mathbf{Q}\mathbf{K}^T$.
3. Prove that $\text{Var}(\mathbf{q}^T\mathbf{k}) = d_k$ under standard initialization and justify scaling by $\sqrt{d_k}$.
4. Write the complete scaled dot-product attention formula and multi-head attention.
5. Apply causal masking for autoregressive generation.
6. Analyze time and memory complexity ($O(n^2 d)$ time, $O(n^2)$ memory).
7. Explain why multi-head attention adds representational capacity without increasing FLOPs.

---

Prerequisites

• Vectors, Spans & Linear Independence — dot products, cosine similarity
• Backpropagation — chain rule, softmax gradients
• Logistic Regression — softmax as multiclass generalization

---

Notation

• $\mathbf{X} \in \mathbb{R}^{n \times d}$ — Input sequence embeddings
• $\mathbf{Q}, \mathbf{K}, \mathbf{V}$ — Query, key, value matrices
• $d_k, d_v$ — Head dimensions for keys and values
• $\mathbf{S} = \mathbf{Q}\mathbf{K}^T$ — Attention score matrix
• $A_{ij}$ — Softmax attention weight
• $h$ — Number of attention heads

---

Core Intuition

A sequence of $n$ tokens must be processed so that each position can incorporate context from all other positions. Recurrent networks pass information through $O(n)$ sequential steps; self-attention provides direct paths between any pair of positions in $O(1)$ computational depth.

The mechanism is elegant: each token asks a query ("what am I looking for?"), advertises a key ("what do I contain?"), and offers a value ("what information do I provide?"). Attention weights measure query–key compatibility; the output is a weighted sum of values.

This is the foundation of Transformers — the architecture behind modern LLMs, vision transformers, and diffusion models. Positional Encoding adds order information; Flash Attention makes long sequences tractable.

Series context. Chapter 8 (Attention & Transformers), Part I in Volume II. Chapter 11 in the full curriculum numbering.

---

The Sequence Modeling Problem

Definition 1 (Sequence Representation). Given token embeddings $\mathbf{X} = [\mathbf{x}_1, \ldots, \mathbf{x}_n]^T \in \mathbb{R}^{n \times d}$, compute output representations $\mathbf{O} \in \mathbb{R}^{n \times d}$ where row $i$ aggregates information from the full sequence relevant to token $i$.

Definition 2 (All-Pairs Interaction). Self-attention computes an $n \times n$ matrix of pairwise interactions, allowing token $i$ to directly access token $j$ for any $j \in \{1, \ldots, n\}$.

Proposition 1 (RNN Path Length). In an RNN, information from position $j$ reaches position $i$ through $|i - j|$ sequential applications of the recurrence. Gradients through this path suffer from vanishing/exploding dynamics (see Backpropagation).

Proposition 2 (Attention Path Length). Self-attention connects any pair $(i, j)$ in a single layer — path length 1.

---

Content-Based Addressing

Definition 3 (Content-Based Retrieval). A database consists of keys $\{\mathbf{k}_j\}$ and values $\{\mathbf{v}_j\}$. A query $\mathbf{q}$ retrieves:

$$\text{output} = \sum_{j=1}^{n} w_j \mathbf{v}_j, \quad w_j = \frac{\exp(\mathbf{q}^T \mathbf{k}_j)}{\sum_l \exp(\mathbf{q}^T \mathbf{k}_l)}. \tag{1}$$

Interpretation. High $w_j$ when $\mathbf{q}$ is compatible with $\mathbf{k}_j$ (large dot product). Softmax ensures $w_j \geq 0$ and $\sum_j w_j = 1$ — a convex combination of values.

Self-attention applies this in parallel for all $n$ queries, with keys, values, and queries all derived from the same input sequence ("self").

---

Query, Key, Value Projections

Definition 4 (Linear Projections). From input $\mathbf{X} \in \mathbb{R}^{n \times d}$:

$$\mathbf{Q} = \mathbf{X}\mathbf{W}_Q, \quad \mathbf{K} = \mathbf{X}\mathbf{W}_K, \quad \mathbf{V} = \mathbf{X}\mathbf{W}_V, \tag{2}$$

where $\mathbf{W}_Q, \mathbf{W}_K \in \mathbb{R}^{d \times d_k}$ and $\mathbf{W}_V \in \mathbb{R}^{d \times d_v}$ are learnable.

Dimensions:

• $\mathbf{Q} \in \mathbb{R}^{n \times d_k}$ — one query per token
• $\mathbf{K} \in \mathbb{R}^{n \times d_k}$ — one key per token
• $\mathbf{V} \in \mathbb{R}^{n \times d_v}$ — one value per token

Proposition 3 (Role Separation). Separate projections allow the same embedding $\mathbf{x}_i$ to simultaneously serve as query (what to seek), key (what is offered), and value (information content) in different subspaces.

This connects to Matrix Operations: each projection is a linear map to a task-specific subspace.

---

Attention Scores and Dot Products

Definition 5 (Attention Score Matrix).

$$\mathbf{S} = \mathbf{Q}\mathbf{K}^T \in \mathbb{R}^{n \times n}, \quad S_{ij} = \mathbf{q}_i^T \mathbf{k}_j. \tag{3}$$

Interpretation. $S_{ij}$ measures compatibility between query $i$ and key $j$. When $\|\mathbf{q}_i\| = \|\mathbf{k}_j\| = 1$, this is proportional to cosine similarity (see Vectors & Linear Independence).

Proposition 4 (Bilinear Form). $\mathbf{S} = \mathbf{X}\mathbf{W}_Q\mathbf{W}_K^T\mathbf{X}^T$ is a learned similarity kernel over input pairs.

---

The Scaling Factor: Variance Analysis

Theorem 1 (Score Variance). Assume entries of $\mathbf{q}_i, \mathbf{k}_j \in \mathbb{R}^{d_k}$ are independent with $\mathbb{E}[q_l] = \mathbb{E}[k_l] = 0$, $\text{Var}(q_l) = \text{Var}(k_l) = 1$. Then:

$$\mathbb{E}[S_{ij}] = 0, \qquad \text{Var}(S_{ij}) = d_k. \tag{4}$$

Proof.

$$\mathbb{E}[S_{ij}] = \sum_{l=1}^{d_k} \mathbb{E}[q_{il} k_{jl}] = \sum_l \mathbb{E}[q_{il}]\mathbb{E}[k_{jl}] = 0. \tag{5}$$ $$\text{Var}(S_{ij}) = \sum_{l=1}^{d_k} \text{Var}(q_{il} k_{jl}) = \sum_l \text{Var}(q_{il})\text{Var}(k_{jl}) = d_k, \tag{6}$$

using independence. $\blacksquare$

Corollary 1 (Softmax Saturation). When $d_k$ is large, $|S_{ij}|$ has standard deviation $\sqrt{d_k}$. Large score magnitudes drive softmax into saturated regions where $\partial \text{softmax} / \partial S_{ij} \approx 0$ (vanishing gradients).

Definition 6 (Scaled Scores).

$$\tilde{S}_{ij} = \frac{S_{ij}}{\sqrt{d_k}}, \qquad \text{Var}(\tilde{S}_{ij}) = 1. \tag{7}$$

Important equation. Scaling by $\sqrt{d_k}$ keeps softmax inputs in a well-conditioned regime regardless of head dimension. This is essential for training stability.

---

Softmax Normalization

Definition 7 (Attention Weights).

$$A_{ij} = \text{softmax}(\tilde{\mathbf{S}}_i)_j = \frac{\exp(\tilde{S}_{ij})}{\sum_{l=1}^{n} \exp(\tilde{S}_{il})}, \tag{8}$$

applied row-wise. Properties: $A_{ij} \geq 0$, $\sum_j A_{ij} = 1$.

Proposition 5 (Gradient of Softmax). For cross-entropy or downstream loss $L$:

$$\frac{\partial L}{\partial \tilde{S}_{ij}} = A_{ij}\left(\frac{\partial L}{\partial A_{ij}} - \sum_l A_{il}\frac{\partial L}{\partial A_{il}}\right). \tag{9}$$

This Jacobian is well-conditioned when $|S_{ij}| = O(1)$ — another reason for scaling.

---

Scaled Dot-Product Attention

Theorem 2 (Attention Formula).

$$\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}}\right)\mathbf{V}. \tag{10}$$

Row $i$ of the output:

$$\mathbf{o}_i = \sum_{j=1}^{n} A_{ij}\, \mathbf{v}_j. \tag{11}$$

A convex combination of value vectors, weighted by query–key compatibility.

---

Multi-Head Attention

Definition 8 (Multi-Head Attention). With $h$ heads, each with projections $\mathbf{W}_Q^{(k)}, \mathbf{W}_K^{(k)} \in \mathbb{R}^{d \times d_k}$, $\mathbf{W}_V^{(k)} \in \mathbb{R}^{d \times d_v}$:

$$\text{head}_k = \text{Attention}(\mathbf{X}\mathbf{W}_Q^{(k)}, \mathbf{X}\mathbf{W}_K^{(k)}, \mathbf{X}\mathbf{W}_V^{(k)}), \tag{12}$$ $$\text{MultiHead}(\mathbf{X}) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\, \mathbf{W}_O, \tag{13}$$

where $d_k = d_v = d/h$ and $\mathbf{W}_O \in \mathbb{R}^{d \times d}$.

Theorem 3 (FLOP Equivalence). Multi-head attention with $h$ heads of dimension $d/h$ costs $O(n^2 d)$ — the same as single-head attention with $d_k = d$.

Proof. Each head: $O(n^2 \cdot d/h)$. Total: $h \cdot O(n^2 d/h) = O(n^2 d)$. $\blacksquare$

Interpretation. Multiple heads learn different attention patterns (syntactic, semantic, positional) in parallel subspaces without extra compute — a representational gain at fixed FLOPs.

---

Causal Masking for Autoregressive Models

Definition 9 (Causal Mask). For autoregressive models (GPT), token $i$ may only attend to tokens $j \leq i$:

$$\tilde{S}_{ij}^{\text{causal}} = \begin{cases} \tilde{S}_{ij} & j \leq i \\ -\infty & j > i \end{cases} \tag{14}$$

After softmax: $A_{ij} = 0$ for $j > i$ (since $e^{-\infty} = 0$).

Definition 10 (Causal Attention).

$$\text{CausalAttention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}} + \mathbf{M}\right)\mathbf{V}, \tag{15}$$

where $M_{ij} = 0$ if $j \leq i$ and $M_{ij} = -\infty$ if $j > i$.

See KV Cache for efficient autoregressive inference.

---

Computational Complexity and Memory

• $\mathbf{Q}\mathbf{K}^T$ — $O(n^2 d_k)$ — $O(n^2)$ score matrix
• Softmax — $O(n^2)$ — $O(n^2)$ weights
• $\mathbf{A}\mathbf{V}$ — $O(n^2 d_v)$ — $O(n d_v)$ output
• Total — $O(n^2 d)$ — $O(n^2)$

Proposition 6 (Quadratic Bottleneck). For sequence length $n = 10^5$, storing $\mathbf{A} \in \mathbb{R}^{n \times n}$ in FP16 requires $\sim 20$ GB — often exceeding model weight memory.

Flash Attention computes exact attention in $O(n)$ memory via tiling and recomputation.

---

Gradients Through Attention

Proposition 7 (Output Layer Gradient). With loss $\mathcal{L}$ and output $\mathbf{O} = \mathbf{A}\mathbf{V}$:

$$\frac{\partial \mathcal{L}}{\partial \mathbf{V}} = \mathbf{A}^T \frac{\partial \mathcal{L}}{\partial \mathbf{O}}, \qquad \frac{\partial \mathcal{L}}{\partial \mathbf{A}} = \frac{\partial \mathcal{L}}{\partial \mathbf{O}} \mathbf{V}^T. \tag{16}$$

Proposition 8 (Score Gradient). Through softmax:

$$\frac{\partial \mathcal{L}}{\partial \tilde{S}_{ij}} = A_{ij}\left(\frac{\partial \mathcal{L}}{\partial A_{ij}} - \sum_l A_{il}\frac{\partial \mathcal{L}}{\partial A_{il}}\right). \tag{17}$$

Proposition 9 (Q/K Gradients).

$$\frac{\partial \mathcal{L}}{\partial \mathbf{Q}} = \frac{\partial \mathcal{L}}{\partial \tilde{\mathbf{S}}} \mathbf{K} \cdot \frac{1}{\sqrt{d_k}}, \qquad \frac{\partial \mathcal{L}}{\partial \mathbf{K}} = \left(\frac{\partial \mathcal{L}}{\partial \tilde{\mathbf{S}}}\right)^T \mathbf{Q} \cdot \frac{1}{\sqrt{d_k}}. \tag{18}$$

Full derivation follows Backpropagation through the computation graph.

---

Worked Examples

Example 1: Single Query, Two Keys

$\mathbf{q} = (1, 0)^T$, $\mathbf{k}_1 = (1, 0)^T$, $\mathbf{k}_2 = (0, 1)^T$, $d_k = 2$. Scores: $S_1 = 1$, $S_2 = 0$. After scaling: $\tilde{S}_1 = 1/\sqrt{2}$, $\tilde{S}_2 = 0$. Weights: $A_1 \approx 0.65$, $A_2 \approx 0.35$.

Example 2: Causal Mask

For $n = 3$, causal mask zeroes upper triangle. Row 1 attends only to position 1; row 2 to positions 1–2; row 3 to all.

Example 3: Memory at Scale

Llama 2 70B: $L = 80$ layers, $H = 64$ heads, $d_k = 128$, seq_len $n = 4096$. Attention matrix per layer: $4096^2 \times 2$ bytes $\approx 33$ MB. Total across layers dominates KV Cache memory.

---

Connection to the Broader Curriculum

• Positional information — Positional Encoding, RoPE
• Efficient attention — Flash Attention
• Inference — KV Cache
• Low-rank adaptation — LoRA modifies Q/K/V projections
• Dot products — Vectors & Linear Independence

---

Common Pitfalls and Misconceptions

Pitfall 1: Omitting $\sqrt{d_k}$ scaling. Causes softmax saturation and training instability for large $d_k$.

Pitfall 2: Confusing self-attention with cross-attention. Self-attention: Q, K, V from same sequence. Cross-attention: Q from one sequence, K/V from another (encoder–decoder).

Pitfall 3: Forgetting causal mask in generation. Without masking, the model sees future tokens — label leakage during training.

Pitfall 4: Assuming attention weights are interpretable. Weights reflect compatibility in learned subspaces, not necessarily human-interpretable relations.

Pitfall 5: Ignoring $O(n^2)$ memory. Long-context models require Flash Attention or sparse attention variants.

---

Research Perspective

Attention originated in machine translation (Bahdanau et al., 2015). Self-attention and Transformers (Vaswani et al., 2017) replaced recurrence entirely. Universal approximation for Transformers (Yun et al., 2020) extends Universal Approximation Theorem to sequence models.

Modern developments: RoPE, Multi-Query Attention, Flash Attention, linear attention approximations, and state-space models (Mamba) as alternatives to quadratic attention.

---

Summary of Takeaways

• Projections — $\mathbf{Q} = \mathbf{X}\mathbf{W}_Q$, etc.
• Scores — $\mathbf{S} = \mathbf{Q}\mathbf{K}^T$
• Scaling — $\tilde{\mathbf{S}} = \mathbf{S}/\sqrt{d_k}$
• Attention — $\text{softmax}(\tilde{\mathbf{S}})\mathbf{V}$
• Multi-head — $h$ heads, $d_k = d/h$, concat + $\mathbf{W}_O$
• Causal mask — $M_{ij} = -\infty$ for $j > i$
• Complexity — $O(n^2 d)$ time, $O(n^2)$ memory

Next article: Positional Encoding →

---

Exercises

Exercise 1. Prove Theorem 1 (variance of dot product) without skipping steps.

Exercise 2. Show that attention output $\mathbf{o}_i$ lies in the convex hull of $\{\mathbf{v}_j\}$.

Exercise 3. Derive (17) — gradient through softmax.

Exercise 4. Prove Theorem 3 (FLOP equivalence of multi-head attention).

Exercise 5. For $n = 4$, write the causal mask matrix explicitly.

Exercise 6. Compute attention output when $\mathbf{q}^T\mathbf{k}_1 = 10$ and $\mathbf{q}^T\mathbf{k}_2 = 0$ with $d_k = 64$, with and without scaling.

Exercise 7. Explain how LoRA modifies the Q/K/V projections without changing attention structure.

Exercise 8 (Conceptual). Why does self-attention alone not encode token order? How does Positional Encoding fix this?