Is Mamba Capable of In-Context Learning?

Riccardo Grazzi, Julien Niklas Siems, Simon Schrodi, Thomas Brox, Frank Hutter
Proceedings of the Third International Conference on Automated Machine Learning, PMLR 256:1/1-26, 2024.

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 at \url{https://anon-github.automl.cc/r/is_mamba_capable_of_in_context_learning-7C49/README.md}.

Cite this Paper

BibTeX
@InProceedings{pmlr-v256-grazzi24a, title = {Is Mamba Capable of In-Context Learning?}, author = {Grazzi, Riccardo and Siems, Julien Niklas and Schrodi, Simon and Brox, Thomas and Hutter, Frank}, booktitle = {Proceedings of the Third International Conference on Automated Machine Learning}, pages = {1/1--26}, year = {2024}, editor = {Eggensperger, Katharina and Garnett, Roman and Vanschoren, Joaquin and Lindauer, Marius and Gardner, Jacob R.}, volume = {256}, series = {Proceedings of Machine Learning Research}, month = {09--12 Sep}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v256/main/assets/grazzi24a/grazzi24a.pdf}, url = {https://proceedings.mlr.press/v256/grazzi24a.html}, 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 at \url{https://anon-github.automl.cc/r/is_mamba_capable_of_in_context_learning-7C49/README.md}.} }
Endnote
%0 Conference Paper %T Is Mamba Capable of In-Context Learning? %A Riccardo Grazzi %A Julien Niklas Siems %A Simon Schrodi %A Thomas Brox %A Frank Hutter %B Proceedings of the Third International Conference on Automated Machine Learning %C Proceedings of Machine Learning Research %D 2024 %E Katharina Eggensperger %E Roman Garnett %E Joaquin Vanschoren %E Marius Lindauer %E Jacob R. Gardner %F pmlr-v256-grazzi24a %I PMLR %P 1/1--26 %U https://proceedings.mlr.press/v256/grazzi24a.html %V 256 %X 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 at \url{https://anon-github.automl.cc/r/is_mamba_capable_of_in_context_learning-7C49/README.md}.
APA
Grazzi, R., Siems, J.N., Schrodi, S., Brox, T. & Hutter, F.. (2024). Is Mamba Capable of In-Context Learning?. Proceedings of the Third International Conference on Automated Machine Learning, in Proceedings of Machine Learning Research 256:1/1-26 Available from https://proceedings.mlr.press/v256/grazzi24a.html.

Related Material