Home

05 Drills: Basic Matrix Math for Transformers

April 2026 (3542 Words, 20 Minutes)

exercise drills matrix linear-algebra lecture-5

A first-encounter drill book for students who want to compute the linear-algebra steps of a transformer by hand before they touch any code. Every exercise reuses the running examples from Lecture 5 — Inside the Transformer: the toy sequence <BOS> dragon drinks with \(T = 3, d \le 6\), and the “the cat sat on it” attention demo with \(T = 5, d_k = 64\).

This document is the bridge between “I can multiply numbers” and the proof-flavoured linear algebra for transformers companion exercises. Once you can finish every entry below in pen-and-paper, the proofs in that companion document become natural.

How to use this document

Table of contents

Chapter 1. Vectors before attention

The geometry of embeddings — dot product, length, cosine — is the only tool you need to interpret what \(QK^\top\) measures in chapter 4. We use 3-dimensional toy embeddings for two of the lecture’s running tokens.

Take

\[x_{\text{dragon}} = \begin{pmatrix} 1 \\ 2 \\ 0 \end{pmatrix}, \qquad x_{\text{drinks}} = \begin{pmatrix} 0 \\ 1 \\ 1 \end{pmatrix}.\]

Exercise 1.1. Compute the dot product \(x_{\text{dragon}}^\top x_{\text{drinks}}\).

Solution

\(x_{\text{dragon}}^\top x_{\text{drinks}} = 1 \cdot 0 + 2 \cdot 1 + 0 \cdot 1 = 2\).

Exercise 1.2. Compute \(\|x_{\text{dragon}}\|_2\), \(\|x_{\text{drinks}}\|_2\), and the cosine similarity \(\cos(x_{\text{dragon}}, x_{\text{drinks}})\).

Solution

\(\|x_{\text{dragon}}\|_2 = \sqrt{1 + 4 + 0} = \sqrt{5}\); \(\|x_{\text{drinks}}\|_2 = \sqrt{0 + 1 + 1} = \sqrt{2}\). Using exercise 1.1,

\[\cos(x_{\text{dragon}}, x_{\text{drinks}}) \;=\; \frac{x_{\text{dragon}}^\top x_{\text{drinks}}}{\|x_{\text{dragon}}\|_2 \, \|x_{\text{drinks}}\|_2} \;=\; \frac{2}{\sqrt{5}\,\sqrt{2}} \;=\; \frac{2}{\sqrt{10}} \;\approx\; 0.632.\]

Exercise 1.3 (MCQ). Two embeddings \(u, v\) point in the same direction (i.e. \(v = \alpha u\) for some \(\alpha > 0\)) but \(v\) is twice as long as \(u\). Which of the following is true?

(a) \(\cos(u, v) = 0\) and \(u^\top v = 0\). (b) \(\cos(u, v) = 1\) and \(u^\top v = 2\|u\|_2^2\). (c) \(\cos(u, v) = 1\) and \(u^\top v = \|u\|_2^2\). (d) \(\cos(u, v) = 2\) and \(u^\top v = 2\|u\|_2^2\).

Solution

(b). Cosine strips length out: same direction always gives \(\cos = 1\), regardless of magnitude. The dot product keeps length: \(u^\top v = u^\top(2u) = 2 u^\top u = 2\|u\|_2^2\). This is exactly why attention scores use the raw dot product \(q^\top k\) — magnitudes carry information that cosine throws away.

Chapter 2. Your first matrix product

The single most important computation in this entire document. We will redo it three times — entry-wise, then as a linear combination of rows of \(B\), then as a linear combination of columns of \(A\) — to land all three views of matrix multiplication in your fingers.

Take

\[A = \begin{pmatrix} 1 & 2 & 0 \\ 0 & 1 & 1 \end{pmatrix} \;\in\; \mathbb{R}^{2 \times 3}, \qquad B = \begin{pmatrix} 1 & 1 \\ 0 & 1 \\ 1 & 0 \end{pmatrix} \;\in\; \mathbb{R}^{3 \times 2}.\]

Exercise 2.1 (entry-wise). Compute \(A B\) by working out each of the four entries via the formula \((A B)_{ij} = \sum_{k=1}^{3} A_{ik} B_{kj}\).

Solution

\((A B)_{11} = 1\cdot 1 + 2\cdot 0 + 0\cdot 1 = 1\). \((A B)_{12} = 1\cdot 1 + 2\cdot 1 + 0\cdot 0 = 3\). \((A B)_{21} = 0\cdot 1 + 1\cdot 0 + 1\cdot 1 = 1\). \((A B)_{22} = 0\cdot 1 + 1\cdot 1 + 1\cdot 0 = 1\).

\[A B = \begin{pmatrix} 1 & 3 \\ 1 & 1 \end{pmatrix}.\]

Exercise 2.2 (row-as-linear-combination-of-rows-of-\(B\)). Let \(b_1^\top = (1, 1)\), \(b_2^\top = (0, 1)\), \(b_3^\top = (1, 0)\) be the rows of \(B\). Verify that

\[(\text{row 1 of } AB)^\top \;=\; A_{11}\, b_1^\top + A_{12}\, b_2^\top + A_{13}\, b_3^\top.\]
Solution

\(A_{11} b_1^\top + A_{12} b_2^\top + A_{13} b_3^\top = 1 \cdot (1, 1) + 2 \cdot (0, 1) + 0 \cdot (1, 0) = (1, 1) + (0, 2) + (0, 0) = (1, 3)\) — exactly row 1 of \(A B\).

This view says: row \(i\) of \(A B\) is a weighted blend of the rows of \(B\), with weights given by row \(i\) of \(A\). This is precisely how \(Z = A V\) works inside attention — chapter 4 of the lecture’s transformer forward-pass formulas and dimensions notes makes this point explicitly.

Exercise 2.3 (column-as-linear-combination-of-columns-of-\(A\)). Let \(a_1 = (1, 0)^\top\), \(a_2 = (2, 1)^\top\), \(a_3 = (0, 1)^\top\) be the columns of \(A\). Verify that

\[\text{column 1 of } A B \;=\; B_{11}\, a_1 + B_{21}\, a_2 + B_{31}\, a_3.\]
Solution

\(B_{11} a_1 + B_{21} a_2 + B_{31} a_3 = 1 \cdot (1, 0)^\top + 0 \cdot (2, 1)^\top + 1 \cdot (0, 1)^\top = (1, 0)^\top + (0, 0)^\top + (0, 1)^\top = (1, 1)^\top\) — exactly column 1 of \(A B\).

This view says: column \(j\) of \(A B\) is a weighted blend of the columns of \(A\), with weights given by column \(j\) of \(B\). The two views — row mixing and column mixing — are the same multiplication seen from different sides.

Chapter 3. Computing the query matrix row by row

Now we apply chapter 2 to the actual attention projection. Each row of X is one token’s embedding; each row of Q is the corresponding query vector. The right-multiplication by W_Q acts on each row independently — that is the whole reason the “row of Q for token \(t\)” is well-defined.

Take three tokens with 2-dimensional embeddings (small for legibility — d = 2 here, but the lecture uses d = 6, and the arithmetic is the same):

\[X = \begin{pmatrix} 1 & 0 \\ 1 & 2 \\ 0 & 1 \end{pmatrix} \;=\; \begin{pmatrix} x_{\langle\text{BOS}\rangle}^\top \\ x_{\text{dragon}}^\top \\ x_{\text{drinks}}^\top \end{pmatrix} \;\in\; \mathbb{R}^{3 \times 2}, \qquad W_Q = \begin{pmatrix} 1 & 1 \\ 0 & 1 \end{pmatrix} \;\in\; \mathbb{R}^{2 \times 2}.\]

Exercise 3.1. Compute the per-token query vector for dragon using the formula \(q_2^\top = x_2^\top W_Q\).

Solution

\(q_2^\top = (1, 2) \cdot \begin{pmatrix} 1 & 1 \\ 0 & 1 \end{pmatrix} = (1\cdot 1 + 2\cdot 0,\; 1\cdot 1 + 2\cdot 1) = (1, 3)\).

So the query vector for dragon is \(q_2 = (1, 3)^\top\). Notice the computation only uses row 2 of X and the entries of W_Q — no other row participates.

Exercise 3.2. Compute \(Q = X W_Q\) directly as a \((3 \times 2) \cdot (2 \times 2)\) matrix product.

Solution

Row by row:

\[Q = \begin{pmatrix} (1,0) \cdot W_Q \\ (1,2) \cdot W_Q \\ (0,1) \cdot W_Q \end{pmatrix} = \begin{pmatrix} 1 & 1 \\ 1 & 3 \\ 0 & 1 \end{pmatrix}.\]

Exercise 3.3. Confirm that the row of \(Q\) you computed in 3.1 matches row 2 of the matrix in 3.2.

Solution

Row 2 of \(Q\) is \((1, 3)\), which is exactly \(q_2^\top\) from 3.1. The takeaway: the matrix-stacked formula \(Q = X W_Q\) is just \(T\) independent per-token formulas \(q_t^\top = x_t^\top W_Q\) written in one line. Right-multiplication by \(W_Q\) acts on rows of \(X\) one at a time.

Chapter 4. The score matrix from Q and K

This is where attention earns the name “all queries dot all keys”. Each entry of S is one inner product; with T = 3 tokens we have 9 inner products — laid out in a \(3 \times 3\) grid.

Take \(Q\) from chapter 3 and a key matrix

\[Q = \begin{pmatrix} 1 & 1 \\ 1 & 3 \\ 0 & 1 \end{pmatrix}, \qquad K = \begin{pmatrix} 2 & 0 \\ 0 & 1 \\ 1 & 1 \end{pmatrix} \;\in\; \mathbb{R}^{3 \times 2}.\]

Exercise 4.1. Compute \(S = Q K^\top\) entry by entry. (Hint: write out \(K^\top\) first, then use chapter 2.)

Solution

\(K^\top = \begin{pmatrix} 2 & 0 & 1 \\ 0 & 1 & 1 \end{pmatrix} \in \mathbb{R}^{2 \times 3}\), and \(S = Q K^\top \in \mathbb{R}^{3 \times 3}\):

\[S = \begin{pmatrix} (1,1)\cdot(2,0) & (1,1)\cdot(0,1) & (1,1)\cdot(1,1) \\ (1,3)\cdot(2,0) & (1,3)\cdot(0,1) & (1,3)\cdot(1,1) \\ (0,1)\cdot(2,0) & (0,1)\cdot(0,1) & (0,1)\cdot(1,1) \end{pmatrix} = \begin{pmatrix} 2 & 1 & 2 \\ 2 & 3 & 4 \\ 0 & 1 & 1 \end{pmatrix}.\]

Exercise 4.2. State in one sentence what \(S_{2,3}\) means in attention terms, given that \(Q\) and \(K\) came from <BOS> dragon drinks.

Solution

\(S_{2,3} = q_2^\top k_3 = 4\) is the raw alignment between the query of token 2 (dragon) and the key of token 3 (drinks) — how much the dragon “asks for” what the drinks “offers”. Higher \(S_{2,3}\) means more attention from dragon to drinks after the row-wise softmax.

Exercise 4.3. Argue that \(S\) need not be symmetric (i.e. \(S_{ij} \ne S_{ji}\) in general). Confirm on the matrix in 4.1 by giving a specific \(i, j\) where they differ.

Solution

\(S_{ij} = q_i^\top k_j\), and \(Q\) and \(K\) are independently parameterised: \(W_Q \ne W_K\) in general, so \(q_i\) and \(k_i\) are different transformations of the same input \(x_i\). Concretely \(S_{1,2} = 1\) but \(S_{2,1} = 2\) in 4.1. Asymmetry matters: it lets attention say “dragon attends to drinks” without forcing “drinks attends to dragon” by the same amount.

Chapter 5. Softmax by hand on the cat-sat-on-it scores

This is the lecture’s headline numerical example. The query is the token it, and we are computing its attention distribution over [the, cat, sat, on, it] — five tokens. The pre-softmax scores are exactly those quoted in the lecture’s chapter 6 walk:

\[s = (3.2,\; 12.8,\; 5.6,\; 2.4,\; 7.2), \qquad d_k = 64.\]

Exercise 5.1 (without scaling). Compute \(\mathrm{softmax}(s)\) directly — exponentiate each entry, divide by the sum. Round each weight to the nearest 0.001.

Solution

\(\exp(3.2) \approx 24.5\), \(\exp(12.8) \approx 362\,217\), \(\exp(5.6) \approx 270.4\), \(\exp(2.4) \approx 11.0\), \(\exp(7.2) \approx 1\,339.4\).

Sum \(\approx 363\,862\). Dividing:

\[\mathrm{softmax}(s) \;\approx\; (0.000,\; 0.996,\; 0.001,\; 0.000,\; 0.004).\]

The mass collapses almost entirely onto cat. The other four tokens together carry less than 0.5% of the attention.

Exercise 5.2 (with the \(\sqrt{d_k}\) scaling). Now divide every score by \(\sqrt{d_k} = \sqrt{64} = 8\) before the softmax, and recompute. Round each weight to the nearest 0.01.

Solution

Scaled scores: \(s / 8 = (0.40,\; 1.60,\; 0.70,\; 0.30,\; 0.90)\).

Exponentials: \(\exp(0.40) \approx 1.49\), \(\exp(1.60) \approx 4.95\), \(\exp(0.70) \approx 2.01\), \(\exp(0.30) \approx 1.35\), \(\exp(0.90) \approx 2.46\).

Sum \(\approx 12.26\). Dividing:

\[\mathrm{softmax}(s / \sqrt{d_k}) \;\approx\; (0.12,\; 0.40,\; 0.16,\; 0.11,\; 0.20).\]

This is the lecture’s reported result: roughly 40% on cat, with the rest spread out. Same query, same keys — only the scaling changed.

Exercise 5.3 (MCQ). Why does the transformer divide by \(\sqrt{d_k}\) before the softmax?

(a) To make the softmax output sum to 1 (which it would not otherwise). (b) To prevent the softmax from saturating onto a single token when scores have large magnitude. (c) To force every entry of the score matrix to lie in \([-1, 1]\). (d) For numerical stability in the floating-point exponential.

Solution

(b). A softmax always sums to 1 by construction (so (a) is wrong); scaling does not bound entries to \([-1, 1]\) (so (c) is wrong); the standard trick for numerical stability is subtracting the row max, not dividing by \(\sqrt{d_k}\) (so (d) is wrong, although there is some indirect benefit). The real motivation is exactly the contrast between exercises 5.1 and 5.2: when scores grow with \(d_k\), the softmax saturates and gradients vanish, so we rescale to keep scores in a moderate range.

Chapter 6. Concatenation as block-stacking

Multi-head attention runs H independent single-head attentions on disjoint slices of the embedding, then concatenates their outputs. “Concatenation” here means horizontal block-stacking: place two matrices next to each other along the column axis.

Take \(H = 2\) heads, each of width \(d_k = 3\), giving total width \(d = H \cdot d_k = 6\). The two head outputs (each one \(T = 3\) tokens × \(d_k = 3\) features) are

\[Z^{(1)} = \begin{pmatrix} 1 & 0 & 0 \\ 0 & 1 & 0 \\ 0 & 0 & 1 \end{pmatrix}, \qquad Z^{(2)} = \begin{pmatrix} 1 & 1 & 1 \\ 2 & 2 & 2 \\ 3 & 3 & 3 \end{pmatrix}.\]

Exercise 6.1. Build the concatenation \(Z = [Z^{(1)} \mid Z^{(2)}] \in \mathbb{R}^{T \times d}\). State its shape explicitly.

Solution \[Z = \begin{pmatrix} 1 & 0 & 0 & 1 & 1 & 1 \\ 0 & 1 & 0 & 2 & 2 & 2 \\ 0 & 0 & 1 & 3 & 3 & 3 \end{pmatrix} \;\in\; \mathbb{R}^{3 \times 6}.\]

The number of rows (\(T = 3\)) is unchanged; the column count goes from \(d_k = 3\) (each head) to \(d = H \cdot d_k = 6\) after concatenating both heads.

Exercise 6.2. What shape must the output projection \(W_O\) have so that \(U = Z W_O\) is a valid matrix product? What shape does \(U\) end up with?

Solution

\(W_O \in \mathbb{R}^{d \times d} = \mathbb{R}^{6 \times 6}\). Then \(U = Z W_O \in \mathbb{R}^{3 \times 6}\) — same shape as the input H to the attention sub-layer. This is a deliberate design choice: the output of multi-head attention has shape \(T \times d\), ready to plug into the next residual connection without further reshaping.

Exercise 6.3 (parameter accounting). With \(d = 6\), how many learnable parameters does single-head attention use for W_Q, W_K, W_V (no W_O)? How many does multi-head attention with \(H = 2, d_k = 3\) use, including W_O? Where does the difference come from?

Solution

Single-head: three matrices W_Q, W_K, W_V each in \(\mathbb{R}^{d \times d} = \mathbb{R}^{6 \times 6}\). Total: \(3 \cdot 36 = 108\). No \(W_O\) in the bare single-head version.

Multi-head, \(H = 2, d_k = 3\): three projections per head, each in \(\mathbb{R}^{d \times d_k} = \mathbb{R}^{6 \times 3}\); two heads. Q/K/V total: \(3 \cdot H \cdot d \cdot d_k = 3 \cdot 2 \cdot 18 = 108\) — identical to single-head. Plus \(W_O \in \mathbb{R}^{6 \times 6}\), contributing 36. Grand total: 144.

Difference: exactly \(d^2 = 36\), all of it in W_O. Multi-head is not “more expensive per projection” — it is the same Q/K/V budget rearranged into H parallel subspaces, with one extra W_O to mix them back.

Chapter 7. The causal mask by hand

Decoder-only LLMs prevent each token from attending to future tokens. The mechanism is a single additive mask M injected into the score matrix before the softmax.

\[M_{ij} \;=\; \begin{cases} 0 & \text{if } j \le i, \\ -\infty & \text{if } j > i. \end{cases}\]

Take \(T = 4\) and the score matrix

\[S = \begin{pmatrix} 1 & 2 & 3 & 4 \\ 2 & 4 & 1 & 3 \\ 3 & 1 & 4 & 2 \\ 4 & 3 & 2 & 1 \end{pmatrix}.\]

Exercise 7.1. Write out \(M \in \mathbb{R}^{4 \times 4}\) explicitly.

Solution \[M = \begin{pmatrix} 0 & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 \end{pmatrix}.\]

Lower-triangular zeros, upper-triangular \(-\infty\). The diagonal is included (a token can always attend to itself).

Exercise 7.2. Compute \(\widetilde S = S + M\), then apply softmax row by row. Round each weight to two decimals.

Solution

After adding M:

\[\widetilde S = \begin{pmatrix} 1 & -\infty & -\infty & -\infty \\ 2 & 4 & -\infty & -\infty \\ 3 & 1 & 4 & -\infty \\ 4 & 3 & 2 & 1 \end{pmatrix}.\]

Softmax-by-row: \(\exp(-\infty) = 0\), so the masked entries vanish. Each row is a probability distribution over only the unmasked positions.

  • Row 1: only one finite entry. \(A_{1,:} = (1.00,\; 0,\; 0,\; 0)\).
  • Row 2: \(\exp(2) \approx 7.39\), \(\exp(4) \approx 54.60\); sum \(\approx 61.99\). \(A_{2,:} \approx (0.12,\; 0.88,\; 0,\; 0)\).
  • Row 3: \(\exp(3) \approx 20.09\), \(\exp(1) \approx 2.72\), \(\exp(4) \approx 54.60\); sum \(\approx 77.41\). \(A_{3,:} \approx (0.26,\; 0.04,\; 0.71,\; 0)\).
  • Row 4: all four entries unmasked. \(\exp(4) \approx 54.60\), \(\exp(3) \approx 20.09\), \(\exp(2) \approx 7.39\), \(\exp(1) \approx 2.72\); sum \(\approx 84.79\). \(A_{4,:} \approx (0.64,\; 0.24,\; 0.09,\; 0.03)\).

The full attention matrix:

\[A \;\approx\; \begin{pmatrix} 1.00 & 0 & 0 & 0 \\ 0.12 & 0.88 & 0 & 0 \\ 0.26 & 0.04 & 0.71 & 0 \\ 0.64 & 0.24 & 0.09 & 0.03 \end{pmatrix}.\]

Exercise 7.3. Look at A from 7.2. State what each row tells you about token \(i\)’s attention pattern. In particular, what does \(A_{1,1} = 1\) mean physically?

Solution

Row \(i\) of A is token \(i\)’s attention distribution over the keys it is allowed to see — positions \(1, \ldots, i\) only. Row 1 attends only to position 1 (itself), so \(A_{1,1} = 1\) says the first token has nothing to look at except itself. Each subsequent row gains one more position it can attend to, and the lower-triangular shape of A is the visible signature of causal masking.

Chapter 8. Permutation matrices

A permutation matrix P is just an identity matrix whose rows have been shuffled. Left-multiplying X by P shuffles X’s rows. This is the algebraic object behind the lecture’s permutation-equivariance lemma in chapter 9.

Exercise 8.1. Write the \(3 \times 3\) permutation matrix P that swaps rows 1 and 2 of any \(3 \times d\) matrix and leaves row 3 alone.

Solution \[P = \begin{pmatrix} 0 & 1 & 0 \\ 1 & 0 & 0 \\ 0 & 0 & 1 \end{pmatrix}.\]

Row 1 of P is the second standard basis vector \(e_2^\top\), row 2 is \(e_1^\top\), row 3 is \(e_3^\top\). Left-multiplying X by P therefore makes new-row-1 = old-row-2, new-row-2 = old-row-1, new-row-3 = old-row-3.

Exercise 8.2. Verify that \(P^\top P = I_3\) for the matrix in 8.1.

Solution

This P is symmetric (\(P^\top = P\)), so \(P^\top P = P P\). Compute:

\[P P = \begin{pmatrix} 0 & 1 & 0 \\ 1 & 0 & 0 \\ 0 & 0 & 1 \end{pmatrix} \begin{pmatrix} 0 & 1 & 0 \\ 1 & 0 & 0 \\ 0 & 0 & 1 \end{pmatrix} = \begin{pmatrix} 1 & 0 & 0 \\ 0 & 1 & 0 \\ 0 & 0 & 1 \end{pmatrix} = I_3.\]

Every permutation matrix is orthogonal: \(P^\top P = I\) and \(P^\top = P^{-1}\). Applying P then \(P^\top\) (or twice for an involution like this swap) brings you back where you started.

Exercise 8.3. Take

\[X = \begin{pmatrix} 1 & 2 & 3 \\ 4 & 5 & 6 \\ 7 & 8 & 9 \end{pmatrix}.\]

Compute \(P X\) using the matrix from 8.1, and confirm it really is X with rows 1 and 2 swapped.

Solution

Row \(i\) of \(P X\) is row \(i\) of P dotted with the columns of X:

\[P X = \begin{pmatrix} (0,1,0) \cdot \text{cols of } X \\ (1,0,0) \cdot \text{cols of } X \\ (0,0,1) \cdot \text{cols of } X \end{pmatrix} = \begin{pmatrix} 4 & 5 & 6 \\ 1 & 2 & 3 \\ 7 & 8 & 9 \end{pmatrix}.\]

Rows 1 and 2 are indeed swapped; row 3 is unchanged. This is the algebraic content of “swap two tokens in the input sequence” — the permutation lives entirely in P, and the rest of attention will see this swapped X exactly as if you had typed the tokens in a different order. That is why permutation equivariance (chapter 9 of the lecture) becomes a one-line algebraic identity once you write the swap as a matrix.

Where to next?

You can now compute every linear-algebra step of attention by hand. Three places to go from here:

© 2026 Dominik Müller   •  Powered by Soopr   •  Theme  Moonwalk