IBMリサーチがカーネギーメロン大学やプリンストン大学、イリノイ大学と共同で、オープンソースの大規模言語モデル「Bamba」を構築し、そのバージョン2をオープンソースとしてリリースしました。

Meet Bamba, IBM’s new attention-state space model - IBM Research

https://research.ibm.com/blog/bamba-ssm-transformer-model



Bambaは97.8億のパラメーターを持つ大規模言語モデルで、ベースとなるアーキテクチャが一般的な大規模言語モデルと少し違う点が特徴です。

IBMリサーチによると、一般的な大規模言語モデルはTransformerというアーキテクチャを利用していますが、応答の際に実行中のシーケンスをメモリに保持する関係上、プロンプトが長くなるにつれて生成のコストが指数関数的に増大するとのこと。たとえばコンテキストウィンドウのサイズが2倍になると、それを処理して応答を生成するコストは2倍どころか4倍になるそうです。

この問題は「2次ボトルネック」と呼ばれ、ユーザーがAIに質問をしてから答えを得るまでのタイムラグの原因の1つになっているといいます。

新しく登場したBamba-9Bは、Transformerアーキテクチャと、状態空間モデル(SSM)というアーキテクチャを組み合わせつつ、メモリに当たるKVキャッシュの管理をTransformerアーキテクチャから根本的に変えたモデルです。通常、Transformerが応答を出力する際、コンテキストウィンドウ内のすべての単語に注意を払うのに対し、SSMは過去の情報を要約した「隠れ状態」を維持するとのこと。情報を選択的に保持するこの手法を使うことで、メモリのオーバーヘッドが少なくなり、推論速度が速くなるそうです。

詳しくは以下のサイトに記載されています。

Bamba-9B-v2 - Fast and powerful!

https://huggingface.co/blog/ibm-ai-platform/bamba-9b-v2



IBMリサーチによると、Bamba-9BはKVキャッシュのメモリ要件を大幅に削減することで、同サイズのTransformerベースのモデルと比べて、同等の精度を保ちながら少なくとも2倍の速度で動作できるとのこと。Transformerの能力と、SSMの実行速度を組み合わせることで、ボトルネックを解消しつつ応答精度を維持したモデルとなっています。

BambaはApache 2.0 ライセンスの下、オープンソースで公開されています。

GitHub - foundation-model-stack/bamba: Train, tune, and infer Bamba model

https://github.com/foundation-model-stack/bamba