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

1 INTRODUCTION . Structured linear maps are fundamental and ubiquitous in modern machine learning . Their efficiency in speed ( fast algorithms ) and space ( few parameters ) can reduce computation and memory usage . The class of structured linear maps includes fixed specialized transforms such as the discrete Fourier transform ( DFT ) and Hadamard transform used in signal processing ( Cooley et al. , 1969 ) , convolutions for image , language , and speech modeling ( Gu et al. , 2018 ) , and low-rank and sparse matrices for efficient storage and inference on edge devices ( Yu et al. , 2017 ) . Forms of structure such as sparsity have been at the forefront of recent advances in ML ( Frankle & Carbin , 2019 ) , and are critical for on-device and energy-efficient models , two application areas of tremendous recent interest ( Tsidulko , 2019 ; Schwartz et al. , 2019 ) . There are a plethora of classes of structured linear maps , each with a significantly different representation , algorithm , and implementation . They have different tradeoffs in terms of inference speed , training speed , and accuracy , and the conventional wisdom is that no one class works uniformly well across all applications . As a result , ML practitioners currently hand-pick specific classes of structured linear maps for each of their applications . This is a difficult and labor-intensive task . ∗These authors contributed equally . Ideally , these problems should be addressed with a universal representation for structured linear maps : ( i ) Such a parameterization should be expressive enough to capture important classes of structure , with a nearly tight parameter count and runtime : the space required to represent the linear map should be close to optimal , and the resulting algorithm for matrix vector multiplication should be close to the fastest possible algorithm . ( ii ) The parameterization should be differentiable in order to be learned as a component of end-to-end ML pipelines , enabling it to easily be used as a drop-in replacement for manually engineered structured components . ( iii ) The parameterization should admit practically efficient algorithms for training and inference , in terms of both speed and memory . Currently , no class of structured linear maps satisfies all of these criteria . Most existing classes of structured matrices—such as the class of low-rank matrices—fail to tightly capture other important types of structure . For example , the DFT has an efficient structured representation of size O ( n log n ) , yet can not be well-approximated by a low-rank transform of size n2 . Another important type of structure is sparsity ; lots of exciting recent work has focused on the design of sparse neural networks . For instance , sparse networks of comparable quality to their dense counterparts—yet an order of magnitude fewer parameters—may be created via pruning ( Han et al. , 2016 ) or by identifying “ winning lottery tickets ” ( Frankle & Carbin , 2019 ) . In parallel , recent theoretical results by De Sa et al . ( 2018 ) show that sparsity and the notion of structure in linear maps are fundamentally linked : any given matrix can be factored into a product of sparse matrices with total parameter count equal to the efficiency ( i.e . minimum arithmetic circuit complexity ) of the matrix . In other words , the representation of linear maps as products of sparse matrices tightly captures all forms of structure . Unfortunately , it is difficult to actually learn these sparse factorizations , because it requires finding the sparsity patterns of the factors—a discrete , nondifferentiable search problem . Thus , current methods for training sparse neural networks are either expensive ( Frankle & Carbin , 2019 ) or rely on highly hand-tuned heuristics for evolving the sparsity patterns throughout training ( Dettmers & Zettlemoyer , 2019 ) . By contrast , we propose a representation of linear maps as products of sparse matrices with specific predefined sparsity patterns ( Section 2 ) , and show that it does satisfy our desiderata : it retains the expressiveness of unstructured sparsity , while being differentiably learnable and efficient like other structured representations . Concretely , our representation is based on products of a particular building block known as a butterfly matrix ( Parker , 1995 ; Dao et al. , 2019 ) ; we term such products kaleidoscope matrices ( K-matrices for short ) .1 ( i ) Our main theoretical contribution ( Section 2.3 ) concerns the expressiveness of this representation : we show that any structured linear map ( i.e . one that can be applied using s n2 arithmetic operations ) can be represented as a K-matrix , with a nearly tight number of parameters and algorithmic complexity ( both on the order of s up to logarithmic factors ) . ( ii ) The kaleidoscope representation is fully differentiable ; thus , all the parameters of a K-matrix can be learned using standard optimization algorithms such as SGD . ( iii ) Because of their simple , regular structure , K-matrices are practical and easy to use . We provide memory- and runtime-efficient implementations of K-matrix multiplication on CPU and GPU for training and inference , with a simple PyTorch interface . We empirically validate that , due to their expressiveness , learnability , and efficiency , we can use K-matrices as a drop-in replacement for linear components in deep learning models . In Section 3.1 , we use K-matrices to replace hand-crafted structure in two different settings . We simplify the six steps of filter bank computation in speech preprocessing into a single learnable K-matrix step , with only an 0.4 % accuracy drop on the TIMIT speech recognition task . We use K-matrices to replace channel shuffles in ShuffleNet , improving ImageNet classification accuracy by up to 5 % . In Section 3.2 , we show that K-matrices can successfully recover latent structure ; a K-matrix is used to learn latent permutations in a permuted image dataset ( Permuted CIFAR ) , resulting in 9 points higher accuracy in a downstream CNN model . In Section 3.3 , we show that our efficient K-matrix multiplication implementation can be applied to speed up real-world tasks : we replace linear layers with K-matrices in a DynamicConv-Transformer network to attain 36 % faster end-to-end inference speed with a 1.0 drop in BLEU score on the IWSLT14 German→English translation task . 1A group of butterflies is known as a kaleidoscope . 2 A NEARLY-TIGHT PARAMETERIZATION OF ALL STRUCTURED MATRICES . We first present some background on the characterization of all structured matrices ( i.e . those with subquadratic multiplication algorithms ) as products of sparse factors , along with the definition of butterfly matrices . We then propose a differentiable family of kaleidoscope matrices , composed of products of butterfly matrices , and prove their expressivity : all structured matrices can be represented in this form , with almost optimal parameter count and runtime . 2.1 BACKGROUND : SPARSE FACTORIZATION , BUTTERFLY MATRICES . Sparse factorization One method of constructing matrices with theoretically fast matrix-vector multiplication algorithms is as a product of sparse matrices , so that multiplication by an arbitrary vector has cost proportional to the total number of nonzeros ( NNZ ) of the matrices in the product . Surprisingly , the converse is also true . De Sa et al . ( 2018 ) introduce the concept of sparse product width ( SPW ) , which roughly corresponds to the total NNZ in a factorization of a matrix , and show that it is an asymptotically optimal descriptor of the algorithmic complexity of matrix-vector multiplication ( Bürgisser et al. , 2013 ) . We use a similar argument in the proof of our main theorem ( Section 2.3 ) . However , attempting to learn such a factorization of a given matrix is difficult , as the sparsity constraint is not continuous . Moreover , because of the possibly irregular sparsity patterns , it is difficult to realize the theoretical speedups in practice ( Gray et al. , 2017 ; Gahvari et al. , 2007 ) . Butterfly matrices Butterfly matrices , encoding the recursive divide-and-conquer structure of the fast Fourier transform ( FFT ) algorithm , have long been used in numerical linear algebra ( Parker , 1995 ; Li et al. , 2015 ) and machine learning ( Mathieu & LeCun , 2014 ; Jing et al. , 2017 ; Munkhoeva et al. , 2018 ; Dao et al. , 2019 ; Choromanski et al. , 2019 ) . Here we define butterfly matrices , which we use as a building block for our hierarchy of kaleidoscope matrices . Definition 2.1 . A butterfly factor of size k ≥ 2 ( denoted as Bk ) is a matrix of the form Bk = [ D1 D2 D3 D4 ] where each Di is a k2 × k 2 diagonal matrix . We restrict k to be a power of 2 . Definition 2.2 . A butterfly factor matrix of size n with block size k ( denoted as B ( n ) k ) is a block diagonal matrix of nk ( possibly different ) butterfly factors of size k : B ( n ) k = diag ( [ Bk ] 1 , [ Bk ] 2 , . . . , [ Bk ] nk ) Definition 2.3 . A butterfly matrix of size n ( denoted as B ( n ) ) is a matrix that can be expressed as a product of butterfly factor matrices : B ( n ) = B ( n ) n B ( n ) n 2 . . .B ( n ) 2 . Equivalently , we may define B ( n ) recursively as a matrix that can be expressed in the following form : B ( n ) = B ( n ) n [ [ B ( n 2 ) ] 1 0 0 [ B ( n 2 ) ] 2 ] ( Note that [ B ( n 2 ) ] 1 and [ B ( n 2 ) ] 2 may be different . ) 2.2 THE KALEIDOSCOPE HIERARCHY . Using the building block of butterfly matrices , we formally define the kaleidoscope ( BB∗ ) hierarchy and prove its expressiveness . This class of matrices serves as a fully differentiable alternative to products of sparse matrices ( Section 2.1 ) , with similar expressivity . In Appendix J , we show where various common structured matrix classes are located within this hierarchy . The building block for this hierarchy is the product of a butterfly matrix and the ( conjugate ) transpose of another butterfly matrix ( which is simply a product of butterfly factors taken in the opposite order ) . Figure 1 visualizes the sparsity patterns of the butterfly factors in BB∗ , where the red and blue dots represent the allowed locations of nonzero entries . Definition 2.4 ( Kaleidoscope hierarchy , kaleidoscope matrices ) . • Define B as the set of all matrices that can be expressed in the form B ( n ) ( for some n ) . • Define BB∗ as the set of matrices M of the form M = M1M∗2 for some M1 , M2 ∈ B . • Define ( BB∗ ) w as the set of matrices M that can be expressed as M = Mw . . .M2M1 , with each Mi ∈ BB∗ ( 1 ≤ i ≤ w ) . ( The notation w represents width . ) • Define ( BB∗ ) we as the set of n × n matrices M that can be expressed as M = SEST for some en × en matrix E ∈ ( BB∗ ) w , where S ∈ Fn×en = [ In 0 . . . 0 ] ( i.e . M is the upper-left corner of E ) . ( The notation e represents expansion relative to n. ) • M is a kaleidoscope matrix , abbreviated as K-matrix , if M ∈ ( BB∗ ) we for some w and e. The kaleidoscope hierarchy , or ( BB∗ ) hierarchy , refers to the families of matrices ( BB∗ ) 1e ⊆ ( BB∗ ) 2e ⊆ . . . , for a fixed expansion factor e. Each butterfly matrix can represent the identity matrix , so ( BB∗ ) we ⊆ ( BB∗ ) w+1e . We show that the inclusion is proper in Appendix E. This hierarchy generalizes the BP hierarchy proposed by Dao et al . ( 2019 ) , as shown in Appendix J . Efficiency in space and speed Each matrix in ( BB∗ ) we is a product of 2w total butterfly matrices and transposes of butterfly matrices , each of which is in turn a product of log ( ne ) factors with 2ne nonzeros ( NNZ ) each . Therefore , each matrix in ( BB∗ ) we has 4wne log ( ne ) parameters and a matrix-vector multiplication algorithm of complexity O ( wne log ne ) ( by multiplying the vector with each sparse factor sequentially ) . We prove this more formally in Appendix E. For the applications in Section 3 , w and e are small constants ( up to 2 ) , so those K-matrices have O ( n log n ) parameters and runtime .

Summary: The authors introduce kaleidoscope matrices (K-matrices) and propose to use them as a substitute for structured matrices arising in ML applications (e.g. circulant matrix used for the convolution operation). The authors prove that K-matrices are expressive enough to capture any structured matrix with near-optimal space and matvec time complexity. The authors demonstrate that learnable K-matrices achieve similar metrics compared to hand-crafted features on speech processing and computer vision tasks, can learn from permuted images, achieve performance close to a CNN trained on unpermuted images and demonstrate the improvement of inference speed of a transformer-based architecture for a machine translation task.