[論文レビュー] 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

ひとこと要約

本稿では、学習可能スケーリング因子とクエリ-キーの共鳴的振動によって引き起こされる量子化意識学習（QAT）における重みの振動を解消する、低ビット用のOscillation-free Quantization（OFQ）を提案する。統計的重み量子化（StatsQ）、信頼度誘導型冷却（CGA）、クエリ-キー再パラメータ化（QKR）を導入することで、学習の安定化を実現し、最先端の精度を達成した。2ビットの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.

研究の動機と目的

• 低ビットのビジョントランスフォーマー（ViT）における量子化意識学習（QAT）の学習不安定要因の根本的要因を解明すること。
• 学習可能スケーリング因子とクエリ-キー重みの共鳴的振動が、重みの振動を引き起こす主な要因であることを特定すること。
• 低ビットViTモデルにおける精度と収束性を向上させる、安定で振動のない量子化フレームワークを開発すること。
• 2ビットモデルを含め、最小限のビット幅でImageNet上で最先端の性能を達成すること。3ビットモデルでは完全精度性能に匹敵することを確認すること。
• 異なるViTアーキテクチャ（DeiT、Swin）およびビット幅（2〜4ビット）に一般化可能で安定した性能を発揮すること。

提案手法

• 学習可能スケーリング因子の代わりに、重み分布からの統計的スケーリングを用いる統計的重み量子化（StatsQ）を提案し、振動を低減する。
• 信頼度誘導型冷却（CGA）を導入し、高い信頼度を持つ重みは凍結し、振動する重みは境界範囲閾値を用いて選択的に更新することで、学習を安定化する。
• クエリ-キー再パラメータ化（QKR）を設計し、量子化されたクエリとキー重みの間の相互干渉を分離することで、勾配の誤推定を軽減する。
• StatsQ、CGA、QKRを統合したフレームワークを構築し、学習の安定化と振動のないモデルへの収束を可能にする。
• 最適化の速度と効果のバランスを図るため、CGAで動的境界範囲を採用し、平均勾配の大きさに基づいて実験的に選定する。
• スケーリング因子の軌道解析を用いて、CGAが学習全体で振動を抑制する有効性を検証する。

実験結果

リサーチクエスチョン

• RQ1 量子化意識学習における学習可能スケーリング因子の使用が、なぜビジョントランスフォーマーにおける重みの振動を引き起こすのか？
• RQ2 自己注意層における量子化クエリとキー重みの相互依存性が、学習中にどのように振動を悪化させるのか？
• RQ3 重みの統計に基づくスケーリングが、学習可能スケーリングを上回り、振動の低減とモデル精度の向上に寄与するのか？
• RQ4 信頼度誘導型冷却（CGA）が、低ビットViT量子化の学習をどれほど安定化させ、振動を排除できるのか？
• RQ5 クエリ-キー再パラメータ化が、量子化クエリとキー重みの間の悪影響を及ぼす相互干渉を効果的に分離できるのか？

主な発見

• 学習可能スケーリング因子は、ノイズの多い勾配更新と閾値の不安定性により、重みの振動を顕著に悪化させる。
• 統計的重み量子化（StatsQ）は、2ビットのDeiT-Sにおいて、ベースラインのLSQに比べて6.3%の精度向上を達成し、振動を低減した。
• 境界範囲0.005を用いた信頼度誘導型冷却（CGA）が最良の性能を示し、他の設定に比べて0.72%の精度向上を達成した。
• StatsQ、CGA、QKRの組み合わせにより、2ビットのDeiT-Sでは従来の最先端手法に比べ7.72%の精度向上を達成した。2ビットのDeiT-Tでは9.88%の向上を達成した。
• 3ビットのDeiT-TおよびDeiT-SはOFQを用いることで、完全精度モデルと同等の精度を達成し、振動制御を施した低ビット量子化の有効性を示した。
• 4ビット設定では、OFQは完全精度モデルをそれぞれ3.24%、1.2%、0.68%上回り、DeiT-T、DeiT-S、Swin-Tで優れた性能を発揮した。