[論文レビュー] Jamba-1.5: Hybrid Transformer-Mamba Models at Scale

Jamba Team, Barak Lenz+59 more|arXiv (Cornell University)|Aug 22, 2024

Fluid Dynamics and Vibration Analysis被引用数 4

ひとこと要約

Jamba-1.5は、ハイブリッドなTransformer-Mamba混合専門家（MoE）アーキテクチャに基づき、2つの指令微調整済みの大規模言語モデル—Jamba-1.5-Mini（12Bのアクティブパラメータ）およびJamba-1.5-Large（94Bのアクティブパラメータ）—を導入した。256Kトークンのコンテキスト長で、最先端の長文コンテキストベンチマーク性能を達成し、9GBのKVキャッシュという低メモリ使用量と高いスループットを実現。新規の量子化技術「ExpertsInt8」を用いて効率的な推論が最適化され、品質損ないなしに8×80GB GPU上でのデプロイが可能となった。

ABSTRACT

We present Jamba-1.5, new instruction-tuned large language models based on our Jamba architecture. Jamba is a hybrid Transformer-Mamba mixture of experts architecture, providing high throughput and low memory usage across context lengths, while retaining the same or better quality as Transformer models. We release two model sizes: Jamba-1.5-Large, with 94B active parameters, and Jamba-1.5-Mini, with 12B active parameters. Both models are fine-tuned for a variety of conversational and instruction-following capabilties, and have an effective context length of 256K tokens, the largest amongst open-weight models. To support cost-effective inference, we introduce ExpertsInt8, a novel quantization technique that allows fitting Jamba-1.5-Large on a machine with 8 80GB GPUs when processing 256K-token contexts without loss of quality. When evaluated on a battery of academic and chatbot benchmarks, Jamba-1.5 models achieve excellent results while providing high throughput and outperforming other open-weight models on long-context benchmarks. The model weights for both sizes are publicly available under the Jamba Open Model License and we release ExpertsInt8 as open source.

研究の動機と目的

• 長文入力を処理できる大規模で効率的かつ高性能な指令微調整済み言語モデルの開発。
• 品質を損なわず、大規模モデルにおけるメモリ使用量の低減と推論効率の向上。
• 新規の量子化およびアーキテクチャ的革新を通じて、Jamba-1.5-Largeのような大規模モデルを標準的なハードウェア（8×80GB GPU）にデプロイ可能にする。
• ハイブリッドなTransformer-Mamba-MoEアーキテクチャが、純粋なTransformerやMambaモデルを上回ることを、特に長文コンテキスト環境で実証すること。
• OECD AI原則に準拠し、完全な透明性を確保したJamba Open Model Licenseの下で、オープンウェイトモデルをリリースすること。

提案手法

• Jamba-1.5アーキテクチャは、Transformer層、Mamba状態空間層、および16の専門家と2つのゲーティングを有する混合専門家（MoE）モジュールを統合している。
• 各ブロックでは、アテンション層とMamba層の比率を1:7に設定し、2層ごとにMoEを適用、隠れ層次元は8192である。
• ExpertsInt8量子化が導入され、MoE層の85％の重みを対象としており、品質劣化なしに8ビット量子化を実現している。
• モデルは1000億トークンのデータで学習され、多様なベンチマークを対象に指令従いおよび会話能力のファインチューニングが実施された。
• Mambaの効率的なシーケンスモデリングとMoEのスパース活性化の両方を活用することで、高いスループットと低KVキャッシュ使用量（256Kトークン時で9GB）を達成した。
• 評価は学術ベンチマーク（例：MMLU、GSM8K）、チャットボットベンチマーク、およびRULERのような長文コンテキストタスクで実施され、LLaMA-3、Mistral-Large-2、Mixtralと比較された。

(a) 350M.

実験結果

リサーチクエスチョン

• RQ1 ハイブリッドなTransformer-Mamba-MoEアーキテクチャは、94Bのアクティブパラメータにスケールアップしても、品質を維持または向上させることができるか？
• RQ2 スケールアップした状態で、MambaとTransformer層の組み合わせが、純粋なMamba-2やTransformerオンリーのモデルを上回る性能を示すか？
• RQ3 ExpertsInt8量子化により、MoE層の完全な8ビット量子化が品質劣化なしに可能となり、8×80GB GPU上でのデプロイが可能になるか？
• RQ4 他のオープンウェイトモデルと比較して、256Kトークンのコンテキスト長での長文コンテキストタスクにおけるモデルの性能はいかがなものか？
• RQ5 Jamba-1.5モデルは、特にリソースが限られた言語において、多言語タスクにどの程度一般化できるか？

主な発見

• Jamba-1.5-Largeは94Bのアクティブパラメータを実現し、ExpertsInt8量子化を用いることで、256Kコンテキスト長で8×80GB GPUの1台のマシンに収容可能であり、品質損ないなしに実現した。
• 256Kトークン時におけるKVキャッシュ使用量は9GBにまで低減され、LLaMA-3.1-70BやMistral-Large-2と比較して10倍の削減を達成した。
• RULERベンチマークでは、Jamba-1.5モデルが256Kトークンの有効コンテキスト長を持つ唯一のオープンウェイトモデルであった。
• Jamba-1.5-Miniは多言語MMLUで64.3％を達成し、LLaMA-3.1-8BやGemma-9Bを上回った。Jamba-1.5-Largeは同ベンチマークで73.94％を記録した。
• 長文コンテキスト評価において、Jamba-1.5-LargeはLLaMA-3.1-70BやMistral-Large-2を上回り、優れたスループットとレイテンシを示した。
• ExpertsInt8量子化がMoE層に対して有効であることが実証され、モデル品質を保持したまま、一般消費者向けハードウェアへの効率的デプロイを可能にした。

(b) 1.3B.