Scaled Dot-Product Attention

The previous chapter scored stored embedding vectors with cosine and dot product. A retriever compares one query vector to many keys.

A transformer layer turns that same routing idea inward: each token supplies a query and a key, while a mask decides which positions it may read.

Take the line "The cache grew quickly because it was full." You know "it" is the cache, not the growth. If that reference reads from the wrong position, later layers receive the wrong context.

Before naming any matrices, predict the operation: each token should choose source positions, turn those scores into weights, and mix the chosen content. Scaled dot-product attention learns that routing with matrix multiplication.

The mask answers a separate question from relevance: which positions are legal to read? The $-\infty$ cells are a causal mask. Compute the unmasked two-token slice by hand first, then put the mask back.

Every active query still needs at least one allowed key. An entirely blocked row has no probability distribution, because subtracting its maximum evaluates $-\infty-(-\infty)$ and produces NaN.

---

Three views of one token

Self-attention computes a weighted mix of token representations. Each query gets its weights from learned relevance scores against the keys.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} To do that, the model turns each hidden state into three roles:

Each token hidden state $x_i$ is projected three ways. $W_Q$ produces what this position is searching for (query), and $W_K$ produces when others should read this position (key). $W_V$ produces the payload that gets mixed if selected (value):

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

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}$. $d_{\text{model}}$ is the residual width; $d_k$ and $d_v$ are the query/key and value widths.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

For the running example, token cache uses these toy vectors after projection:

Predict the split before reading the diagram: Q and K should meet to produce a routing score, while V should travel past that score and remain content. The self-score for cache will therefore use $q$ and $k$, not $v$.

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.

In one sentence, which tensors choose where attention goes, and which tensor carries content?

Answer

Q and K choose where attention goes. Their dot products produce routing scores that softmax converts into weights. V carries the content that gets mixed by those weights.

---

The core formula

Read the formula as a lookup followed by a weighted read. Compatibility of queries with keys becomes weights; those weights mix values.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} Softmax turns each query row into non-negative weights that sum to 1:

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^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 order is the contract: scores choose locations, softmax makes each choice a distribution, and only then do those weights touch V. Predict the shapes before following the picture: two queries compared with two keys make a $2 \times 2$ routing matrix, while the value width survives into the output.

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

Walk through the first two tokens, cache and grew, before adding the third token or a mask. Each lives in a 2-dimensional toy space:

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 "grew" vector = $0.42 \cdot [2.0, 1.0] + 0.58 \cdot [1.0, 2.0] = [1.42, 1.58]$

Each original vector is replaced by a blend of the sequence, weighted by relevance. Cache pulls slightly more from its own value (0.53) than from grew (0.47), while grew mixes both nearly evenly. A real head does this across 64 or 128 dimensions, not two.

In this toy trace, skipping the scaling step barely changes softmax. Why does the same omission hurt when $d_k$ grows to 64 or 512?

Answer

At $d_k = 2$, raw scores stay small. Under the initialization assumptions used below, increasing $d_k$ increases the dot-product spread like $\sqrt{d_k}$. That makes saturated, low-gradient softmax rows more likely unless scores are scaled.

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

Add the causal mask

The mask changes who may be read, not how the allowed weights are computed. With a lower-triangular mask, this function becomes the two-token slice of the opening figure: token cache can only read itself, while token grew can read cache and grew, so its weights match the unmasked walkthrough.

Two masking details matter in real code. First, softmax must run along the key axis 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.

Predict the failure before running this case: with a finite fill value, a fully blocked row gives every blocked key the same score, so softmax invents a uniform distribution. The suppression step must handle that row explicitly.

This short failure case makes the second rule concrete:

In PyTorch, F.scaled_dot_product_attention is the fused entry point for the same math. Current docs list three CUDA implementations (FlashAttention-2, memory-efficient attention, and a C++ math fallback) and an experimental enable_gqa path.

A fused kernel isn't guaranteed for every shape and dtype. Two API traps matter: True in a boolean attn_mask means the position participates, which is the inverse of nn.MultiheadAttention's boolean key_padding_mask; and dropout_p is applied whenever it's greater than zero, so pass 0.0 at eval.^{[2]Reference 2torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html[3]Reference 3FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.https://arxiv.org/abs/2307.08691}

Common shape mistakes

When attention breaks, start with two questions: do the dimensions produce one score per query-key pair, and did the score scale stay stable? The two mistakes below produce distinct symptoms.

Forgetting to transpose K

If you multiply $Q$ by $K$ instead of $K^T$, the inner dimensions won't align. With $Q$ shape $(B, h, n, d_k)$ and $K$ shape $(B, h, n, d_k)$, you need $(n \times d_k)(d_k \times n)$. In PyTorch that's Q @ K.transpose(-2, -1), not Q @ K.

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

The two-token trace hides the problem because its head is only two numbers wide. Each attention score adds one feature product per dimension 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.

The calculation still 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:

So at $d_k=64$, predict a raw standard deviation of 8 and a scaled standard deviation of 1. The table checks that prediction across widths.

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.

Why divide by $\sqrt{d_k}$ instead of by $d_k$?

Answer

The raw dot product has variance $d_k$, so its standard deviation is $\sqrt{d_k}$. Softmax sees the scale of logits through their standard deviation, so dividing by $\sqrt{d_k}$ brings the variance back to 1. Dividing by $d_k$ would over-shrink the logits as dimensions grow.

---

Three types of attention

The formula doesn't decide who may read whom. The mask and the sources of $Q$, $K$, and $V$ do.

Bidirectional attention lets every visible position read every other visible position. Causal attention processes left to right, so 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, and the Vision Transformer you'll meet next, use this pattern when the full input is already available:

2. Causal self-attention (decoder)

Each token can only attend to itself and previous tokens. Future positions are masked with $-\infty$. Decoder architectures such as GPT-style and other autoregressive language models use this pattern for generation tasks where the model must predict the next token without seeing the future.

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.

Why must token 0 be unable to attend to token 1 in causal self-attention?

Answer

During next-token training, token 0 is only allowed to use itself and earlier context. If it can attend to token 1, the model can leak the future answer during training. Loss can look good, but generation fails because that future token isn't available at inference time.

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. The formula describes $h$ parallel heads, each with $d_k = d_{\text{model}} / h$:^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

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.

$\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.^{[4]Reference 4Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting.https://arxiv.org/abs/1905.09418}

Michel et al. found that many heads in the models they tested could be removed at inference with limited quality loss.^{[5]Reference 5Are 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.^{[6]Reference 6In-context Learning and Induction Heads.https://arxiv.org/abs/2209.11895}

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.

Predict before running the arithmetic: with fixed $d_{\text{model}}$, one 512-wide head and eight 64-wide heads should do the same leading score-and-value work. Splitting a fixed width into more heads doesn't change the total number of score-and-value multiply-adds in the attention core:

The usual layout $d_v = d_k = d_{\text{model}} / h$ is a reshape, not extra width. After a dense projection to $D$, slice into $h$ heads and attend. Then concatenate back to $D$ before $W^O$:

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

There are two bottlenecks to keep separate: quadratic query-key pairs and width-heavy projections. 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, 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.

If sequence length grows from 512 to 2048, why does naive attention memory grow by 16x instead of 4x?

Answer

The attention score matrix is $n \times n$. Increasing $n$ by 4x makes the matrix $(4n) \times (4n)$, which has 16x as many entries. This is the core quadratic memory wall.

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. Each new token has only one query row, but the model still has to repeatedly read the accumulated KV cache (past keys and values). That makes incremental inference heavily constrained by memory bandwidth, not FLOPs alone.^{[7]Reference 7Fast 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.^{[7]Reference 7Fast Transformer Decoding: One Write-Head is All You Need.https://arxiv.org/abs/1911.02150}

Grouped-query attention (GQA) uses fewer key/value heads than query heads.^{[8]Reference 8GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245}

Both shrink KV-cache bytes. The Llama 3 paper uses GQA with 8 key-value heads.^{[9]Reference 9The Llama 3 Herd of Models.https://arxiv.org/abs/2407.21783} Whether that improves 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, while MQA/GQA cut persistent KV-cache size.

Multi-head latent attention (MLA) compresses KV state further, as in DeepSeek-V2. That's a later serving layout, not a change to $\mathrm{softmax}(QK^T/\sqrt{d_k})V$.^{[10]Reference 10DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Modelhttps://arxiv.org/abs/2405.04434}

Which optimization reduces temporary attention-matrix I/O, and which reduces persistent KV-cache bandwidth during decoding?

Answer

FlashAttention reduces temporary attention-matrix I/O by tiling attention and avoiding materializing the full $n \times n$ matrix. MQA and GQA reduce persistent KV-cache bandwidth by sharing or grouping key/value heads during autoregressive decoding.

---

Softmax numerical stability

The two-token walkthrough already subtracted the row max before exp. That isn't a style choice.

Naive softmax overflows: $e^x$ grows so fast that a large positive logit becomes Inf, and Inf / Inf is NaN. Attention kernels need a stable form so large but valid scores don't corrupt routing weights.

The max-shift technique

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})}}$

Before running the example, predict what stays invariant: shifting $[1000, 1001, 999]$ by its maximum changes the numbers sent to exp, not the relative probabilities. The largest shifted logit should receive the largest weight, and the row should still sum to 1.

Subtracting $\max(\mathbf{x})$ shifts every logit by the same constant. That constant factors out of the exponentials and cancels between numerator and denominator, so the probabilities don't change.

The largest exponent is $e^0 = 1$, which blocks positive overflow. Very negative shifted values may underflow toward zero, which is harmless.^{[11]Reference 11Online normalizer calculation for softmax.https://arxiv.org/abs/1805.02867}

Online softmax (Milakov and Gimelshein, 2018) extends the same identity to a streaming pass. It tracks a running maximum and a running sum of exponentials, then rescales the sum when a larger maximum appears. That's the algebra FlashAttention uses to tile a query row without storing every score.^{[11]Reference 11Online normalizer calculation for softmax.https://arxiv.org/abs/1805.02867}

---

FlashAttention: tiled exact attention

The $n \times n$ score matrix is the memory problem. If it's temporary, the useful question is whether we need to write the whole thing to HBM at all.

FlashAttention (Dao et al., 2022) keeps the definition exact and changes the execution order. Auxiliary attention state is linear in sequence length; the algorithm still does quadratic work.^{[12]Reference 12FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135}

Instead of writing the full score matrix to GPU HBM, it computes attention in tiles that fit in on-chip SRAM and uses online softmax so each tile can update a running row max, sum, and output accumulator.^{[12]Reference 12FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135[11]Reference 11Online normalizer calculation for softmax.https://arxiv.org/abs/1805.02867}

FlashAttention-2 keeps that math and changes how work is partitioned across GPU thread blocks.^{[3]Reference 3FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.https://arxiv.org/abs/2307.08691} FlashAttention-3 adds asynchrony and low-precision paths on Hopper-class GPUs.^{[13]Reference 13FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision.https://arxiv.org/abs/2407.08608}

Later kernels still implement $\mathrm{softmax}(QK^T/\sqrt{d_k})V$; they don't introduce a new attention formula.

You don't write a custom CUDA kernel to use this in PyTorch. F.scaled_dot_product_attention dispatches among FlashAttention-2, memory-efficient attention, and a C++ math fallback when the CUDA backend can. It may enable grouped-query attention through enable_gqa.

The public contract is the numeric result, not which kernel ran.^{[2]Reference 2torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html[3]Reference 3FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.https://arxiv.org/abs/2307.08691}

---

Attention in modern architectures

Once Q/K/V roles and mask semantics are separate, architecture differences become a small lookup: where do Q, K, and V come from, and which score cells remain visible?

BERT and ViT use bidirectional self-attention because the full input is already there. Decoder-only language models use causal self-attention because they generate left to right; looking at future tokens during training would leak the answer.

Encoder-decoder models such as T5 and Whisper mix both: a bidirectional encoder reads the source, a causal decoder writes the target, and cross-attention is the bridge.

ViT is the same attention operation after a different tokenizer. An image becomes patch tokens, then bidirectional self-attention mixes those patches with no causal mask. That's the next chapter.

---

Derivation checklist

What happens if you remove the scaling factor for a large head dimension such as $d_k=512$?

Answer

Under the independent unit-variance calculation, the variance of $q \cdot k$ grows with $d_k$, so its standard deviation grows as $\sqrt{d_k}$. At $d_k=512$, unscaled logits are much more likely to saturate softmax. That can weaken gradients through Q and K and make routing behave like a brittle hard lookup rather than a smooth weighted average.

Why use multiple attention heads instead of one large head?

Answer

Multiple heads give the model several separately projected routing spaces in parallel. Studies have measured recognizable patterns in some trained heads, but you must inspect a particular model before naming what its heads do. With fixed $d_{\text{model}}$, each head is narrower, so the leading attention-core FLOPs stay in the same asymptotic class.

How can transformer systems handle very long contexts such as 100K tokens?

Answer

There isn't one trick. FlashAttention reduces temporary memory and HBM traffic for exact attention, but exact attention still has quadratic work. MQA and GQA shrink KV-cache bandwidth during decoding. Long-context systems may also use sliding-window attention, distributed cache placement, retrieval, or linear/state-space alternatives when they accept a different quality-efficiency tradeoff.

Mistakes that break attention

Going deeper

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

Bahdanau et al. (2015)^{[14]Reference 14Neural 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.