Lesson 9 of 73
Self-Attention Step by Step
Self-attention is the operation that lets each token use information from the tokens before it. Every other layer of a transformer works on one position at a time. Attention is the only place where positions exchange information. Without it, the model would read each word in isolation, and "bank" in "river bank" would mean the same as "bank" in "money bank".
This lesson builds attention in small steps. It starts from the simplest possible mixing, an average, and ends with the full mechanism used in Llama and Qwen models. The matrices are small, so every number can be read. The lesson uses the dot product, softmax, and matrix product from lesson 8. If any of those feel unfamiliar, read that lesson first.
After the lesson you can write attention in Java, explain every term in an attention configuration (n_head, n_head_kv), and compute how attention cost and cache size grow as the context gets longer.
| Step | Example |
|---|---|
| Why context is needed, and the simplest mixing | 1 and 2 |
| Scores from dot products, then query, key, value | 3 and 4 |
| Scaling and the causal mask | 5 and 6 |
| The full attention matrix, in loops and in matrix form | 7 and 8 |
| Several heads at once | 9 |
| Grouped-query and multi-query attention | 10 and 11 |
| Cost as the text gets longer | 12 |
| A hand-built head | 13 |
The programs use float[][] for readability. Engines use flat arrays (lesson 8, Example 6), and the arithmetic is the same. Random matrices come from a seeded java.util.Random, so every run prints the same numbers. These numbers have no meaning, because a trained model learns its matrices from data. Example 13 builds a meaningful one by hand.
Example 1: a token's vector must depend on its context
A model begins with an embedding table (lesson 8, Example 12). The table maps a token id to one fixed vector. The word bank gets the same vector in river bank and in money bank, although the two phrases use it with different meanings. Something must change the vector of bank using the words around it.
This program shows the problem and a first, crude fix. It prints the fixed lookup for bank, then mixes the bank vector half-and-half with the vector of its neighbor. The three numbers of each embedding loosely mean "nature", "finance", and "place" (invented for the lesson). Look at how the two mixed vectors differ, and at the cosine similarities (lesson 8, Example 3) at the end.
import java.util.Arrays;
import java.util.Locale;
import java.util.Map;
public class WhyContext {
static float cosine(float[] a, float[] b) {
float d = 0, na = 0, nb = 0;
for (int i = 0; i < a.length; i++) { d += a[i] * b[i]; na += a[i] * a[i]; nb += b[i] * b[i]; }
return d / (float) Math.sqrt(na * nb);
}
static float[] mix(float[] a, float[] b) {
float[] out = new float[a.length];
for (int i = 0; i < a.length; i++) out[i] = 0.5f * a[i] + 0.5f * b[i];
return out;
}
public static void main(String[] args) {
Map<String, float[]> emb = Map.of(
"river", new float[] {0.9f, 0.0f, 0.3f},
"money", new float[] {0.0f, 0.9f, 0.1f},
"bank", new float[] {0.3f, 0.3f, 0.9f});
float[] bank = emb.get("bank");
System.out.println("bank in 'river bank' : " + Arrays.toString(bank));
System.out.println("bank in 'money bank' : " + Arrays.toString(bank));
float[] withRiver = mix(bank, emb.get("river"));
float[] withMoney = mix(bank, emb.get("money"));
System.out.println("after mixing with river: " + Arrays.toString(withRiver));
System.out.println("after mixing with money: " + Arrays.toString(withMoney));
System.out.println(String.format(Locale.ROOT, "cosine(river-bank, river) = %.2f, cosine(river-bank, money) = %.2f",
cosine(withRiver, emb.get("river")), cosine(withRiver, emb.get("money"))));
System.out.println(String.format(Locale.ROOT, "cosine(money-bank, river) = %.2f, cosine(money-bank, money) = %.2f",
cosine(withMoney, emb.get("river")), cosine(withMoney, emb.get("money"))));
}
}bank in 'river bank' : [0.3, 0.3, 0.9]
bank in 'money bank' : [0.3, 0.3, 0.9]
after mixing with river: [0.6, 0.15, 0.6]
after mixing with money: [0.15, 0.6, 0.5]
cosine(river-bank, river) = 0.88, cosine(river-bank, money) = 0.25
cosine(money-bank, river) = 0.38, cosine(money-bank, money) = 0.82The lookup gives bank the vector (0.3, 0.3, 0.9) in both phrases, so the table alone cannot tell the two meanings apart.
After mixing, the river bank vector is (0.6, 0.15, 0.6), which is the average of (0.3, 0.3, 0.9) and (0.9, 0.0, 0.3). The money bank vector is (0.15, 0.6, 0.5). The two vectors now differ, and each is closer to its own neighbor. The river version has cosine 0.88 with river and 0.25 with money. The money version has 0.38 with river and 0.82 with money.
A fixed 50-50 average is too crude, because it gives every neighbor the same weight whether or not that neighbor helps. The next examples let the model choose the weights.
Example 2: a uniform average of the previous vectors
A language model reads text from left to right. A token may use only the tokens before it and itself. The reason comes from generation: while the model generates text, the later tokens do not exist yet, so nothing can depend on them.
The simplest way to give a position information from earlier positions is the average of all vectors up to and including that position. This program computes that running average for four vectors with 2 numbers each. Read the output one line at a time and notice how many vectors each position uses.
import java.util.Arrays;
public class UniformAverage {
public static void main(String[] args) {
float[][] x = {{1, 0}, {0, 1}, {1, 1}, {2, 0}};
for (int i = 0; i < x.length; i++) {
float[] avg = new float[2];
for (int j = 0; j <= i; j++) {
for (int d = 0; d < 2; d++) avg[d] += x[j][d] / (i + 1);
}
System.out.println("position " + i + ": average of " + (i + 1) + " vector(s) = " + Arrays.toString(avg));
}
}
}position 0: average of 1 vector(s) = [1.0, 0.0]
position 1: average of 2 vector(s) = [0.5, 0.5]
position 2: average of 3 vector(s) = [0.6666667, 0.6666667]
position 3: average of 4 vector(s) = [1.0, 0.5]Position 0 sees only itself, so its average is its own vector, (1, 0). Position 1 averages (1, 0) and (0, 1) to get (0.5, 0.5). Position 3 averages all four vectors to get (1.0, 0.5).
Each output uses only vectors at or before its position. This means appending a new token never changes the earlier outputs. That property is what makes the KV cache of lesson 11 possible: the engine can store earlier work and reuse it.
The weakness is that the weights are fixed at 1 divided by the count. A distant, useless vector counts as much as a close, useful one.
Example 3: a weighted average from dot-product scores
A uniform average treats every earlier token as equally useful. A better mix gives larger weights to the tokens that matter for the current one. A cheap measure of "matters" is the dot product (lesson 8, Example 2): vectors that point the same way score high.
The recipe has three steps:
- Score: take the dot product of the current vector with each vector up to the current position.
- Weights: apply softmax to the scores, so the weights are positive and add up to 1.
- Mix: add up the vectors, each multiplied by its weight.
The program applies the recipe to the last of the four vectors from Example 2. Compare the printed mixed vector with the uniform average of [1.0, 0.5].
import java.util.Arrays;
import java.util.Locale;
public class WeightedAverage {
public static void main(String[] args) {
float[][] x = {{1, 0}, {0, 1}, {1, 1}, {2, 0}};
int i = 3;
double[] scores = new double[i + 1];
double sum = 0;
for (int j = 0; j <= i; j++) {
scores[j] = x[i][0] * x[j][0] + x[i][1] * x[j][1];
sum += Math.exp(scores[j]);
}
double[] w = new double[i + 1];
float[] mixed = new float[2];
for (int j = 0; j <= i; j++) {
w[j] = Math.exp(scores[j]) / sum;
for (int d = 0; d < 2; d++) mixed[d] += (float) (w[j] * x[j][d]);
}
System.out.println("scores = " + Arrays.toString(scores));
StringBuilder sb = new StringBuilder();
for (double v : w) sb.append(String.format(Locale.ROOT, "%.3f ", v));
System.out.println("weights = " + sb.toString().trim());
System.out.println("mixed = " + Arrays.toString(mixed) + " (uniform average was [1.0, 0.5])");
}
}scores = [2.0, 0.0, 2.0, 4.0]
weights = 0.105 0.014 0.105 0.776
mixed = [1.7615942, 0.11920292] (uniform average was [1.0, 0.5])The last vector is (2, 0). Its scores against the four vectors are 2, 0, 2, and 4. For example, (2, 0) with (0, 1) gives 2*0 + 0*1 = 0, because the two are at a right angle. The vector matches itself best: 2*2 + 0*0 = 4.
Softmax turns the scores into the weights 0.105, 0.014, 0.105, and 0.776. The mixed vector (1.76, 0.12) is dominated by the last vector itself, with equal smaller shares for (1, 0) and (1, 1). The uniform average was (1.0, 0.5).
The important change is that the weights now depend on the content. A different set of vectors gives different weights.
Example 4: queries, keys, and values
Example 3 uses the same vector for three different jobs: asking the question, being compared, and being mixed into the result. Attention separates these jobs. Each token's input vector x is multiplied by three learned matrices:
Wqgives a queryq, which describes what this position looks for.Wkgives a keyk, which describes what this position offers for matching.Wvgives a valuev, which is the content that gets mixed into other positions.
A library analogy helps. The query is the question you bring, the key is the label on a book's spine, and the value is the content inside the book. You compare your question with the labels, then read the books whose labels match. The analogy stops at one point: attention reads a weighted blend of all books, not one book.
The score between position i (the asker) and position j is q_i . k_j, and the mix uses the values v_j. Separate matrices let a model score by one property and pass along another. The size of q and k is the head dimension, written d_head.
This program projects four inputs of 4 numbers each into queries, keys, and values of 2 numbers each, using seeded random matrices. It also shows that scores are not symmetric.
import java.util.Locale;
import java.util.Random;
public class QueryKeyValue {
static float[][] random(Random r, int rows, int cols) {
float[][] m = new float[rows][cols];
for (float[] row : m) for (int c = 0; c < cols; c++) row[c] = (float) (r.nextGaussian() * 0.5);
return m;
}
static float[][] project(float[][] x, float[][] w) { // x: n x d_model, w: d_model x d_head
float[][] out = new float[x.length][w[0].length];
for (int i = 0; i < x.length; i++)
for (int c = 0; c < w[0].length; c++)
for (int k = 0; k < w.length; k++) out[i][c] += x[i][k] * w[k][c];
return out;
}
static void print(String name, float[][] m) {
for (int i = 0; i < m.length; i++) {
StringBuilder sb = new StringBuilder(name + i + " =");
for (float v : m[i]) sb.append(String.format(Locale.ROOT, " %6.2f", v));
System.out.println(sb);
}
}
public static void main(String[] args) {
Random r = new Random(1);
float[][] x = random(r, 4, 4);
float[][] q = project(x, random(r, 4, 2)), k = project(x, random(r, 4, 2)), v = project(x, random(r, 4, 2));
print("q", q);
print("k", k);
print("v", v);
float s21 = q[2][0] * k[1][0] + q[2][1] * k[1][1];
float s12 = q[1][0] * k[2][0] + q[1][1] * k[2][1];
System.out.println(String.format(Locale.ROOT, "score(2 asks 1) = q2.k1 = %.3f, score(1 asks 2) = q1.k2 = %.3f", s21, s12));
}
}q0 = 0.12 -0.46
q1 = 0.43 -1.10
q2 = -0.16 0.59
q3 = -0.18 0.26
k0 = 0.57 0.12
k1 = 0.64 0.86
k2 = -0.10 0.34
k3 = 0.12 -0.51
v0 = 1.44 0.91
v1 = 1.07 -0.33
v2 = 0.05 -0.07
v3 = -0.68 -0.49
score(2 asks 1) = q2.k1 = 0.407, score(1 asks 2) = q1.k2 = -0.417Each input vector produces one query, one key, and one value, all of size 2, computed by three different matrices.
The two score lines at the end show that scores are not symmetric. Position 2 asking about position 1 scores 0.407, while position 1 asking about position 2 scores -0.417. Queries and keys play different roles, so "how much should I look at you" and "how much should you look at me" are different questions.
Real models learn the three matrices during training. Wq and Wk have shape d_model x d_head, where d_model is the width of the input vector. The attention layer has these three matrices for each head, plus the output matrix of Example 9.
Example 5: why scores are divided by the square root of the head dimension
A dot product adds up d_head products. Suppose q and k hold random numbers with variance 1 (variance measures how spread out numbers are). Then the sum has variance d_head, so its standard deviation is sqrt(d_head). The scores spread out more as the head dimension grows.
Large scores are a problem for softmax. They make it almost one-hot: one weight near 1 and all the others near 0. During training, the small weights then receive almost no gradient (the signal that tells the model how to improve), so learning stalls. Dividing every score by sqrt(d_head) brings the spread back to about 1.
The program measures this. For three head sizes, it prints the standard deviation of scores, then the average largest softmax weight over 8 keys, with and without the division. Compare the last two columns.
import java.util.Locale;
import java.util.Random;
public class Scaling {
public static void main(String[] args) {
Random r = new Random(5);
System.out.println("d_head std of q.k sqrt(d_head) largest weight, unscaled largest weight, scaled");
for (int d : new int[] {4, 64, 1024}) {
double sumSq = 0;
int pairs = 2000;
for (int p = 0; p < pairs; p++) {
double dot = 0;
for (int i = 0; i < d; i++) dot += r.nextGaussian() * r.nextGaussian();
sumSq += dot * dot;
}
double std = Math.sqrt(sumSq / pairs);
double maxRaw = 0, maxScaled = 0;
int trials = 500;
for (int t = 0; t < trials; t++) {
double[] scores = new double[8];
for (int j = 0; j < 8; j++) {
double dot = 0;
for (int i = 0; i < d; i++) dot += r.nextGaussian() * r.nextGaussian();
scores[j] = dot;
}
maxRaw += largestWeight(scores, 1.0);
maxScaled += largestWeight(scores, Math.sqrt(d));
}
System.out.println(String.format(Locale.ROOT, "%6d %10.2f %12.2f %24.3f %22.3f", d, std, Math.sqrt(d), maxRaw / trials, maxScaled / trials));
}
}
static double largestWeight(double[] scores, double divisor) {
double max = Double.NEGATIVE_INFINITY;
for (double s : scores) max = Math.max(max, s / divisor);
double sum = 0;
for (double s : scores) sum += Math.exp(s / divisor - max);
return 1.0 / sum;
}
}d_head std of q.k sqrt(d_head) largest weight, unscaled largest weight, scaled
4 2.00 2.00 0.576 0.369
64 8.23 8.00 0.875 0.359
1024 31.67 32.00 0.965 0.358The measured standard deviation of the scores follows sqrt(d_head): 2.00 for 4, 8.23 for 64, and 31.67 for 1,024.
Without scaling, the largest softmax weight among 8 keys averages 0.576 at d_head = 4, 0.875 at 64, and 0.965 at 1,024. Softmax becomes nearly one-hot, and one key takes almost everything. With the division, the largest weight stays near 0.36 at every size, which leaves room for several keys to share the weight.
Many real models use heads of 64 or 128 dimensions. Qwen2-0.5B uses 64, Llama 3 8B uses 128, and the tiny test model of lesson 4 uses 32. The scaling factor is therefore 1/8 for 64 dimensions and about 1/11 for 128. The original transformer paper defines attention with this division, softmax(Q Kᵀ / sqrt(d_head)) V. llama.cpp uses 1/sqrt(d_head) as the scale unless the model file sets a different one.
Example 6: the causal mask
Position i may attend only to positions 0 through i. Section 2 gave the reason for generation. Training has the same need: the model must not be able to read the answer it is trying to predict.
The program computes the full table of scores for all pairs of positions. It then sets every score above the diagonal (the pairs where j > i, a later position) to minus infinity. After softmax, exp(-infinity) is 0, so those weights are exactly 0. This is the causal mask. It makes the output at position i independent of later tokens.
The program prints the raw scores, the masked scores, and the weights after softmax. Look at how the zeros form a triangle.
import java.util.Locale;
public class CausalMask {
public static void main(String[] args) {
int n = 4;
double[][] scores = new double[n][n];
for (int i = 0; i < n; i++) for (int j = 0; j < n; j++) scores[i][j] = 1.0 + 0.5 * ((i * 3 + j * 5) % 4);
System.out.println("raw scores:");
print(scores);
for (int i = 0; i < n; i++) for (int j = i + 1; j < n; j++) scores[i][j] = Double.NEGATIVE_INFINITY;
System.out.println("after the mask:");
print(scores);
System.out.println("weights (softmax per row):");
double[][] w = new double[n][n];
for (int i = 0; i < n; i++) {
double sum = 0;
for (int j = 0; j < n; j++) sum += Math.exp(scores[i][j]);
for (int j = 0; j < n; j++) w[i][j] = Math.exp(scores[i][j]) / sum;
}
print(w);
}
static void print(double[][] m) {
for (double[] row : m) {
StringBuilder sb = new StringBuilder();
for (double v : row) sb.append(String.format(Locale.ROOT, "%9.3f", v));
System.out.println(sb);
}
}
}raw scores:
1.000 1.500 2.000 2.500
2.500 1.000 1.500 2.000
2.000 2.500 1.000 1.500
1.500 2.000 2.500 1.000
after the mask:
1.000-Infinity-Infinity-Infinity
2.500 1.000-Infinity-Infinity
2.000 2.500 1.000-Infinity
1.500 2.000 2.500 1.000
weights (softmax per row):
1.000 0.000 0.000 0.000
0.818 0.182 0.000 0.000
0.331 0.547 0.122 0.000
0.167 0.276 0.455 0.102The raw score table has a value for every pair. After the mask, the upper triangle holds minus infinity, so row 0 can see only column 0, row 1 only columns 0 and 1, and so on.
Softmax on each row gives weights that add up to 1 over the visible positions, and exactly 0 for the hidden ones. Row 0 is (1, 0, 0, 0), and row 1 is (0.818, 0.182, 0, 0).
Without the mask, the model could read the token it is supposed to predict. With the mask, one forward pass over a text of n tokens trains n predictions at once, and no prediction cheats. During generation, the mask allows the engine to reuse stored work (lesson 11).
Example 7: full attention for a 4-token sequence
This program combines every step so far: queries, keys, and values; scores scaled by sqrt(d_head); the causal mask; softmax; and the weighted sum of values. For each position i, the computation is:
score(i, j) = (q_i . k_j) / sqrt(d_head) for j <= i
weights(i) = softmax over j <= i of score(i, j)
output(i) = sum over j <= i of weights(i, j) * v_jThe first line scores each earlier position. The second line turns the scores into weights. The third line mixes the value vectors with those weights. The program uses the same seeded matrices as Example 4 and prints the attention weight matrix and the output vectors.
import java.util.Locale;
import java.util.Random;
public class FullAttention {
static float[][] random(Random r, int rows, int cols) {
float[][] m = new float[rows][cols];
for (float[] row : m) for (int c = 0; c < cols; c++) row[c] = (float) (r.nextGaussian() * 0.5);
return m;
}
static float[][] project(float[][] x, float[][] w) {
float[][] out = new float[x.length][w[0].length];
for (int i = 0; i < x.length; i++)
for (int c = 0; c < w[0].length; c++)
for (int k = 0; k < w.length; k++) out[i][c] += x[i][k] * w[k][c];
return out;
}
public static void main(String[] args) {
Random r = new Random(1);
float[][] x = random(r, 4, 4);
float[][] q = project(x, random(r, 4, 2)), k = project(x, random(r, 4, 2)), v = project(x, random(r, 4, 2));
int n = 4, dHead = 2;
float[][] weights = new float[n][n];
float[][] out = new float[n][dHead];
for (int i = 0; i < n; i++) {
double[] e = new double[i + 1];
double sum = 0;
for (int j = 0; j <= i; j++) {
double score = (q[i][0] * k[j][0] + q[i][1] * k[j][1]) / Math.sqrt(dHead);
e[j] = Math.exp(score);
sum += e[j];
}
for (int j = 0; j <= i; j++) {
weights[i][j] = (float) (e[j] / sum);
for (int d = 0; d < dHead; d++) out[i][d] += weights[i][j] * v[j][d];
}
}
System.out.println("attention weights (row = asking position, column = position looked at):");
for (float[] row : weights) {
StringBuilder sb = new StringBuilder();
for (float w : row) sb.append(String.format(Locale.ROOT, "%8.3f", w));
System.out.println(sb);
}
System.out.println("outputs:");
for (int i = 0; i < n; i++) System.out.println(String.format(Locale.ROOT, "position %d: %7.3f %7.3f", i, out[i][0], out[i][1]));
}
}attention weights (row = asking position, column = position looked at):
1.000 0.000 0.000 0.000
0.635 0.365 0.000 0.000
0.282 0.383 0.335 0.000
0.237 0.270 0.269 0.224
outputs:
position 0: 1.438 0.907
position 1: 1.305 0.457
position 2: 0.833 0.108
position 3: 0.490 0.000The weight matrix is lower triangular, because of the mask. Position 0 has weight 1.000 on itself. Position 1 gives 0.635 to position 0 and 0.365 to itself. Position 3 spreads its weight almost evenly (0.237, 0.270, 0.269, 0.224), because the random matrices carry no meaning. Every row adds up to 1.
The output at each position is the weighted sum of the value vectors. Position 0 outputs its own value, (1.438, 0.907). Later positions blend in the values of the earlier positions.
This is the complete single-head computation. The remaining examples change how it is organized (matrices, heads, shared heads), and not what it does.
Example 8: the same computation in matrix form
Engines do not loop over positions one at a time. They stack the queries, keys, and values of all positions into matrices Q, K, and V, with one row per position. Attention then becomes four steps on whole matrices:
S = Q Kᵀ / sqrt(d_head) all scores at once, an n x n matrix
S = mask(S) minus infinity above the diagonal
A = softmax(S) per row the attention matrix
Output = A VHere Kᵀ ("K transposed") is K with its rows and columns swapped. The product Q Kᵀ places the dot product of query i and key j at row i, column j, so one matrix product computes every score.
The program runs this form with flat row-major arrays and compares the result with the loop version from Example 7.
import java.util.Locale;
import java.util.Random;
public class MatrixForm {
static float[] random(Random r, int n) {
float[] m = new float[n];
for (int i = 0; i < n; i++) m[i] = (float) (r.nextGaussian() * 0.5);
return m;
}
static float[] matMul(float[] a, float[] b, int m, int k, int n) {
float[] c = new float[m * n];
for (int i = 0; i < m; i++)
for (int j = 0; j < n; j++) {
float sum = 0;
for (int p = 0; p < k; p++) sum += a[i * k + p] * b[p * n + j];
c[i * n + j] = sum;
}
return c;
}
static float[] transpose(float[] a, int rows, int cols) {
float[] t = new float[a.length];
for (int i = 0; i < rows; i++) for (int j = 0; j < cols; j++) t[j * rows + i] = a[i * cols + j];
return t;
}
public static void main(String[] args) {
Random r = new Random(1);
int n = 4, dModel = 4, dHead = 2;
float[] x = random(r, n * dModel);
float[] wq = random(r, dModel * dHead), wk = random(r, dModel * dHead), wv = random(r, dModel * dHead);
float[] q = matMul(x, wq, n, dModel, dHead), k = matMul(x, wk, n, dModel, dHead), v = matMul(x, wv, n, dModel, dHead);
float[] s = matMul(q, transpose(k, n, dHead), n, dHead, n);
for (int i = 0; i < n; i++) {
float max = Float.NEGATIVE_INFINITY;
for (int j = 0; j < n; j++) {
s[i * n + j] = j > i ? Float.NEGATIVE_INFINITY : s[i * n + j] / (float) Math.sqrt(dHead);
max = Math.max(max, s[i * n + j]);
}
float sum = 0;
for (int j = 0; j < n; j++) { s[i * n + j] = (float) Math.exp(s[i * n + j] - max); sum += s[i * n + j]; }
for (int j = 0; j < n; j++) s[i * n + j] /= sum;
}
float[] out = matMul(s, v, n, n, dHead);
// The loop version of Example 7 on the same numbers.
double maxDiff = 0;
for (int i = 0; i < n; i++) {
double[] e = new double[i + 1];
double sum = 0;
for (int j = 0; j <= i; j++) {
double dot = 0;
for (int d = 0; d < dHead; d++) dot += q[i * dHead + d] * k[j * dHead + d];
e[j] = Math.exp(dot / Math.sqrt(dHead));
sum += e[j];
}
for (int d = 0; d < dHead; d++) {
double o = 0;
for (int j = 0; j <= i; j++) o += e[j] / sum * v[j * dHead + d];
maxDiff = Math.max(maxDiff, Math.abs(o - out[i * dHead + d]));
}
}
for (int i = 0; i < n; i++) System.out.println(String.format(Locale.ROOT, "position %d: %7.3f %7.3f", i, out[i * dHead], out[i * dHead + 1]));
System.out.println("largest difference to the loop version below 1e-6: " + (maxDiff < 1e-6));
}
}position 0: 1.438 0.907
position 1: 1.305 0.457
position 2: 0.833 0.108
position 3: 0.490 0.000
largest difference to the loop version below 1e-6: trueThe output equals Example 7 to three decimals, and the check confirms that the difference to the loop version is below 1e-6.
The matrix form does one large product Q Kᵀ and one large product A V, in place of many small dot products. Large matrix products are what GPUs and CPU vector units do efficiently.
For a prompt of n tokens, S has n * n entries. A long prompt therefore needs an n x n table for each head. Fast kernels such as FlashAttention (covered in the attention-optimizations lesson later in this course) avoid storing the whole table by computing it in blocks.
Example 9: multi-head attention
One attention head scores positions by one learned notion of similarity. A model runs several heads in parallel so that different heads can look for different things. For example, one head may look at the previous token, another at the matching bracket, and another at the subject of a verb.
With n_head heads and model width d_model, each head has its own Wq, Wk, and Wv, and works on vectors of size d_head = d_model / n_head. The steps are:
- Project
xwith matrices of shaped_model x d_model, then split the result inton_headslices ofd_headcolumns each. - Run attention in each head independently.
- Concatenate the head outputs (place them side by side) so each position again has
d_modelnumbers. - Multiply by an output projection
Wo(d_model x d_model) that mixes the heads together.
The program uses d_model = 8 and 2 heads on four positions. It prints the weights each head gives to the last position.
import java.util.Locale;
import java.util.Random;
public class MultiHead {
static float[][] random(Random r, int rows, int cols) {
float[][] m = new float[rows][cols];
for (float[] row : m) for (int c = 0; c < cols; c++) row[c] = (float) (r.nextGaussian() * 0.5);
return m;
}
static float[][] matMul(float[][] a, float[][] b) {
float[][] c = new float[a.length][b[0].length];
for (int i = 0; i < a.length; i++)
for (int j = 0; j < b[0].length; j++)
for (int k = 0; k < b.length; k++) c[i][j] += a[i][k] * b[k][j];
return c;
}
public static void main(String[] args) {
Random r = new Random(3);
int n = 4, dModel = 8, heads = 2, dHead = dModel / heads;
float[][] x = random(r, n, dModel);
float[][] q = matMul(x, random(r, dModel, dModel)), k = matMul(x, random(r, dModel, dModel)), v = matMul(x, random(r, dModel, dModel));
float[][] wo = random(r, dModel, dModel);
float[][] concat = new float[n][dModel];
for (int h = 0; h < heads; h++) {
int off = h * dHead;
StringBuilder lastRow = new StringBuilder();
for (int i = 0; i < n; i++) {
double[] e = new double[i + 1];
double sum = 0;
for (int j = 0; j <= i; j++) {
double dot = 0;
for (int d = 0; d < dHead; d++) dot += q[i][off + d] * k[j][off + d];
e[j] = Math.exp(dot / Math.sqrt(dHead));
sum += e[j];
}
for (int j = 0; j <= i; j++) {
for (int d = 0; d < dHead; d++) concat[i][off + d] += (float) (e[j] / sum * v[j][off + d]);
if (i == n - 1) lastRow.append(String.format(Locale.ROOT, " %.3f", e[j] / sum));
}
}
System.out.println("head " + h + " weights of the last position:" + lastRow);
}
float[][] out = matMul(concat, wo);
System.out.println("concatenated: " + concat.length + " x " + concat[0].length + ", after Wo: " + out.length + " x " + out[0].length);
System.out.println("parameters in Wq, Wk, Wv, Wo: " + 4 * dModel * dModel);
}
}head 0 weights of the last position: 0.322 0.268 0.056 0.355
head 1 weights of the last position: 0.137 0.140 0.217 0.507
concatenated: 4 x 8, after Wo: 4 x 8
parameters in Wq, Wk, Wv, Wo: 256Each head computes its own weights. For the last position, head 0 gives 0.322, 0.268, 0.056, and 0.355, and head 1 favors position 3 more strongly (0.507). The two heads look at the same sequence differently, because they use different slices of the projected vectors.
The concatenated result has 8 numbers per position, the same as d_model, and the output projection Wo keeps that shape.
The four matrices Wq, Wk, Wv, and Wo each have d_model * d_model parameters, so the attention layer has 4 * 8 * 8 = 256 parameters here. A real layer with d_model = 4096 has 4 * 4096 * 4096, about 67 million, in the case of plain multi-head attention.
Example 10: grouped-query attention
During generation, an engine stores the keys and values of every earlier position so it does not have to compute them again. This store is the KV cache (key-value cache, built in lesson 11). Its size is proportional to the number of key and value heads, and for long contexts it can take gigabytes.
Grouped-query attention (GQA) reduces the cache. The model keeps n_head query heads but uses only n_head_kv key/value heads, where n_head_kv divides n_head. Several query heads share one key/value head: query head h uses key/value head h / (n_head / n_head_kv), with integer division. The mix works as before. The cache shrinks by the ratio n_head / n_head_kv, and the query side keeps all its heads.
The program first prints which key/value head each query head uses. It then does the parameter and cache arithmetic for a published configuration: 32 layers, 32 query heads, 8 key/value heads, head size 128, and model width 4096 (Llama 3 8B).
The same two numbers appear under different names in different tools. llama.cpp prints them as n_head and n_head_kv. A GGUF file stores them as <architecture>.attention.head_count and .head_count_kv. A Hugging Face config.json calls them num_attention_heads and num_key_value_heads.
import java.util.Locale;
public class GroupedQuery {
public static void main(String[] args) {
int smallQ = 8, smallKv = 2, group = smallQ / smallKv;
for (int h = 0; h < smallQ; h++) System.out.println("query head " + h + " -> key/value head " + h / group);
// The tiny test model of lesson 1: 4 layers, 8 query heads, 2 key/value heads, head size 32, 16-bit cache.
long tiny = 2L * 4 * 2 * 32 * 2 * 4096;
System.out.println(String.format(Locale.ROOT, "tiny model, 4096 tokens: %,d bytes = %.2f MiB of KV cache", tiny, tiny / Math.pow(1024, 2)));
int layers = 32, heads = 32, kv = 8, head = 128, width = 4096, bytes = 2;
long wq = (long) width * heads * head, wk = (long) width * kv * head, wo = (long) heads * head * width;
System.out.println(String.format(Locale.ROOT, "per layer: Wq %,d Wk %,d Wv %,d Wo %,d parameters", wq, wk, wk, wo));
for (int nkv : new int[] {32, 8}) {
long perToken = 2L * layers * nkv * head * bytes;
System.out.println(String.format(Locale.ROOT, "%2d key/value heads: %,d bytes of KV cache per token (%d KiB); 8192 tokens = %.2f GiB",
nkv, perToken, perToken / 1024, perToken * 8192 / Math.pow(1024, 3)));
}
}
}query head 0 -> key/value head 0
query head 1 -> key/value head 0
query head 2 -> key/value head 0
query head 3 -> key/value head 0
query head 4 -> key/value head 1
query head 5 -> key/value head 1
query head 6 -> key/value head 1
query head 7 -> key/value head 1
tiny model, 4096 tokens: 4,194,304 bytes = 4.00 MiB of KV cache
per layer: Wq 16,777,216 Wk 4,194,304 Wv 4,194,304 Wo 16,777,216 parameters
32 key/value heads: 524,288 bytes of KV cache per token (512 KiB); 8192 tokens = 4.00 GiB
8 key/value heads: 131,072 bytes of KV cache per token (128 KiB); 8192 tokens = 1.00 GiBThe first lines show which key/value head serves each query head. Query heads 0 to 3 use head 0, and query heads 4 to 7 use head 1. Each group has 4 query heads.
The tiny test model from lesson 4 has this shape: 8 query heads, 2 key/value heads, head size 32, and 4 layers. Its KV cache for 4,096 tokens is 4,194,304 bytes (4.00 MiB). This matches the line printed by the pinned llama.cpp when it loads the model. The real log lines are:
print_info: n_head = 8
print_info: n_head_kv = 2
print_info: n_embd_head_k = 32
llama_kv_cache: size = 4.00 MiB ( 4096 cells, 4 layers, 1/1 seqs), K (f16): 2.00 MiB, V (f16): 2.00 MiBIn the 8B configuration, the key and value matrices (Wk and Wv) have a quarter of the parameters of Wq: 4.2 million against 16.8 million per layer. They produce 8 heads and not 32.
The cache needs 524,288 bytes per token with 32 key/value heads and 131,072 bytes with 8. At 8,192 tokens that is 4.00 GiB against 1.00 GiB. The cache formula, per token, is:
2 (K and V) * layers * n_head_kv * head_size * bytes per valueFor the 8B example with 8 key/value heads: 2 * 32 * 8 * 128 * 2 = 131,072 bytes. Lesson 19 uses this formula to plan memory for real models.
Example 11: multi-query attention
Multi-query attention (MQA) is the extreme case of grouped-query attention with n_head_kv = 1. All query heads share one key head and one value head, which gives the smallest cache. Ordinary multi-head attention is the opposite extreme, with n_head_kv = n_head.
The program lists the cache cost of the three choices for the same model shape: 32 query heads of size 128 in 32 layers, with 16-bit cache values. It also prints the group size, which is the number of query heads that share one key/value head.
import java.util.Locale;
public class Variants {
public static void main(String[] args) {
int layers = 32, heads = 32, head = 128, bytes = 2, context = 32768;
String[] names = {"multi-head (MHA)", "grouped-query (GQA)", "multi-query (MQA)"};
int[] kvHeads = {32, 8, 1};
System.out.println("variant kv heads group KiB/token GiB at 32768 tokens");
for (int i = 0; i < names.length; i++) {
long perToken = 2L * layers * kvHeads[i] * head * bytes;
System.out.println(String.format(Locale.ROOT, "%-20s %8d %6d %10d %20.2f", names[i], kvHeads[i], heads / kvHeads[i], perToken / 1024,
perToken * context / Math.pow(1024, 3)));
}
}
}variant kv heads group KiB/token GiB at 32768 tokens
multi-head (MHA) 32 1 512 16.00
grouped-query (GQA) 8 4 128 4.00
multi-query (MQA) 1 32 16 0.50Multi-head attention stores 512 KiB per token, so a 32,768-token context needs 16 GiB for the cache alone. Grouped-query attention with 8 key/value heads cuts it to 4 GiB. Multi-query attention with one head cuts it to 0.5 GiB.
The group size is the number of query heads that read the same cached keys and values: 1, 4, and 32. Llama 3 8B, the configuration used in Example 10, has 8 key/value heads.
The choice is a trade-off between cache size and the number of independent key/value heads. The number of key/value heads is fixed by the shapes of the model's weight matrices, and the model file records it as n_head_kv. An engine cannot change it after training.
Example 12: cost with sequence length
Attention gets more expensive as the text gets longer. This example counts how much.
Each pair of positions (i, j) with j <= i needs one score (a dot product of length d_head) and one weighted addition of a value (also length d_head). That is about 4 * d_head FLOPs per pair, per head, per layer: 2 for the score and 2 for the mix. (A FLOP is one floating-point operation, lesson 8, Example 14.)
The number of pairs depends on the step:
- In prefill (processing a whole prompt of
ntokens at once), there aren * (n + 1) / 2pairs. The work grows with the square ofn. - In a decode step (generating one token with cached keys and values), there are only
npairs. The work grows linearly.
The weight matrices cost 2 * P FLOPs per token regardless of n, where P is the parameter count.
The program compares the two costs for a model with 8 billion parameters and the shape of Example 10.
import java.util.Locale;
public class AttentionCost {
public static void main(String[] args) {
int layers = 32, heads = 32, head = 128;
double perPair = 4.0 * head * heads * layers; // FLOPs per (query, key) pair across all heads and layers
double perToken = 2 * 8e9; // weight FLOPs per token for an 8B model
System.out.println(" n prefill attention prefill weights decode attention decode weights");
for (int n : new int[] {1024, 8192, 32768, 131072}) {
double prefillAttn = perPair * n * (n + 1.0) / 2;
double prefillWeights = perToken * n;
double decodeAttn = perPair * n;
System.out.println(String.format(Locale.ROOT, "%7d %14.1f TFLOP %12.1f TFLOP %12.1f GFLOP %11.1f GFLOP",
n, prefillAttn / 1e12, prefillWeights / 1e12, decodeAttn / 1e9, perToken / 1e9));
}
}
} n prefill attention prefill weights decode attention decode weights
1024 0.3 TFLOP 16.4 TFLOP 0.5 GFLOP 16.0 GFLOP
8192 17.6 TFLOP 131.1 TFLOP 4.3 GFLOP 16.0 GFLOP
32768 281.5 TFLOP 524.3 TFLOP 17.2 GFLOP 16.0 GFLOP
131072 4503.6 TFLOP 2097.2 TFLOP 68.7 GFLOP 16.0 GFLOPThe attention work for a prompt grows with the square of its length, and the work of the weight matrices grows linearly. At 1,024 tokens, attention is 0.3 TFLOP against 16.4 TFLOP for weights, so attention is minor. At 8,192 tokens it is 17.6 TFLOP against 131.1. At 131,072 tokens it is 4,503.6 TFLOP against 2,097.2, so attention has overtaken everything else.
A decode step shows the same shift in a different form. Attention costs 0.5 GFLOP at 1,024 tokens and 68.7 GFLOP at 131,072, while the weights always cost 16 GFLOP.
These are FLOP counts. The speed of a decode step is limited by how many bytes must be read (the weights, plus the whole KV cache), and not by FLOPs. Lesson 17 returns to this with bandwidth numbers, and the attention-optimizations lesson later in this course covers long contexts. The attention counts assume the causal mask halves the pairs, and they skip the cost of softmax. The projections are part of the weight counts.
Example 13: a hand-built head that attends to the previous token
Trained models contain heads with clear jobs. The simplest is the previous-token head: every position attends to the position immediately before it. This example builds such a head without any training, so the mechanism is visible.
Each input vector has 6 numbers: 4 position slots (a one-hot code of the position) and 2 content numbers. The three matrices are set by hand:
Wkcopies the position code into the key.Wqcopies the position code shifted by one slot, so positioniproduces a query that matches the key of positioni - 1. It multiplies by 10 to make the match sharp.Wvcopies the content.
Position 0 has no earlier position, so it attends to itself. The program prints the attention weights and the output of each position. Compare each output with the content of the previous position.
import java.util.Locale;
public class PreviousTokenHead {
public static void main(String[] args) {
int n = 4;
float[][] content = {{1, 0}, {0, 1}, {1, 1}, {2, 0}};
float[][] x = new float[n][6];
for (int i = 0; i < n; i++) { x[i][i] = 1; x[i][4] = content[i][0]; x[i][5] = content[i][1]; }
float[][] wq = new float[6][4], wk = new float[6][4], wv = new float[6][2];
for (int p = 0; p < 4; p++) wk[p][p] = 1; // key = position code
for (int p = 1; p < 4; p++) wq[p][p - 1] = 10; // query of position p has its 10 in slot p - 1
wq[0][0] = 10; // position 0 asks for itself
wv[4][0] = 1;
wv[5][1] = 1;
float[][] weights = new float[n][n], out = new float[n][2];
for (int i = 0; i < n; i++) {
float[] q = new float[4];
for (int c = 0; c < 4; c++) for (int p = 0; p < 6; p++) q[c] += x[i][p] * wq[p][c];
double[] e = new double[i + 1];
double sum = 0;
for (int j = 0; j <= i; j++) {
double dot = 0;
for (int c = 0; c < 4; c++) {
float kj = 0;
for (int p = 0; p < 6; p++) kj += x[j][p] * wk[p][c];
dot += q[c] * kj;
}
e[j] = Math.exp(dot / Math.sqrt(4));
sum += e[j];
}
for (int j = 0; j <= i; j++) {
weights[i][j] = (float) (e[j] / sum);
for (int d = 0; d < 2; d++) {
float vj = 0;
for (int p = 0; p < 6; p++) vj += x[j][p] * wv[p][d];
out[i][d] += weights[i][j] * vj;
}
}
}
System.out.println("attention weights:");
for (float[] row : weights) {
StringBuilder sb = new StringBuilder();
for (float w : row) sb.append(String.format(Locale.ROOT, "%8.3f", w));
System.out.println(sb);
}
for (int i = 0; i < n; i++) {
System.out.println(String.format(Locale.ROOT, "position %d: own content (%.0f, %.0f) -> output (%.2f, %.2f)", i, content[i][0], content[i][1], out[i][0], out[i][1]));
}
}
}attention weights:
1.000 0.000 0.000 0.000
0.993 0.007 0.000 0.000
0.007 0.987 0.007 0.000
0.007 0.007 0.980 0.007
position 0: own content (1, 0) -> output (1.00, 0.00)
position 1: own content (0, 1) -> output (0.99, 0.01)
position 2: own content (1, 1) -> output (0.01, 0.99)
position 3: own content (2, 0) -> output (1.00, 0.99)The weight matrix shows a previous-token head. Position 1 puts 0.993 on position 0, position 2 puts 0.987 on position 1, and position 3 puts 0.980 on position 2. Position 0 can only attend to itself.
The values carry the content, so each output is the content of the previous position. Position 3 has own content (2, 0), and its output is (1.00, 0.99), which is the content (1, 1) of position 2.
Nothing was trained. The query matrix moves the position code by one slot, and the key matrix keeps it, so the dot product is large exactly when j = i - 1. Trained models learn such heads on their own, together with heads for other jobs.
The factor 10 in Wq sets how sharp the choice is. After the division by sqrt(4), the matching score is 5 and the others are 0, which gives a weight of about 0.98 on the match.
In a real model, positions do not come from one-hot codes. They come from RoPE (rotary position embedding), which rotates each query and key by an angle that depends on its position (lesson 10).
Interview points
- Attention is the only layer in which positions exchange information. Every other layer works on each position alone.
- Query, key, and value come from three learned projections of the same input. The score
q_i . k_jdecides how much ofv_jgoes into positioni. - Scores are divided by
sqrt(d_head)because the variance of a dot product ofd_headunit-variance terms isd_head. Without the division, softmax saturates and gradients vanish. - The causal mask sets scores for later positions to minus infinity before softmax. Their weights are then exactly 0, and the output at a position never depends on later tokens.
- In matrix form, attention is
softmax(mask(Q Kᵀ / sqrt(d_head))) V, which has ann x nscore matrix for each head. - Multi-head attention runs
n_headindependent heads of sized_model / n_head, concatenates them, and applies an output projectionWo. - Grouped-query attention keeps
n_headquery heads and uses onlyn_head_kvkey/value heads, so the KV cache shrinks byn_head / n_head_kv. Multi-query attention is the casen_head_kv = 1. - KV cache bytes per token are
2 * layers * n_head_kv * head_size * bytes per value. - Prefill attention cost grows with
n^2. A decode step with a KV cache costs work proportional ton. - The
n x nscore matrix is why long prompts need fast attention kernels that process blocks, and not the whole matrix at once. - A previous-token head, where the query of position
imatches the key of positioni - 1, is a small example of a pattern that trained models learn.
Summary
- An embedding lookup gives a token the same vector everywhere. Attention changes each vector using the other positions.
- A uniform average of earlier vectors is the simplest mix. Dot-product scores passed through softmax give weights that depend on content.
- Queries, keys, and values are three projections,
x Wq,x Wk, andx Wv. Scores are not symmetric. - Attention weights are
softmax(q_i . k_j / sqrt(d_head))overj <= i. The output is the weighted sum of the value vectors. - The standard deviation of unscaled scores equals
sqrt(d_head)(measured: 2.00, 8.23, and 31.67 for 4, 64, and 1,024). Unscaled softmax gives a largest weight of 0.965 at 1,024 dimensions, and scaled softmax stays near 0.36. - The causal mask makes the weight matrix lower triangular, with exact zeros above the diagonal.
- The matrix form equals the loop form:
A = softmax(mask(Q Kᵀ / sqrt(d_head)))andOutput = A V. - Multi-head attention splits
d_modelinton_headslices, attends in each, concatenates the results, and multiplies byWo. The layer has4 * d_model^2parameters. - Grouped-query attention shares each key/value head between
n_head / n_head_kvquery heads. A configuration with 32 layers, 8 key/value heads, and head size 128 caches 128 KiB per token in 16-bit values. - The tiny test model (8 query heads, 2 key/value heads, 4 layers, head size 32) has a 4.00 MiB KV cache for 4,096 tokens, as llama.cpp reports.
- Multi-query attention uses one key/value head and has the smallest cache: 16 KiB per token for the 32-layer example.
- Prefill attention work is quadratic in the prompt length, decode attention work is linear in the context length, and weight work per token is constant.