[論文レビュー] SDQ: Stochastic Differentiable Quantization with Mixed Precision

Xijie Huang, Zhiqiang Shen+6 more|arXiv (Cornell University)|Jun 9, 2022

Image Enhancement Techniques被引用数 8

ひとこと要約

SDQは、微分可能ビット幅パラメータ（DBP）とGumbel-softmax再パrameterizationを用いて、ネットワーク層全体で最適な混合精度ビット幅割り当てを自動的に学習する確率的微分可能量子化手法を提案する。これにより、滑らかでグローバルに最適化された学習が可能となり、ResNetおよびMobileNetにおいて平均ビット幅が4ビット未塔で、完全精度モデルを上回る最先端の精度を達成するとともに、GPUおよびFPGA上で優れたハードウェア効率を示す。

ABSTRACT

In order to deploy deep models in a computationally efficient manner, model quantization approaches have been frequently used. In addition, as new hardware that supports mixed bitwidth arithmetic operations, recent research on mixed precision quantization (MPQ) begins to fully leverage the capacity of representation by searching optimized bitwidths for different layers and modules in a network. However, previous studies mainly search the MPQ strategy in a costly scheme using reinforcement learning, neural architecture search, etc., or simply utilize partial prior knowledge for bitwidth assignment, which might be biased and sub-optimal. In this work, we present a novel Stochastic Differentiable Quantization (SDQ) method that can automatically learn the MPQ strategy in a more flexible and globally-optimized space with smoother gradient approximation. Particularly, Differentiable Bitwidth Parameters (DBPs) are employed as the probability factors in stochastic quantization between adjacent bitwidth choices. After the optimal MPQ strategy is acquired, we further train our network with entropy-aware bin regularization and knowledge distillation. We extensively evaluate our method for several networks on different hardware (GPUs and FPGA) and datasets. SDQ outperforms all state-of-the-art mixed or single precision quantization with a lower bitwidth and is even better than the full-precision counterparts across various ResNet and MobileNet families, demonstrating the effectiveness and superiority of our method.

研究の動機と目的

• 既存の混合精度量子化（MPQ）手法が、コストの高い探索や部分的に最適でないヒューリスティクスに依存する点を是正すること。
• ネットワーク層全体で最適なビット幅をエンドツーエンド学習可能な、微分可能でグローバルに最適化されたフレームワークの開発。
• 勾配逆伝播プロセスにおける線形補間の代わりに確率的量子化を採用することで、学習の安定性と収束性を向上させること。
• エントロピーに配慮した正則化と知識蒸留を用いて、量子化誤差を最小限に抑え、情報保持成分を維持すること。
• 多様なハードウェアプラットフォーム上で、高精度かつ低ビット幅のモデルを実現し、完全精度モデルおよび先行の量子化モデルを上回ること。

提案手法

• 隣接する候補ビット幅間の確率的選択を表す学習可能な確率的要因として、微分可能ビット幅パラメータ（DBP）を導入する。
• 前方伝搬において、各重みまたは活性化がDBPでパrameter化されたカテゴリカル分布からサンプリングされたビット幅に量子化される確率的量子化を採用する。
• 離散的ビット幅選択を微分可能にすることで勾配逆伝播を可能にするため、Gumbel-softmax再パラメータ化を用いる。
• 層固有の精度要件を考慮することで、情報豊富な重み成分を保持し、量子化誤差を低減するエントロピーに配慮したバイン正則化（EBR）を適用する。
• 完全精度の教師モデルからの知識蒸留を統合し、量子化された学生ネットワークの表現能力を強化する。
• 標準的なバックプロパゲーションを用いて、MPQ戦略全体をエンドツーエンド最適化し、すべての層で最適なビット幅割り当てをワンショットで学習可能にする。

実験結果

リサーチクエスチョン

• RQ1 微分可能で確率的なアプローチが、強化学習やヒューリスティクスベースのビット幅探索に比べ、精度と最適化の安定性において優れているか。
• RQ2 Gumbel-softmax再パラメータ化を用いた確率的量子化は、線形補間と比較して、MPQにおける学習安定性と損失関数の滑らかさをどのように向上させるか。
• RQ3 エントロピーに配慮した正則化と知識蒸留は、低ビット幅環境下で量子化誤差をどれほど低減し、モデル精度を維持できるか。
• RQ4 提案されたSDQフレームワークは、ResNetやMobileNetのような多様なアーキテクチャにおいて、平均ビット幅が4ビット未塔で最先端の精度を達成できるか。
• RQ5 FPGAなどの実際のハードウェアアクセラレータにデプロイした場合、SDQの遅延およびエネルギー効率はどの程度か。

主な発見

• ImageNet-1KとResNet18を用いた実験では、SDQは重み/活性化の平均ビット幅が3.85/2ビットで71.7%のトップ-1精度を達成し、完全精度モデルの70.5%を上回った。
• ImageNet-1KとMobileNetV2を用いた実験では、SDQは平均ビット幅4.02/3.85で71.8%のトップ-1精度を達成し、すべての先行単一精度および混合精度量子化手法を上回った。
• COCO検出タスクとYOLOv4-tinyを用いた実験では、SDQは平均ビット幅3.88/4.00で15.9%のmAPを達成し、4ビット均一量子化を上回り、8ビットベースラインに近づいた一方で、エネルギー消費を41.5%削減した。
• FPGA上でのSDQモデルは21.28msの遅延と167.1mJのエネルギー消費を達成し、4ビット均一量子化モデルと同等のハードウェア効率を示した。
• SDQの損失関数の滑らかさは、線形補間ベースの手法よりも顕著に向上しており、より安定的かつ効果的な最適化が可能であることを示した。
• エントロピーに配慮したバイン正則化と知識蒸留の組み合わせにより、特に低精度環境下で情報保持成分の保存が向上し、量子化誤差が低減した。