Differentiable Model Compression via Pseudo Quantization Noise

TL;DR

<3-5 sentence high-level summary> DiffQ addresses the challenge of compressing large neural networks for on-device deployment by replacing gradient-approximation-based quantization (STE) with a differentiable pseudo quantization noise (PQN) framework. It enables end-to-end optimization of per-group bitwidth via a single penalty parameter, yielding automatic mixed-precision quantization without STE bias. Across language modeling, audio source separation, and image classification, DiffQ achieves substantial model-size reductions while maintaining accuracy close to or matching uncompressed baselines, and often outperforming STE-based methods like QAT and LSQ. The approach also provides insight into layer-wise bitwidth distributions and highlights practical considerations such as training overhead and entropy-aware sizing for deployment.

Abstract

We propose DiffQ a differentiable method for model compression for quantizing model parameters without gradient approximations (e.g., Straight Through Estimator). We suggest adding independent pseudo quantization noise to model parameters during training to approximate the effect of a quantization operator. DiffQ is differentiable both with respect to the unquantized weights and the number of bits used. Given a single hyper-parameter balancing between the quantized model size and accuracy, DiffQ optimizes the number of bits used per individual weight or groups of weights, in end-to-end training. We experimentally verify that our method is competitive with STE based quantization techniques on several benchmarks and architectures for image classification, language modeling, and audio source separation. For instance, on the ImageNet dataset, DiffQ compresses a 12 layers transformer-based model by more than a factor of 8, (lower than 4 bits precision per weight on average), with a loss of 0.3% in model accuracy. Code is available at github.com/facebookresearch/diffq.

Figures (4)

• Figure 1: (a) Using STE and SGD to optimize the 1D least-mean-square problem given by equation \ref{['counter_example_ste']} (with $B=4$ and $X=1$ a.s.). $\mathbf{Q}(w_n, B)$ oscillates between the quantized value just above ($w_+$) and just under ($w_-$) the unquantized ground truth $w_*$, while $w_n$ oscillates around the boundary $(w_+ + w_-) / 2$. (b) Model accuracy vs. epochs for ImageNet using EfficientNet-b3. Results are presented for both QAT over 4 bits and DiffQ.
• Figure 2: Model accuracy and size on CIFAR10 (a) and CIFAR100 (b) using MobileNet, ResNet-18, and WideResNet (WRN) models for various penalty levels using DiffQ, QAT, LSQ, and the baseline.
• Figure 3: We group layers of a given architecture into 11 groups (group 0 being closest to the input, and 10 closest to the output), and report for each group its contribution to the model size. We compare the baseline EfficientNet-B3 (above) and DeiT (below) models (floating point 32 bits) and the quantized models with DiffQ ($\lambda{=}5\mathrm{e}{-}3$ for EfficientNet-B3, $\lambda{=}1\mathrm{e}{-}2$ for DeiT). For quantized model, we also report the distribution over each bitwidth within each group of layers. Scale is logarithmic across layers, and linear inside each one. Finally, "overhead" shows the capacity needed to encode the bitwidth used for each group of weights.
• Figure B.1: (a): DiffQ results with various groups sizes ($g \in \{1, 4, 8, \infty \}$). $g=\infty$ refers to a single group for the entire layer. For reference, we report the accuracy of the uncompressed model (42.8 MB). Models are Resnet-18 trained on CIFAR-100. (b): ImageNet results using EfficientNet-B3 model. We plot the model size vs. model accuracy using different penalty levels. We additionally, present the uncompressed models (uncomp.) and Quantization Aware Training (QAT) using 4 and 8 bits.