We develop a new framework, Polyform, for fast approximate matrix multiplication (AMM), through sums of sparse polynomial multiplications, using (variants of) FFT. Using this framework, we obtain new data-dependent speed-accuracy tradeoffs, which often improve on the worst-case accuracy of randomized sketching algorithms. Meanwhile, Polyform can be viewed as a cheap alternative bilinear operator to matrix multiplication in Deep Neural Networks, which is our motivating application.
Our algorithm involves unexpected connections to Additive Combinatorics and Spherical Harmonics. The former one relates the error of Polyform to the purely existential problem of finding Sumsets which minimize intersections with arithmetic progressions. In the latter, we show that optimizing the polynomial’s coefficients leads to a low-rank SDP problem generalizing Spherical Codes, and provide a novel projection-free approximation algorithm for this problem, in close to input-sparsity time ~O(nnz(A,B)). Meanwhile, our experiments demonstrate that, when using SGD to optimize the coefficients of Polyform in stat-of-art Transformer models (ViT), our algorithm provides a significant reduction in FLOPs (x2-x5) compared to vanilla MM, with little to no accuracy loss.