Jamba (language model) — ジャンバ

From Systems analysis wiki
Jump to navigation Jump to search

Jamba(ジャンバ)は、イスラエルのAI研究企業であるAI21 Labsによって開発された大規模言語モデル(LLM)のファミリーです。Jambaは、AI開発における2つの主要なアプローチ、すなわちトランスフォーマーと状態空間モデル(State Space Models, SSM)、特にMambaアーキテクチャ[1]の重要な要素を統合した、史上初のハイブリッドアーキテクチャを特徴としています。

Jambaの主な目的は、現代のLLMが直面する根本的なトレードオフ、すなわち高品質と高性能(トランスフォーマーの特性)と、効率性と超長コンテキスト処理能力(SSMの特性)との間の問題を解決することです。これらのアプローチを組み合わせ、さらにMixture-of-Experts (MoE)によるスパース性を加えることで、Jambaは強力かつ効率的であり、単一のプロンプトで膨大な量のテキストを処理できるモデルを提供します。

Jambaのアーキテクチャ詳細

Jambaは、TransformerレイヤーとMambaレイヤーを単純に交互に配置するのではなく、綿密に設計されたブロック構造を採用しており、各ブロックは8つのレイヤーで構成されています。

Jambaの1ブロックの構造:

  • 1つのTransformerレイヤー: このレイヤーは「深い」理解と複雑な推論を担当します。このレイヤーにはMixture-of-Experts (MoE)アーキテクチャが組み込まれています。
  • 7つのMambaレイヤー: これらのレイヤーはTransformerレイヤーの後に続き、シーケンスの効率的な処理と、長いコンテキストを通じて情報を「伝達する」役割を担います[2]

このような非対称な構造により、モデルは計算リソースを効率的に管理できます。つまり、重いが強力なTransformerの処理は実行頻度が低く、軽量で高速なMambaの処理は実行頻度が高くなっています。

Mixture-of-Experts (MoE)の統合

Jambaでは、効率をさらに向上させるためにMoEアーキテクチャが採用されています。

  • MoEはTransformerレイヤー内の全結合ブロック(FFN)にのみ適用されます[3]。Mambaレイヤーは密な(dense)ままです。
  • 最初のJambaモデルでは16のエキスパートが使用されました。
  • 各トークンに対して、ルーターネットワークが最適な2つのエキスパートを選択します(Top-2ゲーティング)。

これにより、モデルの総パラメータ数は520億と大きいものの、Transformerレイヤーでの各トークン処理ステップでアクティブになるのは16のエキスパートのうち2つのみとなり、計算が非常に高速になります。

Jambaモデルの進化

Jamba-v0.1 (2024年3月)

このファミリーで最初に発表されたモデルは、以下の仕様を持っています:

Jamba-v0.1の技術仕様
仕様
総パラメータ数 520億
アクティブパラメータ数 約120億
エキスパート数 (MoE) 16(うち2つがアクティブ)
コンテキストウィンドウ 256,000トークン
ライセンス Apache 2.0[4]

このハイブリッドアーキテクチャにより、Jambaは256,000トークン(約400ページの小説に相当)のコンテキストを処理でき、80GBのメモリを搭載した単一のコンシューマー向けGPUにデプロイすることが可能です[5]

Jamba-1.5 (2024年)

2024年、AI21 Labsは更新されたJamba 1.5モデルファミリーを発表しました。これには2つのバージョンが含まれます:Jamba 1.5 Mini(総パラメータ数520億のうちアクティブパラメータ数120億)とJamba 1.5 Large(総パラメータ数3980億のうちアクティブパラメータ数940億)です[6]。これらのモデルは、パフォーマンスにおいて大幅な改善を示しています:

  • 競合モデルと比較して、長文コンテキストでの推論が最大2.5倍高速。
  • 英語、スペイン語、フランス語、アラビア語を含む9つの言語をサポート[7]

主な利点とパフォーマンス

  • 巨大なコンテキストウィンドウ: 256,000トークンは、リリース時点で利用可能なすべてのモデル(プロプライエタリなモデルを含む)の中で最大級のウィンドウサイズです。これにより、Jambaは法務契約書、科学論文、コードベース全体、あるいは長い対話など、大規模なドキュメントの分析を必要とするタスクに最適です。
  • 高いパフォーマンスと効率性: テストにおいて、JambaはLlamaやMixtralといった同規模の主要なオープンモデルと同等またはそれ以上の性能を発揮し、特に長文コンテキストでは3倍高いスループットを示します[8]
  • オープン性とアクセス性: Jambaは寛容なApache 2.0ライセンスの下で配布されており、商業目的および研究目的での自由な利用が可能です。モデルの重みはHugging Faceプラットフォームで公開されています。

ベンチマーク結果

Jamba 1.5は、様々なベンチマークで競争力のある結果を示しています[9]

  • Jamba 1.5 MiniはArena Hardで46.1点を獲得し、同カテゴリの公開モデルの中でトップとなりました[10]
  • Jamba 1.5 LargeはArena Hardで65.4点を獲得し、Llama 3.1 70Bおよび405Bを上回りました。

応用と利用可能性

Jambaはビジネス用途に最適化されており、関数呼び出し、JSON形式での構造化出力、ドキュメント処理といった機能をサポートしています。モデルは以下を含む多くのプラットフォームで利用可能です:

  • Hugging Face
  • Google Cloud Vertex AI
  • Microsoft Azure
  • NVIDIA API catalog
  • Amazon Bedrock[9]
  • AI21 Studio

コスト効率の高い推論をサポートするため、AI21 Labsは新しい量子化技術であるExpertsInt8を発表しました。これにより、Jamba 1.5 Largeを80GBのGPU 8基を搭載したマシンにデプロイし、256Kトークンのコンテキストを品質を損なうことなく処理することが可能になります[11]

参考文献

  • Lieber, O.; et al. (2024). Jamba: A Hybrid Transformer‑Mamba Language Model. arXiv:2403.19887.
  • Lieber, O.; et al. (2024). Jamba‑1.5 Models and ExpertsInt8 Quantization. OpenReview JFPaD7lpBD.
  • Gu, A.; Dao, T. (2023). Mamba: Linear‑Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.
  • Gu, A.; et al. (2021). S4: Efficiently Modeling Long Sequences with Structured State Spaces. arXiv:2111.00396.
  • Fedus, W.; et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961.
  • Yun, L.; et al. (2024). Toward Inference‑Optimal Mixture‑of‑Expert Large Language Models. arXiv:2404.02852.
  • Liu, J.; et al. (2024). A Survey on Mixture of Experts in Large Language Models. arXiv:2407.06204.
  • Gupta, V.; et al. (2024). Lynx: Enabling Efficient MoE Inference through Dynamic Batch‑Aware Expert Selection. arXiv:2411.08982.
  • Liu, J.; et al. (2024). A Survey on Inference Optimization Techniques for Mixture of Experts Models. arXiv:2412.14219.
  • Hsieh, C.‑P.; et al. (2024). RULER: What's the Real Context Size of Your Long‑Context Language Models?. arXiv:2404.06654.

脚注

  1. «Announcing Jamba: AI21's Groundbreaking SSM-Transformer Model». AI21 Labs Blog. [1]
  2. Lieber, O., et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. arXiv:2403.19887.
  3. «Jamba Documentation». Hugging Face Transformers. [2]
  4. «ai21labs/Jamba-v0.1». Hugging Face. [3]
  5. «AI21 Labs' Jamba: A New Hybrid LLM Architecture». Gradient Flow. [4]
  6. «Announcing the Jamba-1.5 model family». AI21 Labs Blog. [5]
  7. «ai21labs/AI21-Jamba-Large-1.5». Hugging Face. [6]
  8. «AI21 Labs разбивает новые барьеры с помощью Jamba». ITinAI. [7]
  9. 9.0 9.1 «Вышла Jamba 1.5: гибридная модель от AI21 Labs». Дзен. [8]
  10. «Jamba-1.5 family of models by AI21 Labs is now available in Amazon Bedrock». AWS What's New. [9]
  11. «ExpertsInt8: A new paradigm for efficient inference of MoE-based LLMs». OpenReview. [10]