[Paper Review] Q-ViT: Fully Differentiable Quantization for Vision Transformer

Zhexin Li, Tong Yang+2 more|arXiv (Cornell University)|Jan 19, 2022

CCD and CMOS Imaging Sensors17 citations

TL;DR

Q-ViT proposes a fully differentiable quantization method for Vision Transformers that jointly learns optimal quantization scales and bit-widths, with head-wise bit-width allocation and a switchable scale technique to stabilize training. It achieves state-of-the-art performance at 3-bit quantization, outperforming uniform quantization by 1.5% on DeiT-Tiny with minimal accuracy drop.

ABSTRACT

In this paper, we propose a fully differentiable quantization method for vision transformer (ViT) named as Q-ViT, in which both of the quantization scales and bit-widths are learnable parameters. Specifically, based on our observation that heads in ViT display different quantization robustness, we leverage head-wise bit-width to squeeze the size of Q-ViT while preserving performance. In addition, we propose a novel technique named switchable scale to resolve the convergence problem in the joint training of quantization scales and bit-widths. In this way, Q-ViT pushes the limits of ViT quantization to 3-bit without heavy performance drop. Moreover, we analyze the quantization robustness of every architecture component of ViT and show that the Multi-head Self-Attention (MSA) and the Gaussian Error Linear Units (GELU) are the key aspects for ViT quantization. This study provides some insights for further research about ViT quantization. Extensive experiments on different ViT models, such as DeiT and Swin Transformer show the effectiveness of our quantization method. In particular, our method outperforms the state-of-the-art uniform quantization method by 1.5% on DeiT-Tiny.

Motivation & Objective

• Address the challenge of low-bit quantization in Vision Transformers (ViT), where existing methods suffer from severe accuracy degradation below 4-bit.
• Investigate component-level quantization robustness in ViT to identify sensitive layers and components.
• Develop a fully differentiable quantization framework that jointly optimizes quantization scales and bit-widths under model size and accuracy constraints.
• Enable mixed-precision quantization tailored to ViT's architectural characteristics, particularly in Multi-Head Self-Attention (MSA) and GELU layers.
• Achieve stable training and high performance at 3-bit precision, pushing the limits of ViT quantization.

Proposed method

• Introduce a fully differentiable quantization framework where both quantization scales and bit-widths are differentiable, learnable parameters.
• Propose head-wise bit-width allocation, allowing individual attention heads in MSA to learn distinct bit-widths based on their robustness to quantization.
• Design a switchable scale mechanism to stabilize joint optimization of scales and bit-widths, resolving convergence issues during training.
• Apply mixed-precision quantization by assigning higher bit-widths to sensitive components—especially GELU and certain MSA heads—based on empirical robustness analysis.
• Formulate the optimization objective with a regularization term to balance accuracy and computational constraints (BitOPs).
• Use a temperature-controlled strategy (σ) to gradually transition from uniform to learned bit-widths during training, improving optimization stability.

Experimental results

Research questions

• RQ1 Which components in Vision Transformers are most sensitive to quantization noise, and how does their robustness vary across layers and heads?
• RQ2 Can joint optimization of quantization scales and bit-widths be stabilized in ViT using a differentiable, end-to-end training approach?
• RQ3 Does head-wise bit-width allocation improve performance in low-bit ViT quantization compared to uniform bit-widths?
• RQ4 Can a 3-bit quantized ViT achieve competitive accuracy without significant performance degradation?
• RQ5 How does the switchable scale mechanism improve convergence and performance in mixed-precision ViT quantization?

Key findings

• GELU activation layers are highly sensitive to quantization due to their long-tailed distribution, requiring higher bit-widths for stable performance.
• Different attention heads in MSA exhibit significantly varying robustness to quantization, with some heads (e.g., index 1 and 2 in DeiT-Small) causing severe accuracy drops when quantized to 2-bit.
• Q-ViT achieves 69.15% top-1 accuracy on ImageNet for DeiT-Tiny at 3-bit quantization, outperforming the state-of-the-art uniform quantization method LSQ+ by 1.5%.
• The switchable scale technique is critical for convergence: without it, Q-ViT performs worse than uniform quantization at both 3-bit and 4-bit settings.
• Head-wise bit-width allocation improves performance by 0.5% across ViT variants, demonstrating the importance of per-head precision adaptation.
• GELU layers consistently learn higher bit-widths (often ≥4 bits) even under a 3-bit constraint, confirming their sensitivity and the effectiveness of adaptive allocation.