BitNet: Scaling 1-bit Transformers forLarge Language Models
Hongyu Wang†‡ Shuming Ma† Li Dong† Shaohan Huang† Huaijie Wang§ Lingxiao Ma† Fan Yang† Ruiping Wang‡ Yi Wu§ Furu Wei†⋄
† Microsoft Research ‡ University of Chinese Academy of Sciences § Tsinghua University arxiv.org
Abstract
-
目的：大規模言語モデルの展開における課題と、高エネルギー消費による環境への影響に対処するため、スケーラブルで安定した1ビットTransformerアーキテクチャであるBitNetを導入すること。
- 手法：nn.Linear層の代わりにBitLinearを導入し、1ビット重みをゼロから訓練することで、大規模言語モデル用に設計されたBitNetを開発。
- 結果：言語モデリングにおける実験結果から、BitNetは競争力のある性能を達成し、最先端の8ビット量子化方法およびFP16 Transformerベースラインと比較して、メモリフットプリントとエネルギー消費を大幅に削減することが示された。さらに、BitNetはフルプレシジョンTransformerと同様のスケーリング法則を示し、効率性と性能の利点を維持しながら、さらに大きな言語モデルへの有効なスケーリングの可能性を示唆している。
Introduction
- 大規模言語モデル（LLM）の急速な成長は、さまざまなタスクで顕著な改善をもたらしているが、高い推論コストとエネルギー消費により、これらのモデルをホストすることは費用がかかる。
- モデルのサイズが大きくなるにつれて、モデルパラメータへのアクセスと処理に必要なメモリ帯域幅が主要なボトルネックとなり、全体的な推論性能を制限している。
- 分散システムやマルチデバイスプラットフォーム上でこれらのモデルを展開する際、デバイス間通信のオーバーヘッドが推論の遅延とエネルギー消費に大きな影響を与えている。
- モデル量子化は、大規模モデルのメモリフットプリントと計算コストを大幅に削減しながら、競争力のある性能を維持できる有望な解決策である。
- 既存の量子化アプローチの多くは、トレーニング後に適用可能なため、簡単に使用できる。しかし、モデルは量子化表現で訓練中、最適化されていないため、精度が低下しやすい。一方で、量子化認識トレーニング（quantization-aware training）は、初期段階から量子化を考慮してモデルをトレーニングするため、より良い精度を実現しやすい。
- 本研究は、1ビット大規模言語モデルの量子化認識トレーニングを調査する最初の作業である。BitNetという1ビットTransformerアーキテクチャを提案し、メモリと計算の両方の面で効率的にスケールすることを目指している。
- BitNetは、低精度のバイナリ重みと量子化されたアクティベーションを使用しながら、トレーニング中に最適化器の状態と勾配には高精度を維持する。
- BitNetの実装は単純であり、Transformer内の線形射影（PyTorchのnn.Linearなど）の置換のみを要求する。また、PagedAttention、FlashAttention、推測デコーディングなど、他の大規模言語モデルの加速方法と補完しあう。
- 言語モデリングのベンチマークにおいてBitNetを評価し、最先端の量子化方法およびFP16 Transformerと比較した。実験結果は、BitNetがパープレキシティと下流タスクの精度の両方で競争力のある性能を達成し、メモリフットプリントとエネルギー消費をベースラインと比較して大幅に削減することを示している。
BitNet
- BitNetは、セルフアテンションとフィードフォワードネットワークのブロックを積み重ねる、Transformerと同じレイアウトを使用している。従来の行列乗算の代わりに、1ビットモデル重みを使用するBitLinearを採用している。
- その他のコンポーネントは以下の理由から量子化せずに高精度（例：8ビット）を維持している。
BitLinear
- 重みを符号関数を用いて+1または−1に二値化し、二値化前に重みをゼロ平均にすることで、限られた数値範囲内での容量を増加させる。二値化後には、実数値と二値化された重み間の$l2$エラーを減らすためにスケーリング係数$β$を使用する。重み$W \in \mathbb{R}^{n \times m}$の二値化は以下のように定式化される。
ここで
- アクティベーションはbビット精度にさらに量子化される。absmax量子化を使用し、入力行列の絶対最大値で割ることにより、アクティベーションを範囲$[−Q_b, Q_b] (Qb = 2^{b−1})$にスケールする。
$$ x_e = \text{Quant}(x) = \text{Clip}\left(\frac{x \times Q_b}{\gamma}, -Q_b + \epsilon, Q_b - \epsilon\right), \quad $$
ここで$\epsilon$ははクリッピングを実行する際にオーバーフローを防ぐための小さな浮動小数点数
この研究では、アクティベーションを8ビットに量子化し、より低い精度は将来の研究とする。また、トレーニング中はテンソルごとに、推論中はトークンごとに量子化を行い、安定性と効率性を確保する。
上記の量子化方程式を用いると、行列乗算は以下のように記述できる
$$ y=\tilde{W}\tilde{x} $$
（～は，チルダ（tilde）真値にほぼ等しい）
ここで、$W$と$x$の要素は互いに独立であり、同じ分布を共有していると仮定する。その後、出力$y$の分散は以下のように推定される：
$$ \text{Var}(y) = n\text{Var}(\tilde{w}\tilde{x}) $$
$$ = n\mathbb{E}[\tilde{w}^2]\mathbb{E}[\tilde{x}^2] \quad $$
- フルプレシジョン（全精度）計算では、Kaiming初期化やXavier初期化などの標準的な初期化方法を使用することで、出力の分散Var(y)が1のスケールに保ち、トレーニングを安定させる。
- 量子化による精度の低下を防ぎ、出力の分散を維持するために、活性化量子化の前にLayerNorm（レイヤー正規化）関数を導入する。これにより、出力$y$の分散は
と推定され、これはフルプレシジョンの出力の分散$Var(y)$と同じ大きさになります。 - Transformerモデルの文脈では、このプロセスはSubLN（サブレイヤー正規化）として実装されている。これは、BitLinearの導入により、1ビット重みと量子化されたアクティベーション間の行列乗算を可能にします。
- BitLinearは、SubLNおよび量子化方法を用いて次のように定式化される
$$ y = \tilde{W}\tilde{f} = \tilde{W}Quant(LN(x)) × \frac{βγ}{ Q_b} $$
ここで、$β$はスケーリング係数、$γ$は正規化のスケール、は$Q_b$量子化のビット数を表す。SubLN操作後、活性化はabsmax関数で量子化され、1ビット重みと量子化された活性化間で行列乗算が実行されます。出力活性化は${β, γ}$で再スケールされ、元の精度に逆量子化される。
Model parallelism with Group Quantization and Normalization
大規模言語モデルのスケールアップにはモデル並列性が重要であり、これは複数のデバイス上での行列乗算を分割する技術である。ただし、全てのパラメータα、β、γ、ηはテンソル全体から計算され、これが独立性の前提を破る。
一つの解決策として、各パラメータに対してall-reduce操作を導入することが考えられるが、モデルが深くなるにつれて通信の量が増加し、処理が遅くなる。
この問題を解決するために、重みとアクティベーションをグループに分割し、各グループのパラメータを独立して推定することで、追加の通信なしにパラメータをローカルで計算できる新しいアプローチを提案する。この手法はグループ量子化と呼ばれる。
具体的には、重み行列 $W \in \mathbb{R}^{n \times m}$をパーティション次元に沿って G グループに分割し、各グループが
のサイズを持つようにする。次に、各グループのパラメータを独立して推定する：
ここで、$W^{(g)}$ は重み行列の g 番目のグループを表す。同様に、入力行列 $x\in \mathbb{R}^{n \times m}$ を G グループに分割し、各グループのパラメータを計算する：
- LN（Layer Normalization）については、グループ正規化技術を適用して、各グループの平均と分散を独立に計算できる。
$$ \text{LN}(x^{(g)}) = \frac{x^{(g)} - \mathbb{E}(x^{(g)})}{\sqrt{\text{Var}(x^{(g)}) + \epsilon}} $$
- この方法により、追加の通信を必要とせずに、より効率的なモデル並列性が実現される。
Model Training
Straight-through estimator
- ビットモデルのトレーニングには、バックプロパゲーション中の勾配を近似するためにストレートスルー推定器（STE）を使用する。この方法は、バックワードパス中の非微分可能な関数（例：Sign関数やClip関数）を回避し、量子化モデルのトレーニングを可能にする。
Mixed precision training
- 重みと活性化は低精度に量子化されるが、勾配とオプティマイザの状態はトレーニングの安定性と精度を保証するために高精度のまま保存される。学習可能なパラメータのために、高精度フォーマットの潜在重みを保持し、パラメータ更新を蓄積する。潜在重みはフォワードパス中に二値化され、推論プロセスには使用されない。
Large learning rate
- 最適化の課題の一つは、潜在重みの小さな更新が1ビット重みにほとんど違いを生じさせないことである。これはバイアスのかかった勾配と更新を生じさせ、特にトレーニングの初期段階で問題となる。この課題に対処するため、学習率を増加させることが最適化を加速する最も単純で最良の方法であることを発見した。BitNetは大きな学習率から収束において利益を得るが、同じ学習率でFP16 Transformerはトレーニングの開始時に発散する。
Computational Efficiency
- BitNetの計算効率は、算術演算のエネルギーとメモリフットプリントの両方の観点で評価される。
- [Hor14, ZZL22]におけるエネルギーモデルによれば、異なる算術演算のエネルギー消費は以下のように推定される。
バニラTransformerにおけるエネルギー消費
$m×n$ と$n×p$ の次元を持つ行列乗算では、エネルギー消費は加算と乗算で次のように計算される
BitNetにおけるエネルギー消費
BitNetでは、1ビットの重みを使用するため、行列乗算のエネルギー消費は加算操作によって支配される。乗算操作は出力をスケーラー$β$と$γ/Q_b$ でスケーリングするためにのみ適用されるため、乗算のエネルギー消費は E${\text{mul}} = (m \times p + m \times n) \times \hat{E}_{\text{mul}}$として計算でき、これはTransformerに比べて著しく小さい。
Comparison with FP16 Transformers
Setup
- BitNetを用いた様々なスケールの自己回帰言語モデルを125Mから30Bの範囲でトレーニングする。モデルはPileデータセット、Common Crawlスナップショット、RealNews、CC-Storiesデータセットから構成される英語コーパス上でトレーニングされる。データの前処理にはSentencpieceトークナイザーを使用し、語彙サイズは16Kである。BitNetに加えて、公平な比較のために同じデータセットと設定でTransformerベースラインもトレーニングする。
Inference-Optimal Scaling Law
- ニューラル言語モデルはバニラTransformerアーキテクチャで予測可能にスケールすることが証明されており、損失はトレーニングに使用される計算量のべき乗則に従ってスケールする。これにより、計算予算の最適な割り当てを決定し、小さいモデルから大規模言語モデルの性能を予測できる。
- バイナライズドTransformerのスケーリング法則を調査するため、BitNetとFP16 Transformerベースラインのパラメータカウントに対するスケーリング曲線をプロットする。BitNetの損失スケーリングはFP16 Transformerに似ており、べき乗則（$L(N) = aNb + c$）に従う(下図で、125Mから6.7Bのモデル結果でべき乗則をフィッティングをしたが13Bおよび30B、の。また、BitNetとFP16 Transformer間のギャップはモデルサイズが大きくなるにつれて小さくなる。
- Inference-Optimal Scaling Lawを導入し、エネルギー消費に対する損失を予測する。これはモデルの使用量に応じてスケールする推論エネルギーコストに焦点を当て、トレーニングコストよりも効率的なスケーリングを提供する。BitNetは固定された計算予算で顕著に良い損失を達成し、FP16モデルと同じ性能を得るための推論コストは大幅に小さい。
Results on Downstream Tasks
- BitNetのスケーリングに伴う能力についても、損失と同様に関心がある。Hellaswag、Winogrande、Winograd、Storyclozeを含む4つの下流タスクで、0ショットと4ショットの結果をテストし、解釈可能な指標でスケーリングに伴う能力を評価する（下図）。BitNetとFP16 Transformerの平均結果を報告し、計算予算が増えるにつれて下流タスクの性能もスケールすることが示される。
- 低ビットTransformerのトレーニングにおける主要な課題は最適化の安定性であるため、BitNetとFP16ベースラインの安定性テストを、異なるピーク学習率でのモデルシリーズのトレーニングによって行う。BitNetは大きな学習率で収束できるが、FP16 Transformerはできないことが示され（下図）、BitNetのトレーニング安定性がより優れていることを示している。
<
- BitNetは学習率の増加から恩恵を受け、PPL(perplexity)の観点でより良い収束を達成できることが示される。
Comparison with Post-training Quantization
- BitNetをAbsmax、SmoothQuant、GPTQ、QuIPを含む最先端の量子化方法と比較する。これらの方法はFP16 Transformerモデル上でのトレーニング後量子化であり、BitNetと同じトレーニング設定およびデータに従う。AbsmaxとSmoothQuantは重みとアクティベーションの両方を量子化し、GPTQとQuIPは重みの精度のみを下げる。
- 重みのみの量子化（GPTQとQuIP）については、W4A16とW2A16で実験を行い、重みとアクティベーションの量子化（AbsmaxとSmoothQuant）には、FP16 TransformerをW8A8、W4A4、W1A8に量子化する。BitNetの実装はバイナリ重み8ビットアクティベーション（W1A8）であり、ベースラインよりも低いか同等のビット数である。
- Winogrande、Winograd、Storycloze、Hellaswagの4つのベンチマークデータセットにおける様々なベースラインアプローチに対するBitNetの提案方法のゼロショット性能の詳細な比較分析を提示する。
- 公平な比較のために、すべてのモデルは6.7Bのモデルサイズを持ち、16から1に至るまでのいくつかの重みビットレベルで評価される。評価指標には、下流タスクのゼロショット精度に加えて、各方法の性能を包括的に理解するための検証セット上の言語モデルパープレキシティが含まれる。
- BitNetは、特に低ビットレベルで、ベースラインアプローチと比較して近い性能レベルを達成している。BitNetのゼロショットスコアは8ビットモデルに匹敵するが、推論コストははるかに低い。
- 4ビットモデルにおいては、重みのみを量子化する方法が重みとアクティベーションの量子化器よりも性能が良い。これは、アクティベーションを量子化することがより困難であるためである。1ビットモデルであるBitNetは、重みとアクティベーションの量子化方法および重みのみの方法よりも著しく優れた結果を達成している。
- 低ビットモデルに関して、BitNetはすべてのベースラインに対して一貫して優れたスコアを持っている。これは、トレーニング後の量子化方法よりも量子化認識トレーニングアプローチの利点を証明している。1.3Bから6.7Bまでモデルサイズをスケールアップする際の、私たちの方法とベースラインのゼロショット精度とフューショット精度の両方を要約する図6は、この利点が異なるスケールで一貫していることを証明している。
Ablation Studies
- BitNetは活性化の量子化にabsmaxを使用し、トレーニングの安定性のためにSubLNを使用する。量子化の代替案として、学習可能なパラメータでスケールを動的に調整するelastic関数がある。実験では、absmaxがelastic関数よりも優れた性能を示すことがわかる。
- さらに、absmax関数はより安定したトレーニングをもたらし、BitNetに対してより大きな学習率を可能にする。SubLNをPre-LNおよびBMTアーキテクチャと比較する。Pre-LNはGPTのデフォルトアーキテクチャであり、BMTはバイナライズドモデルの安定性を改善することが証明されている。実験では、SubLNがPre-LNとBMTの両方を上回ることを示す。したがって、BitNetの実装にはabsmaxとSubLNを選択する。
Conclusion and Future Work
- BitNet、大規模言語モデル用の新しい1ビットTransformerアーキテクチャを紹介した。このアプローチは、大規模言語モデルを効率的に扱うことができるスケーラブルで安定した設計を目指している。
- 実験結果は、BitNetがパープレキシティと下流タスクのパフォーマンスの両方において競争力のある性能を達成し、ベースラインと比較してメモリフットプリントとエネルギー消費を大幅に削減することを示している。さらに、BitNetはフルプレシジョントランスフォーマーと同様のスケーリング法則に従い、パフォーマンスと効率の面で潜在的な利点を持ってさらに大きな言語モデルに効果的にスケールアップできることを示している。
- 将来的には、モデルサイズとトレーニングステップの面でBitNetをスケールアップすることを目指している。また、大規模言語モデルのトレーニングにおいてBitNetを他のアーキテクチャ（例：RetNet）に適用することにも関心がある。
この論文の発展版1.58bitの要約はこちら reseachpaper-matome.hatenablog.com
