[論文レビュー] Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks

Jong-Ho Park, Jaeseung Park+6 more|arXiv (Cornell University)|Feb 6, 2024

Poverty, Education, and Child Welfare被引用数 10

ひとこと要約

本論文は、状態空間モデル（SSMs）、特に Mamba における in-context learning (ICL) を、Transformer と比較し、多様な ICL タスクにわたって評価し、Mamba を注意機構ブロックと組み合わせたハイブリッドな MambaFormer を導入して、タスク全体での強力な ICL パフォーマンスを実現する。

ABSTRACT

State-space models (SSMs), such as Mamba (Gu & Dao, 2023), have been proposed as alternatives to Transformer networks in language modeling, by incorporating gating, convolutions, and input-dependent token selection to mitigate the quadratic cost of multi-head attention. Although SSMs exhibit competitive performance, their in-context learning (ICL) capabilities, a remarkable emergent property of modern language models that enables task execution without parameter optimization, remain underexplored compared to Transformers. In this study, we evaluate the ICL performance of SSMs, focusing on Mamba, against Transformer models across various tasks. Our results show that SSMs perform comparably to Transformers in standard regression ICL tasks, while outperforming them in tasks like sparse parity learning. However, SSMs fall short in tasks involving non-standard retrieval functionality. To address these limitations, we introduce a hybrid model, MambaFormer, that combines Mamba with attention blocks, surpassing individual models in tasks where they struggle independently. Our findings suggest that hybrid architectures offer promising avenues for enhancing ICL in language models.

研究の動機と目的

• 状態空間モデル（SSMs）、特に Mamba が様々なタスクに渡って最初から ICL を学習できるかを評価する。
• Mamba と Transformer モデルを、回帰、パリティ、リトリーバル、チェイン・オブ・思考を含む標準的な ICL タスクで比較する。
• ハイブリッドアーキテクチャが SSM と注意機構の双方の強みを活用して ICL を改善できるかを調査する。
• Mamba のリトリーバルと決定木タスクにおける限界を探り、それらのギャップを緩和するアーキテクチャを特定する。

提案手法

• 回帰、外れ値、離散関数、チェイン・オブ・思考、リトリーバルを横断する ICL タスクのスイートに対して、ゼロからの訓練を行う。
• Mamba、S4、S4-Mamba、Transformer（位置埋め込みあり・なし）、および二つのハイブリッド：Standard Hybrid と MambaFormer を評価する。
• Mamba ブロックを注意ブロックと交互に配置してハイブリッド・アーキテクチャを形成し、タスクごとのパフォーマンスを比較する。
• 固定の訓練目的を使用する：タスクごとにランダムに生成されたプロンプトと入力に対するタスク特異的損失を最小化する。
• 1,280 件のテストプロンプトごとに、経験的損失と収束挙動の観点でパフォーマンスを報告する。

Figure 1 : MambaFormer is a hybrid architecture that replaces MLP blocks within the transformer with Mamba blocks. Importantly, the architecture also starts with a Mamba block and does not use positional encoding. In our ICL evaluations, we find that MambaFormer consistently achieves a best-of-both-

実験結果

リサーチクエスチョン

• RQ1 本研究で検討されたスケールで、Mamba および他の SSM が in-context タスクを学習できるか？
• RQ2 標準的な ICL タスク（線形、疎、2NN、決定木、外れ値、疎パリティ、CoT I/O、MQAR）において、Mamba は Transformer モデルと比べてどうか？
• RQ3 ハイブリッドアーキテクチャ（Standard Hybrid、MambaFormer）は、タスク全体でベスト・オブ・ボース・ワールドの性能を提供するか？
• RQ4 Mamba および Transformer の制約を露呈するタスクは何で、ハイブリッドがそれを緩和できるか？

主な発見

• Mamba はゼロから訓練してさまざまな ICL タスクを実行でき、標準的な回帰タスクで Transformer の性能としばしば同等である。
• Mamba は、Transformer が苦戦する一部の複雑な ICL タスク（例：疎パリティ、長いシーケンス外れ値回帰）で優れている一方、いくつかのリトリーバルおよび決定木タスクでは失敗する。
• Transformer は決定木およびベクトル値 MQAR リトリーバルタスクで Mamba を上回るが、Mamba は特定のリトリーバルおよびパリティ関連タスクで Transformer を上回る。
• ハイブリッドアーキテクチャ、特に MambaFormer（Mamba ブロックの後に注意、位置埋め込みなし）は、タスク全体の性能において強力な結果を示し、両モデルの強みを組み合わせる。
• Mamba と注意層を交互に配置し、Mamba を初期層とすることは、パリティ学習を可能にし、リトリーバル機能を改善するために重要である。
• MambaFormer はほとんどのタスクで Transformer および Mamba を上回るか、同等であり、ノイズの多い設定でも線形回帰を効率的に学習することができ、タスク全体で堅牢な ICL 能力を示す。