(Approximate) Matrix Multiplication via Convolutions

TL;DR

This paper introduces FFT-based algorithms for exact and approximate matrix multiplication, breaking previous speed-accuracy tradeoffs and extending group-theoretic frameworks to improve efficiency, with potential applications in large language model inference.

Contribution

It presents a novel FFT-based exact matrix multiplication algorithm and extends the group-theoretic framework for approximate matrix multiplication, including bounds for abelian and non-abelian groups.

Findings

Exact algorithm runs in $O(n^{2.89})$ time using FFTs.

Breaks the linear speed-accuracy tradeoff in approximate matrix multiplication.

Proposes practical approaches for large language model inference.

Abstract

We study the capability of the Fast Fourier Transform (FFT) to accelerate exact and approximate matrix multiplication without using Strassen-like divide-and-conquer. We present a simple exact algorithm running in $O(n^{2.89})$ time, which only sums a few convolutions (FFTs) in $\mathbb{Z}_{m}^{k}$, building on the work of Cohn, Kleinberg, Szegedy and Umans (2005). As a corollary, combining this algorithm with linear sketching breaks the longstanding linear speed-accuracy tradeoff for "combinatorial" approximate matrix multiplication (AMM, Pagh'13, Sarlos'06, Clarkson-Woodruff'13), achieving error $\frac{1}{r^{1.1}}\left\lVert \mathbf{A} \right\rVert_{F}^{2}\left\lVert \mathbf{B}\right\rVert_{F}^{2}$ in $O(rn^{2})$ time, using nothing but FFTs. Motivated by the rich literature for approximating polynomials, our main contribution in this paper is extending the group-theoretic framework of Cohn and Umans (2003) to approximate matrix multiplication (AMM). Specifically, we introduce and study an approximate notion of the Triple Product Property, which in the abelian case is equivalent to finding a Sumset which minimizes (multi-)intersections with an arithmetic progression. We prove tight bounds on this quantity for abelian groups (yielding a simple and practical AMM algorithm via polynomial multiplication), and establish a weaker lower bound for non-abelian groups, extending a lemma of Gowers. Finally, we propose a concrete approach that uses low-degree approximation of multi-variate polynomials for AMM, which we believe will lead to practical, non-asymptotic AMM algorithms in real-world applications, most notably LLM inference.