[Paper Review] 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 Welfare10 citations

TL;DR

This paper evaluates in-context learning (ICL) in state-space models (SSMs), especially Mamba, against Transformers across diverse ICL tasks, and introduces a hybrid MambaFormer that combines Mamba with attention blocks to achieve strong ICL performance across tasks.

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.

Motivation & Objective

• Assess whether state-space models (SSMs), particularly Mamba, can learn in-context learning from scratch across varied tasks.
• Compare Mamba to Transformer models on standard ICL tasks including regression, parity, retrieval, and chain-of-thought.
• Investigate whether hybrid architectures can leverage strengths of both SSMs and attention mechanisms to improve ICL.
• Explore the limitations of Mamba in retrieval and decision-tree tasks and identify architectures that mitigate these gaps.

Proposed method

• Train from scratch on a suite of ICL tasks spanning regression, outliers, discrete functions, chain-of-thought, and retrieval.
• Evaluate Mamba, S4, S4-Mamba, Transformer (with and without positional encoding), and two hybrids: Standard Hybrid and MambaFormer.
• Interleave Mamba blocks with attention blocks to form hybrid architectures and compare performance across tasks.
• Use a fixed training objective: minimize task-specific loss over randomly generated prompts and inputs per task.
• Report performance in terms of empirical loss and convergence behavior across 1,280 test prompts per task.

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-

Experimental results

Research questions

• 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?

Key findings

• 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.