K ALEIDOSCOPE : A N E FFICIENT , L EARNABLE R EPRESENTATION F OR A LL S TRUCTURED L INEAR M APS
A BSTRACT
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 shufﬂes in ShufﬂeNet improves classiﬁcation accuracy on ImageNet by up to 5%.
K-matrices can also simplify hand-engineered pipelines—we replace ﬁlter bank feature computation in speech data preprocessing with a learnable kaleidoscope layer, resulting in only 0.4% loss in accuracy on the TIMIT speech recognition task.
In addition, K-matrices can capture latent structure in models: for a challenging permuted image classiﬁcation task, a K
-matrix based representation of permutations is able to learn the right latent structure and improves accuracy of a downstream convolutional model by over 9%.
We provide a practically efﬁcient 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.
I NTRODUCTION
Structured linear maps are fundamental and ubiquitous in modern machine learning.
Their efﬁciency in speed (fast algorithms) and space (few parameters) can reduce computation and memory usage.
The class of structured linear maps includes ﬁxed 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 efﬁcient 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-efﬁcient 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 signiﬁcantly 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 speciﬁc classes of structured linear maps for each of their applications.
This is a difﬁcult and labor-intensive task.
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 efﬁcient algorithms for training and inference, in terms of both speed and memory.
Currently, no class of structured linear maps satisﬁes 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 efﬁcient structured representation of size O ( n log n ) , yet cannot be well-approximated by a low-rank transform of size (cid:28) n 2 .
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 efﬁciency (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 difﬁcult to actually learn these sparse factorizations, because it requires ﬁnding 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 speciﬁc predeﬁned sparsity patterns (Section 2), and show that it does satisfy our desiderata: it retains the expressiveness of unstructured sparsity, while being differentiably learnable and efﬁcient like other structured representations.
Concretely, our representation is based on products of a particular building block known as a butterﬂy 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 (cid:28) n 2 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- efﬁcient 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 efﬁciency, 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 ﬁlter 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 shufﬂes in ShufﬂeNet, improving ImageNet classiﬁcation 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 efﬁcient 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.
drop in BLEU score on the IWSLT14 German → English translation task.
A NEARLY - TIGHT PARAMETERIZATION OF ALL STRUCTURED MATRICES
We ﬁrst 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 deﬁnition of butterﬂy matrices.
We then propose a differentiable family of kaleidoscope matrices, composed of products of butterﬂy matrices, and prove their expressivity: all structured matrices can be represented in this form, with almost optimal parameter count and runtime.
2. B ACKGROUND : 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 difﬁcult, as the sparsity constraint is not continuous.
Moreover, because of the possibly irregular sparsity patterns, it is difﬁcult to realize the theoretical speedups in practice (Gray et al., 2017; Gahvari et al., 2007).
Butterﬂy matrices Butterﬂy 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 deﬁne butterﬂy matrices, which we use as a building block for our hierarchy of kaleidoscope matrices.
Deﬁnition 2.1.
A butterﬂy factor of size k ≥ 2 (denoted as B k ) is a matrix of the form B k
= (cid:20)
D
[Equation]
D(cid:21) where each D i is a k 2 × k 2 diagonal matrix.
We restrict k to be a power of 2.
Deﬁnition 2.2.
A butterﬂy factor matrix of size n with block size k (denoted as B ( n ) k ) is a block diagonal matrix of nk (possibly different) butterﬂy factors of size k :
[Equation]
Deﬁnition 2.3.
A butterﬂy matrix of size n (denoted as B ( n ) ) is a matrix that can be expressed as a product of butterﬂy factor matrices: B ( n ) =
B ( n ) n B ( n ) n 2 . . .
B ( n ) 2 .
Equivalently, we may deﬁne B ( n ) recursively as a matrix that can be expressed in the following form:
[Equation]
(Note that [ B ( n 2 ) ] 1 and [ B ( n 2 ) ] 2 may be different.)
2. T HE KALEIDOSCOPE HIERARCHY
Using the building block of butterﬂy matrices, we formally deﬁne 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 butterﬂy matrix and the (conjugate) transpose of another butterﬂy matrix (which is simply a product of butterﬂy factors taken in the opposite order).
Figure 1 visualizes the sparsity patterns of the butterﬂy factors in BB ∗ , where the red and blue dots represent the allowed locations of
nonzero entries.
Deﬁnition 2.4 (Kaleidoscope hierarchy, kaleidoscope matrices) .
• Deﬁne B as the set of all matrices that can be expressed in the form B ( n ) (for some n ).
• Deﬁne BB ∗ as the set of matrices M of the form M = M 1 M ∗ 2 for some M 1 , M 2 ∈ B .
• Deﬁne ( BB ∗ ) w as the set of matrices M that can be expressed as M = M w . . .
M 2 M 1 , with each M i ∈ BB ∗ ( 1 ≤ i ≤ w ).
(The notation w represents width .)
• Deﬁne ( BB ∗ )
we as the set of n × n matrices M that can be expressed as M = SES T for some en × en matrix E ∈ ( BB ∗ ) w , where S ∈ F n × en =
[ I n 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 ∗ ) 1 e ⊆
( BB ∗ ) 2 e ⊆ . . .
, for a ﬁxed expansion factor e .
Each butterﬂy matrix can represent the identity matrix, so ( BB ∗ ) we ⊆
( BB ∗ )
w +1 e .
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.
Efﬁciency in space and speed Each matrix in ( BB ∗ ) we is a product of 2 w total butterﬂy matrices and transposes of butterﬂy matrices, each of which is in turn a product of log( ne ) factors with 2 ne nonzeros (NNZ) each.
Therefore, each matrix in ( BB ∗ ) we has 4 wne 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.
2. A LL LOW - DEPTH STRUCTURED MATRICES ARE IN THE KALEIDOSCOPE HIERARCHY
We now present our main theoretical result: the fact that general linear transformations, expressed as low-depth linear arithmetic circuits, are captured in the BB ∗ hierarchy with low width.
Arithmetic circuits are commonly used to formalize algebraic algorithmic complexity (Bürgisser et al., 2013); we include a primer on this in Appendix M.
The quantities of interest are the total number of gates in the circuit, representing the total number of steps required to perform the algorithm for a serial processor, and the depth , representing the minimum number of steps required for a parallel processor.
Theorem 1.
Let M be an n × n matrix such that multiplication of M times an arbitrary vector v can be represented as a linear arithmetic circuit with s total gates and depth d .
Then, M ∈ ( BB ∗ ) O
( d ) O ( s n ) .
The representation of such a matrix M in the BB ∗ hierarchy has O ( ds log s ) parameters and yields a
O ( ds log s ) multiplication algorithm, compared to the O ( s ) parameters and runtime of the circuit representation.
To the best of our knowledge, the most general classes of efﬁcient matrices that have been studied (De Sa et al., 2018) have depth d on the order of log n or poly log n .
In these cases, the representation with K-matrices matches the best known bounds up to polylogarithmic factors.
The crux of the proof of Theorem 1 (shown in Appendix F) is the construction of an almost tight representation of any sparse matrix as a K-matrix (i.e. a product of butterﬂy matrices): speciﬁcally, we show that any n × n sparse matrix with s nonzeros is in ( BB ∗ )
O ( (cid:100) sn (cid:101) )
O (1) (Theorem 3, Appendix I).
We then leverage the expressivity result of products of sparse matrices to represent all arithmetic circuits (similar to the sparse product width result of De Sa et al.
(2018) referenced in Section 2.1) to complete the proof of Theorem 1.
This intermediate result is also a novel characterization of sparse matrices.
For a matrix with s
NNZ, the kaleidoscope representation has O ( s log n ) parameters and runtime, instead of the optimal
O ( s ) parameters and runtime; so, we trade off an extra logarithmic factor in space and time for full differentiability (thanks to the ﬁxed sparsity patterns in the representation).
The intuition behind the result is as follows: a sparse matrix with s NNZ can be written as a sum of (cid:100) s/n (cid:101) matrices each with at most n NNZ.
Any n × n matrix with at most n NNZ, up to permuting the rows and columns, is a product of two butterﬂy matrices (Lemma I.1).
Sorting networks (Knuth, 1997) imply that permutation matrices are in ( BB ∗ )
O (log n ) , but we tighten the result to show that they are in fact in BB ∗ (Theorem 2, Appendix G).
We thus obtain a kaleidoscope representation for each summand matrix with O ( n log n ) parameters.
By the addition closure property of the BB ∗ hierarchy (Lemma H.5), each sparse matrix with s NNZ then has a kaleidoscope representation with O ( s log n ) parameters.
Tight representation for structured linear maps common in ML Even though Theorem 1 suggests that the kaleidoscope representation can be loose by logarithmic factors, many structured linear maps common in ML can be represented in this hierarchy with an optimal number of parameters and runtime compared to the best known parameterizations, up to constant factors.
Appendix J includes several examples such as discrete transforms (the DFT, discrete cosine transform (DCT), discrete sine transform (DST), and Hadamard transform), convolution (i.e. circulant matrices), Toeplitz matrices (Gray, 2006), structured matrices for kernel approximation ( ( HD ) 3 (Yu et al., 2016)) and compact neural network design (Fastfood (Le et al., 2013), ACDC (Moczulski et al., 2016)).
There have been other large classes of structured matrices proposed in the machine learning literature, such as Toeplitz-like (Sindhwani et al., 2015) and low displacement rank (LDR) (Thomas et al., 2018), but they are not known to be able to capture these common structures as tightly as K-matrices can.
More detailed discussions are in Appendix A.
2. E XTENSIONS
ReLU networks with low-depth structured weight matrices In Appendix L, we prove that ﬁnding an efﬁcient circuit for a ReLU network can be reduced to ﬁnding efﬁcient circuits for each of its weight matrices, with at most a constant factor greater size and run-time (i.e. number of gates).
We also show that ReLU networks with kaleidoscope weight matrices have near-linear VC dimension in the number of parameters, matching the bound for networks with unconstrained weight matrices (Bartlett et al., 1999; Harvey et al., 2017) and LDR (Thomas et al., 2018).
This yields a corresponding sample complexity bound.
Orthogonal kaleidoscope hierarchy Orthogonal butterﬂy matrices are one commonly used variant due to their improved stability (Parker, 1995), where each butterﬂy factor is constrained to be orthogonal:
[Equation]
with C , S being diagonal and C 2 + S
=
I .
Similar to the BB ∗ hierarchy, in
Appendix K, we deﬁne the OBB hierarchy consisting of products of orthogonal butterﬂy matrices and diagonal matrices, and show that this hierarchy has the same expressiveness as the BB ∗ hierarchy.
E MPIRICAL E VALUATION
We validate three claims that suggest that kaleidoscopes are a promising technique to learn different types of structure in modern architectures.
Section 3.1: for applications in speech and lightweight computer vision relying on highly hand-crafted structured transformations, we show that we can recover—and even improve—the quality of such architectures by simply replacing existing hand-structured components with K-matrices, with only a small overhead in memory and computation.
In Section 3.2, for a challenging task with latent structure (Permuted CIFAR-10), a K-matrixbased relaxation of permutations is able to learn the right latent permutation, yielding 9 points better accuracy in a downstream CNN compared to standard RNN and CNN baselines used on such permuted image classiﬁcation tasks.
In Section 3.3, we show that, although not yet highly optimized, our current implementation of K-matrices can improve the inference throughput of DynamicConv Transformer, a state-ofthe-
art fast machine translation model, by 36%, with only a relatively small drop in translation quality.
In all of the above applications, as K-matrices are fully differentiable, we simply train them jointly with the rest of the model using standard learning algorithms (such as SGD).
Full details for all of the experiments (precise architectures, hyperparameters, etc.) are in Appendix B.
3. R EPLACING HAND - CRAFTED STRUCTURES
We validate that kaleidoscope matrices can recover or improve on the performance of hand-crafted structure in ML models.
For example, a single learnable kaleidoscope layer can be used to replace the hand-engineered ﬁlter bank speech preprocessing pipeline with only 0.4% loss in accuracy on the TIMIT speech recognition task (Section 3.1.1).
Replacing channel shufﬂes in ShufﬂeNet with learnable K-matrices improves classiﬁcation accuracy on ImageNet by up to 5.0% (Section 3.1.2).
3.1. S PEECH PREPROCESSING
We show that K-matrices can remove the need for hand-tuning by signiﬁcantly simplifying speech recognition data preprocessing pipelines.
In particular, we can entirely replace the complex handcrafted MFSC featurization commonly used in speech recognition tasks with a fully learnable kaleidoscope layer, with only 0.4% drop in accuracy on the TIMIT speech recognition benchmark.
Results are presented in Table 1.
Our approach is competitive with the accuracy of standard models that use hand-crafted features, and signiﬁcantly outperforms current approaches for learning from raw audio input.
-the-art for learning from raw audio), and is only 0.8% less accurate than the overall state-of-the-art on TIMIT.
2 Additional comparisons are given in Appendix B.1.
[Table]
Modern speech recognition models currently rely on carefully hand-crafted features extracted from the audio, which are then fed into an acoustic model .
By contrast, learning directly from the raw audio—i.e. end-to-end learning from the audio waveform without any manual featurization—obviates the need for this complicated and often expensive preprocessing step.
There have been recent attempts to learn directly from raw audio, such as SincNet (Ravanelli & Bengio, 2018); however, they often rely on specialized architectures designed by domain experts.
Instead, we use a standard RNN speech recognition architecture, but use a learnable kaleidoscope layer to replace the featurization steps.
The baseline architecture takes as input ﬁlter bank (MFSC) features, which are a popular standard featurization for speech recognition (Paliwal, 1999) and involve several steps hand-crafted speciﬁcally for this domain.
These features are extracted from the raw audio waveform, and fed as the input into a Bi-LSTM model.
We signiﬁcantly simplify this pipeline by replacing the featurization step with a trainable kaleidoscope layer that is trained end-to-end together with the Bi-LSTM.
The original pipeline and our modiﬁed kaleidoscope version are depicted in Figure 2.
The computation of MFSC features involves a series of painstakingly hand-designed steps (further described in Appendix B.1), each involving their own hyperparameters: (i) the waveform is framed (split into chunks), (ii) the waveform is dithered (noise is added), (iii) pre-emphasis is applied, (iv) the Hamming window is applied, (v) the FFT is applied and the power spectrum is computed, (vi) the result is mapped to the mel scale (which involves applying a particular linear transformation and then taking the logarithm of the result), (vii) cepstral mean and variance normalization is applied.
We replace the last six steps (ii-vii) of this featurization process with a learnable kaleidoscope layer; speciﬁcally, after windowing, we multiply the input by a K-matrix, and then compute the logarithm of the power spectrum; the output is fed into the Bi-LSTM model.
3.1. R EPLACING CNN CHANNEL SHUFFLE
We evaluate how K-matrices can improve the quality of hand-crafted, lightweight architectures for computer vision tasks, without the need for hand-tuning.
We select ShufﬂeNet (Zhang et al., 2018), which is a state-of-the-art lightweight CNN architecture that uses a manually designed “channel shufﬂe” permutation matrix to improve performance.
By replacing this ﬁxed permutation with a learnable K-matrix, we achieve up to 5% further improvement in classiﬁcation accuracy, without hand-tuned components and with a modest space penalty of up to 10%.
Results are given in Table 2.
[Table]
Grouped convolution (Krizhevsky et al., 2012) is often used to reduce parameter count and speed up inference compared to standard convolution, but, by default, channels in different groups cannot exchange information.
To remedy this, ShufﬂeNet uses a permutation matrix to shufﬂe the channels after each grouped convolution.
Zhao et al.
(2019) propose to instead use the Hadamard transform before and after each grouped convolution to mix the channels.
In place of these hand-engineered solutions, we use a K-matrix before and after each grouped convolution, and learn these end-to-end together with the rest of the network.
As shown in Table 2, across a range of sizes, replacing the channel shufﬂes with K-matrices results in improved performance at comparable parameter counts.
3. L EARNING A LATENT PERMUTATION
We show that K-matrices can be used in a challenging task for which existing classes of structured linear maps have not been found suitable.
We investigate the problem of image classiﬁcation on a permuted image dataset (Permuted CIFAR-10).
This problem is challenging due to the discrete nature of learning the latent permutation of the dataset; we present a differentiable relaxation for this using a K-matrix as a key component.
Results are presented in Table 3; compared to methods that do not have a permutation learning step, our approach gets 9 points higher accuracy (84.4% to 93.6%), coming within 2 points of the accuracy on the un-permuted dataset (94.9%).
[Table]
In this task, we use a permuted image classiﬁcation dataset (Permuted CIFAR-10), wherein a ﬁxed global permutation is applied to the pixels of every image in the original input set.
Typically, only fully-connected (FC) and recurrent models are applied to such datasets (Le et al., 2015), because the permutation destroys locality in the image, presenting a difﬁculty for CNNs.
However, CNNs are much better-suited for standard image tasks.
We thus expect that learning the permutation and then applying a standard CNN should outperform these baselines.
As mentioned in Section 2, the kaleidoscope hierarchy provides a nearly tight parameterization of permutations; this makes them a natural ﬁt for the permutation learning step.
Experimentally, we use a K-matrix to represent a distribution over permutations, which converges to a single permutation at the end of training.
The correct latent structure is learned by applying samples from this distribution to the permuted training images, and minimizing an auxiliary smoothness-based loss that encourages the reconstructed images to be more “natural” (i.e. vary smoothly pixel-to-pixel).
The learned permutation is evaluated by training a ResNet18 with the K-matrix permutation layer inserted at the beginning.
Full details of our approach are provided in Appendix B.3.
In Table 3, we compare our approach to a ResNet18 without this extra K-matrix layer, a ResNet with an extra dense matrix at the beginning instead of a K-matrix, and other baselines.
As generic representations such as unstructured matrices do not have the requisite properties to ﬁt in the pipeline, these baselines fail to effectively learn the latent permutation.
We emphasize that a K-matrix provides this ability to recover latent structure despite not being specialized for permutations.
Figuredescribes the pipeline and displays examples of permuted and unpermuted images.
3. S PEEDING UP I NFERENCE
We evaluate the inference speed beneﬁt of using K-matrices on a real language translation model.
We choose the state-of-the-art DynamicConv Transformer translation model (Wu et al., 2019), which offers 20% inference speedup over the standard Transformer model and replace dense matrices in the decoder’s linear layers with K-matrices, which leads to a further 36% inference speedup (Table 4).
As outlined in Section 2.
K-matrices admit a simple and fast O ( n log n ) matrix-vector multiplication algorithm.
We provide fast implementations of this algorithm in C++ and CUDA, with an interface to
PyTorch (Paszke et al., 2017) and use this implementation in our experiments.
[Table]
We use K-matrices to replace all the linear layers in the decoder of DynamicConv (since 90% of inference time is spent in the decoder).
As shown in Table 4, on the IWSLT-14 German-English translation task, this yields a 25% smaller model with 36% faster inference time on CPU, at the cost of 1.0 drop in BLEU score.
4 (Our model also nearly matches the state-of-the-art BLEU performance of 2 years ago obtained by the Transformer model (Vaswani et al., 2017), despite being over 60% faster for inference than the Transformer.)
The majority (55%) of inference time is spent in matrix-vector multiplication; our implementation of K-matrix-vector multiplication is about 2 times faster than the optimized implementation of dense matrix-vector multiplication in the Intel MKL library.
Direct comparisons of K-matrix multiplication with this and other highly-optimized routines such as the
FFT are further detailed in Appendix C.
C ONCLUSION
We address the problem of having to manually choose among the numerous classes of structured linear maps by proposing the universal (expressive, efﬁcient, and learnable) family of kaleidoscope matrices.
We prove that K-matrices can represent any structured linear maps with near-optimal space and time complexity.
Empirical validations suggest that K-matrices are a promising and ﬂexible way to employ structure in modern ML; they can be used to reduce the need for hand-engineering, capture challenging latent structure, and improve efﬁciency in models.
We are excited about future work on further hardware-optimized implementations of K-matrices, to fully realize the size and speed beneﬁts of structured matrices on a broad array of real-world applications.
A CKNOWLEDGMENTS
We thank Avner May and Jian Zhang for their helpful feedback.
We gratefully acknowledge the support of DARPA under Nos.
FA87501720095 (D3M),
FA86501827865 (SDH), and FA86501827882 (ASED); NIH under No. U54EB020405 (Mobilize), NSF under Nos. {+CCF
+} (Beyond Sparsity), CCF1563078 (Volume to Velocity), [-and 1937301 (RTML); ONR under No. N
-]and(RTML);
{+ONR under No. N
+}
(Unifying Weak Supervision); the Moore Foundation,
NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, the Okawa Foundation, American Family Insurance,
Google Cloud, Swiss Re, and members of the Stanford DAWN project: Teradata, Facebook, Google, Ant Financial, NEC, VMWare, and Infosys.
The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon.
Any opinions, ﬁndings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reﬂect the views, policies, or endorsements, either expressed or implied, of DARPA, NIH, ONR, or the U.S. Government.
Matthew Eichhorn and Atri Rudra’s research is supported by NSF grant CCF-1763481.
A R ELATED W ORK A. S TRUCTURED MATRICES IN MACHINE LEARNING
Structured linear maps such as the DFT, the Hadamard transform and convolution are a workhorse of machine learning, with diverse applications including data preprocessing, random projection, featurization, and model compression.
For example, the DFT is a crucial step in the standard ﬁlter bank speech preprocessing pipeline (Jurafsky & Martin, 2014), and is commonly used when dealing with time series data in general (Panaretos & Tavakoli, 2013).
Fast random projection and kernel approximation methods rely on the fast Hadamard transform (Le et al., 2013; Yu et al., 2016) and convolution (Yu et al., 2015), and convolution is a critical component of modern image processing architectures (Krizhevsky et al., 2012) as well as being useful in speech recognition (Zeghidour et al., 2018) and natural language processing (Wu et al., 2019).
Large learnable classes of structured matrices such as Toeplitz-like matrices (Sindhwani et al., 2015) and low-displacement rank (LDR) matrices (Thomas et al., 2018) have been used for model compression.
However, despite their theoretical speedup, these structured matrix classes lack efﬁcient implementations, especially on
GPUs.
Therefore, their use has largely been conﬁned to small models (e.g. single hidden layer neural nets) and small datasets (e.g. CIFAR-10).
Butterﬂy matrices encode the recursive divide-and-conquer structure of the fast Fourier transform (FFT) algorithm.
They were ﬁrst used in numerical linear algebra for fast preconditioning (Parker, 1995).
The butterﬂy factorization is then generalized to encompass complementary low-rank matrices commonly encountered in solving differential and integral equations (Rokhlin & Tygert, 2006; Tygert, 2008; 2010b;a; Li et al., 2015; 2018).
In machine learning, butterﬂy matrices have been use to approximate the Hessian for fast optimization (Mathieu & LeCun, 2014), and to perform fast random projection (Jing et al., 2017; Munkhoeva et al., 2018; Choromanski et al., 2019).
Dao et al. (2019) show that butterﬂy matrices can be used to learn fast algorithms for discrete transforms such as the Fourier transform, cosine/sine transform, Hadamard transform, and convolution. A. S PARSE MATRICES
Several classes of structured linear transforms are ubiquitous in modern deep learning architectures; particularly widespread examples include convolution and multiheaded attention.
Recently, attempts to impose sparsity on the neural network weights have been gaining traction.
State-of-the art approaches of this type typically accomplish this by pruning small weights (either gradually during training (Zhu & Gupta, 2017), or post-training (Han et al., 2016)) or by training a dense network and then identifying “winning lottery tickets”—sparse
subnetworks which may then be retrained from scratch with appropriate initialization (Frankle & Carbin, 2019).
Importantly, these approaches start from a dense network, and therefore training is expensive.
There is also a more nascent line of work that aims to train unstructured sparse neural networks directly (Mocanu et al., 2018; Mostafa & Wang, 2019; Dettmers & Zettlemoyer, 2019;
Evci et al., 2019).
These approaches maintain a constant network sparsity level throughout training, and use heuristics to evolve the sparsity pattern during training.
One drawback is that the indices of the nonzero entries need to be stored in addition to the entry values themselves, which increases the memory required to store the sparse weight tensors.
Another drawback is that these approaches to learn the sparsity pattern are based on intricate heuristics, which can be brittle.
We note that these heuristic sparsiﬁcation techniques could potentially be combined with our approach, to further sparsify the K-matrix factors.
A. S PEECH RECOGNITION FROM RAW AUDIO
Numerous works focus on the problem of speech recognition from raw audio input, i.e. without manual featurization.
SincNet (Ravanelli & Bengio, 2018) is a CNN-based architecture parameterized with sinc functions, designed so that the ﬁrst convolutional layer imitates a band-pass ﬁlter.
Zeghidour et al.
(2018) formulate a learnable version of a ﬁlter bank featurization; their ﬁlters are initialized as an approximation of MFSC features and then ﬁne-tuned jointly with the rest of the model.
Sainath et al.
(2015) proposed a powerful combined convolutional LSTM (CLDNN)-based model for learning from raw audio, using a large amount of training data.
The WaveNet generative architecture (van den
Oord et al., 2016), based on dilated convolutions, has been adapted to speech recognition and can be trained on raw audio.
Other approaches that can learn from raw audio can be found in (Palaz et al., 2013; Collobert et al., 2016; Ghahremani et al., 2016).
To our knowledge, the 14.6% PER achieved by our kaleidoscope + LSTM model on the TIMIT test set is the lowest error rate obtained by a model trained directly on the raw audio.
A. L EARNING PERMUTATIONS
Permutation matrices ﬁnd use in tasks such as matching and sorting (among many others).
Techniques to obtain posterior distributions over permutations have been developed, such as the exponential weights algorithm
(Helmbold & Warmuth, 2009) and the Gumbel-Sinkhorn network (Mena et al., 2018).
Classifying images with permuted pixels is a standard task to benchmark the ability of RNNs to learn long range dependencies.
Le et al.
(2015) propose the Permuted MNIST task, in which the model has to classify digit images with all the pixels permuted.
Many new RNN architectures, with unitary or orthogonal weight matrices to avoid gradient explosion or vanishing, have been proposed and tested on this task (Le et al., 2015; Arjovsky et al., 2016;
Wisdom et al., 2016; Mhammedi et al., 2017;
Trinh et al., 2018).
Standard gated RNN architectures such as LSTM and GRU have also been found to be competitive with these new RNN architectures on this task (Bai et al., 2018).
B A DDITIONAL E XPERIMENTAL D ETAILS B. S PEECH PREPROCESSING
In this section, we fully describe our settings and procedures for the speech preprocessing experiments in Section 3.1.1, and present additional auxiliary baselines and results.
B.1. E XPERIMENTAL SETUP
We evaluate our speech recognition models on the TIMIT speech corpus (Garofolo et al., 1993), a standard benchmark for speech recognition.
The input is audio (16-bit, 16 kHz .wav format), and the target is the transcription into a sequence of phonemes (units of spoken sound).
Our evaluation metric is the phoneme error rate (PER) between the true phoneme sequence and the phoneme sequence predicted by our model.
We use PyTorch (Paszke et al., 2017), the Kaldi speech recognition toolkit (Povey et al., 2011), and the PyTorch-Kaldi toolkit (Ravanelli et al., 2019) for developing PyTorch speech recognition models for all our experiments and evaluations.
B.1. M ODEL AND EVALUATION
Our baseline Bi-LSTM architecture is taken from the PyTorch-Kaldi repository.
5 This is a strong baseline model that, to the best of our knowledge, matches state-of-the-art performance for models that use a single type of input featurization (Ravanelli et al., 2019).
The original Bi-LSTM model takes as input ﬁlter bank features.
These are computed as follows: (i) the waveform is framed (split into chunks of 25 ms each that overlap by 10 ms each), (ii) the waveform is dithered (zero-mean
Gaussian random noise is added), (iii) pre-emphasis is applied to amplify high frequencies, (iv)
the
Hamming window function (Harris, 1978) is applied, (v) the FFT is applied, and the power spectrum of the resulting (complex-valued) output is computed, (vi) the power spectrum (which has dimension 512) is mapped to the “mel scale” (which is a scale intended to mimic human auditory perception (Stevens et al., 1937)) by multiplication with a speciﬁc banded matrix of dimension 512 × 23 , and the entrywise logarithm of the output is taken (the 23 outputs are called the ﬁlters ), and (vii) cepstral mean and variance normalization (Liu et al., 1993) is applied.
Numerical hyperparameters of this procedure include the dither noise scale, the pre-emphasis coefﬁcient, the Hamming window size, the number of mel ﬁlters, and more; we kept all these the same as the Kaldi/PyTorch-Kaldi defaults.
In contrast, our “K-matrix version” of the model takes as input the raw waveform, split into chunks the same way as before but with no normalization, dithering, or other preprocessing, which is then fed into a complex-valued kaleidoscope [ ( BB ∗ ) 2 ] matrix.
Similarly to the nonlinear steps in computing ﬁlter bank features, the logarithm of the power spectrum of the output (which has dimension 512) is then computed.
This output is fed into the Bi-LSTM; the Bi-LSTM and kaleidoscope layer are trained together in standard end-to-end fashion.
The Bi-LSTM architecture is not modiﬁed aside from changing the input dimension from 23 to 512; this (along with the ≈ 75 K parameters in the kaleidoscope layer itself) results in approximately a 1.1M increase in the total number of parameters compared to the model that takes in MFSC features (a modest 8% relative increase).
Total training time for our kaleidoscope-based architecture is 7% greater than that required for the model that uses
MFSC features, not counting the time required to precompute the MFSC features; the FLOPs for inference-time are approximately 15% greater (mostly due to the larger dimension of the input to the
Bi-LSTM; the kaleidoscope layer accounts for less than 0.5% of the total FLOPs).
As baselines, we also compare to inserting other types of linear transformations before the Bi-LSTM: ﬁxed linear transformations (such as the ﬁxed FFT, or no transform at all [i.e. the identity]), other trainable structured layers (low-rank, circulant, and sparse [using the sparse training algorithm of Dettmers & Zettlemoyer (2019)]), and a trainable unstructured (dense) linear layer.
The kaleidoscope layer performs the best out of all such approaches.
The fact that it outperforms even a dense linear layer with more parameters is particularly notable, as it suggests that the structural bias imposed by the
K-matrix representation is beneﬁcial for performance on this task.
Full results are given in Table 5.
[Table]
In our experiments, we grid search the initial learning rate for the “preprocessing layer” (if applicable) in {5e-5, 1e-4, 2e-4, 4e-4, 8e-4, 1.6e-3}, and ﬁx all other hyperparameters (including the initial learning rates for the other parts of the network) to their default values in the PyTorch-Kaldi repository.
The model and any preprocessing layers are trained end-to-end with the RMSProp optimizer forepochs (as per the defaults in PyTorch-Kaldi).
For each model, we use the validation set to select the best preprocessing learning rate, while the ﬁnal error rates are reported on the separate held-out test set.
For all structured matrix baselines except circulant (which always has n parameters for an n × n matrix), the number of parameters in the structured matrices is set to equal the number of parameters in the butterﬂy layer, while the unconstrained matrix is simply a standard dense complexvalued square matrix.
For all experiments with a trainable “preprocessing layer,” we initialize the preprocessing matrix to represent the FFT (or approximate it as closely as possible [i.e. minimize the
Frobenius error to the true FFT matrix], in the case of low-rank, sparse, and circulant), which we found to outperform random initialization.
B.1. E XTENSION : C OMBINING MFSC AND KALEIDOSCOPE
As an additional experiment, we sought to investigate whether combining the hand-engineered MFSC featurization pipeline and a learnable kaleidoscope layer (instead of replacing the former with the latter) could lead to accuracy gains.
Speciﬁcally, in this experiment we ﬁrst used the standard ﬁlter bank featurization pipeline described above, and trained end-to-end as usual.
Then, we replaced the
FFT step with a K-matrix initialized to the FFT, and made the weights of the Hamming window function and the mel ﬁlter bank matrix learnable as well (similarly to (Zeghidour et al., 2018)).
We ﬁne-tuned the resulting architecture for an additional 10 epochs.
The ﬁnal test PER% attained by this “hybrid” model is 13.9 ± 0 .
2 ; the model has 14.4M parameters—a negligible increase over the 14.3M in the original architecture.
Thus, by combining the manually encoded domain knowledge in the ﬁlter bank featurization and allowing this structure to be learnable rather than ﬁxed, we are able to nearly match the state-of-the-art 13.8% accuracy on TIMIT.
While this “hybrid” model certainly involves some hand-engineering, the state-of-the-art results use a concatenation of three different speech audio featurizations—MFSC, MFCC, and fMLLR—as the neural network input, along with a customized RNN architecture (LiGRU) speciﬁcally designed for speech recognition, and thus require a more complicated pipeline that is arguably even more hand-crafted.
B. R EPLACING CNN CHANNEL SHUFFLE B.2. M ODEL ARCHITECTURES
ShufﬂeNet is a convolutional neural network with residual (skip) connections that uses a permutation matrix to shufﬂe the channels after each grouped 1x1 convolution, sending the i -th channel to the ( i mod g ) -th group, where g is the total number of groups.
The architecture for each residual block in ShufﬂeNet is: 1x1 group conv → Batch norm, ReLU → Permutation → 3x3 depthwise conv →
Batch norm → 1x1 group conv.
The permutation is ﬁxed.
Zhao et al.
(2019) propose to instead use the Hadamard transform before and after each grouped 1x convolution to mix the channels.
Note that the Hadamard transforms are placed before the batch normalization and ReLU layer (unlike the permutation matrix in the original ShufﬂeNet design).
In particular, the architecture for each block is: Hadamard → 1x1 group conv → Hadamard → Batch norm, ReLU → 3x3 depthwise conv → Batch norm → 1x1 group conv.
The Hadamard transform is ﬁxed.
In our architecture, we use a kaleidoscope matrix in OBB (product of an orthogonal butterﬂy matrix, a diagonal matrix, and the transpose of another butterﬂy matrix) before and after each grouped 1x convolution.
We place the second K-matrix after the batch norm and ReLU, to more closely mimic the original ShufﬂeNet design.
The structure for each block is: K-matrix → 1x1 group conv → Batch norm, ReLU → K-matrix → 3x3 depthwise conv → Batch norm → 1x1 group conv.
The K-matrices are trained along with the rest of the network, rather than being ﬁxed.
B.2. E XPERIMENTAL SETUP
We evaluate the CNN architectures on the image classiﬁcation task of the standard ImageNet dataset (Russakovsky et al., 2015).
We use the standard data augmentation, training, and evaluation pipeline as in (Xie et al., 2017).
We train with SGD on 8 GPUs for 90 epochs, with a total batch size of 2048 and initial learning rate 0.8.
For the 1.0 ShufﬂeNet g8 architecture, we reduce the total batch size to 1792 to ﬁt into GPU memory, and correspondingly linearly scale the initial learning rate to 0.7.
Other hyperparameters (e.g. learning rate schedule, weight decay, etc.) are kept the same as in the ShufﬂeNet paper (Zhang et al., 2018).
We use the training script from NVIDIA’s deep learning examples repository.
B.2. A DDITIONAL RESULTS
In Table 6, we report top-5 classiﬁcation accuracy on ImageNet, to complement the top-1 accuracies in Table 2.
We report results for different network width multipliers (number of channels), and for different kinds of matrices used for channel mixing.
Using a learnable K-matrix in place of each ﬁxed permutation (shufﬂe) or Hadamard matrix improves top-5 accuracy by up to 4.8%.
Parameter counts are the same as in Table 2.
[Table]
In each setting, the total training time of our K-matrix approach is within 20% of the total training time of vanilla ShufﬂeNet.
In Figure 4, we plot the loss and accuracy on the training set and validation set when we train 1.
ShufﬂeNet g8, with either a ﬁxed permutation (Shufﬂe) or a K-matrix for channel shufﬂing.
Even though each K-matrix is a product of multiple (sparse) matrices, the model with K-matrices takes about the same number of training steps to converge as the baseline model does.
One possible reason is that we constrain the K-matrices to be orthogonal (Section 2.4), thus avoiding vanishing or exploding gradients.
[Figure]
B. L EARNING PERMUTATIONS B.3. D ATASET
The permuted CIFAR-10 dataset is constructed by applying a ﬁxed permutation to every input.
We choose to use the 2-D bit-reversal permutation, 7 i.e., the bit reversal permutation on 32 elements is applied to the rows and to the columns.
This permutation was chosen because it is locality-destroying: if two indices i, j are close, they must differ in a lower-order bit, so that the bit-reversed indices i (cid:48) , j (cid:48) are far.
This makes it a particularly challenging test case for architectures that rely on spatial locality such as “vanilla” CNNs.
B.3. M ODEL AND T RAINING
We describe the model architectures used in Section 3.1 (those reported in Table 3).
Our model (K + CNN)
The model represents a ﬁxed permutation P , parametrized as a K-matrix, to learn to recover the true permutation, followed by a standard ResNet18 architecture (He et al., 2016).
Because of the simple decomposable nature of the butterﬂy factors (Section 2.1), our parameterization is easily extensible with additional techniques:
(i)
We constrain each butterﬂy factor matrix in the K-matrix to be doubly-stochastic.
For example, each 2 × 2 block in the butterﬂy factor matrix of block size 2 has the form (cid:20) a 1 − a 1 − a a (cid:21) , where a ∈ [0 , 1] .
We treat this block as a distribution over permutations, generating the identity
[Equation]
with probability a and the swap
[Equation]
with probability 1 − a .
Butterﬂy factor matrices with larger block sizes are constrained to be doubly-stochastic in a similar manner.
In this way, a permutation is sampled for each butterﬂy factor matrix, and these permutations are composed to get the ﬁnal permutation that is applied to the image.
(ii) For each minibatch, the examples Px by applying permutation samples on the (permuted) inputs are fed into an additional unsupervised reconstruction loss
[Equation]
measuring total variation smoothness of the de-noised inputs.
Such loss functions are often used in image denoising (Rudin et al., 1992).
A ﬁnal regularization loss was placed on the entropy of P , which was annealed over time to encourage P to converge toward a sharper doubly-stochastic matrix (in other words, a permutation).
The model is trained with just the reconstruction loss to convergence before the standard
ResNet is trained on top.
These techniques are applicable to the K-matrix as well as specialized methods for representing permutations such as Gumbel-Sinkhorn (Mena et al., 2018) and are important for recovering the true permutation.
However, they are not applicable to a general linear layer, which showcases the ﬂexibility of K-matrices for representing generic structure despite not being specially tailored for this task.
We also remark that other classes of structured linear maps such as low-rank, circulant, and so on, are even less suited to this task than dense matrices, as they are incapable of representing all permutations.
Baseline architectures
Fully connected (FC):
This is a 3-layer MLP, with hidden size 1024 and ReLU nonlinearity in-between the fully connected layers.
Recurrent neural network (RNN): We use a gated recurrent unit (GRU) model (Cho et al., 2014), with hidden size 1024.
Many RNN architectures have been proposed to capture long-range dependency on permuted image dataset such as Permuted MNIST (Arjovsky et al., 2016).
Standard gated architectures such as LSTM and GRU have shown competitive performance on the Permuted MNIST dataset, and we choose GRU as a baseline since it has been reported to slightly outperform LSTM (Bai et al., 2018).
CNN:
We use the standard ResNet18 architecture, adapted to smaller image size of the CIFAR-10 dataset (changing stride from 2 to 1 of the ﬁrst convolutional layer, and removing max-pooling layer that follows).
Dense + CNN: We add an additional linear layer (i
e. a dense matrix) of size 1024 ×before the ResNet18 architecture.
This dense layer can in theory represent a permutation, but cannot beneﬁt from the additional techniques described above.
Baseline CNN (unpermuted): We use the standard ResNet18 architecture applied to the unpermuted CIFAR-10 dataset.
All models are trained for 200 total epochs, with the Adam optimizer.
We use the standard learning rate schedule and weight decay from Mostafa & Wang (2019).
We use Hyperband (Li et al., 2017) to tune other hyperparameters such as the initial learning rate and annealing temperature.
B. S PEEDING UP D YNAMIC C ONV ’ S INFERENCE B.4. M ODEL ARCHITECTURE
We start with the DynamicConv Transformer architecture (Wu et al., 2019), which is a variant of the
Transformer architecture (Vaswani et al., 2017) where the self-attention in each layer is replaced with a light-weight DynamicConv module
We the implementation from the library(Ott et al
The architecture of each layer of the decoder is: Linear → DynamicConv → Linear → LayerNorm → Encoder-decoder attention → LayerNorm → Linear → ReLU → Linear → ReLU → LayerNorm.
In every layer of the decoder, we replace the dense weight matrix in each of the four Linear layers with a K-matrix from the B class (i.e. a butterﬂy matrix).
B.4. T RAINING AND EVALUATION
The models are trained from scratch using the training script from the Fairseq repository, with the same hyperparameters (optimizer, learning rate, number of updates, etc.) used in the DynamicConv paper (Wu et al., 2019).
We note that the DynamicConv model with K-matrices in the decoder trains slightly faster than the default DynamicConv model (both models are trained for 50,000 updates, which requires approximately 7% less time for the K-matrix model than for the default model).
To evaluate inference speed, we run the decoding script on the IWSLT-14 De-En test set in singlethreaded mode on a server Intel Xeon CPU E5-2690 v4 at 2.60GHz, and measure wall-clock time.
The test set contains 6750 sentences, with 149241 tokens.
Following Wu et al.
(2019), we set the batch size to 1 and beam size to 1 for this evaluation.
B.4. A DDITIONAL COMPARISON WITH OTHER STRUCTURED MATRICES
We additionally compare the speed-quality tradeoff of K-matrices with other classes of structured matrices, when used to replace the fully-connected layers of DynamicConv’s decoder.
We consider the following additional classes of structured matrices: low-rank, circulant, Toeplitz-like (Sindhwani et al., 2015), ACDC (Moczulski et al., 2016), Fastfood (Le et al., 2013), and sparse.
For classes with a variable number of parameters (e.g. low-rank, sparse), we set the number of parameters to match that of K-matrices.
For sparse matrices, besides the result for an ensemble of 10 models (the default setting in the Fairseq repository), we also report the result for a single model, as that could have faster inference time (since ensembling/averaging sparse matrices produces a less sparse matrix).
In Figure 5, we plot the tradeoff between translation quality (measured by BLEU score) and inference speed (sentences per second).
Most classes of structured matrices produce similar translation quality (between 34.1 and 34.4 BLEU score).
K-matrices have the second fastest inference time, only 7% slower than low-rank matrices.
We note that low-rank matrices beneﬁt from very well-tuned BLAS routines (matrix-matrix multiplication).
Even though our implementation of K-matrix multiplication is not yet highly optimized, it is already quite close to the speed of low-rank matrix multiplication at an equivalent parameter count.
C S PEED BENCHMARK AND IMPLEMENTATION DETAILS
Each K-matrix (for ﬁxed width and expansion), has an O ( n log n ) matrix-vector multiplication algorithm: sequentially multiply the input vector with each of the sparse factors.
Our implementation of this simple algorithm is surprisingly competitive with optimized subroutines, both on GPU (e.g. for training) and on CPU (e.g. for inference).
In Figure 6, we compare the speed of multiplying by a K-matrix in class B (i.e. a butterﬂy matrix) against a specialized implementation of the FFT.
We normalize the speed by the speed of dense matrix-matrix multiply (on GPU) or dense matrix-vector multiply (on CPU).
On GPU, with input sizes n = 1024 and batch size 2048, the training time (forward and backward) of K-matrices matrix is about 3x faster than dense matrix multiply (GEMM from cuBLAS).
For inference on CPU, the kaleidoscope fast multiplication can be one or two orders of magnitude faster than GEMV.
Over a range of matrix sizes, our implementation is within a factor of 2-4x of specialized implementations of the FFT, a highly optimized kernel.
Our implementation is also memory-efﬁcient.
In the forward pass through the O (log n ) sparse factors, we do not store the intermediate results, but recompute them during the backward pass.
Therefore the activation memory required is O ( bn ) for an input batch size of b .
D S YNTHETIC MATRIX RECOVERY
We directly validate Theorem 1 on well-known types of structured matrices used in machine learning.
Given a structured matrix M , we attempt to represent M as closely as possible using K-matrices as BLEU 2. S e n t e n c e s / s Transformer
DynamicConv
Kaleidoscope
Low-rank Circulant Toeplitz-like ACDC
Fastfood Sparse
Sparse (ensemble)
[Figure]
well as the standard classes of structured matrices:
sparse and low-rank.
In Table 7, we quantify the expressivity of each of these three methods, as measured by their ability to approximate a range of different structures.
Results for “global minimum” of kaleidoscope matrices are obtained from the theoretical expressiveness results in Section I and Section J. Low-rank and sparse approximation have closed form solutions: truncating the SVD and keeping the largest-magnitude entries, respectively.
We also report the results using SGD for kaleidoscope matrices to validate that good approximation with K-matrices can be obtained even from standard ﬁrst-order optimization algorithms.
Even with imperfect optimization, kaleidoscope matrices can still capture out-of-class target matrices better than low-rank and sparse matrices.
The target matrices are kaleidoscope, low-rank, sparse, convolution (i.e. circulant matrices), Fastfood (Le et al., 2013), and entrywise random IID Gaussian matrix (to show the typical magnitude of the error).
All target matrices M were randomly initialized such that E [ M T M ]
= I .
[Table] To ﬁnd a kaleidoscope approximation with SGD
[Table]
we used Hyperband to tune its learning rate (between 0.001 and 0.5).
E P ROPERTIES OF THE BB ∗ H IERARCHY
Here, we justify why the deﬁnitions in Section 2.2 give rise to a hierarchy.
We ﬁrst make some basic observations about the parameterization.
Observation E.1.
An n × n matrix M ∈ BB ∗ has 4 n log n parameters.
Proof.
M can be expressed as a product of 2 log n butterﬂy factor matrices of size n × n .
Each of these factor matrices has 2 parameters per row, for a total of 2 n parameters each.
Hence, the total number of parameters is 4 n log n .
Observation E.2.
Let M be an n × n matrix in ( BB ∗ ) we .
Then, given an arbitrary vector v of length n , we can compute Mv with O ( wne log( ne )) ﬁeld operations.
Proof.
Since M ∈ ( BB ∗ ) we , we can decompose it as SE 1 E 2 . . .
E w S T , where S is as given in
Deﬁnition 2.4, and each E i is an en × en matrix in BB ∗ .
Therefore, to compute Mv , we can use associativity of matrix multiplication to multiply the vector by one of these matrices at a time.
Since all of these factors are sparse, we use the naïve sparse matrix-vector multiplication algorithm (begin with a 0-vector and perform the corresponding multiplication and addition for each nonzero matrix entry).
S (and thus S T ) have n NNZ.
Therefore, matrix-vector multiplication by S or S T requires O ( n ) operations, which is dominated by the butterﬂy matrix-vector multiplication.
Each
E i can be further decomposed into 2 log( ne ) matrices with at most 2 ne non-zero entries each (by
Observation E.1).
Therefore, matrix vector multiplication by each E i requires O ( ne log( ne )) .
Since there are w such E i , we require a total of O ( wne log( ne )) operations.
Now, we are ready to show that our deﬁnition of classes ( BB ∗ ) we forms a natural hierarchy.
First, we must argue that all matrices are contained within the hierarchy.
Lemma E.3.
Let M be an arbitrary n × n matrix.
Then M ∈ ( BB ∗ ) (2 n − 2) .
Proof.
Corollary E.3 in Appendix K shows that any n × n matrix can be written in the form
M 1 M (cid:48) 1 ∗ . . .
M n − 1 M (cid:48) n − 1 ∗ MM n M (cid:48) n ∗ . . .
M 2 n − 2 M (cid:48)
n − 2 ∗ , where M i , M (cid:48)
i are orthogonal butterﬂy matrices and M is a diagonal matrix.
We can combine D with M n to form another (possibly not orthogonal) butterﬂy matrix.
This yields a decomposition of M as products of (possibly not orthogonal) butterﬂy matrices and their (conjugate) transposes, completing the proof.
Next, we argue that, up to a certain point, this hierarchy is strict.
Lemma E.4.
For every ﬁxed c ≥ 1 , there is an n × n matrix M n (with n sufﬁciently large) such that
M n ∈ ( BB ∗ ) c +1 but M n (cid:54)∈ ( BB ∗ ) c .
Proof.
Given c , ﬁx n to be a power of 2 such that c < n 4log 2 n .
For sake of contradiction, assume that every n × n matrix in ( BB ∗ ) c
+1 is also in ( BB ∗ ) c .
Let A be an arbitrary n × n matrix.
From
Lemma E.3, A ∈ ( BB ∗ ) (2 n − 2) .
From our assumption, we can replace the ﬁrst c + 1 BB ∗ factors of A with c (potentially different) BB ∗ factors and still recover A .
We can repeat this process until we are left with c BB ∗ factors, implying that A ∈ ( BB ∗ ) c .
From Observation E.1, we require 4 cn log n < n 2 (by our choice of n ) parameters to completely describe A .
This is a contradiction since A is an arbitrary n × n matrix, and therefore has n 2 arbitrary parameters.
Hence, there must be some n × n matrix in ( BB ∗ ) c
+1 that is not in ( BB ∗ ) c .
F A RITHMETIC C IRCUITS IN BB ∗ H IERARCHY
In this appendix, we prove our main theoretical result, namely, our ability to capture general transformations, expressed as low-depth linear arithmetic circuits, in the BB ∗ hierarchy.
This result is recorded in Theorem 1.
Theorem 1.
Let M be an n × n matrix such that matrix-vector multiplication of M times an arbitrary vector v can be represented as a be a linear arithmetic circuit C comprised of s gates (including inputs) and having depth d .
Then, M ∈ ( BB ∗ ) O
( d ) O ( sn ) .
To prove Theorem 1, we make use of the following two theorems.
Theorem 2.
Let P be an n × n permutation matrix (with n a power of 2).
Then P ∈ BB ∗ .
Theorem 3.
Let S be an n × n matrix of s NNZ.
Then S ∈ ( BB ∗ ) 5 (cid:100) [Equation]
Theorem 2 is proven in Appendix G, and Theorem 3 is proven in Appendix I.
Proof of Theorem 1.
We will represent C as a product of d matrices, each of size s (cid:48)
× s (cid:48) , where s (cid:48) is the smallest power of 2 that is greater than or equal to s .
To introduce some notation, deﬁne w 1 , . . .
w d such that w k represents the number of gates in the k ’th layer of C (note that s = n + (cid:80) dk
=1 w k ).
Also, deﬁne z 1 , . . .
z d such that z 1 = n and z k = w k − 1 + z k − 1 ( z k is the number of gates that have already been used by the time we get to layer k ).
Let g i denote the i ’th gate (and its output) of C ( 0 ≤ i < s ), deﬁned such that:
[Equation]
where i 1 , i 2 are indices of gates in earlier layers.
For the k ’th layer of C , we deﬁne the s (cid:48) × s (cid:48) matrix M k such that it performs the computations of the gates in that layer.
Deﬁne the
i ’th row of M k to be:
[Equation]
For any 0 ≤ k ≤ d , let v k be vector
[Equation]
We’d like to argue that v d contains the outputs of all gates in C (i.e, the n values that make up
Mv ).
To do this we argue, by induction on k , that v k is the vector whose ﬁrst z k +1 entries are g 0 , g 1 , . .
, g ( z k − 1) , and whose remaining entries are 0 .
The base case, k = 0 is trivial.
Assuming this holds for the case k − 1 , and consider multiplying v k − 1 by M k .
The ﬁrst z k rows of M k duplicate the ﬁrst z k entries of v k −
The next w k rows perform the computation of gates g z k , . .
, g ( z k +1 − 1) .
Finally, the remaining rows pad the output vector with zeros.
Therefore, v k is exactly as desired.
The ﬁnal matrix product will contain all n elements of the output.
By left multiplying by some permutation matrix P , we can reorder this vector such that the ﬁrst n entries are exactly Mv .
Hence, we are left to argue the position of PM d . . .
M 2 M 1 within the BB ∗ hierarchy.
Each M k is a matrix with total 2 w k + z k < 2 s (cid:48) NNZ.
From Theorem 3, we can, therefore, represent M k as a product of O (1) matrices (of size 2 s (cid:48) ) in BB ∗ .
From Theorem 2, P ∈ BB ∗ .
Note that s ≤ s (cid:48) < 2 s , so s (cid:48)
=
Θ( s ) .
Our ﬁnal decomposition will have O ( d ) BB ∗ factors, and requires an expansion from size n to size 2 s (cid:48) , or an expansion factor of O ( sn ) .
Therefore, M ∈ ( BB ∗ ) O
( d ) O ( sn ) , as desired.
Remark F.1.
By applying Observation E.2, we see that Theorem 1 gives an O ( sd log s ) matrix vector multiplication algorithm for M .
G P ERMUTATIONS IN BB ∗
In this appendix, we prove Theorem 2.
To do this, we decompose permutation matrix P into P = LR , with L ∈ B and R ∈ B ∗ .
Throughout the proof, we make use of the following deﬁnition.
Deﬁnition G.1.
Let L be an n × n permutation matrix ( n a power of 2).
We say that L meets the 2 j balance condition if L can be divided into chunks of 2 j (with each chunk having all columns i such that (cid:4) i 2 j (cid:5) has the same value) such that for every 0 ≤ m < 2 j , each chunk has exactly one L [: , k ] =
e π k with π k ≡ m ( mod 2 j ) .
We say that L is modular-balanced if it meets the 2 j balance condition for each 2 ≤ 2 j ≤ n .
[Equation]
Lemma G.1.
Let L be an n × n modular-balanced matrix.
Then L ∈ B .
Proof.
We proceed by induction on n .
The base case n = 2 is trivial.
As our inductive hypothesis, we assume that all modular-balanced matrices of size n 2 × n 2 are butterﬂy matrices of size n 2 .
From
Deﬁnition 2.3, it is sufﬁcient to show that L can be decomposed as:
[Equation]
where B n is a butterﬂy factor of size n and each L j is an n 2 × n 2 modular-balanced matrix.
Deﬁne L 1 and L 2 such that:
[Equation]
Note that since L is a permutation matrix (and thus has exactly one non-zero entry per column), at most one term of each of these sums can be non-zero.
For sake of contradiction, assume L 1 is not modular-balanced.
Then, for some 2 j ≤ n 2 , there are two columns c 1 , c 2 such that (cid:4) c 1 2 j (cid:5) = (cid:4) c
2 2 j (cid:5) and such that indices of the non-zero entries of L 1 in columns c 1 and c 2 are the same modulo 2 j .
However, from the deﬁnition of L 1 , this implies that the indices of the non-zero entries of L in columns c 1 and c 2 are also the same modulo 2 j , contradicting L being modular-balanced.
Hence, L 1 is modular-balanced.
An analogous argument (that instead considers columns c 1 + n 2 , c 2 + n 2 of L ) shows that L 2 is also modular-balanced.
To complete the proof, we must argue that B n is a butterﬂy factor of size n .
Since each L i is modular-balanced, it is a permutation matrix.
Therefore, L (cid:48) has exactly 1 non-zero entry in each of the ﬁrst n 2 rows and columns from L 1 and exactly 1 non-zero entry in each of the second n 2 rows and columns from L 2 .
Hence, L (cid:48) is a permutation matrix.
Since both L and L (cid:48) are permutation matrices,
B = L ( L (cid:48) )
− 1 must also be a permutation matrix.
Therefore, we can view B as performing a permutation of the rows of L (cid:48) to get L .
Consider the i ’th row of L (cid:48) , with 0 ≤ i < n 2 .
There are two possible cases.
Case 1: L (cid:48) [ i, :] = L [ i, :]
In this case, the column of L with a non-zero entry in row i is in the left n 2 columns.
The column of L with a non-zero entry in row i + n 2 must, therefore, be in the right n 2 columns, otherwise
L would not satisfy the n 2 balance condition.
Therefore, L (cid:48) (cid:2) i + n 2 , : (cid:3) =
L (cid:2) i + n 2 , : (cid:3) , so we set
B [ i, i ] = B (cid:2)
i + n 2 , i + n 2 (cid:3) = 1 .
Case 2: L (cid:48)
[ i, :] (cid:54) =
L [ i, :]
By the deﬁnition of L (cid:48) , L (cid:48)
[ i, :] =
vL (cid:2) i + n 2 , : (cid:3) .
In this case, the column of L with a non-zero entry in row i + n 2 must be in the left n 2 columns.
By the n 2 balance condition of L , the column of L with a non-zero entry in row i must be in the right n 2 columns.
Therefore, L (cid:48) (cid:2) i + n 2 , : (cid:3) =
L [ i, :] , so we set
B (cid:2) i, i + n 2 (cid:3) = B (cid:2)
i + n 2 , i (cid:3)
= 1 .
In both cases, the non-zero entries of B fall into the correct diagonal bands (the main diagonal, and the bands n 2 away).
Hence, B is a butterﬂy factor of size n .
Now, we consider the process of transforming P into a modular-balanced matrix.
We make use of the following lemma.
[Equation]
P
G
B n
Lemma G.2.
Let M be a k × k matrix with 1 non-zero entry per column, such that for each 0 ≤ m < k 2 , there are exactly 2 columns with non-zero entry in a row with index ≡ m (cid:0) mod k 2 (cid:1) .
Then, there is a butterﬂy factor B k such that MB k
= M (cid:48) , where M (cid:48) meets the k 2 balance condition.
Proof.
We construct a directed graph G with nodes in (cid:2) k 2 (cid:3) .
For each 0 ≤ i < k
we add a directed edge from node (cid:0) s mod k 2 (cid:1) to node (cid:0) t mod k 2 (cid:1) if M [: , i ] =
e s and M (cid:2) : , i + k 2 (cid:3)
=
e t .
Each node has (undirected) degree exactly 2 by the structure of M .
Hence, G is a union of disjoint (undirected) cycles.
If M met the k 2 balance condition, then each node would additionally have in-degree exactly 1 and out-degree exactly 1.
By reversing edges of G such that each (undirected) cycle becomes a directed cycle, we can achieve this.
However, reversing edges corresponds to swapping columns of M that are k 2 apart.
Let B k be the permutation matrix that performs all such swaps.
B k has non-zero entries only along the main diagonal and the diagonal bands k 2 away, and thus is a butterﬂy factor of size k .
We are ready to present the decomposition of P .
Lemma G.3.
Let P be an n × n permutation matrix.
Then we can decompose P into P = LR , where L is modular-balanced and R ∈ B ∗ .
Proof.
We repeatedly apply Lemma G.2.
First, we conclude that there is a butterﬂy factor B n such that
[Equation]
where P (cid:48) meets the n 2 balance condition.
Now, we consider the ﬁrst and last n 2 columns of P (cid:48) independently.
We can again apply Lemma G.2 (twice) to conclude that there are butterﬂy factors (cid:2)
B n 2 (cid:3) 1 , (cid:2) B n 2 (cid:3) 2 such that
[Equation]
where P (cid:48)(cid:48) meets the n 2 and n 4 balance conditions.
We continue this process until we obtain a matrix that meets all of the balance conditions.
Our ﬁnal equation is of the form:
[Equation]
where B is a butterﬂy matrix and L is a modular-balanced matrix.
Let R =
B − 1 =
B ∗ (since B is a permutation matrix, and thus is orthogonal) and hence R ∈ B ∗ .
Then P = LR , as desired.
Theorem 2 follows immediately from Lemmas G.3 and G.1.
H BB ∗ C LOSURE L EMMAS
Here, we present some basic facts of the BB ∗ hierarchy that will be useful for later constructions.
For simplicity, we assume (WLOG via 0-padding) that all matrices are square matrices with size that is a power of 2.
Lemma H.1.
If M ∈ B (or M ∈ B ∗ ), then DM , MD ∈ B ( B ∗ resp.)
for any diagonal matrix D .
Proof.
Left multiplication by a diagonal matrix scales the rows of M by the corresponding diagonal entries.
The same can be achieved by scaling all entries the leftmost butterﬂy factor matrix.
Similarly, right multiplication by a diagonal matrix scales the columns of M , which can be achieved by scaling all entries in the columns of the rightmost butterﬂy factor matrix.
Lemma H.2.
Let A , B ∈ F n × n .
If A ∈ ( BB ∗ ) w 1 e and B ∈ ( BB ∗ ) w 2 e then AB ∈ ( BB ∗ ) w 1 + w 2 e .
Proof.
Let E A , E B ∈ F en × en be deﬁned such that A = SE A S T , B = SE B S T (with S as in
Deﬁnition 2.4).
Then
[Equation]
en E A (cid:20) I n   (cid:21) (cid:124) (cid:123)(cid:122) (cid:125) en × en E B S T
(cid:20) I n   (cid:21)
E A ∈ ( BB ∗ )
w 1 , (cid:20) I n   (cid:21)
E B ∈ ( BB ∗ )
w 2 by Lemma H.1.
Hence, AB ∈ ( BB ∗ ) w 1 + w 2 e by
Deﬁnition 2.4.
Lemma H.3.
Let A 1 , . . .
, A m ∈ F k × k .
If A 1 , . . . , A m ∈ ( BB ∗ )
we then Diag( A 1 , . . .
, A m ) ∈ ( BB ∗ )
w +2 e .
Proof.
For each 1 ≤ i ≤ m , let E A i ∈
F ek × ek be deﬁned such that A i = SE
A i S T (with S as in
Deﬁnition 2.4).
Then  
A . . . A. . .
A m   = SP  
E A . . . E A. . .
E
A m 
 P T S T where P is a permutation that that moves the ﬁrst k rows of each E A
i (in order) into the top mk rows.
From Theorem 2, P ∈ BB ∗ , (and so is P T , also a permutation).
Within the RHS block matrix, the decompositions of each E A i can be done in parallel, requiring total width w .
Hence,
Diag( A 1 , . . . ,
A m ) ∈
( BB ∗ ) w +2 e , as desired.
Remark H.4.
If e = 1 in Lemma H.3, then P is unnecessary.
Hence, Diag( A 1 , . . .
, A m ) ∈ ( BB ∗ ) w .
Lemma H.5.
Let A 1 , . . .
A m be k × k matrices in ( BB ∗ )
we then (cid:80) mi =1 A i ∈ ( BB ∗ ) mw 4 e .
Proof.
For each 1 ≤ i ≤ m , let E A i ∈
F ek × ek be deﬁned such that A i = SE
A i S T (with S as in
Deﬁnition 2.4).
Note that E A i ∈ ( BB ∗ ) w .
Consider matrices of the form:  
[Equation]
[Equation] [Equation]   [Equation]
[-E A i 
-][Equation] [Equation]
  (cid:124) (cid:123)(cid:122) (cid:125)
[Equation]
 
[Equation]
[Equation]
[Equation]
  (cid:124) (cid:123)(cid:122) (cid:125)
[Equation]
(cid:20)
[Equation] I[-ek0 I ek        
 I ek-]
[Equation]
(cid:21) (cid:124) (cid:123)(cid:122) (cid:125) R .
Here, L and R compute the sum of the 2 ek × 2 ek matrices on the diagonal of SP 1 , where P 1 is a permutation swapping E A i to the 4 th ek -block column.
Note that S is the diagonalization of four matrices in ( BB ∗ ) w , so S ∈ ( BB ∗ ) w by Remark H.4.
In addition, since each block in S is a butterﬂy matrix of size ek , S only uses butterﬂy factors up to size ek , so the outer factor matrices of sizes 4 ek and 2 ek in S are unused.
Also note that L and R are butterﬂy factor matrices of size 4 ek (or B (4 ek ) 4 ek ), and P 1 is a butterﬂy factor matrix of size 2 ek (or B (4 ek ) 2 ek ).
This allows us to fold the surrounding matrices L , P 1 , R into S , so M i ∈ ( BB ∗ ) w .
Through repeated application ( m times) of the identity
[Equation]
we see that
[-=1 E A i 
 I ek 0  0  0  (cid:124) (cid:123)(cid:122) (cid:125) M ∈ F 4 en × 4 en = m (cid:89) i = M i .-][Equation]
From Lemma H.2, M ∈ ( BB ∗ ) mw .
Finally, note that (cid:80) mi =1 A
i =
SMP 2 S T , where P 2 is a permutation that moves the ﬁrst k columns of the second block-column of M to the left.
P 2 can be folded into the ﬁnal summation factor M m as follows:  
I ek0 I ek
I ek
I ek  (cid:124) (cid:123)(cid:122) (cid:125) P(cid:20)
I 2 ek
I 2 ek (cid:21) (cid:124) (cid:123)(cid:122) (cid:125) R  
I ek
I ek I ek
I ek 
 (cid:124) (cid:123)(cid:122) (cid:125) P=  
I ek
I ek
I ek0 I ek   (cid:124) (cid:123)(cid:122) (cid:125) P (cid:48)(cid:20)
I 2 ek
I 2 ek (cid:21) (cid:124) (cid:123)(cid:122) (cid:125) R (3)
Hence, (cid:80) mi =1 A i ∈ ( BB ∗ ) mw 4 e , as desired.
Lemma H.6.
Let M be an invertible n × n matrix such that M ∈ B .
Then M − 1 ∈ B ∗ .
Proof.
We prove this in a series of steps.
First, let B k be an invertible butterﬂy factor of size k .
Consider the method of computing B − 1 k by performing Gaussian elimination on the matrix [ B k | I k ] to obtain the matrix (cid:2)
I k | B − 1 k (cid:3) .
By the form of B , non-zero entries within a row or column are always exactly k 2 positions apart.
Therefore, the only row operations needed for this Gaussian elimination are:
• Scaling a row by a constant factor c (cid:54) =
• Addition of a row to another row exactly k 2 rows apart
Performing these operations on I k will only allow non-zeros on the main diagonal and k 2 diagonals away from the main diagonal.
Hence, B − 1 k is also a butterﬂy factor of size k .
Next, let B ( n ) k be an invertible butterﬂy factor matrix of size n and block size k .
Its inverse is the block diagonal matrix formed by the inverses of each of its constituent butterﬂy factors.
From above, (cid:16)
B ( n ) k (cid:17) − 1 is also a butterﬂy factor matrix of size n and block size k .
Finally, consider M ∈ B .
[Equation]
Finally, we include a closure result for the Kronecker product, another common matrix composition operation.
Although Lemma H.7 is not directly used in the subsequent proofs, it allows for examples the results for the DFT to be lifted to higher-dimensional Fourier transforms.
We also note that the closure bound in Lemma H.7 can be tightened in such cases (cf ˙ Remark H.4).
Lemma H.7.
Let A , B ∈ F n × n .
If A ∈ ( BB ∗ ) w 1 e and B ∈ ( BB ∗ ) w 2 e then A ⊗ B ∈ ( BB ∗ ) w 1 + w 2 +6 e .
Proof.
Note that
[Equation]
, for some permutation P .
By Lemma H.3, I ⊗ A
and I ⊗ B are in ( BB ∗ )
w 1 +2 e , ( BB ∗ ) w 2 +
e respectively.
The result follows from combining with P ∈ BB ∗ and Lemma H.2.
I S PARSE M ATRICES IN BB ∗ H IERARCHY
In this appendix, we prove Theorem 3.
First, we consider matrices with at most n NNZ.
Lemma I.1.
let S be an n × n matrix with at most n NNZ.
Then, S ∈ ( BB ∗ ) 5 .
We use this lemma and the addition closure lemma to prove Theorem 3.
Proof of Theorem 3.
We note that any s sparse matrix is the sum of (cid:6)
sn (cid:7) matrices of at most n NNZ, and we appeal to Lemma H.5.
In the rest of the section we will prove Lemma I.1.
We begin by deﬁning two classes of matrices that will be used in our decomposition.
Deﬁnition I.1.
An n × n matrix H is a horizontal step matrix if for every 0 ≤ i, i (cid:48) < n and 0 ≤ j ≤ j (cid:48) < n , if H [ i, j ] (cid:54) = 0 and H [ i (cid:48) , j (cid:48) ] (cid:54) = 0 , then j (cid:48) − j ≥
( i (cid:48) − i ) mod n .
An n × n matrix V is a vertical step matrix if V ∗ is a horizontal step matrix.
With this deﬁnition, the horizontal step matrix obeys a “Lipschitz-like" condition.
Each column of a horizontal step matrix can have at most one non-zero entry, and given two non-zero columns k apart, the non-zero entry in the right column must be between 0 and k rows below the non-zero entry in the left column.
Note that to show that a matrix is a horizontal step matrix, it is sufﬁcient to argue that this condition holds for each pair of neighboring non-zero columns.
Similarly, each row of a vertical step matrix can have at most one non-zero entry, and given two non-zero rows k apart, the non-zero entry in the lower row must be between 0 and k columns to the right of the non-zero entry in the upper row.
Lemma I.2.
Let H be an n × n horizontal step matrix.
Then H ∈ B .
Proof.
We proceed by induction on n .
The base case n = 2 is trivial.
As our inductive hypothesis, we assume that all horizontal step matrices of size n 2 × n 2 are butterﬂy matrices of size n 2 .
From
Deﬁnition 2.3, it is sufﬁcient to show that H can be decomposed as:
[Equation]
where H 1 , H 2 are n 2 × n 2 horizontal step matrices and each D k is a n 2 × n 2 diagonal matrix.
Denote the four, n 2 × n 2 corner submatrices of H by:
[Equation]
Then, deﬁne H 1 and H 2 by:
[Equation]
For sake of contradiction, assume that H 1 is not a horizontal step matrix.
Then, there are 0 ≤ i, i (cid:48) < n 2 , 0 ≤ j ≤ j (cid:48) < n 2 such that H 1 [ i, j ] (cid:54) = 0 ,
H
[ i (cid:48) , j (cid:48) ] (cid:54) = 0 , and j (cid:48) − j <
( i (cid:48) − i ) mod n 2 .
From our deﬁnition of H 1 , the non-zero entries in columns j and j (cid:48) of H are either (cid:0) ( i (cid:48) − i ) mod n 2 (cid:1) or (cid:0) n
+ ( i (cid:48) − i ) mod n 2 (cid:1) , both of which are greater than j (cid:48)
− j , rows apart.
This contradicts H being a horizontal step matrix.
Hence, H 1 must be a horizontal step matrix, as must H 2 from an analogous argument.
Next, we deﬁne D 1 , D 2 , D 3 , D 4 by:
D 1 [ k, k ] = (cid:26) 1 H 21 [ k, :] =0 otherwise
D 2 [ k, k ] = (cid:26) 1 H 22 [ k, :] =0 otherwise
D 3 [ k, k ] = (cid:26) 1 H 11 [ k, :] =0 otherwise .
D 4 [ k, k ] = (cid:26) 1 H 12 [ k, :] =0 otherwise .
To ﬁnish the proof, we argue the correctness of the decomposition by equating arbitrary entries of each of the 4 corner submatrices.
We begin with the upper left submatrix.
[Equation] ] by deﬁnition of matrix multiplication = D [Equation] [ i, i ]
[Equation]
H
[Equation]
[-, j ] = n 2 (cid:88) k = D 1 [ i, k ] · H 1 [ k, j ] by deﬁnition of matrix multiplication = D 
-][Equation]
, j ]
D 1 is a diagonal matrix =[-
-]
[Equation] ( H 21 [ i, :]= 0 ) · ( H 11 [ i, j ] + H 21 [ i, j ]) by deﬁnition of D 1 and H
Here, we consider two cases:
Case 1: H 21 [ i, j ] (cid:54) =
Since H is a horizontal step matrix (and hence may have at most one non-zero entry per column), it follows that H
[ i, j ] = 0 .
In this case, the indicator function evaluates to 0 , so D 1 H 1 [ i, j ] =
=
H 11 [ i, j ] , as desired.
Case 2: H 21 [ i, j ] =
If H 11 [ i, j ] = 0 , then D 1 H 1 [ i, j ] =
= H 11 [ i, j ] .
Otherwise, for sake of contradiction, suppose that
H 21 [ i, :] (cid:54) = 0 .
Then, two of the ﬁrst n 2 columns of H would have non-zero entries n 2 rows apart, contradicting H being a horizontal step matrix.
Hence, H
[ i, :] = 0 , so D 1 H 1 [ i, j ] =
H 11 [ i, j ] , as desired.
In all cases, D 1 H 1 [ i, j ] =
H 11 [ i, j ] , so our decomposition correctly recovers the upper left corner of
H .
Analogous arguments show that the other three corners are also correctly recovered.
Hence, our decomposition is correct, and by induction, H ∈ B .
Corollary I.1.
Let V be a vertical step matrix.
Then V ∈ B ∗ .
S =
P
S (cid:48)
P||
H
V (cid:48)
||
P
V
Now, we use step matrices to prove Lemma I.1.
Proof of Lemma I.1.
Given S , we decompose it as S = P 1 HP 2 VP 3 , where each P (cid:96) is a permutation matrix, H is a horizontal step matrix, and V is a vertical step matrix.
For an example of this, see
Figure
We ﬁrst decompose S as S = P 1 S (cid:48) P 3 , where P 1 is the permutation that moves all 0 rows of S to the bottom and P 3 is the permutation that moves all 0 columns of S to the right.
Next, we further decompose S (cid:48) into S (cid:48) =
HV (cid:48) as follows.
Since S (cid:48) has s ≤ n NNZ, we can parameterize S (cid:48) by θ
= { ( c k , i k , j k ) : 0 ≤ k < s } such that S (cid:48)
[ i k , j k ] =
c k , with the non-zero entries indexed in row-major order.
Deﬁne matrix H by:
H [: , k ] = (cid:26) [Equation] k < s otherwise .
Deﬁne matrix V (cid:48) by:
V [Equation]
[ k, :] = (cid:26) e Tj k 0 ≤ k < s otherwise .
To show that S (cid:48) =
HV (cid:48) , we consider an arbitrary entry: HV (cid:48) [ i, j ] = n (cid:88) [Equation]
[Equation]
[Equation]
[Equation]
j ] by deﬁnition of matrix multiplication
[Equation]
H is 0 in all but ﬁrst s columns
[Equation]
by deﬁnition of H and V (cid:48)
Here, we note that ( i, j ) can equal ( i k , j k ) for at most one value of k since the locations in θ are unique.
Hence, HV (cid:48)
[ i, j ] =
c k only if ( i, j ) =
( i k , j k ) for some k , which is exactly the deﬁnition of
S (cid:48) .
Hence, S (cid:48) =
HV (cid:48) .
We argue that H is a horizontal step matrix through a series of assertions.
First, note that H has exactly one non-zero entry in each of its ﬁrst s columns.
Also, note that since θ is in row-major order, these non-zero entries are sorted (any column to the right cannot have a non-zero entry in a higher row).
Hence, to show that H is a horizontal step matrix, it is sufﬁcient to argue that adjacent columns of H have non-zero entries at most one row apart.
This is equivalent to S (cid:48) having no zero rows between two non-zero rows, which is guaranteed by P 1 .
Hence, H is a horizontal step matrix.
Since V (cid:48) has at most one non-zero entry per row, we may permute the rows of V (cid:48) to obtain a matrix
V , where the non-zero entries of V are sorted (any lower row below cannot have a non-zero entry in an earlier column).
Hence, for some permutation matrix ( P 2 ) − 1 , V =
( P 2 ) − 1 V (cid:48) , which implies that V (cid:48) =
P 2 V .
It has exactly one non-zero entry in each of its ﬁrst s columns.
From the action of P 2 , these non-zero entries are sorted.
Therefore, by the same argument as for H above, V T is a horizontal step matrix.
Hence, V is a vertical step matrix.
In all, we have found a decomposition S
= P 1 HP 2 VP 3 , where each P (cid:96) is a permutation matrix ( ∈ BB ∗ by Theorem 2), H is a horizontal step matrix ( ∈ BB ∗ by Lemma I.2), and V is a vertical step matrix ( ∈ BB ∗ by Corollary I.1).
By Lemma H.2, S ∈ ( BB ∗ ) 5 .
Corollary I.2.
Let R be an n × n matrix of rank r .
Then R ∈ ( BB ∗ ) 10 r 4 .
Proof.
We can decompose R as R =
GH ∗ where G , H are n × r matrices.
With appropriate zero-padding, both of these can be made into n × n matrices with at most rn NNZ.
The proof follows immediately from Theorem 3 and Lemma H.2.
J C OMPARISON TO BP H IERARCHY
In this appendix, we draw comparisons between the BB ∗ hierarchy and the BP hierarchy introduced by Dao et al.
Lemma J.1.
Let F n be the Discrete Fourier Transform of size n .
Then F n ∈ ( BB ∗ ) 2 .
Proof.
From Parker (1995), we can express F n as F n = B P , where B ∈ B and P is a permutation (the bit reversal permutation).
From Theorem 2, P ∈ BB ∗ .
Hence, by Lemma H.2, F n ∈ ( BB ∗ ) 2 .
Lemma J.2.
Let H n be the Hadamard Transform of size n .
Then H n ∈ BB ∗ .
Proof.
H n ∈ B , so trivially H n ∈ BB ∗ .
Lemma J.3.
Let S n be the Discrete Sine Transform of size n .
Then S n ∈ ( BB ∗ ) 2 .
Proof.
As described in Makhoul (1980), S n can be performed as a scaled permutation (separating the even and odd indices of the input, and reversing and negating the odd indices) composed with F n .
Therefore, we may decompose S n as S n = B P 2 D P 1 , where P 1 , P 2 are permutations, B ∈ B , and D is a diagonal matrix.
P 2 D P 1 is simply a permutation matrix with scaled entries, which can be equivalently expressed as D (cid:48) P (cid:48) for some diagonal matrix D (cid:48) and permutation P (cid:48) .
By Lemma H.1,
B D (cid:48)
∈ BB ∗ .
By Theorem 2, P (cid:48) ∈ BB ∗ .
Hence, by Lemma H.2, S n ∈ ( BB ∗ ) 2 .
Remark J.4.
An analogous argument shows that the Discrete Cosine Transform is also in ( BB ∗ ) 2 .
Lemma J.5.
Let C n be an n × n circulant (convolution) matrix.
Then C n ∈ BB ∗ .
Proof.
Using Theorem 2.6.4 of Pan (2001), we can express C n as C n = ( F n ) − 1 DF n where F n is the Discrete Fourier Transform and D is a diagonal matrix.
( F n ) − 1 =
B P (with B ∈ B , P a permutation), which implies that F n = ( P ) − 1 ( B ) − 1 .
Therefore
[Equation]
The middle three factors have the effect of performing a permutation, scaling each element, and undoing the permutation, which is equivalent to simply scaling by some diagonal matrix D (cid:48) .
Hence, we are left with
[Equation]
By Lemma H.1, B D (cid:48) ∈ B .
By Lemma H.6, ( B ) − 1 ∈ B ∗ .
Hence, C n ∈ BB ∗ .
Remark J.6.
We can expand any n × n Toeplitz matrix T n into a 2 n × 2 n circulant matrix (with upper left n × n submatrix equal to T n ).
Hence, T n ∈ ( BB ∗ ) 12 by Lemma J.5.
The Fastfood matrix class (Le et al., 2013) can be tightly captured in the BB ∗ hierarchy:
Lemma J.7.
The product SHDPHB where S , D , B are diagonal matrices, H is the Hadamard transform, and P is a permutation matrix, is in ( BB ∗ ) 3 .
Proof.
We have shown in Lemma J.2 that H ∈ BB ∗ , and in Theorem 2 that P ∈ BB ∗ .
Since BB ∗ is closed under diagonal multiplication (Lemma H.1), we conclude that SHDPHB ∈ ( BB ∗ ) 3 .
The two classes of matrices introduced in Moczulski et al. (2016), called AFDF and ACDC, are also tightly captured in the BB ∗ hierarchy:
Lemma J.8.
Let AF − 1 DF be a product of a diagonal matrix A , the inverse Fourier transform F − 1 , another diagonal matrix D , and the Fourier transform F .
Then AF − 1 DF ∈ BB ∗ .
Let AC − 1 DC be a product of a diagonal matrix A , the inverse cosine transform C − 1 , another diagonal matrix D , and the cosine transform C .
Then AC − 1 DC ∈ ( BB ∗ ) 4 .
Proof.
We have argued in Lemma J.5 that F − 1 DF ∈ BB ∗ .
Since BB ∗ is closed under diagonal multiplication (Lemma H.1), we conclude that AF − 1 DF ∈ BB ∗ .
We have shown that C ∈ ( BB ∗ ) 2 , so C − 1 ∈ ( BB ∗ ) 2 as well.
Since BB ∗ is closed under diagonal multiplication (Lemma H.1), we conclude that AC − 1 DC ∈ ( BB ∗ ) 4 .
Remark J.9.
Within each butterﬂy factor matrix of the DFT (excluding the bit reversal permutation) and the Hadamard transform, the columns are pairwise orthogonal and have norm 2.
Hence, we can divide all factors by √ 2 to make orthogonal factor matrices.
To counteract this scaling, we can add a diagonal matrix with √ 2 log 2 ( n ) = √ n in all entries to the factorization.
By doing this we can place all of the above transforms in the OBB hierarchy (deﬁned in Appendix K) with the same width and expansion factor.
J. M ULTI - DIMENSIONAL TRANSFORMS
Here, we show that, using larger matrices, we are able to similarly capture multi-dimensional versions of the above transforms.
Lemma J.10.
Let F 2 n be the 2-dimensional Discrete Fourier Transform (represented as an n 2 × nmatrix). Then F 2 n ∈ ( BB ∗ ) 2 .
Proof.
The separation property of the 2-D DFT allows us to express its action on an n × n matrix as the composition of a 1-D DFT on each of its rows and a 1-D DFT on each of its columns.
If we view the 2-D DFT as an n 2 × n 2 matrix, its input and outputs will both be column vectors of size n 2 .
As our convention, we list the entries of the input vector in the row-major order corresponding to the n × n input matrix.
Then, we consider the 2-D DFT in four steps, where the ﬁrst two steps perform the 1-D DFT row-wise, and the second two steps perform the 1-D DFT column-wise:
Step 1: Permute the columns:
We permute the columns (with a bit reversal permutation), which performs a bit reversal permutation on each row.
Viewing the input as a vector, this step corresponds to left multiplication by a permutation matrix P c that permutes the entries of each chunk of size n of the input vector.
Step 2:
Multiply each row by a butterﬂy matrix
Since the entries of the input were listed in row major order, this step is achieved through multiplication by a block diagonal matrix of n butterﬂy matrices of size n , which can be viewed as a product of butterﬂy factor matrices B [Equation] n
[Equation]
Step 3: Permute the rows:
We permute the rows (with a bit reversal permutation), which performs a bit reversal permutation on each column.
This corresponds to left multiplication by a permutation matrix P r .
Since we are permuting the rows, P r permutes the entries at the granularity of each n -chunk.
Since Steps 1 andeach performed an identical computation to each n -chunk we can move this row permutation before
Step 2, combining P c and P r into a single permutation P .
Step 4:
Multiply each column by a butterﬂy matrix
Consider multiplication by the ﬁrst factor matrix.
In each row, this matrix is taking linear combinations of adjacent column entries.
In our length- n 2 vector, these entries will be exactly n indices apart.
Therefore this multiplication can be handled by a butterﬂy factor matrix B ( n 2 ) 2 n .
Similarly, we ﬁnd that this butterﬂy multiplication can be expressed as multiplication by a product of butterﬂy factor matrices B ( n 2 ) n 2 . . .
B ( n 2 ) n 22 B ( n 2 ) 2 n .
Combined with the factor matrices from Step 2, these form a butterﬂy matrix B of size n 2 .
In all, we see that the 2-D DFT may be realized as multiplication by a permutation matrix P followed by multiplication by a butterﬂy matrix B .
The same argument as Lemma J.1 shows that
F 2 n ∈ ( BB ∗ ) 2 .
Remark J.11.
An analogous argument (using the separation property of the respective transforms) can be used to argue that 2-D Discrete Sine and Discrete Cosine transforms are in ( BB ∗ ) 2 , and that 2-D Hadamard Transforms are in BB ∗ .
Lemma J.12.
Let C 2 n be a 2-dimensional convolution matrix.
Then C 2 n ∈ BB ∗ .
Proof.
We can express a 2-D convolution matrix as C 2 n = ( F 2 n ) − 1 DF 2 n , where D is diagonal, F 2 n is the 2-D Fourier transform and ( F 2 n ) − 1 is the inverse 2-D Fourier transform.
From the proof of Lemma J.10, we see that that we can express F 2 n (and similarly ( F 2 n ) − 1 ) as the product of a butterﬂy matrix and a permutation matrix.
The rest of the argument is analogous to the proof of
Lemma J.5.
Remark J.13.
Using an inductive argument, we can show that all k-dimensional ( k ∈ Z ) variants of the above transforms, expressed as n k × n k matrices are contained in BB ∗ or ( BB ∗ ) 2 .
To do this, we use the separation property of the transforms to break them into a k − 1 -dimensional transform (the inductive hypothesis) followed by a 1-dimensional transform.
K T HE O RTHOGONAL K ALEIDOSCOPE H IERARCHY
Through practical application of the butterﬂy matrices, it has been found useful to constrain them in orthogonality.
In Section K.1 we will modify the existing kaleidoscope hierarchy to create the orthogonal kaleidoscope hierarchy OBB .
Then, in Section K.2, we will argue that all orthogonal matrices, and as a result all matrices, can also be expressed in this hierarchy in O ( n ) width.
Lastly, in
Section K.3, we will argue that permutation matrices and sparse matrices also exist in this hierarchy in O (1) width, which in turn implies a corresponding result for matrices with low-depth arithmetic circuits.
K. D EFINITION
The deﬁnition of the orthogonal butterﬂy is identical to the original butterﬂy, with the constraint that all butterﬂy factors are orthogonal.
We specify this deﬁnition below:
Deﬁnition K.1 (Analog of Deﬁnition 2.1) .
An orthogonal butterﬂy factor of size k ≥ 2 (denoted as (cid:101) B k ) is a butterﬂy factor that is also orthogonal.
Deﬁnition K.2 (Analog of Deﬁnition 2.3) .
An orthogonal butterﬂy matrix of size n (denoted as (cid:101) B ( n ) ) is a butterﬂy matrix with all butterﬂy factor matrices being orthogonal.
Note that the above deﬁnition implies that an orthogonal butterﬂy matrix, as well as its conjugate transpose, is orthogonal.
The orthogonal hierarchy deﬁnition nearly mimics the original hierarchy Deﬁnition 2.4, as follows:
Deﬁnition K.3.
• We say that an n × n matrix M ∈ (cid:101) B if we can express M = (cid:101) B ( n ) .
• We say that an n × n matrix M ∈ (cid:101) B ∗ if we can express M = (cid:104)(cid:101) B ( n ) (cid:105) ∗ .
• We say that an n × n matrix M ∈ OBB if we can express M = M 1 DM 2 for some M 1 ∈ (cid:101) B , M 2 ∈ (cid:101) B ∗ , and diagonal matrix D .
Note that D need not be full rank.
• Width w and expansion e in ( OBB ) we mimic the same deﬁnition as in the original hierarchy, using OBB instead of BB ∗ , such that E ∈ ( OBB ) w .
By padding if necessary, we will assume that n is a power of 2 .
K. E XPRESSIVITY
In this subsection we prove that all orthogonal (resp.
unitary) matrices are contained in OBB n .
To do this, we consider the class of Householder reﬂections , given by I − 2 uu ∗ for any unit vector u (Householder, 1958):
Lemma K.1.
All Householder reﬂections are in OBB with inner diagonal matrix I .
We will prove this lemma shortly.
First, we use this lemma to present a decomposition for all orthogonal (resp.
unitary) matrices.
Lemma K.2.
Let M be an n × n orthogonal/unitary matrix.
Then M ∈ ( OBB ) n − 1 .
Proof.
We consider the QR decomposition of M .
It is known that we can compose M into a product of n − 1 Householder reﬂections and an orthogonal/unitary diagonal matrix (Householder, 1958).
From Lemma K.1, each Householder reﬂection is in OBB .
To complete the proof, we argue that R can be folded into the rightmost butterﬂy matrix.
Let Q 1 be the rightmost butterﬂy factor matrix in Q ( ∈ (cid:101) B ( n ) n ).
Right multiplication of Q 1 by R scales each columns of Q 1 by some c ∈ C with || c ||
= 1 ( R is unitary diagonal).
This preserves both the sparsity 9 Q is the (orthogonal/unitary) product of n − 1 Householder reﬂections. R , the remaining upper triangular matrix after performing these reﬂections, is itself orthogonal/unitary, and therefore diagonal. pattern of Q 1 and the orthogonality of its columns.
Moreover, the norm of each column of Q 1 R is 1.
Therefore, Q 1 R is an orthogonal butterﬂy factor matrix, so M = QR ∈ ( OBB ) n − 1 , as desired.
We now return to the proof of Lemma K.
Proof of Lemma K.1.
Given u ∈ C n ( n a power of 2 ), let u 0 = u
[: n/ 2] ∈ C n/ 2 , u 1 = u
[ n/ 2 :] ∈
C n/ 2 denote the ﬁrst and second halves of u .
To show that H ∈ OBB with inner diagonal matrix I , we proceed by induction.
The base case for n = 2 is trivial.
It sufﬁces to show that there exist unitary butterﬂy factors L , R such that LHR has the form
[Equation]
for some unit vectors v 0 , v 1 ∈ C n/
[Equation]
Deﬁne
[Equation]
otherwise .
It is easily checked that
[Equation]
We choose
[Equation]
Diag ( v ∗ 1 ) Diag ( v 1 ) Diag ( − v 0 ) (cid:21) and R =
L ∗ .
L , R are (permuted) direct sums of blocks of the form {+(cid:20) v 
+}
[ i ] ∗ v {+
 [ i ] ∗ v 
+} [ i ] − v
[ i ] (cid:21) , which are orthogonal by construction (via (5)).
Hence, L ∈ (cid:101) B ( n ) n and R ∈ ( (cid:101) B ∗ ) ( n ) n .
Further,
LHR = (cid:20) [Equation] [Equation] [Equation]
) (cid:21) (cid:18)
I
[Equation]
(cid:20)
[Equation]
u
[Equation]
(cid:21) (cid:20) uu
[Equation]
(cid:21)
[Equation]
(cid:19) (cid:20) Diag ( v
[Equation]
Diag ( v
[Equation]
Diag ( v 1 )
Diag ( − v 0 ) (cid:21) ∗ = I −(cid:20) Diag ( v ∗ 0 ) Diag ( v ∗ 1 )
Diag ( v 1 )
Diag ( − v 0 ) (cid:21) (cid:20) uu
[Equation]
(cid:21) (cid:20) uu(cid:21) ∗ (cid:20) Diag ( v ∗ 0 )
Diag ( v ∗ 1 )
Diag ( v 1 )
Diag ( − v 0 ) (cid:21) ∗ =
I
[Equation]
(cid:20)
[Equation]
(cid:21) (cid:124)
[Equation]
(cid:125)
[Equation] v [Equation]
[-(cid:26)(cid:112) | u 
-](cid:21) (cid:124) (cid:123)(cid:122)
[Equation] where ◦ denotes the Hadamard product. From (6)
[Equation]
Denoting the ﬁrst half of this vector by w 0 ∈ C n/ 2 , we have
[Equation]
where (cid:107) w 0 (cid:107)
= (cid:107) u (cid:107) 2 = 1 .
The result follows inductively.
As an immediate corollary, we can use Singular Value Decomposition to obtain a factorization for an arbitrary n × n matrix.
Corollary K.1.
Let M be an arbitrary n × n matrix.
Then, M ∈ ( OBB ) 2 n − 1 , where all but one matrix in the decomposition is orthogonal (unitary).
Proof.
By employing Singular Value Decomposition, we can decompose M as M = UΣV ∗ , where
U , V ∗ are orthogonal and Σ is diagonal.
By Lemma K.2, U , V ∗ ∈ ( OBB ) n − 1 , and trivially Σ ∈ OBB .
Hence, M ∈ ( OBB ) 2 n − 1 .
Note that Σ is the only matrix in the decomposition that is not orthogonal (unitary).
K. C ONSTRUCTIONS
We show that we can construct s -sparse matrices in the OBB hierarchy with the same width as the
BB ∗ hierarchy.
The proof follows a structure to that of Theorem 3.
We begin by arguing about permutation and step matrices, then using the same factorization to argue that matrices with at most n NNZ are contained in ( BB ∗ ) 5 .
Then, we will appeal to a modiﬁed sum closure lemma to extend the argument to matrices of general s NNZ.
Similar to Appendix F, we can use these results to place all matrices with low-depth circuits for matrix vector multiplication in the OBB hierarchy.
K.3. P ERMUTATIONS
We begin by presenting the argument that permutations are included in OBB as a corollary to
Theorem 2.
Corollary K.2.
Let P be a permutation matrix.
Then P ∈ (cid:101) B (cid:101) B ∗ .
Proof.
We appeal to the decomposition from Theorem 2, noting that all butterﬂy factor matrices constructed in the proofs of Lemmas G.3 and G.1 are permutation matrices, and thus are orthogonal.
Hence, P ∈ OBB where the inner diagonal matrix is I .
To prove the containment of sparse matrices within the OBB hierarchy, we make use of the following lemma.
Lemma K.3.
Let P be a permutation matrix and D a diagonal matrix.
Then there exist diagonal matrices D (cid:48) and D (cid:48)(cid:48) such that:
[Equation]
Proof.
Let σ be the permutation such that P [ i, j ] = δ i,σ ( j ) .
Deﬁne D (cid:48) such that D (cid:48)
[ σ ( j ) , σ ( j )]
= D [ j, j ] .
Then, if i = σ ( j ) : ( PD )[ i, j ] =
P
[ i, j ] D [ j, j ]
= D (cid:48)
[ σ ( j ) , σ ( j )] P [ σ ( j ) , j ] =
( D (cid:48) P )[ σ ( j ) , j ] =
( D (cid:48) P )
[ i, j ] .
Otherwise, if i (cid:54) =
σ ( j ) , then ( PD )
[ i, j ] =
= ( D (cid:48) P )
[ i, j ] .
Hence, PD = D (cid:48) P .
Deﬁne D (cid:48)(cid:48) such that D (cid:48)(cid:48)
[ j, j ]
=
D [ σ ( j ) , σ ( j )] .
An analogous argument to above shows that
DP = PD (cid:48)(cid:48) .
K.3. S TEP M ATRICES
In the BB ∗ hierarchy (Lemma I.2), we were able to show that horizontal step matrices are butterﬂy matrices.
Here, we present a similar result for the OBB hierarchy.
Lemma K.4.
Let H be an n × n horizontal step matrix.
Then we can decompose H =
DO , where
D is a diagonal matrix and O ∈ (cid:101) B .
Proof.
Throughout the proof, we make reference to the original horizontal step matrix construction given in Lemma I.2 and its proof.
To begin, we show that an arbitrary 2 k × 2 k butterﬂy factor H 2 k in the decomposition of H can be expressed as the product of a diagonal matrix and an orthogonal butterﬂy factor.
Since a butterﬂy factor is direct sum of 2 × 2 matrices, there is a permutation matrix P 2 k such that conjugation of
H 2 k by P 2 k gives a block diagonal matrix H (cid:48)
2 k of n 2 2 × 2 matrices, i.e.
[Equation]
(See Figure 10 for an illustration.)
Speciﬁcally, P 2 k is the permutation where:
[Equation]
P
H
P ∗=
H (cid:48)
We argue that each of these 2 × 2 blocks can be decomposed into a diagonal matrix times an orthogonal matrix.
Note that the butterﬂy factor matrices constructed in the proof of Lemma I.2 each have at most one non-zero entry per column.
Hence, there are 4 cases to consider.
Note that matrices with at most one non-zero entry are exhausted by Cases 1 and 2.
Case 1: (cid:20) a b (cid:21) =
(cid:20) a b (cid:21) (cid:124) (cid:123)(cid:122) (cid:125) D
(cid:20) 10(cid:21) (cid:124) (cid:123)(cid:122) (cid:125)
O
Case 2: (cid:20)
a b (cid:21) =
(cid:20) a b (cid:21) (cid:124) (cid:123)(cid:122) (cid:125) D
(cid:20) 01(cid:21)
(cid:124) (cid:123)(cid:122) (cid:125)
O
Case 3: (cid:20) a b  (cid:21) =
(cid:20) √ a 2 + b   (cid:21) (cid:124) (cid:123)(cid:122) (cid:125) D (cid:34) a √ a 2 + bb √ a 2 + b 2 b √ a 2 + b− a √ a 2 + b(cid:35) (cid:124) (cid:123)(cid:122) (cid:125)
O , a,
[Equation]
In the last two cases, O is a 2 × 2 rotation matrix, which is commonly known to be orthogonal.
Assume that we perform the above decomposition on all of the blocks of H (cid:48) 2 k in parallel, therefore expressing H (cid:48) 2 k = D (cid:48)
O (cid:48) .
We now have
[Equation]
By Lemma K.3, we can rewrite this as
[Equation]
Note that P ∗ 2 k O (cid:48)
P 2 k is the product of three orthogonal matrices, and thus orthogonal.
Additionally, the construction of P 2 k ensures that P ∗ 2 k O (cid:48)
P 2 k is butterﬂy factor.
Hence, H 2 k can be expressed as the product of a diagonal matrix and an orthogonal butterﬂy factor, as desired.
Now, we show that this decomposition of butterﬂy factors implies Lemma K.4.
By performing this decomposition in parallel on each butterﬂy factor, we conclude that any butterﬂy factor matrix H ( n ) 2 k of H can be decomposed as H ( n ) 2 k
= D
2 k O ( n ) 2 k .
[Table]
We complete the argument by induction on n .
The base case n = 2 holds by the observation about butterﬂy factor matrices above.
Assume that any horizontal step matrix of size n 2 × n 2 can be expressed as a diagonal matrix times an orthogonal butterﬂy matrix.
Now, consider the n × n horizontal step matrix H .
From Lemma I.2, H can be expressed as
[Equation]
where H 1 , H 2 are n 2 × n 2 horizontal step matrices.
By our inductive hypothesis,
[Equation]
, where D 1 is diagonal and O 1 , O 2 are n 2 × n 2 matrices in (cid:101) B .
However, B ( n ) n D 1 is a butterﬂy factor, and therefore can be expressed as D n O ( n ) n .
Therefore,
[Equation]
with O ∈ (cid:101) B , as desired.
Just as with the BB ∗ hierarchy, the decomposition of vertical step matrices falls out as an immediate corollary to the horizontal step matrix proof.
Corollary K.3.
Let V be a vertical step matrix.
Then we can decompose V
=
O ∗ D , where D is a diagonal matrix and O ∗ ∈ (cid:101) B ∗ .
K.3. S PARSE M ATRICES
Now that we have argued about the decomposition of permutation and step matrices in the OBB hierarchy, we can leverage the construction from Lemma I.1 to argue about matrices with at most n
NNZ.
Corollary K.4.
Let S be an n × n matrix with at most n NNZ.
Then, S ∈ ( OBB ) 5 .
Proof.
We use the construction from Lemma I.1, along with Lemma K.4 and Corollary K.3, to express S as:
[Equation]
(cid:123)(cid:122) (cid:125) P
D
[Equation]
O(cid:124) (cid:123)(cid:122)
[Equation]
V
O
[Equation]
(cid:123)(cid:122) (cid:125) P, with each O i ∈ (cid:101) B , each O (cid:48)
j ∈ (cid:101) B ∗ , and each D k diagonal.
Noting that O (cid:48) 1 and O 5 are permutations, we make use of Lemma K.3 to re-express S as:
[Equation]
Note that each M (cid:96) ∈ OBB .
Hence, S ∈ ( OBB ) 5 , as desired.
Just as in Appendix I, we would like to extend this orthogonal-based construction to capture matrices of general sparsity.
To accomplish this, we introduce an addition closure lemma analogous to
Lemma K.5 for the OBB hierarchy.
Lemma K.5.
Let A 1 , . . .
A m be k × k matrices in ( OBB ) we then (cid:80) mi =1 A i ∈ ( OBB ) mw 4 e .
With Lemma K.5, we arrive at the following Corollary on general orthogonal sparsity.
Corollary K.5.
Let S be an n × n matrix with s NNZ.
Then, S ∈ ( OBB ) 5 (cid:100) sn (cid:101).
Proof.
Just as in the proof of Theorem 3, we accomplish this using a sum of (cid:6)
sn (cid:7) matrices of at most n NNZ.
For handling the sum of matrices, we need to appeal to Lemma K.5.
To conclude the argument, we give the proof of Lemma K.5.
Proof of Lemma K.5.
For each 1 ≤ i ≤ m , let E A i ∈
F ek × ek be deﬁned such that A i = SE
A i S ∗ (with S as in Deﬁnition 2.4).
Note that E A i ∈ ( OBB ) w .
Consider matrices of the form:  
[Equation]
 (cid:124)
[Equation]
[Table]
(cid:125)
[Table]
Note that K , a block diagonal matrix composed of matrices in ( OBB ) w , is itself in ( OBB ) w since
K =
[Equation]
(cid:89)
[Equation]
 
[Equation]
[Table] [Table]
[Equation]
[Table]
 (cid:124) (cid:123)(cid:122) (cid:125)
L j ∈ (cid:101) B  
I ek I ek
D
j     
 (cid:124) (cid:123)(cid:122) (cid:125)
Diagonal  
I ek I ek
O (cid:48)
j
I ek 
 (cid:124) (cid:123)(cid:122) (cid:125)
R j ∈ (cid:101) B ∗ , where each O j is a ek × ek matrix in (cid:101) B , and each O (cid:48)
j is a ek × ek matrix in (cid:101) B ∗ .
L w (the leftmost factor) is a block diagonal matrix composed of 4 ek × ek matrices in (cid:101) B .
Therefore, we can fold O into this factor (since a butterﬂy factor in (cid:101) B (4 ek ) 4 ek was not yet used in L w ) to conclude that OL w ∈ (cid:101) B .
Similarly, since no btterﬂy factor from (cid:101) B (4 ek ) 2 ek has been used in R 1 , we may fold P into R 1 to conclude that R 1 P ∈ (cid:101) B ∗ .
Finally, we address the scalar multiple of √ 2 by multiplying all entries of any diagonal matrix in the decomposition of K by √ 2 .
Hence, we may conclude that M i ∈ ( OBB ) w .
Through repeated application ( m times) of the identity  
I
A0 B
I I
A0 B
I  
  
I 0
C
I
I 0
C
I  
 =  
I
A 1 + B0 C
I I
[Equation] [Equation]
we see that
[Equation]
M ∈ F
[Equation]
Therefore, M ∈ ( OBB ) mw .
Next, we note that
[Equation]
We would like to show that we can fold Q into the rightmost OBB factor of M .
The rightmost matrix in the decomposition of M is P .
Note that
[Equation]
Just as earlier, the factor of √ 2 can be multiplied through any diagonal matrix.
Also, these two orthogonal butterﬂy factor matrices can be folded into the the rightmost R matrix (the decomposition of K above does not use these two, rightmost butterﬂy factors).
Hence, (cid:80) mi =1 A i ∈ ( OBB ) mw 4 e , as desired.
K.3. A RITHMETIC C IRCUITS
Just as in Theorem 1, we can use the sparsity result in Lemma K.5 to place matrices with low-depth (linear) arithmetic circuits for matrix vector multiplication in the OBB hierarchy.
Corollary K.6.
Let M be an n × n matrix such that matrix-vector multiplication of M times an arbitrary vector v can be represented as a be a linear arithmetic circuit C comprised of s gates (including inputs) and having depth d .
Then, M ∈ ( OBB ) O
( d ) O ( sn ) .
Proof.
We use the construction given in the proof of Theorem 1.
Corollaries K.4 and K.2 allow us to recover the same width and expansion factor with the OBB hierarchy.
L R E LU NETWORK WITH STRUCTURED WEIGHT MATRICES
We show that for any neural network with ReLU nonlinearities and whose weight matrices have arithmetic circuits with few gates, its linear network counterpart (obtained by removing all the
ReLU’s) also has an arithmetic circuit with not too many more gates.
This implies that in trying to ﬁnd the smallest arithmetic circuit augmented with ReLU gates to represent a ReLU network, one might as well try to ﬁnd the smallest arithmetic circuits that represent the matrix-vector multiplication of each weight matrix.
Proposition 2.
Consider a neural network architecture consisting of L layers with weight matrices
W 1 , . .
, W L ∈ F n × n and ReLU nonlinearity in between.
Suppose that matrix-vector multiplication of W i times an arbitrary vector v can be represented as a linear arithmetic circuit with s i gates (including inputs).
Then there exists an arithmetic circuit augmented with ReLU gates with (cid:80) Li
=
1 s i + Ln total gates that computes the output
ReLU( W L ( . . .
ReLU( W 1 v ))) of the network for an arbitrary input vector v .
Conversely, if there is an arithmetic circuit augmented with ReLU gates with s total gates that computes all the activations of the network ReLU( W 1 v ) , . . . , ReLU( W L . . .
ReLU( W 1 v )) for an arbitrary input v , then there exists an arithmetic circuit augmented with ReLU gates with 2 s + 2 Ln total gates that computes the activations of the network without ReLU W 1 v , . . .
, W L . . .
W 1 v .
Proof of Proposition 2.
To compute the output of the network ReLU( W L ( . . .
ReLU( W 1 v ))) , we ﬁrst compute the matrix-vector product W 1 v with an arithmetic circuit of s 1 gates by assumption, and use n other ReLU gates to compute the pointwise ReLU.
Then we repeat the process for layer 2 , 3 , . . .
, L , using the arithmetic circuits of W 1 , . . .
, W L and Ln additional gates for ReLU.
In total we obtain an arithmetic circuit augmented with ReLU gates with (cid:80) Li =
1 s i + Ln total gates.
Conversely, to build an arithmetic circuit augmented with ReLU gates to compute
W 1 v , . . .
, W L . . .
W 1 v , we pass v and then − v through the circuit that computes ReLU( W 1 x ) for an arbitrary x to get ReLU( W 1 v ) and ReLU( − W 1 v ) .
Noting that x = ReLU( x ) − ReLU( − x ) , we can use n additional gates to compute W 1 v from ReLU( W 1 v ) and ReLU( − W 1 v ) .
Repeat the process for layer 2 , 3 , . . .
, L (for example, pass W 1 v and − W 1 v to the circuit that computes W 2 x for an arbitrary x on layer 2).
Overall we need to double the circuits that computes all the activations of the network ReLU( W 1 v ) , . .
, ReLU( W L . . .
ReLU( W 1 v )) , requiring 2 s gates.
We also need n additional gates per layer to compute the negation of the input to that layer (e.g. computing − v from v ), and n additional gates per layer to subtract the output of the ReLU circuit (e.g. computing W 1 v from ReLU( W 1 v ) and ReLU( − W 1 v ) .)
Therefore we can construct an arithmetic circuit augmented with ReLU gates with 2 s + 2 L total gates that computes the activations of the network without ReLU W 1 v , . . .
, W L . . .
W 1 v .
We now prove an asymptotic bound on the VC dimension of a ReLU network whose weight matrices are kaleidoscope matrices with bounded width and expansion.
Proposition 3.
Let F be the class of ReLU neural networks consisting of L layers, where each layer is a K-matrix with width and expansion bounded by some constant C .
Suppose that the network has W total parameters.
Let sign F denote the corresponding classiﬁcation functions: { x (cid:55)→ sign f ( x ) : f ∈ F} .
Then this class has VC dimension:
[Equation]
We leverage the result from Thomas et al.
(2018) for the case where the entries of the weight matrices interact multiplicatively, but with polynomially bounded degrees.
This proof is similar to the VC bound for ReLU networks whose weight matrices are butterﬂy matrices (Dao et al., 2019).
Proof.
To use Theorem 3 of Thomas et al.
(2018), we simply need to check that the entries of the linear layer, as polynomials of the parameters, has degree at most c 1 m c 2 l for some universal constant c 1 , c 2 > 0 , where m l is the size of output of the l -th layer.
If the network weight matrices are
K-matrices with bounded width and expansion, each weight matrix is a product of at most c 3 log m l sparse factors, for some universal constant c 3 > 0 .
This means that the degree is polynomially bounded, which satisﬁes the condition of the theorem.
Therefore the VC dimension is bounded to be almost linear in the number of parameters:
[Equation]
M A RITHMETIC C IRCUIT P RIMER
We give a quick overview of arithmetic circuits.
This is a model of computation that has been studied for numerous computational problems (and is the basic model for algebraic complexity theory ).
For our purposes, we will exclusively focus on arithmetic circuits for the matrix-vector multiplication problem.
For a more detailed exposition, the reader is referred to the standard book on this topic (Bürgisser et al., 2013).
Deﬁnition M.1 (Arithmetic Circuits) .
An arithmetic circuit that computes y
=
Ax (for A ∈ F m × n ) has n input gates (corresponding to x [0] , . . .
, x [ n − 1] ) and m output gates (corresponding to y [0] , . . .
, y [ m − 1] ).
All the internal gates correspond to addition, subtraction, multiplication and division 12 over the underlying ﬁeld F .
The circuit is also allowed to use constants from F for ‘free.’
The deﬁnition of the internal gates can depend on A (as well as x of course).
In other words, one can ‘bake’ the knowledge about A into the circuit.
The size s of a circuit is n plus the number of addition, multiplication, subtraction and division gates used in the circuit.
The depth d of a circuit is the minimum number of layers such that all gates in a given layer take as its input gates from previous layers.
One drawback of arithmetic circuits (especially for inﬁnite ﬁelds e.g. F = R , which is our preferred choice in this work) is that they assume operations over F can be performed exactly .
In particular, it ignores precision issues involved with real arithmetic.
Nonetheless, this model turns out to be a very useful model in reasoning about the complexity of doing matrix-vector multiplication for any family of matrices.
Perhaps the strongest argument in support of arithmetic circuits is that a large (if not an overwhelming) majority of matrix-vector multiplication algorithm also imply an arithmetic circuit of size comparable to the runtime of the algorithm (and the depth of the circuit roughly correponds to the time taken to compute it by a parallel algorithm).
For example consider the obvious algorithm to compute Ax (i.e. for each i ∈ [ m ] , compute y
[ i ] as the sum (cid:80) n − 1 i =
A [ i, j ] x [ j ] ).
It is easy to see that this algorithm implies an arithmetic circuit of size O ( nm ) and depth O (log n ) .
[Table]
[Table] [Table]
[Table] [Table]
[Table]
One thing to note about the arithmetic circuit above is that all the multiplications involve at least one input that is a constant from F (recall that we can assume that the entries of A are constants that can be used to build the circuit).
This leads to the following important sub-class of arithmetic circuits:
Deﬁnition M.2 (Linear Arithmetic Circuits) .
An arithmetic circuit is called a linear arithmetic circuit if it only uses addition, subtraction and multiplication.
Further, every multiplcation has a ﬁxed constant from F as at least one of its two inputs.
In other words, all gates in the circuit are linear functions of their inputs (i.e. of the form ax + by for ﬁxed constants a, b ∈ F ).
Intuitively for the matrix-vector multiplication, it makes sense to consider linear arithmetic circuits since the ﬁnal function we want to compute Ax is indeed a linear function of its inputs.
For iniﬁnite ﬁelds (e.g. F = R or F = C ), it turns out that this is essentially without loss of generality:
Theorem 4 ((Bürgisser et al., 2013)) .
Let F be an inﬁnite ﬁeld.
Any (general) arithmetic circuit to compute Ax over F of size s and depth d can be converted into a linear arithmetic circuit of size
O ( s ) and depth O ( d ) .
The above result implies that for asymptotic considerations, linear arithmetic circuits for matrix-vector multiplication are equivalent to general arithmetic circuits.
One important property of linear arithmetic circuits of depth d , which we will use in our arguments, is that such a circuit can be equivalently represented as product of d sparse matrices (see the proof of
Theorem 1 for the precise derivation 16 ).
As mentioned earlier, a vast majority of efﬁcient matrix vector multiplication algorithms are equivalent to small (both in size and depth) linear arithmetic circuit.
For example the FFT can be thought of as an efﬁcient arithmetic circuit to compute the Discrete Fourier Transform (indeed when one converts the linear arithmetic circuit for FFT into a matrix decomposition, 17 then each matrix in the decomposition is a butterﬂy factor, with each block matrix in each factor being the same).
For an illustration of this consider the DFT with n = 4 as illustrated in Figure 11.
- ii - -
i -- i
Figure 12 represent the arithmetic circuit corresponding to FFT with n = 4 .
[Equation] [Equation]
+ + + + + +   1− −   1−− i i wwww+ x a y b ax + by
Semanticsofagate
Finally, Figure 13 is representation of the arithmetic circuit of Figure 12 as a product of a butterﬂy matrix and (the bit-reversal) permutation.
We note that our generic arithmetic circuit to decomposition into BB ∗ is not as tight as in Figure 13.
- ii
B (4)    -   -
B (4)
P
One reason for the vast majority of existing efﬁcient matrix vector algorithms leading to (linear) arithmetic circuits is that they generally are divide and conquer algorithms that use polynomial operations such as polynomial multiplication or evaluation (both of which themselves are divide and conquer algorithms that use FFT as a blackbox) or polynomial addition.
Each of these pieces are well known to have small (depth and size) linear arithmetic circuits (since FFT has these properties).
Finally, the divide and conquer structure of the algorithms leads to the circuit being of low depth.
See the book of Pan (Pan, 2001) for a more elaborate description of this connection.
In fact, the recent work of De Sa et al.
(De Sa et al., 2018) makes this fact explicit and presents the most general known structure on matrices that imply near-linear size linear arithmetic circuits for the corresponding matrix vector multiplication.
Their work combines two separate classes of structures matrices– orthogonal polynomial transforms (Driscoll et al., 1997; Szegö, 1967) as well as matrices with low displacement rank (Kailath et al., 1979; Olshevsky & Shokrollahi, 2000)– and presents a linear class of linear arithmetic circuits to solve their matrix vector multiplication problem.
We note that structured matrices with low displacement rank have been used to replace fully connected layers in some neural network architectures (Sainath et al., 2013; Thomas et al., 2018).