%matplotlib inline
from d2l import torch as d2l
import torch
from torch import nn10.3 Multi-Head and Cross-Attention
Scaled dot product attention, as built in Section 10.2, gives each query exactly one probability distribution over the keys, and returns the correspondingly weighted average of the values. One distribution per query is a real restriction: a token in a sentence participates in several relations at once — it agrees with a verb, refers back to an entity, sits inside a phrase — and a single weighted average must blend all of them into one vector. This section first makes the restriction precise: we construct a task on which a single attention head provably loses half of what it is asked to report, however it allocates its weights, and we verify the bound numerically. Multi-head attention (Vaswani et al. 2017) removes the restriction at essentially no extra cost, and we implement it with one batched matrix multiplication for all heads — a common implementation strategy (fused attention kernels keep an explicit head axis instead). Finally we separate what is often conflated: the attention function is one piece of code, and self-attention and cross-attention are merely two ways of wiring its inputs.
%matplotlib inline
from d2l import jax as d2l
from flax import nnx
import jax
from jax import numpy as jnp10.3.1 One Head Must Average
Here is the smallest task that defeats a single head. 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. For independent isotropic Gaussian values, the best achievable relative squared error of a single head on the copy-both task is \(1/2\) — for every choice of \((\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 add to exactly \(1\) per dimension, against the target’s total of \(2\): half the variance is gone, no matter how the head splits its attention. \(\blacksquare\)
The knob that softmax controls — where to put the weight — cannot help, because the failure is structural: one head hands the readout a single mixture of the values, and the task needs two. Let’s watch the bound hold.
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 collapse the error by orders of magnitude, down to the least-squares solver’s floating-point noise — and note that the second row does it with soft attention, weights \((0.88, 0.12)\). Sharpness is not the point. The two heads supply two different mixtures, the readout inverts the \(2 \times 2\) mixing matrix, and both values come out 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.
This is the design brief for multi-head attention: give each query several independently parametrized attention distributions — several views of the same key–value set — and let a learned output projection recombine them.
10.3.2 Multi-Head Attention
10.3.2.1 From One Head 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 Same FLOPs, More Views
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 a sequence of \(n\) tokens the cost is then
\[ \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 is a free dial that buys representational diversity, not extra compute — realized cost is not quite as flat, since more heads mean more, smaller matmuls and softmaxes for the kernels to keep efficient. 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 trick: 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, values — and nothing in its code cares where they come from. The two wirings that dominate practice differ only in what we pass.
10.3.3.1 One Sequence Querying Itself
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 is assembled from the whole sequence in a single step — this is the layer that transformers stack, and its output shape equals its input shape, which is what makes stacking possible.
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 One Sequence Querying Another
In cross-attention, one sequence asks and another answers: the queries come from sequence \(\mathbf{A}\), while keys and values come from sequence \(\mathbf{B}\). The output has the length of \(\mathbf{A}\) (one answer per question) and the information of \(\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 attention grew out of (Section 10.2). Same function, one changed argument:
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 An Alignment You Can Read
To see cross-attention doing something rather than merely reshaping, we strip it to its core. 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: softmax must hand out probability mass somewhere even when nothing matches. Trained models show all three regimes too, which is one reason reading attention maps as explanations requires care — a diffuse row may mean “nothing relevant”, not “everything relevant”.
10.3.4 Summary
A single attention head gives each query one softmax distribution over the keys and hence one convex mixture of the values; on a task that asks one query to report two values through position-only keys — so the weights cannot adapt to the values — any single head loses half the target variance, however it allocates its weights. Multi-head attention fixes this by running \(h\) independently projected attention heads and linearly recombining their concatenated outputs. With the per-head width set to \(d/h\), parameters and FLOPs are independent of \(h\) — heads buy diversity of views, not extra compute — and the implementation folds the head axis into the batch axis so that all heads run as one batched matmul. The attention function is indifferent to where its arguments come from: self-attention feeds one sequence as queries, keys, and values alike (output length = input length, the stackable case), while cross-attention lets one sequence query another (output length = query length), which is the encoder–decoder and multimodal wiring.
10.3.5 Exercises
- In the copy-both construction, the best linear readout loses half the variance. Does a nonlinear readout help? For Gaussian values the conditional mean is linear, so it cannot; check this numerically by training a small MLP readout on \(\mathbf{m}\) with weights \((0.5, 0.5)\). Then repeat with values drawn uniformly from \(\{-1, +1\}^d\) — explain what changes and why.
- Two heads with weights \((0.5 + \epsilon, 0.5 - \epsilon)\) and \((0.5 - \epsilon, 0.5 + \epsilon)\) still invert exactly in theory. Add Gaussian noise of standard deviation \(0.01\) to the mixtures and measure the readout error as \(\epsilon \to 0\). Relate what you see to the condition number of the \(2 \times 2\) mixing matrix.
- Feed the letter-alignment example through a freshly initialized
MultiHeadAttentionwith 4 heads instead of plain dot product attention, and plot each head’s attention map. Are the maps still readable? What does this suggest about interpreting the attention maps of a network whose projections you have not inspected? - Derive Equation 10.3.3 by counting the matrix multiplications in
MultiHeadAttention. For \(d = 512\), at which sequence length \(n\) does the quadratic score-and-mix term overtake the projection term? Later sections of this chapter take this crossover as their starting point. - Multi-query attention shares a single key and value head across all \(h\) query heads, and grouped-query attention shares them within groups. How do the parameter count and the FLOPs count of Equation 10.3.3 change under each scheme? Why is the saving most valuable during autoregressive generation?
- Suppose you want to prune the least important heads of a trained multi-head attention layer to speed up inference. Design an experiment to measure how much each head matters. How would you guard against two heads that are individually prunable but not jointly?