Scaled Dot-Product Attention

The previous chapter showed how to compare embedding vectors for retrieval. Transformer layers do a related job inside a sequence: each token compares itself to other tokens, decides which positions matter, and blends information from those positions.

When you read "The package missed the dock because it was closed," you instantly know that "it" refers to the dock, not the package. Your brain routes focus to the relevant word. Scaled dot-product attention is the transformer operation that learns that kind of routing with matrix multiplication.

Attention is the core information-routing mechanism in modern transformers. Given a sequence of tokens, it answers: "For each token, which other positions should it use, and how much?" Once that routing idea is clear, the math and code become direct translations of it.

---

The attention mechanism

Intuition

Attention answers a simple question: "For each token in a sequence, which other tokens should it focus on?"

Think of reading a sentence: when you see "it" in "The package missed the dock because it was closed", your brain attends back to "dock" to resolve what "it" means. Self-attention does this computationally: every token computes a weighted combination of other token representations based on learned relevance.^[1]

Setup: Q, K, V projections

To perform attention, the model transforms each token representation into three different roles. Think of it like a support search system:

The model creates these three vectors by multiplying each token representation through learned weight matrices:

$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$

From one embedding to three views

Each token representation $X$ is multiplied by three different weight matrices ($W_Q$, $W_K$, $W_V$) to produce three different views of the same token. One view highlights what this token is searching for (Query), another highlights what this token offers to others (Key), and the third carries the actual information (Value).

Here $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}$ are the learned projections (where $d_{\text{model}}$ is the model's overall hidden dimension size, and $d_k, d_v$ are the dimensions of the queries/keys and values respectively).^[1]

Q and K determine where to look. V determines what information to extract. The separation of "routing" (Q, K) from "content" (V) is fundamental to why attention is so powerful.

---

The core formula

The scaled dot-product attention formula computes a weighted combination of value vectors based on the compatibility between queries and keys:^[1]

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

When some key positions aren't allowed, add a mask matrix $M$ after scaling. An allowed location has $M_{ij}=0$; a blocked location has $M_{ij}=-\infty$:

$\text{Attention}(Q, K, V; M) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V$

Step by step

The shape flow is compact enough to keep beside the formula. Queries and keys build one routing matrix; values join only after softmax:

A trace with real numbers

Before running PyTorch code, walk through one step with actual vectors. Imagine a two-word sequence, "order delayed." Each word lives in a 2-dimensional space for this toy example, and we have learned tiny projection matrices that give:

Step 1: raw scores

Compute the dot product of every query with every key:

$$S = QK^T = \begin{bmatrix} 1.0 \cdot 0.8 + 0.5 \cdot 0.2 & 1.0 \cdot 0.3 + 0.5 \cdot 0.9 \\ 0.5 \cdot 0.8 + 1.0 \cdot 0.2 & 0.5 \cdot 0.3 + 1.0 \cdot 0.9 \end{bmatrix} = \begin{bmatrix} 0.90 & 0.75 \\ 0.60 & 1.05 \end{bmatrix}$$

Step 2: scale

With $d_k = 2$, divide by $\sqrt{2} \approx 1.414$:

$$L = \begin{bmatrix} 0.64 & 0.53 \\ 0.42 & 0.74 \end{bmatrix}$$

Step 3: softmax

Normalize each row so it sums to 1:

$$\alpha = \begin{bmatrix} 0.53 & 0.47 \\ 0.42 & 0.58 \end{bmatrix}$$

Step 4: weighted values

Multiply the weights by the value vectors:

• New "order" vector = $0.53 \cdot [2.0, 1.0] + 0.47 \cdot [1.0, 2.0] = [1.53, 1.47]$
• New "delayed" vector = $0.42 \cdot [2.0, 1.0] + 0.58 \cdot [1.0, 2.0] = [1.42, 1.58]$

The output replaces each original embedding with a blend of the whole sequence, weighted by relevance. Even in this toy example, "order" pulls slightly more from its own value (0.53) than from "delayed" (0.47), while "delayed" mixes both values nearly evenly. In a real model with 512 or 4096 dimensions, this blending happens across thousands of numbers at once.

The same arithmetic is easy to verify without a tensor library:

PyTorch implementation

Here's how the scaled dot-product attention formula translates into practical PyTorch code. The function takes the Query, Key, and Value tensors along with an optional mask (used for causal or padding purposes). It computes the attention scores, scales them by the square root of the dimension, normalizes them into probabilities, and returns both the final weighted output and the attention weights themselves.

Two masking details matter in real code. First, softmax must run along the key dimension (dim=-1) so each query row sums to 1; normalizing along the query axis silently changes the operation. Second, a row with no permitted key has no valid attention distribution. -inf makes that mistake visible as NaN. Replacing it with a finite negative value hides the bug: softmax assigns weight to blocked keys because every blocked logit ties. Ensure each active query has at least one allowed key (causal attention includes its own position), or explicitly suppress outputs for padded query rows.

This short failure case makes the second rule concrete:

Common shape mistakes

The most frequent debugging task in transformer work is tracing a dimension mismatch. Here are two typical mistakes and the exact symptoms they produce.

Forgetting to transpose K

If you write torch.matmul(Q, K) instead of torch.matmul(Q, K.transpose(-2, -1)), the inner dimensions won't align. With Q shape (B, h, n, d_k) and K shape (B, h, n, d_k), the operation attempts to multiply two arrays whose final matrix axes don't form $(n \times d_k)(d_k \times n)$, so it fails before softmax.

The fix is always the same: the last two dimensions of K must be swapped so the matrix multiplication becomes $(n \times d_k) \cdot (d_k \times n)$, yielding the $(n \times n)$ score matrix.

Forgetting to scale

If you skip attn_scores / math.sqrt(d_k), the code won't crash. Under the independent unit-variance setup below, a $d_k=64$ raw dot product has standard deviation 8 rather than 1. That larger spread can saturate softmax and reduce routing gradients. Inspect score statistics and attention entropy when debugging, then restore the scale factor unless you're intentionally testing a different attention formulation.

---

Why scale by $\sqrt{d_k}$? The variance proof

Imagine combining many warehouse sensor signals into one routing score. Without scaling, adding more signals (higher $d_k$) makes raw dot products larger and larger until softmax saturates. Dividing by $\sqrt{d_k}$ normalizes the score so the model can still distinguish several plausible routes instead of locking onto one too early.

Here's the derivation that motivated the transformer's scale factor. For this calculation, assume entries of $q$ and $k$ are independent with mean 0 and variance 1. Learned activations won't satisfy those assumptions exactly, but the calculation explains why unscaled logits begin with dimension-dependent spread.^[1]

Assume

$q_i, k_i$ are independent components, each with mean $0$ and variance $1$. The proof needs only those two moments, not a specific distribution, which is why it holds for any reasonable initialization. The dot product is:

$q \cdot k = \sum_{i=1}^{d_k} q_i k_i$

Each term $q_i k_i$ has $\mathbb{E}[q_i k_i] = 0$ (since $\mathbb{E}[q_i] = \mathbb{E}[k_i] = 0$ and they're independent). The variance is:

$\text{Var}(q_i k_i) = \mathbb{E}[(q_i k_i)^2] - (\mathbb{E}[q_i k_i])^2 = \mathbb{E}[q_i^2]\mathbb{E}[k_i^2] - 0 = 1 \cdot 1 = 1$

By the sum of independent variances:

$\text{Var}(q \cdot k) = d_k$

So the standard deviation of the raw dot product is $\sqrt{d_k}$. As $d_k$ grows under this model, logits spread more widely before softmax:

After dividing by $\sqrt{d_k}$

$\text{Var}\!\left(\frac{q \cdot k}{\sqrt{d_k}}\right) = \frac{d_k}{d_k} = 1$ under these assumptions. Scaling doesn't promise a particular learned attention pattern; it removes a predictable source of width-dependent logit growth.

See it in code. This experiment samples independent unit-variance query/key vectors and checks the standard-deviation calculation rather than choosing one dramatic softmax row:

The sampled values won't be exactly the theoretical values, but their trend should match: raw spread grows with width while scaled spread stays near one.

---

Three types of attention

The transformer architecture uses attention in three distinct patterns. Bidirectional attention is like a support dashboard where every event in an order timeline can see every other event. Causal attention processes that timeline left to right, where each new event can only use earlier events. Cross-attention is like a reply generator reading from a separate policy document while writing the customer response.

1. Bidirectional self-attention (encoder)

Every non-padding token may attend to every other non-padding token; there is no future-token mask. Encoder architectures such as BERT use this pattern for tasks where the full input is available. For example, given a complete sentence, every word's representation can use every other word:

2. Causal self-attention (decoder)

Each token can only attend to itself and previous tokens. Future positions are masked with $-\infty$. Used in decoder architectures such as GPT-style and other autoregressive language models for generation tasks where the model must predict the next token without seeing the future. For instance, when processing a sequence step-by-step, the model progressively builds context but remains strictly blind to upcoming words:

The causal mask is a lower-triangular matrix that ensures each position only looks at itself and the positions before it. This minimal Python version takes the sequence length as input and outputs a boolean matrix where True indicates an allowed connection and False indicates a masked one.

3. Cross-attention (encoder-decoder)

Queries come from one sequence, Keys and Values from another. The original Transformer uses this pattern in its decoder: encoder outputs provide keys and values, while current decoder states provide queries.^[1] The resulting weights choose which source positions contribute to each target representation:

In practice, cross-attention usually applies a source padding mask so decoder tokens don't attend to padded encoder positions. It doesn't use a causal mask over the source sequence, because the encoder has already seen the whole input.

Self-attention produces a square query-by-key matrix. Cross-attention doesn't have to: two target queries reading three source positions produce a 2 x 3 routing matrix.

---

Multi-head attention

Imagine giving several analysts separate learned views of an order timeline, then letting a final projection combine their reports. Multi-head attention makes that capacity available: each head has its own query, key, and value projections. A trained head may acquire a recognizable routing pattern, but the architecture doesn't assign jobs such as "policy clause head" in advance.

Instead of one large attention operation, we run $h$ parallel heads, each with $d_k = d_{\text{model}} / h$:^[1]

$\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O$

Reading the formula

Instead of running one big attention operation, we project into $h$ narrower heads that attend independently. After each head produces its output, we concatenate them and multiply by a final matrix $W^O$ so information from those routes rejoins the residual stream.

Each head computes: $\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$

What do attention heads learn?

Interpretability results here are evidence about specific trained models, not a promise about every transformer. Voita et al. found positional and syntactic patterns among heads in neural machine translation encoders, and Michel et al. found that many heads in the models they tested could be removed at inference with limited quality loss.^[2][3] Olsson et al. studied induction heads, circuits that support copying patterns in autoregressive models under their experiments.^[4]

Critical nuance: same asymptotic attention FLOPs

Multi-head attention doesn't increase the asymptotic cost of the attention core when you keep $d_{\text{model}}$ fixed. It restructures the work. Single-head attention on $d_{\text{model}}=512$ uses roughly the same leading-order FLOPs (Floating Point Operations) for score computation and value mixing as 8-head attention with $d_k=64$ each, because $h \times d_k = d_{\text{model}}$. The dense Q/K/V and output projections still cost $O(n d_{\text{model}}^2)$ either way.

The arithmetic also makes the compute comparison testable. Splitting a fixed width into more heads doesn't change the total number of score-and-value multiply-adds in the attention core:

Here's a runnable PyTorch shape implementation. To keep shapes simple, it uses the common case $d_v = d_k = d_{\text{model}} / h$. The three dense projection layers each produce all head slices at once; reshaping exposes those slices to batched attention.

From-scratch checklist

Before relying on torch.nn.MultiheadAttention, make sure you can implement the pieces above by hand. High-level modules are useful later, but they hide the exact shape and masking mistakes that break production attention code.

For a decoder-only block, your implementation should do each step explicitly:

1. Project x into Q, K, and V.
2. Reshape (B, N, D) into (B, h, N, d_k).
3. Compute Q @ K.transpose(-2, -1) / sqrt(d_k).
4. Apply the causal mask before softmax.
5. Use numerically stable softmax along the key dimension.
6. Multiply attention weights by V.
7. Concatenate heads back to (B, N, D).
8. Apply the output projection.

Two shape assertions catch many bugs:

One mask assertion catches leakage: assert not causal[0][1]. Token 0 must not see token 1.

If a model can read future tokens during training, the loss can look excellent while generation fails. That's why causal masking isn't a detail. It's the contract that makes next-token prediction honest.

Gradient flow through attention

Attention has three gradient paths:

The scale factor helps keep the Q/K path trainable. If scores become too large, softmax saturates, attention weights become nearly one-hot, and the gradient through the routing path becomes tiny. If the causal mask is wrong, gradients flow through illegal future positions and the model learns a shortcut it can't use at inference time.

When debugging attention, don't only print the final output. Inspect the score range, the mask, one row of attention weights, and the gradient norm on W_q and W_k. Those four checks tell you whether the model is learning routing or only moving values through a broken router.

---

Complexity analysis

Concrete memory example

Batch of 8, 32 heads, $n = 8192$, FP16 (16-bit floating point format):

$\text{Attention matrix} = 8 \times 32 \times 8192^2 \times 2 \text{ bytes} \approx 34.4 \text{ GB} \approx 32 \text{ GiB}$

How the numbers work

8 sequences in the batch x 32 attention heads x 8192² entries per attention map x 2 bytes per number (FP16) = about 34.4 GB of raw storage, or about 32 GiB, just for one attention-score tensor before values needed for backpropagation, model weights, or optimizer state.

In the naive formulation, this is why $O(n^2)$ temporary memory becomes a bottleneck for long sequences. Doubling sequence length makes one materialized score tensor four times larger. Fused kernels can avoid storing that full tensor, but exact dense attention still computes interactions across all query-key pairs.

Training vs. inference bottlenecks

During training, or in a naive implementation, the temporary $n \times n$ score matrix is the obvious $O(n^2)$ memory problem. During autoregressive decoding, optimized kernels often avoid materializing that matrix, but the model still has to repeatedly read the accumulated KV cache (past keys and values) for every new token. That makes incremental inference heavily constrained by memory bandwidth, not just FLOPs.^[5]

Architectural variants attack that persistent cache cost directly. Multi-query attention (MQA) shares one key/value head across all query heads, while grouped-query attention (GQA) uses a smaller number of key/value heads than query heads.^[5][6] Both shrink KV-cache bytes. Whether they improve latency enough for a workload is a measurement question because kernel choice, batch size, and quality requirements also matter.

For a simplified decoder cache, the storage count is proportional to layers x tokens x kv_heads x head_dim x 2 (the final factor stores both K and V). Keeping 32 query heads but reducing KV heads changes this count directly:

FlashAttention and MQA/GQA solve different problems. FlashAttention cuts temporary attention I/O. MQA/GQA cut persistent KV-cache size.

---

FlashAttention: tiled exact attention

FlashAttention (Dao et al., 2022)^[7] computes exact dense attention with an IO-aware algorithm. Its attention working memory is linear in sequence length rather than storing a quadratic score matrix; it changes the execution order, not the attention definition.

Core idea

Instead of materializing the full $n \times n$ attention matrix in GPU HBM, FlashAttention computes attention in tiles that fit in on-chip SRAM, using online softmax to avoid storing the full matrix while preserving the exact result.^[7][8]

FlashAttention is IO-aware. It minimizes traffic between large off-chip HBM and small on-chip SRAM by restructuring the computation order, not by changing the mathematical result.

Using fused attention in practice doesn't require writing custom CUDA kernels. PyTorch exposes scaled dot-product attention; the backend chosen for a given run depends on device, dtype, shapes, masks, and framework version. On CUDA, the function attempts to select an enabled implementation based on its inputs, but a fused kernel isn't guaranteed for every call.^[9] The public contract is the result, so start by checking it against the explicit computation:

Two API details prevent quiet bugs when you replace explicit attention with PyTorch's fused primitive. For F.scaled_dot_product_attention, True in a boolean attn_mask means the position participates in attention; that's the inverse of nn.MultiheadAttention's boolean key_padding_mask. Also, dropout_p is always applied when it's greater than zero, so pass 0.0 during evaluation.^[9]

---

Softmax numerical stability

A subtle but critical implementation detail in attention mechanisms is numerical stability. The naive mathematical formulation of softmax overflows for large inputs because the exponential function $e^{x}$ grows exceptionally fast. When dealing with dot products from large vector spaces, even scaled ones can occasionally produce large positive values, causing the exponentiation to result in Inf (infinity) floating-point values and completely breaking the model's training process.

The max-shift technique

To prevent this overflow, deep learning frameworks implement a numerically stable version of softmax. They subtract the maximum value from the input vector before exponentiating:

$\text{softmax}(x_i) = \frac{e^{x_i - \max(\mathbf{x})}}{\sum_j e^{x_j - \max(\mathbf{x})}}$

Why shifting doesn't change the answer

Subtracting $\max(\mathbf{x})$ shifts all logits by the same constant. Because of the properties of exponentials, this constant factors out and cancels between the numerator and denominator. The resulting probabilities don't change, but the highest value being exponentiated is exactly $0$ (since $\max(\mathbf{x}) - \max(\mathbf{x}) = 0$), so the maximum exponential evaluated is $e^0 = 1$. This keeps all numbers numerically safe and avoids overflow.^[8]

Building on this, online softmax (Milakov & Gimelshein, 2018)^[8] extends the max-shift technique to compute softmax in a single streaming pass. Instead of needing to read the entire vector into memory to find the maximum, online softmax tracks the running maximum and the running sum of exponentials simultaneously, dynamically correcting the sum as a new, larger maximum is found. This streaming capability is exactly the mathematical foundation that allows FlashAttention's memory-efficient block tiling.

---

Attention in modern architectures

The mask shape and the source of Q, K, and V reveal which information an architecture permits each output to use:

Encoder-only models like BERT and ViT use bidirectional self-attention because their goal is to build a comprehensive understanding of the entire input at once. They process the full text or image simultaneously to generate rich embeddings. In contrast, decoder-only language models use causal self-attention because they must generate output step-by-step. If they could look ahead at future tokens during training, they would simply "cheat" instead of learning to predict.

Encoder-decoder models (like T5 and Whisper) mix these approaches. They use a bidirectional encoder to fully understand the input (text or audio), and a causal decoder to generate the output. Cross-attention acts as the bridge between the two, allowing the decoder to continuously refer back to the encoder's rich representation of the source material.

Decoder-only language models exercise causal self-attention; encoder-decoder models add cross-attention; ViT uses bidirectional attention over image tokens. Those mechanics matter more than a market survey here. The next chapter will reuse the same attention operation after turning an image into patch tokens.

---

What you should be able to defend

Mistakes that break attention

Going deeper

"Why not use additive attention instead of dot-product?"

Bahdanau et al. (2015)^[10] used additive attention: $\text{score}(q, k) = v^T \tanh(W_q q + W_k k)$. In the Transformer paper, Vaswani et al. note that additive and dot-product attention behave similarly at small dimensions, but dot-product maps much better to batched matrix multiplication and is much faster at the larger dimensions used in transformers.^[1] The $\frac{1}{\sqrt{d_k}}$ scaling is what keeps dot-product attention stable as $d_k$ grows.

"Can attention attend to nothing / everything equally?"

If all scores are equal, softmax produces a uniform distribution $1/n$, and the output is the average of all value vectors. If one score dominates, softmax approximates an argmax, so the output is approximately one value vector. The sharpness of the logits controls where on this spectrum you land.