datarekha

Matrix multiplication, the right mental model

Element-by-element formulas don't stick. Rows-times-columns and shape arithmetic do. Here's how matmul actually works in ML code.

7 min read Intermediate Math for ML Lesson 5 of 37

What you'll learn

  • The "row dot column" intuition for each output entry
  • The shape rule (m×n)(n×p) = (m×p) and why it constrains everything
  • Why AB is generally not BA, and what the identity matrix does
  • How a linear layer `Wx + b` becomes a batched matmul `XW^T + b`

Before you start

The last lesson ended on a puzzle: stack two transforms — rotate, then scale — and there is a single matrix that does both in one stroke. Building that matrix is matrix multiplication, and the formula most textbooks lead with, C[i,j] = Σ_k A[i,k]·B[k,j], is both perfectly correct and the worst way to hold it in your head. Here is the version that sticks: each entry of the product is a dot product — one row of the left matrix against one column of the right. Once that picture clicks, everything else — the shape rule, why the order cannot be swapped, why a batched layer is written the way it is — falls straight out of it.

Row dot column

To compute C = A @ B, read off each entry as a single dot product:

C[i, j] = (row i of A) · (column j of B)

That is the whole rule. The element at position (i, j) takes one row of A and one column of B, multiplies them element-wise, and sums.

A1234×B5678=A·B1·5+2·71·6+2·83·5+4·73·6+4·8=19=22=43=50row i of A dotted with column j of B → entry (i, j)
Each output cell is a row-times-column dot product.
import numpy as np

A = np.array([[1, 2],
              [3, 4]])
B = np.array([[5, 6],
              [7, 8]])

# @ is matrix multiply in NumPy
C = A @ B
print(C)

# Each entry is a dot product. Sanity-check C[0, 1]:
print("C[0, 1] =", A[0, :] @ B[:, 1])   # row 0 of A . col 1 of B
[[19 22]
 [43 50]]
C[0, 1] = 22

The shape rule controls everything

A matrix multiply is legal only when the inner dimensions match:

(m, n)@(n, p)=(m, p)inner dims must match → cancelouter dims survive → result shape
The shape rule: inner dims cancel, outer dims survive.

The inner n cancels; the outer m and p become the result shape. This one rule is what you will lean on for the rest of your ML career to catch bugs before they run.

import numpy as np

A = np.zeros((4, 3))      # 4 rows, 3 cols
B = np.zeros((3, 5))      # 3 rows, 5 cols
print("A @ B shape:", (A @ B).shape)   # inner 3s match

# Now a mismatch
C = np.zeros((4, 3))
D = np.zeros((5, 3))      # inner dims (3 vs 5) don't match
try:
    C @ D
except ValueError as e:
    print("Failure:", e)
A @ B shape: (4, 5)
Failure: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 5 is different from 3)

When your code crashes with a “shapes not aligned” error, the fix is almost always a transpose somewhere.

Why matmul is composing transformations

Now the payoff promised last lesson. Each matrix is a function on vectors: A @ x applies transform A, and B @ (A @ x) applies A and then B. Matrix multiplication collapses that two-step journey into one: (B @ A) @ x. The product B @ A is the single matrix that performs both transforms at once — which is the whole reason it is called composition, and the whole reason two stacked linear layers W2 @ (W1 @ x) are identical to the single layer (W2 @ W1) @ x.

Order matters: AB ≠ BA

Ordinary numbers commute (3 · 5 == 5 · 3). Matrix multiplication generally does not: even when both products exist, they are usually different matrices (and often different shapes).

import numpy as np

A = np.array([[1, 2],
              [0, 1]])
B = np.array([[1, 0],
              [3, 1]])

print("A @ B:")
print(A @ B)
print("B @ A:")
print(B @ A)
A @ B:
[[7 2]
 [3 1]]
B @ A:
[[1 2]
 [3 7]]

Geometrically, “shear then rotate” is simply not “rotate then shear.” In ML this is why layer order is part of the model — you cannot shuffle a stack of linear layers and expect the same function back.

The identity matrix: matmul’s “1”

I @ A == A for any matrix A — the identity is the do-nothing transform, the 1 of matrix multiplication. It also serves as a starting point: ridge regression’s (XᵀX + λI)⁻¹ adds a multiple of I to nudge a near-singular matrix back to invertible.

import numpy as np

I = np.eye(3)
A = np.array([[2.0, 1.0, 0.0],
              [0.0, 3.0, 1.0],
              [1.0, 0.0, 4.0]])

print("I @ A == A?", np.allclose(I @ A, A))
print("A @ I == A?", np.allclose(A @ I, A))
I @ A == A? True
A @ I == A? True

A linear layer, the way frameworks actually write it

In math you meet y = Wx + b for a single sample: with W of shape (d_out, d_in) and x of shape (d_in,), it produces y of shape (d_out,). But a real training loop processes a batch of n samples at once. Stack them into X of shape (n, d_in) and the layer becomes:

Y = X @ W.T + b        # shape (n, d_out)

The transpose appears because we kept W in the (d_out, d_in) convention while the batch X puts features in the last axis — and one matmul now does the work of n separate Wx calls.

import numpy as np

# A linear layer: 4 input features -> 3 outputs
W = np.array([
    [ 0.5, -0.2,  0.1,  0.3],
    [ 0.0,  0.4,  0.2, -0.1],
    [-0.3,  0.1,  0.0,  0.5],
])   # shape (3, 4) = (d_out, d_in)
b = np.array([0.1, -0.1, 0.05])

# A batch of 5 samples, 4 features each (seeded for reproducibility)
X = np.random.default_rng(0).normal(size=(5, 4))

Y = X @ W.T + b        # (5,4) @ (4,3) + (3,) -> (5,3)
print("X:", X.shape, "W.T:", W.T.shape, "Y:", Y.shape)
print(Y.round(3))
X: (5, 4) W.T: (4, 3) Y: (5, 3)
[[ 0.285 -0.035  0.052]
 [ 0.174  0.211  0.72 ]
 [-0.049 -0.735  0.155]
 [-1.363 -0.363  0.359]
 [ 0.245 -0.248  0.703]]

The + b rides on broadcastingb of shape (3,) aligns to (1, 3) and is added to every row. One line, the whole layer.

In one breath

Matrix multiplication builds the single matrix that does two transforms at once, and the rule that makes it memorable is row dot column: entry (i, j) of A @ B is row i of A dotted with column j of B. The shape rule (m, n) @ (n, p) = (m, p) — inner dims cancel, outer dims survive — catches most bugs before they run; the order is load-bearing (A @ B ≠ B @ A, because “shear then rotate” ≠ “rotate then shear”); the identity is matmul’s 1; and a real linear layer over a batch is just X @ W.T + b, one matmul doing the work of n separate Wx calls.

Practice

Quick check

0/3
Q1A has shape (32, 128) and B has shape (128, 64). What is the shape of A @ B?
Q2Why do we write a batched linear layer as `X @ W.T + b` instead of `W @ X + b`?
Q3You compute A @ B and B @ A. Both succeed but give different matrices. Is that a bug?

A question to carry forward

Everything so far has lived in two dimensions: a matrix is a grid, a vector is a row, and a product is a grid of dot products. But the X @ W.T + b you just wrote was already straining at that frame — the + b quietly added a length-3 vector to a whole batch without a loop, and a real network juggles batches of images with height, width, and colour channels all at once. Those objects are not matrices any more; they carry three, four, even five axes.

So here is the thread onward: what is the object that climbs the scalar → vector → matrix ladder one more rung — the tensor — and what is the single permissive, endlessly-tripped-over rule (broadcasting) that lets a small array act on a big one with no loop at all? It is the rule behind that quiet + b, and behind the silent shape bug that has wrecked more training runs than any crash ever has.

Sign in to track your progress

Completed lessons, your XP, level, and streak save to your account — it's free and takes a few seconds.

Practice this in an interview

All questions
Return an array where each element is the product of all other elements, without using division.

Make two passes: a left pass where each position accumulates the product of everything to its left, then a right pass (using a rolling variable) that multiplies in everything to its right. The two passes together give the complete product-of-all-others in O(n) time and O(1) extra space (excluding output).

When should you use apply, map, or applymap versus vectorized pandas operations, and what are the performance implications?

Vectorized pandas and NumPy operations operate on entire arrays in compiled C/Fortran code and should always be your first choice. apply runs a Python function row- or column-wise in a Python loop, map transforms a single Series element-by-element, and applymap (DataFrame.map in pandas 2.1+) applies a function to every scalar — all three are orders of magnitude slower than vectorized equivalents.

What is SwiGLU and why did modern LLMs replace the ReLU/GELU MLP with it?

SwiGLU is a gated feed-forward layer: it projects the input into two paths, passes one (the gate) through SiLU, multiplies the two elementwise, then projects down — SiLU(x·W_gate) ⊙ (x·W_up) · W_down. The elementwise gate is a smooth, learned, per-feature volume control, so the network can decide how much of each feature to pass rather than a hard ReLU on/off. It gives better quality per parameter, which is why LLaMA, Mistral, and Qwen use it; to keep the parameter budget equal it uses a smaller (~2/3) hidden width since it has three weight matrices instead of two.

How do you attribute and control ML spend across teams and models (FinOps for ML)?

Apply FinOps to ML by tagging every workload (training jobs, endpoints, GPU pools) by team, model, and environment so cost is attributable, then track unit-economics metrics like cost per prediction or per training run rather than just total spend. Set budgets and alerts, identify idle GPUs and overprovisioned endpoints, and enforce guardrails like autoscaling and instance-type policies. The goal is continuous visibility and accountability so teams optimize cost without killing experimentation.

Related lessons

Explore further

Skip to content