[Paper Review] Quantization Variation: A New Perspective on Training Transformers with Low-Bit Precision

Xijie Huang, Zhiqiang Shen+2 more|arXiv (Cornell University)|Jul 1, 2023

CCD and CMOS Imaging Sensors4 citations

TL;DR

This paper introduces a variation-aware quantization framework for Vision Transformers (ViTs) that addresses training instability caused by parameter fluctuations (variation) during low-bit quantization-aware training (QAT). By employing multi-crop knowledge distillation, module-dependent scaling, and an oscillation-aware bin regularization, the method stabilizes training and achieves 77.66% Top-1 accuracy on ImageNet-1K with 2-bit Swin-T, outperforming prior SOTA by 3.35%.

ABSTRACT

Despite the outstanding performance of transformers in both language and vision tasks, the expanding computation and model size have increased the demand for efficient deployment. To address the heavy computation and parameter drawbacks, quantization is frequently studied in the community as a representative model compression technique and has seen extensive use on ConvNets. However, due to the unique properties of transformers, the low-bit quantization applications are still limited and underexplored. In this paper, we identify the difficulty of transformer low-bit quantization-aware training on its unique variation behaviors, which significantly differ from ConvNets. Based on comprehensive quantitative analysis, we observe variation in three hierarchies: various module quantization sensitivities, outliers in static weight and activation distribution, and oscillation in dynamic parameter fluctuations. These variations of transformers bring instability to the quantization-aware training (QAT) and negatively influence the performance. We explore the best practices to alleviate the variation's influence during low-bit transformer QAT and propose a variation-aware quantization scheme for both vision and language transformers. We extensively verify and show our scheme can alleviate the variation and improve the performance of transformers across various models and tasks. Our solution substantially improves the 2-bit Swin-T and binary BERT-base, achieving a 3.35% and 1.4% accuracy improvement over previous state-of-the-art methods on ImageNet-1K and GLUE. Codes and models are available at https://github.com/HuangOwen/Quantization-Variation.

Motivation & Objective

• To identify and address the root cause of training instability in low-bit quantized Vision Transformers.
• To analyze the unique variation behaviors in ViTs that differ from CNNs and hinder quantization performance.
• To improve the efficiency and stability of quantization-aware training (QAT) for ViTs at extremely low bitwidths (2–4 bits).
• To develop a module-aware, variation-resilient quantization framework that enhances convergence and accuracy.
• To demonstrate state-of-the-art performance on ImageNet-1K with 2-bit ViT models.

Proposed method

• Introduces a multi-crop knowledge distillation scheme to reduce data variance within mini-batches, stabilizing and accelerating QAT.
• Proposes a module-dependent scaling strategy that applies different scale factors per module to accommodate diverse weight distribution variances.
• Designs an oscillation-aware bin regularization (OBR) that reduces weight oscillation by dynamically adjusting the regularization coefficient using a cosine schedule.
• Employs a variation-aware regularization term to suppress unstable weight updates during QAT, improving convergence.
• Uses a differentiable quantization scheme with LSQ+ as the base, enhanced by gradient scaling and distribution-aware normalization.
• Visualizes attention maps to validate that the quantized model preserves representational capacity, especially at 2-bit precision.

Experimental results

Research questions

• RQ1 Why do Vision Transformers exhibit worse performance under low-bit quantization compared to CNNs?
• RQ2 What causes training instability during quantization-aware training of ViTs?
• RQ3 How do parameter variations in ViTs affect weight oscillation and convergence during QAT?
• RQ4 Can knowledge distillation and module-specific scaling improve the stability and accuracy of low-bit ViT quantization?
• RQ5 To what extent can oscillation-aware regularization suppress weight oscillation and improve model accuracy in low-bit ViT training?

Key findings

• The proposed method achieves 77.66% Top-1 accuracy on ImageNet-1K with 2-bit Swin-T, surpassing the previous SOTA by 3.35%.
• The oscillation-aware bin regularization (OBR) reduces the percentage of oscillating weights from 7.33% (baseline) to 0.23% at training end, correlating with improved accuracy.
• Multi-crop knowledge distillation stabilizes training and reduces data variance, contributing to faster convergence and higher accuracy.
• The loss landscape of the proposed method is smoother and more centralized than the baseline, indicating improved training stability.
• Attention map visualization confirms that the 2-bit quantized Swin-T maintains strong attention to salient objects, unlike the baseline which shows uniform attention.
• The method achieves state-of-the-art performance across multiple ViT architectures (DeiT-T, SReT-T, Swin-T) at 2–4 bit precision, demonstrating generalization.