Scalable Model Compression by Entropy Penalized Reparameterization

1 INTRODUCTION . Artificial neural networks ( ANNs ) have proven to be highly successful on a variety of tasks , and as a result , there is an increasing interest in their practical deployment . However , ANN parameters tend to require a large amount of space compared to manually designed algorithms . This can be problematic , for instance , when deploying models onto devices over the air , where the bottleneck is often network speed , or onto devices holding many stored models , with only few used at a time . To make these models more practical , several authors have proposed to compress model parameters ( Han et al. , 2016 ; Louizos , Ullrich , et al. , 2017 ; Molchanov et al. , 2017 ; Havasi et al. , 2019 ) . While other desiderata often exist , such as minimizing the number of layers or filters of the network , we focus here simply on model compression algorithms that 1. minimize compressed size while maintaining an acceptable classification accuracy , 2. are conceptually simple and easy to implement , and 3. can be scaled easily to large models . Classical data compression in a Shannon sense ( Shannon , 1948 ) requires discrete-valued data ( i.e. , the data can only take on a countable number of states ) and a probability model on that data known to both sender and receiver . Practical compression algorithms are often lossy , and consist of two steps . First , the data is subjected to ( re- ) quantization . Then , a Shannon-style entropy coding method such as arithmetic coding ( Rissanen and Langdon , 1981 ) is applied to the discrete values , bringing them into a binary representation which can be easily stored or transmitted . Shannon ’ s source coding theorem establishes the entropy of the discrete representation as a lower bound on the average length of this binary sequence ( the bit rate ) , and arithmetic coding achieves this bound asymptotically . Thus , entropy is an excellent proxy for the expected model size . The type of quantization scheme affects both the fidelity of the representation ( in this case , the precision of the model parameters , which in turn affects the prediction accuracy ) as well as the bit ∗Work performed during the Google AI Residency Program . ( http : //g.co/airesidency ) rate , since a reduced number of states coincides with reduced entropy . ANN parameters are typically represented as floating point numbers . While these technically have a finite ( but large ) number of states , the best results in terms of both accuracy and bit rate are typically achieved for a significantly reduced number of states . Existing approaches to model compression often acknowledge this by quantizing each individual linear filter coefficient in an ANN to a small number of pre-determined values ( Louizos , Reisser , et al. , 2019 ; Baskin et al. , 2018 ; F. Li et al. , 2016 ) . This is known as scalar quantization ( SQ ) . Other methods explore vector quantization ( VQ ) , closely related to kmeans clustering , in which each vector of filter coefficients is quantized jointly ( Chen , J. Wilson , et al. , 2015 ; Ullrich et al. , 2017 ) . This is equivalent to enumerating a finite set of representers ( representable vectors ) , while in SQ the set of representers is given by the Kronecker product of representable scalar elements . VQ is much more general than SQ , in the sense that representers can be placed arbitrarily : if the set of useful filter vectors all live in a subregion of the entire space , there is no benefit in having representers outside of that region , which may be unavoidable with SQ ( fig . 1 , left ) . Thus , VQ has the potential to yield better results , but it also suffers from the “ curse of dimensionality ” : the number of necessary states grows exponentially with the number of dimensions , making it computationally infeasible to enumerate them explicitly , hence limiting VQ to only a handful of dimensions in practice . One of the key insights leading to this paper is that the strengths of SQ and VQ can be combined by representing the data in a “ latent ” space . This space can be an arbitrary rescaling , rotation , or otherwise warping of the original data space . SQ in this space , while making quantization computationally feasible , can provide substantially more flexibility in the choice of representers compared to the SQ in the data space ( fig . 1 , right ) . This is in analogy to recent image compression methods based on autoencoders ( Ballé , Laparra , et al. , 2017 ; Theis et al. , 2017 ) . The contribution of this paper is two-fold . First , we propose a novel end-to-end trainable model compression method that uses scalar quantization and entropy penalization in a reparameterized space of model parameters . The reparameterization allows us to use efficient SQ , while achieving flexibility in representing the model parameters . Second , we provide state-of-the-art results on a variety of network architectures on several datasets . This demonstrates that more complicated strategies involving pretraining , multi-stage training , sparsification , adaptive coding , etc. , as employed by many previous methods , are not necessary to achieve good performance . Our method scales to modern large image datasets and neural network architectures such as ResNet-50 on ImageNet . 2 ENTROPY PENALIZED REPARAMETERIZATION . We consider the classification setup , where we are given a dataset D = { ( x1 , y1 ) , ... ( xN , yN ) } consisting of pairs of examples xi and corresponding labels yi . We wish to minimize the expected negative log-likelihood on D , or cross-entropy classification loss , over Θ , the set of model parame- ters : Θ∗ = arg min Θ ∑ ( x , y ) ∼D − log p ( y | x ; Θ ) , ( 1 ) where p ( y | x ; Θ ) is the likelihood our model assigns to a dataset sample ( x , y ) . The likelihood function is implemented using an ANN with parameters Θ = { W 1 , b1 , W 2 , b2 , . . . , WN } , where W k and bk denote the weight ( including convolutional ) and bias terms at layer k , respectively . Compressing the model amounts to compressing each parameter in the set Θ . Instead of compressing each parameter directly , we compress reparameterized forms of them . To be precise , we introduce the reparameterizations Φ = { W̃ 1 , b̃1 , W̃ 2 , b̃2 , . . . , W̃N } and parameter decoders fconv , fdense , fbias such that W k = fconv ( W̃ k ) if layer k is convolutional , ( 2 ) W k = fdense ( W̃ k ) if layer k is fully connected , ( 3 ) bk = fbias ( b̃k ) if layer k has a bias . ( 4 ) We can think of each parameter decoder f as a mapping from reparameterization space to parameter space . For ease of notation , we write F = { fconv , fdense , fbias } and Θ = F ( Φ ) . The parameter decoders themselves may have learnable parameters , which we denote Ψ . Our method is visually summarized in figs . 2 and 3 . 2.1 COMPRESSING Φ WITH ENTROPY CODING . In order to apply a Shannon-style entropy coder efficiently to the reparameterizations Φ , we need a discrete alphabet of representers and associated probabilities for each representer . Rather than handling an expressive set of representers , as in VQ , we choose to fix them to the integers , and achieve expressivity via the parameter decoders F instead . Each reparameterizationφ ∈ Φ ( i.e . a W̃ or b̃ representing a weight or bias , respectively ) is a matrix in Zd× ` interpreted as consisting of d samples from a discrete probability distribution producing vectors of dimension ` . We fit a factorized probability model q ( φ ) = d∏ j=1 ∏̀ i=1 qi ( φj , i ) ( 5 ) to each column i of φ , using ` different probability models qi for each corresponding parameter decoder ( the form of qi is described in the next section ) . Fitting of probability models is often done by minimizing the negative log-likelihood . Assumingφ follows the distribution q , Shannon ’ s source coding theorem states that the minimal length of a bit sequence encoding φ is the self-information of φ under q : I ( φ ) = − log2 q ( φ ) , ( 6 ) which is identical to Shannon cross entropy up to an expectation operator , and identical to the negative log likelihood up to a constant factor . By minimizing I over q and φ during training , we thus achieve two goals : 1 ) we fit q to the model parameters in a maximum likelihood sense , and 2 ) we directly optimize the parameters for compressibility . After training , we design an arithmetic code for q , and use it to compress the model parameters . This method incurs only a small overhead over the theoretical bound due to the finite length of the bit sequence ( arithmetic coding is asymptotically optimal ) . Practically , the overhead amounts to less than 1 % of the size of the bit sequence ; thus , self-information is an excellent proxy for model size . Further overhead results from including a description of Ψ , the parameters of the parameter decoders , as well as of q itself ( in the form of a table ) in the model size . However , these can be considered constant and small compared to the total model size , and thus do not need to be explicitly optimized for . The overall loss function is simply the additive combination of the original cross-entropy classification loss under reparameterization with the self-information of all reparameterizations : L ( Φ , Ψ ) = ∑ ( x , y ) ∼D − log p ( y | x ; F ( Φ ) ) + λ ∑ φ∈Φ I ( φ ) . ( 7 ) We refer to the second term ( excluding the constant λ ) as the rate loss . By varying λ across different experiments , we can explore the Pareto frontier of compressed model size vs. model accuracy . To compare our method to other work , we varied λ such that our method produced similar accuracy , and then compared the resulting model size . 2.2 DISCRETE OPTIMIZATION . Since Φ is discrete-valued , we need to make some further approximations in order to optimize L over it using stochastic gradient descent . To get around this , we maintain continuous surrogates Φ̂ . For optimizing the classification loss , we use the “ straight-through ” gradient estimator ( Bengio et al. , 2013 ) , which provides a biased gradient estimate but has shown good results in practice . This consists of rounding the continuous surrogate to the nearest integer during training , and ignoring the rounding for purposes of backpropagation . After training , we only keep the discretized values . In order to obtain good estimates for both the rate term and its gradient during training , we adopt a relaxation approach previously described by Ballé , Minnen , et al . ( 2018 , appendix 6.1 ) ; the code is provided as an open source library1 . In a nutshell , the method replaces the probability mass functions qi with a set of non-parametric continuous density functions , which are based on small ANNs . These density models are fitted to φ̂j , i + nj , i , where nj , i ∼ U ( − 12 , 1 2 ) is i.i.d . uniformly distributed additive noise . This turns out to work well in practice , because the negative log likelihood of these noise-affected variates under the continuous densities approximates the self-information I : I ( φ ) ≈ d∑ j=1 ∑̀ i=1 − log2 q̃i ( φj , i + nj , i ) , ( 8 ) 1https : //github.com/tensorflow/compression where q̃i denote the density functions . Once the density models are trained , the values of the probability mass functions modeling φ are derived from the substitutes q̃i and stored in a table , which is included in the model description . The parameters of q̃i are no longer needed after training .

Summary: The authors propose a novel, information theoretic approach to learning compressed neural networks, whereby model parameters $\theta$ are expressed in an (approximately discrete) latent space as $\theta = f_{\psi}(\phi)$. During training, "reparameterizations" $\Phi$ are treated as draws from a generative distribution $q(\phi)$. As per the source coding theorem, penalizing the entropy of $q$ (more accurately, of $\hat{q}$) directly promotes compressible choices of $\Phi$ and, therefore, of $\Theta$. After training, arithmetic coding of $q$ is used to compress $\Phi$. Empirical results on a suite of common benchmarks demonstrates the proposed algorithm gracefully trades off between model performance and size.