[論文] エネルギー効率の良いAI : 計算資源を最小限に抑える言語モデルの開発

[Microsoft]『The Era of 1-bit LLMs : All Large Language Models are in 1.58 Bits 』を読む
まとめ

[Microsoft]『The Era of 1-bit LLMs : All Large Language Models are in 1.58 Bits 』を読む

論文の概要

論文は、BitNet b1.58がどのようにして従来の全精度モデルと同等の性能を、より効率的に達成しているかを説明しています。具体的には、1ビットの計算を用いることで、メモリ使用量、計算速度、エネルギー消費を大幅に削減しながら、言語理解タスクにおける高い精度を維持する技術について論じています。このアプローチは、大規模言語モデルの将来の方向性を示唆し、特にリソースが限られた環境での応用において、その有効性を強調しています。

論文の構成

1bit LLMの時代
BitNet b1.58
結果
議論と将来の研究

1. 1bit LLMの時代

1ビットLLMの時代

このセクションは、大規模言語モデル（LLM）の急速な進化と、それに伴うエネルギー消費とデプロイメントの課題に焦点を当てています。LLMのサイズ拡大により、環境および経済への影響が懸念される中、推論用の低ビットモデルを作成するためのトレーニング後の量子化が1つの解決策として提示されています。この技術は、メモリと計算要件を削減することで、16ビットから4ビットなどのより低いビットへの移行を促進します。

特に、1ビットモデルアーキテクチャであるBitNetの開発は、性能を維持しつつLLMのコストを削減する有望な進歩を示しています。BitNetは、行列乗算で整数加算のみを使用することで、エネルギーコストを節約し、より速い計算を可能にします。また、1ビットLLMはDRAMからのモデルパラメータ転送コストを削減し、推論をより高速かつ効率的にします。

この作業で紹介されるBitNet b1.58は、すべてのパラメータが三値（-1、0、1）を取るという重要な1ビットLLM変種です。この新しいモデルは、乗算演算をほぼ不要とする新しい計算パラダイムを採用し、FP16 LLMベースラインに比べてメモリ消費、スループット、およびレイテンシーの面で効率的です。さらに、モデル重みに0を含めることで機能フィルタリングをサポートし、1ビットLLMの性能を向上させると同時に、同じ構成でフルプレシジョンベースラインと一致する性能を実現しています。

2. BitNet b1.58

BitNet b1.58とは？

トランスフォーマーモデルに基づいたBitNetアーキテクチャを採用しており、nn.LinearをBitLinearに置き換えています。このモデルは、1.58ビットの重みと8ビットのアクティベーションを使用して、ゼロからトレーニングされます。オリジナルのBitNetに比べて、いくつかの変更が加えられています。

量子化関数

重みを-1、0、+1に制限するため、absmean量子化関数を採用しています。この関数は、まず重み行列をその平均絶対値でスケーリングし、次に各値を{-1、0、+1}の中で最も近い整数に丸めます。

アクティベーションの量子化

アクティベーションの量子化関数はBitNetと同様の実装を採用していますが、非線形関数の前にアクティベーションを[0、Qb]の範囲にスケーリングすることはありません。代わりに、アクティベーションはすべてトークンごとに[−Qb、Qb]にスケーリングされ、ゼロポイントの量子化を排除します。これにより、実装とシステムレベルの最適化がより便利で簡単になり、実験では性能にほとんど影響を与えません。

LLaMA風コンポーネント

LLaMAアーキテクチャはオープンソースLLMの事実上のバックボーンです。オープンソースコミュニティを受け入れるため、BitNet b1.58の設計はLLaMA風のコンポーネントを採用しています。具体的には、RMSNorm、SwiGLU、ロータリーエンベディングを使用し、すべてのバイアスを取り除きます。このようにして、BitNet b1.58は、例えばHuggingface vLLMやllama.cppなどの人気のあるオープンソースソフトウェアに、最小限の努力で統合することができます。

3. 結果

BitNet b1.58は、再現されたFP16 LLaMA LLMと様々なモデルサイズで比較し、性能、GPUメモリ使用量、レイテンシー、およびエネルギー消費において顕著な効率と性能の利点を示しています。

図2：モデルサイズを変化させたBitNet b1.58のデコード待ち時間（左）とメモリ消費量（右）
論文：『The Era of 1-bit LLMs : All Large Language Models are in 1.58 Bits 』より抜粋

具体的には、BitNet b1.58はFP16 LLaMA LLMと同等のパープレキシティを達成しつつ、高速でメモリ使用量が少ないことが特徴です。エネルギー消費に関しては、顕著な削減を実現し、スループットにおいてもLLaMA LLMに比べて大幅に高い効率を示しています。

図3：BitNet b1.58のエネルギー消費量とLLaMA LLMの比較（7nmプロセスノード）。左は演算エネルギーの構成要素。右側は、異なるモデル・サイズにわたるエンドツーエンドのエネルギー・コスト。
論文：『The Era of 1-bit LLMs : All Large Language Models are in 1.58 Bits 』より抜粋

これらの結果から、BitNet b1.58はフルプレシジョンモデルに匹敵するかそれを超える性能を、はるかに低い計算およびエネルギーコストで提供することが実証されています。

4-1. 議論と将来の作業：MoE（Mixture-of-Experts）モデルの課題

MoEモデルの課題にどのように対処するか？

MoE（Mixture-of-Experts）モデルの課題には、主に高いメモリ消費とチップ間通信のオーバーヘッドがあります。これらの課題は、MoEモデルのデプロイメントとアプリケーションを制限しており、特に大規模なモデルや複雑なタスクを扱う場合に顕著です。メモリ消費は、多数の専門家を持つMoEモデルが大量の計算リソースを必要とするために発生します。一方、チップ間通信のオーバーヘッドは、これらの専門家間でのデータの転送と処理が複雑になることに起因します。このような課題に対処するため、1.58ビットLLMのような新しいアプローチが提案されています。

MoEモデルの課題に対する1.58ビットLLMの新しいアプローチとは？

1.58ビットLLMによるMoE（Mixture-of-Experts）モデルの課題への新しいアプローチは、主にメモリフットプリントの削減とチップ間通信のオーバーヘッドの削減に焦点を当てています。このアプローチは、メモリ消費を大幅に削減し、MoEモデルをデプロイするために必要なデバイスの数を減少させることができます。さらに、ネットワークを介したアクティベーションの転送に関連するオーバーヘッドも大幅に削減します。理想的には、モデル全体を単一のチップ上に配置することが可能になり、これによってチップ間通信のオーバーヘッドを完全に排除することができます。このように、1.58ビットLLMは、MoEモデルの課題を効果的に解決するための有望なソリューションを提供します。

4-2. 議論と将来の作業：長いシーケンスの推論における主要な課題

LLMの長いシーケンスのネイティブサポート

LLM（Large Language Models）の時代において、長いシーケンスを扱う能力は重要な要求となっています。長いシーケンスの推論における主要な課題の一つは、KVキャッシュによって導入されるメモリ消費です。

BitNet b1.58は、アクティベーションのビットサイズを16ビットから8ビットに削減することで、大規模言語モデル（LLM）における長いシーケンスのネイティブサポートを向上させる貢献をしています。この削減により、同じメモリと計算リソースでコンテキストの長さを2倍にできるため、メモリ消費を抑えつつ長いテキストを効率的に処理できるようになります。このアプローチは、特に長い文書や文章の解析が必要なタスクにおいて、LLMの適用範囲を広げることを可能にします。

4-3. 議論と将来の作業：1.58ビットLLMの潜在的な可能性

1.58ビットLLMの潜在的な可能性

1.58ビットLLMの潜在的な可能性には、メモリと計算能力に制約があることが多いエッジおよびモバイルデバイスでの言語モデルの性能を向上させる可能性があることが強調されている。
BitNet b1.58によって導入された計算パラダイムに触発され、1ビットLLM用に最適化された特定のハードウェアとシステムの開発を呼びかけている。