Loss Aware Post-training Quantization

TL;DR

This paper investigates the impact of quantization on neural network loss landscapes and introduces a joint quantization method that improves accuracy, especially for aggressive INT4 quantization, by leveraging landscape structure insights.

Contribution

It provides a novel analysis of loss landscape changes due to quantization and proposes a joint quantization approach that enhances post-training quantization accuracy.

Findings

Loss landscape remains flat and separable for mild quantization.

Aggressive quantization leads to non-separable, steep loss landscapes.

Joint quantization of layer parameters improves accuracy significantly.

Abstract

Neural network quantization enables the deployment of large models on resource-constrained devices. Current post-training quantization methods fall short in terms of accuracy for INT4 (or lower) but provide reasonable accuracy for INT8 (or above). In this work, we study the effect of quantization on the structure of the loss landscape. Additionally, we show that the structure is flat and separable for mild quantization, enabling straightforward post-training quantization methods to achieve good results. We show that with more aggressive quantization, the loss landscape becomes highly non-separable with steep curvature, making the selection of quantization parameters more challenging. Armed with this understanding, we design a method that quantizes the layer parameters jointly, enabling significant accuracy improvement over current post-training quantization methods. Reference implementation is available at https://github.com/ynahshan/nn-quantization-pytorch/tree/master/lapq