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 — 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

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/1

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