[論文レビュー] Jamba: A Hybrid Transformer-Mamba Language Model

Opher Lieber, Barak Lenz+20 more|arXiv (Cornell University)|Mar 28, 2024

Language, Discourse, Communication Strategies被引用数 40

ひとこと要約

JambaはTransformerとMambaの混成Mixture-of-Expertsアーキテクチャを導入し、TransformerとMambaレイヤーをMoEと組み合わせて交互に配置することで、高い性能と長文脈能力を実現しつつ、単一の80GB GPUに収めます。

ABSTRACT

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.

研究の動機と目的

• Attentionベースのモデルと状態空間モデルの強みを、TransformerとMambaレイヤを相互に組み合わせることで結びつけられるかを調査する。
• ハイブリッドアーキテクチャへのMoE統合が容量、スループット、メモリに与える影響を評価する。
• 標準ベンチマークと長文脈タスク（256Kトークン文脈を含む）でのパフォーマンスを評価する。
• 7B-12Bパラメータ規模のハイブリッドモデルをコモディティハードウェアで訓練・実用化する上での安定性と実用性を示す。

提案手法

• Jambaブロックは、TransformerまたはMambaレイヤに続くMLPまたはMoEモジュールを組み合わせて定義する。
• a:mのAttention-to-Mamba比率でブロックを交互に配置し、eレイヤごとにMoEを適用して総エキスパート数n、トークンごとのtop-Kルーティングを適用する。
• MambaレイヤでRMSNormを使用し、位置エンベディングを明示的には省略し、ハイブリッド構造による暗黙の位置情報に依存する。
• スループットとメモリ効率を最適化するために64K語彙とBPEトークナイザーで大規模データ上で訓練し、80GB GPU環境での実用性を追求する。
• 学術ベンチマーク、長文脈QAデータセット、異なるハードウェア規模でのスループット測定を評価する。

実験結果

リサーチクエスチョン

• RQ1 ハイブリッドAttention-Mambaアーキテクチャは同程度サイズの純粋なTransformerモデルと同等以上の性能を標準ベンチマークで発揮できるか。
• RQ2 MoEをハイブリッドアーキテクチャに組み込むことで、計算コストを大幅に増やさずに容量を改善できるか。
• RQ3 AttentionとMambaの層の比率が、メモリ使用量、スループット、および長文脈性能にどのような影響を与えるか。
• RQ4 Jambaは非常に長い文脈（最大256Kトークン）を、KVキャッシュの要件を合理的に保ちながら処理できるか。
• RQ5 大規模ハイブリッドモデルの訓練安定性に関する実践的考慮事項は何か。

主な発見

Available params	Active params	KV cache (256K context, 16bit)
LLAMA-2	6.7B	6.7B	128GB
Mistral	7.2B	7.2B	32GB
Mixtral	46.7B	12.9B	32GB
Jamba	52B	12B	4GB

• JambaはMixtralやLlama-2 70Bといった同程度の公開モデルと比較して、標準ベンチマークで競合的または優れた精度を達成する。
• ハイブリッドAttention-Mambaアーキテクチャは256K文脈でKVキャッシュ要件を4GBに削減し、単一の80GB GPUで長文処理を可能にする。
• MoE バリアントは大規模スケール（7Bパラメータ、50Bトークンで訓練）で非MoEハイブリッドより性能を向上させる。
• Attention-Mambaのハイブリッドは、純粋なMambaと比較していくつかのタスクで上回り、Transformer のような文脈内学習をサポートする、補完的な長所を示唆する。
• 明示的な位置情報はJambaには必要なく、Mamba先行構造が暗黙の位置情報を提供する。