Matrix multiplication is a fundamental computation with wide-ranging applications across various fields, including neural networks and scientific computing. The efficiency of matrix multiplication algorithms significantly impacts the overall speed of many computational tasks. While human ingenuity has led to advancements, such as Strassen's algorithm, the search for optimal algorithms remains a challenging problem. The sheer size of the space of possible algorithms makes exhaustive search intractable. This research explores the potential of automating algorithm discovery using machine learning. Specifically, the authors propose a deep reinforcement learning (DRL) approach to discover provably correct and efficient matrix multiplication algorithms. The underlying principle is that matrix multiplication algorithms can be represented as low-rank decompositions of a three-dimensional tensor. Finding such decompositions is an NP-hard problem, making automated discovery a significant challenge. Prior attempts at algorithm discovery relied heavily on human-designed heuristics, possibly leading to suboptimal solutions. This study aims to overcome these limitations using DRL to learn patterns in tensors and predict efficient decompositions, offering a potential breakthrough in the quest for faster matrix multiplication.

Existing literature on matrix multiplication algorithms extensively explores the theoretical complexities and asymptotic bounds. Significant work focuses on characterizing the complexity of the asymptotically optimal algorithm, but this rarely yields practical algorithms. Previous approaches to discovering practical algorithms have employed human-guided search, continuous optimization, and combinatorial search methods, frequently relying on heuristics which may not be optimal. These methods have limitations in scalability and the ability to explore the vast space of possible algorithms. This paper builds upon the success of AlphaZero, a general reinforcement learning algorithm that has demonstrated superhuman performance in games like chess and Go. The authors leverage the AlphaZero framework to tackle the problem of matrix multiplication algorithm discovery, adapting it to handle the significantly larger action space inherent in this mathematical problem.

The authors formulate the matrix multiplication algorithm discovery problem as a single-player game called TensorGame. In this game, the agent (AlphaTensor) aims to decompose a three-dimensional tensor representing the matrix multiplication operation into a sum of rank-one tensors. Each step in the game corresponds to selecting a rank-one tensor to subtract from the current tensor. The objective is to reach the zero tensor in the fewest possible steps. A reward function is used to guide the search, penalizing longer sequences of steps and incentivizing shorter, more efficient decompositions. The reward also incorporates a penalty for failing to reach the zero tensor within a predefined number of steps. To efficiently search the enormous action space, AlphaTensor utilizes a Monte Carlo Tree Search (MCTS) guided by a deep neural network. The neural network architecture is a transformer-based model designed to handle the structure of tensor inputs and incorporates inductive biases, such as invariance to reordering of slices. Training uses a mixture of supervised learning (using synthetically generated tensor decompositions) and reinforcement learning from actual game play. Several techniques were employed to enhance performance: generating synthetic demonstrations of tensor factorizations to aid training; exploiting symmetries of the problem; introducing a change of basis to diversify the games played; and employing data augmentation techniques by creating new training examples from finished games. The agent is trained to discover algorithms over different arithmetics (modular arithmetic in ℤ₂ and standard arithmetic in ℝ).

AlphaTensor successfully rediscovered several well-known matrix multiplication algorithms, demonstrating the validity of the approach. More significantly, it discovered new algorithms that outperform the state-of-the-art for various matrix sizes. A notable achievement is the discovery of an algorithm for multiplying 4x4 matrices using 47 multiplications in ℤ₂, improving upon Strassen's two-level algorithm (49 multiplications). This translates to an asymptotic complexity of O(N^2.778) for multiplying N x N matrices in ℤ₂. In standard arithmetic, AlphaTensor improved the state-of-the-art for several matrix sizes (e.g., finding a rank-76 decomposition for T_4,5,5, improving on the previous best of 80 multiplications). The agent generated a large number of algorithms for each matrix size, allowing the exploration of various combinations and recursive applications to tackle larger matrix multiplications. This led to state-of-the-art results for over 70 matrix multiplication tensors (n, m, p ≤ 12). The agent demonstrates knowledge transfer between different tensor sizes, improving performance overall. AlphaTensor scales to larger problem sizes than existing methods; it successfully tackled tensors (like T₇) which were previously computationally prohibitive. Furthermore, AlphaTensor was applied beyond standard matrix multiplication to discover efficient algorithms for structured matrix multiplication, like skew-symmetric matrix-vector multiplication, and was shown to outperform prior art algorithms. AlphaTensor's adaptability was further proven by tailoring matrix multiplication algorithms to specific hardware (GPU and TPU) by modifying the reward function to incorporate runtime performance. These hardware-specific algorithms outperformed Strassen's square algorithm in practice, highlighting the potential for creating optimized algorithms for different computing environments. The results underscore the richness and diversity of the space of matrix multiplication algorithms, exceeding previous understanding.

AlphaTensor's success in discovering more efficient matrix multiplication algorithms than those previously known significantly advances the field of algorithmic discovery. The ability to improve upon Strassen's algorithm, a landmark result in the field, after 50 years demonstrates the power of the reinforcement learning approach. The flexibility to adapt the reward function and apply the method to different arithmetics and structured matrix multiplications expands the potential applications of AlphaTensor. The discovery of numerous algorithms for each matrix size highlights a previously unexplored richness in the space of possible solutions, potentially leading to further mathematical insights. Optimizing algorithms for specific hardware demonstrates the practical applicability of AlphaTensor beyond theoretical improvements. These findings open new avenues for research into algorithmic discovery, suggesting that DRL could be a powerful tool for solving complex mathematical problems and potentially assisting mathematicians in making new discoveries. While the pre-defined set of factor entries (F) limits the search space, future research could address this limitation.

AlphaTensor demonstrates the potential of reinforcement learning to automate the discovery of efficient algorithms for fundamental mathematical problems. It successfully discovered novel matrix multiplication algorithms that outperform existing state-of-the-art methods for numerous matrix sizes and various contexts (different arithmetics, hardware architectures, structured matrices). The flexible framework allows for optimization across multiple criteria, from minimizing the number of scalar multiplications to improving runtime performance on specific hardware. Future work could focus on relaxing the discretization of the search space and extending the approach to other complex mathematical problems.

A key limitation is the need to pre-define a finite set of potential factor entries (F) for the tensor decomposition. This discretization of the search space might prevent the discovery of certain optimal algorithms that involve coefficients outside this set. While AlphaTensor demonstrated strong transfer learning between different tensor sizes, the approach still requires significant computational resources for training, limiting its immediate scalability to extremely large matrix sizes. Finally, although hardware-specific algorithms were discovered, optimizing for a single hardware configuration may not translate directly to optimal performance across different architectures. Further research is needed to enhance the generality and efficiency of the proposed methodology.