Scaled Dot Product Attention — Module 10.1

How This Lesson Fits the Module & Volume

You can now name Q, K, V, assemble attention, and compute scores. Scaled dot-product attention is the specific scoring + weighting recipe used throughout Transformers (Vaswani et al., 2017).

This lecture is the last of the Module 10.1 “foundation eight.” After it, you move to structural variants—self-attention, cross-attention, multi-head attention—that reuse this same formula inside encoders and decoders.

Learning Objectives

By the end of this lesson, students should be able to:

Write the scaled dot-product formula softmax((QK^T)/√d_k) V.
Explain why dividing by √d_k stabilizes softmax.
Implement the algorithm with masking in PyTorch (manual and F.scaled_dot_product_attention).
Trace shapes from Q, K, V through scores, weights, and output.
Relate this op to multi-head attention (heads run SDPA in parallel).
Identify numerical issues (overflow, all-masked rows) and mitigations.

Definition

Scaled dot-product attention is attention with score function score(q, k) = (q · k) / √d_k. In matrix form: Attention(Q, K, V) = softmax(QK^T / √d_k) V.

Why Scale?

Assume query/key components are independent with mean 0 and variance 1. Then q · k = ∑_i=1^d_k q_i k_i has variance d_k, so magnitudes grow like √d_k. Softmax on huge logits saturates. Dividing by √d_k keeps score variance near 1 regardless of head dimension—gradients stay healthy.

Step	Operation	Shape (batched)
1	Form Q, K, V	(B, T_q, d_k), (B, T_k, d_k), (B, T_k, d_v)
2	QK^T	(B, T_q, T_k)
3	Divide by √d_k	same
4	Mask (optional)	same
5	softmax over keys	(B, T_q, T_k)
6	weights @ V	(B, T_q, d_v)

Manual Implementation vs. Fused API

Manual (Learning)

Explicit matmuls + softmax.
Easy to print weights.
Great for homework/debug.

F.scaled_dot_product_attention

Fused kernels (FlashAttention, etc.).
Faster / lower memory.
Production default in modern PyTorch.

Multi-Head Wrap

Split d_model into h heads.
Run SDPA per head.
Concat + W_O.

PyTorch: From Scratch and Built-In

import math
import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))
    weights = F.softmax(scores, dim=-1)
    return weights @ V, weights

B, T, d_k = 2, 5, 16
Q = K = V = torch.randn(B, T, d_k)
causal = torch.tril(torch.ones(T, T))           # (T, T) lower-triangular
out, w = scaled_dot_product_attention(Q, K, V, mask=causal)
print(out.shape, w[0].round(decimals=2))

# Fused API (PyTorch 2+); attn_mask True=keep / float mask variants exist
out2 = F.scaled_dot_product_attention(Q, K, V, is_causal=True)
print(out2.shape)   # torch.Size([2, 5, 16])

Masking Pitfalls

If an entire score row is −∞ (every key masked), softmax becomes NaN. Always ensure at least one valid key per query, or handle empty rows explicitly. Causal masks for length T are lower-triangular; PAD masks are usually key-side boolean tensors broadcast across queries.

Strengths

Simple, parallel, GPU-friendly.
Scaling fixes softmax saturation.
Foundation for all Transformer attention.

Tradeoffs

Still O(T²) for dense attention.
Needs careful mask semantics.
Fused kernels can hide weights (harder to inspect).

Common Misconception

“Scaling changes which key wins.” Dividing every score in a row by the same positive constant does not change the argmax, but it does change the softmax distribution—making it less peaky—and that is what preserves gradient flow. Scaling is about soft weights and training dynamics, not about reordering preferences under a hard max.

Knowledge Check

Short Answer: Write the scaled dot-product attention formula. Answer: softmax(QK^T / √d_k) V.
True/False: The scale factor is √d_k in the denominator. Answer: True.
Multiple Choice: Scaling exists mainly to: (a) reduce vocabulary size, (b) keep softmax from saturating as d_k grows, (c) remove the need for V. Answer: (b).
Short Answer: What is the shape of QK^T for batched Q, K? Answer: (B, T_q, T_k).
True/False: A causal mask is typically lower-triangular for autoregressive decoding. Answer: True.
Multiple Choice: F.scaled_dot_product_attention is: (a) deprecated, (b) a fused optimized implementation, (c) only for CNNs. Answer: (b).
Short Answer: What happens if all keys in a row are masked to −∞? Answer: Softmax can produce NaNs; ensure a valid key or handle empty rows.
True/False: Uniform positive scaling of a score row changes the argmax of that row. Answer: False—argmax is unchanged; softmax softness changes.
Multiple Choice: Multi-head attention runs SDPA: (a) once globally only, (b) per head then concatenates, (c) only on values. Answer: (b).
Short Answer: Which lecture comes next in Module 10.1 after this foundation? Answer: Self Attention.

Key Takeaways

Scaled dot-product attention is softmax(QK^T/√d_k) V.
Scaling controls softmax sharpness as head dimension grows.
Masks + stable softmax are part of a correct implementation.
Next: Self Attention—SDPA when Q, K, V share one sequence.

Trainer’s Guide

Hands-on idea: Match manual SDPA outputs to F.scaled_dot_product_attention within tolerance; then enable is_causal=True.

Discussion prompt: Derive the variance argument for /√d_k on the board with a tiny numerical example.

Recap: Scaled dot-product attention is the Transformer’s core attention op. Continue with Self Attention.