Softmax Bias Correction for Quantized Generative Models
Nilesh Prasad Pandey Marios Fournarakis Chirag Patel Markus Nagel
Post-training quantization (PTQ) is the go-to compression technique for large
generative models, such as stable diffusion or large language models. PTQ
methods commonly keep the softmax activation in higher precision as it has been
shown to be very sensitive to quantization noise. However, this can lead to a
significant runtime and power overhead during inference on resource-constraint
edge devices. In this work, we investigate the source of the softmax
sensitivity to quantization and show that the quantization operation leads to a
large bias in the softmax output, causing accuracy degradation. To overcome
this issue, we propose an offline bias correction technique that improves the
quantizability of softmax without additional compute during deployment, as it
can be readily absorbed into the quantization parameters. We demonstrate the
effectiveness of our method on stable diffusion v1.5 and 125M-size OPT language
model, achieving significant accuracy improvement for 8-bit quantized softmax.