PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations

Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, Christopher Ré

arXiv:1903.05895v2cs.LGstat.ML

TL;DR

Hand-crafted fast transforms and platform-specific implementations motivate the search for methods that can learn efficient algorithms with limited prior knowledge. The paper uses butterfly factorizations to parameterize recursive transforms, recovering several known transforms to machine precision and reporting higher CIFAR-10 accuracy with fewer parameters in a compressed network.

  • Problem

    Fast transforms require carefully designed algorithms and specialized implementations, motivating the question of how to learn them for realistic sizes with limited prior knowledge.

    Source: §1 · p. 1

  • Method

    The method parameterizes recursive fast algorithms using butterfly factorizations, a relaxed product of special block-diagonal matrices.

    Source: §1 · p. 2, §5 · p. 11

  • Results

    DFT, Hadamard, DCT, and convolution are recovered to machine precision through N = 1024; on CIFAR-10, the single hidden-layer model exceeds a fully connected baseline by 3.9 points with 40X fewer parameters.

    Source: §1 · p. 2, §5 · p. 11

  • Takeaways & Limitations

    Butterfly factorizations provide O(N)-parameter transforms with O(N log N) multiplication and can serve as components of end-to-end machine-learning models.

    Source: §1 · p. 2, §5 · p. 11

    Orthogonal-polynomial transforms such as the DLT use a more general divide-and-conquer decomposition, and the butterfly parameterization does not perfectly capture the DLT.

    Source: §4.1 · pp. 8–9

Abstract

from arXiv · show

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

1 Introduction

The paper aims to learn fast transforms and their implementations with limited prior knowledge, addressing the need for hand-designed algorithms and platform-specific code. It introduces butterfly factorizations and evaluates transform recovery, neural-network performance, and implementation speed.

1 Introduction

  • The motivating challenge is to learn effective transforms and efficient implementations without hand-crafting each algorithm and platform-specific implementation.The desired approach should recover fast transforms at realistic sizes with limited prior knowledge and use differentiable primitives and common operations.Source: §1 · p. 1
  • Butterfly factorizations parameterize recursive algorithms as products of special block-diagonal matrices, with O(N) parameters and O(N log N) multiplication.The representation uses a relaxed butterfly structure that can learn from data and capture a larger class of structures than specific prior butterfly uses.Source: §1 · p. 2
  • 3.9 points higher CIFAR-10 classification accuracy than a fully connected baseline, with 40X fewer parameters, is reported for the single hidden-layer benchmark.The paper also reports a 0.43-point accuracy improvement from adding a lightweight butterfly layer to a modern ResNet architecture.Source: §1 · p. 2
  • DFT, Hadamard, DCT, and convolution are recovered to machine precision at dimensions up to N = 1024, unlike standard sparse and low-rank baselines.The method also recovers the FFT algorithm to machine precision, as described in the paper’s abstract.Source: §1 · p. 2
  • The generic implementation runs within 3–5X of specialized DFT and DCT implementations while representing a richer class of transforms.The paper compares training and inference speed with specialized implementations of discrete transforms.Source: §1 · p. 2

2 Related Work

Related work includes established structured matrix classes and efforts to learn fast algorithms or compress neural networks. The paper distinguishes its approach by learning recursive sparse structure rather than sparsity alone or searching a large discrete space.

2 Related Work

  • Structured matrix classes include Toeplitz, Hankel, Vandermonde, and Cauchy matrices [35]; low displacement rank and displacement structure generalized these families.Orthogonal-polynomial transforms also generalize the DFT and DCT and are used in areas including differential equations and optics.Source: §2 · p. 2
  • De Sa et al. generalized these structured classes, many of which have recursive structure and superfast algorithms; the characterization is also known through arithmetic circuits.The related-work discussion connects these classes to fast matrix multiplication.Source: §2 · pp. 2–3
  • Unlike sparse matrix factorization focused on factor sparsity, this approach uses recursive transform structure to avoid the discreteness problem inherent in learning sparsity.Prior work considered sparse PCA, dictionary learning, products of random sparse matrices, and multi-layer sparse approximations.Source: §2 · p. 3
  • Earlier fast-transform learning methods searched combinatorially large discrete spaces and were limited to matrices up to 8 × 8; this approach learns simpler permutations for realistic dimensions.Differentiable approaches to general algorithms such as matching, sorting, and traveling salesman were described as limited to problems of size at most 100.Source: §2 · p. 3
  • Compressed-model approaches replace unconstrained weight matrices with structured classes; examples use Fourier-related matrices [40] or specialized recursive algorithms [42].The paper presents its own linear-parameter, efficiently multiplied matrix class as a drop-in replacement for matrices in end-to-end models.Source: §2 · p. 3

3 Recovering Fast Transforms

The section connects fast matrix-vector multiplication to sparse matrix products and develops recursive butterfly factorizations with learnable permutations. It also characterizes the transforms and matrix classes these parameterizations can represent.

3.1 Preliminaries

  • Sparse matrix products provide fast multiplication, and sparse product width corresponds to the shortest linear straight-line program; the DFT’s even/odd recursion yields the Cooley-Tukey FFT.The DFT recursion applies transforms to the even and odd indices before combining their results.Source: §3.1 · pp. 3–4

3.2 Recovering Fast Transform Algorithms

  • The transform factorization T_N = B^(N)P^(N) uses a butterfly matrix and a product of log2(N) block permutations; composing two such modules yields BPBP, which captures one-dimensional convolutions.The paper relates convolution to an FFT, a component-wise product, and an inverse FFT.Source: §3.2 · p. 5
  • At each recursive step, the method models structured permutation choices as a categorical distribution, learned through three logits passed through sigmoid functions.The allowed operations include keeping halves intact or separating even and odd indices, with optional reversals of either half.Source: §3.2 · pp. 5–6
  • Initialization targets near-unitary or orthogonal butterfly factors to preserve input and output magnitudes and avoid exploding or vanishing gradients across products.The passage gives N(0, 1/2) entries as an example for the real case.Source: §3.2 · p. 6
  • Compared with earlier butterfly parametrizations, this method learns a permutation, does not enforce orthogonality, orders interactions from nearby elements outward, and ties blocks within factors.These choices give the parametrization fewer parameters and a tighter recursive interpretation than some prior schemes.Source: §3.2 · p. 6

3.3 Expressivity and the butterfly hierarchy

  • The BP hierarchy ranges from structured matrices with linearly many parameters to all square matrices, and can represent all matrices with linear depth.The hierarchy is built by composing BP modules.Source: §3.3 · p. 7
  • BP captures FFT, Hadamard, and their inverses exactly; BP squared captures DCT, DST, and convolution, while BP^(4N+10) contains all N × N matrices.The section also states that the hierarchy can represent all matrices with linear depth.Source: §3.3 · p. 7
  • Networks with BP and BPBP layers have VC dimension almost linear in their parameter count, with a corresponding sample-complexity bound.The passage compares this property to fully connected and LDR networks.Source: §3.3 · p. 7

4 Empirical Evaluation

The evaluation tests whether butterfly parameterizations recover fast discrete transforms, improve neural-network compression, and retain competitive speed. Results include a limitation for transforms whose recursive structure is more general than the parameterization captures exactly.

4.1 Discrete Transforms

  • Butterfly recovers fast algorithms through N=1024 for transforms and N=512 for convolution, including FFT factorizations with bit-reversal and other exact permutations.Source: §4.1 · p. 8, Figure 3 · p. 9
  • The parameterization does not perfectly capture DLT, whose algorithms use a slightly more general divide-and-conquer decomposition, but reconstructs it better than baselines.Source: §4.1 · pp. 8–9

4.2 Neural Network Compression

  • BPBP achieves higher test accuracy than fully connected layers on all datasets, with the real version using more than 56X fewer parameters; complex BPBP performs even better with a slight parameter increase.Source: §4.2 · pp. 9–10, Table 1 · p. 10
  • On CIFAR-10 ResNet18, adding BPBP before the final FC/softmax improves performance more than adding FC, with a 0.07% increase in parameters.Source: Table 2 · p. 10, ResNet · p. 10

4.3 Training and Inference Speed Comparison

  • Training butterfly is 15% faster than GEMM and within 40% of FFT; inference is one or two orders faster than GEMV, within a factor of 5 of FFT and 3 of DCT/DST.Source: §4.3 · pp. 10–11, Figure 4 · p. 11

5 Conclusion

The paper concludes that recursive algorithms parameterized by butterfly factorizations can learn several fast transforms to machine precision at dimensions up to 1024. The appendix shows how this structure covers specific transforms and matrix classes, while noting an open representation problem for general orthogonal polynomial matrices.

5 Conclusion

  • Butterfly factorizations learn DFT, DCT, Hadamard, and convolution transforms to machine precision for dimensions up to N = 1024.The conclusion also reports consistent performance improvements, compression, and speed increases when the method is used in end-to-end ML models.Source: §5 · p. 11

A.1 Discrete Cosine Transform (DCT) Matrix

  • The DCT is computed by permuting even and reversed odd indices, taking an FFT, phase-scaling its outputs, and taking the real part.Its factorization uses the left butterfly-permutation component for the FFT and final scaling, with an identity right butterfly matrix.Source: §A.1 · p. 16

A Matrix Factorizations of Linear Transforms

  • Recursive algorithms are parameterized through butterfly factorizations, the basis of the paper’s transform representation.The appendix develops matrix factorizations for specific transforms using this structure.Source: §A.1 · p. 16

A.2 Discrete Sine Transform (DST) Matrix

  • The DST is computed from an FFT after permuting and reversing odd indices, negating the second half, phase-scaling, and taking the real part.Its factorization is described as a (BP)2 factorization with the left component performing the FFT and final scaling.Source: §A.2 · p. 17

A.3 Hadamard Matrix

  • For powers of 2, the Hadamard matrix has a recursive factorization with butterfly factors whose entries may be real.The factorization uses the identity permutation matrix.Source: §A.3 · p. 18

A.4 Convolution

  • Convolution by a fixed vector is represented by a circulant matrix and computed through DFT, elementwise multiplication, and inverse Fourier transform.The resulting factorization is (BP)2; the appendix also states that skew-circulant matrices lie in (BP)2.Source: §A.4 · p. 18

A.5 Toeplitz Matrices

  • Toeplitz matrices are embedded in a 2N × 2N circulant matrix, yielding a factorization contained in (BP)2.The inner circulant matrix is decomposed using the convolution factorization.Source: §A.5 · p. 19

A.6 Orthogonal Polynomial Matrices

  • The appendix leaves open whether general orthogonal polynomial matrices can be represented with butterfly matrices, while presenting an alternate sparse factorization.The recursive decomposition can be applied to the submatrices, and the discrete Legendre transform is stated to admit the resulting factorization.Source: §A.6 · pp. 19–21

B Proofs

For butterfly-layer networks with fixed permutations and piecewise linear activations, the VC dimension is bounded by O(LW log W). The proof applies a result for bounded-degree parameter interactions because the layer entries have degree at most 2 log2 n.

B.1 VC Dimension Bound for Neural Network with Butterfly Layers

  • VCdim(sign F) = O(LW log W) for networks with L butterfly layers and W total parameters.The proposition covers BP or BPBP parameterizations with fixed permutations and piecewise linear activations.Source: §B.1 · p. 21
  • The standard VC-dimension bound for fully connected layers does not directly apply because parameters within a butterfly layer interact multiplicatively.The proof instead uses a variant for cases where the degree of multiplicative interaction is not too high.Source: §B.1 · p. 21
  • Each BP or BPBP layer has parameter degree at most 2 log2 n, satisfying the condition used to bound VC dimension.The degree follows from each B(n) being a product of log2 n matrices.Source: §B.1 · p. 21

B.2 Proposition 1

Proposition 1 establishes that several transforms fit BP or BP2 factorizations and that arbitrary N × N matrices fit a longer BP factorization. The supplied table caption specifies an RMSE stopping threshold but gives no cell values.

B.2 Proposition 1

  • The DFT and Hadamard transform are included in (BP)1, while the DCT, DST, and convolution are included in (BP)2.The cited sections provide the respective factorizations.Source: §B.2 · p. 21
  • Table 4 reports RMSE for learned fast algorithms across common transforms and N values, stopping early when RMSE < 1e-4.The caption supplies the stopping threshold, not individual table values.Source: Table 4 · p. 22
  • Every N × N matrix can be expressed in (BP)^(4N+10) using a product of Toeplitz matrices, each contained in (BP)2.The construction appends the BP modules for the Toeplitz factors and uses a diagonal matrix to select the required upper-left entries.Source: §B.2 · p. 22

C.1 Recovering Fast Transforms

The experiment uses BP or BPBP parameterizations to recover fast algorithms for given transform matrices and reports RMSE across transforms and sizes. Hyperband tunes learning rate, initialization, and permutation-logit sharing.

C.1 Recovering Fast Transforms

  • BP or BPBP parameterizations are used to recover a fast algorithm from each transform’s matrix representation.The passage identifies this as the Section 4.1 procedure.Source: §C.1 · p. 22
  • RMSE is reported for different transforms and values of N.The supplied passages do not give the table’s individual values.Source: §C.1 · p. 22
  • Hyperband tunes learning rate, initialization, and whether logits in the permutation block are shared.The learning-rate search ranges from 0.0001 to 0.5.Source: §C.1 · p. 22

C.2 Fully connected network

The fully connected and ResNet experiments describe CIFAR-10 network setups and how model selection or training settings were handled.

C.2 Fully connected network

  • For the fully connected experiment, a grayscale CIFAR-10 input of size 1024 feeds a single N × N hidden layer and a fully connected softmax layer.Table 1 reports test accuracy for the model with the highest validation accuracy; batch size is 50, with a 15% validation set and momentum 0.9.Source: §C.2 · p. 22
  • For the fully connected experiment, learning rates tested for BP methods ranged from 0.005 to 0.2; values outside that range were ineffective.The listed rates were 0.005, 0.01, 0.02, 0.05, 0.1, and 0.2.Source: §C.2 · p. 22

C.3 Resnet

  • The ResNet18 experiment modifies the initial convolution to kernel size 3 and stride 1, and removes the first max-pool layer.Weight decay is 0.0002; training lasts 100 epochs, and Table 2 reports mean and standard deviation of test accuracies for hyperparameters with the highest average validation accuracy.Source: §C.3 · pp. 22–23

C.4 Speed Comparison

The section benchmarks butterfly-factorization training and inference against dense multiplication and specialized transform implementations, using different implementations and hardware for each.

C.4 Speed Comparison

  • Training compares the CUDA butterfly fast algorithm with cuBLAS GEMM and cuFFT, timing forward and backward passes at batch size 256 on a Tesla P100 GPU.The GPU has 16GB of memory.Source: §C.4 · p. 23
  • Inference compares a Python BP implementation, defined as B(N) times a fixed permutation P(N), with GEMV, FFT, DCT, and DST implementations in single-threaded mode on an Intel Xeon E5-2690 v4 CPU.The comparison uses NumPy for GEMV and FFT, and SciPy FFTPACK for DCT and DST.Source: §C.4 · p. 23
  • Figure 4 presents the speed-comparison results for the butterfly-factorization benchmarks.The section benchmarks both training and inference speed.Source: §C.4 · p. 23

D BP Hierarchy

The BP hierarchy is examined for expressiveness: the section establishes strict increases in its depth levels and states a broader expressiveness conjecture.

D BP Hierarchy

  • For every fixed c ≥ 1, some sufficiently large N admits an N × N matrix in (BP)c+1 but not in (BP)c.Thus, the hierarchy has strictly increasing levels for sufficiently large dimensions.Source: D BP Hierarchy · p. 23
  • The proof argues by contradiction: embedding an arbitrary N/2 × N/2 matrix as a block would yield a representation with too few parameters.The argument uses the assumed collapse of successive BP-factor counts and compares the resulting parameter count with the embedded matrix's N² entries.Source: D BP Hierarchy · p. 23
  • Conjecture 1 proposes that an N × N matrix computable by an arithmetic circuit of size N poly log(N) and depth poly log(N) belongs to (BP)poly log N.The passage presents this as a conjecture about the hierarchy's expressiveness.Source: D BP Hierarchy · p. 23

Related papers

Loading 1903.05895v2…