Neural Learning of Fast Matrix Multiplication Algorithms: A StrassenNet Approach

TL;DR

This paper introduces StrassenNet, a neural architecture that learns fast matrix multiplication algorithms, successfully recovering known algorithms for 2x2 and providing insights into the minimal rank for 3x3 multiplication.

Contribution

The paper presents a neural network approach to discover and analyze low-rank tensor decompositions for fast matrix multiplication, including recovering Strassen's algorithm and exploring minimal ranks.

Findings

Successfully reproduces Strassen's algorithm for 2x2 multiplication.

Identifies a numerical threshold at rank 23 for 3x3 multiplication.

Preliminary results on border-rank decompositions align with known bounds.

Abstract

Fast matrix multiplication can be described as searching for low-rank decompositions of the matrix--multiplication tensor. We design a neural architecture, \textsc{StrassenNet}, which reproduces the Strassen algorithm for $2\times 2$ multiplication. Across many independent runs the network always converges to a rank-$7$ tensor, thus numerically recovering Strassen's optimal algorithm. We then train the same architecture on $3\times 3$ multiplication with rank $r\in\{19,\dots,23\}$. Our experiments reveal a clear numerical threshold: models with $r=23$ attain significantly lower validation error than those with $r\le 22$, suggesting that $r=23$ could actually be the smallest effective rank of the matrix multiplication tensor $3\times 3$. We also sketch an extension of the method to border-rank decompositions via an $\varepsilon$--parametrisation and report preliminary results consistent with the known bounds for the border rank of the $3\times 3$ matrix--multiplication tensor.