We discovered faster way to compute product of matrix by its...
Dmitry Rybin@DmitryRybin1
9 views
May 22, 2025
~2 min read
1
We discovered faster way to compute product of matrix by its transpose!
This has profound implications for data analysis, chip design, wireless communication, and LLM training!
paper: arxiv.org/abs/2505.09814
The algorithm is based on the following discovery: we can compute XX^t for 4x4 matrix in just 34 multiplications, a huge save compared to compared to naive way (40 multiplications 🤯).
We can apply this algorithm to any m x n matrix X (with n, m >= 4) by dividing it into 16 blocks X_1, ..., X_16.
- Estimated energy save: 5-10% ✅
- Estimated time save: 5% ✅
The discovery was made by combining Machine Learning-based Search and Combinatorial Optimization. We used RL to sample bilinear expressions. We then used combinatorial solvers (Gurobi) to enumerate relations between these expressions and combine these expressions together into one algorithm for XX^t. One way think of it is modification of AlphaTensor approach - We reduced the action space by a factor of a million (x1000000) at the expense of relying on combinatorial solvers.
The matrix XX^t is used everywhere:
- Data Analysis: linear regression
- Finance: covariance matrix for asset returns
- LLM training: Muon, SOAP, Shampoo
- Wireless Communication: 5G, MIMO channel capacity
This operation is performed trillions of times every minute globally. Imagine if we can save 5% of energy used for these computations!
Coauthors: Yushun Zhang @ericzhang0410, Zhi-Quan Luo.
This has profound implications for data analysis, chip design, wireless communication, and LLM training!
paper: arxiv.org/abs/2505.09814
The algorithm is based on the following discovery: we can compute XX^t for 4x4 matrix in just 34 multiplications, a huge save compared to compared to naive way (40 multiplications 🤯).
We can apply this algorithm to any m x n matrix X (with n, m >= 4) by dividing it into 16 blocks X_1, ..., X_16.
- Estimated energy save: 5-10% ✅
- Estimated time save: 5% ✅
The discovery was made by combining Machine Learning-based Search and Combinatorial Optimization. We used RL to sample bilinear expressions. We then used combinatorial solvers (Gurobi) to enumerate relations between these expressions and combine these expressions together into one algorithm for XX^t. One way think of it is modification of AlphaTensor approach - We reduced the action space by a factor of a million (x1000000) at the expense of relying on combinatorial solvers.
The matrix XX^t is used everywhere:
- Data Analysis: linear regression
- Finance: covariance matrix for asset returns
- LLM training: Muon, SOAP, Shampoo
- Wireless Communication: 5G, MIMO channel capacity
This operation is performed trillions of times every minute globally. Imagine if we can save 5% of energy used for these computations!
Coauthors: Yushun Zhang @ericzhang0410, Zhi-Quan Luo.
2
Addendum:
For general matrix product A*B people use Strassen-Winograd algorithm based on 2x2 block-matrix case. They don't use other asymptotically (n -> inf) faster algorithms because of big overhead for small matrix sizes.
In fact, the smallest known algorithm with exponent better than Strassen occurs for 12x12 matrices! It has big overhead because of extra additions and overall complexity.
Recent work AlphaEvolve by DeepMind showed a rank 48 solution for 4x4 matrix using complex numbers (compared to Strassen rank 49). Because of complex numbers we cannot use it to accelerate real 4x4 matrix multiplication. It is open question if we can get a real solution with rank 48.
RXTX is very unique in that sense:
- strictly better than Strassen-based algorithm for 4x4 matrix product XX^t✅
- has relatively small overhead ✅
- works with real numbers ✅
For general matrix product A*B people use Strassen-Winograd algorithm based on 2x2 block-matrix case. They don't use other asymptotically (n -> inf) faster algorithms because of big overhead for small matrix sizes.
In fact, the smallest known algorithm with exponent better than Strassen occurs for 12x12 matrices! It has big overhead because of extra additions and overall complexity.
Recent work AlphaEvolve by DeepMind showed a rank 48 solution for 4x4 matrix using complex numbers (compared to Strassen rank 49). Because of complex numbers we cannot use it to accelerate real 4x4 matrix multiplication. It is open question if we can get a real solution with rank 48.
RXTX is very unique in that sense:
- strictly better than Strassen-based algorithm for 4x4 matrix product XX^t✅
- has relatively small overhead ✅
- works with real numbers ✅