[论文解读] Q-ViT: Fully Differentiable Quantization for Vision Transformer

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

CCD and CMOS Imaging SensorsEngineering被引用 17

一句话总结

Q-ViT 提出了一种用于视觉变换器（Vision Transformers）的完全可微量化方法，联合学习最优的量化比例和位宽，采用逐头位宽分配和可切换比例技术以稳定训练。该方法在 3 位量化下实现了最先进性能，在 DeiT-Tiny 上相比均匀量化方法准确率提升 1.5%，且准确率下降极小。

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.

研究动机与目标

• 解决视觉变换器（ViT）中低比特量化的问题，现有方法在 4 位以下时面临严重准确率下降。
• 研究 ViT 中各组件的量化鲁棒性，以识别敏感层和组件。
• 开发一种完全可微的量化框架，联合优化量化比例和位宽，同时满足模型大小与准确率约束。
• 实现针对 ViT 架构特性的混合精度量化，尤其针对多头自注意力（MSA）和 GELU 层。
• 在 3 位精度下实现稳定训练与高性能，推动 ViT 量化极限。

提出的方法

• 提出一种完全可微的量化框架，其中量化比例和位宽均为可微、可学习的参数。
• 提出逐头位宽分配机制，使 MSA 中的各个注意力头可根据其对量化敏感性的不同，学习独立的位宽。
• 设计一种可切换比例机制，以稳定比例与位宽的联合优化，解决训练过程中的收敛问题。
• 通过经验鲁棒性分析，为敏感组件（尤其是 GELU 层和部分 MSA 头）分配更高位宽，实现混合精度量化。
• 在优化目标中引入正则化项，以平衡准确率与计算约束（BitOPs）。
• 采用温度控制策略（σ），在训练过程中逐步从均匀位宽过渡到学习得到的位宽，提升优化稳定性。

实验结果

研究问题

• RQ1 视觉变换器中哪些组件对量化噪声最敏感，其鲁棒性在不同层和头之间如何变化？
• RQ2 能否通过可微的端到端训练方法，在 ViT 中稳定联合优化量化比例与位宽？
• RQ3 与均匀位宽相比，逐头位宽分配是否能提升低比特 ViT 量化中的性能？
• RQ4 3 位量化 ViT 是否能实现具有竞争力的准确率，且无显著性能下降？
• RQ5 可切换比例机制如何提升混合精度 ViT 量化中的收敛性与性能？

主要发现

• GELU 激活层因具有长尾分布，对量化极为敏感，需更高位宽以保证稳定性能。
• MSA 中的不同注意力头对量化表现出显著不同的鲁棒性，部分头（如 DeiT-Small 中的索引 1 和 2）在 2 位量化时会导致严重准确率下降。
• Q-ViT 在 3 位量化下于 ImageNet 上实现 DeiT-Tiny 的 69.15% 顶级-1 准确率，优于最先进均匀量化方法 LSQ+ 的 1.5%。
• 可切换比例技术对收敛至关重要：若无该机制，Q-ViT 在 3 位与 4 位设置下的表现均劣于均匀量化。
• 逐头位宽分配在所有 ViT 变体中提升性能 0.5%，证明了逐头精度自适应的重要性。
• 在 3 位约束下，GELU 层始终学习更高位宽（通常 ≥4 位），证实其敏感性及自适应分配的有效性。