Bi-Mamba: Towards Accurate 1-Bit State Space Models

要約

Mamba の典型的な選択的状態空間モデル (SSM) は、シーケンス長による 2 次計算の複雑さや、キーと値のキャッシュによる大幅な推論時間のメモリ要件など、Transformer のいくつかの制限に対処しています。
しかし、Mamba モデルのサイズが増大することにより、トレーニングと展開に課題が生じ続け、エネルギー消費量が大幅に増加するため環境への懸念も生じています。
この研究では、780M、1.3B、および 2.7B にわたる複数のサイズを持つ、より効率的な大規模言語モデル向けに設計された、スケーラブルで強力な 1 ビット Mamba アーキテクチャである Bi-Mamba を紹介します。
Bi-Mamba モデルは、自己回帰蒸留損失を使用する通常の LLM として、データ量に基づいて最初からトレーニングされます。
言語モデリングに関する広範な実験結果は、Bi-Mamba がメモリ フットプリントとエネルギーを大幅に削減しながら、完全精度の対応物 (FP16 や BF16 など) に匹敵するパフォーマンスと、トレーニング後二値化 (PTB) Mamba ベースラインよりもはるかに高い精度を達成することを実証しています。
オリジナルの Mamba モデルと比較した消費量。
私たちの研究は、低ビット表現における新しい線形計算複雑性 LLM フレームワークの先駆者であり、効率的な 1 ビット Mamba ベースの LLM に合わせた特殊なハードウェアの将来の設計を促進します。

要約(オリジナル)

The typical selective state-space model (SSM) of Mamba addresses several limitations of Transformers, such as quadratic computational complexity with sequence length and significant inference-time memory requirements due to the key-value cache. However, the growing size of Mamba models continues to pose training and deployment challenges and raises environmental concerns due to considerable energy consumption. In this work, we introduce Bi-Mamba, a scalable and powerful 1-bit Mamba architecture designed for more efficient large language models with multiple sizes across 780M, 1.3B, and 2.7B. Bi-Mamba models are trained from scratch on data volume as regular LLM pertaining using an autoregressive distillation loss. Extensive experimental results on language modeling demonstrate that Bi-Mamba achieves performance comparable to its full-precision counterparts (e.g., FP16 or BF16) and much better accuracy than post-training-binarization (PTB) Mamba baselines, while significantly reducing memory footprint and energy consumption compared to the original Mamba model. Our study pioneers a new linear computational complexity LLM framework under low-bit representation and facilitates the future design of specialized hardware tailored for efficient 1-bit Mamba-based LLMs.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google