Sparse Linear Networks with a Fixed Butterfly Structure: Theory and Practice

1 INTRODUCTION . A butterfly network ( see Figure 6 in Appendix A ) is a layered graph connecting a layer of n inputs to a layer of n outputs withO ( log n ) layers , where each layer contains 2n edges . The edges connecting adjacent layers are organized in disjoint gadgets , each gadget connecting a pair of nodes in one layer with a corresponding pair in the next layer by a complete graph . The distance between pairs doubles from layer to layer . This network structure represents the execution graph of the Fast Fourier Transform ( FFT ) ( Cooley and Tukey , 1965 ) , Walsh-Hadamard transform , and many important transforms in signal processing that are known to have fast algorithms to compute matrix-vector products . Ailon and Chazelle ( 2009 ) showed how to use the Fourier ( or Hadamard ) transform to perform fast Euclidean dimensionality reduction with Johnson and Lindenstrauss ( 1984 ) guarantees . The resulting transformation , called Fast Johnson Lindenstrauss Transform ( FJLT ) , was improved in subsequent works ( Ailon and Liberty , 2009 ; Krahmer and Ward , 2011 ) . The common theme in this line of work is to define a fast randomized linear transformation that is composed of a random diagonal matrix , followed by a dense orthogonal transformation which can be represented via a butterfly network , followed by a random projection onto a subset of the coordinates ( this research is still active , see e.g . Jain et al . ( 2020 ) ) . In particular , an FJLT matrix can be represented ( explicitly ) by a butterfly network followed by projection onto a random subset of coordinates ( a truncation operator ) . We refer to such a representation as a truncated butterfly network ( see Section 4 ) . Simple Johnson-Lindenstrauss like arguments show that with high probability for any W ∈ Rn2×n1 and any x ∈ Rn1 , Wx is close to ( JT2 J2 ) W ( JT1 J1 ) x where J1 ∈ Rk1×n1 and J2 ∈ Rk2×n2 are both FJLT , and k1 = log n1 , k2 = log n2 ( see Section 4.2 for details ) . Motivated by this , we propose to replace a dense ( fully-connected ) linear layer of size n2 × n1 in any neural network by the following architecture : JT1 W ′J2 , where J1 , J2 can be represented by a truncated butterfly network and W ′ is a k2 × k1 dense linear layer . The clear advantages of such a strategy are : ( 1 ) almost all choices of the weights from a specific distribution , namely the one mimicking FJLT , preserve accuracy while reducing the number of parameters , and ( 2 ) the number of weights is nearly linear in the layer width of W ( the original matrix ) . Our empirical results demonstrate that this offers faster training and prediction in deployment while producing results that match and often outperform existing known architectures . Compressing neural networks by replacing linear layers with structured linear transforms that are expressed by fewer parameters have been studied extensively in the recent past . We compare our approach with these related works in Section 3 . Since the butterfly structure adds logarithmic depth to the architecture , it might pose optimization related issues . Moreover , the sparse structure of the matrices connecting the layers in a butterfly network defies the general theoretical analysis of convergence of deep linear networks . We take a small step towards understanding these issues by studying the optimization landscape of a encoder-decoder network ( two layer linear neural network ) , where the encoder layer is replaced by a truncated butterfly network followed by a dense linear layer in fewer parameters . This replacement is motivated by the result of Sarlós ( 2006 ) , related to fast randomized low-rank approximation of matrices using FJLT ( see Section 4.2 for details ) . We consider this replacement instead of the architecture consisting of two butterfly networks and a dense linear layer as proposed earlier , because it is easier to analyze theoretically . We also empirically demonstrate that our new network with fewer parameters performs as well as an encoder-decoder network . The encoder-decoder network computes the best low-rank approximation of the input matrix . It is well-known that with high probability a close to optimal low-rank approximation of a matrix is obtained by either pre-processing the matrix with an FJLT ( Sarlós , 2006 ) or a random sparse matrix structured as given in Clarkson and Woodruff ( 2009 ) and then computing the best low-rank approximation from the rows of the resulting matrix1 . A recent work by Indyk et al . ( 2019 ) studies this problem in the supervised setting , where they find the best pre-processing matrix structured as given in Clarkson and Woodruff ( 2009 ) from a sample of matrices ( instead of using a random sparse matrix ) . Since an FJLT can be represented by a truncated butterfly network , we emulate the setting of Indyk et al . ( 2019 ) but learn the pre-processing matrix structured as a truncated butterfly network . 2 OUR CONTRIBUTION AND POTENTIAL IMPACT . We provide an empirical report , together with a theoretical analysis to justify our main idea of using sparse linear layers with a fixed butterfly network in deep learning . Our findings indicate that this approach , which is well rooted in the theory of matrix approximation and optimization , can offer significant speedup and energy saving in deep learning applications . Additionally , we believe that this work would encourage more experiments and theoretical analysis to better understand the optimization and generalization of our proposed architecture ( see Future Work section ) . On the empirical side – The outcomes of the following experiments are reported : ( 1 ) In Section 6.1 , we replace a dense linear layer in the standard state-of-the-art networks , for both image and language data , with an architecture that constitutes the composition of ( a ) truncated butterfly network , ( b ) dense linear layer in smaller dimension , and ( c ) transposed truncated butterfly network ( see Section 4.2 ) . The structure parameters are chosen so as to keep the number of weights near linear ( instead of quadratic ) . ( 2 ) In Sections 6.2 and 6.3 , we train a linear encoder-decoder network in which the encoder is replaced by a truncated butterfly network followed by a dense linear layer in smaller dimension . These experiments support our theoretical result . The network structure parameters are chosen so as to keep the number of weights in the ( replaced ) encoder near linear in the input dimension . Our results ( also theoretically ) demonstrate that this has little to no effect on the performance compared to the standard encoder-decoder network . ( 3 ) In Section 7 , we learn the best pre-processing matrix structured as a truncated butterfly network to perform low-rank matrix approximation from a given sample of matrices . We compare our results 1The pre-processing matrix is multiplied from the left . to that of Indyk et al . ( 2019 ) , which learn the pre-processing matrix structured as given in Clarkson and Woodruff ( 2009 ) . On the theoretical side – The optimization landscape of linear neural networks with dense matrices have been studied by Baldi and Hornik ( 1989 ) , and Kawaguchi ( 2016 ) . The theoretical part of this work studies the optimization landscape of the linear encoder-decoder network in which the encoder is replaced by a truncated butterfly network followed by a dense linear layer in smaller dimension . We call such a network as the encoder-decoder butterfly network . We give an overview of our main result , Theorem 1 , here . Let X ∈ Rn×d and Y ∈ Rm×d be the data and output matrices respectively . Then the encoder-decoder butterfly network is given as Y = DEBX , where D ∈ Rm×k and E ∈ Rk× ` are dense layers , B is an ` × n truncated butterfly network ( product of log n sparse matrices ) and k ≤ ` ≤ m ≤ n ( see Section 5 ) . The objective is to learn D , E and B that minimizes ||Y − Y ||2F . Theorem 1 shows how the loss at the critical points of such a network depends on the eigenvalues of the matrix Σ = Y XTBT ( BXXTBT ) −1BXY T 2 . In comparison , the loss at the critical points of the encoder-decoder network ( without the butterfly network ) depends on the eigenvalues of the matrix Σ′ = Y XT ( XXT ) −1XY T ( Baldi and Hornik , 1989 ) . In particular , the loss depends on how the learned matrix B changes the eigenvalues of Σ′ . If we learn only for an optimal D and E , keeping B fixed ( as done in the experiment in Section 6.3 ) then it follows from Theorem 1 that every local minima is a global minima and that the loss at the local/global minima depends on howB changes the top k eigenvalues of Σ′ . This inference together with a result by Sarlós ( 2006 ) is used to give a worst-case guarantee in the special case when Y = X ( called auto-encoders that capture PCA ; see the below Theorem 1 ) . 3 RELATED WORK . Important transforms like discrete Fourier , discrete cosine , Hadamard and many more satisfy a property called complementary low-rank property , recently defined by Li et al . ( 2015 ) . For an n×n matrix satisfying this property related to approximation of specific sub-matrices by low-rank matrices , Michielssen and Boag ( 1996 ) and O ’ Neil et al . ( 2010 ) developed the butterfly algorithm to compute the product of such a matrix with a vector inO ( n log n ) time . The butterfly algorithm factorizes such a matrix into O ( log n ) many matrices , each with O ( n ) sparsity . In general , the butterfly algorithm has a pre-computation stage which requires O ( n2 ) time ( O ’ Neil et al. , 2010 ; Seljebotn , 2012 ) . With the objective of reducing the pre-computation cost Li et al . ( 2015 ) ; Li and Yang ( 2017 ) compute the butterfly factorization for an n × n matrix satisfying the complementary low-rank property in O ( n 3 2 ) time . This line of work does not learn butterfly representations for matrices or apply it in neural networks , and is incomparable to our work . A few works in the past have used deep learning models with structured matrices ( as hidden layers ) . Such structured matrices can be described using fewer parameters compared to a dense matrix , and hence a representation can be learned by optimizing over a fewer number of parameters . Examples of structured matrices used include low-rank matrices ( Denil et al. , 2013 ; Sainath et al. , 2013 ) , circulant matrices ( Cheng et al. , 2015 ; Ding et al. , 2017 ) , low-distortion projections ( Yang et al. , 2015 ) , Toeplitz like matrices ( Sindhwani et al. , 2015 ; Lu et al. , 2016 ; Ye et al. , 2018 ) , Fourier-related transforms ( Moczulski et al. , 2016 ) and matrices with low-displacement rank ( Thomas et al. , 2018 ) . Recently Alizadeh et al . ( 2020 ) demonstrated the benefits of replacing the pointwise convolutional layer in CNN ’ s by a butterfly network . Other works by Mocanu et al . ( 2018 ) ; Lee et al . ( 2019 ) ; Wang et al . ( 2020 ) ; Verdenius et al . ( 2020 ) consider a different approach to sparsify neural networks . The works closest to ours are by Yang et al . ( 2015 ) , Moczulski et al . ( 2016 ) , and Dao et al . ( 2020 ) and we make a comparison below . Yang et al . ( 2015 ) and Moczulski et al . ( 2016 ) attempt to replace dense linear layers with a stack of structured matrices , including a butterfly structure ( the Hadamard or the Cosine transform ) , but they do not place trainable weights on the edges of the butterfly structure as we do . Note that adding these trainable weights does not compromise the run time benefits in prediction , while adding to the expressiveness of the network in our case . Dao et al . ( 2020 ) replace handcrafted structured subnetworks in machine learning models by a kaleidoscope layer , which consists of compositions of butterfly matrices . This is motivated by the fact that the kaleidoscope hierarchy captures a structured matrix exactly and optimally in terms of multiplication operations required to perform the matrix 2At a critical point the gradient of the loss function with respect to the parameters in the network is zero . vector product operation . Their work differs from us as we propose to replace any dense linear layer in a neural network ( instead of a structured sub-network ) by the architecture proposed in Section 4.2 . Our approach is motivated by theoretical results which establish that this can be done with almost no loss in representation . Finally , Dao et al . ( 2019 ) show that butterfly representations of standard transformations like discrete Fourier , discrete cosine , Hadamard mentioned above can be learnt efficiently . They additionally show the following : a ) for the benchmark task of compressing a single hidden layer model they compare the network constituting of a composition of butterfly networks with the classification accuracy of a fully-connected linear layer and b ) in ResNet a butterfly sub-network is added to get an improved result . In comparison , our approach to replace a dense linear layer by the proposed architecture in Section 4.2 is motivated by well-known theoretical results as mentioned previously , and the results of the comprehensive list of experiments in Section 6.1 support our proposed method .

Summary: The paper studies “butterfly networks”, where, a logarithmic number of linear layers with sparse connections resembling the butterfly structure of the FFT algorithm, along with linear layers in smaller dimensions are used to approximate linear layers in larger dimensions. In general, the paper follows the idea of sketching to design new architectures that can reduce the number of trainable parameters.  In that regard, the paper is very appealing, as it shows that replacing linear layers with the butterfly networks does not result in any loss in performance.