[論文レビュー] BitNet: Scaling 1-bit Transformers for Large Language Models

Hongyu Wang, Shuming Ma+8 more|arXiv (Cornell University)|Oct 17, 2023

Topic Modeling被引用数 26

ひとこと要約

BitNetは1ビットTransformerとBitLinearを導入し、ビナリ化された重みでスクラッチから訓練可能にしつつ、パープレキシティで競争力を維持し、FP16および事後訓練8ビットベースラインと比べて顕著なメモリとエネルギーの節約を実現する。さらに、フル精度のTransformerと類似のスケーリング法則を示す。

ABSTRACT

The increasing size of large language models has posed challenges for deployment and raised concerns about environmental impact due to high energy consumption. In this work, we introduce BitNet, a scalable and stable 1-bit Transformer architecture designed for large language models. Specifically, we introduce BitLinear as a drop-in replacement of the nn.Linear layer in order to train 1-bit weights from scratch. Experimental results on language modeling show that BitNet achieves competitive performance while substantially reducing memory footprint and energy consumption, compared to state-of-the-art 8-bit quantization methods and FP16 Transformer baselines. Furthermore, BitNet exhibits a scaling law akin to full-precision Transformers, suggesting its potential for effective scaling to even larger language models while maintaining efficiency and performance benefits.

研究の動機と目的

• 性能を犠牲にせず、Large Language Models（大規模言語モデル）におけるメモリとエネルギーコストの削減を動機づける。
• トレーニング可能なスクラッチから訓練できるドロップイン型の1ビットTransformerコンポーネント（BitLinear）を導入する。
• BitNet の安定性、スケーラビリティ、効率を FP16 および事後訓練量子化ベースラインと比較して示す。
• BitNet はフル精度のTransformerと類似のスケーリング法則に従い、より大きなモデルへ効果的にスケールできることを示す。

提案手法

• nn.Linear を BitLinear に置換して、最適化子の状態と勾配を高精度のままにしつつ、ビナリ化された重みで動作させる。
• 重みを ±1 に二値化する。sign(W - alpha) を用いて中央化をゼロ平均にし、二値化後にスケーリング beta を適用。
• 訓練時には absmax で 8-bit へ活性化を量子化し、テンソル単位量子化を行う；推論時はトークンごとに量子化。
• 活性化量子化前に LayerNorm を適用する（SubLN）ことで、量子化経路を通じた分散を安定化。
• 追加の通信なしに効率的なモデル並列性を可能にするため、 Group Quantization と Group Normalization を実装。
• 非微分可能なステップには straight-through estimator (STE) を用い、混合精度（重み/活性化は低精度、勾配/最適化子は高精度）で訓練する；収束を改善するために大きな学習率を使用。

実験結果

リサーチクエスチョン

• RQ1 量子化認識訓練を経て訓練された1ビット重みは、大規模言語モデルにおいて競争力のある perplexity およびダウンストリームタスクパフォーマンスを達成できるか。
• RQ2 モデルサイズが拡大するにつれて、エネルギー消費とメモリフットプリントの観点で、BitNetはFP16 Transformerおよび事後訓練量子化手法とどう比較されるか。
• RQ3 BitNet はフル精度Transformerに類似のスケーリング法則を示すか、推論エネルギーのスケーリング挙動はどうなるか。
• RQ4 1ビット重みでの訓練を達成するために、どのような訓練安定化手段（例：SubLN、absmax 活性化量子化）が重要か。

主な発見

• BitNet は FP16 Transformer および 8-bit 量子化ベースラインと比較して、メモリとエネルギー使用を大幅に削減しつつ、競争力のある perplexity およびダウンストリームタスク精度を達成する。
• BitNet のエネルギー節約は、1ビット重みとスケール因子への乗算制限のため、行列乗算で特に顕著。
• BitNet はフル精度Transformerに類似のスケーリング法則に従い、FP16 との推定性能ギャップはモデルサイズが大きくなるにつれて縮小する。
• 同一学習率設定下で FP16 より訓練安定性が高く、より大きな学習率とより速い収束を可能にする。
• ゼロショットおよびフewショットのダウンストリームタスク性能は計算予算とともにスケールし、より大規模なスケールでの BitNet の推論効率の優位性を示す。