Jamba: A Hybrid Transformer-Mamba Language Model

要約

Jambaは、TransformerとMambaのハイブリッドMoEアーキテクチャに基づく新しい大規模言語モデルである。具体的には、JambaはTransformer層とMamba層のブロックをインターリーブし、両モデルの利点を享受する。MoEはこれらのレイヤーの一部に追加され、アクティブパラメーターの使用量を管理可能な状態に保ちながら、モデル容量を増加させる。この柔軟なアーキテクチャは、リソースや目的に特化した構成を可能にする。我々が実装した特別な構成では、80GBのGPU1個に収まるパワフルなモデルが完成した。大規模に構築されたJambaは、バニラTransformerと比較して高いスループットと小さなメモリフットプリントを提供し、同時に標準的な言語モデルのベンチマークとロングコンテキストの評価において最先端の性能を発揮する。驚くべきことに、このモデルは最大256Kトークンのコンテキスト長に対して強力な結果を示す。我々は、Transformer層とMamba層をどのように組み合わせるか、エキスパートをどのように混合するかなど、様々なアーキテクチャ上の決定を研究し、そのうちのいくつかが大規模モデリングにおいて重要であることを示す。また、Jambaの学習と評価によって明らかになった、これらのアーキテクチャのいくつかの興味深い性質について説明し、この新しいアーキテクチャのさらなる探求を促すために、様々なアブレーション実行のチェックポイントを公開する予定である。我々は、Jambaの実装の重みを、寛容なライセンスの下で公開する。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL