Matrix Multiplication Without Matrices

Linear algebra without matrices sounds bizarre[1]. What if we told you it’s true? What if we told you, it (matmuls without mats) is in fact SOTA in our search for practical sub-cubic matrix multiplication algorithms? AGI’s alignment is correlated with compute cost and finding new matrix multiplication algorithms democratizes the cost of intelligence. AGI will go well only if you and your mother can afford a matmul. Let me explain.

The Tensor Decomposition Dilemma

First, we give context: GPUs multiply matrices using the vanilla, cubic-complexity matrix multiplication algorithm. GPU (and CPU) vendors avoid Strassen’s sub-cubic algorithm because it doesn’t cache well, and at FP8, floating-point errors accumulate significantly[2]. Other algorithms like Coppersmith-Winograd have ‘galactic complexity’. That is, gains are only observed when multiplying extremely large matrices ~2,000,000*2,000,000 dimension. For context, a GPU tile holds 256*256 matrices, and GPT-2’s largest matrix is an embedding table of dimension 50,257 * 768 [3]. All this algorithms have one thing in common: they work by finding low-rank tensor decompositions. Tensor decompositions are not useful. If they were, then Google would have made TPUs a consumer product and we’d all be using Jax, not Pytorch, to train our models.

A Group-Theoretic Approach To Matrix Multiplication

Microsoft and Caltech researchers [4] found a discrete-fourier-transform-esque matrix multiplication algorithm in 2003 but it’s a theoretical computer science artifact. It works in three steps:

1. Find a group algebra that supports a thing called the ‘Triple Product Property’.

2. Embed your matrices into the group algebra and perform a Discrete Fourier Transform convolution (like in signal processing), not a matmul.

3. Read off your matrix product from the convolution’s result.

Here’s a fully annotated Python version[5] for step-by-step details.

This algorithm stands out because:

1. No tensor decompositions are involved. Thus, we completely avoid the ‘galactic complexity’ shortcomings of other matrix multiplication algorithms like Coppersmith-Winograd.

2. Every Discrete Fourier Transform admits a Fast Fourier Transform!

We need to search for this algorithm’s FFTs.

The Challenge and Our Approach

We’ve established our northstar: finding a Fast Fourier Transform (FFT )for the Triple Product Property matmul algorithm’s Discrete Fourier Transform (DFT). I must warn you, the last time someone found an FFT for a DFT, they were awarded a presidential medal of science[6]. It’s not easy.

We propose modifying Google’s AlphaTensor to search for DFTs, not tensor decompositions. We believe we’ll find something great, much similar to how Google found new (but pretty useless) high-dimensional tensor decompositions.

[1]. Meme of Dua Lipa holding a mathematical textbook with the title, “Linear Algebra Without Matrices” with a million views. Twitter.

[2]. Why Compiler Engineers Rarely Use Strassen’s Algorithm for Fast Matrix Multiplications. Substack. https://leetarxiv.substack.com/p/why-compilers-rarely-use-strassens-algorithm

[3] GPT-2 Model Card. HuggingFace. https://huggingface.co/openai-community/gpt2?show_file_info=model.safetensors

[4]. Cohn, H., & Umans, C. (2003). A Group Theoretic Approach to Fast Matrix Multiplication. In Proceedings of the 44th Annual IEEE Symposium on Foundations of Computer Science (FOCS 2003) (pp. 438–449). IEEE Computer Society. ArXiv.

[5] The Annotated Triple Product Property Matrix Multiplication Algorithm. GitHub. https://github.com/MurageKibicho/The-Annotated-Triple-Product-Property-Matrix-Multiplication-Algorithm