[論文レビュー] The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits

Shuming Ma, Hongyu Wang+8 more|arXiv (Cornell University)|Feb 27, 2024

Natural Language Processing Techniques被引用数 37

ひとこと要約

本論文は BitNet b1.58 を紹介します。1.58ビットLLM で、全てのパラメータが三値 {-1, 0, 1}。同程度のサイズで FP16/LLaMA の性能と同等でありながら、速度・メモリ・エネルギーを大幅に向上させ、1ビットLLMの新しいスケーリング則と潜在的なハードウェア最適化を可能にします。

ABSTRACT

Recent research, such as BitNet, is paving the way for a new era of 1-bit Large Language Models (LLMs). In this work, we introduce a 1-bit LLM variant, namely BitNet b1.58, in which every single parameter (or weight) of the LLM is ternary {-1, 0, 1}. It matches the full-precision (i.e., FP16 or BF16) Transformer LLM with the same model size and training tokens in terms of both perplexity and end-task performance, while being significantly more cost-effective in terms of latency, memory, throughput, and energy consumption. More profoundly, the 1.58-bit LLM defines a new scaling law and recipe for training new generations of LLMs that are both high-performance and cost-effective. Furthermore, it enables a new computation paradigm and opens the door for designing specific hardware optimized for 1-bit LLMs.

研究の動機と目的

• 性能を犠牲にすることなく、LLMの推論コストとエネルギーを削減する動機付け。
• モデリング能力を高めるゼロ状態を含む 1.58-bit 重み量子化方式を提案する。
• BitNet b1.58 が 3B サイズから、困惑度（perplexity）とエンドタスク指標の両方で全精度ベースラインと同等であることを示す。
• ハードウェアとエネルギーへの影響を示し、より大規模モデルやエッジ機器へのスケーラビリティを議論する。

提案手法

• BitNet b1.58 を導入します。Linear 層を BitLinear に置換し、重みを {-1, 0, 1} に量子化した Transformer。
• weight を制約するために absmean 量子化を用い、gamma を平均絶対値として定義する。
• RMSNorm、SwiGLU、rotary embeddings など LLaMA 風のアーキテクチャ要素を採用し、バイアスなし。
• RedPajama データ上で、1.58-bit 重みと 8-bit アクティベーションでゼロから学習する。
• モデルサイズ（700M–3.9B）を横断して、FP16 LLaMA ベースラインと比較した困惑度とゼロショットエンドタスク性能を評価する。
• FasterTransformer と 2-bit/Ladder カーネルを用いて、メモリ、レイテンシ、スループット、推定エネルギーを測定する。スケーラビリティとエネルギー節約を分析する。

実験結果

リサーチクエスチョン

• RQ1 同じモデルサイズと同じ学習設定で、1.58-bit LLM が perplexity およびエンドタスクで全精度 LLM の性能と同等になり得るか？
• RQ2 モデルサイズが大きくなるにつれて、1.58-bit LLM のメモリ、レイテンシ、エネルギーの影響はどうなるか？
• RQ3 サイズを問わず、FP16 ベースラインとの性能差を縮め、3B 以降でエンドタスクの結果が同等となるか？
• RQ4 2T トークンでの訓練は、標準ベンチマーク上で一般化を維持または改善するよう BitNet b1.58 をスケールさせるか？

主な発見

• BitNet b1.58 は同じ設定で、3B サイズから perplexity において FP16 LLaMA ベースラインと一致する。
• BitNet b1.58 は 3B で FP16 LLaMAと比較して 2.71x 高速、GPU メモリは 3.55x 少なくなる。
• 3.9B では BitNet b1.58 は 2.4x 高速、メモリ効率は 3.32x 向上し、エンドタスクでは LLaMA 3B を上回る。
• サイズを問わず、BitNet b1.58 は FP16 ベースラインとの性能差を縮め、3B 以降でエンドタスクの結果が同等となる。
• BitNet b1.58 70B は LLaMA 70B に対して最大 4.1x のスピードアップを達成し、2 台の A100 GPU で最大 11x のスループット向上、メモリも削減。
• エネルギー分析では、7nmチップ上の行列乗算における演算エネルギーを FP16 ベースラインと比較して約 71.4x 節約し、より大きなモデルで利点が拡大します。