Scaled dot-product attention

Bahdanau attention computes a context vector as a weighted sum of encoder states, with weights determined by a learned alignment score between the decoder state and each encoder position. The Transformer (Vaswani et al., 2017) generalizes this to a single unified mechanism: scaled dot-product attention, which any sequence -- not just encoder-decoder pairs -- can apply to itself[1].

The query-key-value abstraction

Frame attention as a soft lookup table: you have a table of key-value pairs, and a query you want to look up. Instead of returning the value for an exact key match, you return a weighted sum of all values, with weights proportional to how well each key matches the query.

For a set of $n$ queries $Q \in \mathbb{R}^{n \times d_k}$, $m$ keys $K \in \mathbb{R}^{m \times d_k}$, and $m$ values $V \in \mathbb{R}^{m \times d_v}$:

Steps:

1. Compute similarity scores: $Q K^\top \in \mathbb{R}^{n \times m}$ (one score per query-key pair).
2. Scale by $1/\sqrt{d_k}$ to prevent saturation.
3. Apply softmax row-wise to get weight matrix $A \in \mathbb{R}^{n \times m}$ (each row sums to 1).
4. Return $A V \in \mathbb{R}^{n \times d_v}$ (each query gets a weighted sum of values).

Why scale by $1/\sqrt{d_k}$?

If the components of $q$ and $k$ are i.i.d. with mean 0 and variance 1, then $q \cdot k = \sum_{i=1}^{d_k} q_i k_i$ has mean 0 and variance $d_k$, so the standard deviation grows as $\sqrt{d_k}$.

Without scaling, the dot products grow large in magnitude as $d_k$ increases. The softmax of large values approaches a one-hot distribution: almost all weight is placed on the largest score. This pushes the attention into a nearly-argmax regime where gradients are tiny (the softmax is saturated). Dividing by $\sqrt{d_k}$ keeps the inputs to softmax at unit variance, maintaining informative gradients.

Put numbers on it at $d_k = 64$, a typical head dimension, where the score standard deviation is $\sqrt{64} = 8$. Take two keys whose raw scores differ by one standard deviation, a gap of $8$, with everything else held fixed so only the scaling changes:

• Unscaled. $\text{softmax}(8, 0)$ puts weight $0.9997$ on the larger key. The distribution is one-hot in all but name; the second key gets $0.0003$ and its gradient is starved.
• Scaled by $1/\sqrt{d_k} = 1/8$. The same gap becomes $1$, and $\text{softmax}(1, 0)$ puts weight $0.731$ on the larger key. The attention stays soft and both keys carry gradient.

A wider gap of two standard deviations ($16$ raw, $2$ after scaling) saturates even harder unscaled ($\approx 1.0000$) but stays soft once scaled ($0.881$). The scaling does not change which key wins; it changes whether the loser is heard at all.

Self-attention

When $Q$, $K$, $V$ are all derived from the same sequence $X \in \mathbb{R}^{T \times d}$ via linear projections:

the result is self-attention: every position attends to every other position in the same sequence. Each output position is a weighted sum of all value vectors, with weights reflecting how relevant each key position is to the current query position.

Causal masking

For autoregressive language modeling, position $t$ must not attend to positions $> t$ (it cannot see the future). Causal masking sets the pre-softmax scores to $-\infty$ for all $j > t$:

After softmax, $-\infty$ becomes 0, so future positions receive zero weight.

Implementation

import numpy as np

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: (n, d_k)  K: (m, d_k)  V: (m, d_v)  mask: (n, m) boolean (True = masked out)
    Returns: (n, d_v) output, (n, m) attention weights
    """
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)          # (n, m)
    if mask is not None:
        scores[mask] = -1e9                   # set masked positions to -inf
    attn = np.exp(scores - scores.max(axis=-1, keepdims=True))
    attn /= attn.sum(axis=-1, keepdims=True)  # softmax row-wise
    out = attn @ V                            # (n, d_v)
    return out, attn

# Demo: 5-token sequence, d_k = d_v = 4
T, d_k, d_v = 5, 4, 4
rng = np.random.default_rng(0)
X  = rng.normal(0, 1, (T, 8))   # input embeddings (d_model=8)
Wq = rng.normal(0, 0.1, (8, d_k))
Wk = rng.normal(0, 0.1, (8, d_k))
Wv = rng.normal(0, 0.1, (8, d_v))

Q, K, V = X @ Wq, X @ Wk, X @ Wv

# Causal mask: True where j > i
mask = np.triu(np.ones((T, T), dtype=bool), k=1)

out, attn = scaled_dot_product_attention(Q, K, V, mask=mask)
print("Output shape:", out.shape)
print("Attention weights (row = query, col = key):")
print(np.round(attn, 3))
print("Each row sums to:", attn.sum(axis=1).round(4))

Verify that the upper triangle of the attention matrix is zero (masked), and that each row in the lower-triangular region sums to 1.

Complexity

For a sequence of length $T$:

• Computing $Q K^\top$: $O(T^2 d_k)$ -- quadratic in sequence length.
• Computing $A V$: $O(T^2 d_v)$.
• Total: $O(T^2 d)$ time and $O(T^2)$ memory (the attention matrix).

This quadratic scaling in $T$ is the core limitation of standard attention, motivating FlashAttention (D7) and linear attention variants.

Exercises

Exercise 1. A single query $q = (1, 0)$ attends to two keys $k_1 = (1, 0)$ and $k_2 = (0, 1)$ with values $v_1 = (1, 0)$ and $v_2 = (0, 1)$, head dimension $d_k = 2$. Compute the scaled scores, the attention weights, and the output vector by hand.

Exercise 2. Two keys produce raw scores that differ by $8$ (everything else fixed). Compute the softmax weight on the larger key with no scaling, and again after dividing by $\sqrt{d_k} = 8$. Explain why the scaled case keeps gradient flowing to the smaller key.

Exercise 3 (code). Implement scaled dot-product attention and verify two structural facts on a random causal example: every attention row sums to $1$, and the strictly upper triangle (future positions) is exactly zero.

import numpy as np

def attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(d_k)
    if mask is not None:
        scores = np.where(mask, -1e9, scores)
    attn = np.exp(scores - scores.max(axis=-1, keepdims=True))
    attn /= attn.sum(axis=-1, keepdims=True)
    return attn @ V, attn

rng = np.random.default_rng(0)
T, d_k, d_v = 6, 8, 8
Q, K, V = rng.normal(size=(T, d_k)), rng.normal(size=(T, d_k)), rng.normal(size=(T, d_v))
mask = np.triu(np.ones((T, T), dtype=bool), k=1)   # True above diagonal = future

out, attn = attention(Q, K, V, mask=mask)
assert np.allclose(attn.sum(axis=1), 1.0)          # rows are distributions
assert np.allclose(attn[mask], 0.0)                # no weight on the future
print("rows sum to 1, upper triangle zeroed; output shape", out.shape)

The next lesson stacks $h$ independent attention heads in parallel -- multi-head attention -- which allows different heads to attend to different representation subspaces.