Oscillations Make Neural Networks Robust to Quantization

TL;DR

This paper reveals that weight oscillations during Quantization Aware Training are crucial for model robustness, and introduces a regularizer to induce such oscillations, improving post-training quantization performance.

Contribution

It provides a theoretical analysis of oscillations in QAT, introduces a regularizer to induce oscillations, and demonstrates improved quantization results on neural networks.

Findings

Oscillations are essential for effective QAT.

Training with induced oscillations matches QAT performance.

Regularizer improves post-training quantization results.

Abstract

We challenge the prevailing view that weight oscillations observed during Quantization Aware Training (QAT) are merely undesirable side-effects and argue instead that they are an essential part of QAT. We show in a univariate linear model that QAT results in an additional loss term that causes oscillations by pushing weights away from their nearest quantization level. Based on the mechanism from the analysis, we then derive a regularizer that induces oscillations in the weights of neural networks during training. Our empirical results on ResNet-18 and Tiny Vision Transformer, evaluated on CIFAR-10 and Tiny ImageNet datasets, demonstrate across a range of quantization levels that training with oscillations followed by post-training quantization (PTQ) is sufficient to recover the performance of QAT in most cases. With this work we provide further insight into the dynamics of QAT and contribute a novel insight into explaining the role of oscillations in QAT which until now have been considered to have a primarily negative effect on quantization.