Quantized sparse PCA for neural network weight compression

1 INTRODUCTION . Deep neural networks have achieved state-of-the-art results in a wide variety of tasks . However , deployment remains challenging due to their large compute and memory requirements . Neural networks deployed on edge devices such as mobile or IoT devices are subject to stringent compute and memory constraints , while networks deployed in the cloud do not suffer such constraints but might suffer excessive latency or power consumption . To reduce neural network memory and compute footprint , several approaches have been introduced in literature . Methods related to our approach , as well as their benefits and downsides , are briefly introduced in this section and described in more detail in the related works section . Tensor factorization approaches ( Denil et al. , 2013 ) replace a layer in a neural network with two layers , whose weights are low-rank factors of the original layer ’ s weight tensor . This reduces the number of parameters and multiply-add operations ( MACs ) , but since the factorizations are by design restricted to those that can be realized as individual layers , potential for compression is limited . By pruning a neural network ( Louizos et al. , 2017 ; He et al. , 2017 ) , individual weights are removed . Pruning has shown to yield moderate compression-accuracy trade-offs . Due to the overhead required to keep track of which elements are pruned , real yield of ( unstructured ) pruning is lower than the pruning ratio . An exception to this is structured pruning , in which entire neurons are removed from a network , and weight tensors can be adjusted accordingly . However , achieving good compression ratios at reasonable accuracy using structured pruning has proven difficult . Scalar quantization ( Jacob et al. , 2018 ; Nagel et al. , 2021 ) approximates neural network weights with fixed-point values , i.e. , integer values scaled by a fixed floating point scalar . Scalar quantization has shown to yield high accuracy at reasonable compression ratios , e.g. , 8 bit quantization yields a 4x compression ratio and virtually no accuracy degradation on many networks . However , scalar quantization does not yield competitive compression vs accuracy trade-offs at high compression ratios . Vector quantization ( Stock et al. , 2019 ; Martinez et al. , 2021 ) approximates small subsets of weights , e.g. , individual 3 × 3 filters in convolutional weight tensors , by a small set of codes . This way , the storage can be reduced by storing the codebook and one code index for each original vector , instead of the individual weights . While vector quantization can achieve high compression with moderate accuracy loss , these methods usually struggle to reach the accuracy of uncompressed models in low compression regimes . In this paper , we provide a novel view on tensor factorization . Instead of restricting factorization to those that can be realized as two separate neural network layers , we show that much higher compression ratios can be achieved by shifting the order of operations . In our method , we find a factorization C , Z for an original weight tensor W , such that the matrix product CZ closely approximates the Weight tensor ! ! `` # × ! $ % ×ℎ× $ Tile and reshape Weight tensor tiled and reshaped to % × & ; & ≫ % Sparse PCA C matrix % × ) ; ) < % * Z matrix ( sparse ) ) × & ; & ≫ ) & = ! ! `` # × ! $ % ×ℎ× $ / % Weight tensor ! ! `` # × ! $ % ×ℎ× $ original weight tensor . The compr ssion is then pushed further by obtaining quantized factors and additionally sparse Z . During inference , the product CZ is computed first , and its result is reshaped back into the original weight tensor ’ s shape , and used for the following computations . This approach allows t e u e of an arbitrary factorization of the original weight tensor . We show that this approach outperforms or is on par with vector quantization in high compression regimes , yet extends to scalar quantization levels of compression-accuracy trade-offs for lower compression ratios . Our contributions in this paper are as follows : • We show that the problem of tensor factorization and vector quantization can be formulated in a unified way as quantized sparse principle component analysis ( PCA ) problem . • We propose an iterative projected gradient descent method to solve the quantized sparse PCA problem . • Our experimental results demonstrate the benefits of this approach . By simultaneously solving tensor factorization and vector quantization problem , we can achieve better accuracy than vector quantization in low compression regimes , and higher compression ratios than scalar quantization approaches at moderate loss in accuracy . 2 RELATED WORK . SVD-based methods and tensor decompositions SVD decomposition was first used to demonstrate redundacy in weight parameters in neural networks in Denil et al . ( 2013 ) . Later several methods for reducing the inference time based on SVD decomposition were suggested ( Denton et al. , 2014 ; Jaderberg et al. , 2014 ) . A similar technique was proposed for gradients compression in dataparallel distributed optimization by Vogels et al . ( 2019 ) . The main difference between these methods is the way 4D weights of a convolutional layer is transformed into a matrix which leads to different shapes of the convolutional layers in the resulting decomposition . Following a similar direction , several works focus on higher-order tensor decomposition methods which lead to introducing of three or four convolutional layers ( Lebedev et al. , 2014 ; Kim et al. , 2015 ; Su et al. , 2018 ) . Weight pruning A straightforward approach to reducing neural network model size is removing a percentage of weights . A spectrum of weight pruning approaches of different granularity has been introduced in the literature . Structured pruning approaches such as He et al . ( 2017 ) kill entire channels of the weights , while unstructured pruning approaches ( Louizos et al. , 2017 ; Zhu & Gupta , 2017 ; Neklyudov et al. , 2017 ; Dai et al. , 2018 ) focus on individual values . A recent survey on unstructured pruning is provided in Gale et al . ( 2019 ) . Scalar quantization and mixed precision training By quantizing neural network weights to lower bitwidths , model footprint can be reduced as well , as each individual weight requires fewer bits to be stored . For example , quantizing 32 bit floating points weight to 8 bit fixed point weights yields a 4x compression ratio . Most quantization approaches use the straight-through estimator ( STE ) for training quantized models ( Bengio et al . ( 2013 ) ; Krishnamoorthi ( 2018 ) ) . One way to further improve the accuracy of quantized models is learning the quantization scale and offset jointly with the network parameters ( Esser et al . ( 2019 ) ; Bhalgat et al . ( 2020 ) ) . A recent survey on practical quantization approaches can be found in Nagel et al . ( 2021 ) . In order to improve the accuracy of quantized models , several methods suggest using mixed precision quantization . The work by Uhlich et al . ( 2019 ) introduced an approach on learning integer bit-width for each layer using STE . Using non-uniform bit-width allows the quantization method to use lower bit-width for more compressible layers of the network . Several works ( van Baalen et al. , 2020 ; Dong et al. , 2019 ; Wang et al. , 2019 ) improve upon the approach by Uhlich et al . ( 2019 ) by using different methods for optimization over the bit-widths . Vector quantization . Several works use vector quantization approach for compression of weights of convolutional and fully connected layers ( Gong et al . ( 2014 ) ; Martinez et al . ( 2021 ) ; Fan et al . ( 2020 ) ; Stock et al . ( 2019 ) ; Wu et al . ( 2016 ) ) . The convolutional weight tensors are reshaped into matrices , then K-means methods is applied directly on the rows or columns . Besides weight compression , the work by Wu et al . ( 2016 ) suggests using vector quantization for reducing the inference time by reusing parts of the computation . Recently several works suggested improvements on the basic vector quantization approach . Dataaware vector quantization which improves the clustering method by considering input activation data is demonstrated to improve the accuracy of the compressed models by Stock et al . ( 2019 ) . Another direction is introducing a permutation to the weight matrices which allows to find subsets of weights which are more compressible ( Martinez et al. , 2021 ) . An inverse permutation is applied at the output of the corresponding layers to preserve the original output of the model . Besides weights compression problem , various improved vector quantization methods were applied in image retrieval domain in order to accelerate scalar product computation for image descriptors ( Chen et al . ( 2010 ) ; Ge et al . ( 2013 ) ; Norouzi & Fleet ( 2013 ) ) . In particular , additive quantization method which we will show is related to quantized sparse PCA was introduced Babenko & Lempitsky ( 2014 ) . Surveys on vector quantization methods are provided in Matsui et al . ( 2018 ) ; Gersho & Gray ( 2012 ) . Sparse PCA . Introduced in Zou et al . ( 2006 ) , sparse PCA can be solved by a plethora of algorithms . The method proposed in this paper can be considered as an instance of thresholding algorithms ( Ma , 2013 ) . Although soft-thresholding methods are prevalent in the literature , we adopt an explicit projection step using hard thresholding to have direct control over the compression ration . Note that sparse PCA can be extended to include additional structures with sparsity ( Jenatton et al. , 2010 ; Lee et al. , 2006 ) . 3 METHOD . In this section , we describe the main algorithm , which can be considered as sparse quantized principle component analysis ( PCA ) . 3.1 QUANTIZED SPARSE PCA . Consider a weight tensor of a convolutional layer W ∈ Rfout×fin×h×w , where fin , fout are the number of input and output feature maps , and h , w are the spatial dimensions of the filter . We reshape it into a matrix W̃ ∈ Rd×n , where d is the tile size and n is the number of tiles . For the reshape , we consider the dimensions in the order fout , fin , h , w in our experiments . The goal is to factorize W̃ into the product of two matrices as follows : W̃ = CZ , ( 1 ) where C ∈ Rd×k is the codebook , and Z ∈ Rk×n is the set of linear coefficients , or a latent variable ( k < d < n ) . Following the standard PCA method , we factorize the zero mean version of W̃ . With this decomposition , every column of W̃ , denoted by W̃ : ,i , is a linear combination of k codebook vectors ( columns of C ) : W̃ : ,i = k∑ j=1 ZjiC : ,j , ( 2 ) where C : ,j is the j-th column of C. This decomposition problem is an instance of sparse PCA methods ( Zou et al . ( 2006 ) ) . For network compression , we are additionally interested in quantized matrices C and Z . The quantization operation for an arbitrary C and Z is defined as follows : Cq = Qc ( C ; sc , bc ) , ( 3 ) Zq = Qz ( Z ; sz , bz ) , ( 4 ) where bc , bz are the quantization bit-widths , and sc and sz are the quantization scale vectors for C and Z , respectively . We consider per-channel quantization , i.e. , the quantization scale values are shared for each column of C and each row or column of Z ( see section 3.2 for details on which is used when ) : Cq , ij = clamp ( ⌊ Cij si ⌉ , 0 , 2bc − 1 ) si , ( 5 ) Zq , ij = clamp ( ⌊ Zij sj ⌉ , 0 , 2bz − 1 ) sj , ( 6 ) where b·e is rounding to nearest integer , and clamp ( · ) is defined as : clamp ( x , a , b ) = a x < a x a ≤ x ≤ b b x > b . ( 7 ) We refer to the problem of finding quantized factors C and Z with sparse Z as quantized sparse PCA problem . Once we obtain the factors C and Z , we get the matrix W̃ = ( CZ ) , which can be reshaped back to a convolutional layer . The reshaped convolutional layer for factors C , Z is denoted by [ CZ ] . It is well known in network compression literature that it is better to find the best factorization on the data manifold ( Zhang et al . ( 2015 ) ; He et al . ( 2017 ) ; Stock et al . ( 2019 ) ) . Therefore , we solve the following optimization problem : C∗ , Z∗ = argmin Cq , Zq E ( X , Y ) ∼D ( ‖Y − [ CqZq ] ∗X‖2F ) ( 8 ) s.t . Cq = Qc ( C ; sc , bC ) Zq = Qz ( Z ; sz , bZ ) ‖Zq‖0 ≤ S , Z ∈ R k×n , C ∈ Rd×k , where the parameter S controls the sparsity ratio of Z , X and Y are the stored input output of the target layer in the original model , D is the data distribution , and ∗ is the convolution operation . L0 norm is used for the constraint on the number of nonzero elements in the matrix Zq . We approximate the expected value of above optimization problem using a subset of the training data : E ( X , Y ) ∼D ( ‖Y − [ CqZq ] ∗X‖2F ) ≈ 1 m m∑ i=1 ‖Yi − [ CqZq ] ∗Xi‖2F , where m is the number of samples used for optimization . Following the compression method introduced in Zhang et al . ( 2015 ) , we use output of the previous compressed layer as X instead of the stored input to the layer in the original model . This approach aims at compensating for the error accumulation in deep networks .

Summary: The paper proposes a method for compression of neural network weights. The proposal turns weight tensors into matrices, factorizes these matrices into a rank-k factorization via PCA, applies quantization to the factor matrices, and additionally makes the right (latent) matrix sparse. An algorithm is presented, with two different thresholding options (per-iteration, or one-shot). In experiments, these ideas lead to an accuracy/compression tradeoff which is either competitive with or better than previous state-of-the-art across all compression ratios investigated.