[Paper Review] Robust Quantization: One Model to Rule Them All

Moran Shkolnik, Brian Chmiel+5 more|arXiv (Cornell University)|Feb 18, 2020

Advanced Neural Network ApplicationsComputer Science27 references34 citations

TL;DR

Proposes Kurtosis Regularization (KURE) to sculpt weights toward a uniform-like distribution, yielding a single model robust to varying quantization bit-widths and policies and compatible with PTQ and QAT.

ABSTRACT

Neural network quantization methods often involve simulating the quantization process during training, making the trained model highly dependent on the target bit-width and precise way quantization is performed. Robust quantization offers an alternative approach with improved tolerance to different classes of data-types and quantization policies. It opens up new exciting applications where the quantization process is not static and can vary to meet different circumstances and implementations. To address this issue, we propose a method that provides intrinsic robustness to the model against a broad range of quantization processes. Our method is motivated by theoretical arguments and enables us to store a single generic model capable of operating at various bit-widths and quantization policies. We validate our method's effectiveness on different ImageNet models.

Motivation & Objective

• Motivate the need for quantization-robust DNNs that work under diverse quantization settings.
• Develop a theoretical framework comparing uniform vs normal weight distributions for quantization robustness.
• Introduce KURE to steer tensor distributions toward uniformity without harming full-precision performance.
• Demonstrate empirical robustness gains across PTQ and QAT on ImageNet CNNs.

Proposed method

• Formulate quantization sensitivity and MSE under uniform and normal distributions to justify robustness advantages of uniformity.
• Derive and prove analytic results showing uniform tensors have lower quantization sensitivity than normal tensors.
• Introduce KURE by adding a kurtosis-based regularization term to the training loss to drive weights toward a uniform-like distribution.
• Show that weight uniformization does not hurt convergence or full-precision accuracy prior to quantization.
• Empirically validate KURE on ImageNet models with PTQ and QAT, across multiple architectures.
• Compare KURE to gradient L1 regularization approaches in prior work, highlighting computational benefits and broader robustness.

Experimental results

Research questions

• RQ1 How does weight/activation distribution affect robustness to quantization step-size and bit-width variations?
• RQ2 Do uniformly distributed tensors exhibit lower quantization sensitivity than normally distributed tensors?
• RQ3 Can a regularization term targeting kurtosis steer model weights toward a uniform-like distribution without harming full-precision performance?
• RQ4 Is a single KURE-regularized model robust across different quantization policies and bit-widths under PTQ and QAT?
• RQ5 How does KURE compare to prior robustness methods in both performance and computational cost?

Key findings

• Uniform weight distributions yield lower quantization sensitivity than normal distributions across step-size variations.
• Kurtosis Regularization (KURE) can enforce a uniform-like weight distribution without hurting convergence or full-precision accuracy before quantization.
• KURE improves robustness to changes in quantization step size and to varying bit-widths in PTQ and QAT settings.
• KURE achieves stronger robustness than prior L1-gradient regularization methods while having lower computational overhead.
• Empirical results on ImageNet with ResNet-18/50 and MobileNet-V2 show robustness gains across multiple quantization scenarios.