Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations

By Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, Christopher R'e. 2019

131 citations recorded by the source corpus.

Fast linear transforms are ubiquitous in machine learning, including the discrete Fourier transform, discrete cosine transform, and other structured transformations such as convolutions. All of these transforms can be represented by dense matrix-vector multiplication, yet each has a specialized and highly efficient (subquadratic) algorithm. We ask to what extent hand-crafting these algorithms and implementations is necessary, what structural priors they encode, and how much knowledge is required to automatically learn a fast algorithm for a provided structured transform. Motivated by a characterization of fast matrix-vector multiplication as products of sparse matrices, we introduce a parameterization of divide-and-conquer methods that is capable of representing a large class of transforms. This generic formulation can automatically learn an efficient algorithm for many important transforms; for example, it recovers the O(N log N) Cooley-Tukey FFT algorithm to machine precision, for dimensions N up to 1024. Furthermore, our method can be incorporated as a lightweight replacement of generic matrices in machine learning pipelines to learn efficient and compressible transformations. On a standard task of compressing a single hidden-layer network, our method exceeds the classification accuracy of unconstrained matrices on CIFAR-10 by 3.9 points—the first time a structured approach has done so—with 4× faster inference speed and 40× fewer parameters.

Subfields: Computational Mathematics, Computational Mechanics, Hardware and Architecture, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computational Mathematics

In Computational Mechanics

In Hardware and Architecture

In Optimization & Training

Citations in the corpus

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, methodology): “The butterfly pattern we use for our block-sparse attention is motivated by the fact that butterfly matrices [14, 61] and their products have been shown to be able to express any structured matrices with almost optimal runtime and number of parameters [15, 19].”
HiPPO: Recurrent Memory with Optimal Polynomial Projections (2020, methodology): “[20] Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré.”
Scatterbrain: Unifying Sparse and Low-rank Attention Approximation (2021, background): “De Sa et al. [24] shows that any structured matrix can be written as product of sparse matrices, and products of sparse matrices even with ﬁxed sparsity pattern have been shown to be eﬀective at parameterizing compressed models [1, 20, 21].”
Low-Memory Neural Network Training: A Technical Report (2019, background): “…as low-rank matrices (Sainath et al., 2013), Toeplitz-like matrices (Sindhwani et al., 2015), block-circulant matrices (Cheng et al., 2015, Ding et al., 2017), Fastfood transforms (Yang et al., 2015), low displacement rank matrices (Thomas et al., 2018), and butterﬂy matrices (Dao et al., 2019).”
Fast Sparse ConvNets (2019, methodology): “Alternatively, factorizations of the 1×1 convolutions have been considered in ShuffleNet [47] and Learnable Butterfly Factorizations [6].”
Randomized numerical linear algebra: Foundations and algorithms (2020, background): “3) or Davidson and Szarek (2001, Sec. 2.3) or Vershynin (2018, Secs. 7.2–7.3). More recently, it has been observed that Gordon’s inequality can be reversed in certain settings (Thrampoulidis, Oymak and Hassibi 2014). In several instances, we require more detailed information about Gaussian random ma”
Kaleidoscope: An Efficient, Learnable Representation For All Structured Linear Maps (2020, methodology): “This proof is similar to the VC bound for ReLU networks whose weight matrices are butterﬂy matrices (Dao et al., 2019).”
Hyena Hierarchy: Towards Larger Convolutional Language Models (2023, methodology): “…of (8) to an arbitrary order is inspired by fast evaluation algorithms for structured dense matrices based on butterfly decompositions (Li et al., 2015; Dao et al., 2019; 2022c), with length of the decomposition closely tied to its expressivity (in the classes of matrices it can represent).”
Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models (2021, methodology): “Butterfly, Sparse + Low-rank Matrices: Butterfly matrices have been used in numerical linear algebra [Parker, 1995, Li et al., 2015] and machine learning [Mathieu and LeCun, 2014, Jing et al., 2017, Munkhoeva et al., 2018, Dao et al., 2019, Choromanski et al., 2019].”
Monarch: Expressive Structured Matrices for Efficient and Accurate Training (2022, methodology): “Dao et al. (2019) introduced the notion of a butterfly matrix as a certain product of permuted block-diagonal matrices, inspired by the Cooley-Tukey fast Fourier transform algorithm (Cooley & Tukey, 1965).”