【BitNet】1ビットで推論する大規模言語モデル
3つの要点
✔️ ±1の２種類の重みだけを使うBitNetを提案
✔️ 従来より遥かに高いエネルギー効率とメモリ効率を達成
✔️ 通常のTransformerと同様のスケーリング則に従うことを発見
BitNet: Scaling 1-bit Transformers for Large Language Models
written by Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Huaijie Wang, Lingxiao Ma, Fan Yang, Ruiping Wang, Yi Wu, Furu Wei
(Submitted on 17 Oct 2023)
Comments: Work in progress
Subjects: Computation and Language (cs.CL)code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
概要
大規模言語モデルは高精度達成のためにますます巨大化していますが、巨大化につれその展開に課題が生じ、計算量やエネルギー消費量の増大が懸念されています。この研究では、重みが$±1$の1ビットで構成されるTransformerを提案し、それが従来の 16ビットモデルと同等の性能を、遥かに少ない計算リソースと遥かに高いエネルギー効率で達成できることを示しました。また、興味深いことに、モデルのサイズ増大に伴って、従来のTransformerと同様のスケーリング則に従うことを発見しました。この革新的な手法は、軽量で扱いやすい1ビット大規模言語モデルの基礎であると言えます。
背景
大規模言語モデルの巨大化とモデル量子化への期待
大規模言語モデルは巨大化の一途を辿り、今後も更なる巨大化が予想されます。しかし、その推論コストとエネルギー消費の大きさは、大規模言語モデルの使用コストを著しく増加させます。そのため、メモリ使用量や計算負荷軽減を目的に、大規模言語モデルのパフォーマンスを維持しながら、モデルを軽量化する方法論としてモデルの量子化（パラメータの整数化）に注目が集まっています。
モデル量子化の従来研究
現状の多くのモデル量子化は、学習後の行われるものです。これは、モデルの学習パイプラインの変更や再学習を必要としないため、単純で容易に適用することができます。しかし、学習中に量子化のために最適化されているわけではないため、量子化の結果として著しく精度が損なわれる可能性があります。また、モデル量子化のもう一つの方法は、学習中に量子化を考慮することです。これによって、大規模言語モデルには不可欠な、モデルの継続学習や微調整が可能になります。しかし、学習中にモデルの量子化を考慮する上での課題は、その最適化にあります。つまり、モデルの精度が低くなるにつれその収束が悪くなる傾向があることです。加えて、量子化を考慮した学習が言語モデルのスケーリング則に従うかは非自明です。
これらを背景として、筆者らは量子化を考慮した学習機構を伴ったBitNetを提案しました。また、その学習におけるスケーリング則に関してもこの研究中では議論されています。
提案手法
BitNet
BitNetとは端的に言えば、従来のTransformerのLiner層を1ビットの重みによって表現されるBitLiner層に置換したものです。図１にBitNetの簡単な概要を示しています。その他に関しては、通常のTransformerと同様の機構を持っています。
モデルの訓練
この研究では、BitNetの訓練のためにいくつかの工夫を導入しています。以下に、それらを簡単に説明します。
・Straught-through estimator (STE)
BitNetはいくつかの微分不可能な関数を含んでいます。そのため、勾配計算の際はそれらの関数を迂回するようにバイパスすることで、それらの関数の直接的な微分を回避しています。これをこの研究中では、STEと呼んでいます。
・Mixed Precision tranding
BitNet中の重みは量子化されていますが、BitNetの訓練の安定性と精度を確保するために勾配の計算や最適化の状態の管理には高い精度の変数を準備し、それらを用いてパラメータのアップデートを行います。その際に、計算用の潜在的な重みを用います。ただし、それらの潜在的な重みは、フォワードパス中に２値化し、推論プロセスには使用されません。
・Mixed Precision tranding
BitNetを訓練する上での、いくつかの課題は潜在的な重みの小さな更新が、1ビットの変数に違いを与えないことが多い点です。その結果、1ビットの重みの基づいて計算される勾配と更新に偏りが生じてしまいます。この課題を解決するために筆者たちは、学習率を大きくし最適化を加速させることが、最も簡単で最良の方法であると結論しています。筆者らの実験によると、BitNetは学習率を大きくすることで収束の面で有利になる一方で、FP16型の通常のTransformerは同じ学習率でも学習開始時に発散することが確認されました。これは、BitNetの学習の効率性と高い安定性を示唆するものです。
実験結果
計算の効率性
表１に通常のTransformerとBitNetのエネルギーの観点から計算効率性の比較を示しています。この結果から、モデルのサイズに関わらず、BitNetは通常のTransformerに比べて遥かに高い効率性を示していることが分かります。
ロスに関するスケーリング則の検討
図２にBitNetと通常のTransformerのモデルのサイズ増大に伴ったスケーリング則の比較を示しています。この結果から、重要なことにBitNetが通常のTransformerと同様に、モデルのサイズの増大に伴ってスケーリング則に従ってロスが小さくなることを示しています。また、モデルサイズもしくはエネルギー消費を一定とした時に、BitNetと通常のTransformerを比較すると、BitNetの方がロスが小さくなっていることが分かります。これは、エネルギー消費量やモデルのサイズに関して、BitNetが通常のTrasformerよりも高い効率性を持っていることを示しています。また、もう一つの重要な点として、モデルのサイズの増大に伴ってBitNetと通常のTransformerの誤差が小さくなっていることが確認できます。
精度に関するスケーリング則の検討
図３に精度に関するBitNetとTransformerの比較を示しています。この結果からも、BitNetが高い精度を示しており、その高い効率性が示唆されます。
学習の安定性に関する検証
図４(a)にBitNetとTransformerの学習の安定性を比較したものを示しています。この結果から、通常のTransformerと比較した時の、BitNetの学習の高い安定性が示唆されます。また、図４(b)にいくつかの学習率に対するBitNetの学習の履歴を可視化したものを示しています。この結果からも、学習率によらず安定的に学習を行うことができるBitNetの高い安定性と効率性が確認できます。
従来の学習後に行うモデル量子化との比較
図５にBitNetといくつかのベースラインの精度の比較をまとめたものを示している。この結果からも、その他の量子化の手法と比較してBitNetが高い精度を示していることが確認できます。
まとめ
この研究では、1ビットのTransformerをベースとした大規模言語モデルを提案しました。また、BitNetと従来の量子化手法や通常のTransformerを包括的に比較しその違いを議論しました。その結果、BitNetは高い効率性と従来の量子化手法を凌駕する精度を実現し得ることが示唆されました。加えて、通常のTransformerと比較しても遜色ない精度を示している点は驚くべき結果です。今後、大規模言語モデルの更なる巨大化に伴って、モデル量子化の手法の検討の重要性は増すことが予想されBitNetはその筆頭になり得るポテンシャルを秘めていると考えられます。その意味で、このモデルの汎用性や適用限界の更なる議論が期待されます。
この記事に関するカテゴリー
