Kaleidoscope: An Efficient, Learnable Representation For All Structured Linear Maps

By Tri Dao, Nimit Sharad Sohoni, Albert Gu, Matthew Eichhorn, Amit Blonder, Megan Leszczynski, Atri Rudra, Christopher R'e. 2020

56 citations recorded by the source corpus.

Modern neural network architectures use structured linear transformations, such as low-rank matrices, sparse matrices, permutations, and the Fourier transform, to improve inference speed and reduce memory usage compared to general linear maps. However, choosing which of the myriad structured transformations to use (and its associated parameterization) is a laborious task that requires trading off speed, space, and accuracy. We consider a different approach: we introduce a family of matrices called kaleidoscope matrices (K-matrices) that provably capture any structured matrix with near-optimal space (parameter) and time (arithmetic operation) complexity. We empirically validate that K-matrices can be automatically learned within end-to-end pipelines to replace hand-crafted procedures, in order to improve model quality. For example, replacing channel shuffles in ShuffleNet improves classification accuracy on ImageNet by up to 5%. Learnable K-matrices can also simplify hand-engineered pipelines---we replace filter bank feature computation in speech data preprocessing with a kaleidoscope layer, resulting in only 0.4% loss in accuracy on the TIMIT speech recognition task. K-matrices can also capture latent structure in models: for a challenging permuted image classification task, adding a K-matrix to a standard convolutional architecture can enable learning the latent permutation and improve accuracy by over 8 points. We provide a practically efficient implementation of our approach, and use K-matrices in a Transformer network to attain 36% faster end-to-end inference speed on a language translation task.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures

How researchers see this paper

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

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

Citations in the corpus

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, methodology): “[15] Tri Dao, Nimit Sohoni, Albert Gu, Matthew Eichhorn, Amit Blonder, Megan Leszczynski, Atri Rudra, and Christopher Ré.”
Initialization and Regularization of Factorized Neural Layers (2021, background): “The success of SI and FD for both low-rank and tensor decomposition also suggests that these schemes may be useful for other types of factorized neural layers, such as ACDC (Moczulski et al., 2015) or K-matrices (Dao et al., 2020).”
HiPPO: Recurrent Memory with Optimal Polynomial Projections (2020, methodology): “[21] Tri Dao, Nimit Sohoni, Albert Gu, Matthew Eichhorn, Amit Blonder, Megan Leszczynski, Atri Rudra, and Christopher Ré.”
Monarch: Expressive Structured Matrices for Efficient and Accurate Training (2022, methodology): “(Note that this analysis is not tight: a more careful analysis can show the inclusion is strict even for smaller values of n.) We end this section with a theorem on the expressivity of the “monarch hierarchy” (products of monarch matrices), which follows from Theorem 1 of (Dao et al., 2020).”
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].”
Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models (2021, methodology): “Vanilla Butterfly is 5× slower than Pixelfly as expected, because (1) it does not take advantage of the hardware property – not structured sparsity(2) it is a series of products.”