Multi-Head Attention
Multiple attention patterns in parallel
21.1 What You Will Learn
Chapter 20 gave every position a query, a key and a value, and the resulting attention row was a single distribution over the sequence. That is one relationship per position, and one is often not enough. Take the sentence “the cat sat on the mat because it was tired”, where the word “it” has two jobs at once. It has to find “cat” to resolve what the pronoun refers to, and it has to find “was tired” to know what is being said about it. A single distribution has to split its weight between those two, and whichever it favors, the other is degraded.
Multi-head attention refuses the compromise by running several complete attention mechanisms side by side. Each head has its own W_Q, W_K and W_V and therefore its own idea of what relevance means, so one can specialize in syntax while another tracks position and a third handles coreference. Their outputs are concatenated and passed through one more matrix that mixes them back into a single representation.
The surprise, which we measure later on, is that this costs nothing extra. Splitting the model dimension across heads means each head works in a smaller space, and the arithmetic works out so that eight heads use exactly as many parameters as one. The chapter builds the split, runs two heads that visibly disagree, joins them, counts the parameters, and finishes by designing two heads whose specializations are legible in an ASCII map.
21.2 The Idea
Rather than one attention operating across all d_model dimensions, the model dimension is divided among h heads so that each head works in d_head = d_model / h dimensions. Every head carries its own three projection matrices, runs the complete score, normalize and blend from Chapter 20 independently, and produces a d_head wide output for every position. Those h outputs are laid end to end to recover d_model, then multiplied by an output matrix W_O.
The division is what makes the whole thing free. Each head is smaller by exactly the factor that there are more of them, so the total width of the projections is unchanged. What is gained is independence, since nothing forces two heads to agree and the loss has no reason to make them.
Figure 21-1 draws the arrangement with the sizes the programs in this chapter use, so d_model is 6 and two heads of 3 divide it. Read it bottom up. One input reaches both heads, each head owns its own Q, K and V projections and runs the attention function from Chapter 20 unchanged, and the two 3 wide results are laid end to end to make a 6 wide row again.
The values in the two head rows and in the concatenation are what the third program prints for the word The, and the concatenation is literally the two rows placed side by side with nothing combined. Only W_O at the top mixes them, which is why it is needed at all.
What each head ends up specializing in is decided by training rather than by design. In practice, heads in trained transformers have been found attending to the previous token. Others lock onto the first token of the sequence, or onto matching brackets and quotation marks, or onto syntactic dependents. None of that is built in. It emerges because several heads with nothing tying them together will drift apart if drifting apart lowers the loss.
21.3 Splitting into Heads
The word split invites a picture that is almost right and misleading in one important way. With d_model of 8 and two heads, it is tempting to imagine the vector cut in half, with head 0 taking the first four numbers and head 1 the last four.
That is not what happens, and this program shows both so the difference is visible rather than asserted.
/* 110_Split.c */
#include <stdio.h>
#define D_MODEL 8
#define H 2
#define D_HEAD (D_MODEL / H)
/* The naive picture, slice the vector into h pieces */
static void slice(const float x[D_MODEL], int head,
float out[D_HEAD])
{
int i;
for (i = 0; i < D_HEAD; i++)
out[i] = x[head * D_HEAD + i];
}
/* What actually happens, project the FULL vector
down to d_head with this head's own matrix */
static void project(const float x[D_MODEL],
const float W[D_HEAD][D_MODEL],
float out[D_HEAD])
{
int i, j;
for (i = 0; i < D_HEAD; i++) {
out[i] = 0;
for (j = 0; j < D_MODEL; j++)
out[i] += W[i][j] * x[j];
}
}
int main(void)
{
float x[D_MODEL] = { 0.1f, 0.5f, -0.3f, 0.8f,
-0.2f, 0.4f, 0.7f, -0.1f };
/* Head 0 reads mostly the front of the vector,
head 1 mostly the back, but both see all of it */
float W0[D_HEAD][D_MODEL] = {
{ 0.9f, 0.1f, 0.0f, 0.0f,
0.2f, 0.0f, 0.0f, 0.0f },
{ 0.0f, 0.8f, 0.2f, 0.0f,
0.0f, 0.1f, 0.0f, 0.0f },
{ 0.0f, 0.0f, 0.7f, 0.3f,
0.0f, 0.0f, 0.1f, 0.0f },
{ 0.1f, 0.0f, 0.0f, 0.9f,
0.0f, 0.0f, 0.0f, 0.2f },
};
float W1[D_HEAD][D_MODEL] = {
{ 0.2f, 0.0f, 0.0f, 0.0f,
0.9f, 0.1f, 0.0f, 0.0f },
{ 0.0f, 0.1f, 0.0f, 0.0f,
0.0f, 0.8f, 0.2f, 0.0f },
{ 0.0f, 0.0f, 0.1f, 0.0f,
0.0f, 0.0f, 0.9f, 0.3f },
{ 0.0f, 0.0f, 0.0f, 0.2f,
0.1f, 0.0f, 0.0f, 0.8f },
};
float s0[D_HEAD], s1[D_HEAD];
float p0[D_HEAD], p1[D_HEAD];
int i;
slice(x, 0, s0);
slice(x, 1, s1);
project(x, W0, p0);
project(x, W1, p1);
printf("d_model=%d, %d heads, d_head=%d\n\n",
D_MODEL, H, D_HEAD);
printf("Full vector: [");
for (i = 0; i < D_MODEL; i++)
printf("%+.1f%s", x[i],
i < D_MODEL-1 ? ", " : "]\n\n");
printf("The naive picture, slicing the vector\n");
printf(" head 0: [");
for (i = 0; i < D_HEAD; i++)
printf("%+.2f%s", s0[i],
i < D_HEAD-1 ? ", " : "]\n");
printf(" head 1: [");
for (i = 0; i < D_HEAD; i++)
printf("%+.2f%s", s1[i],
i < D_HEAD-1 ? ", " : "]\n\n");
printf("What really happens, each head projects\n");
printf("the whole vector through its own matrix\n");
printf(" head 0: [");
for (i = 0; i < D_HEAD; i++)
printf("%+.2f%s", p0[i],
i < D_HEAD-1 ? ", " : "]\n");
printf(" head 1: [");
for (i = 0; i < D_HEAD; i++)
printf("%+.2f%s", p1[i],
i < D_HEAD-1 ? ", " : "]\n\n");
printf("The projected values differ from the "
"slice\n");
printf("because every head can read every input\n");
printf("component. What is split is the output\n");
printf("width, not the input.\n");
return 0;
}
Figure 21-2 sets a slice of the input against what a head really does. The upper pair is the naive picture. Head 0 gets [+0.10, +0.50, −0.30, +0.80], which is literally the first half of the input, and head 1 gets the second half. Nothing has been computed, the numbers are just copied.
The lower pair is what a real head does. Head 0 produces [+0.10, +0.38, +0.10, +0.71] and head 1 produces [−0.12, +0.51, +0.57, +0.06], and neither matches its slice. The third component of head 0 is +0.10 where the slice said −0.30, because the projection matrix for that row reads component 6 as well as component 2, and component 6 of the input is +0.7.
The distinction matters because it decides what a head is able to see. Under the slicing picture, head 1 has no access to the first four components at all and could never learn anything that depends on them. Under the real arrangement, every head reads the entire input vector and merely writes a narrower output. What gets divided is the output width, not the input.
The confusion is worth clearing up because implementations often do slice, but they slice after a single wide projection rather than before. One matrix of d_model by d_model produces the full width and the result is then carved into h pieces, which is mathematically identical to h separate matrices of d_head by d_model and faster on hardware that likes big matrix multiplies.
21.4 One Head at a Time
A head is a function. It takes the sequence and one set of projection matrices and returns an attention distribution per position, which means running two heads is calling that function twice with different weights.
The two sets below are chosen rather than trained, since random weights produce near uniform rows that demonstrate nothing. Head 0 was built to link words to syntactic partners and head 1 to look one position back, which are two patterns that genuinely occur in trained transformers.
/* 111_Single_Head.c */
#include <stdio.h>
#include <math.h>
#define SEQ_LEN 4
#define D_MODEL 6
#define D_HEAD 3
static void softmax(float *x, int n)
{
float mx = x[0], s = 0;
int i;
for (i = 1; i < n; i++) if (x[i] > mx) mx = x[i];
for (i = 0; i < n; i++) {
x[i] = expf(x[i] - mx);
s += x[i];
}
for (i = 0; i < n; i++) x[i] /= s;
}
/* One attention head. Called once per head with that
head's own projection matrices. */
static void head(const float X[SEQ_LEN][D_MODEL],
const float W_Q[D_HEAD][D_MODEL],
const float W_K[D_HEAD][D_MODEL],
float weights[SEQ_LEN][SEQ_LEN])
{
float Q[SEQ_LEN][D_HEAD], K[SEQ_LEN][D_HEAD];
float scale = 1.0f / sqrtf((float)D_HEAD);
int i, j, k, m;
for (i = 0; i < SEQ_LEN; i++)
for (k = 0; k < D_HEAD; k++) {
Q[i][k] = 0;
K[i][k] = 0;
for (m = 0; m < D_MODEL; m++) {
Q[i][k] += W_Q[k][m] * X[i][m];
K[i][k] += W_K[k][m] * X[i][m];
}
}
for (i = 0; i < SEQ_LEN; i++) {
for (j = 0; j < SEQ_LEN; j++) {
weights[i][j] = 0;
for (k = 0; k < D_HEAD; k++)
weights[i][j] += Q[i][k] * K[j][k];
weights[i][j] *= scale;
}
softmax(weights[i], SEQ_LEN);
}
}
static void show(const char *title, const char *w[],
float W[SEQ_LEN][SEQ_LEN])
{
int i, j;
printf("%s\n", title);
printf(" ");
for (j = 0; j < SEQ_LEN; j++) printf("%-6s", w[j]);
printf("\n");
for (i = 0; i < SEQ_LEN; i++) {
printf(" %-6s ", w[i]);
for (j = 0; j < SEQ_LEN; j++)
printf("%.2f ", W[i][j]);
printf("\n");
}
printf("\n");
}
int main(void)
{
float X[SEQ_LEN][D_MODEL] = {
{ 1.0f, 0.2f, -0.3f, 0.5f, -0.1f, 0.8f },
{ 0.3f, 0.9f, 0.1f, -0.2f, 0.6f, 0.4f },
{ -0.2f, 0.4f, 0.8f, 0.3f, -0.5f, 0.1f },
{ 0.6f, -0.1f, 0.2f, 0.9f, 0.3f, -0.4f },
};
const char *words[SEQ_LEN] = {
"The", "cat", "sat", "down"
};
/* Head 0, chosen to pick out syntactic links */
float W_Q0[D_HEAD][D_MODEL] = {
{ -0.29f, +0.54f, +1.15f, +0.43f,
-0.91f, +0.28f },
{ +1.26f, +1.03f, +2.02f, +2.22f,
-3.58f, +2.71f },
{ +1.40f, +2.66f, +1.56f, +1.52f,
+4.53f, -2.21f },
};
float W_K0[D_HEAD][D_MODEL] = {
{ +0.54f, -0.09f, -0.11f, +0.47f,
-0.13f, +0.24f },
{ +0.06f, +0.63f, +0.14f, -0.19f,
+0.60f, +0.02f },
{ -0.01f, +0.28f, +0.76f, +0.55f,
-0.26f, -0.15f },
};
/* Head 1, chosen to look at the previous
position */
float W_Q1[D_HEAD][D_MODEL] = {
{ +2.02f, +2.52f, -0.38f, -0.18f,
+1.01f, +2.57f },
{ -0.73f, +1.36f, +2.86f, +1.06f,
-2.26f, +0.70f },
{ +1.14f, -0.16f, +0.91f, +2.38f,
+1.84f, -2.30f },
};
float W_K1[D_HEAD][D_MODEL] = {
{ +0.57f, -0.09f, -0.09f, +0.52f,
-0.09f, +0.19f },
{ +0.23f, +0.60f, +0.28f, +0.18f,
+0.88f, -0.34f },
{ -0.16f, +0.30f, +0.64f, +0.24f,
-0.50f, +0.15f },
};
float W0[SEQ_LEN][SEQ_LEN], W1[SEQ_LEN][SEQ_LEN];
head(X, W_Q0, W_K0, W0);
head(X, W_Q1, W_K1, W1);
show("Head 0 attention weights", words, W0);
show("Head 1 attention weights", words, W1);
printf("Same input, same code, "
"different weights.\n");
printf("Head 0 links each word to a syntactic\n");
printf("partner. Head 1 looks "
"one position back.\n");
printf("A single head has to pick one of these.\n");
return 0;
}
Figure 21-3 has both heads over the same four words. The two matrices disagree about almost everything, which is the point. Head 0 sends “The” to “cat” at 0.82, sends “cat” forward to “sat” at 0.67, and sends “sat” back to “cat” at 0.70. Read those three rows together and they trace determiner to noun, subject to verb, verb to subject.
Head 1 does something completely different and completely regular. “The” attends 0.62 to itself since there is nothing before it, “cat” attends 0.62 to “The”, and “sat” attends 0.62 to “cat”. Every row points one position to the left. That is a positional pattern with no interest in what the words are.
Now consider what a single head would have to do with both jobs. To satisfy head 0 it would need “sat” attending to “cat”, and to satisfy head 1 it would need “sat” attending to “cat” as well, which happens to agree here, but “The” would need to attend to “cat” for one job and to itself for the other, and those are incompatible. One distribution cannot hold both, so a single head must pick one relationship and abandon the other.
Note that both heads run the same function over the same input. Nothing in the code distinguishes them and only the weights differ, which is what makes adding a head a matter of allocating three more matrices rather than writing new logic.
21.5 Concatenate and Project
Two heads produce two narrow outputs per position, and the layer has to return one vector of the original width. Concatenation restores the width, since h outputs of d_head each lay end to end to give exactly d_model. A final matrix W_O then mixes across the join so the heads are not merely stacked but combined.
/* 112_Concat_Project.c */
#include <stdio.h>
#include <math.h>
#define SEQ_LEN 4
#define D_MODEL 6
#define D_HEAD 3
#define H 2
static void softmax(float *x, int n)
{
float mx = x[0], s = 0;
int i;
for (i = 1; i < n; i++) if (x[i] > mx) mx = x[i];
for (i = 0; i < n; i++) {
x[i] = expf(x[i] - mx);
s += x[i];
}
for (i = 0; i < n; i++) x[i] /= s;
}
/* One head, returning its output vectors */
static void head(const float X[SEQ_LEN][D_MODEL],
const float W_Q[D_HEAD][D_MODEL],
const float W_K[D_HEAD][D_MODEL],
const float W_V[D_HEAD][D_MODEL],
float out[SEQ_LEN][D_HEAD])
{
float Q[SEQ_LEN][D_HEAD], K[SEQ_LEN][D_HEAD];
float V[SEQ_LEN][D_HEAD], w[SEQ_LEN];
float scale = 1.0f / sqrtf((float)D_HEAD);
int i, j, k, m;
for (i = 0; i < SEQ_LEN; i++)
for (k = 0; k < D_HEAD; k++) {
Q[i][k] = 0;
K[i][k] = 0;
V[i][k] = 0;
for (m = 0; m < D_MODEL; m++) {
Q[i][k] += W_Q[k][m] * X[i][m];
K[i][k] += W_K[k][m] * X[i][m];
V[i][k] += W_V[k][m] * X[i][m];
}
}
for (i = 0; i < SEQ_LEN; i++) {
for (j = 0; j < SEQ_LEN; j++) {
w[j] = 0;
for (k = 0; k < D_HEAD; k++)
w[j] += Q[i][k] * K[j][k];
w[j] *= scale;
}
softmax(w, SEQ_LEN);
for (k = 0; k < D_HEAD; k++) {
out[i][k] = 0;
for (j = 0; j < SEQ_LEN; j++)
out[i][k] += w[j] * V[j][k];
}
}
}
int main(void)
{
float X[SEQ_LEN][D_MODEL] = {
{ 1.0f, 0.2f, -0.3f, 0.5f, -0.1f, 0.8f },
{ 0.3f, 0.9f, 0.1f, -0.2f, 0.6f, 0.4f },
{ -0.2f, 0.4f, 0.8f, 0.3f, -0.5f, 0.1f },
{ 0.6f, -0.1f, 0.2f, 0.9f, 0.3f, -0.4f },
};
const char *words[SEQ_LEN] = {
"The", "cat", "sat", "down"
};
/* The same two heads as Step 2 */
float W_Q0[D_HEAD][D_MODEL] = {
{ -0.29f, +0.54f, +1.15f, +0.43f,
-0.91f, +0.28f },
{ +1.26f, +1.03f, +2.02f, +2.22f,
-3.58f, +2.71f },
{ +1.40f, +2.66f, +1.56f, +1.52f,
+4.53f, -2.21f },
};
float W_K0[D_HEAD][D_MODEL] = {
{ +0.54f, -0.09f, -0.11f, +0.47f,
-0.13f, +0.24f },
{ +0.06f, +0.63f, +0.14f, -0.19f,
+0.60f, +0.02f },
{ -0.01f, +0.28f, +0.76f, +0.55f,
-0.26f, -0.15f },
};
float W_Q1[D_HEAD][D_MODEL] = {
{ +2.02f, +2.52f, -0.38f, -0.18f,
+1.01f, +2.57f },
{ -0.73f, +1.36f, +2.86f, +1.06f,
-2.26f, +0.70f },
{ +1.14f, -0.16f, +0.91f, +2.38f,
+1.84f, -2.30f },
};
float W_K1[D_HEAD][D_MODEL] = {
{ +0.57f, -0.09f, -0.09f, +0.52f,
-0.09f, +0.19f },
{ +0.23f, +0.60f, +0.28f, +0.18f,
+0.88f, -0.34f },
{ -0.16f, +0.30f, +0.64f, +0.24f,
-0.50f, +0.15f },
};
/* Each head keeps a different slice of content */
float W_V0[D_HEAD][D_MODEL] = {
{ 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f },
{ 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f },
{ 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f },
};
float W_V1[D_HEAD][D_MODEL] = {
{ 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f },
{ 0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f },
{ 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 1.0f },
};
/* W_O mixes the two heads back to d_model */
float W_O[D_MODEL][H * D_HEAD] = {
{ 0.8f, 0.1f, 0.0f, 0.2f, 0.0f, 0.1f },
{ 0.1f, 0.7f, 0.1f, 0.0f, 0.3f, 0.0f },
{ 0.0f, 0.2f, 0.9f, 0.1f, 0.0f, 0.2f },
{ 0.3f, 0.0f, 0.1f, 0.8f, 0.1f, 0.0f },
{ 0.0f, 0.2f, 0.0f, 0.1f, 0.7f, 0.2f },
{ 0.1f, 0.0f, 0.2f, 0.0f, 0.1f, 0.9f },
};
float h0[SEQ_LEN][D_HEAD], h1[SEQ_LEN][D_HEAD];
float cat[SEQ_LEN][H * D_HEAD];
float out[SEQ_LEN][D_MODEL];
int i, j, k;
head(X, W_Q0, W_K0, W_V0, h0);
head(X, W_Q1, W_K1, W_V1, h1);
for (i = 0; i < SEQ_LEN; i++) {
for (k = 0; k < D_HEAD; k++) {
cat[i][k] = h0[i][k];
cat[i][D_HEAD + k] = h1[i][k];
}
for (j = 0; j < D_MODEL; j++) {
out[i][j] = 0;
for (k = 0; k < H * D_HEAD; k++)
out[i][j] += W_O[j][k] * cat[i][k];
}
}
printf("Per head outputs, "
"d_head=%d each\n\n", D_HEAD);
for (i = 0; i < SEQ_LEN; i++) {
printf(" %-5s h0=[%+.2f,%+.2f,%+.2f] ",
words[i], h0[i][0], h0[i][1], h0[i][2]);
printf("h1=[%+.2f,%+.2f,%+.2f]\n",
h1[i][0], h1[i][1], h1[i][2]);
}
printf("\nConcatenated, "
"d_model=%d\n\n", H * D_HEAD);
for (i = 0; i < SEQ_LEN; i++) {
printf(" %-5s [", words[i]);
for (k = 0; k < H * D_HEAD; k++)
printf("%+.2f%s", cat[i][k],
k < H*D_HEAD-1 ? "," : "]\n");
}
printf("\nAfter the W_O projection, d_model=%d\n\n",
D_MODEL);
for (i = 0; i < SEQ_LEN; i++) {
printf(" %-5s [", words[i]);
for (j = 0; j < D_MODEL; j++)
printf("%+.2f%s", out[i][j],
j < D_MODEL-1 ? "," : "]\n");
}
printf("\nInput and output are both %d wide, "
"which\n",
D_MODEL);
printf("is what lets these layers stack.\n");
return 0;
}
Figure 21-4 concatenates the head outputs and mixes them with W_O. Read the per head outputs first and one detail jumps out. Head 1 produces [+0.57, +0.03, +0.41] for “The” and the identical vector for “cat”, because head 1′s attention rows for those two positions are identical, both attending 0.62 to “The”. On head 1′s evidence alone the two words are indistinguishable.
Head 0 separates them cleanly. It gives [+0.33, +0.77, +0.12] for “The” and [+0.07, +0.30, +0.57] for “cat”, which are nothing alike. So after concatenation the two positions differ, and they differ only because of head 0. That is a direct demonstration of why more than one head helps. A head that collapses a distinction costs nothing as long as another head preserves it, whereas a single head making the same collapse would lose the information for good.
The concatenated vectors are six wide, matching d_model, and W_O then produces the final output of the same width. Look at “The” going from [+0.33, +0.77, +0.12, +0.57, +0.03, +0.41] to [+0.50, +0.59, +0.40, +0.57, +0.31, +0.43] and notice that the third component moved from +0.12 to +0.40. W_O has mixed material from head 1 into a position that head 0 owned, which is its job. Without it the two halves of the output would remain in separate lanes and never interact.
Input width and output width both being d_model is not a coincidence but a requirement. It is what allows one of these layers to feed the next, which is what Chapter 24 does when it stacks them.
21.6 Parameter Count
The claim that heads are free deserves arithmetic rather than assurance, so this program counts the parameters across a range of head counts at a fixed model dimension.
/* 113_Params.c */
#include <stdio.h>
int main(void)
{
int d_model = 512;
int heads[] = { 1, 2, 4, 8, 16 };
int n = 5;
int i;
printf("Multi-head attention parameters, "
"d_model=%d\n\n", d_model);
printf(" heads d_head W_Q+W_K+W_V W_O "
" total\n");
printf(" ----- ------ --------------- -------"
"---- --------\n");
for (i = 0; i < n; i++) {
int h = heads[i];
int d_head = d_model / h;
/* Per head W_Q is [d_head x d_model],
and the same for W_K and W_V */
/* Across h heads that comes to
h * 3 * d_head * d_model,
which is 3 * d_model * d_model */
int qkv_params = h * 3 * d_head * d_model;
/* W_O: [d_model x d_model] (always the same) */
int wo_params = d_model * d_model;
int total = qkv_params + wo_params;
printf(" %3d %4d %10d %8d "
"%8d\n",
h, d_head, qkv_params, wo_params, total);
}
printf("\nThe total is ALWAYS 4 * d_model^2 "
"= %d.\n", 4 * d_model * d_model);
printf("The head count does not change it.\n\n");
printf("Each head has smaller projections, "
"since d_head = d_model/h.\n");
printf("More heads with smaller projections "
"costs the same total\n");
printf("as fewer heads with "
"larger projections.\n\n");
printf("What changes is the diversity of "
"attention patterns.\n");
printf("8 heads can learn 8 relationship "
"types, 1 head learns 1.\n");
return 0;
}
Figure 21-5 counts parameters against head count at a fixed model dimension. Every row of the table shows the same total. At d_model of 512 the answer is 1,048,576 whether there is one head or sixteen, and the two component columns do not move either, with the three projections always costing 786,432 and W_O always costing 262,144.
The cancellation is exact rather than approximate. One head’s W_Q is d_head by d_model, so h of them come to h * d_head * d_model, and since d_head is d_model / h the h cancels and leaves d_model squared. Three such matrices give 3 * d_model^2, W_O adds another d_model^2 because it maps d_model to d_model, and the total is 4 * d_model^2 with no h anywhere in it.
So the head count is free in parameters and in multiply count for the same reason, and costs only a little bookkeeping. What it buys is the independence we saw earlier, where two heads held two incompatible patterns that one head could not have held together. Doubling the heads halves the width each one works in, which is the only real trade, and a head with too few dimensions cannot express much, which is what exercise 2 explores.
The original transformer used d_model of 512 with eight heads, giving d_head of 64, and that ratio has largely stuck. Models have grown their model dimension and their head count together, keeping d_head near 64 rather than letting it drift.
21.7 Two Heads, Two Specializations
We already saw two heads disagreeing on a four word sentence. This one makes the specialization unmistakable by using a six token sequence whose types and positions are deliberately in tension, and printing each attention matrix as a map rather than as numbers.
The tokens are A0, B1, A2, C3, B4 and D5, where the letter is a type and the digit is a position. Two A tokens sit at positions 0 and 2, so a head that groups by type and a head that groups by proximity must disagree about them.
/* 114_Head_Patterns.c */
#include <stdio.h>
#include <math.h>
#include <float.h>
#include <stdlib.h>
static float randf(void)
{
return (float)rand() / RAND_MAX;
}
static float dot(const float *a, const float *b,
int n)
{
float s = 0;
int i;
for (i = 0; i < n; i++) s += a[i] * b[i];
return s;
}
static void softmax(float *x, int n)
{
float mx = -FLT_MAX, s = 0;
int i;
for (i = 0; i < n; i++) if (x[i] > mx) mx = x[i];
for (i = 0; i < n; i++) {
x[i] = expf(x[i] - mx);
s += x[i];
}
for (i = 0; i < n; i++) x[i] /= s;
}
#define SL 6
#define DM 4
#define NH 2
#define DH (DM / NH)
/* Get attention weights for one head */
static void head_weights(const float X[SL][DM],
const float WQ[DH][DM],
const float WK[DH][DM],
float w[SL][SL])
{
float scale = 1.0f / sqrtf((float)DH);
int i, j, k, m;
for (i = 0; i < SL; i++) {
float q[DH];
for (k = 0; k < DH; k++) {
q[k] = 0;
for (m = 0; m < DM; m++)
q[k] += WQ[k][m] * X[i][m];
}
float scores[SL];
for(j = 0;j<SL;j++) {
float key[DH];
for (k = 0; k < DH; k++) {
key[k] = 0;
for (m = 0; m < DM; m++)
key[k] += WK[k][m] * X[j][m];
}
scores[j] = dot(q, key, DH)*scale;
}
softmax(scores, SL);
for(j = 0;j<SL;j++) w[i][j] = scores[j];
}
}
int main(void)
{
/* A 6-position sequence */
float X[SL][DM] = {
/* pos 0: type A */
{ 1.0f, 0.0f, 0.0f, 0.0f },
/* pos 1: type B */
{ 0.0f, 1.0f, 0.0f, 0.0f },
/* pos 2: type A */
{ 1.0f, 0.0f, 0.0f, 0.0f },
/* pos 3: type C */
{ 0.0f, 0.0f, 1.0f, 0.0f },
/* pos 4: type B */
{ 0.0f, 1.0f, 0.0f, 0.0f },
/* pos 5: type D */
{ 0.0f, 0.0f, 0.0f, 1.0f },
};
const char *labels[] = { "A0", "B1", "A2",
"C3", "B4", "D5" };
/* Hand designed so the two heads differ */
/* Head 0 attends to the same token type */
float WQ0[DH][DM] = {{ 1, 0, 0, 0 },
{ 0, 1, 0, 0 }};
float WK0[DH][DM] = {{ 1, 0, 0, 0 },
{ 0, 1, 0, 0 }};
/* Head 1: attends to adjacent positions
(positional) */
float WQ1[DH][DM] = {{ 0, 0, 1, 0 },
{ 0, 0, 0, 1 }};
float WK1[DH][DM] = {{ 0, 0, 1, 0 },
{ 0, 0, 0, 1 }};
float w[SL][SL];
int h, i, j;
/* Modify X to include positional info in dims
2-3 */
for (i = 0; i < SL; i++) {
X[i][2] = sinf(i * 0.5f);
X[i][3] = cosf(i * 0.5f);
}
printf("Multi-head attention, 2 specialized "
"heads\n\n");
/* Head 0 */
head_weights(X, WQ0, WK0, w);
printf("Head 0, type matching, A attends to A\n");
printf(" ");
for (j = 0; j < SL; j++) printf("%-5s", labels[j]);
printf("\n");
for(i = 0;i<SL;i++) {
printf(" %-4s ",labels[i]);
for(j = 0;j<SL;j++){
int bar = (int)(w[i][j] * 20);
if (bar > 4) printf("#### ");
else if (bar > 2) printf("## ");
else printf(". ");
}
printf("\n");
}
/* Head 1 */
head_weights(X, WQ1, WK1, w);
printf("\nHead 1, positional, nearby wins\n");
printf(" ");
for (j = 0; j < SL; j++) printf("%-5s", labels[j]);
printf("\n");
for(i = 0;i<SL;i++) {
printf(" %-4s ",labels[i]);
for(j = 0;j<SL;j++){
int bar = (int)(w[i][j] * 20);
if (bar > 4) printf("#### ");
else if (bar > 2) printf("## ");
else printf(". ");
}
printf("\n");
}
printf("\nHead 0 groups same-type tokens, "
"A with A.\n");
printf("Head 1 focuses on positional neighbors.\n");
printf("Together they capture type and "
"position.\n");
printf("One head could not do both "
"simultaneously.\n");
return 0;
}
Figure 21-6 draws both attention maps, one head keyed on type and the other on position. Head 0 produces a map with a clear off diagonal structure. A0 and A2 attend strongly to each other and weakly to everything else, B1 and B4 do the same across four positions of separation, and C3 and D5, having no partner of their own type, spread their weight evenly across all six. Type is the only thing this head responds to, and distance is irrelevant to it, which is why B1 reaching B4 looks the same as A0 reaching A2.
Head 1 produces a band along the diagonal. Every token attends most strongly to itself and its immediate neighbors, and the weight falls away with distance regardless of what the tokens are. A0 and A2 are two apart and get a middling weight from each other, exactly the same as A0 and B1 which are one apart get, because this head cannot tell an A from a B.
Put the two maps side by side and the incompatibility is visible. On head 0, A0 and A2 are a strong pair while A0 and B1 are nothing. On head 1, A0 and B1 are closer than A0 and A2. A single head has to produce one matrix, so it would have to choose which of those two statements to make, and the other relationship would be unavailable to every layer downstream.
Real transformers develop this diversity without being asked. Heads have been observed attending to the previous token, and others to the first token of the sequence, the matching bracket or quotation mark, or the syntactic head of the current word. Nothing in the architecture specifies any of it, and it emerges because h independent heads with no term tying them together will separate if separating lowers the loss.
21.8 Key Takeaways
Multi-head attention runs h complete attention mechanisms in parallel, each with its own W_Q, W_K and W_V of size d_head by d_model, where d_head is d_model divided by h.
Heads do not slice the input. We saw a slice giving [+0.10, +0.50, −0.30, +0.80] where the real projection gave [+0.10, +0.38, +0.10, +0.71], because every head reads the whole input vector and only the output width is divided.
Two heads can hold incompatible patterns. We built one head sending “The” to “cat” at 0.82 and another sending “The” to itself at 0.62, which no single distribution could express at once.
Head outputs are concatenated back to d_model and mixed by W_O. Without that final matrix the heads would occupy separate lanes of the output and never interact.
A head that collapses a distinction is harmless if another head keeps it. We watched head 1 produce identical vectors for two positions while head 0 kept them clearly apart.
Parameters come to 4 * d_model^2 regardless of the head count, because h heads of width d_model/h cancel exactly. At d_model of 512 that is 1,048,576 whether h is 1 or 16.
The trade is width per head rather than total cost. More heads means each one works in fewer dimensions, and a head with too few dimensions cannot express much.
Input and output are both d_model wide, which is what allows these layers to stack. Chapter 24 relies on that.
The standard configuration from the original transformer is d_model of 512, h of 8 and d_head of 64, and models have since grown d_model and h together rather than letting d_head drift far from 64.
Specialization is learned, not designed. Heads in trained models have been found attending to the previous token, the first token, matching punctuation and syntactic dependents, none of which is built into the architecture.
21.9 Exercises
Run 112_Concat_Project.c with 4 heads instead of 2. The d_head becomes d_model/4. Does the output change significantly?
What happens if you use only 1 head with d_head = d_model? This is equivalent to single-head Q/K/V attention from Chapter 20. Verify the parameter count matches.
Implement multi-head attention where the heads run sequentially (loop) instead of in parallel. The output should be identical. On a GPU, parallel execution is faster; on a microcontroller, sequential is the only option.
For d_model = 64 and h = 8, compute the total memory needed for all weight matrices (W_Q, W_K, W_V per head, plus W_O) in bytes (float32).
Add biases to the projection matrices. How many additional parameters does this add per head? Is it significant compared to the weight matrices?
In 114_Head_Patterns.c, design a third head that attends to tokens two positions ahead. What W_Q and W_K matrices would achieve this?