Learn: Scaled Dot-Product Attention

What you're learning

The single operation the entire Transformer is built on. By the end you’ll be able to write scaled dot-product attention from memory, explain why it’s scaled by $1/\sqrt{d_k}$, and compute its FLOPs and $O(s^2)$ cost. Everything in this curriculum is a variation on what you learn here.

First lesson of the Attention Mechanisms curriculum. Vocabulary follows the attention-glossary.

---

1. Learning map

graph TD
    A["Problem: mix information<br/>across a sequence,<br/>content-dependently"] --> B["Idea: soft key–value lookup"]
    B --> C["Q, K, V projections"]
    C --> D["scores = QKᵀ"]
    D --> E["scale by 1/√d_k"]
    E --> F["softmax → weights"]
    F --> G["weighted sum of V"]
    G --> H["output"]
    D --> I["cost: O(s²·d)<br/>the quadratic wall"]
    style B fill:#44a,color:#fff
    style I fill:#a44,color:#fff

---

2. Intuition: attention is a soft dictionary lookup

Forget matrices for a moment. A Python dict is a hard lookup: you have a query key, you find the one entry whose key matches exactly, and you return its value. Attention is the soft, differentiable version of that:

• Every token emits a query (“what am I looking for?”),
• every token also exposes a key (“what do I offer?”) and a value (“what I’ll actually hand over”),
• a query compares itself against all keys, producing a similarity score for each,
• those scores become weights (they sum to 1), and the output is the weighted average of all values.

So instead of retrieving one value, each token retrieves a blend of everyone’s values, weighted by how well its query matches each key. It’s content-based, learned, and — crucially — differentiable, so the network learns what to look for.

The one-sentence version

Attention lets every position pull in information from every other position, with the amount pulled decided by learned query–key similarity. It is the mechanism for content-dependent, long-range information routing.

A little history. The idea entered deep learning as additive attention for machine translation (Bahdanau et al., 2014), where an MLP scored a decoder query against encoder states.¹ The Transformer replaced that MLP with a plain dot product (cheaper, parallelizable) and made attention the entire model rather than an add-on to an RNN.²

---

3. The math

Pack the sequence’s queries, keys, and values into matrices. For $s$ tokens, with query/key dimension $d_k$ and value dimension $d_v$:

$Q\in {\mathbb{R}}^{s\times d_k},K\in {\mathbb{R}}^{s\times d_k},V\in {\mathbb{R}}^{s\times d_v}$

obtained by linear projections of the input $X\in {\mathbb{R}}^{s\times d}$: $Q=X{W}^Q$, $K=X{W}^K$, $V=X{W}^V$.

Scores. Every query dotted with every key:

$S=Q{K}^{\top }\in {\mathbb{R}}^{s\times s},S_{ij}=q_i\cdot k_j$

$S_{ij}$ measures how much token $i$‘s query matches token $j$‘s key.

Scale, then softmax row-wise to get weights that sum to 1 across keys:

$A=\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)\in {\mathbb{R}}^{s\times s}$

Output = weighted sum of values:

$\boxed{\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V\in {\mathbb{R}}^{s\times d_v}}$

Why divide by $\sqrt{d_k}$?

Suppose the entries of $q$ and $k$ are independent with mean 0 and variance 1. Their dot product $q\cdot k={\sum }_{l=1}^{d_k}{q}_l{k}_l$ is a sum of $d_k$ independent zero-mean, unit-variance terms, so it has variance $d_k$ and standard deviation $\sqrt{d_k}$.² As $d_k$ grows (say 128), raw scores swing by $\pm \sqrt{128}\approx 11$. Feed scores that large into softmax and it saturates: almost all mass lands on one key, and the gradient through the near-flat softmax tails $\to 0$. Dividing by $\sqrt{d_k}$ renormalizes the scores back to unit variance, keeping softmax in its responsive, high-gradient regime. It’s a variance-control trick, nothing more — but training breaks without it.

Causal masking (for autoregressive LMs)

In a decoder LM, token $i$ must not see the future. Before softmax, set $S_{ij}=-\infty$ for all $j>i$ (a lower-triangular mask); softmax then assigns those positions zero weight. This mask is what makes the same attention op “causal” — and, later, what lets you cache past keys/values (KV cache).

Quiz (retrieval): From memory, write the full SDPA formula and state in one line why the $1/\sqrt{d_k}$ factor is there.

Answer

$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(Q{K}^{\top }/\sqrt{d_k})V$. The scaling is there because $q\cdot k$ has variance $d_k$; without dividing by $\sqrt{d_k}$ the scores grow with dimension, softmax saturates, and gradients vanish.

---

4. Illustration: a 3-token pass

flowchart LR
    X["X (s×d)"] -->|W^Q| Q["Q (s×d_k)"]
    X -->|W^K| K["K (s×d_k)"]
    X -->|W^V| V["V (s×d_v)"]
    Q --> S["S = QKᵀ (s×s)"]
    K --> S
    S -->|"÷√d_k, softmax"| A["A (s×s)"]
    A --> O["O = A·V (s×d_v)"]
    V --> O

Concretely, with $s=3$: $S$ is $3\times 3$; row $i$ of $A$ is a probability distribution over the 3 tokens; output row $i$ is $A_{i1}{v}_1+A_{i2}{v}_2+A_{i3}{v}_3$ — token $i$‘s personalized blend of all values. If $A_{i2}$ is large, token $i$ is “attending to” token 2.

---

5. FLOPs & complexity

Count multiply-adds as 2 FLOPs. For batch $b$, sequence $s$, model dim $d$ (take $d_k=d_v=d$ for a single-head view):

Step	Shape	FLOPs
Scores $Q{K}^{\top }$	$(s\times d)(d\times s)$	$2b{s}^{2}d$
Softmax	$s\times s$	$\approx 3b{s}^2$ (small)
$A\cdot V$	$(s\times s)(s\times d)$	$2b{s}^{2}d$

So the attention core is $\approx 4b{s}^{2}d$ FLOPs, and — the headline — it is quadratic in sequence length $s$. Double the context, quadruple the attention compute (and the $s\times s$ score matrix’s memory). The projections $Q,K,V,O$ add $O(bs{d}^2)$, which dominates at short $s$ but is overtaken by the $s^2$ term as context grows. This single $O(s^2)$ fact is the reason Tiers 2–4 of this curriculum exist.

Try it: plug numbers into the calculator

Open the interactive Attention FLOPs & KV-cache calculator and set $d=4096$, one head, $b=1$. Sweep $s$ from 2k → 8k → 32k and watch the prefill FLOPs grow ~quadratically while the KV cache grows linearly. This is the intuition you’ll formalize in 04-quadratic-bottleneck-and-io-aware.

Quiz (retrieval): You increase context length from 4k to 16k. By what factor does the $Q{K}^{\top }$ term grow, and by what factor does the KV cache grow?

Answer

$Q{K}^{\top }$ scales as $s^2$, so $4\times$ longer → 16× the score FLOPs. The KV cache scales as $s$ (it stores one K and one V per token), so 4×. Attention compute explodes faster than attention memory as context grows — but at inference the linear-in-$s$ KV cache is usually the binding constraint (lesson 05, and 12-llm-inference-systems).

---

6. Practice

Problem: single-head attention cost

A layer processes $b=2$ sequences of $s=1024$ tokens with $d=1024$ (single head). Estimate the FLOPs of the attention core (scores + A·V), ignoring projections and softmax.

Solution. Core $\approx 4b{s}^{2}d=4\cdot 2\cdot 1024^2\cdot 1024=4\cdot 2\cdot 2^{20}\cdot 2^{10}=8\cdot 2^{30}\approx 8.6\times 10^9$ FLOPs $=8.6$ GFLOP. (Check against the calculator with $n\_heads=1$.)

---

Read next / primary source

The primary source is Vaswani et al., “Attention Is All You Need” (2017)² — read §3.2 (scaled dot-product & multi-head attention). For the pre-Transformer origin, skim Bahdanau et al. (2014)¹.

Ask your teacher

Stuck on the variance argument, the masking, or the FLOPs count? Ask — I’m your teacher for this course and can go deeper on any piece, or make you a harder practice problem.

What’s next

• 02-multi-head-attention-and-the-block — run many of these in parallel and wrap them in the full Transformer block.

---

References

Reference: attention-glossary | Topic: attention-mechanisms | Filed: 2026-09-10

Footnotes

1. Bahdanau, Cho, Bengio, “Neural Machine Translation by Jointly Learning to Align and Translate” — arXiv:1409.0473 (2014). The additive (MLP-scored) attention that preceded dot-product attention. ↩ ↩²

2. Vaswani et al., “Attention Is All You Need” — arXiv:1706.03762 (2017). Introduces scaled dot-product attention (§3.2.1), the $1/\sqrt{d_k}$ scaling and its variance rationale, multi-head attention, and the Transformer. ↩ ↩² ↩³