Compressing Neural Networks: Towards Determining the Optimal Layer-wise Decomposition

1 Introduction Neural network compression entails taking an existing model and reducing its computational and memory footprint in order to enable the deployment of largescale networks in resource-constrained environments . Beyond inference time efficiency , compression can yield novel insights into the design ( Liu et al. , 2019b ) , training ( Liebenwein et al. , 2021a , b ) , and theoretical properties ( Arora et al. , 2018 ) of neural networks . Among existing compression techniques – which include quantization ( Wu et al. , 2016 ) , distillation ( Hinton et al. , 2015 ) , and pruning ( Han et al. , 2015 ) – low-rank compression aims at decomposing a layer ’ s weight tensor into a tuple of smaller low-rank tensors . Such compression techniques may build upon the rich literature on low-rank decomposition and its numerous applications outside deep learning such as dimensionality reduction ( Laparra et al. , 2015 ) or spec- ⇤denotes authors with equal contributions . Code : https : //github.com/lucaslie/torchprune 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . tral clustering ( Peng et al. , 2015 ) . Moreover , low-rank compression can be readily implemented in any machine learning framework by replacing the existing layer with a set of smaller layers without the need for , e.g. , sparse linear algebra support . Within deep learning , we encounter two related , yet distinct challenges when applying low-rank compression . On the one hand , each layer should be efficiently decomposed ( the “ local step ” ) and , on the other hand , we need to balance the amount of compression in each layer in order to achieve a desired overall compression ratio with minimal loss in the predictive power of the network ( the “ global step ” ) . While the “ local step “ , i.e. , designing the most efficient layer-wise decomposition method , has traditionally received lots of attention ( Denton et al. , 2014 ; Garipov et al. , 2016 ; Jaderberg et al. , 2014 ; Kim et al. , 2015b ; Lebedev et al. , 2015 ; Novikov et al. , 2015 ) , the “ global step ” has only recently been the focus of attention in research , e.g. , see the recent works of Alvarez and Salzmann ( 2017 ) ; Idelbayev and Carreira-Perpinán ( 2020 ) ; Xu et al . ( 2020 ) . In this paper , we set out to design a framework that simultaneously accounts for both the local and global step . Our proposed solution , termed Automatic Layer-wise Decomposition Selector ( ALDS ) , addresses this challenge by iteratively optimizing for each layer ’ s decomposition method ( local step ) and the low-rank compression itself while accounting for the maximum error incurred across layers ( global step ) . In Figure 1 , we show how ALDS outperforms existing approaches on the common ResNet18 ( ImageNet ) benchmark ( 60 % compression compared to ⇠20 % for baselines ) . Efficient layer-wise decomposition . Our framework relies on a straightforward SVD-based decomposition of each layer . Inspired by Denton et al . ( 2014 ) ; Idelbayev and Carreira-Perpinán ( 2020 ) ; Jaderberg et al . ( 2014 ) and others , we decompose each layer by first folding the weight tensor into a matrix before applying SVD and encoding the resulting pair of matrices as two separate layers . Enhanced decomposition via multiple subsets . A natural generalization of low-rank decomposition methods entails splitting the matrix into multiple subsets ( subspaces ) before compressing each subset individually . In the context of deep learning , this was investigated before for individual layers ( Denton et al. , 2014 ) , including embedding layers ( Chen et al. , 2018 ; Maalouf et al. , 2021 ) . We take this idea further and incorporate it into our layer-wise decomposition method as additional hyperparameter in terms of the number of subsets . Thus , our local step , i.e. , the layer-wise decomposition , constitutes of choosing the number of subsets ( k ` ) for each layer and the rank ( j ` ) . Towards a global solution for low-rank compression . We can describe the optimal solution for low-rank compression as the set of hyperparameters ( number of subspaces k ` and rank j ` for each layer in our case ) that minimizes the drop in accuracy of the compressed network . While finding the globally optimal solution is NP-complete , we propose ALDS as an efficiently solvable alternative that enables us to search for a locally optimal solution in terms of the maximum relative error incurred across layers . To this end , we derive spectral norm bounds based on the Eckhart-YoungMirsky Theorem for our layer-wise decomposition method to describe the trade-off between the layer compression and the incurred error . Leveraging our bounds we can then efficiently optimize over the set of possible per-layer decompositions . An overview of ALDS is shown in Figure 2 . 2 Method . In this section , we introduce our compression framework consisting of a layer-wise decomposition method ( Section 2.1 ) , a global selection mechanism to simultaneously compress all layers of a network ( Section 2.2 ) , and an optimization procedure ( ALDS ) to solve the selection problem ( Section 2.3 ) . 2.1 Local Layer Compression . We detail our low-rank compression scheme for convolutional layers below and note that it readily applies to fully-connected layers as well as a special case of convolutions with a 1⇥ 1 kernel . Compressing convolutions via SVD . Given a convolutional layer of f filters , c channels , and a 1 ⇥ 2 kernel we denote the corresponding weight tensor by W 2 Rf⇥c⇥1⇥2 . Following Denton et al . ( 2014 ) ; Idelbayev and Carreira-Perpinán ( 2020 ) ; Wen et al . ( 2017 ) and others , we can then interpret the layer as a linear layer of shape f ⇥ c12 and the corresponding rank j-approximation as two subsequent linear layers of shape f ⇥ j and j ⇥ c12 . Mapped back to convolutions , this corresponds to a j ⇥ c⇥ 1 ⇥ 2 convolution followed by a f ⇥ j ⇥ 1⇥ 1 convolution . Multiple subspaces . Following the intuition outlined in Section 1 we propose to cluster the columns of the layer ’ s weight matrix into k 2 separate subspaces before applying SVD to each subset . To this end , we may consider any clustering method , such as k-means or projective clustering ( Chen et al. , 2018 ; Maalouf et al. , 2021 ) . However , such methods require expensive approximation algorithms which would limit our ability to incorporate them into an optimization-based compression framework as outlined in Section 2.2 . In addition , arbitrary clustering may require re-shuffling the input tensors which could lead to significant slow-downs during inference . We instead opted for a simple clustering method , namely channel slicing , where we simply divide the c input channels of the layer into k subsets each containing at most dc/ke consecutive input channels . Unlike other methods , channel slicing is efficiently implementable , e.g. , as grouped convolutions in PyTorch ( Paszke et al. , 2017 ) and ensures practical speed-ups subsequent to compressing the network . Overview of per-layer decomposition . In summary , for given integers j , k 1 and a 4D tensor W 2 Rf⇥c⇥1⇥2 representing a convolution the per-layer compression method proceeds as follows : 1 . PARTITION the channels of the convolutional layer into k subsets , where each subset has at most dc/ke consecutive channels , resulting in k convolutional tensors { Wi } ki=1 where Wi 2 Rf⇥ci⇥1⇥2 , and Pk i=1 ci = c. 2 . DECOMPOSE each tensor Wi , i 2 [ k ] , by building the corresponding weight matrix Wi 2 Rf⇥ci12 , c.f . Figure 3 , computing its j-rank approximation , and factoring it into a pair of smaller matrices Ui of f rows and j columns and Vi of j rows and ci12 columns . 3 . REPLACE the original layer in the network by 2 layers . The first consists of k parallel convolutions , where the ith parallel layer , i 2 [ k ] , is described by the tensor Vi 2 Rj⇥ci⇥1⇥2 which can be constructed from the matrix Vi ( j filters , ci channels , 1 ⇥ 2 kernel ) . The second layer is constructed by reshaping each matrix Ui , i 2 [ k ] , to obtain the tensor Ui 2 Rf⇥j⇥1⇥1 , and then channel stacking all k tensors U1 , · · · , Uk to get a single tensor of shape f ⇥ kj ⇥ 1⇥ 1 . The decomposed layer is depicted in Figure 3 . The resulting layer pair has jc12 and jfk parameters , respectively , which implies a parameter reduction from fc12 to j ( fk + c12 ) . 2.2 Global Network Compression . In the previous section , we introduced our layer compression scheme . We note that in practice we usually want to compress an entire network consisting of L layers up to a pre-specified relative reduction in parameters ( “ compression ratio ” or CR ) . However , it is generally unclear how much each layer ` 2 [ L ] should be compressed in order to achieve the desired CR while incurring a minimal increase in loss . Unfortunately , this optimization problem is NP-complete as we would have to check every combination of layer compression resulting in the desired CR in order to optimally compress each layer . On the other hand , simple heuristics , e.g. , constant per-layer compression ratios , may lead to sub-optimal results , see Section 3 . To this end , we propose an efficiently solvable global compression framework based on minimizing the maximum relative error incurred across layers . We describe each component of our optimization procedure in greater detail below . The layer-wise relative error as proxy for the overall loss . Since the true cost ( the additional loss incurred after compression ) would result in an NP-complete problem , we replace the true cost by a more efficient proxy . Specifically , we consider the maximum relative error `` : = max ` 2 [ L ] `` ` across layers , where `` ` denotes the theoretical maximum relative error in the ` th layer as described in Theorem 1 below . We choose to minimize this particular cost because : ( i ) minimizing the maximum relative error ensures that no layer incurs an unreasonably large error that might otherwise get propagated or amplified ; ( ii ) relying on a relative instead of an absolute error notion is preferred as scaling between layers may arbitrarily change , e.g. , due to batch normalization , and thus the absolute scale of layer errors may not be indicative of the increase in loss ; and ( iii ) the per-layer relative error has been shown to be intrinsically linked to the theoretical compression error , e.g. , see the works of Arora et al . ( 2018 ) and Baykal et al . ( 2019a ) thus representing a natural proxy for the cost . Definition of per-layer relative error . Let W ` 2 Rf ` ⇥c ` ⇥ ` 1⇥ ` 2 and W ` 2 Rf ` ⇥c `` 1 ` 2 denote the weight tensor and corresponding folded matrix of layer ` , respectively . The per-layer relative error `` ` is hereby defined as the relative difference in the operator norm between the matrix Ŵ ` ( that corresponds to the compressed weight tensor Ŵ ` ) and the original weight matrix W ` in layer ` , i.e , . '' ` : = kŴ ` W ` k/kW ` k. ( 1 ) Note that while in practice our method decomposes the original layer into a set of separate layers ( see Section 2.1 ) , for the purpose of deriving the resulting error we re-compose the compressed layers into the overall matrix operator Ŵ ` , i.e. , Ŵ ` = [ U ` 1V ` 1 · · ·U ` k ` V ` k ` ] , where U ` i V ` i is the factorization of the ith cluster ( set of columns ) in the ` th layer , for every ` 2 [ L ] and i 2 [ k ` ] , see supplementary material for more details . We note that the operator norm k · k for a convolutional layer thus signifies the maximum relative error incurred for an individual output patch ( “ pixel ” ) across all output channels . Derivation of relative error bounds . We now derive an error bound that enables us to describe the per-layer relative error in terms of the compression hyperparameters j ` and k ` , i.e. , `` ` = `` ` ( k ` , j ` ) . This will prove useful later on as we have to repeatedly query the relative error in our optimization procedure . The error bound is described in the following . Theorem 1 . Given a layer matrix W ` and the corresponding low-rank approximation Ŵ ` , the relative error `` ` : = kŴ ` W ` k/kW ` k is bounded by '' ` p k/↵1 ·max i2 [ k ] ↵i , j+1 , ( 2 ) where ↵i , j+1 is the j + 1 largest singular value of the matrix W ` i , for every i 2 [ k ] , and ↵1 = kW ` k is the largest singular value of W ` . Proof . First , we recall the matrices W ` 1 , · · · , W ` k and we denote the SVD factorization for each of them by : W ` i = Ũ ` i ⌃̃ ` i Ṽ ` i . Now , observe that for every i 2 [ k ] , the matrix Ŵ ` i is the j-rank approximation of W ` i . Hence , the SVD factorization of Ŵ ` i can be writen as Ŵ ` i = Ũ ` i ⌃̂ ` i Ṽ ` T i , where ⌃̂ ` i 2 Rf⇥d is a diagonal matrix such that its first j-diagonal entries are equal to the first j-entries on the diagonal of ⌃̃ ` i , and the rest are zeros . Hence , W ` Ŵ ` = [ W ` 1 Ŵ ` 1 , · · · , W ` k Ŵ ` k ] = [ Ũ ` 1 ( ⌃̃ ` 1 ⌃̂ ` 1 ) Ṽ ` 1 , · · · , Ũ ` k ( ⌃̃ ` k ⌃̂ ` k ) Ṽ ` k ] = [ Ũ ` 1 · · · Ũ ` k ] diag ⇣ ( ⌃̃ ` 1 ⌃̂ ` 1 ) Ṽ ` 1 , . . . , ( ⌃̃ ` k ⌃̂ ` k ) Ṽ ` k ⌘ . ( 3 ) By ( 3 ) and by the triangle inequality , we have that W ` Ŵ ` h Ũ ` 1 · · · Ũ ` k i diag ⇣ ( ⌃̃ ` 1 ⌃̂ ` 1 ) Ṽ ` 1 , . . . , ( ⌃̃ ` k ⌃̂ ` k ) Ṽ ` k ⌘ . ( 4 ) Now , we observe that h Ũ ` 1 · · · Ũ ` k i 2 = h Ũ ` 1 · · · Ũ ` k ih Ũ ` 1 · · · Ũ ` k iT = kdiag ( k , . . . , k ) k = k. ( 5 ) Finally , we show that diag ⇣ ( ⌃̃ ` 1 ⌃̂ ` 1 ) Ṽ ` 1 , . . . , ( ⌃̃ ` k ⌃̂ ` k ) Ṽ ` k ⌘ = max i2 [ k ] ( ⌃̃ ` i ⌃̂ ` i ) Ṽ ` i ( 6 ) = max i2 [ k ] ( ⌃̃ ` i ⌃̂ ` i ) = max i2 [ k ] ↵i , j+1 , ( 7 ) where the second equality holds since the columns of V are orthogonal and the last equality holds according to the Eckhart-Young-Mirsky Theorem ( Theorem 2.4.8 of Golub and Van Loan ( 2013 ) ) . Plugging ( 7 ) and ( 5 ) into ( 4 ) concludes the proof . Resulting network size . Let ✓ = { W ` } L ` =1 denote the set of weights for the L layers and note that the number of parameters in layer ` is given by |W ` | = f ` c `` 1 ` 2 and |✓| = P ` 2 [ L ] |W ` | . Moreover , note that |Ŵ ` | = j ` ( k ` f ` + c `` 1 ` 2 ) if decomposed , ✓̂ = { Ŵ ` } L ` =1 , and |✓̂| = P ` 2 [ L ] |Ŵ ` | . The overall compression ratio is thus given by 1 |✓̂|/|✓| where we neglected other parameters for ease of exposition . Observe that the layer budget |Ŵ ` | is fully determined by k ` , j ` just like the error bound . Global Network Compression . Putting everything together we obtain the following formulation for the optimal per-layer budget : '' opt = min { j ` , k ` } L ` =1 max ` 2 [ L ] '' ` ( k ` , j ` ) ( 8 ) subject to 1 |✓̂ ( k1 , j1 , ... , kL , jL ) |/|✓| CR , where CR denotes the desired overall compression ratio . Thus optimally allocating a per-layer budget entails finding the optimal number of subspaces k ` and ranks j ` for each layer constrained by the desired overall compression ratio CR .

Summary: The paper proposes a one-shot approach to compress already trained neural networks using low-rank compression. The authors propose a new scheme of applying low-rank: instead of applying it directly to a layer, authors suggest partitioning the filters into K blocks, and apply low-rank to each of these blocks separately. Haven defined this scheme, authors then propose a one-shot alternating-optimization algorithm to select ranks and to partition matrix columns/filters into K blocks. Once ranks and partitions are determined, the matrices are decomposed and finetuned. Optionally, this entire compress=>finetune process can be repeated multiple times. Authors present numerical results where the suggested approach almost always wins, however some of the comparisons are questionable. In terms of the algorithm itself, the details of alt.opt are missing, as well as theoretical analysis of the proposed compression scheme is missing.