Source-linked AI summary

Training with Quantization Noise for Extreme Model Compression

Angela Fan, Pierre Stock, Benjamin Graham, Edouard Grave, Remi Gribonval, Herve Jegou, Armand Joulin

arXiv:2004.07320v3cs.LGstat.ML

TL;DR

The paper addresses how to preserve model accuracy under extreme compression, where standard QAT and STE are inadequate. It introduces Quant-Noise, which randomly quantizes weight subsets during training, and reports strong accuracy–size trade-offs across language and vision benchmarks.

  • Problem

    Extreme compression methods such as Product Quantization make STE approximations severe, limiting standard Quantization Aware Training.

  • Method

    Quant-Noise randomly quantizes different subsets of weights during each forward, allowing unbiased gradients through unaffected weights and supporting multiple quantization methods.

  • Results

    The method establishes new state-of-the-art accuracy–model-size trade-offs across NLP and image classification, including 82.5% MNLI accuracy for RoBERTa at 14 MB and 80.0% ImageNet top-1 accuracy for EfficientNet-B3 at 3.3 MB.

  • Takeaways & Limitations

    Quant-Noise supports extreme compression with int4, int8, Product Quantization, and combined PQ–int8 schemes, and can also improve already trained networks after quantization.

Abstract

from arXiv · show

We tackle the problem of producing compact models, maximizing their accuracy for a given model size. A standard solution is to train networks with Quantization Aware Training, where the weights are quantized during training and the gradients approximated with the Straight-Through Estimator. In this paper, we extend this approach to work beyond int8 fixed-point quantization with extreme compression methods where the approximations introduced by STE are severe, such as Product Quantization. Our proposal is to only quantize a different random subset of weights during each forward, allowing for unbiased gradients to flow through the other weights. Controlling the amount of noise and its form allows for extreme compression rates while maintaining the performance of the original model. As a result we establish new state-of-the-art compromises between accuracy and model size both in natural language processing and image classification. For example, applying our method to state-of-the-art Transformer and ConvNet architectures, we can achieve 82.5% accuracy on MNLI by compressing RoBERTa to 14MB and 80.0 top-1 accuracy on ImageNet by compressing an EfficientNet-B3 to 3.3MB.

1 INTRODUCTION

Large neural networks are difficult to deploy because their parameter memory remains substantial. Quant-Noise extends training-time quantization to severe compression by randomly quantizing weight subsets, preserving unbiased gradients and enabling strong accuracy–size trade-offs.

  • Large Transformer and EfficientNet models can require millions of parameters and dozens to hundreds of megabytes, limiting deployment in robotics and virtual assistants.
  • Quantization reduces bits per weight, but approximation errors can accumulate through the forward pass and significantly degrade performance.
  • QAT quantizes all weights during training and uses STE, but its gradient approximation is insufficient when compression errors are severe.
  • Quant-Noise randomly quantizes only a subset of weights at each forward, allowing unbiased gradients through most unaffected weights.
  • 82.5% accuracy on MNLI is achieved by compressing RoBERTa to 14 MB, while ImageNet reaches 80.0% top-1 accuracy with EfficientNet-B3 compressed to 3.3 MB.

2 RELATED WORK

Related compression work reduces model size through pruning, weight sharing, lightweight architectures, distillation, or quantization. Quantization is distinctive because it reduces bits per weight and can accelerate inference on suitable hardware.

  • Pruning removes individual weights or structures and can reduce both memory footprint and inference time.
  • Lightweight architectures modify existing models or design new networks such as MobileNet, ShuffleNet, and EfficientNet.
  • Knowledge distillation has been applied to sentence representation and to reducing the size of BERT models.
  • Scalar quantization studies low-precision weights and activations, often benefiting from specialized hardware that improves inference runtime.
  • Vector Quantization and Product Quantization quantize groups of weights jointly to achieve higher compression.

3 QUANTIZING NEURAL NETWORKS

Quantization replaces weight blocks with codebook entries and indices, trading representation precision for compact storage. The section covers scalar, vector, Product, iterative, and fixed-point Product Quantization methods.

  • Quantization stores a codebook and block indices instead of real weights, then reconstructs each block from its assigned codeword during inference.
  • Scalar quantization operates on single weights, whereas vector quantization jointly quantizes groups of weights.
  • FIXED-POINT SCALAR QUANTIZATION: Uniform fixed-point quantization rounds weights and activations to N-bit representations, considering int4 and int8 in this work.
  • PRODUCT QUANTIZATION: Product Quantization splits columns into subvectors, shares one codebook across them, and assigns each subvector to a nearest centroid.
  • PRODUCT QUANTIZATION: PQ shares subvector representations, allowing higher compression rates than intN.
  • ITERATIVE PRODUCT QUANTIZATION: Iterative PQ quantizes layers sequentially and fine-tunes upper layers under an uncompressed teacher to reduce activation drift.
  • FIXED-POINT PRODUCT QUANTIZATION: Combining PQ with int8 compresses centroids and activations into fixed-point representations while retaining PQ’s compression benefits.

4 METHOD

Quant-Noise trains compressed networks by applying target quantization distortions to randomly selected weight blocks during each forward pass, preserving unbiased gradients through unaffected weights. The framework supports scalar, product, combined, and pruning-related noise, with experiments assessing compression schemes on language modeling and image classification.

  • Training networks with quantization noise: Quant-Noise randomly selects weight blocks for distortion during each forward pass instead of quantizing the entire network.It applies noise associated with a target quantization method to selected blocks during training.
  • Training networks with quantization noise: Unbiased gradients flow through unaffected blocks, while each weight regularly receives gradients without the distortion function.The method uses STE during backpropagation but preserves undistorted gradient paths for randomly unselected blocks.
  • Composing quantization noises: Quantization noise operators can be composed to combine multiple quantization methods or incorporate pruning during training.The paper explicitly discusses combining scalar quantization with product quantization and composing quantization with pruning.
  • Experimental evaluation: Table 1 compares quantization schemes with and without Quant-Noise using Transformer language modeling and EfficientNet-B3 image classification.Reported quantities include perplexity or top-1 accuracy, model size in MB, and compression ratio.
  • Fixed-point scalar quantization: For intN quantization, the noise function rounds atomic weights using trainable scale and zero-point parameters.The application to int8 is described as a stochastic amelioration of QAT.
  • Product quantization: For PQ, the noise function assigns selected blocks to nearest learned codewords, while a faster proxy zeroes selected subvectors.Online codebook learning is supported, and the proxy avoids the computational cost of exact PQ noise.

5 RESULTS

Quant-Noise improves quantized-model performance across language and vision settings, including high-compression iPQ, and remains useful for finetuning and combinations with pruning or sharing.

  • 5 RESULTS: Quant-Noise improves performance across quantization schemes in Transformer language modeling and EfficientNet image classification settings.It targets int8/int4 and iPQ, with particularly strong benefits for high-performance iPQ quantization.
  • 5 RESULTS: iPQ provides the best trade-off between model performance and compression, with little negative impact on FLOPS.The comparison indicates that QAT can degrade iPQ performance, even relative to postprocessing quantization.
  • 5 RESULTS: Combining Quant-Noise with pruning and weight sharing yields additional performance-versus-size operating points across language modeling, MNLI, and ImageNet.The conclusions are consistent across the reported tasks and benchmarks.
  • 5 RESULTS: ×20−94 memory reductions are achieved on competitive models when the best quantization setup is combined with weight sharing and pruning.The reported comparisons cover language modeling, text classification, and image classification against prior compression methods.
  • 5 RESULTS: 78.5% top-1 accuracy is retained after compressing EfficientNet-B3 from 46.7Mb to 3.3Mb, a ×14 reduction.The quantized model is reported as smaller and more accurate than MobileNet or ShuffleNet in Figure 2.
  • 5 RESULTS: Finetuning with Quant-Noise nearly matches training from scratch, with a 0.2 PPL difference in language modeling and effective RoBERTa compression during MNLI finetuning.The procedure applies Quant-Noise after starting from existing models.

6 CONCLUSION

The paper concludes that random-subset quantization preserves performance in high-quantization regimes and supports diverse quantization, compression, and post-processing settings.

  • 6 CONCLUSION: Quant-Noise maintains performance in the high quantization regime by quantizing a random subset of weights during training.The framework is validated across text and vision applications and multiple quantization schemes.
  • 6 CONCLUSION: Quant-Noise combines iPQ and int8 to provide extreme compression with fixed-point arithmetic for weights and activations.The method can also prepare already trained networks for later quantization through post-processing.
  • 6 CONCLUSION: Post-processing with Quant-Noise improves the performance of compressed models derived from already trained networks.This extends the method beyond training models from scratch.

7 APPENDIX

Appendix experiments examine architecture-specific noise rates and proxy noise functions, finding limits for iPQ noise rates and practical equivalence of approximations.

  • 7 APPENDIX: Quant-Noise provides consistent additional improvement over iPQ compression for ResNet-50.The comparison uses results from Stock et al. (2019).
  • 7.3 IMPACT OF NOISE RATE: For iPQ, the proxy noise function degrades Transformer performance when the noise rate exceeds 0.5.At rates above 0.5, more than half of the weights pass through the proxy noise function.
  • 7.3 IMPACT OF NOISE RATE: For int8 quantization, higher noise rates are only slightly worse, while rate 1 corresponds to QAT with the full matrix quantized using STE.This comparison demonstrates the potential benefit of partial quantization during training.
  • 7.4 IMPACT OF APPROXIMATING THE NOISE FUNCTION: The exact and approximate iPQ noise functions do not affect the performance of the quantized models in the reported comparisons.The approximation zeroes selected vectors and can optionally account for cluster assignments.
  • 7.4 IMPACT OF APPROXIMATING THE NOISE FUNCTION: The proxy approximation is faster and more practical because it avoids computing cluster assignments and centroids after every epoch.It also requires less involved code.

7.5 EXPERIMENTAL SETTING

Experiments evaluate Quant-Noise on language modeling, sentence representation transfer, and image classification using Transformer, BERT/RoBERTa, and EfficientNet models.

  • 7.5 EXPERIMENTAL SETTING: Language-modeling experiments report test-set perplexity for a 16-layer Transformer trained on WikiText-103.The benchmark contains 100M tokens and a 260k-word vocabulary.
  • 7.5 EXPERIMENTAL SETTING: Sentence-representation experiments pretrain base BERT on BooksCorpus + Wiki and report MNLI accuracy after finetuning.Training follows the parameters in Liu et al. (2019).
  • 7.5 EXPERIMENTAL SETTING: The Classy Vision EfficientNet-B3 baseline achieves 81.5% top-1 accuracy, compared with 81.9% reported by Tan and Le.The model is trained on the ImageNet object-classification benchmark.

7.6 TRAINING DETAILS

Training details specify task-specific architectures, optimization settings, quantization configurations, and evaluation metrics for language modeling, RoBERTa, and EfficientNet.

  • Language Modeling: Wikitext-103 language modeling uses adaptive softmax and adaptive input with 1024-dimensional embeddings and three adaptive vocabulary bands.The bands contain 20K, 40K, and 200K entries; training uses cosine learning-rate scheduling and Nesterov acceleration.
  • RoBERTa: RoBERTa uses a 12-layer architecture with embedding size 768, FFN size 3072, large-batch training, Adam, polynomial decay, and LayerDrop 0.2.Quant-Noise is set to 0.1, with block size 8 during training.
  • Evaluation: Evaluation reports Wikitext-103 test perplexity and model size, while classification reports accuracy and model size in megabytes.Lower perplexity is better; higher classification accuracy is better.
  • RoBERTa: RoBERTa finetuning searches learning rates of 1e-5, 2e-5, and 3e-5 with batch sizes of 16 or 32 sentences.Finetuning is single-task, uses no ensembling, and retains LayerDrop and Quant-Noise when applicable.
  • EfficientNet: EfficientNet-B3 uses Quant-Noise 0.1 after searching values 0.05, 0.1, and 0.2.Training block sizes are 4 for 1 × 1 convolutions and the classifier, 9 for depth-wise 3 × 3 convolutions, and 5 for depth-wise 5 × 5 convolutions.
  • Quantization: Scalar quantization emulates PyTorch 1.4 by quantizing weights and activations, with fixed activation scales and zero points estimated from preliminary forward passes.The Histogram method approximately minimizes L2 quantization error and refines MinMax quantization.

7.8 IPQ QUANTIZATION DETAILS

The iPQ quantization setup uses task- and structure-specific block sizes, while the reported classification tables measure accuracy and model size under the compression experiments.

  • Language Modeling: Language modeling uses block size 8 for FFN and embeddings and block size 4 for attention.Attention block size is tuned between 4 and 8, and Quant-Noise is applied to all layers during training.
  • RoBERTa: RoBERTa uses block size 4 for FFN, embeddings, and attention.Block size is tuned between 4 and 8, with Quant-Noise applied to all layers during training.
  • EfficientNet: EfficientNet quantizes blocks sequentially and finishes with the classifier, using block sizes 4, 9, 5, and 4 across its listed structures.Quant-Noise is applied to weights in InvertedResidual blocks except Squeeze-Excitation subblocks, the head convolution, and the classifier.
  • Evaluation: MNLI and ImageNet tables report accuracy and model size in megabytes, with higher accuracy preferred.The MNLI table also marks distillation and data augmentation where applicable.
  • Layer Sharing and Pruning: Layer sharing and pruning can be combined by sharing adjacent layers in chunks of two and pruning selected shared chunks.The described example shares A–B, C–D, E–F, and G–H, then prunes alternating chunks.

7.10 NUMERICAL RESULTS FOR GRAPHICAL DIAGRAMS

The paper reports numerical values from graphical results in separate tables for language modeling, BERT, and ImageNet.

  • Numerical Results: Numerical values displayed in Figure 2 are reported in Table 6 for language modeling, Table 7 for BERT, and Table 8 for ImageNet.

7.11 FURTHER ABLATIONS

Further ablations examine Quant-Noise rates, centroid counts, model size, structure order, structure-specific compression, quantizer variants, and STE use with pruning noise.

  • Quant-Noise Rate: Quant-Noise rate p is evaluated with scalar int8 quantization on EfficientNet for ImageNet classification.The study reports the influence of the noise rate on the experiment.
  • Number of Centroids: Using more centroids improves perplexity, while 256 centroids balance model size and representation capacity.The number of centroids can be adjusted according to practical storage constraints.
  • Initial Model Size: As layers are removed, shallow models do not show an increasing quantized-versus-non-quantized perplexity gap, unlike models with smaller FFN sizes.Smaller FFNs have less redundant weights and are more difficult to quantize with iPQ.
  • Quantization Order: Quantization order causes small performance variations, so Transformer structures are quantized in the order FFN, embeddings, then attention matrices.The ordering is motivated by accumulated reconstruction error and later centroid adjustment during finetuning.
  • Model Structures: FFN and embedding structures tolerate more aggressive compression, whereas attention loses performance as block sizes increase.The comparison varies one structure’s block size from 4 to 32 while keeping two structures fixed.
  • Quantizer Variants: Per-channel quantization learns channel-wise scale and offset vectors, providing more flexibility than a histogram quantizer.The histogram quantizer maintains running minima and maxima while minimizing L2 distance.
  • Pruning Noise: Applying STE to LayerDrop pruning noise produces slightly worse results.The comparison is reported for Wikitext-103.
Loading 2004.07320v3…