Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learning
This provides a rigorous theoretical foundation for Mamba's performance, addressing a gap in understanding for researchers in sequence modeling and in-context learning, though it is incremental as it builds on prior analyses of Transformers.
The paper tackles the theoretical understanding of Mamba's in-context learning by analyzing its ability to learn low-dimensional nonlinear target functions, proving it can achieve efficient in-context learning with a test-time sample complexity comparable to nonlinear Transformers and improving upon linear Transformers.
Mamba, a recently proposed linear-time sequence model, has attracted significant attention for its computational efficiency and strong empirical performance. However, a rigorous theoretical understanding of its underlying mechanisms remains limited. In this work, we provide a theoretical analysis of Mamba's in-context learning (ICL) capability by focusing on tasks defined by low-dimensional nonlinear target functions. Specifically, we study in-context learning of a single-index model $y \approx g_*(\langle \boldsymbolβ, \boldsymbol{x} \rangle)$, which depends on only a single relevant direction $\boldsymbolβ$, referred to as feature. We prove that Mamba, pretrained by gradient-based methods, can achieve efficient ICL via test-time feature learning, extracting the relevant direction directly from context examples. Consequently, we establish a test-time sample complexity that improves upon linear Transformers -- analyzed to behave like kernel methods -- and is comparable to nonlinear Transformers, which have been shown to surpass the Correlational Statistical Query (CSQ) lower bound and achieve near information-theoretically optimal rate in previous works. Our analysis reveals the crucial role of the nonlinear gating mechanism in Mamba for feature extraction, highlighting it as the fundamental driver behind Mamba's ability to achieve both computational efficiency and high performance.