Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learning

TL;DR

This paper analyzes Mamba's in-context learning for Gaussian single-index targets, showing that a gradient-pretrained Mamba can perform test-time feature learning via a nonlinear gating mechanism. The key result is that test-time sample complexity scales with the intrinsic dimension $r$ and the generative exponent $ ext{ge}(g_*)$, rather than the ambient dimension $d$, and can surpass CSQ lower bounds through test-time feature extraction. The authors introduce a two-stage pretraining procedure and provide rigorous arguments and Hermite-based analysis to show gating enables extraction of the low-dimensional feature and that the MLP can learn the link function $g_*$, with empirical results aligning with theory and highlighting competitive performance against Transformers. The work broadens the theoretical understanding of ICL beyond attention-based models and demonstrates that Mamba can achieve efficient adaptation with favorable sample complexity while maintaining computational efficiency.

Abstract

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.

Figures (1)

• Figure 1: Prediction error for in-context learning with Transformer and Mamba models, and kernel regression.

Theorems & Definitions (40)

• Definition 2.1: Gaussian Single-Index Model
• Definition 2.2: ICL Data Distribution
• Remark 2.3
• Remark 2.4
• Definition 3.1
• Lemma 3.2: Proposition 6 in lee2024neural
• Theorem 3.3
• Proposition 4.1: Informal
• Lemma A.1
• proof : Proof of Lemma \ref{['lemma:mamba_output']}