Efficiently Training A Flat Neural Network Before It has been Quantizated

TL;DR

This paper introduces a framework for pre-training neural networks to be more suitable for low-bit quantization, reducing quantization errors and improving post-training quantization efficiency.

Contribution

It proposes a novel approach to pre-condition neural networks by modeling quantization errors as Gaussian noises and optimizing for flat minima, enhancing low-bit PTQ performance.

Findings

Model-agnostic neural networks can be pre-trained for low-bit quantization.

Gaussian noise modeling improves quantization error understanding.

Experimental results show significant reduction in quantization errors.

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.