Discovering Faster Matrix Multiplication Algorithms with Reinforcement Learning

Every time you ask a voice assistant a question, apply a filter to a photo, or run a neural network, the computer is multiplying matrices. Matrix multiplication is to computing what addition is to arithmetic — a primitive operation that everything else builds upon. Speeding it up by even a small fraction compounds across billions of daily computations worldwide.

The schoolbook algorithm for multiplying two n×n matrices performs n³ multiplications: for two 2×2 matrices, that means 2³ = 8 multiplications. Each scalar multiplication is the expensive step — additions are comparatively cheap. So the fundamental question becomes: can we multiply matrices using fewer multiplications?

In 1969, Volker Strassen proved that 2×2 matrix multiplication can be done with 7 multiplications instead of 8. The trick: instead of computing each entry of the result independently, combine rows and columns of the input matrices in clever sums, multiply those sums, then recombine the products to get the final answer. You trade one multiplication for several extra additions — a good deal because multiplications dominate the cost.

Strassen's insight has a profound consequence: by applying the 2×2 trick recursively to larger matrices (divide each n×n matrix into four n/2 × n/2 blocks), the asymptotic complexity drops from O(n³) to roughly O(n^2.807). For large matrices, this is a massive speedup. But for larger base cases — 3×3, 4×4, 5×5 — nobody knew the optimal number of multiplications. The search space grows so rapidly that brute-force enumeration is impossible.

Here is the elegant mathematical connection that makes AlphaTensor possible. Any algorithm for multiplying an m×k matrix by a k×n matrix can be encoded as a decomposition of a 3D tensor of size (mk) × (kn) × (mn). Think of this tensor as a cube of numbers that encodes which entries of the input matrices need to be combined to produce each entry of the output.

Each "step" of the algorithm — one scalar multiplication — corresponds to one rank-1 term in the decomposition: three vectors u, v, w whose outer product u ⊗ v ⊗ w contributes one layer of the tensor. The total number of rank-1 terms is the rank of the decomposition, and it equals the number of scalar multiplications the algorithm uses.

The standard algorithm for 2×2 matrices has rank 8. Strassen's has rank 7. Finding a faster algorithm means finding a lower-rank decomposition of the same tensor — a classic problem in algebraic complexity theory that has resisted human efforts for decades.

The crucial design decision of this paper is reframing tensor decomposition as a single-player game. The game works as follows:

State: a 3D tensor S_t, initialized to the matrix multiplication tensor 𝒯. Action: at each step t, the player chooses three vectors u, v, w and subtracts their outer product from the current state to get the next state: S minus u ⊗ v ⊗ w. Goal: reduce S_t to the all-zeros tensor. Reward: a negative reward at each step (encouraging fewer steps), plus a large negative penalty if the tensor is not zeroed by a maximum number of steps.

The number of steps to reach zero equals the rank R — the number of multiplications. Playing the game optimally means finding the lowest-rank decomposition, which is the most efficient algorithm.

AlphaTensor builds on AlphaZero but introduces critical innovations to handle the unique challenges of TensorGame. Unlike Go where the board is a 2D grid with 361 possible moves, TensorGame involves a 3D tensor with an action space exceeding 10³³ — thirty orders of magnitude larger than Go.

The neural network is an encoder-decoder Transformer. The encoder takes three inputs: the current tensor state (projected along all three axes to capture its 3D structure), the last several actions taken (for temporal context), and a scalar time embedding. It processes the tensor through attention layers that treat all three cyclic transpositions of the axes equally — a key inductive bias that respects the mathematical symmetry of tensor decomposition.

The decoder outputs two things: a policy — a distribution over candidate actions (which u, v, w to choose next) — and a value — an estimate of how many more steps are needed from the current state. Since enumerating all possible actions is infeasible, AlphaTensor uses the Sampled AlphaZero approach: the network proposes a small set of promising candidate actions, and Monte Carlo Tree Search explores among them.

Training combines two data sources. First, synthetic demonstrations: randomly generated tensor decompositions that provide a warm-start curriculum. These are not optimal solutions — just valid ones that teach the network the basic structure of the problem. Second, self-play games: the agent plays TensorGame using MCTS, and the resulting trajectories are added to a replay buffer. The network is trained to predict both the policy (which move was selected by MCTS) and the value (how many steps remained) from each state.

A crucial technique is change-of-basis augmentation: the same tensor can be expressed in many equivalent coordinate systems. By randomly transforming the tensor before each game, AlphaTensor sees diverse views of the same problem, injecting symmetry awareness without hardcoding it. The agent plays all views in parallel and only needs to solve one — if any basis yields a decomposition, it's mathematically valid in all bases.

AlphaTensor was trained on a range of matrix sizes from 2×2 up to 5×5, and its discoveries were remarkable.

Starting from zero knowledge, the agent first rediscovers the standard algorithms, then Strassen's 7-multiplication solution for 2×2 matrices, and then pushes beyond. For 4×4 matrices over the finite field Z/2Z (modular arithmetic mod 2), AlphaTensor discovered an algorithm using 47 multiplications — beating Strassen's two-level recursion of 49 for the first time since 1969.

For the practically important case of multiplying a 4×5 by a 5×5 matrix, where the standard algorithm uses 100 multiplications and prior human-designed algorithms had reached 80, AlphaTensor found algorithms using only 76. Across more than 70 matrix sizes tested, AlphaTensor improved upon the best known algorithm in many cases and matched it in the rest.

Perhaps even more surprising: AlphaTensor discovered over 14,000 non-equivalent algorithms for a single problem (4×4 matrix multiplication). This diversity revealed that the space of efficient algorithms is far richer than mathematicians had suspected.

Minimizing the number of multiplications is not the only objective that matters. Different algorithms with the same number of multiplications can have very different runtimes on real hardware because of factors like memory access patterns, parallelism, and cache utilization.

AlphaTensor exploits the diversity of its discovered algorithms by adding a second optimization stage. After finding many mathematically equivalent algorithms (same rank), it benchmarks each one on a target hardware platform — such as an NVIDIA V100 GPU or a Google TPU v2 — and feeds the runtime back as an additional reward signal.

The result: hardware-tailored algorithms that run 10–20% faster than commonly used implementations on the same hardware. This demonstrates a key advantage of AI-driven algorithm discovery: the ability to optimize for arbitrary, non-mathematical objectives that matter in practice.

AlphaTensor is part of a broader trajectory at DeepMind: extending the AlphaZero framework from board games to real-world problems. AlphaZero mastered chess, Go, and shogi. MuZero learned to plan without even knowing the rules. AlphaTensor shows the next step — applying the same principles to an open problem in pure mathematics.

The approach is generalizable. Any problem that can be expressed as "take a sequence of actions to reach a verifiable goal state" is a candidate for this framework. Potential applications include discovering faster algorithms for other computational primitives (sorting, FFT, convolution), optimizing circuit layouts, finding new chemical synthesis pathways, and more.

The paper also connects to AlphaFold 2 — another instance of DeepMind applying AI to solve a fundamental scientific problem that had resisted decades of human effort. Together, these works suggest a new paradigm where AI is not just a tool for automating known tasks, but a partner in scientific discovery itself.

AlphaTensor's deepest contribution is not any single algorithm it discovered, but the proof of concept: open mathematical problems, long considered the exclusive domain of human genius, can be attacked by AI agents that learn from scratch. The search space may be astronomically large, but a well-designed game formulation combined with powerful neural networks and tree search can find needles in haystacks that humans could not even begin to sift through.