class MultiHeadAttention(nnx.Module):
"""Multi-head attention."""
def __init__(self, num_hiddens, num_heads, dropout, bias=False, rngs=None):
assert num_hiddens % num_heads == 0, 'heads must divide num_hiddens'
rngs = nnx.Rngs(params=0, dropout=1) if rngs is None else rngs
self.num_hiddens, self.num_heads = num_hiddens, num_heads
self.attention = d2l.DotProductAttention(dropout, rngs=rngs)
self.W_q = nnx.Linear(num_hiddens, num_hiddens, use_bias=bias,
rngs=rngs)
self.W_k = nnx.Linear(num_hiddens, num_hiddens, use_bias=bias,
rngs=rngs)
self.W_v = nnx.Linear(num_hiddens, num_hiddens, use_bias=bias,
rngs=rngs)
self.W_o = nnx.Linear(num_hiddens, num_hiddens, use_bias=bias,
rngs=rngs)
def __call__(self, queries, keys, values, valid_lens):
# Shape of queries, keys, or values:
# (batch_size, no. of queries or key-value pairs, num_hiddens)
# Shape of valid_lens: (batch_size,) or (batch_size, no. of queries)
queries = self.transpose_qkv(self.W_q(queries))
keys = self.transpose_qkv(self.W_k(keys))
values = self.transpose_qkv(self.W_v(values))
if valid_lens is not None:
# On axis 0, copy the first item (scalar or vector) for num_heads
# times, then copy the next item, and so on
valid_lens = jnp.repeat(valid_lens, self.num_heads, axis=0)
# Shape of output: (batch_size * num_heads, no. of queries,
# num_hiddens / num_heads)
output, attention_weights = self.attention(
queries, keys, values, valid_lens)
# Shape of output_concat: (batch_size, no. of queries, num_hiddens)
output_concat = self.transpose_output(output)
# NNX idiom: return (output, weights); PyTorch returns output only
return self.W_o(output_concat), attention_weightsMulti-Head and Cross-Attention
Dive into Deep Learning · §10.3
Multi-head and cross-attention
why one head must average · heads for free · self- vs. cross-attention
One head, one mixture
The smallest task that defeats a single head
Two positions hold random values \mathbf{v}_1, \mathbf{v}_2; keys mark positions only. One query must report both: \mathbf{t} = [\mathbf{v}_1; \mathbf{v}_2].
A head can only output \mathbf{W}_o(\alpha_1 \mathbf{v}_1 + \alpha_2 \mathbf{v}_2) — one convex mixture per query.
Proposition. Best achievable relative squared error = 1/2, for every choice of (\alpha_1, \alpha_2).
The knob softmax controls — where the weight goes — cannot help. The failure is structural: one mixture, two requests.
The bound, numerically
single head, weights (0.5, 0.5): relative error 0.706
single head, weights (0.7, 0.3): relative error 0.706
single head, weights (0.9, 0.1): relative error 0.705
single head, weights (1.0, 0.0): relative error 0.705
theory: sqrt(1/2) = 0.707- Best linear readout via least squares, attention split swept from (0.5, 0.5) to (1, 0).
- Error pinned at \sqrt{1/2} \approx 0.707 across the whole range.
Two heads suffice
two heads, weights (1.00, 0.00) and (0.00, 1.00): relative error 6.8e-07
two heads, weights (0.88, 0.12) and (0.12, 0.88): relative error 6.8e-07- Error collapses to solver noise, orders of magnitude down.
- Second row: soft heads, weights (0.88, 0.12) — sharpness is not the point. Two different mixtures let the readout invert the 2 \times 2 mixing matrix.
Multi-head attention
Learn the heads: each gets its own projections, then ordinary scaled dot product attention
\mathbf{h}_i = f\big(\mathbf W_i^{(q)}\mathbf q, \mathbf W_i^{(k)}\mathbf k, \mathbf W_i^{(v)}\mathbf v\big), \qquad \textrm{output} = \mathbf W_o [\mathbf h_1; \ldots; \mathbf h_h].
Heads are free
Per-head width p = d/h ⇒ for n tokens:
\underbrace{8nd^2}_{\textrm{projections}} + \underbrace{4n^2 d}_{\textrm{scores, mixing}} \ \textrm{FLOPs}, \qquad 4d^2 \ \textrm{parameters}
— independent of h. Heads buy diversity of views, not FLOPs or parameters (realized kernel efficiency does vary with head count).
What does change with h: the attention map — h distributions per query instead of one.
Implementation
One d \times d projection each for Q, K, V; reshape into h heads; fold the head axis into the batch axis:
The reshape trick
(batch, len, d) → (batch, len, h, d/h) → (batch·h, len, d/h): the attention code sees heads as extra batch entries — no loop over heads.
@d2l.add_to_class(MultiHeadAttention)
def transpose_qkv(self, X):
"""Transposition for parallel computation of multiple attention heads."""
# Shape of input X: (batch_size, no. of queries or key-value pairs,
# num_hiddens). Shape of output X: (batch_size, no. of queries or
# key-value pairs, num_heads, num_hiddens / num_heads)
X = X.reshape((X.shape[0], X.shape[1], self.num_heads, -1))
# Shape of output X: (batch_size, num_heads, no. of queries or key-value
# pairs, num_hiddens / num_heads)
X = jnp.transpose(X, (0, 2, 1, 3))
# Shape of output: (batch_size * num_heads, no. of queries or key-value
# pairs, num_hiddens / num_heads)
return X.reshape((-1, X.shape[2], X.shape[3]))
@d2l.add_to_class(MultiHeadAttention)
def transpose_output(self, X):
"""Reverse the operation of transpose_qkv."""
X = X.reshape((-1, self.num_heads, X.shape[1], X.shape[2]))
X = jnp.transpose(X, (0, 2, 1, 3))
return X.reshape((X.shape[0], X.shape[1], -1))Shape and cost checks
num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_heads, 0.5)
batch_size, num_queries, num_kvpairs = 2, 4, 6
valid_lens = d2l.tensor([3, 2])
X = d2l.ones((batch_size, num_queries, num_hiddens))
Y = d2l.ones((batch_size, num_kvpairs, num_hiddens))
d2l.check_shape(nnx.view(attention, deterministic=True)(
X, Y, Y, valid_lens)[0],
(batch_size, num_queries, num_hiddens))1 heads: 262144 parameters
2 heads: 262144 parameters
4 heads: 262144 parameters
8 heads: 262144 parametersSelf- vs. cross-attention: two wirings
The function never asks where its arguments come from.
- Self-attention:
attention(X, X, X)— one sequence queries itself; output shape = input shape (the stackable case). - Cross-attention:
attention(A, B, B)— sequence A asks, sequence B answers; output length = query length.
An alignment you can read
Letters of “attention” query letters of “translation” through raw dot product attention on shared random embeddings:
Three regimes in one map
i,o: unique partner → sharp attention.a,t,n: two copies → weight splits — the averaging of the construction, in miniature.e: no partner → diffuse row; softmax must spend its mass somewhere.
A diffuse row may mean “nothing relevant”, not “everything relevant” — one reason attention maps are tricky as explanations.
Recap
- One head = one softmax distribution per query = one value mixture; the copy-both task loses half its variance, whatever the split.
- Multi-head: h independently projected heads, concatenated, recombined by \mathbf{W}_o — with p = d/h, same parameters and FLOPs as one head.
- Implementation folds heads into the batch axis: one batched matmul.
- Self-attention and cross-attention are the same code with different arguments; only the wiring differs.