BayesQ: Uncertainty-Guided Bayesian Quantization

TL;DR

BayesQ introduces an uncertainty-guided Bayesian post-training quantization method that optimizes quantization under the posterior expected loss, leading to improved model accuracy with minimal additional preprocessing.

Contribution

It is the first to optimize quantization using a Bayesian posterior framework, incorporating uncertainty to enhance post-training quantization performance.

Findings

Outperforms strong PTQ baselines on ResNet-50 and BERT-base.

Achieves up to +1.5% top-1 accuracy on ImageNet.

Requires comparable preprocessing to existing methods.

Abstract

We present BayesQ, an uncertainty-guided post-training quantization framework that is the first to optimize quantization under the posterior expected loss. BayesQ fits a lightweight Gaussian posterior over weights (diagonal Laplace by default; optional K-FAC/low-rank), whitens by the posterior covariance, designs codebooks to minimize posterior-expected distortion, and allocates mixed precision via a greedy knapsack that maximizes marginal expected-loss reduction per bit under a global budget. For scalar quantizers, posterior-expected MSE yields closed-form tables; task-aware proxies are handled by short Monte Carlo on a small calibration set. An optional calibration-only distillation aligns the quantized model with the posterior predictive teacher. At matched average bits/weight of 3.0/3.5/4.0, BayesQ improves over strong PTQ baselines on ResNet-50 (ImageNet) and BERT-base (GLUE) e.g., vs. GPTQ by $+1.5/+0.7/+0.3$ top-1 percentage points on RN50 and $+1.1/+0.4/+0.2$ GLUE points on BERT, while requiring one-time preprocessing comparable to a GPTQ pass. BayesQ reframes low-bit quantization as uncertainty-aware risk minimization in a practical, post-training pipeline.