[論文レビュー] Position-based Scaled Gradient for Model Quantization and Pruning

Jangho Kim, KiYoon Yoo+1 more|arXiv (Cornell University)|May 22, 2020

Advanced Neural Network Applications被引用数 14

ひとこと要約

本稿では、重みの位置に基づいて勾配をスケーリングする新しい正則化手法であるPosition-based Scaled Gradient (PSG) を提案する。PSGDは、微調整なしに、後処理による量子化やプルーニングに適した最小解への収束を促進することで、フルプレシジョンモデルの低ビット量子化やプルーニングを容易にする。PSGDは、フルプレシジョンと低プレシジョンの重み分布のギャップを低減し、CIFAR-100およびImageNetにおいて4ビット量子化でも最先端の性能を達成する。

ABSTRACT

We propose the position-based scaled gradient (PSG) that scales the gradient depending on the position of a weight vector to make it more compression-friendly. First, we theoretically show that applying PSG to the standard gradient descent (GD), which is called PSGD, is equivalent to the GD in the warped weight space, a space made by warping the original weight space via an appropriately designed invertible function. Second, we empirically show that PSG acting as a regularizer to a weight vector is favorable for model compression domains such as quantization and pruning. PSG reduces the gap between the weight distributions of a full-precision model and its compressed counterpart. This enables the versatile deployment of a model either as an uncompressed mode or as a compressed mode depending on the availability of resources. The experimental results on CIFAR-10/100 and ImageNet datasets show the effectiveness of the proposed PSG in both domains of pruning and quantization even for extremely low bits. The code is released in Github.

研究の動機と目的

• 後処理による量子化が困難な、フルプレシジョンと量子化済み重み分布の根本的なギャップを解消すること。
• 再訓練やデータアクセスなしに、1つのモデルをフルプレシジョンモードと圧縮モードの両方で柔軟に展開可能にする。
• 収束を特定のターゲットポイントへ正則化することで、低ビット量子化およびプルーニングに適した重みを生み出すトレーニング手法を開発すること。
• 追加のファインチューニングやデータを必要としない、既存の量子化に配慮した訓練（QAT）および後処理量子化（PTQ）手法の理論的裏付けのある代替手法を提供すること。
• PSGDが、量子化やプルーニングで使用される離散的重みグリッドとより相性の良い鋭い最小解に収束することを示すこと。

提案手法

• 重みの現在値とそのターゲット量子化値（例：4ビットグリッドの最近傍点）とのL2距離に基づいて勾配をスケーリングする、位置に基づくスケーリング勾配（PSG）を導入する。
• 標準勾配の代わりにPSGを使用する確率的勾配降下法（PSGD）を定義し、重みを圧縮に適した位置へ正則化する。
• PSGDが、重み空間を歪曲した空間における標準的勾配降下法（SGD）と等価であることを理論的に証明する。歪曲関数は可逆的かつ微分可能であり、重みをターゲットグリッドポイントへ引き寄せるように設計されている。
• 可逆的かつ微分可能な関数を用いて、元の重み空間を歪曲し、歪曲空間における勾配更新が元の空間で標準的GDに対応するようにする。
• プルーニングの場合は、ターゲットポイントをゼロに設定し、位置依存の勾配スケーリングにより重みをゼロへ正則化することでスパーストレーニングを可能にする。
• 初期トレーニング段階でPSGDを適用し、高精度を維持しつつ、直接的に低ビットフォーマットに量子化可能なフルプレシジョンモデルを生成する。

実験結果

リサーチクエスチョン

• RQ1 重みの位置に基づく勾配スケーリングは、フルプレシジョンと低プレシジョンモデル間の重み分布ギャップを低減できるか？
• RQ2 PSGDは、標準的SGDと比較して、量子化やプルーニングに適した最小解に収束するか？
• RQ3 PSGDでトレーニングされたモデルは、追加のファインチューニングなしに、フルプレシジョンおよび低プレシジョンモードの両方で高い精度を達成できるか？
• RQ4 PSGDは、既存のQATおよびPTQ手法と比較して、精度および展開の柔軟性において優れているか？
• RQ5 PSGDは極めて低ビット量子化（例：4ビット）および構造的プルーニングにおいて有効か？

主な発見

• PSGDは、図1(b)に示すように、フルプレシジョンと4ビット量子化重みの間の分布ギャップを低減する。PSGDでトレーニングされたモデルの重み分布は、低プレシジョンモデルの分布と非常に類似している。
• CIFAR-100では、PSGDは4ビット重みおよび活性化量子化（W4A4）で66.5％のトップ1精度を達成し、微調整なしの事前学習モデルとして、標準的SGDおよび同時期のPTQ手法を上回る性能を示した。
• LAPQ PTQ手法と組み合わせた場合、PSGDで事前学習されたモデルはImageNet W4A4で66.5％の精度を達成し、LAPQ単体（60.3％）およびPSGD単体（63.3％）をそれぞれ3.1ポイントおよび6.2ポイント上回った。
• PSGDでトレーニングされたモデルは、標準的SGDと比較して鋭い最小解に収束しており、ヘッセ行列の固有値のうちλ > 10⁻³の正の大きな固有値が9倍多く観察された。これは、より鋭い谷の最小解に収束していることを示唆している。
• 本手法により、リソース制約のある環境でも再訓練やデータアクセスなしに、1つのモデルをフルプレシジョンおよび圧縮モードの両方で直接展開可能となった。
• PSGDは重み量子化およびマグニチュードベースのプルーニングの両方で有効であり、モデル圧縮のための統合的正則化手法としての多様性を示した。