Mambaは学習方法を学べるか？インコンテクスト学習タスクに関する比較研究

要旨

状態空間モデル（SSMs）、例えばMamba（Gu & Dao, 2034）は、言語モデリングにおいてTransformerネットワークの代替として提案されており、ゲーティング、畳み込み、および入力依存のトークン選択を組み込むことで、マルチヘッドアテンションの二次コストを軽減しています。SSMsは競争力のある性能を示すものの、現代の言語モデルが持つ顕著な特性である文脈内学習（ICL）能力、つまりパラメータ最適化なしでタスクを実行する能力については、Transformerと比較して未解明の部分が多く残されています。本研究では、Mambaを中心としたSSMsのICL性能を、さまざまなタスクにおいてTransformerモデルと比較評価します。その結果、SSMsは標準的な回帰ICLタスクではTransformerと同等の性能を示す一方、スパースパリティ学習などのタスクではTransformerを上回る性能を発揮することがわかりました。しかし、非標準的な検索機能を必要とするタスクではSSMsは劣ることが判明しました。これらの制約を克服するため、Mambaとアテンションブロックを組み合わせたハイブリッドモデル「\variant」を提案し、個々のモデルが苦手とするタスクにおいても優れた性能を発揮することを示しました。我々の知見は、ハイブリッドアーキテクチャが言語モデルのICL能力を向上させるための有望な道筋を提供することを示唆しています。

English

State-space models (SSMs), such as Mamba Gu & Dao (2034), 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, \variant, 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.

Mambaは学習方法を学べるか？インコンテクスト学習タスクに関する比較研究

Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks

要旨

Support