Efficiently Training A Flat Neural Network Before It has been Quantizated

TL;DR

The paper tackles the limitation that post-training quantization (PTQ) performance is bottlenecked by the loss landscape of the pre-trained full-precision model. It introduces Differential Noise-driven Quantization (DNQ), a pre-conditioning framework that models Weight Quantization Error (WQE) and Activation Quantization Error (AQE) as independent Gaussian noises and injects differential weight noise and stochastic activation noise during FP training to steer toward flat minima. The authors establish a theoretical connection between quantization-induced degradation and the Hessian norm $\|\mathbf{H}\|$, and demonstrate empirically that pre-conditioned models achieve state-of-the-art PTQ performance with simple quantization, across CIFAR-100 and standard CV architectures. This proactive landscape-shaping perspective offers a general, model-agnostic pathway to reliable low-bit deployment and suggests extensions to broader architectures and non-Gaussian noise models.

Abstract

Post-training quantization (PTQ) for vision transformers (ViTs) has garnered significant attention due to its efficiency in compressing models. However, existing methods typically overlook the relationship between a well-trained NN and the quantized model, leading to considerable quantization error for PTQ. However, it is unclear how to efficiently train a model-agnostic neural network which is tailored for a predefined precision low-bit model. In this paper, we firstly discover that a flat full precision neural network is crucial for low-bit quantization. To achieve this, we propose a framework that proactively pre-conditions the model by measuring and disentangling the error sources. Specifically, both the Activation Quantization Error (AQE) and the Weight Quantization Error (WQE) are statistically modeled as independent Gaussian noises. We study several noise injection optimization methods to obtain a flat minimum. Experimental results attest to the effectiveness of our approach. These results open novel pathways for obtaining low-bit PTQ models.