[论文解读] Oscillation-free Quantization for Low-bit Vision Transformers

Shih-Yang Liu, Zechun Liu+1 more|arXiv (Cornell University)|Feb 4, 2023

Advanced Neural Network Applications被引用 6

一句话总结

该论文提出无振荡量化（OFQ）用于低比特视觉Transformer模型，解决了由可学习缩放因子和查询-键权重共振荡引起的量化感知训练中的权重振荡问题。通过引入统计权重量化（StatsQ）、置信度引导退火（CGA）和查询-键重参数化（QKR），OFQ稳定了训练过程，并实现了最先进（SOTA）的准确率——在2-bit DeiT-T/S和Swin-T上分别较之前方法提升9.88%、7.72%和4.64%。

ABSTRACT

Weight oscillation is an undesirable side effect of quantization-aware training, in which quantized weights frequently jump between two quantized levels, resulting in training instability and a sub-optimal final model. We discover that the learnable scaling factor, a widely-used $ extit{de facto}$ setting in quantization aggravates weight oscillation. In this study, we investigate the connection between the learnable scaling factor and quantized weight oscillation and use ViT as a case driver to illustrate the findings and remedies. In addition, we also found that the interdependence between quantized weights in $ extit{query}$ and $ extit{key}$ of a self-attention layer makes ViT vulnerable to oscillation. We, therefore, propose three techniques accordingly: statistical weight quantization ($ m StatsQ$) to improve quantization robustness compared to the prevalent learnable-scale-based method; confidence-guided annealing ($ m CGA$) that freezes the weights with $ extit{high confidence}$ and calms the oscillating weights; and $ extit{query}$-$ extit{key}$ reparameterization ($ m QKR$) to resolve the query-key intertwined oscillation and mitigate the resulting gradient misestimation. Extensive experiments demonstrate that these proposed techniques successfully abate weight oscillation and consistently achieve substantial accuracy improvement on ImageNet. Specifically, our 2-bit DeiT-T/DeiT-S algorithms outperform the previous state-of-the-art by 9.8% and 7.7%, respectively. Code and models are available at: https://github.com/nbasyl/OFQ.

研究动机与目标

• 探究视觉Transformer（ViT）在低比特量化感知训练（QAT）中训练不稳定的根源原因。
• 识别出可学习缩放因子和查询-键权重共振荡是导致权重振荡的主要驱动因素。
• 开发一种稳健的无振荡量化框架，以提升低比特ViT模型的准确率与收敛性。
• 在极低比特宽度下实现ImageNet上的最先进性能，包括3-bit模型达到与全精度性能相当的水平。
• 确保在不同ViT架构（DeiT、Swin）和比特宽度（2–4比特）下的泛化性与稳定性。

提出的方法

• 提出统计权重量化（StatsQ），用基于权重分布统计的缩放因子替代可学习缩放因子，以减少振荡。
• 引入置信度引导退火（CGA），通过边界范围阈值冻结高置信度权重，并选择性地更新振荡权重。
• 设计查询-键重参数化（QKR），以解耦量化查询与键权重之间的相互影响，缓解梯度误估问题。
• 将StatsQ、CGA和QKR整合为统一框架，稳定训练过程并实现无振荡模型的收敛。
• 在CGA中采用动态边界范围，平衡优化速度与有效性，基于平均梯度幅值得到经验性选择。
• 通过缩放因子轨迹分析，验证CGA在训练过程中抑制振荡的有效性。

实验结果

研究问题

• RQ1 为何在量化感知训练中使用可学习缩放因子会导致视觉Transformer中的权重振荡？
• RQ2 量化查询与键权重在自注意力层中的相互依赖关系如何在训练过程中加剧振荡？
• RQ3 基于权重统计的统计缩放是否能优于可学习缩放，从而更有效地减少振荡并提升模型准确率？
• RQ4 置信度引导退火在多大程度上能稳定低比特ViT量化训练并消除振荡？
• RQ5 查询-键重参数化是否能有效解耦量化查询与键权重之间的负面相互影响？

主要发现

• 可学习缩放因子由于梯度更新噪声和阈值不稳定性，显著加剧了权重振荡。
• 统计权重量化（StatsQ）通过基于权重分布的统计缩放，在2-bit DeiT-S上相比基线LSQ将振荡减少并提升准确率6.3%。
• 采用边界范围为0.005的置信度引导退火（CGA）取得最佳性能，相比其他设置提升准确率0.72%。
• StatsQ、CGA与QKR的联合应用在2-bit DeiT-S上相比之前SOTA提升7.72%准确率，2-bit DeiT-T提升9.88%。
• 使用OFQ的3-bit DeiT-T和DeiT-S模型准确率与全精度模型相当，证明了在振荡控制下低比特量化的有效性。
• 在4-bit设置下，OFQ在DeiT-T、DeiT-S和Swin-T上仍分别优于全精度模型3.24%、1.2%和0.68%的准确率。