State of the art foundation models such as GPT-4 perform surprisingly well at in-context learning (ICL), a variant of meta-learning concerning 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 enable generalizations of in-context learned AutoML algorithms (like TabPFN or Optformer) to long input sequences.
The paper explores the in-context learning ability of Mamba, a state space model, as an effective alternative to transformer-based models.
Mamba demonstrates comparable or superior in-context learning performance to transformer and other models, such as S4 and RWKV, on various tasks.
A probing strategy unravels how Mamba incrementally refines its internal state for task-solving, hinting at similarities to transformers.
Mamba shows potential in natural language processing tasks, with scaling advantages suggesting its suitability for high-complexity language models.
In the field of AI, in-context learning (ICL) stands as a transformative facility exhibited by large neural networks, especially those with transformer architectures, negating the need for explicit retraining or fine-tuning to accommodate new tasks. Recently, there's been a burgeoning interest in models like Mamba—a selective structured state space model—primarily due to its potential advantages in handling longer sequences over transformers. The study under discussion contributes significantly to the current understanding of Mamba's ICL abilities, affirmation of which could present Mamba as a powerful and efficient alternative to transformers for ICL tasks.
One central finding is that Mamba matches or exceeds the performance of (self-supervised) pre-trained transformer models in ICL tasks, overcoming the limitations posed by transformers in processing longer inputs. This result asserts the robustness of the Mamba architecture, as it performs comparably with transformers in tasks ranging from regression to complex language processing. The analysis extends to show Mamba's superiority to its predecessor S4 and other baseline models such as RWKV for these tasks. Importantly, results indicate that Mamba maintains its ICL capabilities across both in-distribution and out-of-distribution examples.
Delving deeper into Mamba's ICL methodology, the study employs a probing strategy to elucidate the model's iterative optimization process for task solving. By examining intermediate representations layer by layer, the analysis suggests that Mamba refines its internal state incrementally to solve ICL tasks. Here, it exhibits an approach somewhat akin to transformers. Yet, some ambiguity remains in different cases such as ReLU networks and decision trees, pointing to areas for future scrutiny.
Further empirical evidence reinforces the efficacy of Mamba when selectively fine-tuned and pre-trained on large datasets for NLP tasks, showing that it compares favorably against contemporary models like RWKV, LLama, Pythia, and even GPT-J at similar or fewer parameters. In this domain, Mamba's scalability with in-context examples and parameter count is particularly noteworthy. The study indicates that as the model size increases, Mamba's ICL accuracy improves substantively, demonstrating its potential for high-complexity NLP.
The paper crystallizes the contention that Mamba is not only capable of ICL but does so with a proficiency that puts it on an even keel with transformer models. Crucially, this capability extends to longer sequence inputs, situating Mamba as a compelling alternative to the transformer paradigm. In essence, for ICL tasks—whether they are function approximations or dense, intricate language modeling—Mamba's architecture represents a promising innovation. This work lays a strong foundation for deepening our understanding of state-of-the-art machine learning architectures and their inherent learning strategies.
GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https://github.com/kingoflolz/mesh-transformer-jax, May 2021.