Discovering faster matrix multiplication algorithms with reinforcement learning

ABSTRACT

Matrix multiplication is a fundamental computational task, heavily utilized in neural networks, scientific computing, and graphics. Despite its ubiquity, finding optimal algorithms for matrix multiplication remains an open problem. We frame the discovery of matrix multiplication algorithms as a single-player game, TensorGame, and train an agent, AlphaTensor, using deep reinforcement learning. AlphaTensor discovered algorithms that outperform the state-of-the-art complexity bounds established over the last 50 years, including Strassen's famous algorithm for 4x4 matrices. This highlights the potential of RL to discover novel, provably correct algorithms for fundamental mathematical operations.