Scaled dot-product attention

What Is Scaled Dot-Product Attention?

Components breakdown:

$QK^T$ : Compatibility scores (how well each query matches each key)
$\frac{1}{\sqrt{d_k}}$ : Scaling to prevent gradient saturation
softmax: Converts scores to probability distribution (weights sum to 1)
Multiply by $V$ : Weighted sum of values based on attention weights

Derivation From First Principles

Step 1: Why Dot Product for Similarity?

Start with the fundamental question: How do we measure if query $q$ matches key $k$ ?

The dot product $q \cdot k = \sum_{i=1}^{d_k} q_i k_i$ measures alignment:

If $q$ and $k$ point in the same direction: large positive value
If orthogonal: zero
If opposite: large negative value

Why this step? We need a differentiable similarity metric. Dot product is computationally cheap ( $O(d_k)$ ) and naturally captures semantic similarity in embedding spaces.

Step 2: Matrix Form for Batch Processing

For $n$ queries and $m$ keys: $S = QK^T \in \mathbb{R}^{n \times m}$

where $S_{ij} = q_i \cdot k_j$ is the compatibility score between query $i$ and key $j$ .

Why this step? Modern hardware (GPUs/TPUs) is optimized for matrix multiplication. Batching all queries at once is100x faster than loping.

Step 3: The Scaling Problem

Observation: If $Q$ and $K$ have entries drawn from $\mathcal{N}(0, 1)$ , then:

$\mathbb{E}[q \cdot k] = \mathbb{E}\left[\sum_{i=1}^{d_k} q_i k_i\right] = \sum_{i=1}^{d_k} \mathbb{E}[q_i]\mathbb{E}[k_i] = 0$

$\text{Var}(q \cdot k) = \sum_{i=1}^{d_k} \text{Var}(q_i k_i) = \sum_{i=1}^{d_k} \mathbb{E}[q_i^2]\mathbb{E}[k_i^2] = d_k$

So the standard deviation grows as $\sqrt{d_k}$ ! For typical $d_k = 64$ , dot products have std $\approx 8$ .

Why this matters: Softmax is $\sigma(z)_i = \frac{e^{z_i}}{\sum_j e^{z_j}}$ . When inputs have large magnitude:

Large positive values: softmax $\approx$ one-hot (gradient $\approx 0$ )
Relative differences shrink: hard to learn which key is slightly better

The fix: Divide by $\sqrt{d_k}$ to normalize variance back to 1: $\text{Var}\left(\frac{q \cdot k}{\sqrt{d_k}}\right) = \frac{d_k}{d_k} = 1$

Step 4: Softmax Normalization

Apply softmax row-wise to convert scores to attention weights: $A = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) \in \mathbb{R}^{n \times m}$

where $A_{ij} = \frac{\exp(S_{ij}/\sqrt{d_k})}{\sum_{j'=1}^m \exp(S_{ij'}/\sqrt{d_k})}$ .

Why this step?

Probabilistic interpretation: Each row of $A$ is a distribution over keys ( $\sum_j A_{ij} = 1$ )
Differentiability: Smooth function for backprop
Competition: Keys compete—attending more to one means less to others

Step 5: Weighted Value Aggregation

Finally, use attention weights to take a weighted average of values: $\text{output}_i = \sum_{j=1}^m A_{ij} v_j$

In matrix form: $\text{Attention}(Q, K, V) = AV \in \mathbb{R}^{n \times d_v}$

Why this step? The query "asks a question," keys "match the question," and values "provide the answer." We retrieve a soft-blended answer based on relevance.

Worked Examples

Worked example Example 1: Single Head, Small Dimensions

Setup: $d_k = 4$ , $d_v = 3$ , one query, two keys.

$Q = \begin{bmatrix}1 & 0 & 1 & 0\end{bmatrix}, \quad K = \begin{bmatrix}1 & 0 & 1 & 0\\0 & 1 & \end{bmatrix}, \quad V = \begin{bmatrix}5 & 0 & 0\\0 & 3 & 0\end{bmatrix}$

Step 1: Compute $QK^T$ $QK^T = \begin{bmatrix}1 & 0 & 1 & 0\end{bmatrix} \begin{bmatrix}1 & 0\\0 & 1\\1 & 0\\0 & 1\end{bmatrix} = \begin{bmatrix}2 & 0\end{bmatrix}$

Why this step? Query matches key 1 perfectly (both have pattern [1,0,1,0]), doesn't match key 2 at all.

Step 2: Scale by $\sqrt{d_k} = 2$ $\frac{QK^T}{\sqrt{d_k}} = \begin{bmatrix}1 & 0\end{bmatrix}$

Why this step? Prevents saturation; now scores are order-1.

Step 3: Softmax $A = \text{softmax}([1, 0]) = \begin{bmatrix}\frac{e^1}{e^1 + e^0} & \frac{e^0}{e^1 + e^0}\end{bmatrix} \approx \begin{bmatrix}0.731 & 0.269\end{bmatrix}$

Why this step? Query attends 73% to key 1, 27% to key 2. Step 4: Weighted sum of values $\text{output} = 0.731 \begin{bmatrix}5\\0\\0\end{bmatrix} + 0.269 \begin{bmatrix}0\\3\\0\end{bmatrix} = \begin{bmatrix}3.655\\0.807\\0\end{bmatrix}$

Interpretation: Output is mostly value 1 (makes sense—query matched key 1 best).

Worked example Example 2: Why Scaling Matters

Same setup but $d_k = 64$ and query/key entries $\sim \mathcal{N}(0, 1)$ . Suppose $QK^T = \begin{bmatrix}8 & 7.5\end{bmatrix}$ .

Without scaling: $\text{softmax}([8, 7.5]) = \begin{bmatrix}\frac{e^8}{e^8 + e^{7.5}} & \frac{e^{7.5}}{e^8 + e^{7.5}}\end{bmatrix} \approx \begin{bmatrix}0.622 & 0.378\end{bmatrix}$ But gradient of softmax w.r.t. input: $\frac{\partial \sigma_i}{\partial z_i} = \sigma_i(1 - \sigma_i) \approx 0.622 \times 0.378 = 0.235$

With scaling ( $\sqrt{d_k} = 8$ ): $\text{softmax}([1, 0.9375]) \approx \begin{bmatrix}0.516 & 0.484\end{bmatrix}$ $\frac{\partial \sigma_1}{\partial z_1} \approx 0.516 \times 0.484 = 0.250$

Why this matters: Scaled version has larger, more balanced gradient—model learns faster and distinguishes subtle differences between keys.

Common Mistakes

Active Recall Flashcards

#flashcards/ai-ml

What is the formula for scaled dot-product attention? :: $\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

Why do we scale by

\sqrt{d_k}

instead of

d_k

Because variance of dot product grows as

d_k

, so std grows as

\sqrt{d_k}

. Dividing by

\sqrt{d_k}

normalizes variance back to 1, keeping softmax in its sensitive gradient region.

What is the shape of $QK^T$ if $Q$ isn \times d_k $and$ K $is$ m \times d_k $? :::$ n \times m $(each element$ (i,j) $is the similarity between query$ i $and key$ j$)

What happens if you don't scale the dot product with large

d_k

Softmax saturates (becomes nearly one-hot), gradients vanish, and the model can't learn subtle differences between keys.

In attention mechanism, what do Q, K, V represent?

Q (query): what you're looking for; K (key): what's available to match; V (value): actual content to retrieve.

What does the softmax output represent in attention?

A probability distribution over keys—attention weights that sum to 1, indicating how much to "look at" each position.

Why use dot product instead of other similarity metrics?

Computationally efficient (

O(d_k)

), differentiable, naturally captures alignment in embedding spaces, and hardware-optimized (matrix multiplication).

What is the output dimension of attention with

n

queries and

m

keys?

n \times d_v

(one output vector per query, each with dimension

d_v

)

Mnemonic

Feynman Explanation

Recall Explain to a 12-year-old

Imagine you're doing homework and you can ask your textbook questions. But your textbook is HUGE—hundreds of pages.

Attention is like a smart search system:

Query (Q): You write down your question: "What is photosynthesis?"
Keys (K): Each page in the textbook has a title/summary (like "Plants," "Cells," "Energy "History")
Comparison: You check how well your question matches each page title. The "Plants" page is a great match (score = 10), "Energy" is okay (score = 5), "History" barely matches (score = 1).
The Scaling Trick: If your textbook was MASSIVE (thousands of pages), the scores would get huge (score = 1000 for best match). That makes your brain say "ONLY look at this ONE page" and ignore everything else. Scaling is like saying "let's keep scores between 1-10 so I consider multiple pages."
Softmax: Convert scores to percentages: 70% Plants, 25% Energy, 5% History.
Values (V): Each page has actual content (paragraphs about plants, energy, etc.)
Output: You read70% of the Plants page, 25% of Energy page, 5% of History page, and blend them into one answer.

Why scaling matters: Without it, you'd become a robot that reads ONLY the top match and ignores everything else—even if the second-best page has useful info too! Scaling keeps your mind open to multiple sources.

Connections

Multi-Head Attention — uses multiple scaled dot-product attention operations in parallel
Self-Attention — when Q, K, V all come from the same sequence
Cross-Attention — when Q comes from one sequence, K/V from another (encoder-decoder)
Softmax Function — the normalization function that converts scores to probabilities
Gradient Descent — why scaling matters for stable learning
Positional Encoding — provides sequence order information (attention is permutation-invariant)
Query-Key-Value Model — the conceptual framework borrowed from information retrieval
Attention Mask — used to prevent attending to certain positions (future tokens, padding)

Last updated: 2026-07-01 | Master this before moving to multi-head attention

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Scaled dot-product attention basically ek smart tarika hai information retrieve karne. Socho tumhare pas ek query hai (jo tum dhondh rahe ho) aur bahut sare keys hain (jo available options hain). Ab simple dot product se tum check karte ho ki tumhari query kis key se match karti hai. Lekin problem ye hai ki jab dimensions bade ho jate hain (jaise d_k = 64 ya 512), tab dot product ki values bhi bahut badi ho jaati hain. Agar tum directly softmax lagao, toh wo almost one-hot ban jata hai—matlab sirf ek hi key pe focus ho jata hai aur baki sab ignore ho jaate hain. Learning ke liye ye bahut bura hai kyunki model subtle differences nahi seekh pata.

Isliye hum scale karte hain √d_k se divide karke. Ye scaling ensure karti hai ki softmax sensitive region mein rahe, jahan gradients healthy rahein aur model properly learn kar sake. Iske baad softmax se hum attention weights nikaalte hain (probability distribution), aur phir un weights se values ka weighted average lete hain. Final output mein humein mil jata hai ek context-aware representation jo query ke liye sabse relevant information contain karta hai.

Ye mechanism Transformers ka heart hai—bilkul jaise tum library mein relevant books dhoondh rahe ho. Query tumhara sawaal hai, keys books ke titles hain, aur values actual content hai. Scaling ensure karti hai ki tum multiple relevant books consider kar sako, na ki sirf ek pe fixate ho jao. Isi wajah se Transformers itne powerful hain—wo flexible attention patterns seekh sakte hain without losing gradient signal during training.

What Is Scaled Dot-Product Attention?

Derivation From First Principles

Step 1: Why Dot Product for Similarity?

Step 2: Matrix Form for Batch Processing

Step 3: The Scaling Problem

Step 4: Softmax Normalization

Step 5: Weighted Value Aggregation

Worked Examples

Common Mistakes

Active Recall Flashcards

Mnemonic

Feynman Explanation

Connections

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Transformer Architecture

Connections