Scaled Dot-Product Attention Formula

Introduction

Now that we have intuition for what attention does, let's derive the mathematical formula precisely. By the end of this section, you'll understand every component of:

This formula is deceptively simple—just a few operations—but each component is carefully chosen. Understanding why each part exists will help you implement attention correctly and debug issues when they arise.

What Does This Formula Actually Do?

Think of attention as: Each token decides how much it should "look at" every other token.

The flow is elegantly simple:

Why Is V Outside the Softmax?

This is a common point of confusion. Let's be crystal clear:

The softmax must operate on scores, not content.

• $QK^T$ gives raw compatibility scores: "How similar am I to each token?"
• Softmax converts these to attention weights $alpha_{ij}$ = how much token $i$ should pay attention to token $j$
• $V$ contains the actual information we want to retrieve

What Context Does the Attention Matrix Encode?

The attention matrix A = ext{softmax}left(rac{QK^T}{sqrt{d_k}} ight) is an $n imes n$ matrix (for n tokens).

Each row $A_i$ tells you: For token $i$, how much it attends to every other token.

Example: Consider the sentence "The cat sat on the mat."

At the token "cat", the attention row might look like:

This row vector = context distribution. It tells us "cat" mostly cares about itself (0.60) and the verb "sat" (0.25).

The Google Search Analogy

Here's the most intuitive way to understand attention:

2.1 The Attention Formula Breakdown

The Complete Formula

Breaking down each component:

Let's identify each component:

Dimensions

For a single sequence:

For batched operations with batch size $b$:

2.2 Dot Product as Similarity Measure

Why Dot Product?

The dot product between two vectors measures their similarity:

Where $heta$ is the angle between vectors.

Properties:

• If vectors point in same direction ($heta = 0$): dot product is high (positive)
• If vectors are orthogonal ($heta = 90°$): dot product is zero
• If vectors point in opposite directions ($heta = 180°$): dot product is negative

Computing QKᵀ

When we compute $QK^T$:

The result is a matrix where:

• $ext{scores}[i][j]$ = dot product of query $i$ and key $j$
• Shape: $[n_q, n_k]$

Example:

1Q = [[1, 0],    # Query 1
2     [0, 1]]    # Query 2
3
4K = [[1, 0],    # Key 1
5     [1, 1],    # Key 2
6     [0, 1]]    # Key 3

Reading the result:

• Query 1 is similar to Key 1 and Key 2 (score = 1), dissimilar to Key 3 (score = 0)
• Query 2 is similar to Key 2 and Key 3 (score = 1), dissimilar to Key 1 (score = 0)

2.3 The Critical Scaling Factor: $sqrt{d_k}$

The Problem Without Scaling

Consider what happens as $d_k$ (dimension of keys/queries) increases:

If $Q$ and $K$ have elements drawn from a distribution with mean 0 and variance 1:

• Each element of $QK^T$ is a sum of $d_k$ terms
• By Central Limit Theorem: $ext{mean} \approx 0$, $ext{variance} \approx d_k$

Why Extreme Values Are Bad

Softmax with extreme values produces nearly one-hot outputs:

Problems:

1. Vanishing gradients: Softmax saturates, $abla ightarrow 0$
2. No learning: Model can't adjust attention patterns
3. Numerical instability: Overflow in $exp()$

The Solution: Scale by $sqrt{d_k}$

Dividing by $sqrt{d_k}$ normalizes the variance:

Mathematical Derivation

Let ec{q}, ec{k} be vectors with independent elements $sim mathcal{N}(0, 1)$:

After scaling:

Empirical Demonstration

1import torch
2
3d_k = 512
4q = torch.randn(1000, d_k)  # 1000 queries
5k = torch.randn(1000, d_k)  # 1000 keys
6
7# Without scaling
8scores_unscaled = (q @ k.T)
9print(f"Unscaled std: {scores_unscaled.std():.2f}")  # ≈ 22.6 (≈ √512)
10
11# With scaling
12scores_scaled = (q @ k.T) / (d_k ** 0.5)
13print(f"Scaled std: {scores_scaled.std():.2f}")  # ≈ 1.0

2.4 Softmax Normalization

Purpose of Softmax

Softmax converts raw scores into a probability distribution:

Properties:

1. All outputs positive: $e^x > 0$ for all $x$
2. Outputs sum to 1: $\sum_i ext{softmax}(x)_i = 1$
3. Preserves ordering: if $x_i > x_j$, then $ext{softmax}(x)_i > ext{softmax}(x)_j$
4. Differentiable: enables gradient-based learning

Applying Softmax to Attention Scores

Each row sums to 1:

• Row $i$ contains attention weights for query $i$
• $A_{ij}$ = how much query $i$ attends to key/value $j$

Softmax Temperature

The scaling factor acts like inverse temperature:

• Lower T (higher scaling): Sharper, more focused attention
• Higher T (lower scaling): Softer, more uniform attention

2.5 Weighted Sum with Values

The Final Computation

After computing attention weights, we use them to aggregate values:

Shape transformation:

What This Computes

For each query position $i$:

This is a weighted average of all value vectors, where weights indicate relevance.

Example

1attention_weights = [[0.7, 0.2, 0.1],    # Query 1: mostly attends to Value 1
2                     [0.1, 0.8, 0.1]]    # Query 2: mostly attends to Value 2
3
4V = [[1, 0],    # Value 1
5     [0, 1],    # Value 2
6     [1, 1]]    # Value 3

2.6 Complete Algorithm

Pseudocode

1def attention(Q, K, V, mask=None):
2    # Step 1: Compute raw attention scores
3    d_k = K.shape[-1]
4    scores = Q @ K.transpose(-2, -1)  # [n_q, n_k]
5
6    # Step 2: Scale scores
7    scores = scores / sqrt(d_k)
8
9    # Step 3: Apply mask (optional)
10    if mask is not None:
11        scores = scores.masked_fill(mask == 0, -infinity)
12
13    # Step 4: Softmax to get attention weights
14    attention_weights = softmax(scores, dim=-1)
15
16    # Step 5: Weighted sum of values
17    output = attention_weights @ V  # [n_q, d_v]
18
19    return output, attention_weights

The Algorithm in Equations

Computational Complexity

For self-attention where $n_q = n_k = n$:

2.7 Gradient Flow Through Attention

Why Gradients Flow Well

The attention mechanism has favorable gradient properties:

1. Softmax gradients: Non-zero gradients as long as attention isn't saturated (thanks to scaling!)
2. Direct paths: Every output position has a direct connection to every input position
3. No vanishing through time: Unlike RNNs, no multiplicative chains over sequence length

Gradient of Softmax

For softmax output $s = ext{softmax}(x)$:

Where $delta_{ij}$ is the Kronecker delta (1 if $i=j$, else 0).

Backpropagation Through Attention

All operations are differentiable with well-behaved gradients!

2.8 Mathematical Notation Summary

Standard Notation

The Formula in Matrix Form

Batched Form

With batch dimension $b$:

With Multi-Head Attention (Preview)

Summary

The Complete Picture

Component purposes:

1. $QK^T$: Compute similarity between each query and all keys
2. $/ sqrt{d_k}$: Normalize variance to prevent softmax saturation
3. $ext{softmax}$: Convert to probability distribution
4. $imes V$: Aggregate values weighted by attention

Key Equations at a Glance

Critical Insights

1. Scaling is essential: Without $sqrt{d_k}$, training fails for large $d_k$
2. Softmax creates focus: Converts scores to probability distribution
3. Output is weighted average: Each output is combination of all values
4. $O(n^2)$ complexity: Quadratic in sequence length, limiting long sequences

Exercises

Mathematical Derivations

1. Variance calculation: Prove that if $q$ and $k$ are vectors with independent $mathcal{N}(0,1)$ elements, then ext{Var}[ec{q} \cdot ec{k}] = d_k.
2. Softmax property: Prove that $\sum_i ext{softmax}(x)_i = 1$ for any input vector $x$.
3. Temperature effect: If we use $ext{softmax}(x/T)$, what happens to the output distribution as $T ightarrow 0$? As $T ightarrow infty$?

Calculation Practice

4. Given:

1Q = [[1, 0]]
2K = [[1, 0], [0, 1], [1, 1]]
3V = [[1, 2], [3, 4], [5, 6]]
4d_k = 2

Compute the attention output step by step:

• Compute $QK^T$
• Apply scaling: $rac{QK^T}{\sqrt{d_k}}$
• Apply softmax
• Compute final output: $A cdot V$

5. What happens to attention weights if one score is much larger than others? Compute $ext{softmax}([0, 0, 10])$.

Conceptual Questions

1. Why can't we just normalize $QK^T$ by dividing by its maximum value instead of $sqrt{d_k}$?
2. If we wanted attention to be "sharper" (more focused on one position), what could we change in the formula?
3. Why do $K$ and $V$ need to have the same sequence length, but $Q$ can have a different length?

Next Section Preview

In the next section, we'll work through a complete numerical example with actual numbers. You'll compute attention by hand (or calculator) on a tiny 3-token, 4-dimensional example, building confidence through concrete verification.