Discovering faster matrix multiplication algorithms with reinforcement learning

By Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Romera-Paredes, Mohammadamin Barekatain, Alexander Novikov, Francisco J. R. Ruiz, Julian Schrittwieser, Grzegorz Swirszcz and 3 more. Nature, 2022

714 citations recorded by the source corpus.

Improving the efficiency of algorithms for fundamental computations can have a widespread impact, as it can affect the overall speed of a large amount of computations. Matrix multiplication is one such primitive task, occurring in many systems—from neural networks to scientific computing routines. The automatic discovery of algorithms using machine learning offers the prospect of reaching beyond human intuition and outperforming the current best human-designed algorithms. However, automating the algorithm discovery procedure is intricate, as the space of possible algorithms is enormous. Here we report a deep reinforcement learning approach based on AlphaZero1 for discovering efficient and provably correct algorithms for the multiplication of arbitrary matrices. Our agent, AlphaTensor, is trained to play a single-player game where the objective is finding tensor decompositions within a finite factor space. AlphaTensor discovered algorithms that outperform the state-of-the-art complexity for many matrix sizes. Particularly relevant is the case of 4 × 4 matrices in a finite field, where AlphaTensor’s algorithm improves on Strassen’s two-level algorithm for the first time, to our knowledge, since its discovery 50 years ago2. We further showcase the flexibility of AlphaTensor through different use-cases: algorithms with state-of-the-art complexity for structured matrix multiplication and improved practical efficiency by optimizing matrix multiplication for runtime on specific hardware. Our results highlight AlphaTensor’s ability to accelerate the process of algorithmic discovery on a range of problems, and to optimize for different criteria. A reinforcement learning approach based on AlphaZero is used to discover efficient and provably correct algorithms for matrix multiplication, finding faster algorithms for a variety of matrix sizes.

Subfields: Artificial Intelligence, Computational Mathematics, Computational Theory and Mathematics

View Paper: https://www.semanticscholar.org/paper/442ab95eb9cfbc03bb17a27b52313b5d25eaa738

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Lichao Sun
2. Philip S. Yu
3. Eduardo C. Garrido-Merchán
4. Richard Ngo
5. Zhiling Yan
6. Siyu Li
7. Yihan Cao
8. Yutong Dai

In Computational Mathematics

1. Philip S. Yu

In Computational Theory and Mathematics

1. Eduardo C. Garrido-Merchán
2. Philip S. Yu
3. Richard Ngo
4. Lichao Sun

Citations in the corpus

• The alignment problem from a deep learning perspective (2022, methodology): “…addition to standard drivers of progress, AI developers are increasingly using ML systems for accelerating programming [OpenAI, 2023a], and developing new architectures [Elsken et al., 2019], algorithms [Fawzi et al., 2022], training data [Huang et al., 2022a], and chips [Mirhoseini et al., 2021].”
• AlphaEvolve: A coding agent for scientific and algorithmic discovery (2025, neutral): “This problem has been tackled with many approaches, from specialized alternating least squares solvers [93] to deep reinforcement learning [26] and custom search algorithms [47]; yet, despite decades of effort, even for the simple case of multiplying two 3 × 3 matrices, the minimum achievable rank…”
• ChatGPT is not all you need. A State of the Art Review of large Generative AI models (2023, background): “Alphatensor, created by the research company Deepmind, is a completely revolutionary model in the industry for its ability to discover new algorithms [15].”
• Neurosymbolic Program Synthesis (neutral): “Such RL methods have been recently scaled in efforts such as AlphaTensor [33] and AlphaDev [34].”
• A Comprehensive Survey of AI-Generated Content (AIGC): A History of Generative AI from GAN to ChatGPT (2023, methodology): “ChatBot Xiaoice Microsoft 2018 [200] Xiaoice ChatBot Meena Google 2020 [201] Meena Blog ChatBot BlenderBot Meta 2022 [202] Blenderbot ChatBot ChatGPT OpenAI 2022 [10] ChatGPT ChatBot Alexa Amazon 2014 Amazon Alexa ChatBot Lex Amazon 2017 Amazon Lex Music AIVA Aiva Tech 2016 AIVA Music Jukebox OpenAI”