Matrix Multiplication

Matrix multiplication combines two grids of numbers into a third in which every entry is a dot product of a row from the first matrix with a column from the second. The dimensions rule is simple: a matrix with m rows and n columns can be multiplied by a matrix with n rows and p columns, because the inner sizes must agree, and the result has m rows and p columns. Order matters, since multiplying A by B generally differs from multiplying B by A and the reverse product may not even exist, but grouping does not, so (A times B) times C equals A times (B times C). Multiplying a matrix by a vector is the special case that applies a linear transformation, and the whole operation is the workhorse of numerical computing.

Matrix multiplication is the computation that deep learning spends most of its time doing. A dense neural network layer multiplies a batch of input vectors by a weight matrix and adds a bias; a batch of 32 inputs with 768 features times a 768-by-3072 weight matrix yields a 32-by-3072 output and costs 32 × 768 × 3072, about 75.5 million multiply-add operations, for a single layer. Convolutions are commonly lowered to matrix multiplications as well, and in transformers the large feed-forward blocks and the projections that produce queries, keys, values and outputs are all matrix products. The standard library interface is called GEMM, general matrix multiply, a routine tuned for decades in BLAS libraries such as cuBLAS, MKL and OpenBLAS.

Attention is two matrix multiplications with a softmax between them. In the transformer of Vaswani and colleagues (2017), the queries, keys and values for all tokens are stacked into matrices, the product of the query matrix and the transposed key matrix gives every token-to-token score at once, and after scaling and softmax a second product with the value matrix produces the outputs. Writing attention this way, instead of as a loop over pairs of tokens, is what allows it to run efficiently on hardware, and the size of the n-by-n score matrix is also why attention memory and compute grow quadratically with sequence length, the cost that techniques such as FlashAttention reorganize but do not remove.

GPUs and tensor cores exist largely to make this one operation fast. Matrix multiplication offers enormous parallelism, since each output entry can be computed independently, and heavy data reuse, since each input value takes part in many products, which lets tiled kernels keep operands in fast on-chip memory. NVIDIA introduced tensor cores with the Volta architecture and the V100 in 2017: dedicated units that multiply small tiles of matrices and accumulate the result in a single operation, in reduced precision such as 16-bit floating point. Mixed-precision training (Micikevicius and colleagues, 2017) pairs these low-precision products with higher-precision accumulation, and Google's TPUs use systolic arrays built for the same purpose. Whether a given multiplication reaches peak speed depends on its arithmetic intensity: small or skinny products are limited by memory bandwidth, not arithmetic.

Faster-than-schoolbook algorithms exist, but practice rarely uses them. The textbook method needs n cubed multiplications for two n-by-n matrices. Volker Strassen showed in 1969, in "Gaussian Elimination is not Optimal," that 2-by-2 blocks can be multiplied using seven multiplications instead of eight, and applying this recursively gives a running time proportional to n to the power log2 of 7, about n^2.81. Later work by Coppersmith and Winograd in 1987 and many successors lowered the exponent theoretically to about 2.37, but the constants make those algorithms impractical. In 2022, DeepMind's AlphaTensor found new multiplication schemes for small matrix sizes by treating the search as a game. Highly tuned cubic-time kernels still dominate real systems, because they are regular, numerically stable and hardware-friendly.

Common pitfalls involve confusing operations and ignoring shape. The element-wise (Hadamard) product multiplies matching entries and is a different operation from the matrix product; in NumPy and PyTorch the star operator is element-wise and the at-sign operator is the matrix product. Shape mismatches are the most frequent bug, and silently broadcasting a vector can produce wrong results with no error. Floating-point addition is not exactly associative, so changing the order of accumulation, which parallel hardware does freely, can change low-order digits of results.

Read matrix multiplication as a grid of dot products organized so a machine can compute all of them at once: the inner-dimension rule defines when two linear maps compose, attention and every dense layer are built from it, and GPUs, tensor cores and TPUs are shaped around doing it as fast as physics allows.