Is Mamba Capable of In-Context Learning?

Riccardo Grazzi, Julien Niklas Siems, Simon Schrodi, Thomas Brox, Frank Hutter

Keywords: In-Context Learning, ICL, Mamba, S4, Transformer, State space models, meta-learning

TL;DR: This study demonstrates that Mamba, a structured state space model, matches transformer performance in in-context learning across various tasks, suggesting its potential as an efficient alternative for processing longer sequences.

Abstract: The surprising generalization capabilities of foundation models have been enabled by in-context learning (ICL), a new variant of meta-learning that denotes the learned ability to solve tasks during a neural network forward pass, exploiting contextual information provided as input to the model. This useful ability emerges as a side product of the foundation model's massive pretraining. While transformer models are currently the state of the art in ICL, this work provides empirical evidence that Mamba, a newly proposed state space model which scales better than transformers w.r.t.\@ the input sequence length, has similar ICL capabilities. We evaluated Mamba on tasks involving simple function approximation as well as more complex natural language processing problems. Our results demonstrate that, across both categories of tasks, Mamba closely matches the performance of transformer models for ICL. Further analysis reveals that, like transformers, Mamba appears to solve ICL problems by incrementally optimizing its internal representations. Overall, our work suggests that Mamba can be an efficient alternative to transformers for ICL tasks involving long input sequences. This is an exciting finding in meta-learning and may also enable generalizations of in-context learned AutoML algorithms (like TabPFN or Optformer) to long input sequences. The anonymous code to reproduce our experiments is available [here](https://anon-github.automl.cc/r/is_mamba_capable_of_in_context_learning-7C49/README.md).

Submission Checklist: Yes

Broader Impact Statement: Yes

Paper Availability And License: Yes

Code Of Conduct: Yes

Optional Meta-Data For Green-AutoML: This blue field is just for structuring purposes and cannot be filled.

CPU Hours: 900

GPU Hours: 900

TPU Hours: 0

Evaluation Metrics: No

Submission Number: 16