2.1 The attention mechanism

Attention is a rule for retrieving information from a set of vectors. A query describes what a position is looking for; keys determine how strongly other positions match; values supply the information that gets combined.

This chapter builds scaled dot-product attention from its shapes, then implements a causal version in Python. Familiarity with matrix multiplication and inner products is enough to follow the derivation.

From tokens to queries, keys, and values

Let $X\in\mathbb{R}^{n\times d_{\mathrm{model}}}$ contain the representations of $n$ tokens, one per row. A single attention head applies three learned projections:

$$\begin{aligned} Q &= XW_Q \in \mathbb{R}^{n\times d_k}, \\ K &= XW_K \in \mathbb{R}^{n\times d_k}, \\ V &= XW_V \in \mathbb{R}^{n\times d_v}. \end{aligned} \tag{2.1}$$(2.1)

Queries and keys have the same feature dimension because their inner product must be defined. Values can have a different dimension. When $d_k = 64$, each comparison sums 64 products.

What the names mean

For position $i$, the vector $\mathbf{q}_i$ is compared with every allowed key $\mathbf{k}_j$. The resulting weights select a mixture of the value vectors $\mathbf{v}_j$. The names are a useful retrieval analogy; all three representations are learned, and their coordinates need not have a simple interpretation.

Scaled dot-product attention

The score matrix is $S=QK^{\mathsf T}/\sqrt{d_k}$. Its entry $S_{ij}$ measures the compatibility of query $i$ with key $j$. Softmax is applied separately across each row:

$$\begin{aligned} A_{ij} &= \frac{\exp(S_{ij})}{\sum_{r=1}^{n}\exp(S_{ir})}, \\ Z_i &= \sum_{j=1}^{n}A_{ij}V_j. \end{aligned} \tag{2.2}$$(2.2)

In matrix notation, the computation is

$$\operatorname{Attention}(Q,K,V) =\operatorname{softmax}\!\left(\frac{QK^{\mathsf T}}{\sqrt{d_k}}\right)V. \tag{2.3}$$(2.3)

In (2.3), each output row is a weighted average of the value rows. The weights are nonnegative and sum to one. This is the scaled dot-product attention formulation from Vaswani et al..

Why divide by the square root?

As a simplifying calculation, suppose the coordinates of a query and key are independent, zero-mean random variables with unit variance. Their dot product has variance $d_k$. Dividing by $\sqrt{d_k}$, as in (2.3), keeps that variance at one under these assumptions. This helps keep initial scores from growing solely because the head dimension is larger; it is not a guarantee that learned representations stay normalized.

Causal masking

An autoregressive language model predicts a token from the prefix before it. During a parallel training pass, position $i$ must therefore be prevented from reading later positions. Add a mask before softmax:

$$M_{ij}=\begin{cases} 0 & j\leq i,\\ -\infty & j>i. \end{cases} \qquad A=\operatorname{softmax}(S+M). \tag{2.4}$$(2.4)

For three positions, the mask is

$$M=\begin{bmatrix} 0 & -\infty & -\infty \\ 0 & 0 & -\infty \\ 0 & 0 & 0 \end{bmatrix}. \tag{2.5}$$(2.5)

The masked softmax in (2.4) assigns zero probability to future entries. The diagonal remains available: the representation at an input position can attend to that input token while being used to predict the following token.

A small Python implementation

The following example implements one head for one sequence using NumPy. It omits dropout, batching, and gradients so the core computation stays visible. It uses random projected vectors rather than learned model weights.

Python implementation

attention.py

import numpy as np


def causal_attention(Q, K, V):
    """Q: (n, d_k), K: (n, d_k), V: (n, d_v)."""
    if Q.ndim != 2 or K.ndim != 2 or V.ndim != 2:
        raise ValueError("Expected two-dimensional arrays")
    if Q.shape != K.shape or Q.shape[0] != V.shape[0]:
        raise ValueError("Incompatible sequence or head dimensions")
    if Q.shape[0] == 0 or Q.shape[1] == 0:
        raise ValueError("Sequence and key dimensions must be nonzero")

    scores = (Q @ K.T) / np.sqrt(Q.shape[-1])
    future = np.triu(np.ones(scores.shape, dtype=bool), k=1)
    scores = np.where(future, -np.inf, scores)
    scores = scores - scores.max(axis=-1, keepdims=True)
    weights = np.exp(scores)
    weights /= weights.sum(axis=-1, keepdims=True)
    return weights @ V, weights


rng = np.random.default_rng(42)
Q = rng.normal(size=(4, 8))
K = rng.normal(size=(4, 8))
V = rng.normal(size=(4, 6))
output, weights = causal_attention(Q, K, V)

assert output.shape == (4, 6)
assert np.allclose(weights.sum(axis=-1), 1.0)
assert np.allclose(np.triu(weights, k=1), 0.0)
print(output.shape)  # (4, 6)

Subtracting the row maximum makes softmax more numerically stable without changing its mathematical result. The mask used here always leaves at least one finite score in each row. A general mask that hides an entire row needs an explicit policy: normalizing an all-masked row is undefined.

Read the axes before the code

Q @ K.T mixes positions into an $n\times n$ score matrix. weights @ V then mixes value vectors across positions. Softmax runs over keys, the final axis of the score matrix.

From one head to many

Multi-head attention runs several learned attention heads in parallel, concatenates their outputs, and applies an output projection. Each head can learn a different compatibility function. The benefit is additional representational flexibility; it does not imply that every head corresponds to a stable linguistic concept.¹

The original transformer describes multi-head attention in Section 3.2.2 of Attention Is All You Need. A full transformer also combines attention with position information, feed-forward networks, residual connections, and normalization.

What this explanation leaves open

The explicit score matrix has $n^2$ entries. That observation describes this straightforward implementation’s intermediate storage, not a requirement that every exact attention implementation materialize the entire matrix. It also does not describe the whole cost of a transformer layer.

At inference time, earlier tokens’ keys and values can be reused rather than recomputed. That changes the execution pattern without changing the causal attention rule. The distinction between a mathematical operation and its execution will return throughout training and inference.

References

• Vaswani, A. et al. (2017). Attention Is All You Need. Sections 3.2.1–3.2.3 introduce scaled dot-product, multi-head, and masked attention.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. A next reading on computing exact attention without storing the full attention matrix in high-bandwidth memory.

Footnotes

1. An attention weight tells us how much one value contributes to a particular head’s weighted sum. Explaining a model’s final prediction also requires accounting for value projections, output projections, other heads, residual paths, and later layers. ↩