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: This paper introduces a novel method of weight compression. Weight tensors are stored as sparse, quantized matrix factors, and the underlying matrix factorization problem can be considered as a quantized sparse PCA problem and be solved through iterative projected gradient descent methods. The authors' method  is applicable to both moderate and extreme compression regimes, and is claimed to achieve or be on par with state-of-the-art trade-offs between accuracy and model size.