A Theoretical Analysis of Mamba’s Training Dynamics: Filtering Relevant Features for Generalization in State Space Models
Authors: Mugunthan Shandirasegaran
Hongkang Li
Songyang Zhang
Meng Wang
Shuai Zhang
Venue: ICLR 2026 Poster (2026)
Keywords: Training Dynamics
Feature Learning
Generalization
Mamba
Learning Theory
Selective State Space Models
Primary area: learning theory

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.

Published as a conference paper at ICLR 2026
A THEORETICAL ANALYSIS OF MAMBA'S TRAINING DYNAMICS: FILTERING RELEVANT FEATURES FOR GENERALIZATION IN STATE SPACE MODELS
Mugunthan Shandirasegaran
Hongkang Li University of Pennsylvania lihk@seas upenn. edu
New Jersey Institute of Technology ms ms3537 705577nnjiteei
Songyang Zhang
Meng Wang
University of Louisiana at Lafayette songyang. wh ana.
Rensselaer Polytechnic Insitute wangm7 wammmerrpiieee edu
Shuai Zhang
New Jersey Institute of Technology s2457@njit.en edu
ABSTRACT
The recent empirical success of Mamba and other selective state space models (SSMs) has renewed interest in non-attention architectures for sequence model- ing, yet their theoretical foundations remain underexplored. We present a first- step analysis of generalization and learning dynamics for simplified but rep- resentative Mamba block: a single-layer, single-head selective SSM with input- dependent gating, followed by 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 feature-selection role similar 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 theoretical counter- point to Transformer-centriic explanations.
INTRODUCTION
Transformers /Vaswani et al., 2017) have become the mainstream framework in large language mod-
els (Achiam al.] 2023; Guo al. 2025 Brown al., 2020 Touvron al. 2023). However, due
to the quadratic time and memory complexity introduced by the attention mechanism with respect to input length (Gu & Dao 2023 Dao & Gu 2024). Transformers are inefficient when handling long
input sequences. Recently State Space Models (SSMs) (Gu & Dao, 2023; Dao & Gu, 2024 Zhu
al. 2024 Wang al. 2024a] Behrouz & Hashemi 2024] Liu et al. [024 Wang et al. (024b)
have shown competitive or superior performance to Transfrmers across domains such as language
(Gu Dao 2023), vision (Zhu et al. 2024 Liu et al. 2024),, graphs (Wang et al. 2024a Behrouz &
Hashemi 2024), audio (Yadav & Tan 2024). and reinforcement learning (Lu et al. 2022)) SSMs has
models, such as linear computational com- properties that enable efficient parallelization. Among these models, Mamba (Gu & Dao, 2023) proposes selection mechanism, which parameterizes the SSM with the
brought many advantages absent in Transformer-bart
plexity and hardware-friente
input, which allows the model to dynamically retain discard relevant and irrelevant information.

---

Published as a conference paper at ICLR 2026
This enables the Mamba model to achieve performance comparable Transformer- based models in long-text modeling as well as tasks such as visual classification and dense prediction (Zhu et al. 2024
Liu et al., 2024), but in more efficient manner.
Although recent work has primarily focused on the empirical performance of Mamba and its ar- chitectural comparisons with other models, the theoretical understanding of Mamba remains less investigated. In addition, recent empirical evidence shows that Mamba's 's success is highly sensitive hyperparameter tuning (Okpekpe & Orvieto 2025). Such dependence on fragile optimization choices raises fundamental questions about why and when Mamba succeeds. These include funda- mental inquiries such as:
* Under what conditions can Mamba be trained to achieve satisfactory generalization?
How is the selection mechanism implemented through Mamba's components?
Existing theoretical studies on Mamba or related SSMs mainly focus on the expressive power and the mechanisms of optimal parameters. Orvieto et al. (2024) and Nishikawa & Suzuki prove SSMs augmented with MI.ss are universal approximators of regular functionals and can mimic token selec- tion dynamically. MMea Cirone et al, (2024) and Huang et al. (2025) show that Mamba has stronger expressive power than its diagonal SSM predecessor, especially in approximating discontinuous functions. Li et al. (2024c) and Li et al. (2025c) respectively prove that two simplified SSMs, H3 and GLA, implicitly perform weighted preconditioned GD at the global minima of in-context learn- ing problems when input with context examples. However, these works do not explain whether the selection mechanisms and advantages of Mamba can actually be obtained through practical training. Moreover, these studies do not analyze the generalization ability of Mamba models.
Contributions of this paper. In this work, we study nonlinear neural model composed of one- layer Mamba block and two-layer perceptron, which is simplified but sufficiently representative reflect the gating structure in Mamba. By assuming the presence of the class-relevant feature that influence the label and class-irrelevant features that do not, we respectively formulate majority- voting and locality- locality of a certain class-relevant feature in the data. To the best of our knowledge, this work provides the first theoretical analysis of Mamba's training dynamics with input-dependent gating, together with generalization guarantees under the two structured data regimes. The highlights of our technical contributions include:
data, whose labels depend on the proportion and the spatial/temporal
First, develop general theoretical framework for analyzing gated architectures trained with gra- dient descent on structured data. Our analysis explains how the selection mechanism within Mamba interacts with data structure to enable efficient learning and guaranteed generalization, complement- ing prior results that focus mainly on attention based models.
Second, we provide theoretical characterization of the gating mechanism in Mamba. We show that the gating parameter vector is trained to amplify delevvat features while ignoring class- irrelevant ones, thereby formalizing the intuition that the gating network dynamically allocates pacity to informative patterns.
Third, we establish the sample complexity and the required number of iterations for two canonical data types: majority-voting and locality-structured data sequences. For majority-voting data, these bounds scale with the gap between the class-relevant and confusion features; for locality-structured data, they depend on the concentration of class-relevant tokens. In both regimes, stronger signal and lower token-level noise yield faster convergence and better generalization.
1.1 RELATED WORK
State Space Models (SSMs). Building upon the early $4 models (Gu et al. 2021 Guppa et al, 2022 Smith et a..)., Mamba (Gu &
Dao & Gu (024) introduced input-dependent gating to dynamically select relevant features, achieving remarkable performance in NLP and CV. Recent works extending SSMs beyond ID sequences have highlighted the importance of input ordering and scanning. For example, Vammba (Liu et al., 2024 introduces SS2D, employing multiple scanning routes to bridge sequential structure with the non-sequential nature of vision inputs, while Graph Mamba (Wang et al., 2024a, Behrouz & Hashemi, 2024) adapts SSMs to non- non-Euclidee Euclidean domains by leveraging graph connectivity. Collectively, these works show that the effectiveness of SSMs is
2

---

Published as a conference paper at ICLR 2026
tightly linked to input ordering and scanning strategies, a challenge that also motivates theoreti- cal analysis.
Theoretical Analysis of SSMs. Theoretical understanding of Mamba is still in its early stages and has so far centered primarily on approximation theory, such as connections to attention like mech-
anisms (Dao & Gu, 2024 Nishikawa & Suzuki), expressive capacity (Cohen-Karlik et al. 2025
Huang et al. [025| Muca Cirone et al. [024]] Bao et al. 2025). long-range dependency modeling ling
(Ma & Najarian 2025] Yu & Erichson 2025), and the comparison with Transformers (Jelassi al,
2024). Beyond approximation theory, several recent works have begun examining optimization and
(2025) provide generalization error bound based on Rademacher complexity; Slutzkyy et al. (2024) study implicit bias under a teacher-student setting and show that gradient flow can converge to a low-rank solution, though their model does not incorporate Mamba's 's input-dependent gating. These analyses provide valuable intuition about the representational strengths and weaknesses of Mamba blocks. However, such results remain largely structural: they establish only the existence of desirable representations, without explaining whether or how these capabilities arise during training, particularly under Mamba's 's unique mechanism. Mo- tivated by this gap. focus studying how Mamba interacts with structured data, with particular emphasis on the role of its gating mechanism in shaping training dynamics and generalization.
generalization aspects of SSMs. Honarpisheh
Feature Learning Framework. Recent theoretical studies of deep learning have shifted focus from the NTK framework (Jacot et al.) 2018: Allen-Zhu et al.) 2019b] Arora et al.] 2019) to the feature-leaming framework, where data is modeled as combination of features and the central question is how neural networks align with these features. Much of the recent work has concentrated
on transformers (Li al. 2023al 2024b 2023b 2025b
a Liao al.) 2025), feedforward neural Arora et al. 2019 Wen & Li. [002]] Sun et al.] (2025), graph neural networks (Zhang et al., 2023 Li al.j 2024a), and of Chowdhury al. (2023).
networks (Bakshi et al.,
2019
Due to the inherent complexity of non-convex optimization and modern architectures, prior works on feature learning have, to the best of our knowledge, focused primarily on shallow networks. In this work, we extend the structural data model to analyze the training dynamics of shallow yet representative Mamba block, with particular emphasis on how its data-dependent gating mechanism shapes learning and generalization.
2 PRELIMINARIES
Structured state space models (S4). For the 1-th token, e.g., at time step t, let æ; E Rd be the input, H, E R denote the corresponding hidden state, and ye E Rd denote the output. Let A E RN XN and b, R be model parameters. The discrete-time SSM is given by
H; = AH,-1 + bot, Y/- =H, c, where A = ex)(AA) and b = A (cxp( I)b with A > 0 as the sampling step.
(1)
Mamba. To overcome the data-independence of S4, recent work introduced selective state space models (Gu & Dao 2023), where key parameters are made input-dependent. Concretely, given input tokens F, E IR0 the recurrence parameters are defined as
b W Dt.
A. = log (1 + e 2
C₁ WT Dz.
(2)
with learnable projections WB, We E Rdx and a gating vector wA E Rdd. The discretization then yields two input-dependent gates,
b = σ (wIw.) b..
ā = 1 - (w)
which respectively control the input update and the carry-over of past states. With hidden state H, E R the recurrence becomes
H: = H,
(4)
Mamba output at token is given by:
y/ (X): := H c = a(w xz) (W xx) (Woxx)) oi + (11 (()]) H H
Σ J=
(1 (01-0wwaa) ((wz)z))) -(wIxs) (WI. xx) (WJx)).. x,. (5)
3

---

Published as a conference paper at ICLR 2026
Connection and Difference with Transformer. The Mamba formulation reveals a natural analogy to attention mechanisms (Dao & Gu 2024 Sieber et al. 2024). In particular, the input-dependent matrices We and WB be interpreted as counterparts to queries and keys in the self-attention, while the gating term a 1 mx)) acts as a dynamic weight controlling how past information con- tributes to the current output (Dao & Gu 2024). This structure yields a formulation closely related 3025c/LLu Lu al. 2025), thereby highlighting a connec- tion between SSM and Transformer models. Meanwhile, Mamba departs from these architectures: its gating mechanism is defined through multiplicativv interactions, effectively involving products of successive terms. This nonlinearity makes the analysis Mamba substantially different and more challenging than that of gated linear attention. Unlike additive attention-style weighting, Mamba's gating introduces input-dependent multiplicative modulation in the selection mechanism. This alters how information is propagated through the model and results in training dynamics that differ from attention- based architectures.
to gated linear attention (Yang et al. [i et 2025c
PROBLEM FORMULATION
Following existing works (Brutzkus & Globerson, 2021; Zhang et al. 2023 Li et al. 2023a),
consider a binary classification problem with training data { (x)),, 2 n= sampled i.i.d. from
an unknown distribution D, where 2(") E [+1,-1) is the label. The goal is to learn a model that
maps X to 2 for any (X; 2) D. Each input takes the form (() [ax (n)
(1)) R L
with L tokens, where each token is d-dimensional. Tokens can be image patches (Dosovitskiy et al., 2021 Touvron et al. 2021 or subwords (Sennrich et al., 2016: Kudo Richardson 2018).
Learning is performed using simplified Mamba block formulated by 5 followed by two-layer ML., Formally, the model output can be expressed as
m 1 Σ
F(X)
ds (W
(6)
where (--) denotes the ReLU function, and Wo E mxd with Wor,)) being the i-th row of Wo. Here, (X) corresponds to the I-th token output of Mamba, as defined in (5) addition, 2 represents the output-layer weight for the i-th hidden unit.
Model Training. Let (v. Wo, WA WB, Wc) denote the set of model parameters. The training process is to minimize the empirical risk fx( (
NV
min (N(
((X (m))
N
==1
where (mX(m), (n);; D is the hinge loss function, i.c.
(( ((x("), (()) v ma 2(1). F(X()))
(8)
The empirical risk minimization problem in 7 is solved via gradient descent (GD). For the theo- retical analysis, we consider the full batch gradient update with learning rate of 1) at each iteration t= -11,3 1,2, .......T. Each entry of Wo E Rm is independently initialized from N N/(0, <g) and WA is initialized 0. Similarly, each entry of IRm is independently sampled from 7m 7m
with equal probability. v is fixed during training, as in other theoretical works (Allen
Li
2022 Arora et al. 2019 Karp et al. 2021 Allen-Zhu al. 2019a Li et al. 2023a 2024b).
Generalization.
The generalization error of the learned model V is evaluated using the population risk f(V), defined
f(w) = f(v, Wo. WA WB, We) = Exx,,)<D ((22)).
THEORETICAL RESULTS
4.1 KEY TAKEAWAYS AND INSIGHTS OF THE FINDINGS
Before formally presenting our data assumptions and theoretical results, we first summarize key insights derived from our theoretical findings. We consider data model where tokens noisy
4

---

Published as a conference paper at ICLR 2026
Table 1: Some important notations
yi
Mamba block output at token position
N Number samples in batch
Embedding dimension
m
The number of neurons Wo
"
Learning rate for gradient descent
L
Length of the squence
A A
Concentration of class-relevant tokens
0r
Average fraction of class-relevant tokens
A o.
Dispersion of the confusion tokens
O2 Average fraction of confusion tokens
versions of class-relevant patterns that determine the data label and class-irrelevan patterns that do not affect the label. Some important parameters are summarized in Table
(TI). Convergence and sample complexity analysis of GD to achieve guaranteed generaliza- tion. We introduce a theoretical framework for analyzing gated architectures with structured data. Compared with existing results on attention- baaedd models, our framework captures the role of the gating mechanism inside the Mamba block and structured weight interactions, explaining how gra- dient descent (GD) exploits data structure to improve learning efficiency. Based on this analysis, we show that model trained with GD achieves guaranteed generalization with high probability over the randomness of the data and the GD updates.
(T2). Theoretical characterization of the gating mechanism in Mamba. We prove that dur- ing training, the gating network learns to prioritize class-relevant features while ignoring irrele- vant ones. In the majority-voting regime, the gating vector wa becomes increasingly aligned with class-relevant directions: gradients along those directions grow, while those along irrelevant fea- tures remain negligible. In the locality-structureed data regime, learning emphasizes the elimination of irrelevant features. Their directions are consistently pushed downward by negative updates, while the directions of relevant features remain nearly unchanged. This occurs because class-relevant and confusion tokens appear in equal proportion, so the model cannot amplify the former and instead reduces the influence of the latter. These dynamics strengthen informative tokens and weaken un- informative ones, inducing effective sparsity in the activations and formalizing the intuition that Mamba allocates capacity to the most important patterns in the data.
(T3). Larger fraction or higher local concentration of class-relevan features accelerates learn- ing. We show that both the number of iterations and the sample complexity required for gen- eralization depend on the discriminative structure of the data and the token-level noise т. For majority-voting data, these quantities scale as (a,- - ac)) so learning is faster when the frac- tion of class-relevant tokens is larger. For pocality-structure data, the number of iterations scales as [(1) (1)= while the sample complexity scales as [(]) - A6, -2 Here, AL,, denotes the separation between class-relevant features 0+ in positive samples (capturing their locality), and AT denotes the separation between confusion features 0+ in negative samples (cap- turing the locality of confusing patterns). Thus, when ALT >> AL, the locality of class-relevant features dominates, which reduces both the number of iterations and the sample complexity needed for convergence, implying faster learning when class-relevant tokens are more concentrated locally. Finally, in both regimes, smaller token-level noise further accelerates learning.
4.2
DATA MODEL
Consider an arbitrary set of orthonormal vectors O = [0+, 0_. .z. in Rd, where o., and o_.... are discriminative features and the remaining vectors jj 2 3, are class-irrelevant (filler) features. Depending on the class label, either 04 o_ serves as the class-relevant pattern, while the other
(m) acts as a confusion pattern. Each token x; in (() is a noisy version of one of the input patterns (features), i.e., (m) o + where O and & is the Gaussian noise. We consider two different data types: majority-voting and
data.
Majority Voting Data. For the majority voting data type, the label is determined by majority vote over the class-reld relevant patterns. Let a, and a, denote the average fractions of class-relevant tokens and confusion tokens over the distribution D, respectively. In positive samples, noisy variants of o
5

---

Published as a conference paper at ICLR 2026
class-relevant, while noisy variants of o_ act confusion tokens. In negative samples, the roles are reversed. All other tokens correspond to -tree-irrelevant features.
Locality-structuu Data. For the Docality-structuree data type, each sequence contains two 0+ tokens and two o tokens, while all other tokens correspond to class-irrelevant features. In positive samples, the two 0+ tokens are close to each other, while the two tokens far apart, formally, AL. & AL where AL; and AL denote the distances between the two 04 and o_ tokens, respectively. In negative samples, the pattern is reversed: AT-__ < '...o.
In addition, we consider a balanced dataset sampled from the unknown distribution D. Let W4 {(x("),
=(1) = +1, E [N]} and N_ = ((x(), ((()) z(() -1, n E [N]} denote the sets of positively and negatively labeled samples, respectively. Then the class balance satisfies N).
2())) (x(())2(m))
[ - [
Interpreting the Data Model in Practice. Our theoretical data models are motivated by common patterns observed in practical machine learning tasks.
On the one hand, the majority-voting data model captures a widely adopted assumption (Li et al, 2023a 2024b) in theoretical analysis, whereby the label is determined by the aggregate contribution through majority votc. For example, in image classification tasks (Krizhevsky et al. 2012) Simonyan Zisserman 2014 He et al. 2016). the class label is often driven by multiple discriminative patches corresponding to foreground objects (class-relevant tokens). In contrast, background patches may contain other objects or patterns that are not associated with the target class (confusing tokens), along with random patches that are entirely unrelated (class-irrelevant tokens) (Dosovitskiy et al., 2021 Touvron et al. 202I).
On the other hand, the Docalit--stuccture data corresponds to tasks where semantic meaning is concentrated in spatially or temporally localized clusters, while background features are more dis- persed. This structure is most familiar in vision tasks such object detection and localization (Ren
et al.
2015
Carion et al. 2020 Zhou et al., 2016) and image captioning (Vinyals et al. 2016 Xu
al.
2013 Radford al. 2021). Ip. where the decisive content is often confined to small region of
the image. For example, in an image labeled "dog in park," the prediction relies primarily on the
contiguous region containing the dog rather than on scattered background textures. A similar prin- ciple holds in audio and speech recognition (Yadav & Tan [2244] Gulati et al.) (2220), where short phonetic segments capture the information needed recogmize words, and in genomics (Alipanahi] et al. 2015; Zhou & Troyanskaya 2015), where functional elements such as sequence motifs and regulatory regions are localized to short windows of DNA. In these settings, the local structure of nearby tokens strongly correlates with the label.
Together, the majority-voting and locality- locality-structured models offer complementary perspectives on when selective recurrence can most effectively support learning from structured real-world data.
4.3 FORMAL THEORETICAL RESULTS
4.3.1 THEORETICAL RESULTS FOR MAJORITY- VOTING DATA
We next present a lemma characterizing how the gating vector aligns with different features under the majority-voting data.
Lemma 4.1 (Gating Vector Alignment for Majority Voting Data). With initialization where each
entry of Wo is drawn independently from NO0, and w 0. With sufficient number of training samples and iterations, we have
nT Olaal -
(
0+
>
812
(10)
w 0 2 O( aa
(
1i
(11)
8L2
( (0 0,) < 5 (1)) Vj 3.
(12)
Lemma 4.1 establishes that after sufficient training, the gating vector wA aligns positively with the class-relevant features 04 (10) and
shown in (IT), while its alignment with irrelevant features

---

Published as a conference paper at ICLR 2026
remains strictly negative shown in 12 In other words, the selection mechanism implicitly acts as a feature selector, amplifying relevant tokens and ignoring irrelevant ones. Lemma serves as an informal version of Lemmas B.5 and B.6
Remark 1: With majority voting data, the gating vector aligns with discriminative features, i.c., o+ and o_.. As result, the model's output focuses primarily these features, giving more weight to tokens that carry discriminative features while reducing the influence less important tokens. Since the number of class-relevant tokens is greater than the number of confusing ones, e.g., in positive sample, the tokens containing O+ outnumber those containing o-, the model can correctly assign the label through this majority effect. Furthermore, as the difference between the counts of class- relevant features and confusing features (i.e.. a,, a.)) increases, the gating vector converges much faster. Overall, this gating mechanism allows the model to use training samples more efficiently because it learns to emphasize the most relevant feature early on and ignore irrelevant features.
We now present the theorem establishing the generalization guarantee for Mamba under the majority-voting data.
Theorem 1 (Generalization for Majority Voting Data). Suppose the model width satisfies m 2 d² log q for some constant q > O, and the token noise level is bounded as < O()). Then, with probability at least N-d, if the number of training samples N satisfies
I2d
N20
(13)
17 ac)²
and the number of iterations' I satisfies
T= m/a,, az)²
(14)
the resulting model achieves guaranteed generalization, i.e.,
(v W (T) w (0)
(T)
w 0.
(15)
Theorem T establishes the sample complexity, as shown in 13 and the convergence rate, as given in (14). that are required to guarantee desirable generalization when training the model in (6) using GD for the majority-voting data type. In other words, the model achieves good generalization once a sufficient number of samples is available, as specified in (13). and training has proceeded for sufficient number of iterations, as specified in (14).
Remark 2: With majority-vtinng data, the Mamba architecture can effectively capture the under- lying data distribution by first identifying discriminative features through its gating mechanism and then aggregating them via data-dependent recurrent mechanism. In this sense, Mamba behaves similarly to the Transformer (Li et al. 2023a), suggesting a close connection between the two mod- els despite their architectural differences. According to the results Lemma 4.1 the model further benefits from a faster convergence rate and reduced sample complexity when the gap between class- relevant and confusing features larger.
4.3.2 THEORETICAL RESULTT FOR SECALITY-STRCCTURT DATA
We next present a lemma characterizing how the gating vector aligns with different features under the Docality-structune data.
Lemma 4.2 (Gating Vector Alignment for Locality-structureed Data). With initialization where each entry of Wo is drawn independently from ((0,e)) and w 0. With sufficient number of training samples and iterations, we have
(w (T)
0+) 2
- Ō (1/poly(d)),
(16)
(w A ) 2 - (1/polyy(d)
(T)
(17)
AL -2
Al. -2
srt
(
0,
1
(18)
16L
2
7

---

Published as a conference paper at ICLR 2026
Lemma 4.2 establishes that after sufficient training, the gating vector W remains close to zero for
class-relevant features 0+ as shown in (16) and o_ as shown in (I7). however its alignment with irrelevant features remains strongly negative as shown in (T8). Through this mechanism, the gating favors class-relevant features to select the most informative feature for learning. Lemma 4.2 serves
as informal version of Lemmas C.3 [and
Remark 3: The gating vector behaves differently from majority voting, though the overall insights remain similar. We can no longer guarantee that wa will always grow in the direction of discrimi- native features, because we assume that the number of class-relevant features can be comparable to the number of confusing features. This assumption is introduced highlight the role of data locality in shaping the gating vector, which is more challenging to analyze in isolation since majority voting can readily reinforce it; however, their combined effect better reflects real-workee data. Although this direct growth no longer holds, the gating vector consistently decreases in the direction of irrelevant features. At a higher level, this can be seen as synergistic interaction: the recurrent mechanism captures locality and suppresses irrelevant features, which pushes the gating vector to decrease along those directions, while the gating itself further amplifies this suppression. From another perspective, by making the model pay less attention to irrelevant features, the gating vector effectively shifts more attention toward discriminative features.
We now present the theorem establishing the generalization guarantee for Mamba under the locality- structured data.
Theorem (Generalization for struuity- Data), Suppose the model width satisfies m 2 dl log for some constant >> O, and the token noise level is bounded as T < (2)). Then, with probability at least - N if the number of training samples N satisfies
Ltd
N 2 Ω
(19)
? [((122)
ALL. (1/2)Alo) }3
and the number of iterations T satisfies
12 n (1122) - (1/2) slo
T T=
(20)
the resulting model achieves guaranteed generalization, i.e.,
(o w w
= 0.
(21)
Theorem 2 shows that good generalization on Docality-structuree data is guaranteed if the sample complexity meets (19) and training proceeds for at least (20) iterations.
Remark 4: We establish that Mamba can also effectively learn this type of data through its ability to exploit locality, in contrast to Transformers, where no such guarantee is provided in (Li al (023a)) In our analysis. A captures the distance between class-relevant tokens, reflecting the locality of class-relevant features, while A captures the locality of confusing features. The effectiveness of learning is governed by the separation between these two quantities. In particular, when A > AL the locality of class-relevant features dominates that of confusing features. In particular, when A A > A AL.,+ the locality of class-relevant features dominates that of confusing ones, which reduces both the sample complexity and the number of iterations required for convergence, allowing Mamba to learn more effectively and efficiently.
4.4 TECHNICAL NOVELTY AND CHALLENGES
Differences with Existing Works. Our work mainly inspired by prior feature-learning analyses of (Bakshi et al., 2019 Arora et al. 2019 Brutzkus & Globerson, 2021 Li et al. 2023a 2025b).
Building on these foundations, we develop a framework specifically tailored to gated architectures with structured data. Unlike these existing models, Mamba introduces an input-dependent gating mechanism, absent from other network architectures, which acts as a dynamic selection operator and requires new analytical techniques to capture its learning dynamics. Moreover, while the majority- voting data model has been previously studied in the context of Transformers (Li et al. 2023a), show that Mamba also learn this type of data with comparable performance. Furthermore,
8

---

Published as a conference paper at ICLR 2026
find that Mamba is particularly effective at capturing the inherent locality of the data, which motivates us to introduce a new generalization guarantees within the framework of selective state space models, thereby advancing our understanding of this class of architectures and clarifying their distinctions from Transformers. A proof sketch be found in Appendix A.2
data model. For both regimes, we establish
Technical Challenges. Our analysis faces several unique technical challenges stemming from the structure of selective SSMs. Unlike attention-based models, where interactions are primarily addi- tive, Mamba's gating mechanism introduces multiplicative recurrences across tokens, with dynam- ics that are explicitly sensitive token order. These multiplicative effects accumulate over time, substantially complicating the training analysis. To capture this behavior, we systematically track the gradient updates of the gating vector wa, decomposing the contributions from different token positions and analyzing how token placement influences training dynamics.
Specifically, in the majority-voting data, the gradient decomposition of the gating includes off- g(t) diagonal terms +, that exhibit additional multiplicative decay due to the recursive gating struc-
ture. whereas the diagonal term a!! is independent of token position. Hence, it is important to carefully consider competing token contributions to prove that the gating vector is indeed aligned with class-relevant feature directions.
Instead, in the Docality-structinn confusion tokens in positive and negative samples is negligible, as our data model assumes an equal number of fels-rel and confusion tokens. Consequently, we need rely on AL and A AL ensure that the lucky neuron Wo(s.) learns the class-relevant feature. Moreover, since the number of class-relevant and confusion tokens is balanced, updates along the class-relevant feature direction for the gating vector remain close zero. To demonstrate how the gate filters information, we show that gradient updates along dtass-irrelevat features are driven strongly negative. To prove this, in addition to the terms considered in the majority-votinn setting, we must also bound the positively contributing terms that hinder the gate's ability to suppress irrelevant features. Specifically, we
data, the variation introduced by the number of class-relevant and
bound these opposing terms as O ((1-o (p2)) AL. + O ((1 (P2)) ensuring that their effect remains minimal. This reveals that the gate effectively suppresses irrelevant features while preserving class-relevant features for this data model.
5 NUMERICAL EXPERIMENTS
We verify our theoretical results through synthetic experiments based the data models described in Section 4.2 Due to the space limit, we defer the experiment details and additional numerical results to Appendix A.3
Faster convergence with larger majority-voting gap. Fig. illustrates that increasing the majority-voting gap ar ae consistently reduces the number epochs across various sizes of training samples. These findings are consistent with our theoretical results in (1) and (14).
Gating mechanism amplifies relevant features majority-voting data. Fig. shows the cosine similarity between the gating vector wA and both class-relevant and class-irrelevant features. The similarity with class-relevant features steadily increases, while that with class-irrelevant features remains essentially unchanged. This empirically confirms Lemma 4.1 demonstrating that the gate prioritizes informative features while ignoring irrelevant ones.
Locality affects the learning. Fig. 3 illustrate the effecct of class-relevant token separation AL on the convergence in the locality-structuree data. Larger slows convergence across different training sample sizes, which is consistent with our results in 19 and (20).
Gating mechanism suppresses irrelevant features. Fig. g illustrates that while the cosine similar- ity is negative for both types of features, it stays close zero for class-relevam features but becomes largely negative for class-irred irrelvannt ones. This contrast drives the gating mechanism to prioritize class-relevant features, consistent with Lemma 4.2

---

Published as a conference paper at ICLR 2026
Epochs to Reach 99% Accuracy
Convergence Behavior of Mamba
Alignment of WA during training
0.6
N=50
N=100
Cosine similarity
0.5
N-200
0.4
5-
0.3
Class-relevant feature
Class-irrelevant feature
0.2
0.1
0.0
0
50
100
150
200
250
0.30
0.35
0.40
0.45
0.50
Iteration
Fraction class relevant patterns
Figure Alignment of wA for majority-voting data.
Figure I: Convergence vs. majority-voting gap.
Epochs to Reach 99% Accuracy
Convergence Behavior of Mamba
Alignment of wa during training
0.00
N=50
Class-relevant feature
800
N-200
Cosine similarity
-0.05
Class-irrelevant feature
N-400
-0.10
600
-0.15
400
-0.20
-0.25
-0.30
-0.35
0
100
200
300 400 500 600 700 Iteration
Distance between class-relevant patterns
Figure 3: Convergence under locality-structured data.
Figure 4: Alignment of ws for locality- structured data.
CONCLUSION
Encouraged by the emergence and successful applications of the Transformer alternative architecture Mamba, this paper provides novel theoretical generalization analysis of Mamba by considering its unique gated selection mechanism. Focusing on a data model with class-relevant and class-irreeewe tokens, we establish the con-asymptotic sample complexity and the convergence rate required to achieve desirable test accuracy. Our analysis further shows that the gating parameter vector filters out the class-relevant features while ignoring irrelevant ones. To the best of our knowledge, this is the first theoretical analysis of Mamba's 's training dynamics, with its input-dependent gating mechanism, together with generalization guarantees.
Finally, we note some limitations of our work. First, our theoretical analysis focuses on a simplified Mamba setting that abstracts away practical components such as depth, multiple heads, residual con- nections, and layer normalization. Second, our data model, while standard in theoretical studies, also simplifies real-world sequence structures. Extending the analysis to more realistic multi-layer and multi-head Mamba architectures, richer data models, and alternative designs such gated Trans- formers or hybrid
frameworks remains an important direction for future work.
ACKNOWLEDGMENTS
This work was supported part by the National Science Foundation (NSF) under Grants #2349879, #2349878, #2425811, and #2430223. Part of Hongkang's work completed while he was Ph.D. student at Rensselaer Polytechnic Institute (RPI) and was supported in part by the Army Research Office (ARO) under Grant W911NF-25-1-002 as well by the Rensselaer-IIM Future
10

---

Published as a conference paper at ICLR 2026
of Computing Research Collaboration (http://airo) rpp.rppi...... edu). We also thank the anonymous reviewers for their constructive and insightful comments.
LLM USAGE DISCLOSURE
We used langg-a models (ChatGPT) aid polishing the writing of this paper. For numer- ical experiments, we employed Al-assisted coding tools (GitHub Copilot and ChatGPT) support code development.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint Xii:23303008774, 2023.
Babak Alipanahi, Andrew Delong, Matthew T Weirauch, and Brendan J Frey. Predicting the se- quence specificities of dna-and rna-binding proteins by deep learning. Nature biotechnology. 33 (8):831-838 2015.
Zeyuan Allen-Zhu and Yuanzhi Li. Feature purification: How adversarial training performs robust deep learning. In 2021 IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS), pp. 977-988. IEEE, 2022.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparame- terized neural networks, going beyond two layers. Advances in neural information processing systems, 32, 2019a.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference Machine Learning, pp. 242-252. PMLR, 2019b.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of op- timization and generalization for overparameterriee two-layer neural networks. In International Conference on Machine Learning, pp. 322-332. PMLR, 2019.
Ainesh Bakshi, Rajesh Jayaram, and David P Woodruff. Learning two layer rectified neural networks in polynomial time. In Conference on Learning Theory, pp. 195-268. PMLR, 2019.
Zeyu Bao, Penghao Yu, Haotian Jiang, and Qianxiao Li. The effect of depth on the expressivity of deep linear state-space models. arXiv preprint arXiv:2506. 19296, 2025.
Ali Behrouz and Farnoosh Hashemi. Graph mamba: Towards learning graphs with state space models. In Proceedings of the 30th ACM SIGKDD conference on knowledge discovery and data mining, pp. 119-130, 2024.
Tom Brown, Benjamin Mann, Nick Ryder, Mclanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Alon Brutzkus and Amir Globerson. An optimization and generalization analysis for max max-pooling networks. In Uncertainty Artificial Intelligence, pp. 1650-1660. PMLR, 2021.
Nicolas Carion. Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference computer vision, pp. 213-229 Springer, 2020.
Mohammed Nowaz Rabbani Chowdhury, Shuai Zhang, Meng Wang, Sijia Liu, and Pin- Yu Chen. Patch-level routing in mixture-of-experts is provably sample-efficient for convolutional neural networks. In International Conference on Machine Learning, pp. 6074-6114, PMLR, 2023.
Edo Cohen-Karlik Itamar Zimerman, Liane Galanti, Ido Atad, Amir Globerson, and Lior Wolf. On the expressivity selective state-space layers: A multivariate polynomial approach. arXiv preprint arXiv:2502.02220 2025.
11

---

Published as a conference paper at ICLR 2026
Tri Dao and Albert Gu. Transformers ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21066, 2024.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby, An image worth 16x scale. ICLR, 2021.
words: Transformers for image recognition at
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint (.Xi::2312.00755. 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint (.Xix:21110003000 2021.
Anmol Gulati, James Qin, Chung- Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, et al. Conformer: Convolution -uggionted transformer for speech recognition. In Proc. Interspeech 2020, pp. 5035550360000 1006-504,2200 2020.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.12945 2025.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces, Advances in neural information processing systems, 229994,,, 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pailern recognition, pp- 770-778, 2016.
Arya Honarpisheh, Mustafa Bozdag, Octavia Camps, and Mario Sznaier. Generalization ror analysis for selective state-space models through the lens of attention. arXiv preprint arXiv:2502.01473. 2025.
Ningyuan Huang, Miguel Sarabia, Abhinav Moudgil, Pau Rodriguez, Luca Zappella, and Federico Danicli. Understanding input selectivity in mamba: Impact on approximation power, memoriza- tion, and associative recall capacity. arXiv preprint arXiv:2506. //811, 2025.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and gener- alization in neural networks. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after me: Trans- formers are better than state space models at copying. arXiv preprint urXiv:2402.01032, 2024.
Stefani Karp, Ezra Winston, Yuanzhi Li, and Aarti Singh. Local signal adaptivity: Provable feature learning in neural networks beyond kernels. Advances in Neural Information Processing Systems, 34:24883-24897, 2021.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convo- lutional neural networks. Advances in neural information processing systems, 25, 2012.
Taku Kudo and John Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. EMNLP 2018, pp. 66, 2018.
Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Yu Chen. A theoretical understanding of shallow vi- sion transformers: Learning, generalization, and sample complexity. In International Conference on Learning Representations (ICLR 2023), 2023a.
Hongkang Li, Meng Wang, Tengfei Ma, Sijia Liu, Zaixi Zhang, and Pin-Yu Chen. What improves the generalization of graph transformers? a theoretical dive into the self-attention and positional encoding. In International Conference on Machine Learning, pp. 28784 28829. PMLR, 2024a.
Hongkang Li, Meng Weng, Songtao Lu, Xiaodong Cui, and Pin- Yu Chen. How do nonlinear trans- formers learn and generalize in in-context learnin?? In International Conference on Machine Learning, 2024b.
12

---

Published as a conference paper at ICLR 2026
Hongkang Li, Songtao Lu, Pin- Yu Chen, Xiaodong Cui, and Meng Wang. Training nonlinear trans- formers for chain-of- thouugt inference: A theoretical generalization analysis. In International Conference Learning Representations, 2025a.
Hongkang Li, Yihua Zhang, Shuai Zhang, Pin- Chen, Sijia Liu, and Meng Wang. When is task vector provably effective for model editing? generalization analysis of nonlinear transformers. In International Conference on Learning Representations, 2025b.
Yingcong Li, Ankit S Rawat, and Samet Oymak. Fine- Fine-grained analysis of in-context linear estima- tion: Data, architecture, and beyond. Advances Neural Information Processing Systems, 37: 138324-138364 138364, 2024c.
Yingcong Li, Davoud Atace Tarzanagh, Ankit Singh Rawat, Maryam Fazel, and Samet Oymak. Gating is weighting: Understanding gated linear attention through in-context learning. arXiv preprint MXXi::2504.04300 2025c.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In International Conference on Machine Learning, pp. 19689-19729. PMLR, 2023b.
Haixu Liao, Yating Zhou, Songyang Zhang, and Shuai Zhang. On the training dynamics of con- trastive learning with imbalanced feature distributions: A theoretical study of feature learning. In UniReps: 3rd Edition of the Workshop on Unifying Representations in Neural Models, 2025.
Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye. Jianbin Jiao, and Yunfan Liu. Vmamba: Visual state space model. Advances in neural information processing systems, 37:103031-10300 2024.
Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. Advances in Neural Information Processing Systems, 36:47016-47031, 2023.
Peng Lu, Ivan Kobyzev, Mehdi Rezagholizadeh, Boxing Chen, and Philippe Langlais. Regla: Re- fining gated linear attention. arXiv preprint arXiv:2502.01577, 2025.
Cong Ma and Kayvan Najarian. Rethinking the long-range dependency in mamba/ssm and trans- former models. arXiv preprint arXiv:2509.04226, 2025.
Nicola Muca Cirone, Antonio Orvieto, Benjamin Walker, Cristopher Salvi, and Terry Lyons. The- oretical foundations of deep selective state-space models. Advances in Neural Information Pro- cessing Systems, 17:122226-1 1121122266-12272, 2024.
Naoki Nishikawa and Taiji Suzuki. State space models are provably comparable to transformers in dynamic token selection. In The Thirteenth International Conference Learning Representa- tions.
Destiny Okpekpe and Antonio Orvieto. Revisiting associative recall in modern recurrent models. In First Workshop on Scalable Optimization for Efficient and Adaptive Foundation Models, 2025.
Antonio Orvieto, Soham De, Caglar Gulcehre, Razvan Pascanu, and Samuel L Smith. Universality of linear recurrences followed by non-linear projections: finite-width guarantees and benefits of complex eigenvalues. In Proceedings of the 41st International Conference Machine Learning, pp. 38837-38863. 2024.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference machine learning, pp- 8748-8763. PMLR, 2021.
Shaoqing Ren, Kaiming He, Ross Girshick, and Jian Sun. Faster r-cnn: Towards real-time object detection with region proposal networks. Advances in neural information processing systems, 28, 2015.
13

---

Published as a conference paper at ICLR 2026
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). pp. 1715-1725, 2016.
Mugunthan Shandirasegaran, Yating Zhou, Songyang Zhang, and Shuai Zhang. Theoretical analysis of the selection mechanism in mamba: Training dynamics and generalization. NeurIPS 2025 Workshop WCTID, 2025.
wwtrrr...tttbb.. www.tt.thh..... io/ wttoo/dff//// pdffff//...
chandirasese
et al pdd
Jerome Sieber, Carmen Alonso, Alexandre Didier, Melanie Zeilinger, and Antonio Orvieto. Understanding the differences in foundation models: Attention, state space models, and recurrent neural networks. Advances in Neural Information Processing Systems, 37:134534-134566 2024.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1556, 2014.
Yonatan Slutzky, Yotam Alexander, Noam Razin, and Nadav Cohen. The implicit bias of structured state space models can be poisoned with clean labels. arXiv preprint arXiv:2410.10477 2024.
Jimmy TH Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations.
Jiawei Sun, Shuai Zhang, Hongkang Li, and Meng Wang. Contrastive learning with data mis- alignment: Feature purity, training dynamics and theoretical generalization guarantees. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025.
Hugo Touvron, Matthicu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and
Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning, pp. 10347-10357. PMLR, 2021.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 30, 2017.
Oriol Vinyals, Alexander Toshev, Samy Bengio, and Dumitru Erhan. Show and tell: Lessons learned from the 2015 mscoco image captioning challenge. IEEE transactions on pultern analysis and machine intelligence, 39(4):652-663, 2016
Chloe Wang, Oleksii Tsepa, Jun Ma, and Bo Wang. Graph-mamba: Towards long-range graph sequence modeling with selective state spaces. arXiv preprint arXiv:2402.00788 2024a.
Junxiong Wang, Tushaar Gangavarapu, Jing Nathan Yan, and Alexander M Rush. Mambabyte: Token-free selective state space model. arXiv preprint arXiv:2401.13660 2024b.
Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self-supervised contrastive learning. In International Conference on Muchine Learning, pp. 11112-11122.
PMLR, 2021.
Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. In International conference on machine learning, pp. 2048-2051 PMLR, 2015.
Sarthak Yadav and Zheng-Hup Tan. Audio mamba: Selective state spaces for self sule-supeeeieed audio representations. arXiv preprint mxXiv24466.02178 2024.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficie training. In Forty-first International Conference on Machine Learning.
14

---

Published as a conference paper at ICLR 2026
Annan Yu and N Benjamin Erichson. Block- biased mamba for long-range sequence processing. arXiv preprint arXiv:2505.09022 2025.
Shuai Zhang, Meng Wang, Pin-Yu Chen, Sijia Liu, Songtao Lu, and Miao Liu. Joint edge-model sparse learning is provably efficient for graph neural networks. The Eleventh International Con- ference on Learning Representations, 2023.
Bolei Zhou, Aditya Khosla, Agata Lapedriza, Aude Oliva, and Antonio Torralba. Learning deep features for discriminative localization. Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2921-2929, 2016.
Jian Zhou and Olga G Troyanskaya. Predicting effects of noncoding variants with deep learning- based sequence model. Nature methods, 12(10):931-934 2015.
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learning with bidirectional state space model. Proceed- ings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 62429 62442. PMLR, 2024.
15

---

Published as a conference paper at ICLR 2026
A NOTATIONS, PROOF SKETCH AND ADDITIONAL EXPERIMENTS
A.1 NOTATIONS
A.1.1 LUCKY NEURON DEFINITION
Let
== {i [m]: : > 0}.
K_ fie [m] : u, < 0}
(22)
denote the sets of neurons with positive and negative output layer weights, respectively. We define the sets of lucky neurons at initialization as:
W(0) {i E K++ Voul))00)04) 0},
(23)
u(0)= e e K_ (0) 0}.
(24)
where 0+ and o_ denote the class-relevant features for the positive and negative classes, respec- tively.
A.1.2 LOSS FUNCTION
The loss function for the nhth sample defined as
E( (n)) ((m)) = max {0, 2(1)_ (a))
1
(WoLt)ei (n)) }
(25)
max 0, - ((r)
vi
The empirical loss is denoted by C and given by
N c E (x(m))
(26)
The population loss is denoted by L and is defined
C Exx,))=D((X,
(27)
With additional important notations can be found in Table 2
A.2 PROOF SKETCH
The major idea of our proof is analyze how GD gradually aligns both the hidden-layer weights and gating vector with class-relevant features while ignoring the irrelevant ones. A key tool in our analysis is the notion of lucky neuron, i.e., hidden layer neuron whose initialization is well aligned with a class-relevant feature. For the majority-voting data model, the signal driving this alignment proportional the gap between the fractions of class-relevant and confusion tokens, Oaa, - ac), as established by Lemmas B. B Lucky neurons move consistently toward their class-relevant feature, while the magnitude of unlucky ones remains small (upper-bounded by the inverse square root of the number of samples). For the locality-structure the update in the relss-tel feature direction for the gating vector remains close to zero because an equal number of class-relevant and confusion tokens are present in the data. We then show that the gating vector consistently decreases along irrelevant feature directions, thereby enabling the gate to effectively select the class-relevant feature.
data model, we prove that
Due to these properties, the training dynamics can be simplified to show that the network output in (6) changes linearly with the iteration number f. In particular, we prove that, for new posi- tive sample (w.lo.g.)) during inference, the learned model's output is strictly positive. From this analysis, derive the sample complexity and the required number of iterations for achieving generalization error for both data types, as shown in (13) and (44) for the majority-voting setting in
Theorem I and similarly in (19) and 20) for the locality locality-struetured setting in Theorem
16

---

Published as a conference paper at ICLR 2026
Table 2: Summary of notations
F(X(m)
The final model output for (
Or
The average fractions of class-relevant tokens
0c
The average fractions of confusion tokens
AL
Separation between class-relevant features 0+ in positive samples
ALo
Separation between confusion features 0+ in negative samples
AL
Separation between class-relevant features o.. in negative samples
AL
Separation between confusion features o in positive samples
O
The set of class-relevant and class-irrelevant patterns
K
The set of lucky neurons with respect to W
K_
The set of lucky neurons with respect to (710)
N
The set training data
N+
The set of training data with positive labels
N_
The set of training data with negative labels
W(I)
Set of lucky neurons for the positive class iteration /
U(t)
Set of lucky neurons for the negative class at iteration
O+), 2(-), 9()
We use the standard convention: f(x) = (g((z) (resp. L(g(x)),
((g(z))) means fz) grows at most (resp. at least, on the order of) g(x).
0(-)
Soft- notation: hides polylog factors
poly(d)
An unspecified polynomial in d
2.
f(x) 2 g(x) (resp. f(x) S g(x)) abbreviates f(x) 2 [g(()) (resp. f(x) S (g(x))).
A.3 ADDITIONAL NUMERICAL EXPERIMENTS
Experiment settings.
The data dimension and token embedding size are both set to = 32, which also corresponds to the number of feature directions. Unless otherwise stated, experiments in the main text exactly the model defined in Eq. (6) to match our theoretical setting. We also use the model without convolution, and keep W B W c I frozen in Eq. (15). The total number of neurons in the hidden layer Wo is set to m = 50. For simplicity, we fix the ratio of different features to be the same across all data. The sequence length is
= 30.
We run 100 independent trials and consider only the successful trials to compute the mean epochs for convergence for given fraction of class-relevant patterns. An experiment is successful the testing loss is smaller than 10 10-3 3 For this experiment, we fixed the fraction of the confusion tokens 0.10 and varied the fraction of class-relevant features.
MLP weights selectively align with only class-relevant features. Fig. tracks the average cosine similarity between each neuron Wog.)) and both class-relevant & class-irrelevant features. The alignment increases for class-relevant features and stays essentially unchanged for irrelevant fea- tures, which is consistent with our findings Lemmas B. and B.3 in the Appendix.
Mamba outperforms Transformer and local attention on locality-structured data. Intuitively, locality-structun data favors models that exploit local biases. Global attention performs only marginally better than random guessing, whereas both local attention and Mamba learn meaningful patterns, with Mamba achieving the best performance.
Additional Results MLP Weight Alignment.
Figure illustrates the alignment of sampled neurons with the class-relevans feature. We observe that, with a good initialization, subset of neurons, denoted as lucky neurons, consistently increases in the direction of the claevatelevant feature, while another subset, denoted as unlucky neurons, remains almost unchanged, which supports our findings in Lemmas B.3 and
17

---

Published as a conference paper at ICLR 2026
Alignment of Woy. during training
0.6
Performance locality-structured data
Cosine similarity
Mamba
0.4
90
Transfomerr
Loca Attri
Test Test Accuracy
0.2
70
0.0
Class-relevant feature
60
-0.2
Class-irrelevant feature
0
50
100
150
200
250
0
10
20
30
40
Iteration
Epochs
Figure 5: Average alignment of (o)))) during training.
Figure 6: Mamba outperforms locality data.
In contrast, Figure shows the alignment of sampled neurons with the class-irrelevant feature. In this case, we observe that all neurons, both lucky and unlucky, remain nearly unchanged in the direction of the class-irrelevant feature, which further supports our findings Lemmas B.4 B.d and C.4
Alignment of Won. during training
Alignment of Wor. during training
0.30
0.8
neg[0]
neg[1
0.25
neq[3]
Cosine similarity
Cosine similarity
0.6
0.20
0.4
neg[o]
neg[a1]
neg[3] neg[8]
neg[18]
0.15
0.2
neg[23]
0.10
-02
0.05
0.00
100
150
200
250
50
100
150
200
250
Iteration
Iteration
Figure 7: Alignment of Wo(l.) with class- relevant feature directions during training on the majority-voting data.
Alignment of Worl.) with class- irrelevant feature directions during training on the majority-voting data.
Figure 8:
A.3.1
ADDITIONAL. EXPERIMENTS
To further strengthen the empirical connection between our theoretical analysis and practical Mamba architectures, we conducted additional experiments using the multi-layer, multi-head Mamba model Dao & Gu (2024). trained on synthetic datasets that follow the same structured data models as in our theory.
We first evaluated the Mamba2 block, which includes residual connections and RMSNorm. We focused 2-block Mamba model with 4 heads and report the cosine similarity of the learned gating vectors and MLP weights with relevantl and class-irrelewe features BnguresPJandJii For deeper 5-block Mamba model with the same configuration, summarize the final alignment values Table which exhibit the same qualitative trends predicted by our analysis.
Table 3: Cosine similarity alignment in the 5-block Mamba model
Component
Class-relevant
Class-irrelevant
Gating vector
0.53
0.00
MLP weights
0.73
0.00
18

---

Published as a conference paper at ICLR 2026
Alignment WA during training
Alignment Wou. .1 during training
0.6
0.4
Cosine similarity
Cosine similarity
0.3
0.4
Class-relevant feature Class-ireeleant feature
0.2
Class-relevant feature
0.3
Class-irrelevant feature
0.2
0.1
0.1
0.0
0.0
0
10
20
30
40
10
20
30
40
Iteration
Iteration
Figure 9: Alignment of the gating vector in the 2-block Mamba model.
Figure 10: Alignment of the MI.P weights the 2-block Mamba model.
Next, we examined the effect of the gating mechanism by comparing models trained with and with- out gating across both structured data regimes. On the majority-voting data, the gated model consis- tently outperforms the ungated variant (Figure IT). On the locality-structe data, gating becomes essential: the ungated model fails to learn the task, whereas the gated model converges reliably (Figure
Performance majority voting data
Performance locality-structured data
100
90
Test Accuracy (96
Test Test Accuracy (96)
80
Model with gattng
70
Model without gating
60
60
Model with gating Model without gating
50
8 Epochs
10
12
14
0
10
20
30
40
Epochs
Figure 11: Test accuracy with and without gat- ing on the majority-voting data.
Figure 12: Test accuracy with and without gat- ing on the locality-structured data.
We also conducted two controlled ablations. First, we varied the feature dimension d E (32, 64, 128) and observed that the qualitative behavior of the model remained consistent across all three settings (Figures 13) Second, we varied the data distribution parameter de, the fraction of confusion tokens the majority-voting data. Across all three choices of ace the empirical results remained closely aligned with the theoretical predictions (Figures 16 18).
Escori ta Reach 59% accuacy
Converpence
Convergence
Convergence
Esccoe Roach 55%
Reach
0.30
0.40
0.45
U.5
0.20
0.55
0.A0
0.30
0.55
0.40
0.43
0.30
Fraction class relevant patterss
Fraction of class relevant, pazterss
Fraction of class relevant parterms
Figure 13: Ablation with fea- ture dimension 32.
Figure 14: Ablation with fea- ture dimension d = 64.
Figure 15: Ablation with fea- ture dimension d 128.
19

---

Published as a conference paper at ICLR 2026
Epachs to Raach 90% Accuracy
Mamba
N-
N--50
ki-00
100
k-000
k00
99%
200
10
Epochs to Reach
1.0
0.50
0.35
0.40
0:45
0.50
0.30
0.35
0.40
0.45
0.50
0..30
0.35
0.40
0.45
0.50
Frection of cass relevant potterms
Fraction of class relevant patterns
Fraction of dass relevant potterns
Figure 16: Ablation with con- fusion fraction a₂ = 0.17.
Figure 17: Ablation with con- fusion fraction ag 0.20
Figure 18: Ablation with fusion fraction a 0.23.
B MAJORITY-VOTING DATA
B.1 USEFUL LEMMAS
Lemma B.I provides bounds on the gradient updates of lucky neurons i E W(t) in the directions of both class-relevant features (0+, O-) and irrelevant features.
Lemma B.1. Suppose p1 S (wo²² o+) S q and p1 I (w/a 0-) q Then, for any lucky neuron i E W(t) at iteration t, the following bounds hold:
(L1.1) A lower bound on the gradient of with respect to Wo(L.) at iteration t, in the direction of 0., is given by
ac
1
(V
N mN OT)..
ow 06.))
0+
2
VmL ( -ao))- -0
(28)
(L1.2) An upper bound on the gradient of & with respect to Wols.) at iteration in the direction of
0+, is given by
ac
(Volog ( mN
N
0+
Vml. (a +0
+(()). (29)
ow O(..)
(2)
(L1.3) A lower bound on the gradient of with respect to Wo(L.) at iteration t, in the direction of
O_, is given by
0c
(dllog N V
O
2
I.
O(T).
(30)
(t)
/mL
O(..))
(L1.4) An upper bound the gradient of 2 with respect to Wo(s.) at iteration in the direction of O_, is given by
owl)) (0) 0- .
ac
ddog N Vallog
mN + O
(31)
o(..)
(L1.5) An upper bound on the gradient of with respect to Wo(i.) at iteration t, in the direction of Oj, given by
Oj 50 < dlog N + O(T), for + 1,2.
V mN
(32)
(w
O(.,)
Lemma B.2 shows that, for unlucky neurons associated with the positive class, the gradients in the directions of both class-relevant and irrelevant features small.
20

---

Published as a conference paper at ICLR 2026
Lemma B.2. For any unlucky neuron E K+ \ W(t at iteration the following bounds hold:
(L..1) An upper bound the gradient of & with respect to W Wots)) at iteration in the direction of
0+,, is given by
OC
10
V WNN
dog N mN
0+
+0(T).
(33)
ow O(.))
(L2.2) An upper bound the gradient of & with respect to W Wots.) at iteration t, in the direction of o_, is given by
ac
VV
a N
o
+
+O(T).
(34)
ow
(0)
O(..)
(L2.3) An upper bound the gradient of & with respect to Wo(s.) at iteration L, in the direction of jj is given by
ac
20(0v dllg N mN + O(T,, for + 1,2.
Oj 50
(35)
(1) O(,,))
B.3 and B.4 by symmetry, state the analogous results for lucky and unlucky neurons associated with the negative class.
Lemmas
Lemma B.3. Suppose pi S
91 and p1 S (w o_) S qi- Then, for any lucky
neuron i (I) at iteration I, the following bounds hold:
(L3.1) A lower bound on the gradient of & with respect to Wo(L.) iteration in the direction of O_, is given by
ac
(pi)0((a,), act.) - /mL
1
(V N mN
O(T).
(36)
Ow
(
O(.))
(L3.2) An upper bound on the gradient of & with respect to Wo(L.) at iteration t, in the direction of O-, is given by
ac
(Volog N mN + +O()).
VmL. d. ac
(37)
(W O6.)
(1)
(L3.3) A lower bound on the gradient of E with respect to WoG,) at iteration the direction of 0+, is given by
OC
1
(0 N mN
0+ 2
a (gi)) -ood)
-cl
O(T).
(38)
ow o(L.)
VmL.
(L3.4) An upper bound on the gradient of E with respect to Wo(L.) at iteration t, in the direction of 0+, is given by
ac
c dlog mN
0+ so
+ +
(39)
(w O0.)
(L3.5) An upper bound the gradient of & with respect to Wo(s.) at iteration in the direction of
Oj. is given by
ac aw(+) O(..))
<
V
/log N
.0) 50
O(T), for j # 1,2.
(40)
mN
21

---

Published as a conference paper at ICLR 2026
Lemma B.4. For any unlucky neuron E K- (u(t) at iteration the following bounds hold:
(L..1) An upper bound the gradient of & with respect to Wots.) at iteration in the direction of O-,, is given by
so(( dlog N mN
V
owio O(i.)
o
+ +
(41)
(I.4.2) An upper bound on the gradient of & with respect to Wo(L.) at iteration 1, in the direction of O+, is given by
ac
50
(v dooog
+ +
(42)
oW
m.
O(a.))
(L4.3) An upper bound on the gradient of with respect to Woti.) at iteration t, in the direction of oj given by
ac (t) Ow O(.,)) oj 50
dlog N mN (
0 0
O(T,, for + 1,
(43)
Lemma B.5 establishes bounds for the gradient updates of wA in the class-relevant feature direc- tions.
( ,04) S st. Ler (W()) ptt and [u((t) Pi Then, at
Lemma B.5. Suppose ri S (Wo)))
o(..)
iteration t, the following bounds hold:
(L5.1) lower bound on the gradient of E with respect to ws at iteration t, in the direction of 0+, is given by
ac 0+
ri pt O(aL))
ymsi
(VV N mN
2
((a -0
O(T). (44)
Ow
(t)
2/mL
4L
(L5.2) A lower bound on the gradient of & with respect to wA at iteration t, in the direction of o given by
Øc
2/mL ((a,L)- - Vmst
-o |dlo (Valos mN
O(a
O(()) (45)
Üru! (t):
4L
Lemma estabbiiss bounds for the gradient updates of wA in the directions of irrelevant features.
Lemma B.6. An upper bound on the gradient of é with respect to ws iteration I, in the direction oj. oj. is given by
Oc
d N mN
Ow 1t)
Oj
so
O(T), for j + 1,2.
(46)
B.2 PROOF OF CONVERGENCE
The proof starts with the base at 0 and proceeds analyze the training dynamics in a deductive manner, providing additional details in deriving the corresponding conver- gence and sample complexity bounds.
Proof of Theorem 7
(S1) Warm-up (Base case): Training dynamics at the first iteration / = 0.
(0) Recall that we set w 2
0. Then, we have
(w A ... = and
(0)
(ur a o
(0)
22

---

Published as a conference paper at ICLR 2026
(S1.1) Training dynamics of (o( at the first iteration 0.
From Lemma B.1 identify p1 0 and gi Let ar and de denote the average fraction of label-relevant tokens and confusion tokens, respectively. Then, for any lucky neuron i E W(0), we obtain
2/mL (
0(a,L acL) -
1 (pols(a)
0c
ow.0) (0)
0+
((1,)
(47)
(podylat)
( .a +
2/mL
poll(d))
ac
0 ( for j + 1.
and
0,
(48)
ow ((+.))
Recall that set the number of samples in a batch N poly(d).
Recall that the initialization is
W W (0) = 8104 + 220 + .... sorg
6j ltd NN(0.²2) 1,2,2 ... ,.,
(49)
Then, after one gradient descent step, we have
81 Bla,L acL) - (polydd) (p(plyaa)
2/mL
2 Woll) ,0+)
(1)
7 - poly(d)
(
S 01 +
(50)
2/ml
and
W Won.))
03) 2
1
for j # 1.
(51)
poly(d)
By applying Lemma
B.3 for any lucky neuron i U(0), obtain
82
2 2/mL - a.T. -
poly(d)
W WoR)) .00))
(1)
< 62 +
"
Oa,l -l,,L-aa. acl) - +
(52)
2/m
poly(d)
and
( O(
(1) .o,) 2ō (polycc)) for j + 2.
(53)
For any unlucky neuron i E K.K. \U(0), Lemma B.4 gives
[ Wow) (1) .0) σ 1
for Vj.
(54)
poly(d))
(S1.2) Training dynamics of ur the first iteration / 0.
23

---

Published as a conference paper at ICLR 2026
Now consider the gradient update for ws- Define:
a = 81 +
7
( - acc)- ppol(dd)
2/mL
" ( - acL) +0
1
+
2/m
poly(d)
Applying Lemma B.5
with rf a, si = b. and Po |W(()), we get
Oc
a
Po ((a - Vmb
()a.L) - - (ppoliaa) =: =:a (55)
iws
(0)*
0+
2
2 m/mL
4L
Let 61
(oly(d)) Since b= Ö
(poly(a)) ((ool)aan) that is sufficiently small,
Vmb
1 e(acL) 3 poly(a))
1
O(a,L)
a
2L
Vm
2
Vma
(0(a, - -o(a -6
2L
2
poly(d)
Vm
T
O((a,L - -aol)) poly(d)
4T,
2/ml.
"
@(a, - aoL) -(arL-aaL))) -
=0
(56)
SL2
poly(d)
From Lemma B.6
we also obtain
ac Ow [o) oj o (p(ottaa) =:y for j + 1,2.
(57)
A
(S2) Induction Step:
Training dynamics at a general iteration
Suppose (ur (t) 0+) = a² 2a.1, (w) . = 3" 2 B 1, and (un ( ,,) = n Sy.5. where
8 a'
Vmb'
()a Ö 1
Po O(a,L)
(pllydd) >0. 0,
(58)
2 /L
4T.
d' = 82 +
n
O(a, acL) - - poly(d))
(59)
2LmL
6' 89 +
" Oa,,L a.L) +
1
(60)
mLL
poly(d)
Following the same approach as in
56 we can simplify and obtain
a.L)3) 0((aLL-aL)))
O((aL
poly(d) 0.
(61)
SL2
For any lucky neuron E W(t) at the (t + 1)-th iteration, we have mLL o(aa (o(a*)0((a (0 acL) - poly(d)
OW()) (1) ,0
O(L.)
a (aa))((a,L... - a.T.) + poly(d))
VmL
24

---

Published as a conference paper at ICLR 2026
and
0C
00
0 poly(d)
for j +
(62)
@w
(1) O(i,)
Next, we have o(a)) > 1 since a* > 0 when t = 1. By simple induction, this further ensures
ac
ac
ac
ac
ow O(C,)
0+
ow (1) O(.,)
S ...
awlf
(t) O(t,)
0+ $
ow (++1) (++1) O(2,)
(63)
Thus, we obtain the following bound after the second gradient descent step:
2/mL O I a.t) [1 20(a")) ( (poly(aa)
"
81
(W))))) 0+)
81 2/ml.
- [1 (1 (a)))
+
(porl(aa))
ppoly(a))
(64)
Similarly, applying Lemma
B.3 to any lucky neuron i U(1) at iteration 2, we get
o(8")00(aa adl) - )
1
MLL
poly(d)
ac
.0_
(65)
ow O(i-)
acl)+ 03 - +
(poty(a))
VLL
ac
poly(d) ( for j # 2.
and
(1()) , o0 5 O(1.)
(66)
Applying Lemma
B.5 with ri = and st = we obtain
Vnv
-0 = X.
Ow (1)
0+
2 mL pt O(a,L)
4L
(67)
Since u v δ
that is sufficiently small, we have X X20 0.
poly(d)
By applying Lemma B.6
get
ac Ow (0) 0 < A
1 poly(d)
for j * 1,2.
(68)
(S3) Induction conclusion: Training dynamics when the algorithm ends.
We proceed by induction on t: the base case t = is established in (S1), and the induction step for general is shown in (S2). For any lucky neuron E W(T), we obtain
(T) Wo .0 aT,
(69)
and (Woie)) (1) ",,)) s 2 1
for j #
(70)
(dly
25

---

Published as a conference paper at ICLR 2026
For any lucky neuron i (TT), obtain
(Wol.)) (2) ,0- 2 aT;
(71)
and ( (Wow.) (T)
50 (T,o)
1 poly(d)
.0j)
for j #2
(72)
Also, we obtain
w .0 2 aT,
(T)
(73)
(T) w .-)) 2 ST,
(74)
(T) (w _ .s) S ST..
and
(75)
(S4) Derivation for the generalization bound.
We will demonstrate that once the weights have converged at iteration T, the model accurately captures the underlying data distribution, which leads to zero generalization error, shown in (44)
Consider (()) +1 as an example. The sequence (n)
(n)
(x)
(n) L
has first a
22
tokens correspond to the feature 0+, while the following L tokens correspond to the feature o...
F(X(1) (n))
Ui Φ (ool,)l))))
T.
I=1
(Wonsel"))
VmL. Σ
① (Wolssl())
/mL
i
iEK
1=1
Σ
Σ
(Won.ee"))
Σ
(Wor.)u)")
/LL.
VmL.
(EW(O) 1=1
(EM((0) 1=1
Σ
Σ Φ (WoLs)b())
(n)
(76)
VmL
EEK (U(0)
The Mamba output y (n) is defined as
yi (n) Σ II
(1-o(w]af)))
o(wIx!") (Ic"') x )o("))
(77)
j=++1
We now derive lower bound for
Σ (EW
Σ 1=1
s0 (Wo..))))
To that end, consider the aggregated projection
Σ W Wo
Σ (EW(D)
ΣΣ 1=1
w Woo.). 0,) 0,).
(78)
(EW
26

---

Published as a conference paper at ICLR 2026
For any E W(0), we know that
Wo 2 aT.
(79)
Hence, let's obtain a lower bound for (y:, 0+).
We only need to consider the cases where 0+ for some in the range 1 s After I' iterations, we know
(wa.0)) aT, (wa.0-) ST, (ws.o,) ST for j + 1,2.
(80)
Therefore, we have
0+) O(a =x)
for 1,2, ...., a,L.
(81)
We now lower bound the objective
Σ
(oo.))u)
(EW(O) 1=1
Note that
Wo(,)) = Σ
W (Wol.)))) (ULt.os)
i=1
and yr has only 04 component for = 1,2,...,,,,. Therefore,
Wo(,) W (oo..)) o+) (y). 0-) 2 aT O((((TT) 0, for = 1,2...... ... aLL
Applying d(z) for positive z, obtain
6 (Woi,)) 2 aT ((a(aT)), for 1,2
Hence,
Σ Σ de (Won.)m)) >
Σ aT O(o(aT)) a.L (EW(O)
(82)
iEW(O) 1m1
Next, we derive an upper bound for
Σ Σ
Wonn)yi
(EU(() 1=1
For any E U(0), we know that
0 ( WOlL),. o-) S bT.
(83)
We now derive an upper bound for o_). We only need to consider the cases where x, = such that < I.
0_
We have,
(0 S
where
W 7 (((aaa--a))))) a.L)) +
(
1
(84)
812
polyydd)
(11 -0a =()
for =1,2,.... 1,2,...,acl TcT.
(85)
27

---

Published as a conference paper at ICLR 2026
Wo br 50 ((o(TT) a.L.
(86)
Σ (0(((0 1=1
E
Woonnv (n) 1
Σ (EU(0)
br ((o(())WT)) acL.
(87)
In addition, we have
E Σ (WOL))ul (")) s0 (00')ssppl)a) (poolyaa)
(88)
iEK (4(0)
By 76..
we can write
((
m 2
aT ((a(aT)) a,L bT -(o(WT)(TT a.L δ
(pml) ) (89)
VmL
with
" O(a a,L) -
a
2/mL
(poly(d)
(90)
and b "
2 mL
0(a,L acL) +ō
poll(d))
(91)
a O((a,L 9((aLLaaoL)) aoL)))- -(cr 1
(
polyy(a)
812
-W - 1
poly(d))
(92)
Therefore, we conclude that
(x)()))
m
aT ((o(aT)) (ay, - -ac) ar)- (d)ly(d)
(ppoty(a))
(93)
2
There, for any positive sample, prove that
((X())) C, where C' is some positive constant.
(94)
Similar to the previous analysis, one can show that the negtive sample X,, leads to
(S4.1) Derivation for the convergence rate. Let's find the number of iterations T required such that F(X (*)) 2 1, since the label is +1. We require
mm
aT aae) 2 1+
(95)
2
Substituting the value of a 2 b O L - L,LL)), the condition becomes
2/ml
VmaT
(ar - are) = Vm
0 -. acT.)T (a, - ac)
2
2 2/mL
T
O((a, ac)3) 2 l+c.
(96)
Solving for T, obtain
4(1 (1)
T > 100(a, ac)2) 2 100( - 2)
(97)
28

---

Published as a conference paper at ICLR 2026
Now, additionally require that the sigmoid activation o(aT) be sufficiently large, i.e.
aaTT) 2 1 -.
(98)
When 2 is sufficiently large we can approximate
(() 1 2 -
Substituting aT, condition
98
becomes:
o(aT) 1--e-- c, ea
c aa S,
aT - In(c)
T In(c)
(99)
Substituting a 37- ((a, - ())) (0,
we get:
812
T2 10--nn6 -In(c)
(100)
8
T2 In(e)
(101)
10((a,,-))) -0o)
Hence, by combining 977 and 101 we obtain
8 n ac) In(c) no((a, -a))))
a
T 2 max
(102)
By combining 995 and (88) with the expression for the model output (x)))) in 93p. we obtain
F(x(1)) 2 (1+t).(11))
21 21-0(e2)
(103)
Hence, for sufficiently small the model output satisfies F(X()) 1.
Similarly, for a negative sample, one can show by symmetry that the model output satisfies F(X(m)) 1.
(S4.2) Derivation for the sample complexity.
Now we derive sample-complexin bound that guarantees zero generalization error. Assuming enough samples, write for sufficiently small À < 1
(V dog N mN (Vales
o
< X
7 Oo, L..))
(104)
2/mL
From this, we can derive a lower bound on the required sample size,
4L2d
N 2 Ω -
(0( 2)
(105)
12d
2 Ω n20 (ar 22
which will be (1) in Theorem
29

---

Published as a conference paper at ICLR 2026
C LOCALITY-STRUCTURED DATA
C.1 USEFUL LEMMAS
Lemma C.I provides bounds on the gradient updates of lucky neurons i W(t) in the directions of both class-relevant features (0+, O-) and irrelevant features.
Lemma C.1. Suppose p1 < (w)') 04) S Q1, P1 S (1 (w 0-) Qı, and p2 S (ww) ( j * 1, Then, for any lucky neuron E W(t) iteration t, the following bounds hold:
o,)) S qa for
(LI.I) A lower bound on the gradient of & with respect to Wotis) iteration in the direction of O4, is given by
ac Ow
opp) (1 (1-o(91))²
o(91))
(1-a(1)) ((22) a(1)) -6 -2 (1- (1-((p2))) A
(0)
0+
/mL
O(s.))
(106)
-0
dlog mN
(L1.2) An upper bound on the gradient of E with respect to W Wo(s.) at iteration t, in the direction of 0+, is given by
oc
((1))- (1 (1--(p))) o(pi))?
[a- (1-r(p2)) 2rd -2 =-11
ald2))
-(1 (1- (1-o(12))
dw 06.))
(0)
o+
VmL
(107)
dlog N mN
+ o
(L1.3) A lower bound on the gradient of 2 with respect to W Wot.. at iteration t, in the direction of O_, is given by
ac
o(q)) (1 op:))) (p-) Al. -2 -26 (1 ((=)) AL -2]
aw (1)
o
) VmL
O(s.))
(o dlog N mN
(108)
(L1.4) An upper bound on the gradient of & with respect to W Wo(s.) at iteration in the direction of O-, is given by
ac
(
Vallog d
o
5
(109)
Ow OlL))
mN
(L..5) An upper bound on the gradient of & with respect to WorL.) at iteration I, in the direction of Oj given by
ac O (0) 0j O(..)
(Viloge N mN
00
for j * 1,2.
(110)
Lemma C.2 shows that, for unlucky neurons associated with the positive class, the gradients in the directions of both class-relevant and irrelevant features are small.
Lemma C.2. For any unlucky neuron E Kc+ W(t) at iteration t, the following bounds hold:
(I.2.1) An upper bound on the gradient of & with respect to Wotl.) at iteration i, in the direction of 0+, is given by
ac
(v
dog N m.N
@w (3)
o+
00
(111)
O(..)
30

---

Published as a conference paper at ICLR 2026
(L2.2) An upper bound the gradient of & with respect W Woti.) at iteration t, in the direction of o_, is given by
ac
( ddog
<
00
(112)
OW O
mN
(L2.3) An upper bound on the gradient of È with respect to Wo(s.) at iteration t, in the direction of Oj given
(v d N m.N
(w (f) so
for j + 1,2.
(113)
O(..)
Lemmas C.3 and C.4 by symmetry, state the analogous results for lucky and unlucky neurons associated with the negative class.
Lemma C.3. Suppose P1 I (w)) o-) S Q1 P1 S (w o+) S 91 and S (w o,) S 92 for j * 1,2. Then, for any lucky neuron E u(t) at iteration t. the following bounds hold:
(L3.1) lower bound on the gradient of & with respect to Wo(L.) at iteration the direction of is given by
ac OW ((..)
opp)). .(1 - o(g))) (1 (1-a(()) o(22) -/ - -(1-o(p2)) A -2]
VmL
(114)
0
dlog.N m.
(L3.2) An upper bound on the gradient of E with respect to WorL.) at iteration 1, in the direction of O_, is given by
ac
(1-) o(pa)) -2 - -(1 -11- -(1-o(gz)) 2 -2] 0(11)- - (1 ((1-)p))))
o(p)))²
Ow O(.))
mLL
(115)
dlog mN
+ O
(L3.3) A lower bound on the gradient of with respect to Wo(i.) at iteration t, in the direction of 44, is given by
(1 (1-0(pz)) (p2))) art -2 - -10
( (1 {pp)))
0W() o(.))
0+
VmL.
(o( ddlog
- 0
mN
(116)
(I3.4) An upper bound on the gradient of E with respect to WorL.) at iteration 1, in the direction of 0+, is given by
ac Ow!!) O(s.)
dlog N mN
0+ 50
(117)
(L3.5) An upper bound the gradient of E with respect to W Wois.) at iteration i, in the direction of Oj. is given by
ac
(V N
( (0)
oj
50
mN )
for j # 1,2.
(118)
O(..)
31

---

Published as a conference paper at ICLR 2026
Lemma C.4. For any unlucky neuron i E K- (U(t) at iteration the following bounds hold:
(L4.1) An upper bound the gradient of & with respect W Wots.) at iteration in the direction of O-,, is given by
ac ow O(i.)
so(v
dd dog N mN
o 50
(119)
(L4.2) An upper bound on the gradient of & with respect to Wo(s.) at iteration in the direction of O+, is given by
ac
{ /dllog
d
0+
00
(120)
ow O(i.))
mN
(L4.3) An upper bound on the gradient of & with respect to Wo(s.) at iteration I, in the direction of jj is given by
ac
<V
Valoog N m. ) for + 1,2.
(0) .0
(121)
OW O(t.)
Lemma C.5 establishes bounds for the gradient updates of wA in the class-relevant feature direc- tions.
Lemma C.5. Suppose Pı 0+) S 9i and rj s (W(++)) and [u(t) = Then, we have:
O(1.) .0+) S sj. Let [W(() pt
ac Ow [t)
o(r)) (1- ogi)) -
pt o(a)) (1 (1-o vmm o (v dlog N
( mN
0+
(122)
V m
2
Suppose P1 S (w)) o_) S 41 and ri (oo) _-)) S si. Let (W(t)] pt and [U((t) Pr.
(t)
(t+++))
Then, we have:
0c Ow (f)
(1- or a(gi) opp))) si V
(Valog
dl N mN
(123)
2
Lemma C.6 establissss bounds for the gradient updates of wA in the directions of irrelevant features.
Lemma C.6. Suppose P1 I (w 04) S 91 Pı I (w o_) qı, (wo²) oo) 92 # 1,2,
and ri S (W!l) (Wol.)) O(..) .0+). Let p IW(1)) and Pt = [(())) Then we have:
02 Ow (1) 0j
ri (pp)) (1 x(9))) (1 (1- (1-n))))))( ((p))
((2))) Aio pt + +1 +1-o((12)) Ali pi
2/m
(124)
(1- o (1-- +0
(
(V
mN
C.2 PROOF OF CONVERGENCE
Similar the proof of Theorem I the proof starts with the base case at and proceeds to analyze the training dynamics in deductive manner, providing additional details in deriving the corresponding convergence and sample complexity bounds.
Proof of Theorem 2
(S1) Warm-up (Base case): Training dynamics at the first iteration / = O.
Recall that we set 3
0. Then, we have
A o_)
(0) 0,
w A o 0, and (w .0,) Vj.
(0)
32

---

Published as a conference paper at ICLR 2026
(S1.1) Training dynamics of (i...)) at the first iteration 0.
From Lemma C.1 identify = 0, qı = 0. P2 = 0 and Ø2 = 0. Let A and AL be the distance between two 0+ and o- features respectively in the positive sample. Similarly, in a negative sample, let the distance between the two o+ tokens AL. and the distance between the two o_ tokens as A Then, for any lucky neuron E W(O), we obtain
c
ALo
-2
AL -2
Ō poly(d)
2/mL
2
2
Ow
(0)
0+
O(1,)
ALo
AL
S
+ + polyydd)
(125)
2//L
and
ac (0) aw ((h.)
oj I Ō
(pllyaa))
1
for j + 1.
(126)
Recall that we set the number of samples in a batch N = poly(d).
Suppose the initialization is
Wo(L, (0) = 010 + ò20 + +... + OpOa,
5j 44 52
j 11
(127)
Then, after one gradient descent step, we have
nd²
1
A
22
Aloo
61 +
2 mL
2
poly(d)
WT. Wot.))
,0+)
ALd
1
Akat
$ 81
+ polly(d))
(128)
2/mL
2
2
and
W (WoG.)
(1) s (,)
(
1
for j # 1.
(129)
poly(d))
By applying Lemma
C.3 for any lucky neuron i U(0)), we obtain
AL5_
AL
82 +
2 2L
2
2
poly(d)
W Wo
.04
( 1
Ar.
1
A
$ 82
+
(plly(a))
(130)
2 mL
2
2
and
(Wo
(1) ,o) < 2
(
poly(d))
for j + 2.
(131)
33

---

Published as a conference paper at ICLR 2026
For any unlucky neuron E K- \u()), Lemma C.4 gives
(1) Wou. 00 2 poly(d))
(poly(a))
1
for Vj.
(132)
(S1.2) Training dynamics of wA at the first iteration / t= O.
Now consider the gradient update for was Define:
ALd, -2
ALo
-2
(potytat)
51 2 2/mL
2
ALt
ALo
b b= 61 +
1
1
+ Ō
2 mL
2
2
poly(d)
Applying Lemma C.5 with p1 = 0, ạ1 = 0, r = a, o = b, and P [W(O]. we get
ac Ow ()) A
bym
m
0+
Pu
(oll((a)
(133)
4ym m
8
We can relax this lower bound and obtain
d' 2a
Ow (0)
o+
2
4
mm
Pt
vmb
5 poly(d) =:
(134)
Recall that 81 poly(d)
Since a
b
①
that is sufficiently small,
poly(d)
Vmb - δ ( (pol(dd) m
(potsica)
a =
4 m 2
[Vma- Vma] - poly(d)
4
0 δ (poly(d)
-
= 0
(135)
poly(d)
From Lemma C.6
we also obtain
ac (0) oj
Aut
1
A1.
Po
+
Po
(136)
Ow A
8mmm
2
where we apply the lemma with the values
Pi 0, q1 0
92
0, and
r
We can relax this upper bound and obtain
ac
-ac' <
ALt
11
Al.
Ow
(o)
4 4/m
2
po
+
2
Po
=:7
(137)
34

---

Published as a conference paper at ICLR 2026
Taking Po- = Pu = m
+ Ö
we can simplify and write
2
puly[a)
ac
m
AL
AL._
=
4 m
2
2
2
d
1
AL5_
--ma-
+
8
2
2
(138)
AL
-2
1
Sba
-2
[
Ab
1
16L.
2
2
poly(d)
(S2) Induction Step: Training dynamics at general iteration
Let ( 0+) a* 2 a. (w)) o-)
3r 2 B t, and (w0²) oj) = where
y
B = c' 2a'
4
V
Po
vmb
5
poly(d)
> 0
(139)
nc²
ALo
ALt.
a' = 82 +
2 2/mL
②
poly(d)
AL5_
AL
b' = 62 + 7
+
2 2/mL
2
poly(d)
Following the same approach as in (135). we can simplify and obtain
== == δ
(140)
poly(d)
For any lucky neuron E W(1) at the (1 1)-th iteration, we have
d²
a(a") [1 -7(")) Ard -2 - (1 - -r()))) A 9] - δ (polltaa))
VmL.
ac
(! (1) ,0
((1-))
(141)
aaa") . (1 - = (103))2(110000 (1-0))))) -2_
_-1 (1 - --((7))) ariz
- -3]
v mLL
+ 5
poly(d))
ac aw O(..)
00 2 (pollycc) for j #
(142)
Note that, aaa)) > 1/2 and o(>*) < 1
Thus, we obtain the following bound after the second gradient descent step.
(Wil).)) , ,,o))
281 n scb, Vml.
aL.
2
(143)
=a(a')) (1 (1- 2 -2 (1
o) (1-o(r"))
-0( [ - -ō (poit)a))
35

---

Published as a conference paper at ICLR 2026
and
(Wo "6, T o.)
1
(1 2
A42
<h1 + 2/L
2
(144)
a(o')) (11 (1-0 - -2 -1 a(>"(()))))))) a -4z - :) δ (poty (d)
+2((a)) (1
-1
For the convenience of presentation, use denote the lower bound in [43 and denote the upper bound in 144 Similarly, applying Lemma
C.3 any lucky neuron i (11) iteration get
ac (3) o(.))
c²
a(8)) (1 -((()))) -2 -1 (1-o(1"))ard_- -2
Ow()
o-
/ml.
(145)
- =)(ot) u
polyydd)
VLL .a(3)) . (1 -33)))
(1 (1-0()1)) -2 2-1 (a-o(r')))
And -2
aw(1) O(1.)
+ poly(d)
(146)
and
σc aw(1) O(..)
00 2 pollya))) )
for j +
(147)
(1)
Applying Lemma
with P1 a²" 91
ax ri = u, and s v, we obtain
dw (I)
ac
0+
o(ax)c'
2
m
2u
pt
me δ
poly(d)
(148)
Since pt
and δ
(poly(a))
that is sufficiently small, we have x 0.
By applying Lemma C.6 with
P1 (= B*),
ai a= 8*), === and ri u, we have
ac
d'u oaa") (1 - (1-ode o)("))) St Pt + (1 - (-) A6. pē] =:6 (149)
dw (1) <
2/m
A
Note that here we assumed the distribution of AL+ is identical to AL to have a* = 8
(S3) Induction conclusion: Training dynamics when the algorithm ends.
We proceed by induction on t: the base case t = 0 is established in (S1), and the induction step for general is shown in (S2). For, any lucky neuron i W(T), we obtain
(T) .0+) aT (oil))
(150)
W WoOl) (1)
5 5 1
.0
poly(d)
for j #
(151)
36

---

Published as a conference paper at ICLR 2026
For any lucky neuron i (TT), obtain
(3) (ol)) ;00) 2 at
(152)
W o(t,) .0 ) s (poly(a))
(T)
for j #
(153)
Also, we obtain
(T) w .++ 2 aT,
(154)
(2) ws . 2 ST,
(155)
and
( ,o) 0j) S TT.
(156)
(S4) Derivation for the generalization bound.
We will demonstrate that once the weights have converged at iteration T,, the model accurately captures the underlying data distribution, which leads to zero generalization error, as shown in (180).
Consider z(n) +1 an example. The sequence X (m) 04 at Li and L and two o at Li and Lz.
22
x (n)
contains two
m
F(X(1)
Σ
(oon))l))) (m)
8/
L
1=1
(Wonnui")
VML. Σ
VLL.
1=1
EEK-
1=1
L
ML Σ
Σ
(Wonneri))
Σ (WoL)))())
VmL.
iEW(O) 1=1
tEN((0)
Σ
Σ
Φ (WoLsDi))) (n)
(157)
VmL
IEK VU(0)
1=1
The Mamba output (n) defined as
yi
(n)
II (1- (I-o(wZo())) j=8+1
1
a(wIx(")
(x!") (n) aa!))
x
(158)
We now derive lower bound for
Σ Σ o (Wou,))) (EW(O)
To that end, consider the aggregated projection
Σ
Σ
W Wo
Σ (EW(D)
ΣΣ j=
w Wou.). 0,) 0,).
(159)
(EW
37

---

Published as a conference paper at ICLR 2026
For any E W(0), we know that
(Wo (Wo))))))) 04) 2 aT.
(160)
Hence, let's obtain a lower bound for 30,
We only need to consider the cases where æ, 0. for some 8 in the range 1 S 8 S l. In particular, we will focus on the following instances:
1i and I E 1 L }
s = Li and / = L
After T iterations, know
(wa. aT, (wa, 0-) (,) STT for j + 1,2.
(161)
Therefore,
(1 =0 04))
(162)
We have,
(ws, 0+55++s550 WIT,
s
(Wa, 0-) S W2T,
where
W1 δ poll(d))
(163)
W2 δ poly(d)
(164)
Then we obtain the following:
2 a(aT) - (1 OLCC0-------.00000-0000000----00
-2 o(aT)
- ((
Alt
(165)
We now lower bound the objective
Σ
Σ
o (Woo.))
(EW(O) E=
We begin with
Σ (EW(O) t=1
D W (on)))))
Σ
[o(WO.))LL) +&Woo,))L))
(EW(O)
Note that
d Woo.) = Σ
W (o) (U)tt..))))
and V2 has only 04 component.
Therefore,
WoR.)y1, (o ( at o(a)) >0.
Similarly, write
WouL; 2 aT aT o(aT) -
-2) >0.
Applying ((z) = for positive 2, we obtain
(Woo.))1)) aT a(TT),
38

---

Published as a conference paper at ICLR 2026
(Woun)LL) aT (6 ((ws,, SK -
Hence,
Σ
o (io(.))
Σ aT' o(al)) (EW(O)
(ew(O)
(166)
+ (1 - (1 ((ws,, o,))) Ariz -2
Next, we derive an upper bound for
Σ Σ (WoL)vll)) (EU(O) 1=1
For any e U(0), we know that
0 < ( (Wol.)) o(, 0-) S VT.
(167)
We now derive an upper bound for (yi, o-). We need to focus on the following instances:
= Li and E
s s= L and / = L
o-) (-)) soW2T). ( = (
(168)
s + (1 -2 o(WTT)
- of S4. -2] (169)
=
Hence,
Σ (E2((0)
o (Wol)))))
Σ bT oWWT))
tEU((0)
o(BT)) 2 (1 o (1- -
o(wws, ( -2]
addition, we have
Σ Σ
W (Wol)ui()) s s0 $0( 1
(polvia))
(170)
iEK (K(0) 1=1
157 By
we can write
F(X (n)) 2
1
m
oT o(aT) ( 2+ -4000-000-00000-00000000000000-0 oz)i) Aut
v /L
m
bT (WWT)
12+ (1- -1- (1- ((()(()(())))))))) 2b) -2
2
-
poly(d))
(171)
39

---

Published as a conference paper at ICLR 2026
with
a nd2
AL
-2
Al -2
- (poly(a))
(172)
2 mL
Alio
Al.o
7
and
b =
+
(173)
mL
2
poly(d)
2a
v vmb
(poly(a)
mm
-
(174)
poly(d)
Therefore, conclude that
F(X(1)) (m) 2 VmL
m
aT o(aa))(1 o00T)(--0WW2)))000WW11
2
(1 (1- 0;))) A42 -2 o(ws,, 0311)
-2 - } poly(d)
(
(175)
If we can show (1 (01-o((wg)) (ows, (;))) -2 -1 -1 o (wWs,
> 0. then we can prove
(((x())) C for some positive constant C.
First define a random variable v11 (Wa, Oj). Then, we have from the definition of our locality- structured data type
En (L ((1-((11))) shod -2 -2 - -(1-)(b1)) Ard -2 =/ >
(176)
for positive constant K..
The random variable 4/2 = (1 o (1111) ALL -2 -2_ (1- (1-0(111))
ALt is bounded above by 1.
Applying Hoeffding's 's bound, for any 9 > O,
q log N V
11/2 1etha
(177)
N
From this we can conclude that,
(0
(Voloe
2i2 = (1 (1- (1_o(((w..)))) -2 -(1-o((wA,)))) Ard -27
2 -0
N
(178)
with probability at most ~9,
Hence, for sufficiently large N. we have from (176
(1 (1-o((ws,, ALo -2 (1 ((WA, -2
-2] 00
(179)
Therefore,
F(x( 2 C, where some positive constant.
(180)
Similarly, for negative sample, one can show by symmetry that the model output satisfies (() 1.
40

---

Published as a conference paper at ICLR 2026
(S4.1) Derivation for the convergence rate. Let's find the number of iterations T required such that F(X (m)) 1, since the label is +1. We require
m aT o(TT) 2 te.
(181)
VmL. 2
(1)
-(2) -b) and a(al)) ~
Substituting the value of a 2/mL condition becomes
since Ro 0, the
/maT
Vm
AL
Alo
T
4L
4T.
2/mL.
AL
Abo
n/T'
2 1 +
(182)
8T,2
2
Solving for T, we obtain
81, +.)
81,2
T2
2
(183)
ALo - ( Abo
ALt
-11)
Abo
2
-
2
By combining (181) with the expression for the model output F(X in
175
we obtain
F(X (x) 2 e)
(184)
Hence, the model output satisfies F(X (n)) 2 1.
(S4.2) Derivation for the sample complexity. Now we derive a sample-complexit bound that guar- antees zero generalization error.
Assuming enough samples, write for sufficiently small X <<
1
ddoo N mN
ALT
1
ALo
o
<
(185)
2/mL
2
2
From this can derive a lower bound on the required sample size,
4L2d
Ω
2
2
1/2
(1)
A
(1)
Alo+
(186)
12d
2
Ω
{ ALo
(2)
which will be (19) in Theorem
D PROOF OF LEMMAS APPENDIX B
Please refer to the supplementary material or preliminary work at (Shandirasegaran et al. 2025) for this section. We defer all proofs to the supplementary material, as the high-level ideas underlying the lemmas overlap with those presented in Appendix C for locality data. However, the case of locality-structured data presents additional challenges. Appendix [ provides the complete proofs for the locality- data, which contain the main technical ideas.
41

---

Published as a conference paper at ICLR 2026
E PROOF OF LEMMAS IN APPENDIX C
E. PROOF LEMMA
Proof. We know that the gradient of the loss function for the nthh sample is
de dE
(x)
aWow)
(F(X(m)
(n))
L
vs o' (WOLL)D(")) m
(187)
L
If we consider the gradient for the population loss,
Oc = -R
of (Wounvrf") "//
(188)
OWoG)
L
1=1
T, Ui ø' (Wor.)ef") yi (1)
--=-=-
+ E E,=-1
Σ
L 1i ø' (Wonnol)))
yi (n)
(189)
1=1
We given that < w 0+ 91
Pa
( .-)
S 92
and
Ps S [ o0) S 9s
(t)
for j + 1,2.
(190)
The Mamba output can be written as
(1 1 (1-0(w0)) (20 (0)) (2))) o(wol)) (1) xx) xx)
(1) =
II Jms+1
(191)
We have to consider FOUR cases.
Case I: L
3G = 07 = =
(192)
Byr 0+)
(1) o(w)) o+) 2 (p1) 0(())
1tePi O()).
1
(193)
Case I1: | = 8- = L
Ey Ev 11 0 = aww T 0+) 0(().
(194)
Case III: - I,t, $= I
Eyy 1 .o))= = 1 1 1 o+4) (1 (1- o-)
AL -2
1 (1) T 0,)
-oww (0) 0+) 0(T).
(195)
Combining (194) and 195p. we obtain
By 0+
a (ur 0+)
(t)
1 - o(r
(1) o.t)) (1-o(u((07 o.))
(196)
Aet -2
- a(w (T) o,)
(w 0+) ++)+0(1)
(1)
42

---

Published as a conference paper at ICLR 2026
Case IV: Others For the other token positions, #i + 04. Since we assume orthogonality among the features, y = 0. From our initialization, for the lucky neuron E W(0), 04 v.m. For E W(0), and 2 +1, we have
.0+) .o.) )
Σ T d' (Wort.)) (m) (n)
E Bz=+1
yi
1=1
VmL "ww o+) 2 1 (-o(w0)) ( (
(1-o olwo
o.))
(1- (1- ( o,)
sKi+ -27
+ OT()
(197)
Similarly for 1, we can obtain
1 11 d' (Wor.)of") w. :00)
E
Σ
L
aww (t)
2 1 1- o+)) (1-o(w5) o.)
o+)
VmL
(1-0(W?) 01)
Ar.o -27
O(T).
(198)
Therfore, combining 197) and (98
Uc
1
Ø (Wors))f")
v/ .04
(w(!) 06.))
(1)
0+
E.=+1
L
Us
T u. oy (Wor.)p(")) (m) 3/ (m)
1
E
.0
1=1
0+) (I (1-0(u))) o4)) (1 (0-o ( o-))
(w (1)
VmL
(1- ((0 (t) 0,))
st -2
-(1-0lug)) o,) slioy -27
1 O(()..
(199)
We aim to bound the deviation between the gradient of the population loss and that of the empirical M 7v
loss. Specifically,
dc owo... o6.)
ac (w
17n
where
12
(n)
7n
Σ
Ui
ø' (Wor.ne/") ()) m.
(200)
L
1=1
Consider a fixed vector with ||a:|12 = 1. We will show that a 7n is sub-Gaussian random variable.
la
7al TaTra
I ||al|2 11 = 11
(201)
By the problem setup, we know that
1 0' (Won.bb)) $1.
|1y|
(202)
mm
43

---

Published as a conference paper at ICLR 2026
Recall the Mamba output,
(0)0) ) olw!"' xx) x))).
(1-o(u0))))
yi (t) = Σ
II
(203)
1=0+1
Since | 12 1, we get
1-811
(x. [1]] |
s=1
a' (where a' denotes constant).
(204)
Therefore, the norm Yn satisfies
111ll2~
lod
0'
L
(1)
<
yi
m
L
d
a
(205)
L
1=1
Hence,
d'
a
T 7n
(bounded).
(206)
This implies that a T 7n is sub-Gaussian with variance proxy
- 0 m
(207)
Now consider the independent Lub--aussian variables a T n...., a YN. each bounded
mm S a 7n
(208)
m
Applying Hoeffding's inequality, for any > O,
MaJ Ex ETJY
q log N mn
7n
Z
SN
(209)
Observe that this can be written as
1 N
N Σ Yn
= aa - Ex nn =n
Eya
= TTC.
(210)
N
N
n=1
n=l
Therefore, by Hoeffding's 's inequality (cf. 209
9 log N MN
a
Ta
5 NN
(211)
To bound ||<|]2- use the dual norm identity
|Cl|
sup a T Ç. ||||==1
(212)
44

---

Published as a conference paper at ICLR 2026
We apply an E-cover argument to obtain
sup a ç = 1- acc. |x||==1
max a
< 2 max a a²Ç. acC1/2
(213)
We have shown that for any fixed az
9 log mN
P a
S N
(214)
Therefore, for all fixed a E C1/2.
log N mN
a T 2 V
with probability at most 4
(215)
Then,
q log N mN
max 2EC112 Recall that the covering number satisfies
a
[TTc/ 2
with probability at most /-
(216)
3B d
|C2|
(217)
For B and 2 have
|C1/2| S 6a.
(218)
We can therefore write
q glog N mN
||5l] 2 V P
S N
(219)
We want this probability to be sufficiently small. Set q = d, that
dog N mN
N -d
P ||51|2 2/
(220)
6
Hence, the deviation is bounded with high probability:
(Velox>
|K||| > o
with probability at most O(V
(221)
Or equivalently, with probability most O(N-4).
1
N
(VValoo
d N mN
Yn
Eyn
0
(222)
N
n=1
That is, with high probability 1 O(N- -),
we have
N
1
so
d mN
Σ
Eyn
(223)
N
n=1
Using the identities
1
8c
(w(1)
(1)
N
Σ
Yns
aw
(2)
Eyy,,
(224)
O(..)
((+))
45

---

Published as a conference paper at ICLR 2026
conclude that, with high probability,
((- ac aw
) (t) O(.))
ac aw (t) OLL.)
0c
0L
d N m.
0
(225)
ow!h
(t)
ow!!). O(1))
(t)
O(1))
122
Using the Cauchy- Schhwaz inequality, we have
7C
OC
0c
Oc
11 l12
@w
(t)
ow
0+
Ow ((..)
ow(() ((+,))
((,-)
((,,))
2
OC
Oc Ow(()
(since ||0+112 1)
ow(()
(2)
(+) ((+,))
O(+,)
dllog N mN
0
(226)
Therefore, obtain
Oc
-VV N
oc
0
aw
(1)
0+
m.N
ow
(0)
0+
O(t.)
O(6.))
(227)
OC
+0 (V
(Valnx>
/dlogg N
0+
ow(() O(.))
(t)
By pairing 1999)
with the given the conditions wa in (90 write
OC
ow (t) 0+
O(1,))
o(p)) - (1- (1--(q11) (1 o(4))) [1 o(gs)) مb. -2 (1 -ppps)) -2 - -(T)
ML
(228)
and
7L
(! (1) 0+
O(i,))
<
(q1) - (1 (1-o op1)) (1-o(())) (1
(1-o(p))) [11 (pp))) -2
o(2)))
-2
- (1- (1-o (1-o(ga)) Al.oz -3 + O(T)
/LL
(229)
(-
Therefore, we can obtain the lower bound and the upper bound of
(-ow!))
ac
0+
O(..))
VmL ((1)) (1 (()))) (1 a(2)))
(11 o(gs))
ALt, -2
(
( N
بجام -2 -
(1 o(p)))
-((7)
ac
@w (1)
0+
O(1,)
(230)
46

---

Published as a conference paper at ICLR 2026
and
ac @w 0+
O(1,)
((11) - o(r))) (1-)(())()))) (1 - op2)) [1 (p=) - -2 - (1 o(2s)) A -2] (231)
VmL
(v dlog mN
+ +0
+ +O()).
This concludes the proof of (16) and 1177) in Lemma C.1
To obtain
ac ow o(.))
we have to consider E Ele 7. oi of (Wors)f")) v
If (1 O(i,)
0- 0,
E B
Σ
1 ui o' (Wor.)ul") vn(h)) .
L
vmL aww (e) o-) 2
1 - 1- 0-1) (1-0(1100)))
(ww (0)
0.))
()) T sLe -2
(1-0
1 OT))
(232)
W (1)
Oli,)) 0- S0.
Σ Ur o' (Wor..s(") y/"") o_ = + O(T).
E
(233)
From 189
We know that
ac
L
I vi ø' (Wor.)(")) v, (m)
Ex=+1
Σ
O(1.))
(234)
Σ L Us & (Wornnef") v/ (m)
E
Hence, combining both cases, we conclude
(w (t) o-) 2+ (1 1++11-(w0) (I-o(Ws)
(1-o(w)) T o.)) (1- a(w ( o+4)
VmL
A15 -27
1 1 1- ( 0,)
- -O(T)
0c
o_ 0T)))
(235)
00 (1)
O(..)
47

---

Published as a conference paper at ICLR 2026
From 225
similar to
227
write
dc
dlog N m.
0
ow (1)
((1,)
ac
Ow 06,)
OC
+0 (
(Vellos
fllggN N mN
V
(236)
aw (2)
((+,)
Hence, we have
0-) 2 1 o(w) T o.)) (1 (1-0(10) oww) ( o4))
(w
(t)
VML
(w
(t)
0,)
Ala -2
d dlog N mN
O(T)
3c
d'log N mN
O()).
(237)
(w (t)
O(..)
This concludes the proof of [10 and 119 in Lemma C.I
Now consider
8W1D O16.)
Oj
for j # 1,2.
ac (1)
Σ Uu o' (Worr)e"") v/ (m)
0g
Ex=+1
.0
aw O(i,))
E Σ L
u. & (Wo (Worssl!") 31 (n) o,
(11,0,) (12, o,)
(238)
Because oj for * 1,2 identical in both I and 12. (1,0,) (12,0,) 0(r). Hence,
@W:0 O(L.))
8C
0j
+ (()). From (225). similar 227h, we can write
ac
(
d log mN
Ow(#) ((+)
(A)
Oj
0
ac
Oj
OW O6.))
OC
+0 (v
d N mN
Sw (2)
Oj
+ +O()).
(239)
((.))
Therefore,
ac
d N
Ow(C) oj < 5 O(L.)
(V
for j + 1,2.
(240)
mn
This concludes the proof of (110) in Lemma C.I
48

---

Published as a conference paper at ICLR 2026
E.2 PROOF OF LEMMA C
Proof. By definition, for any unlucky neuron i E K,+ W(0), we have
W 0.
(241)
We first consider the alignment with 04. That is,
ac @w O
0+
(242)
The gradient given in (88). We only need to consider the cases where (v)") .4
0. However,
since o s have
o (WoLt.el)) = 0.
(243)
Oc ow
0+
Eg=+1
Σ
L
Us ø' (Wor.nvl") vi"") .0+
O(1))
1=1
I.
Ui o' (Worr)))")) ")) v/ (m)
Ez=-1
E
Σ
,0
[==1
0 + O(T).
(244)
We know by (27
ac
ac
+oVa
+ +0
/dlgg N mN
ow O(..))
iw!!)
0+
(245)
O(i.)
Hence,
OC
(v
dog mN
0+
50
+ OT().
(246)
(t) @w
O(1)))
We now analyze the alignment with o To obtain the bound on
85
we consider the
ow!l) O(i.))
expectation E R-=-1 Elen 1 o' (Wor.)uf") (m)
y (n)
If (t)
o(..) o_ > 0, the inner product satisfies
I us o6 (Worssoi") yi (m)
.0_ ..))
E
alw! T o-) 2 1 1-o(w
(t) o-)) (1-0( T o+))
VmL
(1 (1-0(111)
( 0,))
AL. -2
OTT).
(247)
(t)
O(,,) 0_ < 0, then
Σ vi 8 (Wouner")) -v(") o_ = =0+0(1) L
E E-=-1
(248)
49

---

Published as a conference paper at ICLR 2026
From 189
We know that
Oc
L
1
Us o' (WorsBBl"))
y1" (m)
o_
BB=+1
(t) aw
L
O(i.))
Σ L vi o (Wort)ef") m/")) ,0_
E
(249)
Hence, combining both cases, we conclude
(20
T o-) 2 1 (1-o olw (I-O(W5)
T 0.)) (1- alwo o4))
VML
Al -27
- (w (t)
0,)
- O(T).
7c
ow (f)
o
S
(T))
(250)
((..)
From 225
similar to
227
can write
Oc ow
dlog N mN
0
(t)
((.)))
OC
(W (t)
O6+))
ac
(
Vallog N !
(w (1) O_
+0
(251)
O(..))
Hence,
(w (t) 0-) 2 1 1- 1--o(wh" T o.)) (1- (1-0(0))) o o4))
mL
1 (- (t)
Ah. -2
dlog
0j)
- O())
mN
ec
<
/ N m
0
+ O(T).
(252)
(t)
O(..)
Now consider
ac aw o(3.)
,0)
for j 1,2.
0C
0j
Eg=+1
Σ
L
U: Ø' (Woc.)ol")
vnn)) o
o O(..)
L.
Σ u/ 06 (Wors)ef") y/²)) .0
E
[=1
(T,,
(Iz,
(253)
50

---

Published as a conference paper at ICLR 2026
Because oj for # 1,2 is identical in both 4 and 1z. (1,00) (12, 0,) = 0 = O(r). Hence,
OC OW!! D(..)
03 0 (r)) From (225) similar to (27)) we can write
Oc
d log N m.N
Ow()) (1) ((+))
Oj
o
ac
OW O6,))
Oc
+0 (v
d dol N mN
+ O
(254)
aw (1) ((1,)
Therefore,
ac
(Venns
N mN
aw O(L.)
(1)
oj
for j + 1,2.
(255)
E.3 PROOF OF LEMMA
By symmetry, the proof is analogous to that of Lemma C.T
Please see Appendix E.1
E.4 PROOF OF LEMMA C.4
By symmetry, the proof is analogous to that of Lemma
C.2
Please see Appendix E.2
E.5 PROOF OF LEMMA
Proof.
The gradient of the loss with respect to wA for the nhhh sample is given by
de. (n)
m
L
vi me (WoLs)Di"))
Σ W (WDn!" a'"h)) (wJaf")) (Wor)aa!))
Ows
L
1=
o (I
II (1- o1-o(wm) (oTo (wIT("))
r=st1
11 σ (w)e(()) ) 2!(1)
Σ (1-0 (wIa((i) me)"))
j=xt1
(n)
0 (Woissus"))
(n)
(256)
We define the gradient summand
1 (n) = B x!!))
Σ
( (n)
(257)
j=st1
where the coefficients Ba,a and Bag are given by
Ba (W (1)) ( at (fooc
II (1- {( (1 (
(258)
rr=+1
51

---

Published as a conference paper at ICLR 2026
and
Gx.j (W (W =(")) (WW)) (a
II (1-o(w]a!")) (1 (
(259)
Lr=++1
If we consider the gradient of the empirical loss,
1 N (n) N
ΣΣ vi & (Wors)of)) 1=1
(260)
Ows
L
n=l
We are given that
(t) Pi S (w) 2ı
and r (witt)" o(L.) ,04) S s/.
(261)
From initialization, for all E Kt, we have o2
This gives
mm
Ot [m)
m
Owss
0+
('Won,
(2) Σ ( (1)))0)).
(262)
Averaging over the training samples, the inner product of the empirical gradient becomes
ac dws
N
(n)
m
0
2o'( (")
(n) o+)
N
a=1
1=1
L
1
Σ
8 (Won.
(n)
((n)
N
L
m
lis
m=2(n)=41
IEKg
1=1
8=1
L
8' (Wo(s))
o+)
vm
ick_
1=1
Σ
(' (Wol.)) (n)
Σ
1
T,
vm
==2(m)==1
iEK+
Σ
6'(Wo..)3
(n)
(n)
(263)
mml
First, we focus on the contribution from the samples where 2(n) = +1, for which we seek lower bound. We analyze the inner terms by considering four cases.
Case I: / = LT, = Lt Since ( and x. o+, it follows from 257 that
[ 0+ = Bo.e-
(264)
Using 258 with WB We I and क F, = 0+, we obtain
1, 0+
(m)
Br, = (w(++)) T 06.) ..) al(w (t) (4)) (1 (1-0((w))
0+))
. .+.)))
(265)
Given the conditions in 261 we can write
(n)
0+
2 (rì - O(F)) - 0(7)) - (1- + (())))
(266)
52

---

Published as a conference paper at ICLR 2026
We approximate d(1 0(r)) o(p)) O(r) and 1 + = - d(1)) O(T), since O(T) < O
Therefore, obtain
(I.0
2 - 0(7)) - (1- -(F)
(267)
2 (p)) (1-0 -1- -
Case II: 1 L3 L2
This configuration yields the same result as in Case I. We again obtain
1
0+ 2 ri o(pi) (1 (1- o((1) O(r).
(268)
Case III: [ = L2 8 = L Comparing (58 with (259) we see that the two expressions differ only their last term. In this setting, æ; equals o only when j = Li Consequently, Vs = 2j 0+ which implies B.,s Bz.j Hence,
(m) .0 Ba,s - Bn, 0+
(269)
Case IV: Others
For the other token positions, (1(2) ..+) 0 due orthogonality among the features.
Combining the above, the total contribution becomes
Σ (1 0+ 2 - ((pp))(11-1 (1-o(11)) O(r).
(
(270)
We now bound the entire sum over all tokens:
o' (Wo(.)))) Σ ( o-
1 -20 (271)
Let p.- W(t) be the number of contributing neurons. Then the total contribution from the active neurons lower bounded
(R) o+ ) 2r (pp)) (1 (1-)((1)
JΣΣ v (Wo(..))1) TEK+ 1=
Σ sm
-t_ 0(T). (272)
T,
1.s
vm
Next, we consider 2(m) 1 for K4. For (n) 1, the negative sample also contains two 0 features.
Similar to the above, we have consider 4 cases.
Case I: = LT= s == L
Since / 8, it follows from 257 that
(x)
Ba,a FI.
(273)
Since 2; 04, we have
( (I 0+ Bs,.-
(274)
We now seek an upper bound for this contribution. From the initial conditions in 261). we know
(
,o.) < O(T).
(275)
0(1.)
53

---

Published as a conference paper at ICLR 2026
Hence, we obtain
T 0+ (0) + O(T) a(1 + (()) -1-1(110(71))
(276)
We approximate (q + 0(+)) o())) (() and 1 o(pi 0(r)) - opi) +
since (T) < O
Therefore, obtain
(1 1 (si + 0()) (a(() 0()) ((-o(pi) O(T))
(277)
$ s (qi) :11 -1-o(p)))-0(())
Case I1: 11:1=10 Li L2
This configuration yields the same result as in Case I. We again obtain
0 s s o(n) (1 (p))) O(r)-
(278)
Case III: / = 8 Lt In this case, the contribution vanishes:
[ 0+ 0 0(r).
(279)
Case IV: Others
For the other token positions, (r o+) 0 due orthogonality among the features.
The maximum number of such contributing neurons is m bounded above by
Therefore, the total contribution is 2
Σ ((n) 1,s o+) < 2s aai)) . (1 (1- (11-((1)) m + O(T)
vi o' (Woir)
T. KE
vm
2
(280)
1=1
=Vm 87 ((1)) (0-00(p)) O(r).
Thirdly, let consider the contribution for (n)) +1 from E K_. From initialization, for mm For ((n) +1, we seek upper bound on the contribution from such
i E KT, 1/
neurons.
Let (*) +1. To maximize the term Wo( (n) in 258).
consider the token locations which contain o.. features since Wo(s.) has a large component in the o direction. Then ø Y1 contains the feature.
However, in this case, x, = o- = I₁ and due to orthogonality.
(n) 1 0+ = 0.
(281)
Hence, we only need consider time steps = L Lt. where 04 features appear. Recall that
af Ows 0+
m
L
o &'(Wo(.13)
Σ 1,s .04
(282)
1=1
We analyze the inner contributions case by case.
Case I: L $ T
54

---

Published as a conference paper at ICLR 2026
Given that
(Vales N mN
of 5 +0
(283)
we obtain
(m) L. 0+ s
(284)
Case II: 1 8 L
This configuration yields the same bound:
1 0+ s sc.o(qi) (1)o)))))0)))
(285)
Case III: / Lÿ s L In this case, the contribution vanishes:
[ 10 = 0 0(r).
(286)
Case IV: Others
For the other token positions, (1/00) 0+) 0 due orthogonality among the features.
Thus, the total contribution from each E K satisfies
Σ
1. L, ..) 2c ((1)) (1 opi)) + + O(T).
(287)
The maximum number of such contributing neurons is m 2 so the full contribution is bounded by
o' (Won.)y))
Σ {11ra))
(n) o
2c (()) (p-)(p1))
2 O(T)
MLL
Vm
(288)
iEK_ t=1
Vme o(q)) ol-o)(p1) +OT).
Therefore, the overall contribution
o' (Wo
Σ (m) .o.) -vmc a(gi) (1 a(pi)) ((r). (289)
MLL tEK.
Finally, consider 2(n) -1 for i K_.. For 2(n) -1, want lower bound since
We could consider / = 1+ 8: O4, and write
(V N
Won o+) 2 81 -0
mN
(290)
However, the minimum number of such contributing neurons not tractable. Thus, if we consider the worst case where W Woú.)) for E K- does not learn the o. feature, the obvious lower bound is zero:
Σ VEK_ 1=1 Vm Φ (Wo Dt)
(n) 1.8 o+ 2 0.
(291)
We now combine the bounds for the four terms identified in (263). corresponding to the contributions
from: (i) K.,, with (n)) +1 as shown in (72). (ii) K4, with (()) -1 shown in (280). (iii)
K_ with z(n) +1 shown in (889 and (iv) K_ with z(n) -1 shown in 291
We assume
55

---

Published as a conference paper at ICLR 2026
the batch is balanced, the number of positive and negative samples is equal, with each class contributing samples. Then we have
ac
2r o(pi) (1 o(91))
.+
2
(1- -
Ows
Vm
-Vm si (())) (1 (1- (1-r(()))) + - (T))
a(pi) (1 o(()) ri pf a(9i) (1 - (p1) vm
(292)
Vm
2
dlog N mN
(iT).
(293)
(V
(v (alxx N m
where we have used the fact vm o(9)) (1 cp-))
Velog N
since 0
E.6 PROOF OF LEMMA C.6
Proof. The gradient is given in (256
Let's consider the alignment with oμ for k + 1.2.
Owss
OE
Os
(())
m L
Us
& (Woo,))) (m))
Σ
((0) Ok
(294)
8=1
From our initialization, for all E K.t, we have mm
We first consider the case 2nn) = +1 for i E K+ Since WORL),, for i E K+ has a large 0+ component, have consider the token features with 04. For g(n) +1, only when = L 8 = L we have 2 = I, 0+. Therefore, Woti, D, is significant. Hence, we have
1 0k
Σ
Bag(zs)")
(n) ok)
jmst1
S Bs,st1
(Assuming W.L.O.G. - = ok)
(W O(t.)
ALt 0+) o((w o.)) (1-) (1-0((u3)) o+))) (1-0((101 on)))
(295)
Using the the conditions in 261p. we can write
1/ (n) ox) < (-ri (-ri+0(7)) opp + O(F) (1- (4))) - AL..
(296)
We can approximate op + 2 0()) - -0))) 1 + 0(T) and
0
1 の (42 - ((r)) = =1-a (g2) + (T)), since ((T))(())))))))) (
Hence, we obtain
o' (Wou.wsl in)
(n) Lis , ok
m
(297)
1
[ (p))) . o ) - (1 - 0(92)) + O()))
mm
56

---

Published as a conference paper at ICLR 2026
Let Pi [W())] be the number of contributing neurons. Then the total contribution from K+ neurons bounded as
1 L
Σ vs o (WorDDD ( iEK+
1. os
ri opi) (1 d(1))
vm
(298)
(1 0 / (z pt O(r)).
Next, we consider 2(n) = -1 for E K4+. Since ( , on)
< 0, we require a lower bound for this.
I (n)
Ok
Σ Baglar (n) ,0r)
j=st1
2 (Witth) o(..)
0+) o(w ((W 0+)) (1 (1-o((10)) (1 (1- 1--o((u/ .x)))
Ar
(299)
Using the the conditions in 261p. we can write
(n) L. ox 2 of o(qi) -1 (1- (1 (p22))
(300)
Σ Σ
Ui ø (Worsse))) (m) Σ Lo
L
iEK+ 1=1
(301)
o(a) (1 - -f (p2) مملد pt.
si
Since ALo, > ALt this term is negligible which leads to
Σ D; ø (Worsvl") (m) 1=1
Σ
T ox) -
(p2) 0.
(302)
L
Kek.
Thirdly, we consider the case i KC=, Lj the contribution is significant.
"C,,foo for z(m)
1. Similar to 295 and
when / Li s
( --ri) (1-))((1)((1))) (1 -(())) Sho + O(T)).
(303)
Hence, we obtain
(Woo. v)(")) Σ
((m)
1,8 ox
L
m
Iml
(304)
[-nt ( (1 o(gi)) (1-(()))))))
o()))) AL. + O(T))
1.
mm
1=1
Let p. = [((t] be the number of contributing neurons. Then the total contribution from K_
neurons bounded
r 2s o' (Wonsvl") Σ (1 ox) s ri a(pi) (1 - a(9))) (1 0(92)) ماō Pi vm
KEK
l=1
+ +T)))
(305)
57

---

Published as a conference paper at ICLR 2026
Finally, we consider i K_ for 2(a) can write
+1. Following the same approach as in
299 to 301
1
Σ
Σ
vi o (Wois)8)))) Tm))
Σ 1.
L
iEk_ 1=1
(306)
si
(1 - (1- stio Pi.
Since AL. > Alo this term is negligible which leads to
1 ti o' (Worssl))) (m)
(o ok) 2 01(11)(22))211
2/3 0.
(307)
L
iEK
1=1
Putting it together, We know
ac
no ( (Wo (n)
Σ
(n) ox
Ok
dws
N
n=l
1=1
o (Wo(L))
(n)
(n)
N
L
mm"
LEK+
1=1
s=1
8' o(Woo..)
1,8 ON
v
icK.
1=1
Σ
0 (W (oo
1,s , Ok
Vm
==-(")==1
iEK,
1=1
s=1
L
8' (Wo6.13
L. .o)
(308)
vml
YEK_ 1=1
We now combine the bounds for the two terms identified in equation 308
corresponding to the (36).. (ii) K_ with
contributions from: (i) K4 with 2(n) 11 298 (ii) K4 with 2(a)
((n) +1 (01}) and (iv) K_ with 2(n) 1 305 We assume the batch is balanced, so the
number of positive and negative samples is equal, with each class contributing N samples. Then have
Oc
Ow
(r)
0k
ri 2/m (pp)) (1 (())) (1-a((2)) Shi pt + (1 - a(9z)) Ali pi)
si
- (1-)) 10 (11 (122) -1ă_) Pi + +o 0(11 a(p2)) -pt
+ Vm"
(1
(309)
+(T))
(310)
ac
dw (1) 0k
2/m (o)))) (1 - a(11)) (11 (1-o(1))))
ri
(1-0 p (1 o(12)) Alion Pi
5
+ 0 (1 ((p2)) + +1((1o(p22)) ) +0()))
(
(311)
58

---

Published as a conference paper at ICLR 2026
From
302
and
(307]
we conclude
"i
o(1)) o(q:)) (1
(1-o(i))
(1=) (1- pt= (1- o(92)) pi
Ow
(t) Ok
2/m
A
+ O()))
(312)
F EXTENSION TO MULTI-CLASS CLASSIFICATION
Consider the classification problem with four classes, where each example is assigned a label 2 = (21, (2) ++1,-1} representing four distinet classes. Similarly the binary setting, there exist four orthogonal discriminative patterns. In the output layer, the scalar coefficient associated with hidden neuron is replaced by a two-dimensional vector v; R
Hence, we define the model output as
F(X)=
v の W (o.)))
(313)
L
1 Σ
m
(1)
(w:) o W (o
(314)
L
(2(
Σ
(v))2 o W Wo
(315)
L
The dataset can be divided into four groups
D1 {(X (n) (n))
(n) (1, 1),
D2 ((x(") (n)
2 (n) (1, -1),
(316)
D3 ((x(") (())
(0)) (-1,
D1 {(X (N)) x1))) 2x000 (-1, -1)}.
The loss function data (x(*), (n)) is
(n),
Loss X (n), (n))) = max {1 - -1-21 (1)) ((x(1))) 03.
(317)
Since v. E (++ (t=t )2, we divide neurons into four groups:
W1 {i J (1, 1)},
W2 = {i: Uz 7m (1, - -1)}
(318)
W3 {i: 2s J (-1, 1)]:
W4 {i: U. J (-1,-1)
For neuron i, the gradient decomposes as
ÔLoss
-2 (n) OF (n)
(n) OF (x(())))
(319)
dWo
aWoG.
@Wo
Let 01, O2, O3, 04 denote the four discriminative directions. Consider i W2, i.e v, - (1, -1). Projecting the gradient onto O2, for any
(x ((()) DD, we obtain
OLoss
2
022
||22|2 > 0,
(320)
"Wou))
m
59

---

Published as a conference paper at ICLR 2026
showing GD moves Wot).) toward 02- For samples from the other classes:
ØLoss
(x(m)2(0) EDI: )
oWor.)) 01 0,
(())) E DD
@Wo())) 03
ÕLoss
2 vm
00
(321)
ŌLoss
(x("), E D4:
@Wo()) 00 = 0.
Thus, for i E W2, the update direction aligns with O2, and similarly neurons in Wi, W3, W4 align with 01, 03, 04 respectively. Similarly, we can analyze the gradient dynamics of the gating vector WA.
60


Reviews:
[{"review_id": "LfpjsW0XFF", "reviewer": "Reviewer_Sfca", "summary": "This paper theoretically analyzes the training dynamics of a single Mamba layer wrapped with a two-layer MLP trained via gradient descent. Under assumptions like majority-voting or locally structured data, the authors establish generalization guarantees with improved sample complexity and convergence bounds derived from the characterization of the final solution. The analysis centers on the gating mechanism, a key innovation in modern recurrent models, and is supported by empirical evidence validating the theory.", "questions": "Q.1. I think $v_i$​ in Equation 6 is not introduced before it appears in the equation, and the authors should make this clearer.\n\nQ.2. The visibility of Figure 1 should be improved. For example, it would be better to use figures of the same size (the one on the far right is smaller), and the space between the first two figures is too narrow, causing the caption text to almost overlap.\n\nQ.3. Which aspects of the proof are not relevant for gated linear attention but are relevant only to Mamba layers?\n\nQ.4. Are $W_B$​ and $W_C$​ trained, or is there a typo with the superscript 0 in Eq. 15 for those weights? It seems that the generalization guarantees are derived for a model where some parameters are frozen.", "limitations": "", "rating": 2, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "W.1. Expanding the theoretical understanding of Mamba and its gating mechanism is important, as these components are popular in both transformers and modern RNNs.\n\nW.2. The analysis is novel, non-trivial, and sheds light on the training dynamics and the role of critical components in modern architectures.", "weaknesses": "**W.1. Connection between theoretical settings and the real world:**\n\nWhile the analysis of the simplified model is interesting, its connection to standard practices in the domain remains vague. I understand that providing theoretical proofs for the full model is challenging and that simplifications (such as using a single-layer model or specific datasets) are standard. I’m not expecting a full theoretical proof for the complete model, but I believe the authors should better link key aspects of the theoretical analysis to real models.\n\nFor example, they could run synthetic tasks inspired by the theoretical analysis on the full model and investigate whether the behavior in which MLP weights align with class-relevant features also appears in the complete models. Similarly, they could examine whether the cosine similarity between the gating vector and both class-relevant and class-irrelevant features is consistent with the theoretical (or single layers)  findings.\n\n**W.2. Related work is missing.**\n\n Several studies analyzing the training dynamics, optimization issues, and generalization properties of Mamba are not mentioned. Examples include [1, 2, 3, 4, 5]. In this sense, some of the claims made in the paper (see ’To the best of our knowledge, this work presents the first theoretical generalization analysis of the Mamba architecture.’ and ‘To the best of our knowledge, we are the first to theoretically analyze the learning and generalization performance of the trained model with gradient descent (GD)’) may be overstated.\n\n**W.3. The experimental analysis is limited:** \n\n(i) Some information is missing or insufficiently detailed. For example, it is unclear which model was used (simplified, single-layer, with or without convolution, and whether parameters were frozen as in Eq. 15).\n\n(ii) The experiments could be improved, for instance, one could empirically analyze multiple models across different data regimes (varying numbers of features, data distribution parameters, etc.). In addition, several architectural ablations could be performed to better understand the dynamics. For example, training the model without gating to see whether it struggles, as the theory predicts, could strengthen the connection between theory and practice.\n\nW.4. Minor: The visibility of Figure 1 should be improved. For example, it would be better to use figures of the same size (the one on the far right is smaller), and the space between the first two figures is too narrow, causing the caption text to almost overlap.\n\n**W.5. Generality can be improved:**\n\n Some parts of the theoretical analysis could be applied to other modern gated RNNs, such as Gated Linear Attention, Gate Delta-Net, RWKV, and RetNet. It would be good to explain which parts of the analysis are relevant to each architecture, and which components in those architectures improve or negatively impact the optimization dynamics or generalization.\n\n___\n\n**References** \n\n[1] Generalization Error Analysis for Selective State-Space Models Through the Lens of Attention  . Honarpisheh. Nips 2024\n\n[2] The Implicit Bias of Structured State Space Models Can Be Poisoned With Clean Labels. Slutzky∗ et al. NIps 2025. \n\n[3] REVISITING ASSOCIATIVE RECALL IN MODERN RECURRENT MODEL. Okpekpe et al.\n\n[4] Repeat after me transformers are better than state space models at copying. Jelassi et al.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ak28zAMTOW", "reviewer": "Reviewer_8FgG", "summary": "The paper presents a theoretical analysis of learning and generalization for a simplified Mamba block (single-layer selective SSM with input-dependent gating followed by a 2-layer MLP) trained by gradient descent on structured sequence data. Two scenarios are studied: majority-voting and locality-structured sequences with token-level noise. The authors prove non-asymptotic sample complexity and convergence guarantees under width and noise conditions, show that the gating vector aligns with class-relevant features while suppressing irrelevant ones, and provide synthetic experiments supporting the theory. They also position Mamba’s dynamics relative to Transformers, arguing that Mamba can match the majority-voting setting and outperform attention baselines on locality-structured data due to selective recurrence.", "questions": "- What would the theoretical results be if a more complex setting of the Mamba block were used? For example, what happens if a normalization layer is added or if the Mamba block is repeated?\n\n- Are the assumptions in Theorems 1 and 2 regarding width and small noise required only for the proofs to work, or are they fundamentally unavoidable? What happens if these assumptions do not hold?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper contains a clear novel theoretical results. In particular, two formal and non-asymptotic bounds are presented for both sample complexity and iteration complexity in two scenarios (Theorems 1 and 2), with interpretable dependence on signal gap, locality separations, step size, sequence length, and noise.\n\n- In addition, the authors provide a rigorous characterization that the gate prioritizes class-relevant features and suppresses confusing ones; connects to attention’s feature-selection role while highlighting differences arising from multiplicative recurrence.\n\n- Synthetic experiments are closely tied to theory, including comparisons where Mamba outperforms Transformer/local attention for locality.", "weaknesses": "- My first main concern is the simplified assumptions (single layer, single head, simplified recurrence), which make the considered setting far from practical Mamba stacks (with depth, multi-head structure, residuals, normalization, and learned discretization). However, I understand that a simplified setting is commonly used in theoretical analysis, as the practical setting is usually too complicated to allow for a thorough analysis.\n\n- My second concern is the strong assumptions required in Theorems 1 and 2: the width $m$ must be larger than or equal to the square of the dimension $d$, and the noise level must satisfy $\\tau = O(1/d)$. These assumptions weaken the novelty of the theoretical results. A justification for these assumptions is needed.\n\n- The experimental results are quite limited. The empirical comparison to Transformers and local attention is restricted to synthetic settings, with no controlled ablations for optimizer choice, learning rate scaling, or positional encoding. It is unclear whether attention could close the gap with appropriately tuned locality biases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DA2HQIkZBY", "reviewer": "Reviewer_rcjx", "summary": "This paper claims to provide the first theoretical study of the learning and generalization dynamics of the Mamba architecture, a selective state space model (SSM) that has recently emerged as an efficient alternative to Transformers. The authors analyze a simplified single-layer Mamba block with input-dependent gating followed by a two-layer MLP trained via gradient descent under two structured data regimes—majority-voting and locality-structured sequences containing both class-relevant and irrelevant tokens. They prove non-asymptotic sample complexity and convergence guarantees, showing that Mamba generalizes efficiently when class-relevant signals are strong or locally concentrated, and that its gating vector naturally aligns with informative features while suppressing noisy ones, formalizing its feature-selection role. Synthetic experiments align with theoretical insights.", "questions": "Please address weaknesses above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper’s primary strength lies in its original theoretical treatment of Mamba, offering the first formal generalization and convergence analysis for selective state space models—an area that has so far been dominated by empirical studies. In terms of originality, it introduces a novel analytical framework for gated architectures trained with gradient descent, bridging ideas from feature-learning theory and state-space modeling, and extending prior analyses of Transformers to a fundamentally different recurrence-based mechanism. In terms of technical quality, the authors provide rigorous, non-asymptotic sample complexity and convergence bounds under well-defined data models, supported by clear mathematical reasoning and alignment with empirical trends. The paper could benefit from an improved exposition of the theoretical results. In general, the paper is well structured, with intuitive explanations preceding formal results. Regarding significance, the work fills a key theoretical gap by formalizing how Mamba’s gating acts as a feature selector that prioritizes class-relevant signals, thus deepening understanding of why such models generalize efficiently and under what data conditions they outperform attention-based architectures. While it is unclear whether this specific architecture will stand the test of time, this reviewer celebrates the introduction of rigorous analysis of architectural components.", "weaknesses": "While the paper makes a valuable theoretical contribution, several weaknesses limit its impact and clarity. First, the analysis is restricted to a highly simplified single-layer Mamba block and may not generalize to deeper or multi-head architectures used in practice; extending the framework to multi-layer dynamics or coupling between blocks would strengthen its relevance. While this is hard to do, authors could comment on whether these results extend to multi-layer settings. Second, the structured data assumptions (majority-voting and locality-structured) are somewhat idealized and may not capture the complexity of real-world sequences; providing evidence that these regimes approximate practical scenarios (e.g., via ablations on natural datasets) would improve the paper’s applicability. Third, the empirical validation is minimal—limited to synthetic data—and lacks comparisons to other theoretically analyzed architectures (e.g., gated linear attention), making it hard to assess the distinct benefits of Mamba’s gating mechanism. Fourth, while the paper claims to establish generalization guarantees, the proofs rely on strong assumptions such as balanced datasets, small token noise, and specific initialization schemes, which could limit robustness. Finally, the connection to prior theoretical work on Transformers and SSMs could be deepened by discussing how Mamba’s multiplicative gating fundamentally changes training dynamics beyond attention-style weighting. Strengthening these aspects would make the work more comprehensive and impactful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0bRayXUiMQ", "reviewer": "Reviewer_T7xH", "summary": "The paper presents a formal theoretical study of the Mamba architecture. It analyzes a simplified single-layer Mamba followed by a two-layer MLP trained with gradient descent on two synthetic data regimes — a majority-voting model and a locality-structured model — where class labels depend on subsets of relevant features within the input sequence. It derives non-asymptotic convergence and generalization guarantees, proving that Mamba’s gating mechanism learns to amplify class-relevant features while suppressing irrelevant ones. The analysis links learning efficiency and generalization to interpretable data parameters such as signal-to-noise ratio and the fraction of informative tokens. Synthetic experiments confirm these trends qualitatively.", "questions": "1. Have you observed any empirical evidence suggesting these theoretical results hold beyond your simplified setup? Or how do you think these results fair in practice?\n2. Can your analysis or conclusions be generalized to more complex or non-Gaussian data distributions (e.g., correlated, heavy-tailed, or structured inputs as in language or vision)? If not, what aspects of Mamba’s behavior might your theory fail to capture?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written, clearly structured, and theoretical results are supported by rigorous proofs with proof intuitions given in the main text.\n2. The paper is well-situated in the literature. I would only suggest to add [1,2], when discussing the connection of SSMs and Transformers in L147-154, for additional reading on how Mamba relates to attention, altough not critical.\n3. The theoretical analysis of the learning process of Mamba is novel and timely, as Mamba has become an alternative to Transformers.\n4. The results are convincing, however their generality is not assessed yet.\n5. The paper does a good job in linking the empirical findings to the theoretical results, supporting each statement with empirical evidence.\n\n[1] Dao & Gu (2024), \"Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality\", https://arxiv.org/abs/2405.21060\n\n[2] Sieber et al. (2024), \"Understanding the differences in Foundation Models: Attention, State Space Models, and Recurrent Neural Networks\", https://arxiv.org/abs/2405.15731", "weaknesses": "1. The main drawback of the paper is its limited scope, both in its model selection and the application setting. While chosing a single model (in this case Mamba) for such a study is reasonable, I would have expected a broader application setting, i.e., not limiting the study to binary classification.\n2. The investigated data models are purely theoretical and are not representative of more realistic data models, e.g. the Gaussian distribution assumption is not realistic for practical applications. Having a discussion on this point in the paper would be appreciated.\n3. The theoretical results are not novel in itself and are mainly extensions to the Mamba specifics. Altough I think this is not a major problem, but it reduces the originality of the contribution.\n4. The paper could do a better job at clarifying its limitations, e.g. with a paragraph in the conclusion.\n\n\nMinor points:\n- The short paragraphs before lemmas/theorems already discuss the results of the lemmas/theorems before they are stated. It would be better for the flow to reserve this discussion for after the lemma/theorem and only give a brief introduction to which results is shown next.\n- The remarks after the results are nice, but should not be wrapped in a proof environment, i.e., should not have a square at the end.\n- the empirical results provide no quantitative comparison of actual to theoretical convergence rates and are purely qualitative. Including, such a comparison/result would strengthen the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)