%matplotlib inline
from d2l import torch as d2l
import torch
from torch import nn10.3 Multi-Head and Cross-Attention
Scaled dot-product attention assigns each query one distribution over the keys and returns one weighted average of the values. A single average cannot preserve several independently requested values. We demonstrate this limitation on a task for which any single head loses half of the target variance. Multi-head attention (Vaswani et al. 2017) computes several projected attention outputs and combines them linearly. With the usual choice of head dimensions, this changes the representation without changing the leading parameter or FLOP counts. We then distinguish self-attention and cross-attention as two input configurations of the same operation.
%matplotlib inline
from d2l import jax as d2l
from flax import nnx
import jax
from jax import numpy as jnp10.3.1 A One-Mixture Bottleneck with Position-Only Keys
Consider a task that exposes this limitation. Two source positions hold value vectors \(\mathbf{v}_1, \mathbf{v}_2 \in \mathbb{R}^d\), drawn independently at random for every example; their keys mark only the positions, not the content. A query must report both values: the target is the concatenation \(\mathbf{t} = [\mathbf{v}_1; \mathbf{v}_2] \in \mathbb{R}^{2d}\). Think of a pronoun that must simultaneously retrieve its antecedent and its governing verb.
A single head can only produce, for this query, softmax weights \((\alpha_1, \alpha_2)\) with \(\alpha_1 + \alpha_2 = 1\), the mixture \(\mathbf{m} = \alpha_1 \mathbf{v}_1 + \alpha_2 \mathbf{v}_2\), and then a linear readout \(\mathbf{W}_o \mathbf{m}\). Since the keys are fixed position markers, the weights cannot depend on the values: whatever the head learns, it applies the same mixture to every example. How well can it do?
Proposition. Consider one attention layer with two fixed, position-only keys, independent isotropic Gaussian values, no residual path, and a linear readout. For the copy-both task, the best achievable relative squared error of a single head is \(1/2\) for every softmax pair \((\alpha_1,\alpha_2)\).
Proof. Given the mixture \(\mathbf{m} = \alpha_1 \mathbf{v}_1 + \alpha_2 \mathbf{v}_2\), the best estimate of \(\mathbf{v}_1\) is the conditional mean \(\mathbb{E}[\mathbf{v}_1 \mid \mathbf{m}] = \tfrac{\alpha_1}{\alpha_1^2 + \alpha_2^2}\, \mathbf{m}\), which is linear, so a linear readout attains it. Its per-dimension error variance is \(1 - \tfrac{\alpha_1^2}{\alpha_1^2 + \alpha_2^2} = \tfrac{\alpha_2^2}{\alpha_1^2 + \alpha_2^2}\); the error for \(\mathbf{v}_2\) is \(\tfrac{\alpha_1^2}{\alpha_1^2 + \alpha_2^2}\). The two errors sum to exactly \(1\) per dimension, while the target’s total variance is \(2\) per dimension. Thus the relative squared error is \(1/2\) regardless of how the head splits its attention. \(\blacksquare\)
Changing the softmax weights cannot remove the error because one head supplies the readout with only a single mixture of the two values. The following experiment verifies the bound.
torch.manual_seed(0)
N, d = 10000, 16
V1, V2 = torch.randn(N, d), torch.randn(N, d)
target = torch.cat([V1, V2], dim=1)
def readout_error(M, target):
"""Best linear readout of target from M, relative error."""
W = torch.linalg.lstsq(M, target).solution
return ((M @ W - target).norm() / target.norm()).item()
for a in (0.5, 0.7, 0.9, 1.0):
M = a * V1 + (1 - a) * V2
print(f'single head, weights ({a:.1f}, {1 - a:.1f}): '
f'relative error {readout_error(M, target):.3f}')
print(f'theory: sqrt(1/2) = {0.5 ** 0.5:.3f}')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.706
single head, weights (1.0, 0.0): relative error 0.707
theory: sqrt(1/2) = 0.707N, d = 10000, 16
key1, key2 = jax.random.split(jax.random.key(0))
V1, V2 = jax.random.normal(key1, (N, d)), jax.random.normal(key2, (N, d))
target = jnp.concatenate([V1, V2], axis=1)
def readout_error(M, target):
"""Best linear readout of target from M, relative error."""
W = jnp.linalg.lstsq(M, target)[0]
return float(jnp.linalg.norm(M @ W - target) / jnp.linalg.norm(target))
for a in (0.5, 0.7, 0.9, 1.0):
M = a * V1 + (1 - a) * V2
print(f'single head, weights ({a:.1f}, {1 - a:.1f}): '
f'relative error {readout_error(M, target):.3f}')
print(f'theory: sqrt(1/2) = {0.5 ** 0.5:.3f}')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.707The error sits at \(\sqrt{1/2} \approx 0.707\) across the whole range, exactly as computed. Now give the model a second head: two sets of attention weights, their outputs concatenated before the readout.
for a in (1.0, 0.88):
M2 = torch.cat([a * V1 + (1 - a) * V2,
(1 - a) * V1 + a * V2], dim=1)
print(f'two heads, weights ({a:.2f}, {1 - a:.2f}) and '
f'({1 - a:.2f}, {a:.2f}): '
f'relative error {readout_error(M2, target):.1e}')two heads, weights (1.00, 0.00) and (0.00, 1.00): relative error 1.4e-06
two heads, weights (0.88, 0.12) and (0.12, 0.88): relative error 6.0e-07for a in (1.0, 0.88):
M2 = jnp.concatenate([a * V1 + (1 - a) * V2,
(1 - a) * V1 + a * V2], axis=1)
print(f'two heads, weights ({a:.2f}, {1 - a:.2f}) and '
f'({1 - a:.2f}, {a:.2f}): '
f'relative error {readout_error(M2, target):.1e}')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-07Two heads reduce the error to the least-squares solver’s floating-point precision. The second row does so with soft attention weights \((0.88, 0.12)\). The two heads need distinct rather than one-hot weightings: they supply two mixtures, the readout inverts the \(2 \times 2\) mixing matrix, and both values are recovered exactly (the inversion is well conditioned as long as the heads differ appreciably; the exercises probe what happens as they collapse toward each other). Figure 10.3.1 summarizes the geometry.
The bound is a fact about this deliberately restricted interface — one layer, with value-blind keys that mark position only — not a universal law: it lifts once a one-head model gets content-dependent keys, more depth, or a residual connection, and the copy-both separation from two heads narrows accordingly.
Multi-head attention addresses this limitation by assigning each query several independently parameterized attention distributions and recombining their outputs with a learned projection.
10.3.2 Multi-Head Attention
10.3.2.1 Extending to \(h\) Heads
Rather than hand-designing the heads as above, we learn them. Given a query \(\mathbf{q} \in \mathbb{R}^{d}\), keys \(\mathbf{k} \in \mathbb{R}^{d}\), and values \(\mathbf{v} \in \mathbb{R}^{d}\), head \(i\) (\(i = 1, \ldots, h\)) first projects them into its own subspace and then runs ordinary scaled dot product attention \(f\) from Section 10.2:
\[ \mathbf{h}_i = f\big(\mathbf W_i^{(q)}\mathbf q,\; \mathbf W_i^{(k)}\mathbf k,\; \mathbf W_i^{(v)}\mathbf v\big) \in \mathbb R^{p}, \tag{10.3.1}\]
with learnable \(\mathbf W_i^{(q)}, \mathbf W_i^{(k)}, \mathbf W_i^{(v)} \in \mathbb R^{p \times d}\). Because each head scores with its own projected queries and keys, each head attends according to its own notion of relevance, and each moves its own projection of the values — the two independent mixtures of the construction above, now learned and generalized to \(h\). The output is a learned linear recombination of the concatenated heads:
\[ \mathbf W_o \begin{bmatrix}\mathbf h_1\\\vdots\\\mathbf h_h\end{bmatrix} \in \mathbb{R}^{d}, \qquad \mathbf W_o \in \mathbb R^{d \times hp}. \tag{10.3.2}\]
Figure 10.3.2 shows the layout.
10.3.2.2 Parameter and FLOP Counts
The standard dimension choice makes the heads collectively no more expensive than one big head: set the per-head width to \(p = d/h\), so the \(h\) heads together produce \(hp = d\) numbers per token, exactly what a single head of width \(d\) would. For self-attention, write the sequence inputs as \(\mathbf{Q},\mathbf{K},\mathbf{V}\in\mathbb{R}^{n\times d}\); each of the four learned projections is a \(d\times d\) matrix, and each head receives matrices of shape \(n\times p\). The resulting cost is
\[ \underbrace{8nd^2}_{\textrm{projections}} \;+\; \underbrace{4n^2 d}_{\textrm{scores and mixing}} \quad \textrm{FLOPs,} \qquad 4d^2 \ \textrm{parameters,} \tag{10.3.3}\]
independent of \(h\): the four projection matrices are \(d \times d\) regardless of how their outputs are sliced into heads, and each head’s score and mixing matmuls shrink by the factor \(h\) that their count grows by (counting one multiply–add as two FLOPs). In this arithmetic, the number of heads changes the representation but not the leading operation count. Realized cost is not exactly constant because increasing the number of heads produces more, smaller matrix multiplications and softmaxes, whose kernel efficiency may differ. What does change with \(h\) is the shape of the attention map — \(h\) distributions of \(n\) weights per query instead of one.
10.3.2.3 Implementation
The per-head dimension choice also yields a standard implementation: compute one query projection of width \(d\), then reshape it into \(h\) heads of width \(d/h\), and fold the head axis into the batch axis. The attention code from Section 10.2 then runs all heads of all sequences as one big batched matrix multiplication, with no loop over heads. Within the forward pass, valid_lens is repeated \(h\) times so that each head of a sequence sees that sequence’s mask.
class MultiHeadAttention(d2l.Module):
"""Multi-head attention."""
def __init__(self, num_hiddens, num_heads, dropout, bias=False, **kwargs):
super().__init__()
assert num_hiddens % num_heads == 0, 'heads must divide num_hiddens'
self.num_heads = num_heads
self.attention = d2l.DotProductAttention(dropout)
self.W_q = nn.LazyLinear(num_hiddens, bias=bias)
self.W_k = nn.LazyLinear(num_hiddens, bias=bias)
self.W_v = nn.LazyLinear(num_hiddens, bias=bias)
self.W_o = nn.LazyLinear(num_hiddens, bias=bias)
def forward(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 = torch.repeat_interleave(
valid_lens, repeats=self.num_heads, dim=0)
# Shape of output: (batch_size * num_heads, no. of queries,
# num_hiddens / num_heads)
output = self.attention(queries, keys, values, valid_lens)
# Shape of output_concat: (batch_size, no. of queries, num_hiddens)
output_concat = self.transpose_output(output)
return self.W_o(output_concat)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_weightsThe two transpositions do the head bookkeeping: transpose_qkv moves from (batch, length, num_hiddens) to (batch * num_heads, length, num_hiddens / num_heads), and transpose_output reverses it after the attention call.
@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 = X.permute(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 = X.permute(0, 2, 1, 3)
return X.reshape(X.shape[0], X.shape[1], -1)@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))A shape check: with 5 heads and 100 hidden units, batches of 2 sequences with 4 queries against 6 key–value pairs come back as (batch_size, num_queries, num_hiddens), and the per-sequence valid_lens masking flows through unchanged.
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(attention(X, Y, Y, valid_lens),
(batch_size, num_queries, num_hiddens))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))And the accounting of Equation 10.3.3 in the flesh: the parameter count does not move as the head count sweeps from 1 to 8.
X = d2l.ones((2, 10, 256))
for num_heads in (1, 2, 4, 8):
attention = MultiHeadAttention(num_hiddens=256, num_heads=num_heads,
dropout=0)
attention(X, X, X, None) # Materialize the lazily initialized layers
print(f'{num_heads} heads: '
f'{sum(p.numel() for p in attention.parameters())} parameters')1 heads: 262144 parameters
2 heads: 262144 parameters
4 heads: 262144 parameters
8 heads: 262144 parametersfor num_heads in (1, 2, 4, 8):
attention = MultiHeadAttention(num_hiddens=256, num_heads=num_heads,
dropout=0)
num_params = sum(p.size for p in jax.tree.leaves(
nnx.state(attention, nnx.Param)))
print(f'{num_heads} heads: {num_params} parameters')1 heads: 262144 parameters
2 heads: 262144 parameters
4 heads: 262144 parameters
8 heads: 262144 parameters10.3.3 Self-Attention and Cross-Attention
MultiHeadAttention takes three sequence arguments: queries, keys, and values. Self-attention and cross-attention differ only in the sources of these arguments.
10.3.3.1 Self-Attention
In self-attention (Lin et al. 2017; Vaswani et al. 2017), one sequence supplies all three arguments: every token emits a query against every token’s key, and each output position is a value mixture from the same sequence. Given inputs \(\mathbf{x}_1, \ldots, \mathbf{x}_n \in \mathbb{R}^d\), the layer returns an equally long sequence \(\mathbf{y}_1, \ldots, \mathbf{y}_n\) with
\[ \mathbf{y}_i = f\big(\mathbf{x}_i,\; (\mathbf{x}_1, \mathbf{x}_1), \ldots, (\mathbf{x}_n, \mathbf{x}_n)\big) \in \mathbb{R}^d, \tag{10.3.4}\]
where \(f\) is multi-head attention pooling. Each token’s new representation combines information from the whole sequence in a single step. Transformers stack these layers because the output and input shapes agree.
num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_heads, 0.5)
batch_size, num_queries, valid_lens = 2, 4, d2l.tensor([3, 2])
X = d2l.ones((batch_size, num_queries, num_hiddens))
d2l.check_shape(attention(X, X, X, valid_lens),
(batch_size, num_queries, num_hiddens))num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_heads, 0.5)
batch_size, num_queries, valid_lens = 2, 4, d2l.tensor([3, 2])
X = d2l.ones((batch_size, num_queries, num_hiddens))
d2l.check_shape(nnx.view(attention, deterministic=True)(
X, X, X, valid_lens)[0],
(batch_size, num_queries, num_hiddens))10.3.3.2 Cross-Attention
In cross-attention, the queries come from sequence \(\mathbf{A}\), while keys and values come from sequence \(\mathbf{B}\). The output has the length of \(\mathbf{A}\) and aggregates information from \(\mathbf{B}\). This is the wiring of the original encoder–decoder transformer, where each target-language position queries the source sentence; of a vision–language model, where text tokens query image patches; and of the historical alignment models that introduced learned attention (Section 10.2). These cases use the same function with different argument sources:
A = d2l.ones((batch_size, 4, num_hiddens))
B = d2l.ones((batch_size, 9, num_hiddens))
d2l.check_shape(attention(A, B, B, None),
(batch_size, 4, num_hiddens))A = d2l.ones((batch_size, 4, num_hiddens))
B = d2l.ones((batch_size, 9, num_hiddens))
d2l.check_shape(nnx.view(attention, deterministic=True)(
A, B, B, None)[0],
(batch_size, 4, num_hiddens))10.3.3.3 Visualizing Cross-Attention Alignment
To visualize cross-attention, assign every letter of the alphabet a random embedding, and let the letters of one word query the letters of another through plain scaled dot product attention — no learned projections. Matching letters share an embedding, so their score concentrates near \(\sqrt{d}\) after scaling, while unrelated pairs score near zero; softmax turns this gap into a readable alignment.
words = 'attention', 'translation'
letters = sorted(set(''.join(words)))
torch.manual_seed(0)
emb = torch.randn(len(letters), 32)
queries, keys = (emb[[letters.index(c) for c in w]][None] for w in words)
attention = d2l.DotProductAttention(dropout=0)
attention(queries, keys, keys, None)
d2l.show_heatmaps(attention.attention_weights[None],
xlabel='keys: t r a n s l a t i o n',
ylabel='queries: a t t e n t i o n',
figsize=(4, 3.5), cmap='Blues')
words = 'attention', 'translation'
letters = sorted(set(''.join(words)))
emb = jax.random.normal(jax.random.key(0), (len(letters), 32))
queries, keys = (emb[jnp.array([letters.index(c) for c in w])][None]
for w in words)
attention = d2l.DotProductAttention(dropout=0)
output, attention_weights = attention(queries, keys, keys, None)
d2l.show_heatmaps(attention_weights[None],
xlabel='keys: t r a n s l a t i o n',
ylabel='queries: a t t e n t i o n',
figsize=(4, 3.5), cmap='Blues')
Read the map row by row (queries are the letters of “attention”, keys the letters of “translation”). Three regimes appear. The i and o queries find their unique partners and attend sharply. The a, t, and n queries find two copies each and split their weight: a single head cannot choose between identical keys, the averaging of the first section in miniature. And the e query, whose letter does not occur in “translation” at all, spreads its weight diffusely because softmax assigns all probability mass across the available keys even when none matches. Trained models show all three regimes too, which is one reason reading attention maps as explanations requires care. A diffuse row may indicate that no key is relevant rather than that all keys are relevant.
10.3.4 Summary
A single attention head gives each query one distribution over the keys and therefore one value mixture. On the value-blind copy-both task, any single head loses half the target variance. Multi-head attention computes \(h\) projected mixtures and combines them linearly. With per-head width \(d/h\), the parameter and FLOP counts are independent of \(h\); implementations can fold the head axis into the batch axis. Self-attention draws queries, keys, and values from one sequence. Cross-attention draws queries from one sequence and keys and values from another, with output length equal to the query length.
10.3.5 Exercises
- [code] Nonlinear readouts. ● For Gaussian values, the proof of the proposition in Section 10.3.1 shows that the conditional mean \(\mathbb{E}[\mathbf{v}_1 \mid \mathbf{m}]\) is linear in the mixture. Hence no readout, linear or not, improves on the relative error \(\sqrt{1/2}\) that
readout_errormeasures.- Train a small multilayer perceptron to predict
targetfrom the mixture \(\mathbf{m} = a\,\mathbf{v}_1 + (1 - a)\,\mathbf{v}_2\) formed fromV1andV2, for \(a = 0.5\) and \(a = 0.7\). Compare its relative error on held-out samples with that of the linear readout. - Now draw the values uniformly from \(\{-1, +1\}^d\). For \(a = 0.5\), compute \(\mathbb{E}[\mathbf{v}_1 \mid \mathbf{m}]\) coordinate by coordinate, show that it is again linear in \(\mathbf{m}\), and conclude that the error stays at \(\sqrt{1/2}\) for every readout.
- For \(a = 0.7\) and the same discrete values, show that the map \((v_{1j}, v_{2j}) \mapsto m_j\) is injective. Train the perceptron again and report its held-out error. Explain why this escape from the bound does not contradict the proposition.
- Train a small multilayer perceptron to predict
- [code] Nearly identical heads. ● Two heads with weights \((0.5 + \epsilon, 0.5 - \epsilon)\) and \((0.5 - \epsilon, 0.5 +
\epsilon)\) supply two distinct mixtures, so the linear readout inverts them exactly for every \(\epsilon > 0\).
- Write the \(2 \times 2\) mixing matrix, compute its eigenvalues and eigenvectors, and express its condition number as a function of \(\epsilon\).
- Add independent Gaussian noise of standard deviation \(\sigma = 0.01\) to both mixtures and measure the error with
readout_errorfor \(\epsilon \in \{0.4, 0.1, 0.03, 0.01, 0.003, 0.001\}\). Plot the error against \(\epsilon\) on logarithmic axes and read off the slope where the curve is straight. - Rewrite the two noisy mixtures as noisy observations of \(\mathbf{v}_1 + \mathbf{v}_2\) and of \(\mathbf{v}_1 - \mathbf{v}_2\). Use this to explain the slope, to locate the \(\epsilon\) below which the second head stops helping, and to predict how that point moves with \(\sigma\). Check the prediction by repeating the sweep at \(\sigma = 0.001\), extended down to \(\epsilon = 10^{-5}\).
- What the bound counts. ● The proposition of Section 10.3.1 fixes the softmax weights, but its error bound rests on a weaker fact: the readout receives a single vector \(\mathbf{m} \in \mathbb{R}^d\) and must produce the target \(\mathbf{t} \in \mathbb{R}^{2d}\).
- Let \(\mathbf{m} \in \mathbb{R}^d\) be any function of \(\mathbf{v}_1\) and \(\mathbf{v}_2\), including a mixture whose weights depend on the values, and let \(\mathbf{W}_o \in \mathbb{R}^{2d \times d}\) be any matrix. Show that \[\mathbb{E}\,\|\mathbf{W}_o \mathbf{m} - \mathbf{t}\|^2 \geq \tfrac{1}{2}\, \mathbb{E}\,\|\mathbf{t}\|^2.\]
- The two-head construction of Figure 10.3.1 concatenates two mixtures, so its readout receives \(2d\) numbers. A multi-head layer instead returns a vector of width \(d\), by Equation 10.3.2. Show that no number of heads lets such a layer, followed by a linear readout, solve the copy-both task with values in \(\mathbb{R}^d\).
- Now let \(\mathbf{v}_1, \mathbf{v}_2 \in \mathbb{R}^c\) with \(c < d\), padded with zeros to width \(d\) before they enter the layer. Find the largest \(c\) for which a layer with two heads of width \(p = d/2\) solves the task exactly. Give the attention weights of each head, value projections \(\mathbf{W}_i^{(v)}\), and an output projection \(\mathbf{W}_o\) that do so.
- Generalize the proposition to \(k\) source positions with independent isotropic Gaussian values \(\mathbf{v}_1, \ldots, \mathbf{v}_k \in \mathbb{R}^c\) and the target \([\mathbf{v}_1; \ldots; \mathbf{v}_k]\). Each of \(h\) heads applies fixed weights to the values, the weights of head \(i\) form row \(i\) of a matrix \(\mathbf{A} \in \mathbb{R}^{h \times k}\), and a linear readout sees all \(h\) mixtures. Show that the best relative squared error is \[1 - \frac{\operatorname{rank}(\mathbf{A})}{k}, \tag{10.3.5}\] and that softmax weights, whose entries are all positive, can attain \(\operatorname{rank}(\mathbf{A}) = \min(h, k)\).
- [code] Averages under key noise. ● Let \(m = 4\) keys \(\mathbf{k}_i \in \mathbb{R}^8\) be drawn as \(\mathbf{k}_i \sim
\mathcal{N}(\boldsymbol{\mu}_i, \boldsymbol{\Sigma}_i)\) with orthonormal means \(\boldsymbol{\mu}_i\), let the values \(\mathbf{v}_1, \ldots,
\mathbf{v}_4\) be the standard basis vectors of \(\mathbb{R}^4\), and let the score be \(\mathbf{q}^\top \mathbf{k}_i\). When every covariance is \(\beta \mathbf{I}\) with vanishingly small \(\beta\), the query \(c\,(\boldsymbol{\mu}_a + \boldsymbol{\mu}_b)\) with a large scale \(c\) returns an output \(\mathbf{o}\) close to \(\tfrac{1}{2}(\mathbf{v}_a +
\mathbf{v}_b)\).
- Now let \(\boldsymbol{\Sigma}_a = \beta \mathbf{I} + \tfrac{1}{2} \boldsymbol{\mu}_a \boldsymbol{\mu}_a^\top\), so that \(\mathbf{k}_a \approx (1 + \gamma)\, \boldsymbol{\mu}_a\) with \(\gamma \sim \mathcal{N}(0, \tfrac{1}{2})\), and leave the other covariances at \(\beta \mathbf{I}\). Write the ratio of the weights on \(\mathbf{k}_a\) and \(\mathbf{k}_b\) as a function of \(\gamma\), and describe the distribution of \(\mathbf{o}\) across resamples of the keys.
- Replace the single query by two queries \(\mathbf{q}_1\) and \(\mathbf{q}_2\), each with its own softmax and output \(\mathbf{o}_1\) or \(\mathbf{o}_2\), and return \(\tfrac{1}{2} (\mathbf{o}_1 + \mathbf{o}_2)\). Choose the queries so that the output tends to \(\tfrac{1}{2}(\mathbf{v}_a + \mathbf{v}_b)\) as \(c \to \infty\) for most values of \(\gamma\), and identify the values of \(\gamma\) for which it does not.
- Simulate three settings: isotropic covariances with the query \(c\,(\boldsymbol{\mu}_a + \boldsymbol{\mu}_b)\), the anisotropic \(\mathbf{k}_a\) with that query, and the anisotropic \(\mathbf{k}_a\) with the two queries. Use \(\beta = 10^{-4}\), \(c \in \{10, 100\}\), and \(1000\) resamples of the keys. Report the median of \(\|\mathbf{o} - \tfrac{1}{2}(\mathbf{v}_a + \mathbf{v}_b)\|\) and the fraction of resamples on which it exceeds \(0.1\). Explain why raising \(c\) from \(10\) to \(100\) hurts the single query even when the covariances are isotropic.
- Repeat the two single-query settings with the query \(-c \sum_{i \notin \{a, b\}} \boldsymbol{\mu}_i\), which lowers the scores of the other keys instead of raising those of \(\mathbf{k}_a\) and \(\mathbf{k}_b\). Show that its output does not depend on \(\gamma\), and determine which failure of the query \(c\,(\boldsymbol{\mu}_a + \boldsymbol{\mu}_b)\) remains. Explain why the two queries avoid it.
- [code] Head pruning. Removing heads from a trained layer at inference time (Michel et al. 2019) requires a measure of each head’s importance. The copy-both task supplies a case in which importance is not additive.
- Form three heads with weights \((1, 0)\), \((0.88, 0.12)\), and again \((0.88, 0.12)\), concatenate their mixtures, and evaluate
readout_errorafter removing each single head and after removing each pair of heads. - Define a head’s importance as the increase in error when it alone is removed. A one-shot procedure computes every importance once and removes all heads whose importance is below the tolerance \(\tau = 0.01\). Which heads does it remove, and which error results?
- A greedy procedure instead removes the least important head, re-estimates the importance of the remaining heads, and repeats until every remaining head has an importance of at least \(\tau\). Find the heads it keeps and the resulting error. Without the stopping rule, count the evaluations of
readout_errorthat it needs to prune \(h\) heads down to one, and compare the count with the number that the one-shot procedure needs.
- Form three heads with weights \((1, 0)\), \((0.88, 0.12)\), and again \((0.88, 0.12)\), concatenate their mixtures, and evaluate
- [code] Narrow heads. With \(p = d/h\), the leading count of Equation 10.3.3 does not depend on \(h\), so the head count looks free up to \(h = d\). Consider one head with scores \(\mathbf{q}^\top \mathbf{k}_j / \sqrt{p}\) over keys \(\mathbf{k}_1,
\ldots, \mathbf{k}_m \in \mathbb{R}^p\) and no mask. Call a query sharp on key \(j\) when its weights concentrate on \(\mathbf{k}_j\) as the query is scaled by a factor that grows without bound.
- For \(p = 1\) and distinct scalar keys, show that every query’s weights are monotone in \(k_j\). Conclude that the queries of such a head can be sharp only on the largest or the smallest key.
- Show that for \(p = 2\) the keys and queries can be chosen so that each of \(n\) queries is sharp on a key of its own choosing, for every assignment of queries to keys.
- Use these two results to state which attention patterns a layer gives up as \(h\) approaches \(d\).
- Time the forward pass of
MultiHeadAttentionwith dropout 0 for self-attention with \(d = 256\) and \(n = 512\) on a batch of \(8\) sequences for \(h \in \{1, 2, 4, 8, 16, 32, 64\}\). Make a warm-up call first and wait for the result before stopping the clock. Report the ratio of the slowest to the fastest configuration. Identify the cost that grows with \(h\) and that Equation 10.3.3 leaves out, and check that it accounts for the trend.
- Cost of cross-attention. Equation 10.3.3 counts self-attention over \(n\) tokens.
- Redo the count for cross-attention with \(n\) queries from \(\mathbf{A}\) and \(m\) keys and values from \(\mathbf{B}\). Identify which terms grow with \(n\), which with \(m\), and which with \(nm\).
- A decoder generates its output one token at a time and cross-attends from each new token to a fixed source of \(m\) tokens. Count the FLOPs per generated token when the projected source keys and values are recomputed at every step and when they are computed once and stored. Give the ratio for \(d = 512\) and \(m = 1000\).
- Multi-query and grouped-query attention. Multi-query attention (Shazeer 2019) shares one key head and one value head of width \(p = d/h\) across all \(h\) query heads. Grouped-query attention (Ainslie et al. 2023) shares them within each of \(g\) groups of \(h/g\) query heads.
- Redo the parameter count and the FLOP count of Equation 10.3.3 for both schemes. Which of the two terms changes, and by what factor?
- During autoregressive generation every past key and value is cached. Compute the cache size per token and per layer in bytes for \(d = 4096\), \(h = 32\), and 16-bit numbers under the three schemes, with \(g = 8\) for grouped-query attention.
- Show that two query heads that share one key head and one value head can still produce two distinct mixtures of the same projected values, as the copy-both construction requires. Then state the limit that sharing places on the information the layer can move from each source token to its output.
- [code] Heads as slices.
MultiHeadAttentioncomputes one projection of width \(d\) for each of the queries, keys, and values and splits it into heads withtranspose_qkv.- Determine from
transpose_qkvwhich output features ofW_q,W_k, andW_vfeed head \(i\), and which input features ofW_oreceive that head’s output. Write a forward pass that loops over the heads, applies the sliced projections, and callsd2l.DotProductAttentiononce per head. Verify on random inputs withvalid_lensthat it matchesMultiHeadAttentionwith dropout 0. - Replace the repetition of
valid_lensin the forward pass ofMultiHeadAttentionwith a tiling that copies the whole vectornum_headstimes. Find inputs for which the output changes, and explain which mask each head of each sequence receives. Characterize thevalid_lensfor which no test can detect the change.
- Determine from
- Permuting the source. In cross-attention, the queries come from \(\mathbf{A} \in \mathbb{R}^{n \times d}\) and the keys and values from \(\mathbf{B} \in \mathbb{R}^{m \times d}\), neither carrying positional information. Let \(\boldsymbol{\Pi}\) be an \(m \times m\) permutation matrix.
- Show that replacing \(\mathbf{B}\) by \(\boldsymbol{\Pi}\mathbf{B}\) in both the key and the value argument leaves the output of
MultiHeadAttentionunchanged when no mask is applied, and that replacing it in the key argument alone changes the output in general. - Now let
valid_lensmark only the first \(\ell\) rows of \(\mathbf{B}\) as valid. Which permutations leave the output unchanged whenvalid_lensstays fixed? - Self-attention without positional information is permutation equivariant: permuting the rows of its input permutes the rows of its output in the same way. A translation decoder cross-attends to the output of an encoder built from such self-attention layers. Show that the decoder’s outputs are identical for the source sentences “dog bites man” and “man bites dog”, and state where word order must enter the model for cross-attention to use it.
- Show that replacing \(\mathbf{B}\) by \(\boldsymbol{\Pi}\mathbf{B}\) in both the key and the value argument leaves the output of
- [code] Random projections and readable maps. The alignment of Section 10.3.3.3 uses shared random embeddings and no learned projections.
- Pass the letter embeddings of the two words as queries, keys, and values through a freshly initialized
MultiHeadAttentionwith 4 heads, 32 hidden units, and dropout 0, and plot the attention weights of each head. - Which properties of the plain map hold in every head, whatever the projections, and which are lost? Explain both from the score that head \(i\) assigns to a query letter with embedding \(\mathbf{e}\) and a key letter with embedding \(\mathbf{e}'\), \[\mathbf{e}^\top \mathbf{W}_i^{(q)\top} \mathbf{W}_i^{(k)} \mathbf{e}' / \sqrt{p}. \tag{10.3.6}\]
- State what this implies for reading the attention maps of a trained network whose projections have not been inspected.
- Pass the letter embeddings of the two words as queries, keys, and values through a freshly initialized
- [extended] Learned copy-both heads. ● Train
MultiHeadAttentionwith 16 hidden units and dropout 0 on a copy-both task whose target fits the layer’s width. Each example has one query and two key–value pairs, all vectors in \(\mathbb{R}^{16}\). The values are \([\mathbf{v}_1; \mathbf{0}]\) and \([\mathbf{v}_2; \mathbf{0}]\), with \(\mathbf{v}_1, \mathbf{v}_2 \in \mathbb{R}^4\) drawn independently from \(\mathcal{N}(\mathbf{0}, \mathbf{I})\) for every example. The two keys are the first two standard basis vectors, so they mark position only, and the query is the third standard basis vector in every example. The prediction is the first 8 coordinates of the layer’s output, and the target is \([\mathbf{v}_1; \mathbf{v}_2]\). By the proposition of Section 10.3.1, one head cannot bring the relative error below \(\sqrt{1/2}\) on this task, whereas two heads with suitable projections bring it to zero.- Train the layer with Adam on the squared error for \(h \in \{1, 2, 4, 8\}\), and report the relative error on held-out examples.
- Plot the attention weights of each head after training, and sort the heads into those that end near one-hot, those that are soft but distinct from all others, and those that nearly duplicate another head. Every head starts with weights close to \((0.5, 0.5)\), where the mixing matrix of any two heads is nearly singular. Explain what this start implies for the early phase of training and for the duplicates you observe.