Jamba：ハイブリッドTransformer-Mamba言語モデル

Jamba: A Hybrid Transformer-Mamba Language Model

March 28, 2024

著者: Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, Yoav Shoham

cs.AI

要旨

我々は、新たなハイブリッドTransformer-Mamba mixture-of-experts（MoE）アーキテクチャに基づく新しい基盤大規模言語モデル、Jambaを発表します。具体的には、JambaはTransformer層とMamba層を交互に配置し、両モデルファミリーの利点を享受します。これらの層の一部にMoEを追加することで、モデルの容量を増やしつつ、アクティブなパラメータ使用量を管理可能な範囲に保ちます。この柔軟なアーキテクチャにより、リソースと目的に応じた設定が可能です。我々が実装した特定の設定では、単一の80GB GPUに収まる強力なモデルを実現しました。大規模に構築されたJambaは、従来のTransformerと比較して高いスループットと小さなメモリフットプリントを提供し、同時に標準的な言語モデルベンチマークおよび長文脈評価において最先端の性能を発揮します。特に、このモデルは最大256Kトークンの文脈長においても強力な結果を示します。我々は、Transformer層とMamba層をどのように組み合わせるか、エキスパートをどのように混合するかなど、さまざまなアーキテクチャ上の決定を研究し、そのいくつかが大規模モデリングにおいて重要であることを示します。また、Jambaのトレーニングと評価を通じて明らかになったこれらのアーキテクチャの興味深い特性についても説明し、さらなる探求を促すために、さまざまなアブレーション実行のチェックポイントを公開する予定です。我々は、Jambaの実装の重みを寛容なライセンスの下で公開します。

English

We present Jamba, a new base large language model based on a novel hybrid Transformer-Mamba mixture-of-experts (MoE) architecture. Specifically, Jamba interleaves blocks of Transformer and Mamba layers, enjoying the benefits of both model families. MoE is added in some of these layers to increase model capacity while keeping active parameter usage manageable. This flexible architecture allows resource- and objective-specific configurations. In the particular configuration we have implemented, we end up with a powerful model that fits in a single 80GB GPU. Built at large scale, Jamba provides high throughput and small memory footprint compared to vanilla Transformers, and at the same time state-of-the-art performance on standard language model benchmarks and long-context evaluations. Remarkably, the model presents strong results for up to 256K tokens context length. We study various architectural decisions, such as how to combine Transformer and Mamba layers, and how to mix experts, and show that some of them are crucial in large scale modeling. We also describe several interesting properties of these architectures which the training and evaluation of Jamba have revealed, and plan to release checkpoints from various ablation runs, to encourage further exploration of this novel architecture. We make the weights of our implementation of Jamba publicly available under a permissive license.