[論文レビュー] Memory Efficient Optimizers with 4-bit States

Bingrui Li, Jianfei Chen+1 more|arXiv (Cornell University)|Sep 4, 2023

Machine Learning and Data ClassificationComputer Science被引用数 3

ひとこと要約

本論文では、Adamに類する最適化手法の1回目および2回目のモーメントのための新しい量子化技術を導入することで、メモリ効率の高いディープラーニング学習のための4ビット最適化手法を提案する。より小さなブロックサイズ、2回目のモーメントのためのゼロポイントフリー線形量子化、およびランク1正規化を用いることで、多様なNLP、ビジョン、インstructチューニングタスクにおいて、8ビットベースラインと比較して50%のメモリ削減を達成しながら、フル精度の収束精度を維持する。

ABSTRACT

Optimizer states are a major source of memory consumption for training neural networks, limiting the maximum trainable model within given memory budget. Compressing the optimizer states from 32-bit floating points to lower bitwidth is promising to reduce the training memory footprint, while the current lowest achievable bitwidth is 8-bit. In this work, we push optimizer states bitwidth down to 4-bit through a detailed empirical analysis of first and second moments. Specifically, we find that moments have complicated outlier patterns, that current block-wise quantization cannot accurately approximate. We use a smaller block size and propose to utilize both row-wise and column-wise information for better quantization. We further identify a zero point problem of quantizing the second moment, and solve this problem with a linear quantizer that excludes the zero point. Our 4-bit optimizers are evaluated on a wide variety of benchmarks including natural language understanding, machine translation, image classification, and instruction tuning. On all the tasks our optimizers can achieve comparable accuracy with their full-precision counterparts, while enjoying better memory efficiency.

研究の動機と目的

• 高精度の最適化状態（例：32ビットAdam状態）に起因する大規模ニューラルネットワーク学習におけるメモリボトルネックを解消すること。
• 現在の8ビットの最小値からさらに4ビットに最適化状態のビット幅を低減し、メモリ使用量をさらに圧縮すること。
• 既存のブロックワイド量子化の限界を克服し、1回目および2回目のモーメントにおける複雑なアウーターパターンを捉えられない問題に対処すること。
• 2回目のモーメント量子化におけるゼロポイント問題を解消し、パラメータ更新の方向が著しくずれるのを防ぐこと。
• 安定な埋め込み層に依存せずに、収束速度と精度を維持する強力な4ビット最適化手法を開発すること。

提案手法

• 1回目のモーメントにおける局所的なアウーターパターンをよりよく捉えるために、量子化用のより小さなブロックサイズを提案する。
• 非ゼロ値がゼロに量子化されるのを防ぎ、更新方向の歪みを回避するため、2回目のモーメントのためのゼロポイントフリー線形量子化を導入する。
• アウーター構造をより効果的にモデル化することで、2回目のモーメントの近似を向上させるランク1正規化を設計する。
• 4ビット量子化された1回目のモーメントと要因分解された2回目のモーメントを組み合わせ、メモリ効率を向上させたハイブリッド最適化手法を構築する。
• 各イテレーションで最適化状態を圧縮・展開する圧縮ベースの最適化フレームワークに、提案された量子化器を統合する。
• 実験的分析を用いて量子化器設計を検証し、4ビット精度下でも損失なしの収束を保証する。

Figure 1 : Visualization of the first moment in the layers.3.blocks.1.mlp.fc1 layer in a Swin-T model. (a): Magnitude of the first moment. (b): Histogram of the first moment. (c): Moment approximated by B128/DE . (d): Moment approximated by B2048/DE .

実験結果

リサーチクエスチョン

• RQ1 最適化状態を4ビットに圧縮しても、モデルの収束性や精度を損なわずに可能か？
• RQ2 特にアウーターパターンの観点から、1回目および2回目のモーメントの4ビット圧縮における量子化の課題は何か？
• RQ3 2回目のモーメント量子化におけるゼロポイント問題が最適化の安定性に与える影響は何か？そして、これを軽減できるか？
• RQ4 提案された量子化技術は、安定な埋め込み層に依存せずに、フル精度最適化手法と同等の性能を達成できるか？
• RQ5 多様なタスクにおいて、既存の8ビットベースラインと比較して、4ビット最適化手法のメモリおよび学習効率の向上はどの程度か？

主な発見

• 提案された4ビット最適化手法は、NLP、機械翻訳、画像分類、インstructチューニングを含むすべてのベンチマークで、フル精度Adamと同等の収束速度と最終的な精度を達成する。
• 8ビット最適化手法と比較して、メモリ消費量を50%削減し、既存の低精度最適化手法の中で最も低いメモリフットプリントを達成する。
• 2回目のモーメントのゼロポイントフリー線形量子化は、パラメータ更新における大きなずれを効果的に防止し、4ビット精度での安定した学習を可能にする。
• ランク1正規化は、特に複雑なアウーターパターンが存在する状況で、2回目のモーメントの近似を顕著に向上させる。
• 安定な埋め込み層を必要とせず、さまざまなモデルアーキテクチャにわたって、収束性と一般化性能に優れたロバスト性を示す。
• 最適化状態のオフロードを適用することで、4ビット最適化手法は、より低い精度の状態に起因する通信コストの低減により、言語モデル微調整におけるスループットを向上させる。

Figure 2 : Outlier patterns vary across two first moment tensors. (a): outliers lie in fixed rows (dimension 0). (b): outliers lie in fixed columns (dimension 1).