How Can Mamba Learn In Context with Outliers and Generalize Provably?

The Mamba model has gained significant attention for its computational advantages over Transformer-based models, while achieving comparable performance across a wide range of language tasks. Like Transformers, Mamba exhibits in-context learning (ICL) capabilities, i.e., making predictions for new tasks based on a prompt containing input-label pairs and a query, without requiring fine-tuning. Despite its empirical success, the theoretical understanding of Mamba remains limited, largely due to the nonlinearity introduced by its gating mechanism. To the best of our knowledge, this paper presents the first theoretical analysis of the training dynamics of a one-layer Mamba model, which consists of a linear attention component followed by a nonlinear gating layer, and its ICL generalization on unseen binary classification tasks, even when the prompt includes additive outliers. Our analysis shows that Mamba leverages the linear attention layer to select informative context examples and uses the nonlinear gating layer to suppress the influence of outliers. By establishing and comparing to the analysis of linear Transformers under the same setting, we show that although Mamba may require more training iterations to converge, it maintains accurate predictions even when the proportion of outliers exceeds the threshold that a linear Transformer can tolerate. These theoretical findings are supported by empirical experiments.

A recently proposed model called Mamba has attracted significant attention because it can process long sequences more efficiently than Transformers while achieving similar performance on many language tasks. However, unlike Transformers, Mamba contains a nonlinear gating mechanism that makes it much harder to theoretically understand why the model works and how it learns from data. In this work, we provide the first theoretical study of how a simplified Mamba model learns during training and how it performs in in-context learning tasks. We focus on binary classification problems where the prompt may contain noisy or misleading examples, known as outliers. Our analysis shows that Mamba uses one component of the model to identify useful examples from the prompt, while its nonlinear gating helps reduce the influence of outliers. Our theoretical findings are supported by experiments, providing new insight into why Mamba can achieve strong and robust performance in practice.