Scaled Dot-Product Attention

Retrieval compares embedding vectors. 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 parser rejected the patch because it was malformed," you instantly know that "it" refers to the patch, not the parser. 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?"

When you read a sentence, "it" in "The parser rejected the patch because it was malformed" points back to "patch." Self-attention does this computationally: every token computes a weighted combination of other token representations based on learned relevance.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

Setup: Q, K, V projections

To perform attention, the model transforms each token representation into three different roles:

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 represents what this token is searching for (Query), another represents 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]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

Q and K determine where to look. Their dot products produce routing scores, and softmax converts those scores into weights. V determines what information to extract. Separating "routing" (Q, K) from "content" (V) is what lets attention route information flexibly.

---

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]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} Softmax turns each query row of scores into non-negative weights that sum to 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 logit gets $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. Use a two-token sequence, "cache grows." Each token 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 "cache" vector = $0.53 \cdot [2.0, 1.0] + 0.47 \cdot [1.0, 2.0] = [1.53, 1.47]$
• New "grows" 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, "cache" pulls slightly more from its own value (0.53) than from "grows" (0.47), while "grows" 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

The scaled dot-product attention formula translates directly into 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. Two typical mistakes produce these exact symptoms.

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 repair is always the same: swap the last two dimensions of K 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

Each attention score adds many independent feature products into one logit. Without scaling, increasing $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 distribute probability across several plausible tokens instead of locking onto one too early.

The derivation that motivated the transformer's scale factor starts with a simple assumption: entries of $q$ and $k$ are independent across vectors and dimensions, 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]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

Assume

$q_i, k_i$ are independent components, each with mean $0$ and variance $1$, and the product terms are independent across dimensions. The calculation doesn't require a Gaussian distribution; it uses these moments and independence to add the per-dimension variances. 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 lets every position read every other visible position. Causal attention processes a sequence left to right, where each new token can only use earlier tokens. Cross-attention lets a target sequence read from a separate source sequence, such as a decoder reading encoder states.

1. Bidirectional self-attention (encoder)

Every non-padding token may attend to every other non-padding token; there's 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 lets each position look only 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]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} 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

Multi-head attention gives the model several learned routing views of the same token sequence, then lets a final projection combine those views. 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 "syntax head" or "copy head" in advance.

Instead of one large attention operation, we run $h$ parallel heads, each with $d_k = d_{\text{model}} / h$:^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

$\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.

For self-attention with input $X$, each head computes: $\text{head}_i = \text{Attention}(XW_i^Q, XW_i^K, XW_i^V)$. Cross-attention uses target states for the query projection and source states for the key/value projections.

What do attention heads learn?

These interpretability results 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]Reference 2Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting.https://arxiv.org/abs/1905.09418[3]Reference 3Are Sixteen Heads Really Better than One?.https://arxiv.org/abs/1905.10650} Olsson et al. studied induction heads, circuits that support copying patterns in autoregressive models under their experiments.^{[4]Reference 4In-context Learning and Induction Heads.https://arxiv.org/abs/2209.11895}

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

Big-O notation describes how work or storage grows as sequence length $n$ and model width $d_{\text{model}}$ increase. It hides constant factors, so use it for growth trends and use concrete byte/FLOP arithmetic for capacity planning.

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, $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 FLOPs alone.^{[5]Reference 5Fast Transformer Decoding: One Write-Head is All You Need.https://arxiv.org/abs/1911.02150}

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]Reference 5Fast Transformer Decoding: One Write-Head is All You Need.https://arxiv.org/abs/1911.02150[6]Reference 6GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245} 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]Reference 7FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135} 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]Reference 7FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135[8]Reference 8Online normalizer calculation for softmax.https://arxiv.org/abs/1805.02867}

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]Reference 9torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html} 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]Reference 9torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html}

---

Softmax numerical stability

Naive softmax can overflow. The exponential $e^x$ grows quickly, so a large positive logit can become Inf; dividing Inf by a sum containing Inf can produce NaN. Attention kernels use stable softmax so large but valid scores don't corrupt routing weights.

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 prevents positive exponential overflow; very negative shifted values may underflow harmlessly toward zero.^{[8]Reference 8Online normalizer calculation for softmax.https://arxiv.org/abs/1805.02867}

Building on this, online softmax (Milakov & Gimelshein, 2018)^{[8]Reference 8Online normalizer calculation for softmax.https://arxiv.org/abs/1805.02867} 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 they read the entire input before producing representations. They process the full text or image simultaneously to generate contextual 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 "cheat" instead of learning to predict.

Encoder-decoder models (like T5 and Whisper) mix these approaches. They use a bidirectional encoder to read 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 states for 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 to check before moving on

Mistakes that break attention

Going deeper

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

Bahdanau et al. (2015)^{[10]Reference 10Neural Machine Translation by Jointly Learning to Align and Translate.https://arxiv.org/abs/1409.0473} 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]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} 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.

Practice drill

Build a tiny attention-debug notebook or trace for three tokens:

1. Print Q, K, score, mask, softmax row sums, and output shapes after every operation.
2. Add one padded key and one causal mask case, then verify blocked positions receive zero probability.
3. Trigger one wrong-axis softmax failure and record the symptom you would see in row sums or output behavior.
4. Write a checklist for reviewing production attention code: tensor shapes, mask semantics, stable softmax, and fully masked rows.

This makes routing errors observable before they become mysterious model-quality bugs.