Scaled Dot-Product Attention

Learning Objectives

By the end of this chapter, you will be able to:

1. Explain why RNNs and LSTMs hit a fundamental bottleneck for long sequences and why a parallel, all-pairs mechanism was needed.
2. Derive the scaled dot-product attention formula $\text{softmax}\!\bigl(\tfrac{QK^\top}{\sqrt{d_k}}\bigr)V$ from first principles, understanding every symbol and operation.
3. Prove why the $1/\sqrt{d_k}$ scaling factor is necessary and what happens without it.
4. Compute attention weights and outputs by hand for our shared example sentence “The cat sat on mat”.
5. Implement a complete, runnable Python class that you can use to simulate scaled dot-product attention on any input.
6. Connect this foundational mechanism to its extensions: multi-head attention, Flash Attention, KV-cache, and positional encodings.

Where this appears: Scaled dot-product attention, or a close variant of it, is the computational core of most transformer-based models — including GPT-4, Claude, Gemini, LLaMA, Stable Diffusion, AlphaFold, Codex, and ViT. These systems compute this formula (or an optimized equivalent like Flash Attention) billions of times during inference. Understanding it deeply is the foundation for understanding modern AI architectures.

The Shared Example — Used in Every Chapter

Every mechanism in this book operates on the exact same sentence, the same matrices, and the same parameters. This allows you to directly compare what each mechanism does differently — the only thing that changes is the attention computation itself.

The Sentence and Tokens

The Query, Key, and Value Matrices

$Q$, $K$, and $V$ are the three fundamental matrices in every attention mechanism. In practice they are produced by learned linear projections of the input embeddings: $Q = XW^Q$, $K = XW^K$, $V = XW^V$, where $X \in \mathbb{R}^{N \times d_{\text{model}}}$ is the input embedding matrix, $W^Q, W^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ are the learned query and key projections, and $W^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$ is the learned value projection (in practice $d_v = d_k$, but the formalism allows them to differ). Here we fix them so all 15 mechanisms start from the same point.

$Q$ (Query Matrix) — $5 \times 4$ — each row encodes “what is this token looking for?”

$K$ (Key Matrix) — $5 \times 4$ — each row encodes “what information does this token advertise?”

$V$ (Value Matrix) — $5 \times 4$ — each row is the actual content that gets retrieved when a token is attended to

What Every Chapter Shows

For each mechanism you will find four sections:

1. Problem & Intuition — what limitation this mechanism was designed to fix
2. The Math — the exact formula with every variable defined
3. Step-by-Step Calculation — the formula applied to “The” (row 0) using real numbers
4. Python Code — a clean, runnable class implementation using only NumPy

Every chapter ends with the full Attention Weight Matrix ($5 \times 5$) and Output Matrix ($5 \times 4$) so you can immediately see how the pattern changes from one mechanism to the next.

The Real Problem

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L., & Polosukhin, I. (2017). “Attention Is All You Need.” Advances in Neural Information Processing Systems (NeurIPS), 30.

To understand why scaled dot-product attention exists, you must first understand what it replaced and why that predecessor was breaking under the weight of real-world language.

The RNN Bottleneck

Before 2017, the dominant architecture for sequence modeling was the Recurrent Neural Network (RNN) and its variants (LSTM, GRU). These models processed tokens sequentially — one at a time, left to right — accumulating information into a fixed-size hidden state vector $h_t \in \mathbb{R}^d$.

This created three critical problems:

1. The information bottleneck. The entire history of a 10,000-token document had to be compressed into a single vector $h_t$ of, say, 512 dimensions. By the time the model reached the end of a long document, the information from the beginning had been overwritten or diluted beyond recognition. In seq2seq translation (Sutskever et al., 2014), the encoder's final hidden state was the only bridge between the source and target language — a catastrophic bottleneck.
2. No parallelism. Because $h_t$ depended on $h_{t-1}$, every timestep had to wait for the previous one to finish. You could not process tokens 1 through 10,000 in parallel. Training on long sequences was agonizingly slow, even on GPUs designed for massive parallelism.
3. Vanishing and exploding gradients. During backpropagation through time (BPTT), gradients had to flow through $T$ multiplicative steps. Even with LSTM gating, gradients for dependencies spanning hundreds of tokens degraded to near zero, making it nearly impossible to learn long-range patterns like coreference (“The cat... it...” separated by 50 tokens).

The Fundamental Question

Vaswani et al. asked a radical question: what if we removed the recurrence entirely? What if, instead of processing tokens sequentially, every token could directly attend to every other token in a single parallel operation?

The answer required a mechanism with three properties:

• All-pairs comparison — every token should be able to compute its relevance to every other token in the sequence, in parallel.
• Differentiable selection — the model should learn which tokens are relevant via soft weights (probabilities), not hard selection, so gradients can flow.
• Content-based addressing — relevance should be determined by what the tokens mean (their representations), not by their position alone.

Scaled dot-product attention satisfies all three. It is the mechanism that made the Transformer possible.

From Intuition to Mathematics

The Query-Key-Value Intuition

The central metaphor is a library lookup. Imagine you walk into a library with a question in mind (your query). Each book on the shelf has a label on its spine (its key) and content inside (its value). You scan all the spine labels, judge which ones are most relevant to your question, and then read a weighted blend of the most relevant books.

The critical insight is that Q and K live in the same space so their dot product is meaningful as a similarity measure, while V can live in a different space — it carries the content that the query wants to retrieve.

Dot Product as Similarity

Why use the dot product to measure relevance? Consider two vectors $\mathbf{q} = [1, 0, 1, 0]$ and $\mathbf{k} = [1, 0, 1, 0]$:

$\mathbf{q} \cdot \mathbf{k} = 1 \times 1 + 0 \times 0 + 1 \times 1 + 0 \times 0 = 2$

The dot product is large because the vectors “agree” — they are active in the same dimensions. Now compare with $\mathbf{k}' = [0, 1, 0, 1]$:

$\mathbf{q} \cdot \mathbf{k}' = 1 \times 0 + 0 \times 1 + 1 \times 0 + 0 \times 1 = 0$

The dot product is zero because the vectors are orthogonal — they have nothing in common. This is exactly what we want: the dot product naturally measures the alignment between what the query is looking for and what the key advertises. Geometrically, $\mathbf{q} \cdot \mathbf{k} = \|\mathbf{q}\| \|\mathbf{k}\| \cos\theta$, so it captures both magnitude and direction.

The Mathematical Definition

The complete formula for scaled dot-product attention is:

This single line is the most important equation in modern deep learning. Let us dissect every piece.

Symbol-by-Symbol Breakdown

What the Formula Says in Plain English

The formula performs four operations in sequence:

1. Compare every query with every key via dot product ($Q K^\top$), producing an $N \times N$ matrix of raw similarity scores.
2. Scale the scores by $1/\sqrt{d_k}$ to prevent gradient saturation in the softmax.
3. Normalize each row with softmax to convert scores into attention weights (probabilities that sum to 1).
4. Aggregate by multiplying the attention weights with the value matrix, producing a context-aware output for each token.

In short: every token asks “who is relevant to me?” (via Q · K), computes how relevant (via softmax), and collects a weighted blend of answers (via V). The result is that each token's output representation is enriched with contextual information from all other tokens, with more weight given to the most relevant ones.

Why Scale? The Variance Argument

This is the most commonly asked question about the formula: why divide by $\sqrt{d_k}$ and not some other constant? The answer comes from a simple variance analysis.

The Mathematical Proof

Assume each element of $Q$ and $K$ is drawn independently from a distribution with mean 0 and variance 1. The dot product of two such vectors is:

Since each $q_m$ and $k_m$ are independent with mean 0 and variance 1:

• $\mathbb{E}[q_m \cdot k_m] = \mathbb{E}[q_m] \cdot \mathbb{E}[k_m] = 0 \cdot 0 = 0$
• $\text{Var}(q_m \cdot k_m) = \mathbb{E}[q_m^2] \cdot \mathbb{E}[k_m^2] = 1 \cdot 1 = 1$

The dot product is a sum of $d_k$ independent terms, each with variance 1, so by the additivity of variance:

When $d_k = 512$, the standard deviation of the dot product is $\sqrt{512} \approx 22.6$. This means individual scores can easily reach values like $\pm 40$ or more.

Now consider what softmax does with such extreme values. If one score is 40 and the rest are near 0:

The output is essentially a one-hot vector. The gradient of softmax at these saturation points is near zero ($\partial \text{softmax}(z_i) / \partial z_j \to 0$), which means the model cannot learn to adjust these scores — learning stalls.

The fix is elegant: dividing by $\sqrt{d_k}$ rescales the variance back to 1:

Now the scores have a standard deviation of ~1 regardless of the dimension. Softmax receives moderate inputs, produces smooth (non-saturated) probability distributions, and gradients flow properly for learning. This is why we scale by exactly $\sqrt{d_k}$ — not 2, not $d_k$, not a learned parameter. It is the mathematically correct normalization to keep the softmax in its useful operating regime.

Interactive: Scaling Explorer

Use the interactive explorer below to see exactly what happens to softmax outputs as you increase $d_k$. Toggle between “With Scaling” and “Without Scaling” to see how the distribution collapses without the $\sqrt{d_k}$ correction.

Step-by-Step Calculation for “The” (Row 0)

We now trace through every arithmetic step for the first token “The” to see exactly how scaled dot-product attention builds its output vector. Every number here can be verified with the Python class at the bottom of this chapter.

Step 1 — Raw Dot Products: $S_{0,:} = Q_0 \cdot K^\top$

The query vector for “The” is $Q_0 = [1, 0, 1, 0]$. We compute its dot product with every key vector:

Interpretation: “The” has the highest raw similarity with “cat” (score = 2.0) because $Q_{\text{The}} = [1, 0, 1, 0]$ and $K_{\text{cat}} = [1, 0, 1, 0]$ are identical — perfect alignment. The lowest similarity is with itself (score = 0.0) because $K_{\text{The}} = [0, 1, 0, 1]$ is orthogonal to its query.

Step 2 — Scaling: $S_{0,:} / \sqrt{d_k}$

We divide each raw score by $\sqrt{4} = 2.0$:

In our small example ($d_k = 4$), scaling reduces scores by half. In a production model with $d_k = 128$, scores would be divided by ~11.3, a much more dramatic reduction.

Step 3 — Softmax: $A_{0,:} = \text{softmax}(S_{0,:})$

The softmax function converts the scaled scores into a probability distribution. For a vector $z = [z_1, \ldots, z_N]$, the softmax of element $i$ is $e^{z_i} / \sum_{j} e^{z_j}$.

Computing each exponential:

Sum of exponentials: $1.000 + 2.718 + 1.649 + 1.649 + 2.117 = 9.133$

Dividing each exponential by the sum:

Interpretation: “The” pays 29.76% of its attention to “cat” — the token most aligned with its query. But notice that attention is not concentrated on a single token. “mat” also receives significant weight (23.18%) because its key partially overlaps with the query. This soft, distributed attention is what allows the mechanism to capture nuanced relationships.

Step 4 — Weighted Sum of Values: $O_0 = A_{0,:} \cdot V$

Each output dimension is the weighted average of that dimension across all value vectors:

Expanding each value vector:

10.1095 × [1.0, 0.0, 0.0, 0.0] = [0.1095, 0.0000, 0.0000, 0.0000]  (The)
20.2976 × [0.0, 1.0, 0.0, 0.0] = [0.0000, 0.2976, 0.0000, 0.0000]  (cat)
30.1805 × [0.0, 0.0, 1.0, 0.0] = [0.0000, 0.0000, 0.1805, 0.0000]  (sat)
40.1805 × [0.0, 0.0, 0.0, 1.0] = [0.0000, 0.0000, 0.0000, 0.1805]  (on)
50.2318 × [0.5, 0.5, 0.5, 0.5] = [0.1159, 0.1159, 0.1159, 0.1159]  (mat)
6─────────────────────────────────────────────────────────────────────
7Sum:                             [0.2254, 0.4135, 0.2964, 0.2964]  ← Output

Final output for “The”: $O_0 = [0.2254, 0.4135, 0.2964, 0.2964]$

Notice that dim-1 has the largest value (0.4135) because “cat” — the most-attended token — contributes all its weight to dim-1 ($V_{\text{cat}} = [0, 1, 0, 0]$). The output for “The” has been enriched with the semantic content of “cat”, reflecting the strong query-key alignment between these two tokens.

Interactive: Attention Pipeline

Select any token to trace the full attention computation step by step. Click through the four stages — dot products, scaling, softmax, and weighted output — to see how attention flows from the query to the final output.

Full Attention Weights and Output

Below is the complete attention weight matrix and its corresponding output for all five tokens. Hover over any cell in the heatmap to see the full computation for that query-key pair.

Interpreting the Heatmap

Several patterns emerge from the attention weight matrix:

• “cat” attends most to “The” (0.4026) — the determiner that typically precedes it. $Q_{\text{cat}} = [0, 2, 0, 1]$ has strong overlap with $K_{\text{The}} = [0, 1, 0, 1]$ (dot product = 3.0, the largest in the entire matrix).
• “on” attends most to itself (0.3137) — $Q_{\text{on}} = K_{\text{on}} = [0, 0, 1, 1]$. When the query and key of the same token are aligned, self-attention is high. This is common for function words that carry positional rather than semantic weight.
• “mat” distributes attention nearly uniformly (all weights near 0.19-0.24) — its query $[1, 0, 0, 1]$ partially matches many keys without strongly preferring one.

Applications Across Domains

Scaled dot-product attention is not limited to language. The same formula powers breakthroughs across every domain that involves sequences, sets, or structured data.

Natural Language Processing

In GPT-4, Claude, and LLaMA, each layer applies scaled dot-product attention so that every token can gather context from every other token. When the model processes “The cat sat on the mat because it was tired,” the attention mechanism for “it” assigns high weight to “cat” — resolving the pronoun coreference. Without attention, an RNN would need to carry the “cat” information through every intermediate hidden state.

Computer Vision

In Vision Transformers (ViT; Dosovitskiy et al., 2021), an image is split into 16×16 patches. Each patch becomes a token, and scaled dot-product attention lets every patch attend to every other patch. A patch containing a dog's eye can attend to the patch containing its tail — capturing long-range spatial relationships that CNNs can only reach through many stacked convolutional layers.

Code Generation

In Codex and GitHub Copilot, attention enables the model to connect a function's name (token 1) with its return type (token 50), its docstring (tokens 5-20), and the variable names used inside (tokens 30-100). When generating return total_price, the model's attention assigns high weight to the earlier definition total_price = base_price * quantity — even if that line is hundreds of tokens away.

Scientific Sequence Modeling

In AlphaFold2 (Jumper et al., 2021), attention operates over amino acid residues in a protein sequence. Each residue attends to every other residue, learning which pairs of amino acids are likely to be spatially close in the 3D folded structure — even when they are far apart in the linear sequence. The attention weights effectively learn a proxy for physical contact maps, which is remarkable because this structural information is not given as supervision.

Real-World Scale

To appreciate the scale at which this formula operates, the table below shows illustrative parameters from published model architectures. The “Attention Ops” column is a back-of-the-envelope estimate using $\text{Layers} \times \text{Heads} \times N^2$ (the dominant term from $QK^\top$ across all heads and layers).

The last row is our toy example — the same formula, just orders of magnitude smaller. This is why understanding the mechanism at small scale transfers directly to understanding production systems.

Connection to Modern Systems

Scaled dot-product attention is the atomic building block. Every subsequent chapter in this book modifies, extends, or optimizes it. Here is how the major variants relate to this foundation.

Multi-Head Attention (Chapter 2)

Instead of running one attention function with $d_{\text{model}}$-dimensional keys, multi-head attention runs $H$ independent scaled dot-product attention operations in parallel, each on a smaller $d_k = d_{\text{model}}/H$-dimensional subspace. This allows different heads to learn different types of relationships — one head might learn syntactic dependency, another might learn semantic similarity, another might learn positional patterns. The outputs are concatenated and linearly projected back to $d_{\text{model}}$.

Flash Attention (Chapter 13)

Flash Attention (Dao et al., 2022) computes mathematically identical results to scaled dot-product attention, but reorganizes the computation to minimize GPU memory reads/writes. The key insight is that the $N \times N$ attention matrix is never fully materialized in GPU HBM (high-bandwidth memory). Instead, it is computed in tiles that fit in SRAM (on-chip memory), achieving 2-4x speedup and reducing memory from $O(N^2)$ to $O(N)$. The math is the same; only the memory access pattern changes.

KV-Cache Optimization

During autoregressive generation (producing one token at a time), the model computes attention for the new token's query against all previous keys and values. The KV-cache stores the $K$ and $V$ matrices from all previous tokens so they don't need to be recomputed from the input embeddings at each step. The per-token attention cost remains $O(N \cdot d_k)$ (the new query still scores against all $N$ cached keys), but without the cache the model would need to reproject and recompute attention for all previous tokens from scratch — an $O(N^2 \cdot d_k)$ cost per generated token. The cache trades $O(N \cdot d_k)$ memory for this saving. Multi-Query Attention (Chapter 5) and Grouped-Query Attention (Chapter 6) reduce the memory cost by sharing K and V across heads.

Positional Encodings (Chapters 7-9)

Scaled dot-product attention is permutation-invariant — it treats the input as a set, not a sequence. Swapping the order of tokens in the input would produce the same attention weights (just reordered). To inject positional information, the Transformer adds positional encodings to the input embeddings. Chapters 7-9 cover three approaches: learned relative position bias (T5), Rotary Position Embeddings (RoPE, used in LLaMA), and Attention with Linear Biases (ALiBi). Each modifies either the input representations or the attention scores to encode where tokens are, not just what they are.

Complexity Analysis

To understand why attention becomes expensive, do not start with formulas. Start with one simple question:

As the sequence gets longer, how much more work does the model have to do, and how much more memory does it need to hold things while doing that work?

That is all complexity analysis is trying to measure.

• Time complexity tells us how the amount of work grows.
• Space complexity tells us how the amount of memory grows.

In attention, both grow quickly because every token can interact with every other token.

A Simple Picture: Students in a Classroom

Imagine a classroom with $N$ students. Each student wants to decide: “Which other students should I listen to?”

If every student checks every other student, then student 1 checks $N$ students, student 2 checks $N$ students, and so on. The total number of checks is:

That is the main time cost.

Now imagine writing every one of those scores on a big board. The board needs $N$ rows and $N$ columns, so it stores $N^2$ numbers. That is the main space cost. This is exactly what happens in attention.

Where the Cost Comes From in Attention

In scaled dot-product attention, each token asks: “How relevant is every other token to me?” The model answers that by building a giant score table:

If there are $N$ tokens, this table has shape $N \times N$. That means computing it takes a lot of work, and storing it takes a lot of memory. That single all-to-all interaction pattern is the reason attention becomes expensive.

Step-by-Step Intuition

Step 1: Compute $QK^\top$ — the score matrix. This is the “everyone compares with everyone” step. Each token compares itself with every other token. If there are $N$ tokens, there are about $N^2$ token pairs. But each comparison is not just one tiny action — it is a dot product over vectors of size $d_k$. So each pair needs about $d_k$ small multiply-and-add operations.

Time: $O(N^2 d_k)$. Space: $O(N^2)$.

Intuition: a giant spreadsheet where every row is a token and every column is another token, and each cell stores “how much should I pay attention?”

Step 2: Scaling. After the scores are computed, each one is divided by $\sqrt{d_k}$. This keeps the numbers from getting too large. There are $N^2$ scores, so touching all of them costs $O(N^2)$ time. Usually this is done directly on the same score table, so it does not need another giant table — the extra memory cost is $O(1)$.

Intuition: you already filled the board with numbers, and now you go through and slightly shrink each one. That is work, but not extra board space.

Step 3: Softmax. Each row of the score matrix is turned into normalized attention weights. Each row has $N$ numbers, and there are $N$ rows, so the total work is $O(N^2)$. The model still stores the full weight table, so memory is still $O(N^2)$.

Intuition: the raw scores become “attention percentages.” Same giant board, same size, just cleaner numbers.

Step 4: Multiply weights by $V$. The model uses those attention weights to mix the value vectors and produce a new output vector for each token. For each of the $N$ output tokens, it looks across $N$ tokens, and each value vector has size $d_v$. So the work is $O(N^2 d_v)$. The output stores one vector of length $d_v$ per token, so the memory needed is only $O(N d_v)$.

Intuition: each token builds a summary of what it learned from everyone else.

Interactive: Worked Example with 3 Tokens

The steps above describe what happens in words. But to truly understand why each operation costs what it does, nothing beats watching the numbers unfold. The walkthrough below uses 3 tokens with vectors of length 3 and walks through every single multiply-and-add. Click Next to advance.

A Tiny Example

Suppose there are only 4 tokens. The score matrix has $4 \times 4 = 16$ entries. That is small. Easy to compute. Easy to store.

Now suppose there are 1,000 tokens. The score matrix has $1{,}000 \times 1{,}000 = 1{,}000{,}000$ entries.

Now suppose there are 8,192 tokens. The score matrix has $8{,}192 \times 8{,}192 = 67{,}108{,}864$ entries. That is about 67 million scores for just one head in one layer.

This is why attention becomes expensive so quickly: the sequence length grows linearly, but the pairwise interaction table grows quadratically.

What “Bottleneck” Really Means Here

The bottleneck is the part that hurts the most as the input grows. In attention, that bottleneck is the $N^2$ interaction pattern. Both the time to compute all token-to-token interactions and the space to store all those interaction scores become large. So when people say standard attention has an $O(N^2)$ bottleneck, they mean:

The cost blows up because every token talks to every token.

Summary Table

The reason the total is dominated by quadratic behavior is simple: the giant $N \times N$ board keeps showing up.

At $N = 8{,}192$ tokens (a typical context window), the attention matrix has 67 million entries — per head, per layer. At 32 heads and 32 layers, that is ~69 billion score computations per forward pass.

Attention becomes expensive because every token compares itself with every other token. If there are $N$ tokens, that creates an $N \times N$ table of interaction scores. Building that table takes computation, and storing it takes memory. Time complexity measures how the amount of computation grows; space complexity measures how the required memory grows. In standard attention, both are dominated by this all-to-all interaction pattern, which is why the cost scales quadratically with sequence length.

Python Implementation

Below is a complete, runnable implementation of scaled dot-product attention as a Python class. Click any highlighted line on the right to see its detailed explanation on the left. You can copy the full code and run it with python scaled_dot_product_attention.py — the only dependency is NumPy.

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

     The   cat   sat    on   mat
d0  0.0   1.0   1.0   0.0   1.0
d1  1.0   0.0   1.0   0.0   0.0
d2  0.0   1.0   0.0   1.0   0.5
d3  1.0   0.0   0.0   1.0   0.5

      The   cat   sat    on   mat
The  0.00  2.00  1.00  1.00  1.50
cat  3.00  0.00  2.00  1.00  0.50
sat  1.00  2.00  2.00  1.00  1.50
on   1.00  1.00  0.00  2.00  1.00
mat  1.00  1.00  1.00  1.00  1.50

      The   cat   sat    on   mat
The  0.00  2.00  1.00  1.00  1.50
cat  3.00  0.00  2.00  1.00  0.50
sat  1.00  2.00  2.00  1.00  1.50
on   1.00  1.00  0.00  2.00  1.00
mat  1.00  1.00  1.00  1.00  1.50

       The     cat     sat      on     mat
The  0.000   1.000   0.500   0.500   0.750
cat  1.500   0.000   1.000   0.500   0.250
sat  0.500   1.000   1.000   0.500   0.750
on   0.500   0.500   0.000   1.000   0.500
mat  0.500   0.500   0.500   0.500   0.750

       The     cat     sat      on     mat
The  0.000   1.000   0.500   0.500   0.750
cat  1.500   0.000   1.000   0.500   0.250
sat  0.500   1.000   1.000   0.500   0.750
on   0.500   0.500   0.000   1.000   0.500
mat  0.500   0.500   0.500   0.500   0.750

       The      cat      sat       on      mat
The  0.1095   0.2976   0.1805   0.1805   0.2318
cat  0.4026   0.0898   0.2442   0.1481   0.1153
sat  0.1519   0.2505   0.2505   0.1519   0.1951
on   0.1903   0.1903   0.1154   0.3137   0.1903
mat  0.1892   0.1892   0.1892   0.1892   0.2430

       The      cat      sat       on      mat
The  0.1095   0.2976   0.1805   0.1805   0.2318
cat  0.4026   0.0898   0.2442   0.1481   0.1153
sat  0.1519   0.2505   0.2505   0.1519   0.1951
on   0.1903   0.1903   0.1154   0.3137   0.1903
mat  0.1892   0.1892   0.1892   0.1892   0.2430

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

        d0       d1       d2       d3
The  0.2254   0.4135   0.2964   0.2964
cat  0.4602   0.1475   0.3018   0.2058
sat  0.2495   0.3481   0.3481   0.2495
on   0.2854   0.2854   0.2106   0.4089
mat  0.3108   0.3108   0.3108   0.3108

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

      The   cat   sat    on   mat
The  0.00  2.00  1.00  1.00  1.50
cat  3.00  0.00  2.00  1.00  0.50
sat  1.00  2.00  2.00  1.00  1.50
on   1.00  1.00  0.00  2.00  1.00
mat  1.00  1.00  1.00  1.00  1.50

       The     cat     sat      on     mat
The  0.000   1.000   0.500   0.500   0.750
cat  1.500   0.000   1.000   0.500   0.250
sat  0.500   1.000   1.000   0.500   0.750
on   0.500   0.500   0.000   1.000   0.500
mat  0.500   0.500   0.500   0.500   0.750

       The      cat      sat       on      mat
The  0.1095   0.2976   0.1805   0.1805   0.2318
cat  0.4026   0.0898   0.2442   0.1481   0.1153
sat  0.1519   0.2505   0.2505   0.1519   0.1951
on   0.1903   0.1903   0.1154   0.3137   0.1903
mat  0.1892   0.1892   0.1892   0.1892   0.2430

        d0       d1       d2       d3
The  0.2254   0.4135   0.2964   0.2964
cat  0.4602   0.1475   0.3018   0.2058
sat  0.2495   0.3481   0.3481   0.2495
on   0.2854   0.2854   0.2106   0.4089
mat  0.3108   0.3108   0.3108   0.3108

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

      d0   d1   d2   d3
The  0.0  1.0  0.0  1.0
cat  1.0  0.0  1.0  0.0
sat  1.0  1.0  0.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.5  0.5

      d0   d1   d2   d3
The  1.0  0.0  0.0  0.0
cat  0.0  1.0  0.0  0.0
sat  0.0  0.0  1.0  0.0
on   0.0  0.0  0.0  1.0
mat  0.5  0.5  0.5  0.5

       The      cat      sat       on      mat
The  0.1095   0.2976   0.1805   0.1805   0.2318
cat  0.4026   0.0898   0.2442   0.1481   0.1153
sat  0.1519   0.2505   0.2505   0.1519   0.1951
on   0.1903   0.1903   0.1154   0.3137   0.1903
mat  0.1892   0.1892   0.1892   0.1892   0.2430

        d0       d1       d2       d3
The  0.2254   0.4135   0.2964   0.2964
cat  0.4602   0.1475   0.3018   0.2058
sat  0.2495   0.3481   0.3481   0.2495
on   0.2854   0.2854   0.2106   0.4089
mat  0.3108   0.3108   0.3108   0.3108

1import numpy as np
2import math
3
4class ScaledDotProductAttention:
5    """
6    Scaled Dot-Product Attention (Vaswani et al., 2017)
7
8    Computes: Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V
9
10    This is the foundational attention mechanism used in every transformer.
11    All 15 mechanisms in this book build upon or modify this core computation.
12    """
13
14    def __init__(self, d_k: int):
15        """
16        Args:
17            d_k: Dimension of query/key vectors (used for scaling)
18        """
19        self.d_k = d_k
20        self.scale = math.sqrt(d_k)
21
22    def _softmax(self, x: np.ndarray) -> np.ndarray:
23        """Numerically stable softmax along last axis."""
24        x_shifted = x - np.max(x, axis=-1, keepdims=True)
25        exp_x = np.exp(x_shifted)
26        return exp_x / np.sum(exp_x, axis=-1, keepdims=True)
27
28    def compute_scores(self, Q: np.ndarray, K: np.ndarray) -> np.ndarray:
29        """Step 1: Raw dot-product scores Q @ K^T."""
30        return Q @ K.T
31
32    def scale_scores(self, scores: np.ndarray) -> np.ndarray:
33        """Step 2: Divide by sqrt(d_k) to control variance."""
34        return scores / self.scale
35
36    def compute_weights(self, scaled_scores: np.ndarray) -> np.ndarray:
37        """Step 3: Apply softmax to get attention weights."""
38        return self._softmax(scaled_scores)
39
40    def compute_output(self, weights: np.ndarray, V: np.ndarray) -> np.ndarray:
41        """Step 4: Weighted sum of value vectors."""
42        return weights @ V
43
44    def forward(self, Q: np.ndarray, K: np.ndarray, V: np.ndarray):
45        """
46        Full forward pass.
47
48        Args:
49            Q: Query matrix  (N, d_k)
50            K: Key matrix    (N, d_k)
51            V: Value matrix  (N, d_v)
52
53        Returns:
54            weights: Attention weight matrix  (N, N)
55            output:  Context-enriched output  (N, d_v)
56        """
57        raw_scores    = self.compute_scores(Q, K)       # (N, N)
58        scaled_scores = self.scale_scores(raw_scores)    # (N, N)
59        weights       = self.compute_weights(scaled_scores)  # (N, N)
60        output        = self.compute_output(weights, V)  # (N, d_v)
61        return weights, output
62
63    def explain(self, Q: np.ndarray, K: np.ndarray, V: np.ndarray,
64                tokens: list, query_idx: int = 0):
65        """
66        Print a detailed trace of the attention computation
67        for a specific query token.
68        """
69        raw_scores    = self.compute_scores(Q, K)
70        scaled_scores = self.scale_scores(raw_scores)
71        weights       = self.compute_weights(scaled_scores)
72        output        = self.compute_output(weights, V)
73
74        token = tokens[query_idx]
75        print(f"\n=== Attention trace for '{token}' (row {query_idx}) ===")
76        print(f"Q[{query_idx}] = {Q[query_idx]}")
77        print(f"\nStep 1 - Raw dot products (Q @ K^T):")
78        for j, t in enumerate(tokens):
79            print(f"  Q[{token}] . K[{t}] = {raw_scores[query_idx, j]:.4f}")
80
81        print(f"\nStep 2 - Scaled (/ sqrt({self.d_k}) = / {self.scale:.1f}):")
82        for j, t in enumerate(tokens):
83            print(f"  S[{token},{t}] = {scaled_scores[query_idx, j]:.4f}")
84
85        print(f"\nStep 3 - Softmax weights:")
86        for j, t in enumerate(tokens):
87            bar = '#' * int(weights[query_idx, j] * 40)
88            print(f"  A[{token},{t}] = {weights[query_idx, j]:.4f} |{bar}|")
89
90        print(f"\nStep 4 - Output (weighted sum of V):")
91        print(f"  O[{token}] = {output[query_idx]}")
92        print(f"  Sum of weights = {weights[query_idx].sum():.6f}")
93
94
95# ── Shared Example (used in every chapter) ──
96tokens = ["The", "cat", "sat", "on", "mat"]
97
98Q = np.array([
99    [1.0, 0.0, 1.0, 0.0],   # The
100    [0.0, 2.0, 0.0, 1.0],   # cat
101    [1.0, 1.0, 1.0, 0.0],   # sat
102    [0.0, 0.0, 1.0, 1.0],   # on
103    [1.0, 0.0, 0.0, 1.0],   # mat
104])
105
106K = np.array([
107    [0.0, 1.0, 0.0, 1.0],   # The
108    [1.0, 0.0, 1.0, 0.0],   # cat
109    [1.0, 1.0, 0.0, 0.0],   # sat
110    [0.0, 0.0, 1.0, 1.0],   # on
111    [1.0, 0.0, 0.5, 0.5],   # mat
112])
113
114V = np.array([
115    [1.0, 0.0, 0.0, 0.0],   # The
116    [0.0, 1.0, 0.0, 0.0],   # cat
117    [0.0, 0.0, 1.0, 0.0],   # sat
118    [0.0, 0.0, 0.0, 1.0],   # on
119    [0.5, 0.5, 0.5, 0.5],   # mat
120])
121
122# ── Run ──
123attn = ScaledDotProductAttention(d_k=4)
124weights, output = attn.forward(Q, K, V)
125
126print("Attention Weight Matrix (5x5):")
127print(np.round(weights, 4))
128
129print("\nOutput Matrix (5x4):")
130print(np.round(output, 4))
131
132# Detailed trace for "The" (token 0)
133attn.explain(Q, K, V, tokens, query_idx=0)

PyTorch Implementation

The NumPy version above is ideal for understanding every arithmetic step. But in practice, you will use PyTorch — it gives you GPU acceleration, automatic differentiation, and a built-in optimized implementation. Below is the same scaled dot-product attention as an nn.Module, using the identical shared example.

Three things change from NumPy to PyTorch:

1. Tensors replace arrays. torch.tensor tracks computation graphs for gradient computation. np.array does not.
2. Batching is built in. torch.matmul and K.transpose(-2, -1) work for any number of leading batch dimensions. The NumPy K.T only works for 2D matrices.
3. Masking support. The PyTorch version accepts an optional mask tensor that sets masked positions to $-\infty$ before softmax, so $e^{-\infty} = 0$ and those tokens receive zero attention. This is how causal masking (Chapter 3) and padding masks are implemented.

import torch; x = torch.tensor([1.0, 2.0])

attn = ScaledDotProductAttention(d_k=64).cuda()  # ready for GPU

torch.matmul(Q, K.transpose(-2, -1))  # same as Q @ K.mT in PyTorch 2.0+

      The   cat   sat    on   mat
The  0.00  2.00  1.00  1.00  1.50
cat  3.00  0.00  2.00  1.00  0.50
sat  1.00  2.00  2.00  1.00  1.50
on   1.00  1.00  0.00  2.00  1.00
mat  1.00  1.00  1.00  1.00  1.50

       The     cat     sat      on     mat
The  0.000   1.000   0.500   0.500   0.750
cat  1.500   0.000   1.000   0.500   0.250
sat  0.500   1.000   1.000   0.500   0.750
on   0.500   0.500   0.000   1.000   0.500
mat  0.500   0.500   0.500   0.500   0.750

# Causal mask: mask future tokens
mask = torch.triu(torch.ones(5, 5, dtype=torch.bool), diagonal=1)

       The      cat      sat       on      mat
The  0.1095   0.2976   0.1805   0.1805   0.2318
cat  0.4026   0.0898   0.2442   0.1481   0.1153
sat  0.1519   0.2505   0.2505   0.1519   0.1951
on   0.1903   0.1903   0.1154   0.3137   0.1903
mat  0.1892   0.1892   0.1892   0.1892   0.2430

        d0       d1       d2       d3
The  0.2254   0.4135   0.2964   0.2964
cat  0.4602   0.1475   0.3018   0.2058
sat  0.2495   0.3481   0.3481   0.2495
on   0.2854   0.2854   0.2106   0.4089
mat  0.3108   0.3108   0.3108   0.3108

      d0   d1   d2   d3
The  1.0  0.0  1.0  0.0
cat  0.0  2.0  0.0  1.0
sat  1.0  1.0  1.0  0.0
on   0.0  0.0  1.0  1.0
mat  1.0  0.0  0.0  1.0

# Wrong: attn.forward(Q, K, V)
# Right: attn(Q, K, V)

       The      cat      sat       on      mat
The  0.1095   0.2976   0.1805   0.1805   0.2318
cat  0.4026   0.0898   0.2442   0.1481   0.1153
sat  0.1519   0.2505   0.2505   0.1519   0.1951
on   0.1903   0.1903   0.1154   0.3137   0.1903
mat  0.1892   0.1892   0.1892   0.1892   0.2430

        d0       d1       d2       d3
The  0.2254   0.4135   0.2964   0.2964
cat  0.4602   0.1475   0.3018   0.2058
sat  0.2495   0.3481   0.3481   0.2495
on   0.2854   0.2854   0.2106   0.4089
mat  0.3108   0.3108   0.3108   0.3108

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4import math
5
6class ScaledDotProductAttention(nn.Module):
7    """
8    Scaled Dot-Product Attention (Vaswani et al., 2017) — PyTorch
9
10    Computes: Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V
11
12    This nn.Module version supports:
13      - GPU acceleration (just move tensors to CUDA)
14      - Automatic differentiation (gradients flow through)
15      - Batched inputs (batch_size, N, d_k)
16      - Optional attention mask for causal / padding masks
17    """
18
19    def __init__(self, d_k: int):
20        super().__init__()
21        self.d_k = d_k
22        self.scale = math.sqrt(d_k)
23
24    def forward(
25        self,
26        Q: torch.Tensor,
27        K: torch.Tensor,
28        V: torch.Tensor,
29        mask: torch.Tensor | None = None,
30    ) -> tuple[torch.Tensor, torch.Tensor]:
31        """
32        Args:
33            Q: Query tensor   (..., N, d_k)
34            K: Key tensor     (..., S, d_k)
35            V: Value tensor   (..., S, d_v)
36            mask: Optional boolean mask (..., N, S)
37                  True = position is MASKED (ignored)
38
39        Returns:
40            weights: Attention weights  (..., N, S)
41            output:  Contextualized     (..., N, d_v)
42        """
43        # Step 1: Raw scores — matrix multiply Q with K transposed
44        scores = torch.matmul(Q, K.transpose(-2, -1))   # (..., N, S)
45
46        # Step 2: Scale — divide by sqrt(d_k)
47        scores = scores / self.scale
48
49        # Step 3 (optional): Apply mask before softmax
50        if mask is not None:
51            scores = scores.masked_fill(mask, float("-inf"))
52
53        # Step 4: Softmax — normalize each row to probabilities
54        weights = F.softmax(scores, dim=-1)              # (..., N, S)
55
56        # Step 5: Weighted sum of value vectors
57        output = torch.matmul(weights, V)                # (..., N, d_v)
58
59        return weights, output
60
61
62# ── Shared Example (same matrices as NumPy version) ──
63tokens = ["The", "cat", "sat", "on", "mat"]
64
65Q = torch.tensor([
66    [1.0, 0.0, 1.0, 0.0],   # The
67    [0.0, 2.0, 0.0, 1.0],   # cat
68    [1.0, 1.0, 1.0, 0.0],   # sat
69    [0.0, 0.0, 1.0, 1.0],   # on
70    [1.0, 0.0, 0.0, 1.0],   # mat
71])
72
73K = torch.tensor([
74    [0.0, 1.0, 0.0, 1.0],   # The
75    [1.0, 0.0, 1.0, 0.0],   # cat
76    [1.0, 1.0, 0.0, 0.0],   # sat
77    [0.0, 0.0, 1.0, 1.0],   # on
78    [1.0, 0.0, 0.5, 0.5],   # mat
79])
80
81V = torch.tensor([
82    [1.0, 0.0, 0.0, 0.0],   # The
83    [0.0, 1.0, 0.0, 0.0],   # cat
84    [0.0, 0.0, 1.0, 0.0],   # sat
85    [0.0, 0.0, 0.0, 1.0],   # on
86    [0.5, 0.5, 0.5, 0.5],   # mat
87])
88
89# ── Run ──
90attn = ScaledDotProductAttention(d_k=4)
91weights, output = attn(Q, K, V)
92
93print("Attention Weight Matrix (5x5):")
94print(weights.round(decimals=4))
95
96print("\nOutput Matrix (5x4):")
97print(output.round(decimals=4))
98
99# ── Verify: Use PyTorch's built-in SDPA ──
100with torch.no_grad():
101    # F.scaled_dot_product_attention expects (batch, heads, N, d_k)
102    Q_b = Q.unsqueeze(0).unsqueeze(0)   # (1, 1, 5, 4)
103    K_b = K.unsqueeze(0).unsqueeze(0)
104    V_b = V.unsqueeze(0).unsqueeze(0)
105
106    builtin_out = F.scaled_dot_product_attention(Q_b, K_b, V_b)
107    builtin_out = builtin_out.squeeze(0).squeeze(0)  # back to (5, 4)
108
109print("\nBuilt-in SDPA output matches ours?",
110      torch.allclose(output, builtin_out, atol=1e-4))
111
112# ── GPU: just move tensors ──
113if torch.cuda.is_available():
114    Q_gpu, K_gpu, V_gpu = Q.cuda(), K.cuda(), V.cuda()
115    attn_gpu = attn.cuda()
116    w_gpu, o_gpu = attn_gpu(Q_gpu, K_gpu, V_gpu)
117    print("\nGPU output matches CPU?",
118          torch.allclose(output, o_gpu.cpu(), atol=1e-4))