Matrix multiplication

To multiply two matrices, take each row of the first and each column of the second and compute their dot product. The inner sizes must match, and the outer sizes give the shape of the result. This one operation is most of the work a language model does.

Key idea: C = AB: C[i, j] = row i of A · column j of B; (m × n)(n × p) = m × p; inner sizes must match; order matters; matrix × vector = weighted sum of columns
Real world: Three customers’ shopping lists (rows) times a price table with a column for each of two shops gives a 3 × 2 table: what each customer would pay at each shop.
In an LLM: Q = X·W_Q turns 10 tokens × 768 numbers into 10 queries of 64 numbers each: a (10 × 768)(768 × 64) multiplication, 491,520 multiply-adds, for one head of one layer.

Rows meet columns

is many dot products organised in a grid. To compute $C = AB$ :

Cell $(i, j)$ of $C$ is the dot product of row $i$ of $A$ with column $j$ of $B$ .

C_{ij} = \sum_{k=1}^{n} A_{ik}\, B_{kj} = A_{i1}B_{1j} + A_{i2}B_{2j} + \dots + A_{in}B_{nj}

For the dot product to make sense, a row of $A$ and a column of $B$ must have the same length. So the number of columns of $A$ must equal the number of rows of $B$ . The result has as many rows as $A$ and as many columns as $B$ :

(m \times \underline{n}) \cdot (\underline{n} \times p) = m \times p

The two inner sizes must match and disappear; the outer sizes stay.

Watch it, one product at a time

Press play. The highlighted row of $A$ meets the highlighted column of $B$ ; each pair is multiplied, the products are added into a running sum, and the cell fills in. Then the same row meets the next column. Change the sizes with the buttons, and click any finished cell of $C$ to jump back to it.

rows of A

shared size

columns of B

start

B 3 × 2 →

1−1−11−31

1−10103

↑ A 2 × 3↑ C = AB 2 × 2

Laid out like a times table: each cell of C sits in line with the row of A on its left and the column of B above it.

how it works

Each cell of C pairs one row of A with one column of B: multiply the pairs, add them up. Press play.

inner sizes3 = 3 ✓

result shape2 × 2

multiplications2×2×3 = 12

A is 2 × 3 and B is 3 × 2. The inner sizes match (3 and 3), so the product exists and has shape 2 × 2: the outer sizes. It needs 4 dot products of length 3.

Press play, or click any cell of the result to see how it was made. Click a number in A or B to change it (it counts up from −3 to 3).

Things to notice:

$C$ has $m \times p$ cells, each needing $n$ multiplications, so the whole product is $m \times n \times p$ multiply-adds. Double all three sizes and the work goes up 8 times.
The cells don’t depend on each other. A GPU computes thousands of them at the same time, which is why models are built from matrix multiplications.
Order matters. $AB$ and $BA$ are usually different, and often only one of them exists.

example · shopping lists and price tables

Three customers each want some bread, milk and apples: a $3 \times 3$ matrix of quantities, one row per customer. Two shops each have a price for each item: a $3 \times 2$ matrix, one column per shop. Multiply them and you get a $3 \times 2$ table: row = customer, column = shop, entry = that customer’s bill at that shop. Each entry is one shopping-bill dot product; the multiplication does all six at once.

Do the shapes fit?

Most bugs in model code are shape bugs, and most are caught by checking inner sizes. Try the presets, which are real multiplications inside GPT-2 small reading 10 tokens:

(10 × 768) · (768 × 64) → 10 × 64

set your own sizes

A rows2A cols3B rows3B cols2

legal?yes

result10 × 64

multiply-adds491,520

10 tokens of width 768, times the query matrix: every token gets a 64-number query.

Pick a multiplication from inside GPT-2 (10 tokens), or set your own sizes. Rectangles are drawn on a log scale so 50,257 fits on the screen.

example · one attention head in GPT-2

Ten tokens arrive as $X$ , shape $10 \times 768$ . The query matrix $W_Q$ of one head is $768 \times 64$ . Inner sizes 768 and 768 match, so $Q = XW_Q$ is $10 \times 64$ : one 64-number query per token, costing $10 \times 768 \times 64 = 491{,}520$ multiply-adds. Keys come out the same shape. To score every query against every key we need $(10 \times 64)(64 \times 10)$ , so we multiply by $K^\top$ , not $K$ : that is the reason for the transpose in $QK^\top$ . See queries, keys and values for the full pipeline.

A matrix times a vector

When $B$ is a single column $\mathbf{x}$ , the product $M\mathbf{x}$ is a vector. There are two ways to see it:

Row view: each output number is one row of $M$ dotted with $\mathbf{x}$ .
Column view: the output is a weighted sum of $M$ ’s columns, with the numbers in $\mathbf{x}$ as the weights: $M\mathbf{x} = x_1 \cdot (\text{column } 1) + x_2 \cdot (\text{column } 2)$ .

The column view is the key to the next lesson. Drag the columns and slide the weights:

x₁1.00x₂1.00

M =2.0−0.50.51.5×1.001.00=[1.50, 2.00]

column view: 1.00×[2.0, 0.5] + 1.00×[−0.5, 1.5] = [1.50, 2.00]

row view: [2.0×1.00 + (−0.5)×1.00, 0.5×1.00 + (1.5)×1.00] = [1.50, 2.00]

Two ways to get the same answer. Row view: each output number is a row of M dotted with x. Column view: the output is 1.00 lots of column 1 plus 1.00 lots of column 2. The column view explains the next lesson: whatever x is, M x is built from M’s columns, so the columns say where M sends the two axes.

M has two columns, the blue and amber arrows (drag their tips). M times x walks x₁ of the first column, then x₂ of the second. The green arrow is the result.

Transformer code puts vectors in rows instead, and writes $\mathbf{x}W$ . Then the same idea reads: the output is a weighted sum of $W$ ’s rows, with the weights taken from $\mathbf{x}$ . Embedding lookup is the extreme case: a one-hot $\mathbf{x}$ puts weight 1 on a single row.

going deeper

Rules

Associative: $(AB)C = A(BC)$ . You can choose which product to do first; the cost can differ a lot.
Distributive: $A(B + C) = AB + AC$ .
Not commutative: in general $AB \ne BA$ .
Identity: the identity matrix $I$ (1s on the diagonal, 0s elsewhere) changes nothing: $AI = IA = A$ .
Transpose of a product: $(AB)^\top = B^\top A^\top$ .

Cost

$(m \times n)(n \times p)$ takes $mnp$ multiplications and about as many additions, so about $2mnp$ floating-point operations (FLOPs). For a model with $N$ parameters, a forward pass costs about $2N$ FLOPs per token, because nearly every parameter takes part in one multiply-add. See what a forward pass costs.

In code

import numpy as np
A = np.array([[1, 2, -1],
              [0, 3,  2]])        # 2 × 3
B = np.array([[ 2, 1],
              [ 0, 1],
              [-1, 3]])           # 3 × 2
A @ B                             # [[3, 0], [-2, 9]], shape (2, 2)
B @ A                             # also legal here, but shape (3, 3)
 
# C[i, j] by hand
C = np.zeros((2, 2))
for i in range(2):
    for j in range(2):
        for k in range(3):
            C[i, j] += A[i, k] * B[k, j]
 
import torch
X  = torch.randn(10, 768)
Wq = torch.randn(768, 64)
(X @ Wq).shape                    # torch.Size([10, 64])
Q = X @ Wq; K = X @ torch.randn(768, 64)
(Q @ K.T).shape                   # torch.Size([10, 10])
# Q @ K -> RuntimeError: mat1 and mat2 shapes cannot be multiplied (10x64 and 10x64)