Skip to content

Scaled Dot-Product Attention

Overview

Scaled dot-product attention is the exact math that combines queries, keys, and values into context-aware outputs. It is deliberately simple: a dot product for similarity, a scaling factor for stability, a softmax for normalization, and a weighted sum. It is the core computation inside every transformer, repeated thousands of times per forward pass.

 Q · Kᵀ
Attention(Q,K,V) = softmax(──────) · V
 √d_k
  • Paper: "Attention Is All You Need" (Vaswani et al., 2017)
  • Prerequisite: [01 Query, Key & Value](/01-modeling/00-fundamentals/01-attention/(01-query-key-value/)— what Q, K, V mean
  • Output shape: (seq_len, d_v)— same "length" as input, richer meaning
  • Cost: O(seq_len² · d) per layer— see [06 Complexity & Cost](/01-modeling/00-fundamentals/01-attention/(06-complexity-cost/)

-

Why Dot Product = Similarity

The dot product is the simplest trainable "relevance" function. Geometrically:

dot(a, b) = ‖a‖ · ‖b‖ · cos(θ) θ = angle between a and b

Components of relevance:
 - direction (cos θ): aligned vectors get positive scores
 - magnitude (‖a‖·‖b‖): "strong" concepts boost the score

Two tokens with related directions → high score → strong attention
Unrelated / opposite tokens → low / negative score → little attention

Why dot product and not a neural similarity net?
 - one cheap matmul for ALL pairs at once (batch-GEMM friendly)
 - O(n) parameters per token instead of O(n²) pairwise nets
 - differentiable and easy to scale to n = 8K+ tokens

-

The Four Steps

Step 1 — SCORE: S = Q · Kᵀ
 How similar is each query to every key?
 Shape: (seq_len, seq_len)

Step 2 — SCALE: S' = S / √d_k
 Prevent extreme values that push softmax to 0

Step 3 — NORMALIZE: A = softmax(S', dim=-1)
 Turn scores into probabilities (rows sum to 1)

Step 4 — WEIGHT: Output = A · V
 Weighted sum of values — every token is now a
 blend of the tokens it found relevant

Matrix shapes at each step

Q: (seq_len, d_k) Kᵀ: (d_k, seq_len) V: (seq_len, d_v)

S = Q @ Kᵀ (seq_len, seq_len) ← attention logits
S' = S / √d_k (seq_len, seq_len)
A = softmax(S') (seq_len, seq_len) ← attention weights
Out = A @ V (seq_len, d_v) ← context vectors

With batch & heads (real case):
 Q, K, V: (batch, heads, seq_len, head_dim)
 S: (batch, heads, seq_len, seq_len)

-

Why Scale by √d_k?

This is the "scaled" in scaled dot-product attention. Without it, training breaks.

The variance argument

Assume Q, K entries have mean ≈ 0, variance ≈ 1 (typical after training).

Each entry of Q·Kᵀ is the sum of d_k products:
 (q₁k₁ + q₂k₂ +... + q_{d_k}k_{d_k})

Each product qᵢkᵢ has variance ≈ 1
 → the SUM has variance ≈ d_k
 → the standard deviation ≈ √d_k

Concrete numbers:
 d_k = 128 → √d_k ≈ 11.3 → scores typically span ~±11
 d_k = 64 → √d_k ≈ 8 → scores span ~±8

What happens WITHOUT scaling

Score of 11 vs score of 0:
 e¹¹ ≈ 59,874 vs e⁰ = 1

Softmax → [~1.00, 0.00,...] → nearly one-hot

Problem: softmax nearly saturates
 - gradients of softmax → ~0 for the winning entry
 - gradient ≈ 0 → weights stop updating → training stalls
 - the model "locks in" attention decisions too early

What scaling fixes

S' = S / √d_k  scores back to O(1) range

Score of 1.0 vs 0:
 e¹ᵉ  2.72 vs 1  softmax  [0.57, 0.21,...] (smooth!)

 - softmax stays "soft": multiple keys can share attention
 - gradients remain healthy  stable training
 - the model can refine attention decisions over many steps

Key Insight: The scaling keeps the softmax in its "smooth zone" regardless of d_k. It costs nothing (one division) and prevents a catastrophic training failure.

-

Worked Example A: "The cat sat" (toy, clean numbers)

Let's compute attention for the sentence "The cat sat" with hand-picked Q, K, V:

Tokens: "The" "cat" "sat"
Queries: [1, 0] [0, 1] [1, 1]
Keys: [1, 0] [0, 1] [1, 0]
Values: [1, 2] [3, 4] [5, 6]

(Imagine these come from W_Q, W_K, W_V after training.)

Step 1: Score matrix S = Q·Kᵀ

 K_"The" K_"cat" K_"sat"
 [1,0] [0,1] [1,0]
Q_"The" [1,0] → 1 0 1
Q_"cat" [0,1] → 0 1 0
Q_"sat" [1,1] → 1 1 1

 [1 0 1]
S = [0 1 0]
 [1 1 1]

Step 2: Scale by 1/√2 ≈ 0.707 (d_k = 2)

 [0.71 0 0.71]
S' = [0 0.71 0]
 [0.71 0.71 0.71]

Step 3: Softmax each row (e^0.71 ≈ 2.03, e^0 = 1)

Row "The": [2.03, 1, 2.03]  sum 5.06  [0.40, 0.20, 0.40]
Row "cat": [1, 2.03, 1]  sum 4.03  [0.25, 0.50, 0.25]
Row "sat": [2.03, 2.03, 2.03] sum 6.09  [0.33, 0.33, 0.33]

A = [0.40 0.20 0.40]
 [0.25 0.50 0.25]
 [0.33 0.33 0.33]

Read the rows: "The" pays 40% attention to itself, 20% to "cat", 40% to "sat". "cat" pays 50% to itself.

Step 4: Weighted sum: Output = A · V

V = [1 2] (for "The")
 [3 4] (for "cat")
 [5 6] (for "sat")

Output for "cat" (row 2):
 0.25·[1,2] + 0.50·[3,4] + 0.25·[5,6]
= [0.25 + 1.50 + 1.25, 0.50 + 2.00 + 1.50]
= [3.00, 4.00]

Output for "The" (row 1):
 0.40·[1,2] + 0.20·[3,4] + 0.40·[5,6]
= [0.40 + 0.60 + 2.00, 0.80 + 0.80 + 2.40]
= [3.00, 4.00]

Output for "sat" (row 3):
 0.33·[1,2] + 0.33·[3,4] + 0.33·[5,6]
= [3.00, 4.00]

All three outputs happen to coincide here because the toy Q/K/V are symmetric. In a real (trained) model, the Q/K/V matrices are learned so different queries extract different information— see Worked Example B.


Worked Example B: Asymmetric Example (different outputs)

Same sentence setup, richer vectors that produce distinct outputs:

Tokens: "bank" "river" "money"
Queries: [1, 0] [0, 1] [1, 0.5]
Keys: [1, 0] [0.2, 1] [0.8, 0.2]
Values: [2, 0] [0, 3] [1, 1]

Step 1–2: Scores, then scaled by 0.707

Raw scores S: Scaled S':
 bank river money bank river money
 1 0.2 0.8 → 0.71 0.14 0.57 (bank's row)
 0 1 0.2 0 0.71 0.14 (river's row)
 1 0.7 0.9 0.71 0.50 0.64 (money's row)

Step 3: Softmax rows

Row "bank": [2.03, 1.15, 1.76]  [0.41, 0.23, 0.36]
Row "river": [1.00, 2.03, 1.15]  [0.24, 0.49, 0.28]
Row "money": [2.03, 1.64, 1.89]  [0.36, 0.30, 0.34]

Step 4: Weighted sums

Output for "bank":
 0.41·[2,0] + 0.23·[0,3] + 0.36·[1,1] = [1.18, 1.05]

Output for "river":
 0.24·[2,0] + 0.49·[0,3] + 0.28·[1,1] = [0.75, 1.73]
  strongest value contribution comes from "river" itself

Output for "money":
 0.36·[2,0] + 0.30·[0,3] + 0.34·[1,1] = [1.07, 1.23]

Different queries → different context vectors. "river" is most self-focused and keeps its "watery" value [0,3]; "bank" mixes in money's content strongly— exactly the kind of contextual disambiguation attention is known for.

-

Numerical Stability: The Softmax Trick

Naive softmax can overflow: e^1000 → inf. The standard fix subtracts the row max first:

softmax(xᵢ) = e^xᵢ / Σⱼ e^xⱼ

Numerically stable version:
 1. m = max(x) (per row)
 2. subtract: xᵢ − m → all exponents ≤ 0 → no overflow
 3. softmax(xᵢ − m)

Math note: subtracting a constant leaves softmax IDENTICAL:
 e^(xᵢ−m) / Σ e^(xⱼ−m) = e^xᵢ / Σ e^xⱼ (m cancels out)
Example: scores [1000, 999, 998]
Naive: e¹⁰⁰⁰ → OverflowError 
Stable: m = 1000 → [0, −1, −2]
 → [e⁰, e⁻¹, e⁻²] / sum = [0.67, 0.24, 0.09] 

Production kernels (FlashAttention) never materialize the N×N matrix; they recompute the softmax backward online— see Flash Attention.

-

Putting It Together: One Shot, All Positions

X (6 tokens, d_model = 4)
 │ ┌─────────────────────────────
 │ Q = X·W_Q (6×3) │ S = Q·Kᵀ (6×6)
 │ K = X·W_K (6×3) │ S' = S / √d_k (6×6)
 │ V = X·W_V (6×4) │ A = softmax(S') (6×6)
 ▼ │ Out = A·V (6×4)
 [Attention(Q,K,V)] ◄────┘

Every output row is a blend of ALL value rows,
weighted by how much its query matched each key.

-

Key Takeaways

One formula: softmax(QKᵀ/√d_k)·V— score, scale, normalize, weight Scaling matters: ÷√d_k keeps softmax smooth and gradients alive Dot product = cheap similarity: one matmul for all pairs Worked examples: follow the 4 steps by hand before trusting the formula Stability trick: subtract the row max before exp Learned output: different queries extract different context— that IS the magic

-