S CALABLE M ODEL C OMPRESSION BY E NTROPY P ENALIZED R EPARAMETERIZATION
A BSTRACT
We describe a simple and general neural network weight compression approach, in which the network parameters (weights and biases) are represented in a “latent” space, amounting to a reparameterization.
This space is equipped with a learned probability model, which is used to impose an entropy penalty on the parameter representation during training, and to compress the representation using a simple arithmetic coder after training.
Classiﬁcation accuracy and model compressibility is maximized jointly, with the bitrate–accuracy trade-off speciﬁed by a hyperparameter.
We evaluate the method on the MNIST, CIFAR-10 and ImageNet classiﬁcation benchmarks using six distinct model architectures.
Our results show that state-of-the-art model compression can be achieved in a scalable and general way without requiring complex procedures such as multi-stage training.
I NTRODUCTION
Artiﬁcial 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 ﬁlters of the network, we focus here simply on model compression algorithms that 1.
minimize compressed size while maintaining an acceptable classiﬁcation 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 ﬁdelity of the representation (in this case, the precision of the model parameters, which in turn affects the prediction accuracy) as well as the bit [Figure] rate, since a reduced number of states coincides with reduced entropy.
ANN parameters are typically represented as ﬂoating point numbers.
While these technically have a ﬁnite (but large) number of states, the best results in terms of both accuracy and bit rate are typically achieved for a signiﬁcantly reduced number of states.
Existing approaches to model compression often acknowledge this by quantizing each individual linear ﬁlter coefﬁcient 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 ﬁlter coefﬁcients is quantized jointly (Chen, J. Wilson, et al., 2015; Ullrich et al., 2017).
This is equivalent to enumerating a ﬁnite 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 ﬁlter vectors all live in a subregion of the entire space, there is no beneﬁt in having representers outside of that region, which may be unavoidable with
SQ (ﬁg. 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 ﬂexibility in the choice of representers compared to the SQ in the data space (ﬁg. 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 efﬁcient SQ, while achieving ﬂexibility 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, sparsiﬁcation, 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.
E NTROPY PENALIZED REPARAMETERIZATION
We consider the classiﬁcation setup, where we are given a dataset D
= { ( x 1 , y 1 ) , ... ( x N , y N ) } consisting of pairs of examples x i and corresponding labels y i .
We wish to minimize the expected negative log-likelihood on D , or cross-entropy classiﬁcation loss, over Θ , the set of model parame- x (cid:102)
W(cid:101) b(cid:102)
W(cid:101) b(cid:102) W n f conv f bias f conv f bias f dense ∗ + ∗ + · logits
Wb
Wb
W n conv conv classiﬁer
[Figure]
+ ·
[Equation]
[Figure]
+ ·
W
[Figure]
ters:
[Equation]
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 , b 1 , W 2 , b 2 , . . .
, W N } , where
W k and b k 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 Φ = { (cid:102) W 1 , (cid:101) b 1 , (cid:102) W 2 , (cid:101) b 2 , . . .
, (cid:102) W N } and parameter decoders f conv , f dense , f bias such that
[Equation]
We can think of each parameter decoder f as a mapping from reparameterization space to parameter space.
For ease of notation, we write F = { f conv , f dense , f bias } and Θ = F (Φ) .
The parameter decoders themselves may have learnable parameters, which we denote Ψ .
Our method is visually summarized in ﬁgs.
2 and 3.
2. C OMPRESSING Φ WITH ENTROPY CODING
In order to apply a Shannon-style entropy coder efﬁciently 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 ﬁx them to the integers, and achieve expressivity via the parameter decoders F instead.
Each reparameterization φ ∈ Φ (i.e. a (cid:102) W or (cid:101)
b representing a weight or bias, respectively) is a matrix in Z d × (cid:96) interpreted as consisting of d samples from a discrete probability distribution producing vectors of dimension (cid:96) .
We ﬁt a factorized probability model
[Equation]
to each column i of φ , using (cid:96) different probability models q i for each corresponding parameter decoder (the form of q i 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 :
[Equation]
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 ﬁt 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 ﬁnite 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 classiﬁcation loss under reparameterization with the self-information of all reparameterizations:
[Equation]
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. D ISCRETE 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 classiﬁcation 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 library 1 .
In a nutshell, the method replaces the probability mass functions q i with a set of non-parametric continuous density functions, which are based on small
ANNs.
These density models are ﬁtted to ˆ φ j,i + n j,i , where n j,i ∼ U ( − 12 , 12 ) 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 :
[Equation]
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.
2. M ODEL PARTITIONING
A central component of our approach is partitioning the set of model parameters into groups.
For the purpose of creating a model compression method, we interpret entire groups of model parameters as samples from the same learned distribution.
We deﬁne a fully factorized distribution q (Φ) =
(cid:81)
φ ∈
Φ q φ ( φ ) , and introduce parameter sharing within the factors q φ of the distribution that correspond to the same group, as well as within the corresponding decoders.
These group assignments are ﬁxed a priori.
For instance, in ﬁg. 2, (cid:102) W 1 and (cid:102) W 2 can be assumed to be samples of the same distribution, that is q (cid:102) W 1 = q (cid:102) W 2 .
We also use the same parameter decoder f conv to decode them.
Further, each of the reparameterizations φ is deﬁned as a rank-2 tensor (a matrix), where each row corresponds to a “sample” from the learned distribution.
The operations in f apply the same transformation to each row (ﬁg. 3).
As an example, in f conv , each spatial H × W matrix of ﬁlter coefﬁcients is assumed to be a sample from the same distribution.
Our method can be applied analogously to various model partitionings.
In fact, in our experiments, we vary the size of the groups, i.e., the number of parameters assumed i.i.d.
, depending on the total number of parameters of the model ( Θ ).
The size of the groups parameterizes a trade-off between compressibility and overhead: if groups consisted of just one scalar parameter each, compressibility would be maximal, since q would degenerate (i.e., would capture the value of the parameter with certainty).
However, the overhead would be maximal, since F and q would have a large number of parameters that would need to be included in the model size (defeating the purpose of compression).
On the other hand, encoding all parameters of the model with one and the same decoder and scalar distribution would minimize overhead, but may be overly restrictive by failing to capture distributional differences amongst all the parameters, and hence lead to suboptimal compressibility.
We describe the group structure of each network that we use in more detail in the following section.
E XPERIMENTS
For our MNIST and CIFAR-10 experiments, we evaluate our method by applying it to four distinct image classiﬁcation networks: LeNet300-100 (Lecun et al., 1998) and LeNet-5-Caffe 2 on
MNIST (LeCun and Cortes, 2010), as well as VGG-16 3 (Simonyan and Zisserman, 2015) and
ResNet-20 (He et al., 2016b; Zagoruyko and Komodakis, 2016) with width multiplier 4 (ResNet-204) on CIFAR-10 (Zagoruyko and Komodakis, 2016).
For our ImageNet experiments, we evaluate our method on the ResNet-18 and ResNet-50 (He et al., 2016a) networks.
We train all our models from scratch and compare them with recent state-of-the-art methods by quoting performance from their respective papers.
Compared to many previous approaches, we do not initialize the network with pre-trained or pre-sparsiﬁed weights.
We found it useful to use two separate optimizers: one to optimize the variables of the probability models ˜ q i , and one to optimize the reparameterizations Φ and variables of the parameter decoders Ψ .
While the latter is chosen to be the same optimizer typically used for the task/architecture, the former is always Adam (Kingma and Ba, 2015) with a learning rate of 10 − 4 .
We chose to always use
Adam, because the parameter updates used by Adam are independent of any scaling of the objective (when its hyper-parameter (cid:15) is sufﬁciently small).
In our method, the probability model variables only get gradients from the entropy loss which is scaled by the rate penalty λ .
Adam normalizes out this scale and makes the learning rate of the probability model independent of λ and of other hyperparameters such as the model partitioning.
3. MNIST EXPERIMENTS
We apply our method to two LeNet variants: LeNet300-100 and LeNet5-Caffe and report results in table 1.
We train the networks using Adam with a constant learning rate of 0.001 for 200, iterations.
To remedy some of the training noise from quantization, we maintain an exponential moving average (EMA) of the weights and evaluate using those.
Note that this does not affect the quantization, as quantization is performed after the EMA variables are restored.
LeNet300-100 consists of 3 fully connected layers.
We partitioned this network into three parameter groups: one for the ﬁrst two fully connected layers, one for the classiﬁer layer, and one for biases.
LeNet5-Caffe consists of two 5 × 5 convolutional layers followed by two fully connected layers, with max pooling following each convolutional layer.
We partitioned this network into four parameter groups: One for both of the convolutional layers, one for the penultimate fully connected layer, one for the ﬁnal classiﬁer layer, and one for the biases.
As evident from table 1, for the larger LeNet300-100 model, our method outperforms all the baselines while maintaining a comparable error rate.
For the smaller LeNet5-Caffe model, our method is second only to Minimal Random Code Learning (Havasi et al., 2019).
Note that in both of the
MNIST models, the number of probability distributions (cid:96) = 1 in every parameter group, including in the convolutional layers.
To be precise, the (cid:102) W k for the convolutional weights W k will be
H · W · I · O × 1 .
This is a good trade-off, since the model is small to begin with, and having (cid:96) = 5 · 5 = 25 scalar probability models for 5 × 5 convolutional layers would have too much overhead.
For both of the MNIST models, we found that letting each subcomponent of F be a simple dimension-wise scalar afﬁne transform (similar to f dense in ﬁg. 3), was sufﬁcient.
Since each φ is quantized to integers, having a ﬂexible scale and shift leads to ﬂexible SQ, similar to Louizos,
Reisser, et al.
Due to the small size of the networks, more complex transformation functions would lead to too much overhead.
3. CIFAR-10 EXPERIMENTS
We apply our method to VGG-16 (Simonyan and Zisserman, 2015) and ResNet-20-4 (He et al., 2016b; Zagoruyko and Komodakis, 2016) and report the results in table 1.
For both VGG-16 and
ResNet-20-4, we use momentum of 0.9 with an initial learning rate of 0.1, and decay by 0.2 at iterations 256,000, 384,000, and 448,000 for a total of 512,000 iterations.
This learning rate schedule was ﬁxed from the beginning and was not tuned in any way other than verifying that our models’ training loss had converged.
VGG-16 consists of 13 convolutional layers of size 3 × 3 followed by 3 fully connected layers.
We split this network into four parameter groups: one for all convolutional layers and one each all fully connected layers.
We did not compress biases.
We found that the biases in 32-bit ﬂoating point format add up to about 20 KB, which we add to our reported numbers.
ResNet-20-4 consists of 3 ResNet groups with 3 residual blocks each.
There is also an initial convolution layer and a ﬁnal fully connected classiﬁcation layer.
We partition this network into two parameter groups: one for all convolutional layers and one for the ﬁnal classiﬁcation layer.
We again did not compress biases and include them in our results; they add up to about 11 KB.
For convolutions in both CIFAR-10 models, (cid:96) =
O × I = 9 ; f conv and f dense are exactly as pictured in ﬁg. 3.
To speed up training, we ﬁxed ψ W to a diagonal scaling matrix multiplied by the inverse real-valued discrete Fourier transform (DFT).
We found that this particular choice performs much better than SQ, or than choosing a random but ﬁxed orthogonal matrix in place of the DFT (ﬁg. 4).
From the error vs. rate plots, the beneﬁt of reparameterization in the high compression regime is evident.
VGG-16 and ResNet-20-4 both contain batch normalization (Ioffe and Szegedy, 2015) layers that include a moving average for the mean and variance.
Following Havasi et al.
(2019), we do not include the moving averages in our reported numbers.
We do, however, include the batch normalization bias term β and let it function as the bias for each layer ( γ is set to a constant 1 ).
3. I MAGE N ET EXPERIMENTS
For the ImageNet dataset (Russakovsky et al., 2015), we reproduce the training setup and hyperparameters from He et al.
(2016a).
We put all 3 × 3 convolutional kernels in a single parameter group, similar to in our CIFAR experiments.
In the case of ResNet-50, we also group all 1 ×convolutional kernels together.
We put all the remaining layers in their own groups.
This gives a
[Table]
total of 4 parameter groups for ResNet-50 and 3 groups for ResNet-18.
Analogously to the CIFAR experiments, we compare SQ to using random orthogonal or DFT matrices for reparameterizing the convolution kernels (ﬁg. 4a).
D ISCUSSION
Existing model compression methods are typically built on a combination of pruning , quantization , or coding .
Pruning involves sparsifying the network either by removing individual parameters or higher level structures such as convolutional ﬁlters, layers, activations, etc.
Various strategies for pruning weights include looking at the Hessian (Cun et al., 1990) or just their (cid:96)
p norm (Han et al., 2016).
Srinivas and Babu (2015) focus on pruning individual units, and H. Li et al.
(2017) prunes convolutional ﬁlters.
Louizos, Ullrich, et al.
(2017) and Molchanov et al.
(2017), which we compare to in our compression experiments, also prune parts of the network.
Dubey et al.
(2018) describe a dimensionality reduction technique specialized for CNN architectures.
Pruning is a simple approach to reduce memory requirements as well as computational complexity, but doesn’t inherently tackle the problem of efﬁciently representing the parameters that are left.
Here, we primarily focus on the latter: given a model architecture and a task, we’re interested in ﬁnding a set of parameters which can be described in a compact form and yield good prediction accuracy.
Our work is largely orthogonal to the pruning literature, and could be combined if reducing the number of units is desired.
[Figure]
i f i c a t i o n
[Figure]
rr
[Figure]
Quantization involves restricting the parameters to a small discrete set of values.
There is work in binarizing or ternarizing networks (Courbariaux et al., 2015; F. Li et al., 2016; Zhou et al., 2018) via either straight-through gradient approximation (Bengio et al., 2013) or stochastic rounding (Gupta et al., 2015).
Recently, Louizos, Reisser, et al.
(2019) introduced a new differentiable quantization procedure that relaxes quantization.
We use the straight-through heuristic, but could possibly use other stochastic approaches to improve our methods.
While most of these works focus on uniform quantization, Baskin et al.
(2018) also extend to non-uniform quantization, which our generalized transformation function amounts to.
Han et al.
(2016) and Ullrich et al.
(2017) share weights and quantize by clustering, Chen, J. Wilson, et al.
(2015) randomly enforce weight sharing, and thus effectively perform VQ with a pre-determined assignment of parameters to representers.
Other works also make the observation that representing weights in the frequency domain helps compression;
Chen, J. T. Wilson, et al.
(2016) randomly enforce weight sharing in the frequency domain and
Wang et al.
(2016) use K-means clustering in the frequency domain.
Coding (entropy coding, or Shannon-style compression) methods produce a bit sequence that can allow convenient storage or transmission of a trained model.
This generally involves quantization as a ﬁrst step, followed by methods such as Huffman coding (Huffman, 1952), arithmetic coding (Rissanen and Langdon, 1981), etc.
Entropy coding methods exploit a known probabilistic structure of the data to produce optimized binary sequences whose length ideally closely approximates the cross entropy of the data under the probability model.
In many cases, authors represent the quantized values directly as binary numbers with few digits (Courbariaux et al., 2015; F. Li et al., 2016;
Louizos, Reisser, et al., 2019), which effectively leaves the probability distribution over the values unexploited for minimizing model size; others do exploit it (Han et al., 2016).
Wiedemann, Marban, et al.
(2018) formulate model compression with an entropy constraint, but use (non-reparameterized) scalar quantization.
Their model signiﬁcantly underperforms all the state-of-the-art models that we compare with (table 1).
Some recent work has claimed improved compression performance by skipping quantization altogether (Havasi et al., 2019).
Our work focuses on coding with quantization.
Han et al.
(2016) deﬁned their method using a four-stage training process: 1.
training the original network, 2.
pruning
This approach has inﬂuenced many follow-up publications.
In the same vein, many current highperforming methods have signiﬁcant complexity in implementation or require a multi-stage training process.
Havasi et al.
(2019) requires several stages of training and retraining while keeping parts of the network ﬁxed.
Wiedemann, Kirchhoffer, et al.
(2019) require pre-sparsiﬁcation of the network, which is computationally expensive, and use a more complex (context-adaptive) variant of arithmetic coding which may be affected by MPEG patents.
These complexities can prevent methods from scaling to larger architectures or decrease their practical usability.
In contrast, our method requires only a single training stage followed by a royalty-free version of arithmetic coding.
In addition, our code is publicly available 4 .
Our method has parallels to recent work in learned image compression (Ballé, Laparra, et al., 2017;
Theis et al., 2017) that uses end-to-end trained deep models for signiﬁcant performance improvements in lossy image compression.
These models operate in an autoencoder framework, where scalar quantization is applied in the latent space.
Our method can be viewed as having just a decoder that is used to transform the latent representation into the model parameters, but no encoder.
C ONCLUSION
We describe a simple model compression method built on two ingredients: joint (i e., end-to-end) optimization of compressibility and task performance in only a single training stage, and reparameterization of model parameters, which increases the ﬂexibility of the representation over scalar quantization, and is applicable to arbitrary network architectures.
We demonstrate that state-of-theart model compression performance can be achieved with this simple framework, outperforming methods that rely on complex, multi-stage training procedures.
Due to its simplicity, the approach is particularly suitable for larger models, such as VGG and especially ResNets.
In future work, we may consider the potential beneﬁts of even more ﬂexible (deeper) parameter decoders