[論文レビュー] FlexRound: Learnable Rounding based on Element-wise Division for Post-Training Quantization

Jung Hyun Lee, Jeonghoon Kim+2 more|arXiv (Cornell University)|Jun 1, 2023

Advanced Neural Network Applications被引用数 5

ひとこと要約

FlexRoundは、パラメータの大きさに応じた柔軟な量子化を可能にする、学習可能な丸め機構を提案する。Post-training quantization (PTQ) において、要素ごとの除算を用いて共通の量子化グリッドサイズと個々の重みスケーリング係数を同時に最適化することで、柔軟でマグニチュードに配慮した量子化を実現する。画像分類、NLPの理解および生成タスクにおいて最先端の性能を達成し、大規模言語モデルの4ビット量子化でも精度の低下を最小限に抑える。

ABSTRACT

Post-training quantization (PTQ) has been gaining popularity for the deployment of deep neural networks on resource-limited devices since unlike quantization-aware training, neither a full training dataset nor end-to-end training is required at all. As PTQ schemes based on reconstructing each layer or block output turn out to be effective to enhance quantized model performance, recent works have developed algorithms to devise and learn a new weight-rounding scheme so as to better reconstruct each layer or block output. In this work, we propose a simple yet effective new weight-rounding mechanism for PTQ, coined \emph{FlexRound}, based on element-wise division instead of typical element-wise addition such that FlexRound enables jointly learning a common quantization grid size as well as a different scale for each pre-trained weight. Thanks to the reciprocal rule of derivatives induced by element-wise division, FlexRound is inherently able to exploit pre-trained weights when updating their corresponding scales, and thus, flexibly quantize pre-trained weights depending on their magnitudes. We empirically validate the efficacy of FlexRound on a wide range of models and tasks. To the best of our knowledge, our work is the first to carry out comprehensive experiments on not only image classification and natural language understanding but also natural language generation. Moreover, we demonstrate, for the first time, that large language models can be efficiently quantized, with only a negligible impact on performance compared to half-precision baselines, achieved by reconstructing the output in a block-by-block manner. Our code is available at \url{https://github.com/onliwad101/FlexRound_LRQ}.

研究の動機と目的

• 固定で均一な丸め方式によるPost-training quantization (PTQ) における性能劣化を是正すること。
• 重みのマグニチュードに応じてスケーリングを適応させることで、事前学習済みモデルのより効果的な量子化を可能にすること。
• 完全な再トレーニングを必要とせず、グリッドサイズと個々の重みスケーリング係数を同時に最適化する学習可能な丸め機構の開発。
• 自然言語生成や大規模言語モデルを含む多様なタスクにおける本手法の有効性を実証すること。
• ブロック単位の出力再構成を用いることで、4ビット量子化における大規模言語モデル（LLaMA-33Bなど）の性能劣化を最小限に抑えること。

提案手法

• FlexRoundは、重みの丸めにおける従来の要素ごとの加算を、勾配伝播が事前学習済み重みのマグニチュードを尊重するように、要素ごとの除算に置き換える。
• 微分可能丸め機構を用いて、共通の量子化グリッドサイズと各事前学習済み重みの個別スケール係数を同時に学習する。
• 除算の導関数の逆数則を活用することで、事前学習済み重みがスケール更新をガイドするようになり、マグニチュードに配慮した量子化が可能になる。
• 層／ブロック出力の再構成損失を用いた勾配降下法によりスケールを最適化し、量子化に起因する性能劣化の近似を最小化する。
• チャネル単位およびテンソル単位の量子化をサポートし、モデル性能の維持を図るためにブロック単位で出力再構成を実施する。
• 既存のPTQフレームワークと互換性があり、最小限の変更で既存の量子化パイプラインに統合可能である。

実験結果

リサーチクエスチョン

• RQ1 要素ごとの除算に基づく学習可能な丸め機構は、従来のPTQ手法に比べ、精度と柔軟性の面で優れているか？
• RQ2 除算によるマグニチュードに配慮したスケーリングは、特に大きな重みに対して量子化性能を向上させるか？
• RQ3 FlexRoundは、性能劣化を最小限に抑えつつ、4ビット量子化による大規模言語モデルの有効な実装を可能にするか？
• RQ4 画像分類、NLPの理解および生成を含む多様なタスクにおいて、FlexRoundはAdaRoundや他の最先端PTQ手法と比較して優れているか？
• RQ5 ブロック単位の出力再構成は、4ビット量子化下での高パrameter数モデル（例：LLaMA-33B）において性能維持に有効か？

主な発見

• FlexRoundは、画像分類およびNLP理解ベンチマークでSOTAの結果を達成し、AdaRoundや他のPTQ手法を上回る。
• WikiText2では、8ビット重みでLLaMA-7Bに対してFlexRoundが5.73の perplexity を達成したのに対し、AdaRoundは6.19であり、相対的に7.3%の向上を示した。
• 4ビット量子化では、FlexRoundがWikiText2でLLaMA-7Bに対して perplexity を5.63に低下させ、AdaRoundの5.72を上回り、半精度ベースラインに近い性能を維持した。
• PTBでは、8ビット重みでLLaMA-7Bに対してFlexRoundが9.28の perplexity を達成したのに対し、AdaRoundは9.85であり、相対的に6.1%の向上を示した。
• 3ビット重みでも、FlexRoundはWikiText2でLLaMA-7Bに対して perplexity を5.88に維持したが、AdaRoundは6.41にとどまり、低ビット量子化における優れた性能を示した。
• LLaMA-33Bでは、4ビット重みでFlexRoundがWikiText2で perplexity を4.32に低下させ、半精度ベースラインの4.06からわずか0.26の増加に抑え、ほぼ半精度性能に近い結果を達成した。