Matrices and matrix multiplication

Problem

A layer with 4096 outputs computes 4096 dot products for every input vector, and a batch holds thousands of input vectors. Writing that as a loop over dot products is slow and hides the structure. We need one operation that means "every row of this, dotted with every column of that". It should be easy to reason about on paper, cheap to count, and fast on hardware built for it.

That operation is matrix multiplication. Reading its shapes fluently is what lets you debug a model ("expected 512, got 768"), estimate its cost before running it, and see that attention, linear layers and embeddings lookups are the same few operations.

Intuition

A matrix is a grid of numbers with m rows and k columns, shape [m, k]. Think of it as a stack of m row vectors. Multiplying it by a [k, n] matrix takes each of the m rows and dots it with each of the n columns, giving an [m, n] grid of results. The inner sizes (both k) must match because a dot product needs two vectors of the same length, and they disappear from the result because each dot product collapses them to one number.

Each output cell sits where its row of A and its column of B line up. The inner size, 3, is used up by the dot product; the outer sizes, 2 and 2, become the result's shape.

Mechanics

The definition

C_ij = Σ_p=1..k A_ip B_pj,   A: [m, k], B: [k, n] → C: [m, n]

Two other readings of the same sum are worth keeping. Column j of C is A times column j of B, a weighted sum of A's columns. Row i of C is row i of A times B, a weighted sum of B's rows. The second is how to read a linear layer: each output is a mixture of the weight rows, weighted by the input. Matmul is associative, (AB)C = A(BC), but not commutative: AB and BA usually differ, and often only one of them has valid shapes.

Shape rules and torch.matmul

torch.mm takes exactly two 2-D matrices (mm). torch.matmul, which @ calls, dispatches on rank (matmul):

The last two dimensions are the matrices; everything before them is a batch and broadcasts. When shapes don't fit, the error names the two inner sizes, so "mat1 and mat2 shapes cannot be multiplied (32x768 and 512x10)" means a 768 met a 512.

Batched matmul and einsum

torch.bmm multiplies two stacks of matrices of shape [b, m, k] and [b, k, n], with no broadcasting (bmm). addmm and baddbmm fuse "β·C + α·(A @ B)" into one kernel (addmm, baddbmm). torch.einsum names every axis with a letter and sums over any letter missing from the output, so "ik,kj->ij" is matmul, "bik,bkj->bij" is bmm, and "bhqd,bhkd->bhqk" is every attention head's QKᵀ at once (einsum).

A linear layer is a matmul plus a bias

y = x Wᵀ + b,   x: [batch, in], W: [out, in], b: [out] → y: [batch, out]

nn.Linear stores its weight as [out_features, in_features], one row per output neuron, and its forward is a single call to F.linear (Linear.forward, F.linear). So each output is a dot product of the input with one weight row: the neuron from the previous page. The transpose is why the weight is [out, in] while the math reads "in to out". The input can have any number of leading dimensions, [batch, seq, in] included; only the last must equal in_features.

Counting FLOPs

Each of the m·n output cells needs k multiplies and k adds, so

FLOPs(A @ B) = 2 · m · k · n   (times b for a batch of b)

A linear layer applied to T tokens costs 2 · T · in · out. That is where the rule of thumb "a forward pass costs about 2 FLOPs per parameter per token" comes from, and backward costs about twice the forward, giving the 6·N·D estimate for training a model of N parameters on D tokens. Moving the matrices costs only (mk + kn + mn) numbers, so for large matrices compute grows much faster than memory traffic, which is why GPUs run matmuls near peak and most other ops far below it.

Identity, transpose and inverse

• The identity I has ones on the diagonal and zeros elsewhere; AI = A (torch.eye).
• The transpose Aᵀ swaps rows and columns, and (AB)ᵀ = BᵀAᵀ. In PyTorch .T/.t()/.mT return a view with swapped strides, no copy (torch.t).
• The inverse A⁻¹ of a square matrix undoes it: AA⁻¹ = I. For M = [[2, 1], [1, 1]], M⁻¹ = [[1, -1], [-1, 2]]. Only matrices with a non-zero determinant have one (linalg.inv; torch.inverse is an alias). In practice you rarely invert: to solve Ax = b, torch.linalg.solve is faster and more accurate.

Worked example

Implementation

import torch
import torch.nn as nn

A = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
B = torch.tensor([[7., 8.], [9., 10.], [11., 12.]])
A @ B                                   # tensor([[ 58.,  64.], [139., 154.]])
torch.einsum("ik,kj->ij", A, B)         # the same product, axes named

layer = nn.Linear(3, 2)                 # weight [2, 3], bias [2]
with torch.no_grad():
    layer.weight.copy_(torch.tensor([[1., 0., -1.], [.5, .5, .5]]))
    layer.bias.copy_(torch.tensor([0., 1.]))
layer(A)                                # tensor([[-2.0, 4.0], [-2.0, 8.5]])
A @ layer.weight.T + layer.bias         # identical: what F.linear computes

x = torch.randn(32, 10, 3)              # [batch, seq, in]
layer(x).shape                          # torch.Size([32, 10, 2])
torch.matmul(torch.randn(5, 1, 2, 3), torch.randn(7, 3, 4)).shape  # [5, 7, 2, 4]

From scratch, the triple loop every kernel optimizes:

def matmul(a, b):
    m, k, n = len(a), len(b), len(b[0])
    assert len(a[0]) == k, "inner sizes must match"
    c = [[0] * n for _ in range(m)]
    for i in range(m):
        for j in range(n):
            for p in range(k):
                c[i][j] += a[i][p] * b[p][j]
    return c

Tradeoffs

• Order of a chain matters. (AB)C and A(BC) give the same answer at very different costs. With A [1000, 10], B [10, 1000], C [1000, 1]: (AB)C costs 22,000,000 FLOPs, A(BC) only 40,000. LoRA's low-rank update relies on this.
• matmul versus einsum. einsum is clearer for anything with more than two axes and usually lowers to the same kernels, but an unusual equation can fall back to a slower path; profile hot loops.
• Precision. On recent GPUs, float32 matmuls may run in TF32, which keeps float32 range but only about 10 mantissa bits. That is controlled by torch.backends.cuda.matmul.allow_tf32 and torch.set_float32_matmul_precision. bfloat16 and float16 matmuls are faster still and accumulate in float32.
• Shapes that suit the hardware. Tensor cores work on tiles; sizes that are multiples of 8 (or 64, 128) run noticeably faster, which is why vocabularies get padded.
• Don't invert. Explicit inverses are slow and amplify rounding error for nearly singular matrices; use linalg.solve or a factorization.

Check yourself

Connections

References

• Strang. Introduction to Linear Algebra, 6th ed., chapter 2 (matrices and their multiplication). Wellesley-Cambridge Press, 2023.
• Kaplan et al. Scaling Laws for Neural Language Models. 2020, §2.1 on counting compute as 2N FLOPs per token forward. arXiv:2001.08361
• Hu et al. LoRA: Low-Rank Adaptation of Large Language Models. 2021. arXiv:2106.09685
• PyTorch v2.15.0: torch.matmul and nn.Linear.