SSM-Transformerアーキテクチャ採用で従来の約3倍のスループットを実現した大規模言語モデル「Jamba」をAI21 Labsが発表

イスラエルのAIスタートアップであるAI21 Labsが、英語・フランス語・スペイン語・ポルトガル語に対応した大規模言語モデル「Jamba」を発表しました。Jambaは、従来のTransformerモデルにState Space Model(SSM)のアーキテクチャを組み合わせたSSM-Transformerモデルとなっています。

Introducing Jamba: AI21's Groundbreaking SSM-Transformer Model
https://www.ai21.com/blog/announcing-jamba

AI21 Labs Unveils Jamba: The First Production-Grade Mamba-Based AI Model
https://www.maginative.com/article/ai21-labs-unveils-jamba-the-first-production-grade-mamba-based-ai-model/

AI21 Labs’ new AI model can handle more context than most | TechCrunch
https://techcrunch.com/2024/03/28/ai21-labs-new-text-generating-ai-model-is-more-efficient-than-most/

Jambaは最大140Kトークン・約10万5000単語を、少なくとも80GBのメモリを備えた単一のGPUで処理できる大規模言語モデルです。コアモデルの一部に、プリンストン大学とカーネギーメロン大学で開発されたオープンソースモデル「Mamba」が使われています。

Jambaの特徴は、Transformerアーキテクチャに、「State Space Model(SSM)」と呼ばれるアーキテクチャを組み合わせた「SSM-Transformer」アーキテクチャを採用している点です。

SSMは時系列データや連続したデータを扱うための確率モデルの一種で、Transformerと比較して計算効率が高く、長いシーケンスのデータを処理するのに適しています。JambaのSSM-TransformerアーキテクチャはこのSSMとTransformerのハイブリッドであり、同じサイズのTransformerモデルと比較して、長いコンテキストに対して3倍のスループットを実現できるという特徴があります。

さらに、JambaではSSM-Transformerアーキテクチャに加えて、Mixture-of-Experts(MoE)レイヤーを採用しています。このMoEレイヤーは複数のニューラルネットワークを組み合わせてモデル全体の表現力と効率性を向上させる手法で、AI21 LabsはJambaに組み込むことで推論に利用可能な520億パラメータのうち、120億パラメータのみを使用しながら高いパフォーマンスを達成できたと報告しています。

以下はJamba(桃色)・Llama 2-70B(水色)・Mixtral-8x7B(黄色)で1秒当たりの処理トークン数(スループット)を比較したグラフです。Jambaのスループットは入力トークン＝64Kで他2モデルのほぼ2倍、入力トークン128Kで他2モデルの2～3倍を記録していることがわかります。

さまざまなベンチマークで、Llama 2-13B・Llama 2-70B・Gemma-7B・Mixtral-8x7BとJambaのパフォーマンスを計測した表が以下。Jambaは、Reasoning(推論)のベンチマーク4つのうち3つで他モデルよりも高いパフォーマンスを示しました。

JambaはApache 2.0ライセンスの下で、Hugging Faceで公開されています。また、NVIDIA AI Enterpriseソフトウェアプラットフォームで展開されているNVIDIA NIM推論マイクロサービスで、JambaのAPIにアクセス可能となっています。

なお、記事作成時点だとJambaは商用利用に必要なセーフガードが設けられていない研究モデルとしてリリースされていますが、AI21 Labsは2024年4月中には微調整したより安全なバージョンをリリースする予定だと述べています。