Research map / Architectures and efficiency / Attention and transformers

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

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Mamba architectures · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows theoretically that a one-layer Mamba pretrained with two-stage gradient-based training can learn single-index models in-context via test-time feature learning enabled by nonlinear gating, matching nonlinear-Transformer sample complexity and beating CSQ bounds.

Problem

Whether Mamba, a linear-time state space model, provably achieves strong in-context adaptability (feature learning at test time) on low-dimensional nonlinear targets such as Gaussian single-index models, with sample complexity comparable to Transformers.

Approach

A one-layer Mamba with nonlinear input-dependent gating, after one gradient step on the Mamba layer and MLP fitting, extracts the relevant direction from context examples; gating reduces the effective information exponent to the generative exponent, giving test-time sample complexity r^{O(ge)} below CSQ bounds.

Key results

• Theorem 1 gives pretraining context length and task requirements scaling as d^{Theta(ge(g*))} and test prompt length Omega~(r^{3ge(g*)}) achieving test error tau + o(1).
• Table 1 claims Mamba's test sample complexity r^{Theta(ge)} versus r^{Theta(deg)} for linear Transformer and r^{Theta(ge)} for softmax Transformer, with Mamba pretraining d^{Theta(ge)} versus d^{Theta(ie)} for prior Transformers.
• Experiments with g*(z)=He_3(z)/sqrt(6), r=8, d=32 show Mamba and GPT-2 Transformer both outperform kernel ridge regression; with pretraining context length 16, Mamba stays comparable while the Transformer degrades.

Surprising findings

• With pretraining context length 16, the Transformer's performance deteriorates significantly while Mamba's remains comparable to length 64.

Limitations

• Analysis uses simplified Mamba (fixed A=-I, fixed w and b, merged W_B^T W_C) and a single-layer embedding with degree-2 Hermite features.
• Sample complexity bounds are sufficient conditions; whether Mamba is more pretraining-efficient than Transformers remains open.
• Analysis assumes negligible forgetting in the gating mechanism.

Experimental setup

Datasets: synthetic Gaussian single-index model prompts

Benchmarks: synthetic Gaussian single-index ICL tasks

Metrics: squared prediction error vs number of context examples

Baselines: linear Transformer softmax Transformer (GPT-2) kernel ridge regression

Builds on: Mamba linear Transformer in-context single-index learning generative exponent information exponent

Closest papers by meaning

• Trained Mamba Emulates Online Gradient Descent in In-Context Linear Regression (NeurIPS 2025 · similarity 0.91)
• Can Mamba Learn How To Learn? A Comparative Study on In-Context Learning Tasks (ICML 2024 · similarity 0.89)
• Achilles' Heel of Mamba: Essential difficulties of the Mamba architecture demonstrated by synthetic data (NeurIPS 2025 · similarity 0.87)
• MambaLRP: Explaining Selective State Space Sequence Models (NeurIPS 2024 · similarity 0.87)
• Nonlinear transformers can perform inference-time feature learning (ICML 2025 · similarity 0.86)
• Mamba Modulation: On the Length Generalization of Mamba Models (NeurIPS 2025 · similarity 0.86)