AI21 LabsがJambaとのコラボで生成AIトランスフォーマーを強化

2017年に画期的な研究論文「Attention is All You Need」が発表されて以来、トランスフォーマーは生成AIの領域で中心的な役割を果たしています。しかし、トランスフォーマーだけが生成AIの有力なアプローチではありません。AI21 Labsは、新たなフレームワーク「Jamba」を発表しました。これは、従来のトランスフォーマーを超えた進化を目指しています。

Jambaは、構造化状態空間モデル（SSM）に基づくMambaモデルとトランスフォーマーアーキテクチャを組み合わせ、最適化された生成AIソリューションを提供します。「Jamba」という名前は、Joint Attention and Mamba Architectureの略で、SSMとトランスフォーマーの強みを活用するために設計されています。このモデルはApache 2.0ライセンスの下でオープンソースとして提供されています。

Jambaは、既存のトランスフォーマーに基づく大規模言語モデル（LLM）を置き換えるものではありませんが、特定のアプリケーションにおいて有用な補完として期待されています。AI21 Labsによると、JambaはHellaSwagのようなベンチマークによって証明されるように、生成推論タスクで従来のトランスフォーマーモデルを上回ることができます。しかし、問題解決能力を評価するMassive Multitask Language Understanding（MMLU）のような重要なベンチマークでは、まだトランスフォーマーモデルには及びません。

AI21 Labsは企業向けの生成AIを専門としており、2023年8月には1億5500万ドルを調達しました。同社の企業向け製品には、トーンやブランドに適したコンテンツ生成を支援するWordtuneがあります。2023年には、生成AIの巨人OpenAIに対抗して企業クライアントを獲得したと報告しています。

従来、AI21 LabsのLLMテクノロジーはトランスフォーマーアーキテクチャを活用しており、その一部としてAI21 Studio自然言語処理（NLP）プラットフォームのJurassic-2 LLMファミリーがAPIを通じて企業向けに提供されています。しかし、JambaはSSMとトランスフォーマーのハイブリッドモデルへのシフトを象徴しています。

トランスフォーマーは生成AIで重要な役割を果たしていますが、ある種の限界も抱えています。特に、コンテキストウィンドウが拡大するにつれ、推論が遅くなるという問題があります。AI21 Labsの研究者によれば、トランスフォーマーの注意機構はシーケンスの長さに応じてスケールし、各トークンが前のすべてのシーケンスに依存するため、スループットが低下します。このため、長いコンテキストを必要とする応用には非効率です。

さらに、トランスフォーマーのスケーリングには大きなメモリリソースが必要ですが、これが長いコンテキストや複数の並行タスクの処理を複雑にします。SSMアプローチは、これらのコンテキストとメモリの問題を解決することを目指しています。

Mamba SSMアーキテクチャは、カーネギーメロン大学とプリンストン大学の研究者によって元々開発されており、少ないメモリで大きなコンテキストウィンドウを管理するための異なる注意メカニズムを活用していますが、トランスフォーマーモデルと同じ出力品質を達成するのは難しいです。Jambaは、SSMのリソースおよびコンテキスト最適化とトランスフォーマーの出力能力を組み合わせたハイブリッドアプローチです。

AI21 Labsによると、Jambaモデルは256Kのコンテキストウィンドウを持ち、Mixtral 8x7Bと比較して長いコンテキストでスループットが3倍になるとされています。また、Jambaは同サイズクラスの中で140Kのコンテキストを単一のGPUに収容できる唯一のモデルです。

Mixtralと同様に、Jambaはエキスパートの混合（MoE）モデルを組み込んでいますが、JambaではハイブリッドSSMトランスフォーマー枠組み内でMoEを活用し、最適化のレベルを向上させています。具体的には、JambaのMoE層は推論中に52億のパラメータのうち120億のみを活性化し、同規模のトランスフォーマーモデルよりも効率的です。

現段階では、Jambaはまだ初期段階にあり、AI21 Labsの企業向け製品には含まれていませんが、同社はAI21プラットフォーム上での教育版を近日中にベータ版として公開する計画です。