A Theoretical Analysis of Mamba's Training Dynamics: Filtering Relevant Features for Generalization in State Space Models

AI summary

Overview

Research area: Machine learning theory — specifically the training dynamics, optimization, and generalization of selective state space models (SSMs), with Mamba as the focal architecture.

Technical level: Advanced. The paper is a formal learning-theory analysis built on non-asymptotic bounds, gradient decompositions, and structured data assumptions; readers need comfort with empirical risk minimization, convergence rates, and sample complexity.

Scope: A first-step theoretical study of a simplified single-layer, single-head Mamba block with input-dependent gating followed by a two-layer MLP, trained with gradient descent on two structured synthetic data regimes (majority-voting and locality-structured).

What This Paper Is About

Mamba and other selective SSMs match or beat Transformers on long-sequence tasks with linear rather than quadratic complexity, but almost all theory about them concerns expressive power and existence of good parameters, not whether training actually finds them. This paper asks when and why gradient descent on a gated recurrent model provably generalizes, and how Mamba's selection mechanism is realized during training. The authors prove that the gating vector learns to amplify class-relevant features and suppress class-irrelevant ones, and derive sample-complexity and iteration bounds for two canonical data structures.

Key Contributions

1. A general framework for gated architectures on structured data. The authors build an analysis of a network combining a one-layer Mamba block (Equation 5) with a two-layer MLP (Equation 6) trained by gradient descent, capturing the interaction between input-dependent gating and data structure. They state this is the first theoretical analysis of Mamba's training dynamics with input-dependent gating together with generalization guarantees.

2. A formal characterization of the gating mechanism as a feature selector. They prove that the gating vector w_Δ aligns with class-relevant directions and stays negligible or negative along class-irrelevant directions, giving a Mamba analogue to the feature-selection role usually attributed to attention.

3. Sample complexity and iteration bounds for two data regimes. For majority-voting data, both quantities scale with the gap between class-relevant and confusion features, (α_r − α_c)^(−2). For locality-structured data, they depend on the concentration of class-relevant tokens, scaling as [(1/2)^{ΔL_{o_+}^+} − (1/2)^{ΔL_{o_+}^−}]^(−1) for iterations and the same quantity squared for samples.

4. A new locality-structured data model plus synthetic validation. Majority-voting data was studied before in the Transformer setting; the authors introduce locality-structured data to highlight where selective recurrence has an advantage, and support both regimes with synthetic numerical experiments.

Main Findings

• Gating aligns with relevant features under majority voting. Lemma 4.1 shows ⟨wΔ^(T), o+⟩ ≥ (ηT / (8L²)) Θ((α_r L − α_c L)²) and the same bound holds for o_−, while ⟨w_Δ^(T), o_j⟩ ≤ Õ(1/poly(d)) for all j ≥ 3. The paper describes this alignment with irrelevant features as remaining strictly negative and negligible.

• Gating suppresses irrelevant features under locality structure. Lemma 4.2 shows ⟨wΔ^(T), o+⟩ ≥ −Õ(1/poly(d)) and ⟨wΔ^(T), o−⟩ ≥ −Õ(1/poly(d)) — the gating stays close to zero along class-relevant directions — while ⟨wΔ^(T), oj⟩ is bounded above by −ηTc′³/(16L) · [(1/2)^{ΔL{o+}^+ − 2} − (1/2)^{ΔL_{o_+}^− − 2}] · [(1/2)^{ΔL_{o_+}^+} + (1/2)^{ΔL_{o_−}^−}], a strongly negative update along irrelevant features.

• Guaranteed generalization for majority-voting data. Theorem 1: if the width satisfies m ≥ d² log q and token noise satisfies τ < O(1/d), then with probability at least 1 − N^(−d), N ≥ Ω(L²d / (η²(α_r − α_c)²)) samples and T = Θ(L² / (η(α_r − α_c)²)) iterations suffice for the population risk f(v^(0), W_O^(T), w_Δ^(T), W_B^(0), W_C^(0)) to equal 0.

• Guaranteed generalization for locality-structured data. Theorem 2: under the same width and noise conditions and the same probability, N ≥ Ω(L²d / (η²[(1/2)^{ΔL_{o_+}^+} − (1/2)^{ΔL_{o_+}^−}]²)) and T = Θ(L² / (η[(1/2)^{ΔL_{o_+}^+} − (1/2)^{ΔL_{o_+}^−}])) suffice for the population risk to equal 0.

• Larger signal and less noise mean faster learning. In both regimes, a larger fraction of class-relevant tokens (larger α_r − α_c) or a stronger locality gap (ΔL_{o_+}^+ ≫ ΔL_{o_+}^−), together with smaller token-level noise τ, reduces both the required iterations and the sample complexity.

• Locality is where Mamba diverges from the Transformer story. The authors state that for locality-structured data no such guarantee is provided in the prior Transformer analysis they cite (Li et al., 2023a), whereas their Mamba bounds follow from the recurrent mechanism exploiting the locality of class-relevant tokens.

• Synthetic experiments match the theory. Figure 2 is reported to show that increasing the majority-voting gap α_r − α_c consistently reduces the number of epochs across various training-set sizes, consistent with Equations 13 and 14, and to plot cosine similarity between w_Δ and class-relevant versus class-irrelevant features. Experiment details and additional numerical results are deferred to Appendix A.3.

Methodology in Plain English

The authors set up a binary classification problem with sequences of L tokens in dimension d, and they control the data generation themselves. An orthonormal basis contains two discriminative directions, o+ and o−, plus filler directions o_j for j ≥ 3 that carry no class information. Every token is one of these patterns plus Gaussian noise of level τ, so the researchers know exactly which tokens matter.

Labels are assigned in two ways. Under majority voting, the label follows whichever discriminative direction appears more often, with α_r the average fraction of class-relevant tokens and α_c the average fraction of confusing tokens, and class-relevant tokens outnumber confusing ones. Under locality structure, every sequence contains exactly two o+ tokens and two o− tokens; in positive samples the two o+ tokens sit close together and the two o− tokens are far apart, and negatives reverse this. That makes the class-relevant and confusion counts balanced, so the label must be inferred from spacing rather than counting. Datasets are balanced, with ||N_+| − |N_−|| = O(√N).

The model is a single selective SSM layer whose input-dependent gates are built from b_t = W_B^⊤ x_t, Δ_t = log(1 + e^{w_Δ^⊤ x_t}), c_t = W_C^⊤ x_t, producing the gates σ(w_Δ^⊤x_t) and 1 − σ(w_Δ^⊤x_t) that control how much new input is written and how much past state carries over. A two-layer ReLU MLP with output weights v reads the per-token outputs and averages over the sequence. Training minimizes a hinge loss with full-batch gradient descent.

Crucially, only W_O and w_Δ are trained; the output weights v are fixed at initialization and W_B, W_C are held at their initial values in the guarantees, with w_Δ initialized at zero and WO entries drawn from N(0, c₀²) (the lemmas state N(0, ξ²)). The hard technical part is tracking the gradient of wΔ. Because gating is multiplicative, the gradient decomposes into diagonal terms β{s,s}^(l) that do not depend on token position and off-diagonal terms β{s,s+1}^(l) that decay multiplicatively, so the proof must carefully balance competing contributions from different token positions to show the gate ends up pointing at the right features.

Why This Matters

Impact on research. The paper moves the theory of SSMs past expressivity results — universal approximation, comparisons with diagonal SSMs, and implicit-bias studies that omit input-dependent gating — toward training dynamics and generalization. It provides a theoretical counterpoint to Transformer-centric feature-learning explanations by showing that comparable guarantees can be obtained, and that locality-structured data gives Mamba an advantage, without any attention mechanism. It also casts light on why Mamba is reported to be sensitive to hyperparameter tuning by identifying which data properties (feature gap, locality gap, noise level) drive the required sample and iteration counts.

Real-world applications (the settings the authors use to motivate their data models):

• Image classification, where foreground patches act as class-relevant tokens and background patches as confusing or irrelevant ones.
• Object detection and localization, where decisive content is confined to a small spatial region.
• Image captioning, where prediction relies on a contiguous region rather than scattered background textures.
• Audio and speech recognition, where short phonetic segments carry the label, and genomics, where functional elements such as sequence motifs and regulatory regions are localized to short windows of DNA.

Industry relevance. Mamba-style models offer linear-time sequence processing and hardware-friendly parallelization, which matters for long-context language modeling, long video and audio streams, graph and vision backbones, and reinforcement learning. Knowing which data properties make selective recurrence learn efficiently — and how much data and compute are needed — informs architecture and dataset choices for teams deciding between attention and SSM backbones.

Future Directions

• Beyond a single layer and a single head. The analysis covers one selective SSM layer with one head followed by a two-layer MLP; extending the guarantees to deep multi-layer, multi-head Mamba stacks is the obvious next step.
• Relaxing the frozen parameters. The guarantees hold when W_B and W_C stay at initialization (their initial values) and only W_O and w_Δ move; analyzing joint training of all Mamba parameters, including the output projections W_B and W_C, is open.
• Richer noise and data models. The data model uses Gaussian token noise with τ < O(1/d) and either majority voting or locality structure; real sequences combine both effects and may have heavier-tailed noise, so a unified or more general treatment would sharpen the practical takeaways.
• Closing the empirical gap. The experiments are synthetic and confined to what is reported in Figure 2, with details deferred to Appendix A.3; testing whether the predicted scaling in α_r − α_c, ΔL_{o_+}^+, and τ appears in real language, vision, and genomic benchmarks remains open. A related question is characterizing the sensitivity to hyperparameter tuning that prior empirical work reports.

Target Audience

Theory-oriented machine learning researchers working on SSMs, Mamba, and sequence architectures; learning-theory groups studying feature learning, gradient descent dynamics, and generalization beyond the NTK regime; and practitioners with an interest in the theoretical footing of non-attention backbones who are comfortable with asymptotic notation and sample-complexity statements. Readers looking for empirical benchmarks, architecture recipes, or implementation guidance will not find them here, since the numerical results are synthetic and reported in brief.

Authors’ abstract

The recent empirical success of Mamba and other selective state space models (SSMs) has renewed interest in non-attention architectures for sequence modeling, yet their theoretical foundations remain underexplored. We present a first-step analysis of generalization and learning dynamics for a simplified but representative Mamba block: a single-layer, single-head selective SSM with input-dependent gating, followed by a two-layer MLP trained via gradient descent (GD). Our study adopts a structured data model with tokens that include both class-relevant and class-irrelevant patterns under token-level noise and examines two canonical regimes: majority-voting and locality-structured data sequences. We prove that the model achieves guaranteed generalization by establishing non-asymptotic sample complexity and convergence rate bounds, which improve as the effective signal increases and the noise decreases. Furthermore, we show that the gating vector aligns with class-relevant features while ignoring irrelevant ones, thereby formalizing a feature-selection role similar to attention but realized through selective recurrence. Numerical experiments on synthetic data justify our theoretical results. Overall, our results provide principled insight into when and why Mamba-style selective SSMs learn efficiently, offering a theoretical counterpoint to Transformer-centric explanations.