Source-linked AI summary
Soft-to-Hard Vector Quantization for End-to-End Learning Compressible Representations
Eirikur Agustsson, Fabian Mentzer, Michael Tschannen, Lukas Cavigelli, Radu Timofte, Luca Benini, Luc Van Gool
TL;DR
Deep architectures contain redundant features and parameters, but compressing them requires handling nondifferentiable quantization and entropy estimation. The paper introduces a unified end-to-end framework that anneals soft relaxations of quantization and entropy toward discrete compression while jointly learning model parameters and quantization levels. It achieves results competitive with the state of the art for both image compression and DNN compression.
Problem
Compressible deep representations require optimizing rate-distortion objectives despite nondifferentiable quantization and difficult differentiable entropy estimation.
Method
The framework jointly learns model parameters, quantization levels, and symbol-stream entropy using soft assignments annealed gradually to hard quantization.
Results
The method achieves performance competitive with the state of the art for both image compression and DNN compression.
Takeaways & Limitations
One end-to-end soft-to-hard framework captures both feature or image compression and DNN model compression while avoiding fixed quantization outputs.
Takeaways & Limitations
The mini-batch-decomposable soft entropy variant is no longer an upper bound on sample entropy, while the upper-bound variant cannot use mini-batch gradient descent.
Abstract
from arXiv · showhide
We present a new approach to learn compressible representations in deep architectures with an end-to-end training strategy. Our method is based on a soft (continuous) relaxation of quantization and entropy, which we anneal to their discrete counterparts throughout training. We showcase this method for two challenging applications: Image compression and neural network compression. While these tasks have typically been approached with different methods, our soft-to-hard quantization approach gives results competitive with the state-of-the-art for both.
1 Introduction
The paper frames compressible feature and model representations as a rate-distortion problem with differentiability and entropy-estimation challenges. It proposes a unified soft-to-hard, end-to-end framework applied to image and DNN compression.
- Motivation: DNNs contain redundancy in both intermediate feature representations and model parameters, motivating compression for memory-limited inference.Learning compressible feature representations also supports data-adaptive compression algorithms for different data types.
- Problem formulation: Compression requires quantizing parameters or features to representative levels or vectors, then storing assignments and quantization levels.The framework considers scalar and multi-dimensional vector quantization.
- Problem formulation: Minimizing D + βR is difficult because quantization is nondifferentiable and entropy R must be estimated accurately and differentiably.These issues obstruct gradients through entropy and network outputs that depend on quantized variables.
- Proposed framework: The framework jointly optimizes model parameters, quantization levels, and symbol-stream entropy for feature or model compression.It provides a unified treatment of two compression settings previously studied largely independently.
- Proposed framework: Soft assignments with an adjustable hardness parameter gradually become hard during training, making the coding scheme directly differentiable and trainable end-to-end.Unlike rounding-based or stochastic quantization, the approach learns through the coding scheme itself.
- Applications and results: The method learns quantization levels jointly with weights, avoids fixed quantization outputs, and explores vector quantization for learned compression.The authors report benefits over scalar quantization and competitive performance for image and 32-layer ResNet compression.
2 Related Work
Prior work addressed learned image compression and DNN model compression largely as separate problems. Existing approaches include stochastic or autoencoder-based image methods, multi-step model compression pipelines, and continuation-style optimization.
- Image compression: Recent DNN-based image compression methods use convolutional recurrent networks or convolutional autoencoders and outperform JPEG, with some also surpassing JPEG 2000.These approaches include stochastic quantization during training and are evaluated for full-resolution image compression.
- DNN compression: Several DNN compression methods prune pretrained weights, quantize or cluster the remainder, retrain the network, and then apply entropy coding.This is a multi-step pipeline rather than a single unified end-to-end optimization.
- DNN compression: The minimum description length principle offers an alternative DNN compression approach with an information-theoretic foundation.It differs from pruning, retraining, clustering, and entropy-coding pipelines.
- DNN compression: Many low-precision quantization studies target faster hardware evaluation but do not train parameters specifically for information-theoretic compressibility.This separates computational quantization from compression-oriented representation learning.
- Optimization perspective: Soft-to-hard optimization belongs to continuation methods and is formally related to deterministic annealing for maximum-entropy clustering and vector quantization.The paper positions its optimization strategy within these established methodological contexts.
3 Proposed Soft-to-Hard Vector Quantization
The paper formulates compressible feature and parameter representations through a rate–distortion objective and makes quantization and entropy differentiable for end-to-end optimization. Soft assignments, learned quantization centers, and histogram-based entropy estimates are gradually connected to hard coding.
- Problem formulation: Quantization maps continuous features or parameters to discrete symbols and reconstructs them from learned quantization levels or centers.The encoder assigns reshaped feature subvectors to centers, while the decoder maps symbol sequences back to reconstructed vectors.
- Problem formulation: Compressible representations are optimized by balancing expected loss or distortion against the entropy of the encoded symbol stream.For feature representations, β controls the trade-off between compressibility and distortion; for network parameters, the simplified objective remains appropriate.
- Entropy estimation: Entropy is estimated from assignment histograms, with an i.i.d. component assumption for feature symbols and a parameter histogram for DNN compression.For feature outputs, separate sample entropy terms can also be added; for W, histogram probabilities are taken over the parameter vector itself.
- Entropy estimation: The resulting objective remains subject to assumptions and optimization trade-offs: the i.i.d. histogram estimate can upper-bound true entropy, while a mini-batch-friendly soft entropy loses that guarantee.The differentiability and decomposability benefits therefore come with different entropy-bound properties.
- Soft assignments: The method relaxes hard nearest-neighbor encoding and decoding into continuous soft assignments, enabling direct differentiation through the compression pipeline.Soft assignments use a softmax over distances to centers and produce soft quantization before hardening.
- Soft assignments: The hardness parameter σ increases the concentration of assignments so that soft assignments converge to one-hot assignments of the nearest quantization center.This provides the soft-to-hard transition used during training.
4 Image Compression
The image-compression system trains a convolutional autoencoder with soft-to-hard entropy minimization and evaluates its rate–distortion performance against standard codecs. SHA outperforms JPEG and JPEG 2000 at high compression rates and remains competitive with BPG.
- Method: SHA trains a convolutional autoencoder using soft-to-hard entropy minimization rather than the original architecture’s training approach.The method jointly learns network weights and quantization levels while annealing toward hard quantization.
- Method: The bottleneck representation is optimized channel-wise with separate histograms and entropy losses, then encoded as vectorized patches.Each channel uses its own distribution model because bottleneck-channel features are not expected to be identically distributed.
- Evaluation: For rates below 0.4 bpp, SHA outperforms JPEG and JPEG 2000 on MS-SSIM and is competitive with BPG.The evaluation covers Kodak, B100, Urban100, and ImageNET100, with PSNR, SSIM, and MS-SSIM as quality measures.
- Evaluation: SHA-compressed images have fewer artifacts than JPEG 2000 images in the visual examples.Figure 1 combines MS-SSIM-versus-rate curves with a Kodak example reporting rate, MS-SSIM, SSIM, and PSNR.
- Evaluation: The framework achieves competitive image-compression performance using a convolutional autoencoder and simple entropy coding.This contrasts with domain-specific techniques used by several state-of-the-art methods.
5 DNN Compression
For a 32-layer CIFAR-10 ResNet, the method directly learns a compressible weight representation during training. It preserves comparable classification accuracy while achieving roughly 20× compression with standard entropy coding.
- Setup: The experiment compresses all 464,154 trainable parameters of a 32-layer ResNet trained for CIFAR-10.The parameters are concatenated into one vector and scalar-quantized using 75 centers.
- Results: 92.1% test accuracy is retained while compression reaches 19.15× with Huffman coding and 20.15× with arithmetic coding.The original pretrained model achieves 92.6% accuracy, and the selected setting targets approximately 20× compression.
- Comparison: The approach directly minimizes weight entropy during training instead of using pruning, retraining, clustering, and specialized sparse-matrix encoding steps.The resulting model is encoded with standard entropy coding.
- Optimization: During optimization, the sample entropy decays as most weights condense around a few quantization centers.The index histograms develop this concentration during training.
6 Conclusions
The paper presents a unified end-to-end framework that anneals soft entropy and discretization toward hard quantization while optimizing the rate–distortion trade-off. It achieves competitive results for both image and DNN compression.
- Conclusion: The framework jointly optimizes network parameters, quantization levels, and entropy for compressed deep-architecture representations.It targets either feature representations or the model itself.
- Conclusion: Soft-to-hard annealing transfers training from differentiable relaxations of entropy and discretization to nondifferentiable quantization.This enables optimization of the trade-off between network loss and entropy.
- Conclusion: The framework obtains results competitive with the state-of-the-art for image compression and DNN compression.The conclusion presents this as the supported outcome across both applications.
A.1 Architecture
The image-compression architecture uses convolutional encoder and decoder networks connected through a lower-dimensional bottleneck tensor. The decoder mirrors the encoder with upsampling and deconvolutions.
- Encoder: The encoder downsamples an input image to a w/8 × h/8 × c bottleneck tensor.Two initial convolutional layers reduce spatial dimensions by a factor of four, followed by another factor-of-two downsampling layer.
- Encoder: The encoder increases channels from 3 to 128, processes them through three residual blocks, and reduces the final channel count to c.Each residual block uses 128 filters.
- Decoder: The decoder mirrors the encoder by mapping the bottleneck tensor back to a w×h output image.It uses upsampling and deconvolutions instead of the encoder’s downsampling and convolutions.
- Decoder: The decoder uses standard deconvolutions rather than the subpixel layers used in the referenced architecture.The stated reason is simplicity.
A.2 Hyperparameters
The image-compression experiments vary β and c to explore rate-distortion trade-offs, using a two-stage training schedule with specified optimization settings.
- Five representative (c, β) pairs were selected to span approximately 0.1–0.8 bpp.The models use vector quantization with L = 1000 centers and 2 × 2 patches.
- β controls the entropy penalty and directly controls the resulting bits per pixel.
- Training uses an initial 250k-iteration stage followed by 800k iterations of annealing.The learning rates are 1e−4 and 1e−5, respectively, with annealing temperature T = 50k and KG = 100.
A.3 Effect of Vector Quantization and Entropy Loss
The ablation compares vector and scalar quantization, plus training without entropy loss, to assess their effects on image-compression performance.
- The scalar-quantization comparison uses 1 × 1-dimensional centers with L = 6 to approximate the number of assignments from 2 × 2 patches.The figure compares Vector, Scalar, and 2 × 2 centers with β = 0 on ImageNET100.
- The entropy-loss ablation trains 2 × 2-center models with β = 0.
- Vector quantization and entropy loss each produce higher compression rates at a given reconstruction MSE than their respective alternatives.The alternatives are scalar quantization and training without entropy loss.
A.4 Effect of Annealing
The evaluation examines image-quality measures and entropy-related training behavior across rates, datasets, and training iterations.
- Figure 3 tracks entropy loss, soft and hard PSNR, gap(t), and σ as functions of iteration t.
- PSNR, SSIM, and MS-SSIM quantify similarity between compressed and uncompressed images, with SSIM measures designed to correlate better with human perception than PSNR.
- Average MS-SSIM, SSIM, and PSNR are plotted against rate for ImageNET100, Urban100, B100, and Kodak.
- Sample entropy decreases from 4.07 at t = 216 to 2.90 bits per weight at t = 475 and 1.58 bits per weight at t = 520.The accompanying histograms show the corresponding evolution of index frequencies during training.