[Paper Review] Training with Quantization Noise for Extreme Model Compression

Angela Fan, Pierre Stock+5 more|arXiv (Cornell University)|Apr 15, 2020

Advanced Neural Network Applications72 references114 citations

TL;DR

The paper introduces Quant-Noise, a regularization technique that quantizes only a random subset of weights during training to build models resilient to extreme quantization (e.g., int4/int8 and product quantization). This yields state-of-the-art accuracy under heavy compression, including NLP and image tasks.

ABSTRACT

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.

Motivation & Objective

• Motivate extreme model compression without large accuracy loss.
• Develop a training-time mechanism that makes networks robust to various quantization schemes.
• Enable combinations of scalar quantization, product quantization (PQ/iPQ), and fixed-point arithmetic.
• Investigate Quant-Noise as a post-processing step to improve quantized models without full retraining.

Proposed method

• Quant-Noise selects a random subset of weight blocks during each forward pass and applies a distortion that mimics the target quantization.
• Distortion functions include fixed-point scalar quantization and product quantization (with a proxy noise for PQ/iPQ).
• Backpropagation uses straight-through estimation (STE) for gradients on distorted weights, while unnoised blocks receive unbiased gradients.
• Quant-Noise can be composed with pruning or layer-drop to simulate pruning and structured sparsity during training.
• When using PQ, the noise can be implemented via a proxy that zeros selected subvectors, encouraging useful subvector correlations.

Experimental results

Research questions

• RQ1 Can training with random quantization noise yield models robust to extreme quantization (int4/int8, PQ/iPQ) without large accuracy loss?
• RQ2 Does Quant-Noise improve the accuracy vs. standard QAT under extreme compression regimes?
• RQ3 Can Quant-Noise enable effective post-training quantization improvements on already trained models?
• RQ4 What are the best combinations of PQ/iPQ with fixed-point quantization and pruning for NLP and vision tasks?

Key findings

• Quant-Noise improves performance across quantization schemes (int4, int8, PQ/iPQ) on NLP (RoBERTa-based) and vision (EfficientNet-B3) tasks.
• On NLP (MNLI with RoBERTa), RoBERTa compressed to 14 MB achieves 82.5% accuracy without Quant-Noise, and 83.6% with Quant-Noise during training; post-training finetuning with Quant-Noise yields 83.6%.
• On ImageNet with EfficientNet-B3, compression to 3.3 MB achieves 80.0% top-1 accuracy under Quant-Noise, versus 78.5% for the unquantized baseline after compression.
• iPQ + Quant-Noise achieves 80.0% top-1 on ImageNet with 3.3 MB and 79.8% when combining iPQ with int8 and Quant-Noise; PQ-based methods reach strong compression with minimal accuracy loss.
• Quant-Noise enables extreme compression ratios (e.g., up to ×94 with pruning and sharing in NLP) while maintaining competitive perplexities/accuracy compared to uncompressed models.