[论文解读] 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

一句话总结

本文评估在上下文学习（ICL）在状态空间模型（SSMs），特别是 Mamba，与 Transformer 在多样的 ICL 任务中的表现，并引入一种混合的 MambaFormer，将 Mamba 与注意力块结合，以在各任务上实现强大的 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，是否能在从零开始的学习中学到上下文学习，覆盖多样化任务。
• 将 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 Can Mamba and other SSMs learn in-context tasks at scales explored in the study?
• RQ2 How does Mamba compare to Transformer models on standard ICL tasks (linear, sparse, 2NN, decision tree, outliers, sparse parity, CoT I/O, MQAR)?
• RQ3 Do hybrid architectures (Standard Hybrid, MambaFormer) provide a best-of-both-worlds performance across tasks?
• RQ4 What tasks expose limitations of Mamba and Transformers, and can hybrids mitigate them?

主要发现

• Mamba can be trained from scratch to perform a variety of ICL tasks and often matches Transformer performance on standard regression tasks.
• Mamba excels on certain complex ICL tasks (e.g., sparse parity, long-sequence outlier regression) where Transformer struggles, but fails on some retrieval and decision-tree tasks.
• Transformers outperform Mamba on decision-tree and vector-valued MQAR retrieval tasks, while Mamba outperforms Transformers on certain retrieval- and parity-related tasks.
• Hybrid architectures, especially MambaFormer (Mamba blocks followed by attention, no positional encoding), achieve strong performance across the full task suite, combining strengths of both models.
• Interleaving Mamba and attention layers, with Mamba as the initial layer, is crucial for enabling parity learning and improving retrieval capabilities.
• MambaFormer can match or exceed Transformer and Mamba on most tasks, and efficiently learns linear regression even in high-noise settings, illustrating robust ICL capabilities across tasks.