Research map / Optimization / Implicit bias and generalization

Trained Mamba Emulates Online Gradient Descent in In-Context Linear Regression

NeurIPS 2025 · Optimization › Implicit bias and generalization · approach: In-context learning as gradient descent · 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.

Analyzes gradient-descent training dynamics of a one-layer S6 Mamba on in-context linear regression, proving exponential convergence and an O(1/N) loss bound, and showing Mamba implements a variant of online gradient descent.

Problem

Theoretical understanding of how Mamba learns in-context, using the canonical linear regression ICL task, is limited; training dynamics from random initialization are unknown.

Approach

Trains a one-layer selective SSM (S6) with Gaussian-initialized input/output projections by gradient descent, showing hidden-state recurrence acts as online gradient descent on the context with exponential convergence via vector-coupled dynamics, negative feedback convergence, and fine-grained induction.

Key results

• Trained one-layer Mamba converges exponentially to the ICL solution, with projected hidden state aligning with w as tokens are processed.
• Population loss upper bound 3d(d+1)/(2N), an O(1/N) rate comparable to Transformers, verified on synthetic data with d=4, d_h=80, N from 4 to 80.
• Mamba's hidden-state recurrence is shown to act as a variant of online gradient descent, unlike the single-step GD emulation attributed to Transformers.

Measured effects reported in the paper

Limitations

• Analysis restricted to a one-layer Mamba model.
• Dt (timestep) and A are fixed; only W_B, W_C, b_B, b_C are trained.
• Requires large hidden dimension d_h = Omega~(d^2) and small learning rate, with token length N = Omega(d).

Experimental setup

Datasets: synthetic Gaussian linear regression prompts

Benchmarks: synthetic linear regression

Metrics: cosine similarity between hidden state and w test loss vs. N compared to theoretical bound

Baselines: Transformer ICL analyses S4

Builds on: Mamba S4 Transformer in-context learning gradient descent emulation

Closest papers by meaning

• Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learning (NeurIPS 2026 · similarity 0.91)
• Can Mamba Learn How To Learn? A Comparative Study on In-Context Learning Tasks (ICML 2024 · similarity 0.89)
• Mamba Modulation: On the Length Generalization of Mamba Models (NeurIPS 2025 · similarity 0.89)
• Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression (NeurIPS 2024 · similarity 0.88)
• Decision Mamba: Reinforcement Learning via Hybrid Selective Sequence Modeling (NeurIPS 2024 · similarity 0.88)
• In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization (NeurIPS 2024 · similarity 0.87)