[論文レビュー] An Empirical Study of Mamba-based Language Models

Roger Waleffe, Wonmin Byeon+14 more|arXiv (Cornell University)|Jun 12, 2024

Language, Linguistics, Cultural Analysis被引用数 9

ひとこと要約

この研究は 8B パラメータの Mamba、Mamba-2、Transformers（および 8B Mamba-2-Hybrid を含む）を、最大 3.5T トークンまでのデータで訓練し、35 の NLP タスクと長文コンテキストのベンチマークで評価して、スケーリング、コピー、インコンテキスト学習、ハイブリッドアーキテクチャを評価します。

ABSTRACT

Selective state-space models (SSMs) like Mamba overcome some of the shortcomings of Transformers, such as quadratic computational complexity with sequence length and large inference-time memory requirements from the key-value cache. Moreover, recent studies have shown that SSMs can match or exceed the language modeling capabilities of Transformers, making them an attractive alternative. In a controlled setting (e.g., same data), however, studies so far have only presented small scale experiments comparing SSMs to Transformers. To understand the strengths and weaknesses of these architectures at larger scales, we present a direct comparison between 8B-parameter Mamba, Mamba-2, and Transformer models trained on the same datasets of up to 3.5T tokens. We also compare these models to a hybrid architecture consisting of 43% Mamba-2, 7% attention, and 50% MLP layers (Mamba-2-Hybrid). Using a diverse set of tasks, we answer the question of whether Mamba models can match Transformers at larger training budgets. Our results show that while pure SSMs match or exceed Transformers on many tasks, they lag behind Transformers on tasks which require strong copying or in-context learning abilities (e.g., 5-shot MMLU, Phonebook) or long-context reasoning. In contrast, we find that the 8B Mamba-2-Hybrid exceeds the 8B Transformer on all 12 standard tasks we evaluated (+2.65 points on average) and is predicted to be up to 8x faster when generating tokens at inference time. To validate long-context capabilities, we provide additional experiments evaluating variants of the Mamba-2-Hybrid and Transformer extended to support 16K, 32K, and 128K sequences. On an additional 23 long-context tasks, the hybrid model continues to closely match or exceed the Transformer on average. To enable further study, we release the checkpoints as well as the code used to train our models as part of NVIDIA's Megatron-LM project.

研究の動機と目的

• 8B パラメータ、最大 3.5T トークンというスケールで、Mamba ベースの言語モデルが Transformer ベースラインと比較してどのように性能を発揮するかを評価する。
• 標準タスクと長文コンテキストタスクにおける純粋な SSM (Mamba/Mamba-2) の長所と短所を調査する。
• ハイブリッドな Mamba-Transformer アーキテクチャが、効率性の利点を維持しつつ性能格差を縮められるかを検討する。
• 再現性とさらなる研究を可能にするため、公開ベンチマーク、チェックポイント、コードを提供する。

提案手法

• 同じデータ、ハイパーパラメータ、および評価設定で、Mamba、Mamba-2、Mamba-2-Hybrid、および Transformer モデルを訓練して、直接的な apples-to-apples な比較を行う。
• オープンなベンチマークスイーツ（LM Evaluation Harness、LongBench、RULER）を用いて、12 の標準的なショートコンテキストタスクと 23 の長文コンテキストタスクを評価。
• インコンテキスト学習形式を探るため、MMLU を標準、ターゲット内の選択肢テキスト、クロージへ の3形式で分析。
• Mamba-2、自己注意、MLP 層を最適な性能になるように分配するハイブリッドアーキテクチャを設計するアブレーション研究。
• 純粋モデルおよびハイブリッドモデルの双方に対して、最大 128K トークンまでの長文コンテキスト拡張を調査。
• NVIDIA Megatron-LM および Hugging Face を介した訓練コードとモデル重みの公開。

実験結果

リサーチクエスチョン

• RQ1 制御された条件下で、8B パラメータの Mamba および Mamba-2 が、標準的な NLP タスクで Transformer の性能と同等に発揮できるか（最大 3.5T トークンの大規模トークン予算で訓練した場合）？
• RQ2 インコンテキスト学習、コピー、長文推論を要するタスクにおける純粋な SSM モデルの具体的な弱点は何か？
• RQ3 ハイブリッドな Mamba-Transformer アーキテクチャは、推論時の効率性の利点を維持しつつ、純粋な SSM モデルに観察されるギャップを埋められるか？
• RQ4 長文コンテキスト拡張（16K、32K、128K）は、標準および長文ベンチマークにおける純粋 SSM およびハイブリッドモデルの性能にどのように影響するか？
• RQ5 Mamba-2-Hybrid アーキテクチャは、純粋な Transformer に対して実用的な推論速度向上とスケーラビリティの利点を示すか？

主な発見

• 純粋な SSM モデル（Mamba/Mamba-2）は多くの標準タスクで Transformer と同等または上回ることができるが、MMLU（特に短いホライゾンの場合）や Phonebook のようなコピータスクでは遅れる。
• 3.5T トークンで Mamba-2 を訓練すると MMLU の Transformer との差を大幅に埋め、3.5T でショートコンテキストのベンチマークで平均的には Transformer を上回る可能性がある。
• 8B パラメータの Mamba-2-Hybrid (24 Mamba-2, 4 self-attention, 28 MLP) は、評価対象の 12 のショートコンテキストタスクすべてで 8B-parameter Transformer を上回り、平均で +2.65 ポイント、長文コンテキストでは推論速度が最大 8 倍速くなる可能性がある。
• Mamba-2-Hybrid の長文コンテキスト拡張（16K および 32K）は、23 の長文タスクの平均で Transformer ベースラインとほぼ同等、またはそれを上回る。
• Phonebook 型のコピータスクは、純粋な SSM モデルが約 500 トークンを超えるインコンテキストコピーに苦労するのに対し、Transformer は事前学習のコンテキスト長さ 4096 まで対応する。
• 分散自己注意/MLP 層を備えたハイブリッドモデルは高い性能を示し、アブレーションからは自己注意層がおよそ 8%、MLP 層が 30-50% という効果的な構成が示唆され、RoPE の位置埋め込みは大規模ハイブリッドには必須ではなく、長文コンテキストでは省略可能である。
• 推論の高速化: Mamba-2-Hybrid は長文コンテキストで Transformer よりかなり高速にトークンを生成でき、実用的な MFU は強力な Transformer ベースラインに匹敵する。