[Paper Review] Bayesian Bits: Unifying Quantization and Pruning

Mart van Baalen, Christos Louizos+5 more|arXiv (Cornell University)|May 14, 2020

Advanced Neural Network ApplicationsComputer Science39 references77 citations

TL;DR

Bayesian Bits jointly learns mixed-precision quantization and pruning via gradient-based optimization using a novel residual-quantization decomposition with learnable gates, achieving hardware-friendly bit-widths and better accuracy-efficiency trade-offs than static bit-width baselines.

ABSTRACT

We introduce Bayesian Bits, a practical method for joint mixed precision quantization and pruning through gradient based optimization. Bayesian Bits employs a novel decomposition of the quantization operation, which sequentially considers doubling the bit width. At each new bit width, the residual error between the full precision value and the previously rounded value is quantized. We then decide whether or not to add this quantized residual error for a higher effective bit width and lower quantization noise. By starting with a power-of-two bit width, this decomposition will always produce hardware-friendly configurations, and through an additional 0-bit option, serves as a unified view of pruning and quantization. Bayesian Bits then introduces learnable stochastic gates, which collectively control the bit width of the given tensor. As a result, we can obtain low bit solutions by performing approximate inference over the gates, with prior distributions that encourage most of them to be switched off. We experimentally validate our proposed method on several benchmark datasets and show that we can learn pruned, mixed precision networks that provide a better trade-off between accuracy and efficiency than their static bit width equivalents.

Motivation & Objective

• Motivate reducing inference cost via joint pruning and mixed-precision quantization.
• Introduce a hardware-friendly quantization decomposition exposing power-of-two bit widths.
• Develop Bayesian gates and a variational objective to learn bit-widths and pruning jointly.
• Provide a practical optimization scheme with a STE-inspired gradient estimator and thresholding for gates.
• Demonstrate improved accuracy/efficiency trade-offs on benchmarks and show end-to-end and post-training variants.

Proposed method

• Decompose quantization into successive residual quantizations doubling bit width (2, 4, 8, 16, 32).
• Quantize residuals and sum them to form higher-precision quantized values.
• Introduce binary gates z on each residual to control whether higher-bit residuals are added, enabling zero-bit pruning.
• Frame gate learning as variational inference with autoregressive priors and posteriors that favor low bit widths.
• Derive a practical objective resembling an L0-like regularizer that penalizes inclusion of higher-bit residuals; use hard-concrete relaxations for gradient-based optimization.
• Describe training with gradient checkpointing to manage memory, PACT-based clipping for inputs, and STE for backprop through rounding.

Experimental results

Research questions

• RQ1 Can a residual-based, hardware-friendly decomposition expose all power-of-two bit widths for mixed-precision quantization?
• RQ2 Can learnable gating over residuals effectively trade off accuracy against computation by jointly pruning and quantizing?
• RQ3 Does a Bayesian/prior-regularized objective yield better accuracy–efficiency trade-offs than static bit-width baselines across tasks?
• RQ4 Is the approach feasible end-to-end and in post-training settings on standard benchmarks?

Key findings

• Bayesian Bits achieves better trade-offs between accuracy and computational efficiency (BOPs) than several baselines on MNIST and CIFAR-10.
• On ImageNet with ResNet18 and MobileNetV2, Bayesian Bits delivers favorable accuracy–BOP trade-offs compared to fixed-bit baselines and other quantization methods.
• Varying the global regularization parameter mu controls the sparsity/bit-widths, yielding highly compressed networks with maintained accuracy in some regimes.
• The method supports end-to-end fine-tuning and post-training mixed-precision quantization with competitive performance.
• Gates exhibit interpretable behavior, often pruning to low bit widths while preserving essential layers (e.g., first/last layers higher precision).
• The approach unifies pruning and quantization under a single probabilistic framework with a practical optimization path.