利用强化学习发现更快的矩阵乘法算法

Discovering faster matrix multiplication algorithms with reinforcement learning

October 2022

Nature 2022 [Cover]

Alhussein Fawzi*, Matej Balog*, Aja Huang*, Thomas Hubert*, Bernardino Romera-Paredes*, Mohammadamin Barekatain, Alexander Novikov, Francisco J. R. Ruiz, Julian Schrittwieser, Grzegorz Swirszcz, David Silver, Demis Hassabis & Pushmeet Kohli (* Equal contribution)

[DeepMind]

https://github.com/deepmind/alphatensor

https://www.nature.com/articles/s41586-022-05172-4

https://github.com/nebuly-ai/nebullvm/tree/main/apps/accelerate/open_alpha_tensor

http://alhusseinfawzi.info/

https://dpmd.ai/dm-alpha-tensor

提高基础计算算法的效率会产生广泛影响，因为它会影响大量计算的总体速度。矩阵乘法就是这样一项原始任务，在从神经网络到科学计算程序的许多系统中都会发生。使用机器学习自动发现算法提供了超越人类直觉并超越当前最佳人类设计算法的前景。然而，自动化算法发现过程是复杂的，因为可能的算法空间巨大。在这里，我们报告了一种基于AlphaZero1的深度强化学习方法，用于发现任意矩阵乘法的有效且可证明正确的算法。我们的代理AlphaTensor被训练成玩单人游戏，目标是在有限因子空间内找到张量分解。AlphaTensor发现的算法在许多矩阵大小上都优于最先进的复杂度。特别相关的是4 × 据我们所知，自50年前发现以来，AlphaTensor算法首次改进了Strassen的两级算法2。我们通过不同的用例进一步展示了AlphaTensor的灵活性：具有最先进复杂度的结构化矩阵乘法算法，并通过优化特定硬件上运行时的矩阵乘法提高了实际效率。我们的结果突出了AlphaTensor在一系列问题上加速算法发现过程，并针对不同标准进行优化的能力。

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.