Source-linked AI summary
Soft Weight-Sharing for Neural Network Compression
Karen Ullrich, Edward Meeds, Max Welling
TL;DR
Large neural networks are difficult to run on mobile devices because of their computational, memory, and energy demands, motivating compression. The paper applies soft weight-sharing with a learned mixture prior to jointly encourage pruning and quantization during retraining. It reports competitive and state-of-the-art compression rates, while identifying computational cost and scalability as limitations.
Problem
Mobile deployment is challenged by neural networks’ computational, memory, and energy intensity, motivating compression methods that preserve accuracy.
Method
The method learns mixture-prior parameters jointly with network weights, using soft weight-sharing to encourage pruning and quantization in one retraining procedure.
Results
The approach achieves competitive compression rates, including state-of-the-art rates on MNIST networks and a 64.2 compression rate with accuracy changing from 0.9811 to 0.9806.
Takeaways & Limitations
Soft weight-sharing provides a simple, principled compression approach that optimizes an MDL complexity lower bound while working across different models.
Takeaways & Limitations
The method is computationally costly, and for VGG-scale networks with 138M parameters it is too slow to produce usable results without experimental validation of the proposed solution.
Abstract
from arXiv · showhide
The success of deep learning in numerous application domains created the de- sire to run and train them on mobile devices. This however, conflicts with their computationally, memory and energy intense nature, leading to a growing interest in compression. Recent work by Han et al. (2015a) propose a pipeline that involves retraining, pruning and quantization of neural network weights, obtaining state-of-the-art compression rates. In this paper, we show that competitive compression rates can be achieved by using a version of soft weight-sharing (Nowlan & Hinton, 1992). Our method achieves both quantization and pruning in one simple (re-)training procedure. This point of view also exposes the relation between compression and the minimum description length (MDL) principle.
1 INTRODUCTION
Deep networks are often over-parameterized and difficult to deploy on mobile devices, motivating compression methods that reduce resource demands while preserving accuracy. This work learns a compression-friendly prior through soft weight-sharing, jointly encouraging weights to cluster and remain predictive.
- Mobile deployment is constrained by neural networks’ memory and energy demands, motivating compression.Compression may also reduce bandwidth and speed processing while preserving privacy.
- Large networks can be compressed by approximately two orders of magnitude without significant accuracy loss.The paper presents this over-parameterization finding as motivation for compression algorithms.
- The method learns the prior used to encode network parameters rather than fixing it.This is framed as an empirical Bayes direction.
- A mixture-of-Gaussians prior encourages weights to concentrate around K cluster components while cluster centers preserve predictive accuracy.The approach builds on soft weight-sharing and targets compression directly.
- Compression requires encoding K cluster means in full precision plus each weight’s cluster assignment using log(J) bits per weight.The authors report competitive compression rates from this scheme.
2 MDL VIEW ON VARIATIONAL LEARNING
The paper frames compression as an MDL and variational-inference problem: encode both the model and its data misfit, then learn a prior that lowers the weight-encoding cost. It uses a mixture-of-Gaussians prior to encourage compressible weights while retaining a practical coding scheme.
- MDL selects hypotheses by minimizing model complexity cost and data misfit cost.Variational learning can be reinterpreted as an MDL problem.
- The data error cost lower-bounds information needed to communicate targets when inputs and model parameters are known.Its likelihood form depends on the target distribution.
- The model complexity cost lower-bounds the expected cost of communicating parameters under an agreed prior.This cost is connected to the prior and entropy terms in variational coding.
- Bits-back coding can refund information equal to the entropy of a noisy weight distribution, with the Bayesian posterior being optimal.The scheme is theoretically optimal but often impractical in real-world settings.
- The paper optimizes cross-entropy between a delta posterior and a mixture-of-Gaussians prior as a lower bound on average weight-encoding bits.Coarser quantization makes the model more compressible, but practical bits-back coding relies on independence and a known prior.
- The method learns mixture means, variances, and mixing proportions simultaneously with network weights, while a zero component can enforce pruning.The approach naturally encourages quantization through concentration around mixture components.
3 RELATED WORK
Prior compression work uses pruning, quantization, structural sparsity, and low-rank decomposition to reduce neural-network storage or computation. Han et al.’s multi-step pipeline combines pruning, quantization, and Huffman coding, whereas this paper targets its restrictions with a unified approach.
- Pruning exploits parameter redundancy, and neural networks can tolerate up to 99% weight pruning without significant accuracy loss.Structural pruning can remove entire weight groups and speed convolutional prediction.
- Quantization reduces storage by lowering bits per weight; converting 32-bit floats to 1 bit yields a 32× storage improvement.Prior methods include binary, k-means, product, and residual quantization.
- Han et al. combine pruning, quantization, and Huffman encoding, achieving 35−49× final compression for CNNs.Their stages provide 9−13×, 27−31×, and 35−49× compression respectively.
- Han et al.’s pipeline fixes zero assignments after pruning and cluster assignments after quantization, restrictions the proposed approach seeks to overcome.
- Low-rank matrix decomposition compresses networks through factorization, including low-rank filters and combinations with sparsity.
4 METHOD
The method retrains pretrained networks with soft weight-sharing, jointly learning weights and Gaussian-mixture parameters before quantizing and compressing the result.
- Network re-training: Mixture components are initialized around pretrained weights, with one fixed zero component and trainable means, variances, and mixing proportions for the others.The zero component can instead have a trainable mixing proportion constrained by a Beta hyper-prior.
- Post-processing: After retraining, each weight is quantized to the mean of its most responsible mixture component.Redundant components are identified using pairwise KL divergence and merged below a threshold.
- Post-processing: The compressed representation uses learned cluster means and weight assignments, with practical storage based on the format used by Han et al. (2015a).Assignments encode which mixture mean represents each weight.
- Network re-training: The algorithm alternates gradient updates of network weights and mixture parameters during retraining, using an error-plus-complexity objective.The trade-off parameter τ controls the log-prior contribution, and the procedure uses Adam.
- Overall procedure: Algorithm 1 divides soft weight-sharing compression into network re-training followed by post-processing.The optional hyper-prior parameters control additional variance and mixing-proportion constraints.
5 MODELS
The experiments use two MNIST models and a light ResNet architecture, with baseline models trained or specified using reported error rates and parameter settings.
- MNIST models: LeNet-300-100 is an MNIST model for which the authors trained a model with a 1.89% error rate.No pretrained model was available, so the authors trained their own model.
- MNIST models: LeNet-5-Caffe is a modified LeNet-5 MNIST model whose authors’ trained baseline has a 0.88% error rate.The model specification comes from the Caffe MNIST tutorial.
- ResNet model: The ResNet experiment uses a depth-16, width-4 network without dropout, containing 2.7M parameters.The experiments use light augmentation and normalized data; reported errors are 5.02% on CIFAR-10 and 24.03% on CIFAR-100.
6 EXPERIMENTS
Experiments show that soft weight-sharing can prune and quantize networks while preserving accuracy, achieve competitive compression results, and expose optimization and scalability limitations.
- Initial experiment: 96% of the first-layer parameters, 90% of the second-layer parameters, and 18% of the final-layer parameters are pruned.The authors conclude pruning depends mainly on the number of parameters in a layer rather than its position or type.
- Initial experiment: 64.2 is the reported compression rate, while accuracy changes from 0.9811 to 0.9806.The same experiment finds the retrained distribution sharply peaked around zero.
- Initial experiment: Figure 1 compares pretrained and retrained weight distributions and shows each weight’s change through a scatter plot.The retrained distribution is sharply peaked around zero.
- Hyper-parameter tuning: Spearmint optimizes 13 hyper-parameters against an objective balancing accuracy loss and compression rate, with results plotted for LeNet-300-100 and LeNet-5-Caffe.The optimization uses 40 training epochs, or 18K updates, because of computational restrictions.
- Mixture dynamics: Figure 3 shows 17 initial Gaussian-mixture components collapsing into roughly 6 significant components and a final prior that is sharply peaked for quantization.The figure also compares initial and final weights with mixture-component uncertainties.
- Compression results: The MNIST comparisons achieve state-of-the-art compression rates, and the method is demonstrated on a 2.7M-parameter ResNet.For VGG with 138M parameters, the algorithm is too slow for usable results, and the proposed remedy lacks experimental results.
7 DISCUSSION AND FUTURE WORK
The paper applies soft weight-sharing directly to model compression, optimizing an MDL complexity lower bound while working across different models. It identifies computational cost and ease of implementation as current shortcomings, with broader priors and structured pruning proposed for future work.
- Soft weight-sharing is applied directly to model compression while optimizing the MDL complexity lower bound.
- The method works in practice when applied to different models.
- The method’s current shortcomings are computational cost and ease of implementation.
- Future work includes extending beyond Dirac posteriors and learning to prune entire structures using Bernoulli priors.
A REVIEW OF STATE-OF-THE-ART NEURAL NETWORK COMPRESSION
The reviewed compression scheme stores sparse weights using CSR arrays and improves storage by reducing index precision, encoding differences, and optimizing the arrays.
- CSR format stores the weights in three vectors: A, IR, and IC, instead of storing all parameters directly.
- A stores all non-zero entries, while IR recursively records row boundaries and IC stores each entry’s column index.
- The format is illustrated with an example before the compression-rate calculation.
- The compression rate can be significantly improved by optimizing each of the three arrays.
- IR can use reduced precision pprun because its largest value, the number of non-zero weights, is much smaller than 2^porig.
A.2 STORING THE INDEX ARRAY IC
The index array IC is compressed by storing relative index differences and inserting zeros when gaps exceed the chosen span, with Huffman encoding applied to the indexes.
- Relative differences between indexes reduce the range of values that IC must represent.
- When the distance between non-zero weights exceeds the span of 2^pprun, A is filled with zeros to shrink the utilized range.
- Han et al. propose p = 5 for fully connected layers and p = 8 for convolutional layers.
- The indexes are further compressed using Huffman encoding.
- Figure 4 illustrates IC as relative indexes and shows an additional zero when a relative index is larger than 8 (= 2^3).
B.1 GAMMA DISTRIBUTION
The Gamma distribution is used as a prior for Gaussian precision and is defined over positive precision values. The experiments examine how its variance choices affect the distribution.
- The Gamma distribution is the conjugate prior for the precision of a univariate Gaussian and is defined for λ > 0.
- Its mode is characterized by λ∗ = α − 1.
- The experiments set the desired variance of mixture components to 0.05, corresponding to λ∗ = 400.
- Figure 5 shows a Gamma distribution with λ∗ = 100, while α and β represent different variance choices.
B.2 BETA DISTRIBUTION
The section introduces a Beta distribution whose α and β parameters represent different pseudo-count choices while π*_{j=0}=0.9. It also places the visualization alongside stochastic-gradient optimization and compressed filter examples.
- B.2 BETA DISTRIBUTION: π*_{j=0}=0.9 fixes the Beta distribution’s zero-component probability while α and β vary the pseudo-count choice.The figure caption identifies these parameters and the fixed probability.
- B.2 BETA DISTRIBUTION: Batch gradient descent updates network parameters iteratively using minibatches to approximate the posterior derivative.The approximation can then be used in update schemes such as simple gradient descent.
- B.2 BETA DISTRIBUTION: Filter visualizations compare pre-trained and compressed representations for LeNet-5-Caffe and LeNet-300-100.The LeNet-5-Caffe figure separates first- and second-layer convolution filters, while the LeNet-300-100 figure contrasts pre-trained and compressed feature filters.