[论文解读] Q-ViT: Accurate and Fully Quantized Low-bit Vision Transformer

Yanjing Li, Sheng Xu+4 more|arXiv (Cornell University)|Oct 13, 2022

CCD and CMOS Imaging Sensors被引用 31

一句话总结

Q-ViT 引入信息整正模块（IRM）和分布引导蒸馏（DGD），以实现全量化低比特 Vision Transformers，在 ImageNet 上接近或超越全精度精度，同时实现显著的加速。

ABSTRACT

The large pre-trained vision transformers (ViTs) have demonstrated remarkable performance on various visual tasks, but suffer from expensive computational and memory cost problems when deployed on resource-constrained devices. Among the powerful compression approaches, quantization extremely reduces the computation and memory consumption by low-bit parameters and bit-wise operations. However, low-bit ViTs remain largely unexplored and usually suffer from a significant performance drop compared with the real-valued counterparts. In this work, through extensive empirical analysis, we first identify the bottleneck for severe performance drop comes from the information distortion of the low-bit quantized self-attention map. We then develop an information rectification module (IRM) and a distribution guided distillation (DGD) scheme for fully quantized vision transformers (Q-ViT) to effectively eliminate such distortion, leading to a fully quantized ViTs. We evaluate our methods on popular DeiT and Swin backbones. Extensive experimental results show that our method achieves a much better performance than the prior arts. For example, our Q-ViT can theoretically accelerates the ViT-S by 6.14x and achieves about 80.9% Top-1 accuracy, even surpassing the full-precision counterpart by 1.0% on ImageNet dataset. Our codes and models are attached on https://github.com/YanjingLi0202/Q-ViT

研究动机与目标

• 识别为何与全精度对照相比，完全量化的 ViTs 会出现较大性能下降。
• 开发在量化自注意力中纠正信息失真的机制（IRM）。
• 设计一个蒸馏方案，使量化模型与全精度模型之间的分布对齐（DGD）。
• 证明具备 IRM 和 DGD 的完全量化 ViTs 能超越基线量化并接近或超过 ImageNet 上的 FP 性能。
• 在 DeiT 和 Swin 骨干上对 2–4 位量化范围进行验证。

提出的方法

• 诊断完全量化 ViTs 的瓶颈，并量化自注意力映射中的信息失真程度。
• 引入信息整正模块（IRM），最大化量化查询和键的熵，恢复注意力中的信息内容。
• 提出分布引导蒸馏（DGD），通过基于注意力的相似性矩阵和补丁级语义蒸馏来对齐分布。
• 在一个对称权重、非对称激活的量化感知训练（QAT）框架下对 ViT 组件进行量化，使用直通估计器进行反向传播。
• 利用教师–学生蒸馏，在教师端通过来自量化和全精度激活的相似性矩阵提供分布式引导。
• 在 ImageNet 上以 DeiT 和 Swin 骨干进行评估，并将其与 LSQ 和基线完全量化 ViT 配置进行对比。

实验结果

研究问题

• RQ1 IRM 是否能够在低位比特量化的 ViT 的注意力图中恢复信息内容，使其分布匹配全精度？
• RQ2 在量化注意力表示中最大化熵是否能改善 ViT 的前向传播性能？
• RQ3 分布引导蒸馏（DGD) 是否能有效减少量化与全精度 ViTs 之间的反向传播分布不匹配？
• RQ4 具备 IRM 和 DGD 的完全量化 ViTs 在 ImageNet 上能有多接近或超越全精度 ViTs？
• RQ5 IRM 与 DGD 的收益在 DeiT 与 Swin 骨干及 2–4 位量化范围内是否一致？

主要发现

• 完全量化的 ViT 基线在性能上存在显著下降，尤其是在 MHSA 关键组件被量化时。
• IRM 提升注意力表示的熵，并相对于完全量化基线在 Top-1 准确率上有所提升（例如在 2–4 位量化范围内观察到增益）。
• DGD 通过量化与全精度激活之间的补丁级相似性矩阵对齐分布，进一步带来提升。
• 结合 IRM 与 DGD（Q-ViT）显著缩小与全精度 ViTs 的差距，在某些情境下甚至超过 FP 性能（例如 4 位 Q-ViT 在 Swin-T 上超过全精度）。
• Q-ViT 在 DeiT 与 Swin 骨干上实现显著的 FLOPs 降低（如在某些配置中约达 ~21x），同时在 ImageNet 上的准确率具有竞争力或更优。