Trained Mamba Emulates Online Gradient Descent in In-Context Linear Regression
Authors: Jiarui Jiang
Wei Huang
Miao Zhang
Taiji Suzuki
Liqiang Nie
Venue: NeurIPS 2025 poster ()
Keywords: state space model
Mamba
in-context learning
training dynamics
Primary area: theory

Abstract: State-space models (SSMs), particularly Mamba, emerge as an efficient Transformer alternative with linear complexity for long-sequence modeling. Recent empirical works demonstrate Mamba's in-context learning (ICL) capabilities competitive with Transformers, a critical capacity for large foundation models. However, theoretical understanding of Mamba’s ICL remains limited, restricting deeper insights into its underlying mechanisms. Even fundamental tasks such as linear regression ICL, widely studied as a standard theoretical benchmark for Transformers, have not been thoroughly analyzed in the context of Mamba. To address this gap, we study the training dynamics of Mamba on the linear regression ICL task. By developing novel techniques tackling non-convex optimization with gradient descent related to Mamba's structure, we establish an exponential convergence rate to ICL solution, and derive a loss bound that is comparable to Transformer's. Importantly, our results reveal that Mamba can perform a variant of \textit{online gradient descent} to learn the latent function in context. This mechanism is different from that of Transformer, which is typically understood to achieve ICL through gradient descent emulation. The theoretical results are verified by experimental simulation.

Trained Mamba Emulates Online Gradient Descent in In-Context Linear Regression
Jiarui Jiang",, Wei Huang" Miio Zhang Taiji Suzuki Liqiang Nie!
Harbin Institute of Technology, Shenzhen
2RIKEN AIP
University of Tokyo
j jiaruieeuutoos com, wei fring.nng.o vr@riken. jp, zhangmlao@hit. edu. taiji eeiiimmttt.... tou.ookyo..a..... jp. nieliqiang@gmail com
Abstract
State-space models (SSMs), particularly Mamba, emerge as efficient Trans- former alternative with linear complexity for long-sequence modeling. Recent empirical works demonstrate Mamba's 's in-context learning (ICL) capabilities com- petitive with Transformers, a critical capacity for large foundation models. How- ever, theoretical understanding of Mamba's 's ICL remains limited, restricting deeper insights into its underlying mechanisms. Even fundamental tasks such as linear regression ICL, widely studied as a standard theoretical benchmark for Transform- ers, have not been thoroughly analyzed in the context Mamba. To address this gap, we study the training dynamics of Mamba on the linear regression ICL task. By developing novel techniques tackling non-convex optimization with gradient descent related to Mamba's s structure, rate to ICL solution, and derive a loss bound that is comparable to Transformer' ss Importantly, results reveal that Mamba can perform variant of online gradient descent to learn the latent function context. This mechanism different from that of Transformer, which is typically understood achieve ICL through gradient descent emulation. The theoretical results are verified by experimental simulation.
establish exponential convergence
Introduction
State-space models (SSMs), notably Mamba (Gu and Dao 2024), have recently emerged as com- pelling alternatives to Transformer-bass architectures (Vaswani et al.) 2017). Mamba integrates gating, convolutions, and state-space modeling with selection mechanisms, enabling lincar-time complexity. This effectively addresses the quadratic computational costs typically associated with self-attention mechanisms in Transformers. Consequently, Mamba demonstrates superior efficiency in processing long sequences while maintaining
surpassing Transformer performance across
diverse benchmarks (Gu and Dao
2024 Dao and Gu 2024
Patro and Agneeswaran 2024
Liu et al.
2024 Ahamed and Cheng 2024 et al. 2024a
In-context learning (ICL) (Brown et al. 2020) is a powerful paradigm that enables models to generalize to unseen tasks by dynamically leveraging contextual examples (such as input-output pairs) without task-specific fine-tuning. This capability has become a defining characteristic of large foundation models, significantly enhancing their flexibility and adaptability. While extensive research has provided substantial insights into Transformer-base ICL mechanisms (Garg et al. 2022 Gatmiry
et al. 2024: Sander et al. 2024
Zheng et al. 2024: Zhang et al. 2025). the principles underlying
'Equal contribution
"oorespoonding author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Mamba's ability to perform in-context learning remain largely unexplored, highlighting a compelling research gap.
Recent empirical studies have examined Mamba's (ICL) capabilities, showing it matches Trans- formers many standard ICL tasks, while surpassing them in specialized scenarios like sparse parity PPark et al. 2024) Grazzi et al. 2024). Bondaschi et al. (2025) theoretically analyzed its representational capacity for in-context learning of Markov chains, and Li al. (2025a) investigated
binary classification tasks with outliers. (Yang et al. 2024 2025 Behrouz et al., 2 a leverage
the connection between SSMs and online learning to design new architectures. However, even the linear regression model, a canonical setting widely used for theoretical analysis of Transformer- ICL mechanisms, remains theoretically underexplored in the context of Mamba. To fill this gap. we analyze Mamba's 's training dynamics on in-context linear regression tasks. More precisely. following the previous ICI. analysis in Transformers (Garg et al., 2022] Zhang et al. 2024 Ahn et al. 2023). this paper focuses on data generative model with N input-output pairs (1 Va ri=i) and query input (x: satisfying y = fx) w æ, where æ denotes the input and y denotes the output, and w is randomly sampled from Gaussian distribution, termed the context. In this work, we develop a rigorous theoretical framework to analyze how randomly initialized Mamba models, when trained through gradient descent, evolve to implement in-context learning. We demonstrate that the trained Mamba architecture dynamically leverages the input context to perform implicit estimation of the vector un. This estimation is achieved through hidden state updates that mimic online gradient descent steps, finally implementing prediction for yo = f(ar,) w Z.q. We also provide a loss bound that is comparable to Transformers'. Our contributions are summarized as follows:
We construct a Mamba architecture (S6: S4 with selection) capable of ICL, establishing its exponential convergence rate to ICL solution, and further derive the loss bound after convergence. The loss matches that of Transformers.
Technically, we develop novel techniques to address optimization challenges induced by random initialization and gradient descent, rigorously characterizing Mamba's training dynamics when trained from scratch.
We reveal how trained Mamba achieves in-context linear regression by progressively aligning its hidden states with the context through sequential token processing. This finding provides a new perspective for understanding Mamba's ICL mechanism, distinct from Transformer- based approaches. All the above results are verified by experiments.
2 Related Work
In-Context Learning The seminal work of Brown et al. (2020) demonstrated the in-context learning capability in Transformers, showing their ability to infer functional mappings from input-output exemplars without weight updates. Garg et al. (2022) initiated the investigation of ICL from the perspective of learning particular function classes. Following these, line of research analyze this phenomenon through the lens of algorithm imitation: Transformers can be trained to implement various learning algorithms that can mimic the latent functions in context, including: a single step of gradient descent (Von Oswald et al. 2023 Akyürek et al. (023). J statistical algorithms (Bai et al. 2023), reinforcement learning algorithm (Lin et al. 2024), multi-step gradient descent (Gatmiry et al.
2024), ) mesa-optimization (Zheng et al.] 2024), Newton's 's method (Giannou et al.) 2025). weighted
preconditioned gradient descent (Li et al. 20256), in context classification (Bu et al. 2024
Shen
et al. 2024 Bu et al. 2025) among others.
Recent work extends ICL analysis beyond Transformers: (Lee et al. 2024: Park et al., 2024)
empirically compared popular architectures (e.g., RNNs, CNNs, SSMs, Transformers) on synthetic ICL. tasks, identifying capability variations across model types and task demands. TTog and Pehlevan (2024) demonstrate that MLPs can learn in-context series of classical tasks such as regression and classification with less computation than Transformers. [Sushma et al. (2024) show that state space models augmented with local self-attention can learn linear regression in-context. Unlike existing research on ICL, this work focuses on the ICL mechanism of Mamba (specifically S4 with selection) and its training dynamics.
Theoretical Understanding of SSMs As Gu et al. 2022 introduce structured state spaces models in modeling long sequence and further be extended to Mamba (Gu and Dao 2024), which gained

---

significant attention as alternatives to Transformers, extensive research has sought to theoretically understand the mechanisms and capabilitiess of state-spaces models (SSMs). Dao and Gu (2024)
propose the framework of state space duality, which establishes a connection between SSMs and attention variants through the lens of structured matrices. Vankadara et al. (2024) provide a scaling analysis of signal propagation in SSMs through the lens of feature learning. Cirone et al. (2024) draw the link of SSMs to linear CDEs (controlled differential equations) and use tools from rough path theory to study their expressivity. Chen et al. (2025) establish the computational limits of SSMs and Mamba via circuit complexity analysis, questioning the prevailing belief that Mamba possesses superior computational expressivity compared to Transformers. Nishikawa and Suzuki (2025) demonstrate that state space models integrated with nonlinear layers achieve dynamic token selection capabilities comparable to Transformers Different from the above, provide theoretical understanding of Mamba from the perspective of ICI..
Problem Setup
In this section, we outline the ICL data model, the Mamba model, the prediction strategy, and the gradient descent training algorithm.
Data Model. We consider an in-context linear regression task where each prompt corresponds to new function [(x) = w æ with weights w (N Ia) and d > 1. For each task, we generate N i.i.d. input-output pairs {(x,,)) y:) }) and a query zg. where all inputs 24, Ty ~ (O, Ia) are independent Gaussian vectors, and the outputs satisfy yi f(x.). The goal is to predict y4 f (x,) for the query.
To enable sequential processing of prompts in the Mamba model, we implement an embedding strategy where:
1. The i-th context token encoded as e, (ax vs) formed by concatenating input x; with its corresponding label y/s
2. The query token represented as C.g = (Λr T)T, zero placeholder.
T, masking the unknown target value with a
In many theoretical analyses of Transformer- in-context learning, token embeddings are conven- tionally concatenated into single matrix to enable parallel computation of global attention (Zhang) eet al. 2024; Ahn et al. [023] Huang et al. [2023 Mahankali et al. [2224] Wu et al. (0224). In contrast, since Mamba operates as a sequential model, we feed the embeddings context tokens one by one,
and finally the query token (e: Eg
eN eg).
Mamba Model. We consider a S6 layer of Mamba OTl Mamba u1.2) with selection, discretization, and linear recurrence components, where ut, Or Rdd It can be described as follows:
A,60 = (, Bin
(6)
(1a) A. = cxp(A E Xd. 10 E IRdx11,
(2a)
(0) = c h C E
(Ib) B, = (cxp(AA))) - E Rd (2b)
for E [d.]. Here, the superscript (i) denotes the i-th independent processing channel, where each channel operates on unique feature dimension of the input and output or vectors (i.e., u/ and of correspond to the i-th elements of in and oi respectively). The hidden state h (1) is initialized as h (4) 0 and evolves according to A: E Rddx Bi Rd x1 and the input ui Cr E BBdxx1 maps the hidden state h, to the output oft)_ As shown in A and B, are computed using the zero-order hold (ZOH) discretization method applied to A E dR daxd B, E Rddx1 and the timestep A E TR. Next, we describe the selection mechanism.
B WBu - by, Here, softplus(x) = log(1 + exp(x)). WB, We E Rd xd, by, E Rd, x1, ws E R da E IR, along with A Rddx xd, are the parameters of the Mamba model. We use to denote the collection of all the parameters.
(3) C Wow + bc,
(4) A =oftpluus(uu w w bba), (5)
3

---

Unlike previous work (Sushma et al. 2024) that introduce local self-attention component to augment SSMs, which may inherit the Transformer's 's ICL ability, our model adheres to Mamba's original selective state-space framework (Gu and Dao 2024). This alignment ensures us to mechanistically analyze how Mamba's architecture enables in-context learning (ICI.).
Linear Regression Prediction. In this work, we set de d 1, enabling the Mamba model to process the embeddings C1:N, Ey. Given the prompt (e1,..., eN; E the Mamba model will output sequence 01:N+1 Mamba @ e1, ........ eN, eq). The prediction for the linear regression target y/ w 2q is extracted from the terminal position of the output matrix (corresponding to the zero placeholder in the query token eq (x. 0) T) Concretely, 9g N+1
(d=1)
Training Algorithm. To train a Mamba model over the in-context linear regression task, we consider minimizing the following population loss:
(()) Exg 2. 10
(0g Mo)²
(6)
Given a Mamba model, we use gradient descent to minimize population loss L(0), and the update of trainable parameters 0' (WB. Wo, bg, bc} can be written follows:
0' + 1) 0 C(e()t)
4 Main Results
This section presents main theoretical results that characterize the convergence state of Mamba and its final loss. We also compare the results with other models.
Assumption 4.1 (1) Matrix A = -- -Iaa-- (2) The vector wa is fixed as zero O, and ba is fixed as
In (c(epp (In 2)/N) 1). (3) Matrices WB, B, We are initialized with entries drawn i.i.d. from the
standard Gaussian distribution N (0, 1). (4) The hidden state dimension satisfies: da = (dd)). (5) The learning rate satisfies: = ((d 2d, ). (6) Bias vectors bg, by are initialized as zero 0. (7) Token length N (2dd).
(1) The negative-definite matrix A - Id, guarantees the stable convergence of hidden states h (2) Given the zero-mean and symmetric distribution of embeddings, wA can naturally converge to during gradient descent, and we fix as 0 for simplicity. We further fix ba to an appropriate constant to maintain a suitable timestep A,, enabling us to concentrate our theoretical analysis on WB. We. bg. and bc. In prior works on Transformer-basee based in-context learning, merging key-query
weights (e.g., W := WWWWR Wx) and specific initializations (e.g.. Wo Wx I) are often adopted
simplify optimization analysis (Zhang al. 2024] Ahn al. 2023 Huang al., [023] Mahankali
et al., 2024 Wu et al., (024). (3, 4, 5) In contrast, our Gaussian initialization of W B and Wc demonstrates more practicality, which requires sufficiently large hidden state dimension da and sufficiently small learning rate 17 to ensure favorable loss landscape properties. Assumption (6) is intended simplify the analysis. (7) Token length should be larger enough than the dimension of ur to capture sufficient contextual information.
Theorem 4.1 Under Assumption 4.1. if the Mamba is trained with gradient descent, and given
new prompt (ei..... EN, eq), then with probability at least ô for some E (0, 1), the trainable to parameters that satisfies:
parameters (1(1) (WB()) (c))), (l),), bc() converge as
(a) Projected hidden state: (W ][a (t)h (d+1) = =a(WT a(W (d+1) + (1 a3)yya1
(b) Prediction for target: 0g x EN' (1 a)a't! ByN-
(c) Population loss: L(O(t)) < 3d(d+1)
2N
where a eex(((-1n)))) (n 8 2(1+a)
23(1-0))++1-o0)

---

Theorem 4.1 characterizes the in-context learning (ICL) mechanism of Mamba and establishes an
upper bound on population loss. Specifically, (Thm (a) shows how the hidden state updated
according the given prompt ey (ax un) (Thm (b)) presents the final prediction given
prompt (ei...., eN, eg). (Thm 4.1 (c) provides the upper bound for the population loss, which
comparable to that of the Transformer Zhang et al. 2024). Next, we'll Il discuss it in more detail.
Update of Hidden State. If we define h. (W [1ndd (d+1) then (Thm 4.1 (a)) be rewritten as follows:
by - (1 Byya hi-1 + (1- a)) (Byrar h
We observe its intrinsic connection to online gradient descent, which updates the model parameters (hr) with only one currently arriving sample (e: ( (1))
at each step. Specifically, the system gradually updates b2 along the pseudo-gradient direction Byra,, with a fixed step size (1 at).
For a newly defined task f(x) w x, given that E[yaæ w, the direction of h, converges toward w as mamba processes multiple prompts. This demonstrates mamba's 's ability to internalize f(x through prompt processing, which ultimately ensures that predictions for query token eg (2r closely approximate (3xgg)
Previous works have shown that Transformer can mimic a single step of gradient descent to achieve in- context learning ability (Zhang et al. 2024 Mahankali et al. 2024). Concretely, a trained Transformer
can be described as follows
N
trrestrrmmeee....ee::: R Dg
y:
2 x w.
Our theoretical analysis reveals that Mamba and Transformer have different in-context learning mech- anisms. This divergence stems from their inherent architectural biases: Transformers process contexts globally through self-attention, while Mamba enforces local sequential dependencies via recurrent state transitions. These findings provide fundamental insights into the contrasting capabilities of Transformer-basse and Mamba-based models for in-context learning. As experimental work shows, transformers can learn vector-valued MQAR tasks the context which Mamba cannot, while Mamba succeeds in sparse-parity in-context learning tasks where Transformers fail(Park et al. (024).
Prediction Outcome. Comparing equations Thm 4.1 (b) and 9 we found both similarities and distinctions in how Transformer and Mamba implement in-contex, learning (ICL). Both models leverage a weighted aggregation of y; aligning with the intuition that learning f(x) from context reduces to estimating the latent parameter w, since E[yea:] = w. Notably, their token weighting strategies diverge: Transformer's global attention mechanism implicitly assigns nearly uniform weights (~ 1 imposes position-dependern
w
N. where N is the token length) to all yıa,, while Mamba's 's linear recurrence weight variations. This difference arises from Mamba's 's iterative state update rule, where the influence of prompt tokens e, on the hidden state h; depends on their sequential placement, governed by the model's linear recurrence dynamics.
The derived upper bound (Thm 4.1 (c)) establishes an O11/N) convergence rate for the loss (ignoring dimension factor), demonstrating that Mamba matches the sample complexity scaling of Transformers in linear regression ICL tasks (Zhang et al. 2024).
Compare with S4. Mamba extends the structured state space model (S4) (Gu et al. 2022) by integrating a selection mechanism, which is critical for enabling ICL. In S4 model, the matrices A E dd Xdn, and B, C' e [R are static, leading fixed linear combination of inputs:
A j=1
(10)
where the coefficients C. A ' are task-agnostic. This formulation inherently limits S4's ability to adapt to tpec-specific parameters w in ICL scenarios, as the model cannot adjust its inductive bias to match distinct w across different tasks. Therefore, the S4 model cannot truly learn in-context.
In contrast, Mamba's selection mechanism dynamically adjusts B, and C₁ (and optionally A,) based on the input tokens (U,,..., u.). This allows the model to implicitly adapt its hidden state to
5

---

align with the latent w of each task, effectively transforming the linear combination weights into context-dependent functions f(x) w. Such adaptability essential for ICL, as enables Mamba to reconstruct diverse w from input prompts without task-specific fine-tuning.
Proof Sketch
This section outlines the main technical ideas to prove Theorem 4.1 The complete proofs are given in the appendix.
Linear Recurrence.
To start with, we show how the hidden states update when receiving token a 31) By (Eq. (5)) and Assumption 4.1 ()), we have A = (In 2).N. Combining it with
Cj
(Eq.
22 and get:
(d+1) (1 o)yy B1, hi (d+1) ahi
(11)
where a == exp(-A:) = (n the second equality by discretization rule 22) the third equality by Assumption 4.1 2) and exp(-
expp(-ii)
Prediction Output. We next derive the expression of Un- By recurring (Eq.
(1 ). the hidden state
after receiving the first I context prompts e1.t is given by h (d+1) =11- a Ei= a² -
Receiving all the prompt tokens e1NN and the query token eg = (at 0) we have:
h (d+1) (d+1) + (1- BN (1-a) N+1 a
Σ
yN yN-
(12)
Finally, the prediction output as follows
(d+1) 94 C N+1 aW bc)
Σ 1=0
it yN WB WN- bb).
(13)
To handle Woe, and WBeN-1. we further decompose WB [Bb] and Wc = [C'c, where
B,C Rdd xd b, c 1d.x1 Then we write another form of (Eq. 13
N-1
0, = (1 (Cx + bc)
Σ a i+1 (Bx u bB).
(14)
The loss becomes:
N-1
(()) = Fz (11(())) + bc) 2
Σ 1=0
a
i+1
Un
( b bB) w wT,z)
x..)) ) (15)
By computing the gradient of C, bc. B. and bg with respect to L(O(t)), we derive the following update rule according to Eq. 77)
Lemma 5.1 (Update Rule) Let be the learning rate and we use gradient descent to update the weights WB, We. by, be, for 2 0 we have
B(t 1) B(t) + m3,C((1 m3,C((1C(()) B()),
C(++1) = C(t) + +33BB(1 ((t) C(t) ()) CC()),
b((+1) b(1) - -1300())))0) b(())
b(t). bB(l) bb(() =0,
where 31 =E D En
N-1 (1-a))) 41++2 2x
82 E N-1 Σ/
2 - 33 = اا (1

---

are the Technical Challenges. Unlike many prior Transformer- based ICL analyses that simplify dynamics via merged weights or special initializations, our Gaussian initialized Gaussian-iittilizzed WB. B. Wc and discrete- time gradient descent introduces more complexity (cf. assumption 4.1). To solve the optimization problem described in Lemma 5.1 we have the following three questions answer: (1) Convergence Target: Where do the parameters converge? (2) Convergence Proof: How to rigorously establish convergence? (3) Saddle Point Avoidance: How to avoid saddle points? To answer these three questions, we propose two key techniques: Vector-coupled Dynamic, Negative Feedback Convergence, and apply them with a Fime-grained Induction. We next describe them detail.
5.1 Vector-coupled Dynamics
We can verify by Lemma 5.1 that C Diag(a)....,,.u with a, E {0, d C fixed points for the parameters WB. Wc.
Combining the loss function Eq. 15 and (t) = bc(() = 0 in Lemma 5.1 the loss function can be rewritten as
(()) - 11 (1 a)
N-1
Σ
a
it1
a
T
T
C - -w w/xa/) )
1=0
To minimize loss, the term (1 a) Σ Di=0 a (ar C ByN- - C b) should ap- proximate w Eq. Given N = w and = > O, we derive that C' B should I. while C b converges to 0 to minimize the loss. However, mentioned above, converge to A
C = Diagg(a; (a....., ap) with partial a; can also enable convergence, which is an undesirable
scenario.
To analyze the convergence behavior C B and C b, we introduce the Vector-coupled Dynamics technique, which studies the inner product dynamics between decomposed column vectors of B and C. Specifically, we decompose B and into B = b ... bu], C (cI...ca) Then we have another form of Lemma 5.1 for B, C' and b as the following lemma.
Lemma 5.2 (Vectors Update Rule) Let r be the learning rate and we gradient descent to update the weights WB, Wo, bB, for [d]. 2 we have
6 1 b,(1) + ( (0)0 31 Σ (1))))(()) (x))
kyi
4
c(++1)= - (() +
(Bs - ((+))==((++((--0)((((00000)00000 B1 Σ bx(1) 3ec (100))))((()) Right
b(t + 1) b(t) - (B2 of (()((() (c(())
With Lemma 5.2 we can further analyze the dynamics of the inner products e, (1)b,(1), of (1) and c. (t)b(t), precisely characterizing the behavior C B and C. b. This technique helps answer the question "Where do the parameters converge?"
5.2 Negative Feedback Convergence
As we discuss Section 5.1 to minimize loss, the following conditions must be satisfied for all i E [d] with # j: e; (t) > 3. c. (t)b((t) - O, of (t)b(t) -> 0. To establish the convergence, we introduce the Negative Feedback Convergence technique. This technique leverages the negative feedback terms in the dynamical equations of c. (1)bx(!), ef (I)b,()), and c, (1)b(l) to derive exponential convergence rate. Taking c, (t) as example, derive the following
7

---

update rule by Lemma 5.2
(8s- 3ic (1. +11+1060 ()+((+++)))) 83 -1c, (1),,(()))
n/ (3s BIc, (t)b.(())) (t)b,(t) - n (Bs Bic (D0(
negative seedback term
(16)
+ n
c. (1)b(((() 6% (1),(() + n18 Σ (1)b(()) oc (1)((()
kyi
+ 12c (1)(22((()b(() bf (1)b()) - 31 (c(t 1 c(t)) 1) (b(()))
The term (33-8ic (t- (++1),(t++1) decomposes into its previous state (33-31ch (t)) (marked with underline) plus the remaining terms (increment terms). The increment terms includes a negative feedback term, which induces a tendency to drive (Bs 3ic (t)b;(t) to 0 (cf (t) - 81
8a
Intuitively, by (1) and c, (t) are much larger than bj (4)),((, c. (t)cz (1) and b (t)b(t) at Gaussian initialization with high probability. Also, as c, (t)b, (t), bj (t)b(t) + 0 with i + and small enough, the effect of negative feedback feedbbck term is the dominant term in the increment terms.
Therefore, denoting v(() = 33 c (() and E(L) v( 1) v(() negatiie feedbuck term we can model the update rule of (Eq. T6) as follows:
y(t + 1) (1 n (b, (1)b,(1) + c (()((()))) v()) E(1).
Recur this formula from 0 to f, we have:
y(t + 1) = II (1 m31 (b. +c (T([((((s))) v(0)
==0
(17)
Σ II (1- n/81 (b. (s)b,(s) t (()))
==0s'=st
Denoting = min(b; (s)b,(s), cj (s) for E [0,t], t). the first term on the RHS of (Eq. I7 can be upper bounded by (1 21/37)** y0). i(' ) has an exponentially decaying upper bound (it can be proved when c, (t),,(t (t) + 0. o (t) + with an exponential rate), the second term on the RHS of (Eq. (17) has an exponentially decaying upper bound. Therefore, we can establish an exponential convergence rate for cj (t),, + & The similar method can be used on cJ (t) - 0, ct (t)b(t) - O. This technique helps answer the question "How to rigorously establish convergence?"
B
5.3 Fine-grained Induction
The exponential convergence of c, (t) + A under the Negative Feedback Convergence frame- A work requires the following two conditions for all i,j E [d] with i # j:
(1) Do. (1)b.(1) and c. (L) dominate b. (1), cJ (1)),(1) and 6J (1)b(L) in magnitude. (2) at (t) - 0, c. (f)b(t) + 0 at an exponentially decaying rate.
On the one hand, condition (1) at initialization (t 0) can be established via concentration in- equalities, and critically, the preservation of Condition (1) for t > O relies on the rapid decay of c, (t)b,(t), c. (t), and c, (t)b(t) (condition (2)). On the other hand, under the framework of Negative Feedback Convergence, c. (t)b; + 0 in Condition (2) also relies on Condition (1) and &s the rapid decay of c (t) + 3, c. (t)b(t) + O. This implies mutual dependencies among the bounds of these Vector-coupled inner products.
To handle these dependencies and establish stable bounds, we introduce the technique Fine-grained Induction: Divide the inner products into three groups: (1) Squared norms: by (1),, c, (W),,(), 6 (I)b(t). (T)b(()). (2) Target terms: of (I)b;((), cf (1)b, (I)b(t). (3) Cross- interactions: by (1),,()), of (1),,()))) (1)b(1). And then carefully give bounds for them with an induction.
8

---

Specifically, denoting ö(t) = c d2 1) (c) (s)), bb (s)b(s)]} and 5d, < min (t)b, (t) (t)b(t)), we establish the fol- lowing three properties A(t). B()), and C(t) simultaneously fort 2 0:
A()):
da/2 s by (t), (t) (()b((() - 2d,-
B((): - (1)6,()) (x)))xp((3mppp
|c| (1)b)(() - 28(1) exp(-n8811)
|c (1)b(() - 25(1) (0(())xp((pppp)) exp((1882t) + 8(t)
B2 exp(-n/rnt).
C()):
(1)6,(()) c (4)) (6T0)).110(0))(()))))) (L)6(() s 8(1) 3 dn (og(44((2)))))) + 1)/6).
The initial conditions A(0), B()), and C(O) are established with high probability by concentration inequalities. We also provide the following claims to establish A(t). B(t), and C(t) for 2
Claim 5.1 A(O), A(T), ( C(T) A(T + +
Claim 5.2 A(0), ................................ A(T), ((T)) C(T)) = B(T 1). Claim 5.3 A(0), A(T), ( C(T) = C(T + +).
This induction answers the question "How to avoid saddle points?" because B(t) guarantees that B La I and C b 0, preventing stagnation of partial diagonal entries of C' B zero.
Theorem 4.1 can be proved by substituting C B 83 I,C 0, bn bc 0 into (Eq. 111.
15
6
Experimental Results
We present simulation results on synthetic data to verify our theoretical results. More experimental results can be found in Appendix
C W
0.6
loss
0.4
0.2
recurrent step
token length N
(a) Trained parameters
(b) Similarity
(c) Loss curve
Figure I: (a) Post-training visualization of matrix product C W (b) Cosine similarity evolution
between and hy (WT) (.dd)..d.))) h (d+11) across recurrent steps (after processing prompts e1.); (c)
(d+1)
Test loss versus token sequence length N. Blue curve: experimental results; orange curve: theoretical upper bound.
Experiments Setting We follow Section to generate the dateset and initialize the model. Specif- ically, we set dimension d 4, d,, = 80, prompt token length N 50, and train the Mamba model on 3000 sequences by gradient descent. After training, save the model and test it 1000 new generated sequences, tracking the cosine similarity between h₁(:= (WT) (.. h (d+1)) and w. Moreover, we vary the length of the prompt token N from 4 to 80 and compare the test loss with the theoretical upper bound. For each N, we conduct 10 independent experiments and report the averaged results. All experiments performed on NVIDIA A800 GPU.

---

Experiment Result Recalling that we denote WB [BB,, Figure la reveals the convergence of C' B to a diagonal matrix and C to O, confirming the theoretical induction presented in Section also consistent with (Thm 4.1 (b)). Figure shows that the projected hidden state h; gradually aligns with w more prompt tokens processed, consistent with (Thm (.T (a)). Figure Ic demonstrates that the experimental loss has an upper bound 3d(d+11 that decays linearly with N, aligning with
2N
(Thm (c))
Conclusion
This paper study Mamba's s in-context learning mechanism, and rigorously establish its convergence and loss bound. By analysing the Vector-coupled Dynamics, we provide exponential convergence rate with Negative Feedback Convergence technique in a Fine-grained Induction, and finally establish O(1/N) loss bound. The loss bound is comparable to that of Transformer. Our theoretical results reveal the different mechanism between Transformer and Mamba on ICI. where Mamba emulates a variant of online gradient descení to perform in in-context, while Transformers approximate a single step of gradient descent. Furthermore, our comparison with the S4 model demonstrates that the selection components are essential for Mamba to perform ICL.
Limitations and Social Impact Our analysis focuses on one-layer Mamba model, thus the behavior of Mamba with multi-layer or augmented with other components such as MIP is still unclear. We believe that our work will provide insight for those cases and can be used to study more data models such as nonlinear features. This paper is mainly theoretical investigation, and we not see an immediate social impact.
Acknowledgements
We thank the anonymous reviewers for their insightful comments improve the paper. Miao Zhang was partially sponsored by the National Natural Science Foundation of China under Grant 62306084 and U23B2051, Shenzhen College Stability Support Plan under Grant 1XXD22023112810022333 and Shenzhen Science and Technology Program under Grant ZDSYS202306660011000000 and KJZDD20230922115113322
References
Ahamed, M. A. and Cheng, Q. (2024). Timemachine: A time series is worth mambas for long-term forecasting. In ECAI 2024: 27th European Conference on Artificial Intelligence, 19-24 October 2024, Santiago de Compostela, Spain leciinincluding 13th Conference on Prestigious Applications of Intelligent Systems. European Conference on Artificial Intelli, volume 392, page 1688.
Ahn, K., Cheng. X., Daneshmand, H., and Sra, S. (2023). Transformers learn implement pre- conditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36:4561 45650..
Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. (2023). What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations.
Arora, S., Cohen, N., Golowich, N., and Hu, W. (2019). A convergence analysis of gradient descent for deep linear neural networks. In International Conference on Learning Representations.
Bai, Y., Chen, F., Wang, H., Xiong. C., and Mei, S. (2023). Transformers statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36:57125-57211.
Behrouz, A., Li, Z., Kacham, P., Daliri, M., Deng, Y., Zhong, P., Razaviyayn, M., and Mirrokni, V. (2025a). Atlas: Learning to optimally memorize the context at test time. arXiv preprint arXiv:2505.23735.
10

---

Behrouz, A., Razaviyayn, M., Zhong, P., and Mirrokni, V. (2025b). It's all connected: A journey through test-time memorization, attentional bias, retention, and online optimization. arXiv preprint arXiv:2504. 13173.
Bondaschi, M., Rajaraman, N., Wei, X., Ramchandran, K., Pascanu, R., Gulcehre, C., Gastpar, M., and Makkuva, A. V. (2025). From markov to laplace: How mamba in-context learns markov chains. arXiv preprint arXiv:2502. 10/78.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. (2020). Language models are few-shot learners, Advances neural information processing systems, 33:1877-1901.
Bu, D., Huang, W., Han, A., Nitanda, A., Suzuki, T., Zhang, Q. and Wong, H.-S. (2024). Provably transformers harness multi-concept word semantics for efficient in-context learning, Advances in Neural Information Processing Systems, 37:63342-63405.
Bu, D., Huang, W., Han, A., Nitanda, A., Zhang, Q. Wong, H.-S., and Suzuki, T. (2025). Provable in-context vector arithmetic via retrieving task concepts. In Forty-second International Conference Machine Learning.
Chen, Y., Li, X., Liang Y., Shi, Z., and Song, Z. (2025). The computational limits of state-space models and mamba via the lens of circuit complexity. In The Second Conference on Parsimony and Learning (Proceedings Track).
Cirone, N. M., Orvieto, A., Walker, B., Salvi, C.,, and Lyons, T. (2024). Theoretical foundations of deep selective state-space models. In The eightt Annual Conference on Neural Information Processing Systems.
Dao, T. and Gu, A. (2024). Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning (ICML).
Du, S. and Hu, W. (2019). Width provably matters in optimization for deep linear neural networks. In International Conference on Machine Learning, pages 1655-1664. PMLR.
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. (2022). What can transformers learn in-context? case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598.
Gatmiry, K., Saunshi, N., Reddi, S., Jegelka, S., and Kumar, S. (2024). Can looped transformers learn to implement multi-step gradient descent for in-context learning? In Proceedings of the 41st International Conference on Machine Learning, pages 15130-15152.
Giannou, A.. Yang, L., Wang, T., Papailiopoulos, D., and Lee, J. D. (2025). How well can trans- formers emulate in-context newton's 's method? In The 28th International Conference on Artificial Intelligence and Statistics.
Grazzi, R., Siems, N., Schrodi, S., Brox. T., and Hutter, F. (2024). mamba capable of in-context learning? In Proceedings of the Third International Conference on Automated Machine Learning. volume 256 of Proceedings of Machine Learning Research, pages 1/1-26. PMLR.
Gu, A. and Dao, T. (2024). Mamba: Linear-time sequence modeling with selective state spaces, In First Conference on Language Modeling.
Gu, A., Goel, K., and Re, C. (2022). Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations.
Huang, Y., Cheng, Y., and Liang, Y. (2023). In-context convergence of transformers. arXiv preprint 10..2.231..05249.
Lee, I., Jiang, N., and Berg-Kinkpatrick, T. (2024). Is attention required for icl? exploring the relationship between model architecture and in-context learning ability. The Twelfth International Conference Learning Representations.

---

H., H., Lu, S., Cui, X., Chen, P.-Y., and Wang, M. (2025a). Understanding mamba in in-context learning with outliers: theoretical generalization analysis. In dign-dimensional Learning Dynamics.
Li K., Chen, G., Yang, R., and Hu, X. (2024a). Spmamba: State-space model is all you need in speech separation. arXiv preprini arXiv:2404.02063.
Li, K., Li, X., Wang. Y., He, Y., Wang. Y., Wang. L., and Qiao, Y. (2024b). Videomamba: State space model for efficient video understanding. In European Conference on Computer Vision, pages 237-255. Springer.
Li, Y., Tarzanagh, D. A., Rawat, A. S., Fazel, M., and Oymak, (2025b). Gating is weighting: Understanding gated linear attention through in-context learning. arXiv preprint urXiv:2504.044 044.085004.44388.
Lin, L., Bai, Y., and Mei, (2024). Transformers as decision makers: Provable in-context reinforce- ment learning via supervised pretraining, In The Twelfth International Conference on Learning Representations.
Liu, Y., Tian, Y., Zhao, Y., Yu, H., Xie, L., Wang. Y., Ye, Q., Jiao, J., and Liu, Y. (2024). Vmamba: Visual state space model. Advances in neural information processing systems, 37:103031-10306
Mahankali, A. V., Hashimoto, T., and Ma, T. (2024). One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference on Learning Representations.
Nishikawa, N. and Suzuki, T. (2025). State space models are provably comparable to transformers in dynamic token selection. The Thirteenth International Conference on Learning Representations.
Park, J., Park, J., Xiong, Z., Lee, N., Cho, J., Oymak, S., Lee, K., and Papailiopoulos, D. (2024). Can mamba learn how to learn? a comparative study on in-context learning tasks. In Proceedings of the 41st International Conference on Machine Learning, pages 39793-39812.
Patro, B. N. and Agneeswaran, S. (2024). Mamba-360: Survey of state space models transformer alternative for long sequence modelling: Methods, applications, and challenges. arXiv preprint arXiv:2404. 261112
Sander, M. E., Giryes, R., Suzuki, T., Blondel, M., and Peyré, G. (2024). How do transformers perform in-context autoregressive learning? In Proceedings of the 41st International Conference on Machine Learning, pages 43235-43254.
Shen, W., Zhou, R., Yang. J., and Shen, C. (2024). On the training convergence of transformers for in-context classification of gaussian mixtures. arXiv preprint a1011:4410.1177..
Sushma, N. M., Tian, Y., Mestha, H., Colombo, N., Kappel, D., and Subramoney, A. (2024).
State-space models learn in-context by gradient descent. arXiv preprint arXiv:2410. 1667.
Tong, W. L. and Pehlevan, C. (2024). Mlps learn in-context on regression and classification tasks. arXiv preprint arXiv:2405.15611
Vankadara, L. C., Xu, J., Haas, M., and Cevher, V. (2024). On feature learning in structured state space models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, E.., and Polosukhin, I. (2017). Attention all you need. Advances in neural information processing systems, 30.
Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. (2023). Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174, PMLR.
J., Zou, D., Chen, Z., Braverman, V., Gu, Q. and Bartlett, P. (2024). How many pretraining tasks are needed for in-context learning of linear regression? The Twelfih International Conference on Learning Representations.
12

---

Yang. S., Kautz, J., and Hatamizadeh, A. (2025). Gated delta networks: Improving mamba2 with delta rule. In The Thirteenth International Conference on Learning Representations.
Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, (2024). Gated linear attention transformers with hardware-efficient training. In International Conference on Machine Learning, pages 56501-56523, PMLR.
Zhang, R., Frei, S., and Bartlett, P. L. (2024). Trained transformers learn linear models in-context.
Journal of Machine Learning Research, 25(49):1-55.
Zhang. Y. Singh, K., Latham, P. E., and Saxe, (2025). Training dynamics of in-context learning in linear attention. arXiv preprint arXiv:2501.1626.
Zheng, C.. Huang, W. Wang. R., Wu, G., Zhu, J., and Li, C. (2024). On mesa-optimization in autoregressively trained transformers: Emergence and capability. Advances Neural Information Processing Systems, 37:49081-49129.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract reflects the paper's scope, and the introduction reflects the paper's contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss our limitation in the conclusion Section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
14

---

Justification: We provide the assumptions in Assumption 4.1 In the main paper, we provide a proof sketch (Section 5)) The detailed proofs are in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The experiments follows our problem setup and assumption. We also provide the experiments setting in this paper, and the codes are in the supplementary material.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
15

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
The codes are the supplementary material. We also provide readme file.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
praips.cc/ubblic/guuueee CoodeSumeeieeeeee ssi iessonPolll mollicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The experiments follows our problem setup and assumption. We also provide hyperparameters the experiments setting.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We plot the I-sigma error bar. Figure Ib the error bar for experimental loss be found in Figure 2c
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
16

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We report the computer resources in experiments setting. All experiments are performed on NVIDIA A800 GPU within hours.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't ' make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics stttps://neuripst
Answer: [Yes]
This work focuses on theoretical study Mamba's in-context learning. All data is synthesized. We see no ethical or potential harms of our work. We will not violate the Code Of Ethics.
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss the socictal impacts Conclusion Section. We do not see an immediate social impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
17

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper is mainly a theoretical work. poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification:
We perform experiments on synthetic data. It does not use existing assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
18

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, batasets.papersstt paperswi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: We perform experiments synthetic data. No new assets introduced in the paper.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/code/mooel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsoureing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
19

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: This paper mainly a theoretical work. The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
20

---

Appendix
Contents
Basic Calculations
22
A.1 Data Statistics
22
Output, Loss, Gradient
24
A.3 Training Dynamies
25
B Proof of T'hooree 4.1
27
Complete Proof
Proof of Lemma A
32
Proof of Lemma A.3
33
Proof of Lemma
36
4
Proof of claim B.
44
Proof of claim B.2
50
Proof of claim B.3
57
Bounds of 17 terms
63
Discussion
66
Additional Experimental Results
67
21

---

Table 1: Key notations
Symbols
Definitions
R; Pq: Mr, Mq
2.y w are i.i.d sampled from Gaussian distribution (0, Ia). yi
2t. Vo w Zy.
((( (b, 1) (()))
(t))),
(t (c(t+ c(t)).
b(t) (b( + 1))++(()))
Decompose the matrices Wa, We into colums of vectors: WB = [B] = [b...... by b). We = [Cec [c....,, Ca c)
b.C.b.cb..c
where Wa, We E Rdh ((d+))
B, B.C
b, c, b., C2 E Rdxx1
b,(t) (t(t), ct) (t) b(t) (, b,))), (t) b(t), ((t))
b(t)
inner product of the vectors b. b, with i, 1 dl. c.g. (t) (t))) the inner product of b, (t) and (t).
a
factor, == exp( AI) = exp(( In 2)/N).
The factors appearing the gradient equation. Specifically, B1 (a a 2²²²+ + (1-a)
81, 92, Bs
(1+x)
B2 (daa a1 + (2d +ed)a (1-) 83 = a(1 xa
(1+a)
The lower bound of squared norms b. (1),,()) (t),
and 1b (t))((t). Specifically, = ]d..
The upper bound of cross-interactions: 6, (t)b, c (t), and b," (L)b((). Specifically, 5(1) ma (og 1)/0),
6(T)
|b, (s)b,(s)), le (s)), (b. (s)b(s))).
A
Basic Calculations
This Section provide the data statistics related to gaussian distribution, and compute the expressions for the output, loss, gradient, training dynamics (particularly Vector-coupled Dynamics) of the Mamba model. Section presents the Fine-grained Induction with Negative Feedback Convergence technique, and finally establish the results for Theorem 4.1 Section ddttails the complete proofs for Section A and Section B In Section D we discuss about orthogonal initialization and compare our framework with other techniques. In Section E
we give more experimental results.
A.1 Data Statistics
Lemma A.1 (Concentration Inequalities) Ler (,))) be the i-th colum of B(0), c,(0) be the i-th colum of C(0), and suppose that > 0 and dp = 20log(((222))) + 1)/5)), with probability at least
1 - 8, we have:
3dp
(,) 6, c (.) 6(0) b(0) - 5d
b.())
(((0)Tb,0)) (,) b(0) 2 du log( Ad (d(2d 2d + 11/8),
(,)0TT,,000 (, (.) T b(0) < 2 dn log Ad 1)/0)
c,
for E $..
Proof of Lemma A.1 By Bernstein's inequality, with probability at least - 8/2 + 1) we have
b,(0) b,()) - dn (Vdd, (4( +1)/6)).
22

---

Therefore, as long as d, = 2[loo ((log(((2++1)/))), we have 3d,/4 b.(0) b.(0) - 5d,/4. Similarly, we have
b:(0)bb,(0)
3dh (0) b(0) ((0)
5d
For i,j E [d], + j. By Bernstein's inequality, with probability at least - 6/2d(2d 1), we have
(0) (.) by (0) c, c,( c,,0T,,,,0)) b.(0)T(())
(.)( (,))) c.(0) c.(0))TTT) [((0) -
d, dd log(4d(22 +
s $ da log 1)/6).
We can apply union bound to complete the proof.
Lemma A.2 If vectors a and ur are iid generated from ((0, Ia). y x w we have the following expectations:
E
xx
"ww"
x::
(d + 2)I,
EE y²
= d,
E
3d(d + 2).
The proof of lemma A.2 is in Section
C.1
Lemma A.3 If vectors x; and u are iid generated from (0, x. have the following
expectations:
N-1N-1
Un -N- = a (1- N\ (100--]]][0100000000000 a)²
(d+ 1)a² (1-22)))) ) I.
E
Σ
2+ナ+2
(1 a)(1 aa)
1=0
f=0
N-1N-1
++++
UN VN- -
_N =0,
1=0
j=0
da² 2 - a. v² 2
(2d + 6d]aa (1 [ 2N
N-1 N-1
E
a
UN -iN
2
(1 a)²
+
(1- 111000
1= 0=
1 - a
E
UN -N-
1 I,
1=0
N 1
E
i+1 VN W
= 0,
1=0
N
(4
UN
:N- W
= 0,
1=0
UN-1 DN-1
N 1-11
E
Σ
i+5+2 - -j =
1=0
1=0
N-1N-1
da² Q2N
α UN UN-1WNN -1] = (1
345+2
1=0
1=0
E
a i+ WN-
= 0.
1=
The proof of lemma A.3 is in Section
C
23

---

A.2 Output, Loss, Gradient
This section we derive the output of Mamba given sequence fei.N, eg), and establishe the loss function formulation with its gradient expression.
Linear Recurrence. To start with, we show how the hidden states update when receiving token ej (x)".)) By (Eq. (5)) and Assumption 4.1 (22 and get:
have A softplus(I((xpp((M)) (In 2)/N)
1)) = (In 2).
N. Combining with (Eq.
h (2+11) A (d+1)
+ Biy
A 1-1 + y/ (A(A) (exp(A)))) B,
(d+1)
-ex((-Ai) (d+1)
1-1 y1 A (exp( - I,
(18)
exx((- (00(-4,)bb 4,) (d+1)
A (1-exp( (xp(-A1)) B
=ah (d+1)
Bi
where a := exp( ( A.) exp( (- In () second equality by discretization rule 2 the third
equality is by Assumption 4.1 2) and eapl -,,) = exp( I. (Eq. (18) is similar to theorem
in Gu and Dao 2024)
Prediction Output. We next derive the expression of Dq. Based on (Eq. 11 ) the hidden state after receiving the first / context prompts en is given by:
h (d+1) = ah (d+1) + (1- a)
a (d+1) (1 a)ı B + (1 1
(19)
H
= (1 - Σ
o yt y B
1
Receiving the query token eg (x T)T, we have:
(d+1) h N+1 oh (d+1) +1 a1 a) BN = (1-a)
Σ
a ++ yN BN
(20)
Finally, the prediction output is follows
N-1
0g = C h N+1 (1 - a) Σ
UN- BN-t
1=0
(21)
N-1
= (1 - a) + bb))
Σ
a
i+l
un
(WEeN=t
1=0
To handle Woe, and WBeN-i, we further denote Wp B
Bb and W [Ce], where B, Ce
[ddxx xd b,c pdd Then we write another form of (Eq.
N-1
ig = - a + bc)
Σ
-N- BBEN= UN- b bB)
(22)
The loss becomes:
( 2 (1- bc)
w yn + b by) w wTxz) xx)
Σ
(23)
The following lemma provides the gradient of B, C, b, bg, be with respect to loss (Eq. 23

---

Lemma A.4 (Gradient) The gradient of trainable parameters 0' = (B, b,b.b,bbblbb with respect
loss (Eq. (3)) are as follows:
bg ((0) 0,
V bc (()) = 0,
VBC(0) (1=(1-oN)) + (d + 1) (1 - a) - a22)
CC B - 201(11 -a(1-a) c
(1+ a)
VoC(0) (d - (1=22(1-a\)) 2 (2d2 + (1)a - a) aNN)
CC b,
(1+a)
:=
VcL(0) = (12(1- + (d (1-
BB C
(1+a)
= (2d2 + (1) - 22
( N)²
bb T
(1+a)
a(1 - N) ) B.
:=9g
Here, we denote $1 = (a²(1 - 1a + (1)0)"2 32 = (do2(1 ani2=
(Ita)
(2d +6d)a (1-a) (1-a28) B3 a(1 a for simplicity. The proof of lemma A.4 is in Sec-
(I+a)
tion
A.3 Training Dynamics
With the gradient in lemma A.4 we further provide the update rule for Mamba's 's parameters
and the Vector-coupled Dynamics.
Using gradient descent algorithm O' 0'(t) oo with training rate ry. we have the following update rule base on lemma A.4
Lemma A.5 (Update Rule, restatement of lemma 5.1 Let n the learning rate and we use gra- dient descent to update the weights WB. Wc, by, bc, for 2 O we have
B(t +1) = B(t) + 1)3g((t) - n3) B(t), ((++1)=C() + 1) n/ (() C(t) - -182(())6()
b(t + 1) = b(t) (3²CC(t)C(t) b(t).
bb(t) bc(() 0.
We decompose and as B= ... bul, C |c ... ca). and provide the update rule for b,, C; and b with E [1 d as the following lemma.
Lemma A.6 (Vectors Update Rule, restatement of lemma 5.2) Let n be the learning rate and we gradient descent to update the weights Wp. Wc. by, be, for id] t Z 0 we have
bt
+ ( bt(()
Bic, ((),(()))(() B Σ cl (()).(()-cu)) (t)
25

---

=: btt + (t,
c(++1) ct) + 17 (3g BIc (t)) B1 Σ
c - B2c; (t)b(t) b(t)
=: =1)(()) + (1)
b(t + 1) ==++))((()) b(t) - -n(B2 Σ of (t)b(t) -x((t) =: b(t) + nb(t)
k=1
Here, we denote (l) = - (1,)) for simplicity.
- (()))))c(())aaa)) and n)b(1)
b(l 1) b(1)
Next, we provide the dynamics for the inner products of these vectors.
Lemma A.7 (Vector-coupled Dynamics) Let n be the learning rate and use gradient descent to update the weights WB, Wc, ba, bc, we have
b. (1 + 1)b,
=.)(())(()) 2n (Bs Bic, (T0 B1 Σ (<T()))(0)))
K
+ +00lb:
6, (t+ 1)6,
= (t)(()((())+)) + (ic))))))C(())) + (0000000000000000000000000000000 ((),b.(()) of k
2 Σ
(1))())) - + 1 (1)
-
(ta(t+))
d
= =c (t) + 2n 2n( (80-- 3 (+))((t)) - 3, Σ (cT (t))
- -2((T(())(0)) + :/|la.0)|
+(
=()((),((t + - ((()))))) (t)b, +2) +2 (1) (1)))()
[e (1),(() + (1)),()) 23 k
Σ
c (T)b((() oj (1)
282cT(1)6(() (1)(()))))))))))))))) ((6(())
(1++1))b:
(1)6:(()) " Bic (1),(()) (1),(()) - Σ
c (()b((1) d% (1)b(()
- -Bc (f)b(t) (f)b(t + (Bs 3ic (t)b, (t))()
26

---

31 Σ c (t)b.(t) (())b(().cc((())((())))))))) + (1000)))(((())))))(0))))) (06
Right
(t(b+)))((+++)
= -11 -1--131 n/31 (0000()((()==+000000))00 (1),(1)))) (1),(1)
+n( 31c (1)6, - 3 Σ (t)
kzin
/2cT(1)00 bj (1)b(1) + (Bs - - (1)
BI Σ (1)) cf (1)a((() + (1))()
d
67 T ( 1)b(t+ 1) (t)b(() 2n( Σ
(cl (+))(tt))) +n² [(t) 2
R=1
(1++1))(11 11)
d
= 6, (1)((( + (3s Bic. (100)))))0)) - 31 Σ cf (1))
32
c (f)((t) c (()b.(())) + t, (f)b(t)
<. + 1)b(t 1)
= (1 n12 (b_ (t)b(t) + c. (f) (t)b(t)
7 (8g fIc (1) (1)(()) B1 Σ cJ (1)bb((t) bi (1))(1)
k
32
of (1)((() af (1)) + ., (1)((1
Lemma A.7
derive by calculating the inner products of the vectors update rule in lemma A.6 For (t + 1)b(t + 1) is derived as follow:
example, b
b 1)b(( 1) (b()) + nb(1)) T (b(l) + nb((1)
= =b'(())( 2n)((1) T b(() + bco1
(()(((( 2n (12 Σ
(c) (T)(())))) + ,² |(t0)]
k=1
The other equations are similar it.
B Proof of Theorem
In this section, we present the framework of Fine- Fine-graiedd Induction, and establish the results of Theorem 4.1 after convergence.
27

---

Fine-gained Induction Specifically, denoting 8(t) e0. (s)b, (s)c, 6 (s)b(s))}
and r m (t)b,(t), cf (t)c, 6. (t)b(t)), establish the following three properties
A(t), (t), and (t) simultaneously for t 0:
A(t):
dn/2 s by (t)b. (())(t))66(())))) S 2dh
B(()
|33 Bic, (1),.(() S 8() exp(-11174)
1c (1)b S 26(1) (113111))11)))
5(t)
1c (t)b(t) & 25(t) exp(- +
exp( (xp(t))31nt)
B2
C(I):
16, (1)6,())))) (O),((), [9, (O)(() > 5() 3 dn log (d 6max
Here, i,j E [1,], * j. The initial conditions A(0), B(0), and C(O) are established with high probability by concentration inequalities (lemma A.T establish A(t). B(t), and C(t) for t 0:
We also provide the following claims to
Claim B.1 A(0), A(T), B(0),
= A(T + 1)
B(T)), Claim B.2 A(0), (C),) C(T)) = B(T + +1)
Claim B.3 ...0)...... (()
C(T+1)
Remark. Property A(t) establishes the stability of quadratic norms:
(t))))(()))))))))))))))))))))( min (t) bt)t), ca(t) (t),,, b(t) b(t)) 2 d2/2.
This norm lower bound induces two critical effects:
1. Convergence Rate: As
in property B(t), The upper bound of cj (t),, c. (t)),(t) and c, (t)b(t) is related to y (lower bound of the squared norms), thus the stability quadratic norms ensure a stable rapid convergence rate for property B(t).
2. Saddle Point Avoidance: The strict positivity (> 0) of |b, 12 and |c,|² prevents the dynamics collapse to undesirable solutions b, = C; = 0, which would permanently make c, b, = (saddle points).
Property ((() establishes rapid exponential convergence rate:
33 I. B1
C
B
bb+0
The rapid convergence rate ensures that the variations of Squared norms (in property A(t)) and Cross- interactions (in property C(t)) remain bounded, thereby establishing their constraints. For example, at initialization, by bz (0) is bounded by 3d,/4 /4 < 6, (0)b,(0) < 5dx/4. /4. Further, thanks to the exponential convergence rate in property B(t), we can prove that bb. (t)b,(t) b, (0)b;(0) & d,/4, and therefore d < b.T (I)b,()), of (1),, b (1)b(I) S 3d,/2 < 2d,,
Property C(t) establishes the upper bound for the iooss-interactions. As we discuss in section 5.2 the Squared norms ( c, (t)),(t), c. bt)b,(t) (t) and c, (t)b(t) ) are larger enough than the Cross- interactions b. (t), c, (t), and b (t)b(t) ), we can make use of the negative feedback term establish an exponential convergence rate. Thus property C(L) is also important.
The proof of claim B.T. claim B.2 and claim B.3 are section C.4 section C.5 respectively.
and section C.6
28

---

Proof of Theorem 4.1
After convergence (t + 0), we will have C² B
L II,C - 0 (by
property (()),, and by( t) (t) (by lemma A.5
We will restate some equality for ease reference.
Linear Recurrence (restatement of (Eq. (18))
ht ahi-1
(d+1) (d+1)
+ (1 - ai
(24)
Prediction Output (restatement (Eq. (22
N-1
0g (1-a) bc)
Σ
a
yN UiN=iBeN-4 +
(25)
1=0
Loss (restatement of (Eq. (3
N-1
((0) [ (L-a))CCx, ()- bc)
Σ
a
ww-(Bex-+tw-bbb bB) w )² (26)
Based on (Eq. (2 ), we have
(W ((+1) (W + a)nn ( ( (t)..(()B)
(1))
= (4+1) ( (t) ybb(t) + bB()))
= =a(W)) (d+1) + a) (t)B(t) (1- (t)b(t)
(27)
+ (100)++111000 (1-a) 33
a(: (d+1)
y!
31
2(1 + a1 - a)
(d+1)
11
a (33(1 add + 4-- 20x
where the second equality by selection rule B Wnne bn (Eq. 3 and Wn B b].
e: (z) T. The third equality by bB(t) 0. The fourth equality by C B
8a
I and
51
C = 0. (Eq. 27 establish the first equation (Thm 4.1 (a)) of the Theorem.
Based on (Eq. 25 have
N-1 y t bn)
8. = (1 - (CC + bc)
Σ
1=0
x T
(1 - a)a (N
N-1
N-1
Σ
(1 a)a
UN-IC B + x
Σ
(1 a)a 41 2
c (28)
N-1
(1 aja" i+1 Bz
4
31 UN-
2a + a)(1 a)
UN DN-EEN-i
(3(1 )d 4 2o))
where the second equality is by (t) (t) The fourth equality by CTB B b T
8n and
8,
0. (Eq. 28 establish the second equation (Thm 4.1 (b)) of the Thcorem.
29

---

Based on (Eq. 25
we have
N-1
((
1
(1L
()(Cx + bb)]
Σ
+2 (x- bR) w w(zz) )
E
33
N-1
Σ
(1 - alaaaattttte
yN- - w xn)² )
2
31
1
1=0
83
N-1
i+ yN- )
E
(1 a
2
81
1=0
(29)
N
33
E3
x
(1
a)aa
UN -N-- w wax,)
+
E
w
2
d (by lemma A.2
We compute terms and + as follows:
Bs
a) yNN-NNNNN )
E
2
N-1
N-1
B4
(1 v)a
i U/N N
(1 x)a i+
U/N (oN=) T mig)
89
1=0
1=0
(1 a) 2,332
N-1 N-1
Exe, 1g
Bay
E
j+2
yN -1 -9/ - N 22
8
1=0 j=0
(1
a)233
(a²1 (1²
1)a² (1 2²N)
E x; Ixg
B7
(1 (1-a)² a)²
(1 - a) a)
d33
B1
N-1 For the fourth equality, ExN- ΣNO ビ== a a UN-N-N-NNN--- aa² <23 (1(1-)²²²
[1-a) + (d+ 1)+ (1-))2(11020)) (1-a2M)
I by lemma A.3 The last equality by 31
(1-a)(++o)
a² - a 22² + (d+ (1-a) (1aa2
(1+a)
N
=
E
11
83 T
2y
Σ
(1 - 1l-aaatt2mm
-N- +ia
w xgg
(1 33
x. T E E
Σ
a i+1 yn
yN x
B
1=0
(1
83
1 1- N..
x 1xq
B
1 a
d
31
For the third equality, Exy-..w last equality is by B3 a N)
En
Un uN
= =a -1-00² I by lemma A.3 The
30

---

Substituting and into (Eq. 29 and get:
d
d
1
(())
=d
2
31
2
2
31
d(d
+ 1)a² a) a²
2(1 + a)
(30)
a (1- (1-a2N) 22N)
= d(d + 1)(1 a)
31
2(1 +
d(d + 1)
3
.4.
gg
N
3d(d + 1)
2N
Recall B1 (a² (12(1-1))) (1-a-) 1² (d+1)a²(1aa)) (1-2))21))))) (1-a2) (I+a)
333 aa1 aa and exp((-In2)/N))
For the inequality, 1 -
exp((- In 2))
1n S
N
N
8 2 a²\)² =
(1-02x)
1 2N 1 3. thus d(d + 1)(1 a) d(d+1)
a²
44
S 8(1ta) S 3. (Eq. 2(+aa
A1
30 establish the third equation (Thm 4.1 (c)) of the Theorem.
C
Complete Proof
This section presents the complete proof for the above results. To begin with, we provide the exact assumptions for N, 17 and d, as part of Assumption 4.1
Assumption
2 2
3(d + 1) In2
N = R(d) 2 max
1n ln5'
2
In 2
7 = O(d (dld))d%))) -
2d <
32dh
dn ((d)) 2 max . AnA
where
A = (1728 logg8loo(((((200 + 1)/8) + 576(d - 1)31 log (dd(2 /1/
A2- = (576lo + 192
+ 1)/8))/31
Aa 077
1)/8))/B1
\ = 576 17000g//1/2221111//6/40000001010
1)/5)
+ 384 4d (dd + 1)) + 3840 2 1)/8)
As 2448d
24 A6 = 816d 816d 1)/6) 768 log (2d + 1)/6) 48 +))60)
10 63, + 12m33/d) 2
d7
+ +24ylog(4d(2dd 1)/0) 11)
Tog( 4d(2d 1)/8)
8 3 +8 2) 6 + 12) 2
A 36
bg 36 (d((44((dd))))) + 1)/6) [4(d 1) + 56d In 2)²
?
Auo 36log(4d((2 + 1)/0) (6 43 - 1) + 2(d 2(d-2) -1) 11)
31

---

36
32d 31
883 B1
8
4 + 80 In
A11
(l)(3/2)2 (og(4d(2d + 1)/6)
+
32d +
+
31
3182
Note that under the assumption of N (xp((-In2)NN] we have the following:
max
2ln2
mn and combining a
In 1-5-
5 x a² 431 S 82,
5
2 31
S
15
2
$ 82,
83 = (()
6
24
4
These condition will be use to prove some bounds.
C.1 Proof of Lemma A.2
Lemma C.1 (restatement of lemma A.2 If vectors a and w are generated from (0, Ia), y w we have the following expectations:
E xa
www a. = (d+ + 2)
IE y² d,
E
y 3d(d + 2).
Proof. For (i.j)-th element aw "ww" wx
we have:
E
uru ma:
IF
(1) [111]
[R,5]
21
According the distribution of w. have I. wi [1) defined as:
Eut. where dei is the Kronecker delta
[1 if
5kt
0 if +1, 1
By Isserlis Theorem, we have:
E (p IE +
+ E3 2 20
EB (
EE x + EE (
= Eudis + duikk + 5c
Then we have:
uow" ma:
(6shrr + 5ucks + 61g621) 64
(2ō + d1zōkR)
k=1
= (d - 2)0
32

---

Then we have:
E
www xx T
(d 2)
E
y² = -E x:
w x T w
R
(رار)) ا [1)
Σ
ΣΣ5
1=
= d
E
2
E
x
w
x
T w
x
T w x T
E
( (
(2))) wis)
( mn
d
E 11)))
(i
(oudr + Ouóns + desduz) 2
(dudrs + &duōei + 61soml)
Ejjkel
1)
= d +3 d(d 1) 12
= 3d(d +2)
C.2 Proof of Lemma
Lemma C.2 (restatement of lemma A.3) If vectors x, and w are iid generated from N Ia).
y x. we have the following expectations:
N 1
N-1
1a -
(a=(1--0)) N)² (d+ 1)a (1-2(11) )
E
Σ
Σ
c+
-N-
I. (1 - a1=)(a++a)
(1 - a)²
1=0
N-IN-1
IE ΣΣ
NN -3y - =0,
1=0
Ra² )oa²²² 2
(2d2 + 6d]a (1 2N
N-1 N-1
E
Σ
a 2 2
-N--UN-j
+
(1-
(1-a)
0 1=
N-1
=a 1 - N
Σ
a yN "N=
I,
1 a
33

---

N-1
E
a 2+1
UN w
= =0,
1=0
N- 1 N-1 1
E
a +
x
N
2
= 0,
1=0
1=0
E
Σ
Σ
2
WN- =0,
N- N-1
da 2N
E
WN-J
a1 + +a)
0=
1=0
N 1
E
a
3+1
-N-WW w
=0
1=0
Proof.
N-1N-1
E
Σ
44 2
UN UN-
1=
--11
E
++++2 EN -i N. w
-ン T N--
1=0
WN-i
WN-3
Σ
E
N N 2
ww
E
CNN
ifj
++/+2 R
10N-1001 N ww ON NN-jaut j N
T
(d+2) 1, by Jemma A.2
N-1
N-1
(+1)
E Q2o+z)
-
1=0
0
Σ 21+2 (d 2)I
a² (1 aN N) 2
(d+1)a (1-1))(1-00M))
+
(1 a)²
(1- (1-a)(1+ + a))
Here, Eie odts+2=
E otth1)
NN 221t2) for the third equality.
E
a 2+++2
UN Nー voN
i= 0=
N-1 N
2 W IN - w N
x
DN-j
MN-S
0
Notice that appears three (odd) times in the second equality, if we define a function g(w) EN b N- , we can see that g(- g(-w) = -w(),, and further E, [ 0. Therefore, the above expectation equals to O. We will use the similar property in some of the following equations.
34

---

E
UN NJ
1=0
a
R
12 U + Σ
a
B
UN-1
N-1
N-1
N-1 (Ea
=
((
Σ
a
att1)
a 2442 2
a22222
3d(d + 2)
1=
==
1=0
La
1 a
2
(2d + 6d) 2 1 - a2 2N
+
a)2
(1
a) ox)
where the second equality is by E
EF
yN
d and EE |ト
3d(d 2) (lemma A.2)
E
i+1
UN 1 N-
[+
-N- N ww
NN
N-1
E
nN 2 N
E wwp
1=0
N-1
a i+1
I
(=0
I
N
E
y
yN
N=
E
w
w
2
N
w UN-
1=0
MN-1
0
Notice that w appears three (odd) times in the second equality, thus this expectation equals to 0.
++5+2
a
.N NNN-
1=
+
T N N w N 200
Σ a +4++2
xN- JT
E wxx Jy - w
0 + 0
Notice that - appears three (odd) times in E. .N Jo and x T- appears once
(odd) in wx - _ thus this expectation equals to 0.
35

---

N-1N-1
+ UN UN-IUN-J
j=0
N-1
Σ +
E
I nN-, w
NN w a N w
BN-J
Notice that w appears three (odd) times in the second equality, thus this expectation equals to 0.
N-1N N-
UN UN-
1=0
1=0
+
+
IJN 1IN--110 -y
yÑ-i
N-1
Σ
H
NN 2 N
á 21 2T
UÑ
itj
1=0
MN-1 MN-3
2 (1 Q2N
(1 + a))
Notice that v appears once (odd) in E N wx w
where i # j. thus
C yN JNN--NN 0. Moreover, y by lemma A.2
-,
++1 E
w
w
xN
yx
++1 E
ww
E xN--
C.3 Proof of Lemma A.4
Lemma C.3 (restatement of lemma A.4) The gradient of trainable parameters 0'
{B, C, b, bB, bo) with respect to loss (Eq.
(23)) are as follows:
Vb C(O) = 0,
V by (()) 0,
VBC(0) = (a²(1 v022 + (d + 1)a a) 2N
CC B a(1 - c
(1
:=8
:=
VoC(0) (i2 (1-a N2 (2d2 + (1
a2N)
CC
(1+a)
:
36

---

Vcl(0) = (a²(1-a\)2 + (d (1
BB c
(1+
= (1a2(1aa1) (2d + (1-a)
bb C
(1+a)
:=82
a aa)) B.
Proof of lemma
C.3 Recalling the loss:
N-1
(()) (11 (1-a)) + bc)
T
yN yN- yN bR) w wTxx)) )
UN
We will compute the gradient of { BB, C, b,, bo} with respect to L()). Some expectation calcula- tion are detailed in Section
N-1
Wh 0(0) E (1(1)))) a) (C)) + bb)
Σ
a 11
UN (BeN -n- bB) -w [ng)
-N- yN-: (BNN-t + y
(1 2E vo (Co, + bc) (1 aB vw x
+ +(1 az) +11-a) (1 E vo c Tg
2
vu T
bc -(1 a11 vw T B T
It is clear that E 2 = O. Thus, if be = 0, then Db [()) = 0. Notice that we assume bc(() at initialization, so by induction, bc(t) = and Vb (()(t)) = 0 for t > 0. We will consider be when computing other gradients.
N-1
V (()) = - (1) + bc)
Σ 1=
a
i+1
yN- BBE + UN- +y bB) w x.g
yN
N-
(Ca, + bc)
Σ
a
yn-i
-E(U-a)( a) (I=a) (C-)
Σ
n+ (N bb bR) w xc,)
N-1
Cx.
Σ
a
a yN-1
++
1=0
N-1
N-1
= = 2
x
T
T
yn (BxN b8)) Cz.
Σ
a
+
yN-
1=0
1=0
N-1
(1 - a)B
w Bg Cr²
Σ
a yN-
1=0
da (1
CC bg
(1+a)
37

---

The last equality follows from lemma C.4
where we have: E x CT
Limo a yN- (Noi
doi 1-a 2N
yb=;b + br)) - CB. Σi=0
N-1 a y
CC bB, and E BlwTxr Ca. (1-a)(+aa)
En yN- = 0. Similar to be, notice that by is initialized as O, thus by induction, bg(t) 0 and ((0(t) = -0 for t 2 O. We will consider b = b when computing other gradients.
N-1
Vcl(0)= = E[1 -0) (1)-a) (C- + bc)
Σ
a
i+
+ -NN + bB) w T (c,)
==0
N-1
Σ
-N- (BEN-I Yn b by b0
N-1 CT
(1
a)
F
a+ yN (BEN- + uN UN-
N- 1
BBEN- + yN-:bi)arr)
α yN
1=0
N
- (1 - -a)BBB
w T Bq
Σ
UN
(B + -NN-::-NN,byp) ub),))
1=0
(d 1)
(02(1 - aa²²++
BBTC C
(1+a)
+ (&a (1-1111a) N2² (2d + 6d)a (1-)(1 - 22N) -12
bb a
(1+
aa1 - v B
The last equality follows from lemma C.4 where we have:
N-1
N-1
a yN= + yv-
Σ
yB -NN
9
a (1 -1-aa )²
(d+ 1)a² (1 2N
BBC c
a)²
(1- a1 + a
La? 1 - 11aav)
2
(2d Ed]aa 1- 2N
+
bb=c C
(1 a)
(1 (1- a) + ax)
and E w xg. Σm yx =a(1
11-a" 1-a
B
N-1 CL) (110)=B(0 B - (a)(aa)(-))00) +c)
Σ
i+1
yn-
(BaN- + yN bb) w x.
N-1
(Cr, + bc)
Σ a
itl
y
38

---

N-1
N-1 UN- (BoN + un- Cx.
= -1-a)) 2 z CT
Σ
a
++ UN
Σ
a
++ yN- yN-
1=0
N-1
(1 - a)B w
T 2q
Coa
a 1-
UN yu-
1=0
(o (1-- (d+ 1) a1 - -22m)
CC B
(1+a)
- aa -a(1-an)
9
The last equality follows from lemma C.4 where have:
N-1
N-1
E
C T
Σ
a
i+1
UN-i (Bxx yN- Ca. Σ
a it yN-
1=0
-a²(1-a>))
a² - 2²12 a
(d (d+11) 1)22 - (1-aMM)
a1 (1
CC
(1 a1- +a)
and E w Tu xq Cr, En= /اtl UN UN- = a
1-0
I-a C.
N-1
()( + (1 (I- (1-_))-=((- bc)
Σ
a
(BNN - w xzg
i+1
1=0
(Oxy bc).
Σ
a
++1
VN
1=0
N-1
=E((1 -E((1-))(1)
(1-
Σ
+1 (BxNN + b)) -w Zzg
N-1
Cx,"
Σ
a
i+1
VN -1
i=0
N-1
N-1
= -11-a)2p)
T C T
Σ
a
i1
Un UN
uN (B
(BEN
Cz.
Σ
a ++ yÑー
d=0
1=0
(1
a)l w Tyq Cx
Σ
a
i+1
yÑ-1
(2d2 + 2d)aa (1 a) 2N
(1-2
cC
(1 ++a0 a))
The last equality follows from lemma
C.4 where have:
N
N-1
x
O UN
+ Cx.
Σ
a
it y
1=0
a2 2
1 - v
2
(2d2 + od)a (1 (1- 11022)
a2N
(1-a) CC
+
a))
and E w Taa xg Ca. Eie a - 11 =
39

---

C.3.1 Auxiliary Lemma for Lemma A.4
Lemma C.4 If vectors Da By and w are iid generated from (0, Ia), x.. u we have the
following expectations:
N-1
N-1
x
CT
yN yN BBE + uN-ibi)
Σ
a
+1
yN + DN- uN- o
UN
==
a² [ a )
(d+ 1)a² (1
BB C
(1 al a)²
(1 a) a)
La 1 ax)
2
(2d + 6d)a² 2N
bbb C
a1)
a1-
N-1
1- a.
E
Fq
Σ
+ ( tvx
++
=a
B
1
i=
N-1
N-1 UN (BzN- yN- Cz
E
2 9
C
Σ
a
1+1
Σ
a
++
yN-
yN
a22
N) 2
(d+ 1)a²( (1 - -22N)
CC
(1
a)2
(1- +a)
N-1
1 1- N
w
T xq
Cx
Σ
a 1-
yN
yN
=a
C
1
1=0
N-1
N-1
E
a C
Σ
+1
Un (B Caa
Eatt
1=
i=0
d'aa
1 - a
(2d2 + 6d)a² 1 - 2N
+
CC b
(1 a))
(1 a1-a)(1 +a)
N-1
E3 sww xy
Cza
Σ
a yN
N-1
x T
Σ
a
i+
u Cr.
Σ a
i+
JN
1=0
da? 11-1a0 2N
CC by
(1 + +
E w B.
Ca.
Σ
yn
Proof of lemma C.4 We will use the results of lemma
A.3 prove the above equation.
40

---

N-1 1
E
a
CTC
a
i1
yN UN-i( +
1=0
N-1
i+ UN + xz
++ JN- (BNN:
N-1
Σ
1+ UN B
2/3
1=
N-1N-1
E
+ (B -
i=0
(UN-) C T
B + +UN bb b]])] CB 1
N N-1
BB
2 UN- BTC
1=0 1=0
N-1 N-1
+ BE
24j+2
D
N-IN-1
bE
243+2 UN yN-JVN BTC
N-1N-1
+ E
Σ Σ
i+i+2 2 2
UN -N- bb C
1=0 1=0
a a2(1 )²
(d+ 1)a² (1 (1-22N a2N)
BB
(1 a)²
(1-a) a)
pa? (1 - aN
(2d2 + 6d)o 2 - 22N)
bb TC
(1 a1=a)²
(1 (1-a) a
The
last
equality
follows
from
lemma
A.3
where
we
have:
E
Σ i=
===0
N-1
a a+ UN iNN ,,N - N-
2² I-a
(1-
+
1jo (1-228) I.
(1 (1-a)(1+a)
E
Zi=o
N-1
Σj=0
N-1
a ++ 2 - NN
0 and
E
Ei=o
N-1 N-1 aut y -
Zj=0
Pa-(1- (1-"))
Pa
Lodja 2 1-02x)
(1-a)²
(1-a)(++a
N-1
E
W
T
a yN- (Bx
1=0
N-1 1
a
i+1
yn- (BEN= + WN- w Dg 2
1=0
41

---

BE
yN - - T
E' x
N-1
+ bE
a
[+1
VN
E 2 w
T
==
B
The last equality follows from lemma A.3
where we have: IE
En
YN ut
1-ax*
1-a
I,
i=
O. and [. _
I.
N-1
N-1
E
a q
C
1+
yN-, (BoN= + UV Can
Σ
Un UN- N
i=0
N-1
N-1
E
Ca.
TT
Σ
a
i+l
Un (B-
Σ 1=0
a
i+ن
Un
x
i=0
N-1N-1
CC
C
BE
Σ
yN UN-
2
)
+ (d 1)aa a1 C B
22 2N
0)²
The
last
equality
follows
from
lemma
where
we have:
(d+1)a (1-a2N) [1=a]
Σ
i=0 2/=0
N-1 N-1
a 2++ yn II 2 Nー
I, and
Eg
I.
N
E
Fg
UN NNN-
N
CE
Fg
x
a
i+
U/ UN-1 N
N-1
CE
E
Σ
i+1
yN--WeN
y/N
1=0
1 a
)
1-a
The last equality follows from lemma
A.3
where we have: E DN-- attl yn att2 xw
a and E x =I
L-a*
N-1
N-1
T
Σ 1=0
a i+1 yn (BoN- + ub- uy Cx.
Σ
a i-1 yN
i=0
42

---

N-1
N-1
C
T C
a
H+
UN BB-:(BeN+++NN-11 a
Σ
yy
==0
N-IN-1
CE u
BE
BE
a
-ZN- VN - NN
E=
N-1 N-1
+ CE
bE
+
YN Ai---V--
1=0 ==0
da² [ a
(2d + 6d)a² 1 - 1-a2N 2N
JC
a1-
a) + +a)
The
last
equality
follows
from
lemma
A.3
where
we
have:
Ei=o 2=j=0 ao 1-a0
N-1 a auts+2 -iNN- = 0, E
0=
==
a
a
UN-
=
(2)) (1022)
(I-a)([+a) ) and E 1 = I.
(I-0)2
N-1
E
w
Bg
Ca.
i+1 2
1=0
N-1
E
Cang
Te
i+1 2
i=0
N Σ
CE
Ty
x
++1
y 20
The last equality follows from lemma
A.3 where we have E E a u = 0.
N-1
N-1
E
a T
a
i+
(BN
ub .Cag
Σ a yN-
q
==
1=0
N
N-1
E Can
F.
T T
Σ
yn y (B bR)
it1
Σ
2m0
j=0
CE 2
C BE
JN yN-NNN-
1=
N- N-1
+ CE
TTb
bE
++j+2 2 14j42
Uル -N
i=0
1=0
N-1 N Σ
+ CEE 2
CT bEL
i+i+2 2
yN-
1=0
j=0
da² [ 11-18 2N
CC bg
(1-a) 111++a a)
43

---

The
last
equality
follows
from
lemma
A.3
where
we
have:
1= Σy= 2++++ aN :N - W
0, E
1=0
==
N-
a ++
yN
-UN
= 0.
UN-+
WN--
da² 1-a2N
E E
Σn= UN-
and E Tg I.
(1+a)
w
B
Cag
yN-
i=0
N-1
E
Cag
Lq
w
YN
Ck
T
y
The last equality follows from lemma A.3 where we have E Ei= at yN O.
C.4 Proof of claim B.1
This Section presents the bounds for terms b (T + 1)b;(T + 1), cf (Tt 1)a + 1) and b bT + 1)b(T 1), establishing the property A(T + 1).
(t (t 1)c + 1) and Recurring the Vector-coupled Dynamics equations of 6J (t + 1)b, + 1), c
b (1++11 1) + 1) in lemma A.7 we have:
b. (T 1)b, (T + 1) 6 (T) + 2n - (T) (TT
31 Σ (cl (T)): + 7²² [T)
kighi
b. (0)6 Σ 2n
Bic ( -3, (Bs (2s
B1 Σ (c](()(((s)))
kgi
T b. (0) + 2n Σ (Bs Bıe (s)b;(s))c (s)b.(s) 3 Σ
(cf (s))
4=0
term
term
T
+ 172
0=0
term
- (T) (T)):(T)
(T(
32 (T)()))) + + (T
B
T
= =c, (TTcaa + Σ
2n
(3s (3s 3 (c)))))) 31 Σ (c. (s)ba(s))
k
44

---

- Ba(c (1)((s)())) + + +r'(((()\())
T
= (0)a((()) + 2q Σ
(3g & (8) (+)( (s))cc
2n/31
Σ
Σ
(cf (s))
Aggh ==0
term
=term I1
3n Σ c. (s)b(s)) 2
+"²
(s)
term IV
term V
bT + 11)([T =1)= (T))6() - (3) Σ (ck (T))(1)))) + +
d
T
T
b' (0)6(0)) + 2m132
ΣΣ ==1==0
(ct (s))()))
2
b(s)
=term IV
term
To bound terms I- VI, we will use some inequalities from property B(t) and lemma C.7 with j, k E [1, dl, i + j:
as following
18/ - 3ic, (s)b.(s) s 6() exp( exp((n))ntt
|c] (t) S 25(t) exp(-n
5(t)
c. (t)b(t) s 28()) exp (-p(-) + exp(-
B2
(t) b (t) - 8d, ((t) -13mtt))
at) (x((t) s 8d 8(t) exp( n 40 exp( (x
b(t)] 16dn,3d exp( 2d 2 exp( -nnn
Next we begin bounding terms I- VI.
Bound of term I:
By BB - B1c 58(s) (x)()))xp()))))))))))))))))))))) we have:
8(s) exp( (11877)) ,33
c (s)),
31
45(s) + 2x
a²
50(s)
a²
65(s)
The third inequality is by ((s) 2 2/ dr log (2d
2d + 1)/8) 2 2a = =2a
=
For
the last inequality, as long as N 2
2 1n6- we have 5 << 6.
T
Σ
(8s (s)
T
Σ
65())² (1)(0)18178))
45

---

66 max exp(-11817sdds
66 max exp(n)17)
n1177
The second inequality due to monotone decreasing.
Bound of term II:
(s)
T
(2)(s) exp((=1.17s)) 2
$ 48 max
exp(-
20 max
exp(2n
1317Y
The second inequality is due to exp(-2n8175)) is monotone decreasing.
Bound of term III:
(s(s)
8dh dt ((t) exp(-n1.11)
s=U
< 8dl, 62 max
exp(-1)3170
8d 88 max exp(n317)
1/177
The second inequality is due exp(-n1311) 78) monotone decreasing.
Bound of term IV:
(c] (s))(ss))
8(s)
Σ (2(ss) exp(-n)321)) + exppp exp((nB817s)
B2
T
T
T
< 2
(4
expl (2n27)))))))
+
exp( exp(-n((1) + +2)27s) +
Σ
sxp(-2mB1133
83 ==0
3 ==0
8=0
(4
100
< max
exp( 211227)dds +
exp(-n(31 + 32) B2 -1
7s)ds
32
262 max exp(21827)
+ max exp(n((1 + 32h))
462
82 82 exp(2n317)
1127
n1 82)r
,
1762 max
n2
46

---

The second inequality is due to exp(-2n)22st exp(-n(33 + B2)ys) and monotone decreasing. The last inequality is by expin( +
are since
I and
exp(2n811)
(31+B2)
20 P2
exp(2n81n) I exp(n(8 + 82)n) S and B1 + 82 2 1, 11 2 7
Bound of term V:
Σ (+,0 1
T
8dn dl ((1) +0,
/
$ dd 8d 8
o max
exp(-112775))
8d, d' do d exp(m111)
4032dh m exp(n327)
n
7
16d d2 22 exp(n317)
80 Rax exp(n
+
nB17
1
The second inequality is due to
and are monotone decreasing.
Bound of term VI:
b(s)
16dn 16d
exp(-192nt) + (xx__333
exp(
exp((113171)
& 16dn,3.. 22 max [
(xpl-11375))b + 2d R /
expp-m131
16ds 16d 2 max exp(1)227
of max exp(11317)
7
mB17
The second incquality is due to exp(-n8211) and exp( 31)7s) are monotone decreasing.
We next use the bounds of I VI bound b. (T + (T 1), cT(TT (T + I)c,(T 1) and b (T 1)b(T + 1).
Lower bound of 6, + (1),),,((++)
6. (T+ + (Tb, +)
T
d T
b. (0) +2
(3s ( (s)b (8)b.(s) (s)))) -2n Σ
(c
(s)
24
T
+ 17
(x(s)
==0
20
3dh,
662
max exp(11377)
2m (d - 1) 282 exp(21)31) 7)
2n
4
1111
11117
3dp,
1282 max exp(n317)
4(d - 11002x exp(2m
4
B17 7
7
47

---

3dh.
2 12 9dp log [d((d + 1)/6) 2 4(d - 1) 9dn (g((1d(2d + 1)/0)
31 312d0 d/
d/
dn
2
The third inequality is by 8max 3 log( (og(4d(2d+) 1)/8), exp(n917) S exp((288r)) 2 and = d.. The last inequality follows from d, = (2(d²) z (1728 (7728log(44))) (dd(4((2d + 1)/6) + 576(d 1) log(4d(2d + 1)/6))/B1.
Upper bound of b, (T I)b: +
6,- (T 1)b, +1)
T
T Σ Anght
by (0) - 21
(3g Bic (s)b Σ
-n
Σ (ck (s))
11
17"
bi(
term III
5dh
66 max exp(1117)
8d max exp(n117)
2n
+ +12
4
3/1
m317
5d
126 max exp(n3177)
df max exp(1117)
4
B1Y
B17
5dp
2 12 9dh log( (4d(2d + 1)/8)
2 81dnad 9d,, log(4d(2d + 1)/5)
4
B1 /
31 ddh
2dh
The third inequality is by 6max 3 dn log(4d(2d 1)/6), exp(nBin) S 2 and ddh. The last inequality follows from
d = ÿddd)
(576 (odg 1)/6) + 192 (dd 1))00))
(576log(4d 1) + 384nd; of log(4d(2d 1)/6)) /B1
Lower bound of c, (T 1)c; (1)
(T
T
d T
(0)c;; +2n Σ
-2 (3s 3. (
Σ
Σ (cT (s))
==0
kigh ==0
34p
term
t0rm
T
T +n²
2m
Σ
(c. (s)b(ss))²
Σ
(s)
==0
==0
20
3d
662
exp(n)17)
282
max exp(2n017)
1782 max
2m
max
2m,11 (d- (ldd-))..))
-2132
11
11.17
11,27
3d,,
* 12 * 9d,, log Ad(2dd + 1)/6) * 4(d 1) 9dl, log( log(Ad(2d + 1)/5)
B1 B1
5d
34 9dh log(Ad(2d) 1)/8)
1 dd
48

---

da 2
The second inequality is by δmax - 3v log( ( exp((817))) - exp(2n) S 2 and
y ddh. The last inequality follows from d, (((d)) 2 (17281og( (d 1)/6) (576d 3 log( (4d(2d + 1)/8))/5.
Upper bound of c. (T++)),(T+
CT (Tc,
T
d T
c. (0) +2n Σ
(3g Bic (s)b, (s)b,(s) (s)(() 2n181
ΣΣ (cl (s))
==0
loghi ==0
tenn
-trm
T
3n
Σ
(c] (s)b()))²
+np
Co 8
==0
==0
5d
662
max exp(ns17)
26 max exp(2n317) 2 1)
2m
11177
1117
16d
max exp(n317)
8032dh Smax exp(1)277)
v/2
+
MSIY
mn
5dp
* 12 9dh (og(Ad(2d + 1)/6))
2 4(d 1) * 9d1 log (4d(2d log(Ao(22)))))) +
+ 1)/6)
4
B1 31
_d
2 * 16nd, of 9dn log (2d + 1)/0)
2 80m32dn 9dl log (d(d + 1)/8)
31
$ 2dh
The second inequality is by mmax 3 da log(4d(2d exp(18in) (x))) exp(21)317) 7 dh. The last inequality follows from
s 2 and
d, = R(d))
> 576 3770
1)/0)
+ 384 log
+ 2 1)/6)
> 576 (0d 1)/0)/8 + 192(d- 1)/0)
+ 768nd, of ///0///24-//////
+ 1)/6)
Lower bound of 5T (T ( 1)
d T
T
bt
(T + 1)6(T + 1) = T (0)6(0)-2n/5) Σ
Σ
(cT (s)b(s))
+n7
Σ
[(s)
==0
3dg
3dp
1762 max
2n32d
4
nBay
3d,
34d52.. max
4
3d,
34d * 9d, log( 4d(2d + 1)/8)
4
2 d
dh
2
The third inequality is by 8max = 3 dn log 4d(2d +1/2d++))//)) and 7 ]d.. The last inequality follows from dn (d))) 2448d log (dd + 1)/6).
49

---

Upper bound of b (T 1)6(T +
T
r
bT + 1)b(T + 1) = 6% (0)b(0)
2n Σ Σ
(s)b(s))
b(s)
S 5
VI
5dp
176 max
16dh, 16dh,3dd2 32d2
max exp(1)977
df
8 exp(n317) ni1n
2nd
+ 172
11227
m
5dp +
34d62 max
16nd B2d
* 2nd do
max
317
5dp
34d * 9d,, log(4d(22 1)/6)
32nd,, 32d
9dh, log( (dd44d(2d + 1)/8)
2
d2
Andn, * 9d,, (2d(4d(22 - 1)/5)
B1 d1
2dd
The second inequality is by exp(n917) I exp(n927) S 2 The third inequality is by 8max 3 log( log(4d((2 1)/5) and = ddh The last inequality follows from
d. = )
2 log Ad (2 + 2d log(4d 1og 1)/0) + 48 A8log(4d(2d) + 1)/8)/31 1lddlog(4d(22)))) (2d + 768n log( (d(d + 1)/6) 96nd, log(4d 1)/0)/31
C.5 Proof of claim B.2
This Section presents the exponential decay bounds for terms (3s B1c, (T 1)b,(T 1)). c. (T + 1)6, + 1) and of (T' 1)6(T++ + 1), establishing the property B(T + 1).
Bound of (3g Bic (T + (Tb, 1)) Recall the following equation from lemma
c. (1+1)),(t+1) cj (1)b.(1) + o (8s (ic (1),(1)) - Bi Σ
c. (1))2()))b)()))))))))))))))))))) (1)).(1
k
B2c7 (t)b(t) 6. (t)b(t) + (3s Bic (ic (t)c,(tt
81 c (t)b.(t) cf (t)cx(t)) + + (t)
Based on the above equation, we have:
(3g Bic (T 1)b, 11) = B3 (T)
- (135 (1) (T)b, - Σ c ((T))((T)-6(())6((
32c, (T))(T) b. (T))(T) + - (T (T)a())
- 31 Σ cf (T)), cT (TTaa)TTTTT - (T)..(())
(31)
= (1 +
(T)
+ m Σ (T (T)). + m Σ II] of (T)
+ n1110c0 (T)b(T)T)))))))..).))))))))))))) (B)(T)--6((0) (T)D(( - of (
50

---

The term B3 =1c] (T) is highlighted with underline, and we collect its negative feedback terms together. The factor (1 n/1 (b) (T)6:(T) + of (T)a(T)) < will drive (3s 3ic, (T+
1T, 1)) to converge to zero.
By Recurring (Eq.. (31) from 0 to T, we have:
(Ba-3ic) (T + ()b, (I)
T
II 1 ( ==0
3 (0)6.(0)
T
Σ
II (1 (1-
( (b) -c (s')cc((s)))
==0 s'=+1
(32)
c (s)be(>) Dg (s) + ne Σ
ce (s)b:(+) of (s)x(())
31 B2c,
(5 (s)b(s) b7 (s)b(s) 3 (s)b:(s)
Here II'' (1 (b. (s)b,(() + (T()a(ss)) S (1 - 2mBrn T) since
S b (s), c. (s). Besides, from property B(O). B(T) and lemma C.7 we know that c, (s)bx(s), c, (s)b(s) and ट. (s) have bounds with exponential decreasing rate. Therefore, it easy to derive an exponential decreasing upper bound for (3s Bie (T + (T 1))
(s) (s)b,,ss,, (s)cu((() and By substituting the bounds of (s) of (s)b(ss,
b. (s)b(s), we have:
(8s 31c (T 1)b, + 1))
$ - 21(1r))-111 11s Bs 3.c (0)6((0))))
+ Σ (1 2nB1r) 2n 11-11)222260 28(s)² exp( 3)p(-)))n7))
(33)
5(s)2
+ nB (28(s) exp(-18279)
B2 exp(- (1p(-11175))
+ B2 B1
+ d82d) (xcptt))expppp321110
The notations and & highlight the corresponding terms between (Fq. 32 and (Eq. refference.
33 for
51

---

We further have the following:
(Bs Bic, (T + + 1)b; 1))
21-271177) T+1 B3 - _1c (0)b,(0)
T
+ Σ (1 - 2m131)
(2 1) 28(s)) exp(-=n317)
+ n (xp( exp((=1877)) + 8s)2
exp( exp(-n18188)
B2
1 B1 (8d) exp( exp(-n/11d ) expp
S 2100- -
3 (1)0000)00)))
T
exp(2n817(s (4n ( + 817 ex0
(d-
exp(-n311
+ (2 (x)))
$ + + (x0
2 exp(-nB17(T'+ 1))
11117
(27 87/ (
1ep(-n317(TT) (T+ +1)
n
831(d 1oldd118(T) 1)8(T)
+ 25(T)
16nd, de ((())
(3)
24
+ 31
+
((0)
7
7
5(T) exp( n1 +1)
< ((T)
+ +1)
(34)
The second inequality is derived by factoring out the factors exp(-1)119 and exp(-182 The third inequality is due ΣT E cxp(2n317($ T)) exp(-n517s) W n117 exp(-n3: > [T 1))
and D ( T)) ()7(5))275) S n82 exp(-113>77 1)) in lemma C.5 The fourth inequality is by 8(0) v (77), 2m + +1) S exp( 31 +T+ 1)), and we consider
33
5() v 8(0) ((T). The fifth inequality is by proving (10) + 8 (d ()((7)
28(T)
8(0)
d
6() 24n3,d,d6(1) S as follows:
83
(d 1)8(T)
+ 28(T) (6nd).d5(TT) 6315(T)
24m d, dT)
+ B1 +
8(0)
?
Ag
3
8(T).
(idd-1) (83) 1) + 16 dl + 631 24n3,daa)
+
8(0)
4
As
3
2 du log( 4d(2d + 1)/8)
3 dn log( 4d((d + 1)/8)
(d + 16 + +3, 24m3 1
51
The first inequality is by B1 < 3. The second inequality is by 5(0) 2 2 dr log( (d((d + 1)/8).
8(T)) I 3/ dn log( (d(2d + 1)/8) and 7 _dh. The last inequality hold long as de
52

---

(dd)) (1ug(4d(2d-11/8) + 22/06440111111110044444222244140 2
112g(4 + 24 +=210//0404444110000000000--1--0 Therefore, we have
(0s (ic- (T+1)),(T+1) S ((T (xp((=xx(((==-))))00 ( +77 = (35)
where the last inequality is by 8(T) - 6(T 1).
The proof for the bounds of c (T + (T + 1) and c, (T + 1)b(T + 1) are similar to that of
(gs Bio (T + (T + 1)). We presents the calculation as follows.
Bound of c/ (T + (T+))
c (T + (T 1)
(T),, (T) =[]]
1/ (83 3ic (T)) (T),, Σ c (T) (T)
B2c, (T) (T) (B- (Tca(
= - (
+ (
- Bi (TT n/ Σ (T) (TT),,(
kfihk
nB2c (T)b(T) ( (T),
d
- -n
c
- (1)
T
T
[1-
s'
(Ba 3ie (00) Bic, (s)c
d
d
1131 Σ ci (3)bb((s) (s)b,(() n Σ
ct (s)b((s)) c (8)
B22 (s)b(s) bj. (s)b(s) + (
(1 (1-2m317)) 21 T+1 c. (0)
T
> (1 - 2n -- (n) exp( 311s))
==0
+ 21 (d - 2) - 20 exp( - 11))
53

---

8(s)²
+ n
exp(-113178))
$2
1² (8d, ( d32d, (
S exp(-22m31 + 1)) (
T
(d exp exp(2m/1n(s - (2)6())) (n) - 2)8(s) + +81 +n((s)
(2n)(s))²
d(s))) exp(-n
+ (2nB20(s) + 8n/ B2dnd6(s))) expl (npps)))))))
s exp(- 1))
+ (2n8CT)) 1n (d 2))(T) + n)8(T)² + 817 ord ((T))
2
exp(-n817(I + 1))
1/117
+ (2)B28(T)² + 817 (2T)))
exp( 11- (7 + 1))
11327
(1+ 48(T) + 8(d - 2)8(T) s(d-2)6(T))
+ (T)
16
(5(T)
+ 24ndh,d(T))
+
BI7
?
317
B17
5(T) exp( (T + 1))
(T) exp(-11817(7 + +1)
2T 1)) -n +))
(36)
This bound requires 8(T)
+ 8(d 2) + Bi +
16nd,d Bg
+ 24ndnd)
S 1, which can be
31
2x + 6 + 2) 2 2
verified by dn (ddd²)
+ 1)/6) 11 + 8(d 2)
36 log 4d(2d + 1)/5) 9 + 8(d 2) + 3/4 + 16nd d²
+ 6 + 24ndn
Bound of c. (T + 1)b(T + 1)
c (T + 1)b(T +1)
[c] (T)b(T) + (1s - -icc, (T)), (())() - 31 Ee [C](TTbba()) 6y (T))(T)
kgi
-2cT (T))(() b (T))(T)- 32 Σ cI (TT)) (I( + ((T
= =I(1-n8 n3 (bT + ( (((()))
+n(Ba - (IC nB1 Σ a (T))x(T)) 6% (T))(T)
Rif
B Σ c (T))(T) (
- 11m32 T (s)b(s) + +c (I()) c. (0)6(0) (0))6())
==0
Σ
II (1 1 11 6T (b')((()))) + (3T())(())
54

---

d
3ic (s) (ic
(s)b(s) n Σ CsT (s)bb((())) (s))(s)
lighi
E
c (s)b(s) c (s)c + [e
kfi
(1 - 22
c (0)6(0)
Σ (1 2n --
"000s=(((6 n))s)) exp(- 8
n/B1 1). -20(s)² exp( -1)8178))
0())
3/22dd - 1) (25(s) exp(-11875)
exp(-1181
$2
+ 1² (Adn exp((n9111) + do(s)) expp exp(-n18270))
( 1)) (0)b(0)
T
exp(21)827(s - 2)) (1) + 2m (1)32))+2mn3((d 1)8(2)²
+ rid - 1))(7)) + of exp(-
+ (0))_(d-1)8(72+--4 28,/ 33 exp())
s
(1817)2 + 21/31 (d 1)) + ny(d 177
2 exp(-n13>7TT +1)
11327
+ (2 - 1)8(T)2 28 (2d)d6(T))
11327 1np(- +
-((T) [+ 4(d 1))(())
56n 32dh,
+ 1))
7
8(T) (8(T))
(d 1)8(T))
+ 2(d - 1)0(T)
8ndn of (11) exp( exp((1111(( + +1)
82
7
7
8(T)
< 28T(T) exp( exp(-1)(pp1)))))p 11) + 1)) +
exp( 1ap(-1186 1110-1181771111 + 1))
B2
(T + 1)
$ (T + 1) exp( (T + 1)) +
exp( (T + +1)
(37)
B2
This bound requires 5(T) (((d-1))56n32dd.
and
8(T)
1, which can be verified by dn = (())) 2 36 36log(Ad(2d + 1)/5) (4(d
(4(a - 1) + 56d1n2)
36 log (dd + 1)/8) (4(d ((d 32d,)
and
dn = (())) > 36 log( 4d((2d 1)/0) 43 ()(d))+))++)))) 26
(d - 1)+ 2(d 1) + 87
55

---

Property B[T + 1) is established by (Eq.
335 (Eq. 36) and (Eq. 37
C.5.1 Auxiliary lemma
Lemma C,5 As long as 2m817 S 2n,, and 21327 I In 2, we have
2
exp(2n317($ T)) exp(-nB811s) 1 exp(-nBi>(T + 1))
n117
exp(2n/11((5 T)) exp(-n1271s) S
exp(-11517(T 1))
n3227
T
exp(2n/327($ - T)) exp((=n8119) v
exp( 111/311 +1)
11327
==0
2
exp(2n - eex((-n82ys) v exp(-nn33 (T + 1))
11227
Proof of lemma C.5
exp(218177s--77) exp(-11311)
exp(1]]788 21)317T)ds
0
1
(exp(- 1hsp-~1317TT- 1)) exp(-2n11177)
71977
1
exp( -1 1))
n317
exp(2n311)
+ +1)
1B17
2
exp( po(--1817( +1)
The first inequality is due to exp(n817s) is monotone increasing. The last inequality is due to 2m,317 y < 2.
exp(2n817( -T) exp(-x)221)
T
exp(-n(33 - 28:)7s 2n317T)
==0
cxp(-n/ n(2 231 2ns3
1
(exp(1(32 - 28,)r 2n exp(-n)22TT)
113 2,117
exp(1](32 281)7 21
nl32 2/1)7
exp(11827)
1))
n13 - -281)7
56

---

2 2exp(1827)
exp(-2n8177TT + 1))
11827
+ 1))
m3227
The first inequality is due to exp(
cxp(-n(33 (ys) is monotone decreasing. The fourth inequality is
due to B2 2 4,,, thus 1
2--23,
_ 32-
2 The last inequality is due to n222 I (In 2)/2 s In(3/2).
T
Σ
(xp(22 - T)) exp((nn811s)
exp(n(233 - B1)78 27
T+1
exp(n(282 31) (T)
( - 1112 + +1) exp(-2n
(23, - 317 1
exp(21327 1T + 1))
(232 B1)r
exp(21)327)
exp( 1117 + 1))
(2332 -1:17
< 18277
exp(
(-1317(TT +))
The first incquality is due to exp(n(23g - sss) is monotone increasing. The last inequality is due to 32 2 B1 and 27/827 S 2n
The proof of Σ (xp(2n82)(s -T)) exp( 3 M exp(-n + 1)) is similar to E 3xp(2) (s T)) exp(-118m1 N:. exp(-n8in(T + 1)). Just replace 31 with 82.
T))
and consider 2n, S 2n
C.6 Proof of claim B3
This Section presents the bounds for terms b, (T 1)6, (T++ + 1), c. (T+ 1)c,(T' + 1) and by (T 1)b(T' + 1) with i,j E [1, d], i + j. establishing the property C(T + 1).
Recall the Vector-coupled Dynamics equations of by (1 + 1)b, + 1), of + 1)e,(1 + 1) and b. (t++ 1)b(( + 1) in lemma A.7
6T (t
10sa-3ppe -IcT (1 - (t) (t)b.(t) b, (t)b,(t) 7 (2(B3-31c)
(38)
- + (1)) 233
Σ
((0))(()-<I(()b)))))) + (b]](((6,()))
(t+1))
= (f)c(((( + -2c (BC -
(1)b.(t) + - - (1T((b,,)))((0)))0
(39)
- Ao(eT (1))(())) + (1)b((1) - -23
Σ
k
- (1)(((())(() - cj (06(1))++ + (1),()
57

---

(t+ 1)b(++) +1)
4
- b. (8)b(t) + (Bb)) - B1
Σ (t) (t)((t)
(40)
82 c (1)((1)) (x00)).00000000 (T)0:)))))+++++++++ + -1)
x=1
To give bounds for the above three terms, we will use the following bounds from property B(1) and lemma C.7
|03 31c, (t)b.(t) s 6(t) exp(-n98ndd
1c (1)6,(0) s 28(1) (xp0))
8(t)
c (t)b(t) s 28(t) exp(-n92nt) +
B2
(,) [t) s (1)p
exp(-n81nt)
at) (())),(()) s 8d)d2 exp(- + 4003 (00()))(())))))())2m)))))
b. (1)b(1) < (xxp)_dd((()xxx(((pppp) Ad exp(-11311)
Besides, by A3 (ccT(()::l))))))
s 5(t) exp(-nx3rd)) we have:
c (t)b,((t)
5(t) exp( 31) + As
81
45(t) + 2a
a²
55(t)
a
68(t)
The third inequality is by 5(s) 2 2 dy log 4
2a 2a In For
the last inequality, as long as N > 2 2
have 5. < 6.
In6-In5
We will provide the upper bounds for
bT[ (T 1)b, (T 1)
c (T 11c + and
by (T
1) + 1) by substituting the above bounds into (Eq.
38
(Bq. and (Bq. 40 respectively.
58

---

Bound of (, (T+ 11), + +1)
6. (T
= (10000000100000-1000000000000000
- +((([()))(())+]]]1)))) (T))()((T)) - 23,
Σ
k
by (T)),
3cc --
+ (c) (T)b, cf (T)) + 2
Σ
+ 12 by (T)b,
<T(T) + 4n ((T) 1 28(T) exp(- 2m ( exp((n811T) 2
+ 2n- 2) (25(T) expl exp(- + 8d (
S((T) + ((T) exp(-)33))())) +
exp(
+ 87 () ex00))))xx((m8111))
+ &n)"), -cxp(-n/17T)
((T) + 6(T) (81mmaxx + 411333 + 81131 (d - 2)0max + dad d exp((nn317T)
6max))
(41)
The first inequality is derived by triangle inequality. The second inequality is derived by b. (T) 8(T)) and substituting the bounds of |Bg 3 (1)6,(1), |c; (1)b,(4) (4) and
b)) The third inequality is derived by factoring the common factor 5(T). The last incquality is derived by 5(T) < 8max and exp(-n;1nI)) < 1.
59

---

Bound of cT (T))e) +1)
([++)))((++]
= - Bıc, (T)) (T)),
+(C[T1b/())[[((6[( - 281
Σ
- 232c (T))(T) c (T)b((( + (T)
loT (2) (7)
+ (3s -1c (T
n0 (T)b, (cT
(T)) + 27
Σ
(T) (T)))a((T
+ 2n (T)b(( oj (T)b(T) + 1 (T)
s S(T)) + 4n ((T) exp(-np111) 28(T) ((l7))xpp((pp83 1esp--n-1ytn 2m23 28(T) exp(-
(d 2)- 20)( exp(-1111)))))11 2
2n (8) exp(-n1227)) 6(T)
82 expp(=ninnT))
+ (Sdn ((T) - (0
s 5T(T) + 8(T)) +|0 81/1 2)5(T) (x(T))xxp(-(817) || dn
2T)
32 Syp (
exp(-nn817)
+8(T)(8(220(T) 100 (
exp(-n1371))
2170max
s S((T) + +(T) [6m + 4n + 87 22)-m2mmm + 8172 dn 8max + B2
eexp(=n3112
+ ((T) (87320max + 40n² nun233dnn 3gdn cxp((1182TT)
(42)
The first inequality is derived by triangle inequality. The second inequality is derived by b. (T) 8(T)) and substituting the bounds of |Bs 3.ce (1)6,(1), |e; (()) and (1) (1) The third inequality is derived by factoring the common factor 5(T). The last inequality is derived by 8(T) < 8max. exp(-n,1nI) and exp(-n827T) < 1.
60

---

Bound of (T 1 +1)
(TT(T) +1]
6y (T)b((T B (1) - 31 -
- B2
c ((T)6(T)-c(( (T))..(( + (T)6(Z)
k=1
(T(T)(((T) 1/ (Bs 3ic (T) (T)((T) - 31 (T)(()) (T))))
B2
(T)P(() (T))..(( B2c (T)6(T) of (TT)(((()) + t (T)B(()
keit
< 6, (T)b(T) + ( BicT (T)) + n131 Σ
cl (T)), ( (T))((T)
d
- 3 Σ
C -I + n (T)b(T) (7)b, + 1/ b. (T)
8(T)
(T( + ((T) exp(-n81n7)
(xp(-n(22T))
B2 exp( (pp(-111
- n(3 + Bd) 1). exp(()21p))) (28(T) Spp(-1827T) + 5(T)
exp( (11111)177))
B2
8(T)
- n8s 66(T) (28(T)
exp( -
3g
1/2
cxp( ()pp11)3277))) + 4d, 8(T)² cxp( cxp((c(611)) ((i1Tt)
5(T)
< + S(T) (2( (3)d
exp( exp(-11877)
32
8(T)
6n82
Anj d, d d cxp((=n3113
32
-((7) (2 (2(3) + Bd 1)((T) exxp()(11)p(p)177
- 12m835(TT) 817 B2dn dd 8(T)) exp(-11312
<T(T) + +(T))
(nddōmax + 4n d dmxx) (137))
+ -((T) (81 + 12n 811² B2dn dmax) exp(-112277
(43)
The first inequality is derived by triangle inequality. The second inequality is derived by b. (T) s 8(T) and substituting the bounds of 1 3ic] (1)b, (1), lc (1)) and (T(())T(()) B(t) The third inequality is derived by factoring out the common factor 8(T). The last inequality is derived by 6(T) < 6max. exp(-1)8177 < and 2(3, + 32) - 1) + S 4,2dd since 82 2 81 and B₂ > 1.
61

---

We next provide the upper bound for 5(T + +1).
8(T + 1) max [1xx (T (1b, + 1)1, c (T + 1),, 1)1,1b, (T + +1)[(T+
+ 1)13
8max + 21ydmax
8T()) + ((T)) (16rydōmxx 411333 + Sn (d 2)8max 81fdddd dn
-ex((-18817))
182
+ 8(T) (C(T)(8888)))) (8ns2ddmax B3 6max + 121)/20max + 8r7² Bd mmax) +exp(-n)327))
(44)
This inequality can be verified by comparing with (Eq. (1 (Eq. 42
(Eq.
43
To give more
precise bound, we introduce the following lemma:
Lemma C.6 ut + 1) S v()) cy(() exp(-at) dy(t) exp( bt), with a, e, d> 0,1 0 and a, S In 2, then y(t satisfies:
2c + b
2d
v(1) S v00 exp(
Proof of lemma C.6
u(t s vtt + cy((t) att) + (t exp(-bt)
v( + s v + cexp(-al) d dxxp(- (bbt)
ut(++] $ v(0) II (1 + cexp(- + d (xxp(bss)
==0
ln + 1) v In (0) +
In + cexp(-as) + dcxp(-bb))
+ In + - In (0) +
cexp(-as) + dexp(-bs)
==0
Iny(t 1) < (0) + [ (cexp(-as) + d eexp(-bs)dds
>ny(( + < (0) + (eepp(a) - exp(-al)) + d 5 (exp(b) exp(-bl)))
2c +
2d
+ Iy + 11 < In (0) +
b
2c
2d
= y(t 11) < (0) exp(
+
b
The fourth arrow due In(1 + x)
< x for 0. The fifth arrow is due exp(-as), exp bs) are monotone decreasing. the 7-th arrow is due to a, b S In and - exp(-at) - 0, - exp( bt) S 0. Lemma C.6 presents the core idea of establishing property C(T + 1). a c and b > d in the above lemma, will have y(t 1) S v(0) (()). Similarly, as Mamba converges quickly (C B CC - I,
T + 0 ), we can prove that |b, (t)b, IoT (t), b. (t)) hold their magnitudes around their initial states.
We next combine (Eq. 44 and lemma C.6 to give bound for 5(T + 1).
8(T 1)
217mmax
< ((T) + 6(T) [6nddmax + 4n33 + 81131 2)6max + 81) dn 8max +
-exp(=n3112
32
+ +8(T) (8n32d6max - 40r²² B5n + 12m32mmaxx + 8n7 d2 eep((n
2 2166)d6mmx 4n132 + 81/31 (d 2)0max 8r 6max + 21pmmus
< <0 exp
1
62

---

2 8882ddmaax +407 f2dd + 12m + =811² B2d 8maxx)
+
11227
50 exp 3 da log (2 2d 1)/0)
(322)
+ 88s
16ndad2 4
+6(d - 2) 2)
81
B1
31
31
+ 16d + 801S2dh + 24 + 16ndad²)
3
55(0) 2
& 3/ dr log +
The first inequality derived by (Eq. (44). The second inequality is derived by lemma C.6 The third inequality is derived by zdn. The last inequality is derived by
dn = {dd
36
32d
8β3
2 (ln(3/2))2
log log(4d(2d + 1)/5)
B1
11
+ 16(d - 2) +
8
4
+ 16d + 80 2 24 +8
31
31
36
32d
8B3
log (Ad(
+ 1)/5)
B1 +
(ln(3/2))2
31
d2
+ 16d + 80n + 24 16 2
+ 16(d - 2) +
+
B1
B,
C.7 Bounds of rf2 terms
(())) [B00l o (1), (T)(()))
This Section presents the bounds for 6. (t))(t) (these terms usually appear in the Vector-coupled Dynamics equations with r7² factor) with J E [1, d] under the assumption that A(t), B(t), and C(t) hold.
Lemma C.7 Under the assumption that A(t), B(t), and C(t) hold, we have the following bounds:
(t()) b, (t)
s 8d,d (a (xp( -xp(1111))) 32)71))
b(t) 2 S (6t.:355530000mmm1m77))00---2161 (
2
(t)
< (1. (x) -B2ntt,
a. (t)b(t)
< (
2n327d),
b. (1)b(t) < 883d,, ((t)) exp( n
4d, d 5(t) (xp(-))17t)
where i,j E d). Note that this lemma does not require i + j.
Firstly, recall the following dynamics equation in lemma A.6
bt bt(() + ((8 Bic (1)6((()))) B Σ ol (1)6:((t)
kgA
=: b,),(t) + (t,
c, 1) ci( + 7 (Bs Bic, (1) -3 3, Σ
c. (t))x(t)-bbll
63

---

32d (f)b(t) (((())
=: c(( (1((()
b(t + () b(t) - (1(82 Σ cl (1)b(1) cx(0) =: b(t) + (,((t)
Thus we have
(1( (Bs -31c (s) 31 Σ c ())b:((())
kzi
(1) = (3s (1))((()))) 31 Σe
c (1)baa() (l) - B2c, (1)6())) b())
kti
d
b(t)= = 322
(1)b(()) ({))
1=1
Recalling the properties A(1) and B(I):
A(t): :
dn/2 S (1)b(())) of (l),))) 6 (1)b(t) S 2dh
B(t):
18a/3/ (1)b((() S 8()) exp(-nBB10)
(t)) = 26(t) exp(-n817t) lc (()(()) S 25(() exp(-n8274) +
5(t)
exp(-n8171)
B2
We can derive the follow bounds for the norm of ()))),()))) and b(t):
5. B.(0 = - B (1T (t),.)t))) - 31
c (t)b,(t) |
Aight
d
< ||/c0/| (88 ((86 BicT (1))((())0 + Σ |ca(t)| c (1))
S v 2d, exp(-113171) + + 3d 1) 25(1) exp(-n1811)
S 2/ 2d, 2//2d, (t) exp(-n317t)
The last inequality is by 31 1.
<(ta||
3.c Bic (bb(t)))) 31 Σ c, (t)bb((t) bx))) - 32cT (t))(t) b(t)
kyj
d
< |b,(t)|
- 3ic -80 (t)),(t))) + Σ
[ba(t)|
|c| (t)) + 82 [b(b
cc (T))((())
kid
< (22h)(())xp(-nBB1) exp(-n3nnt) + / 2dh (d - 1) 28(t) exp(-n3
+ V2dn32 : (25(1) (o(1)exp(-n3311 + 5(t)
exp(-n18111)
32
< d (t) exp(- nnnnn11tt 7 + 2 2d2, 320(t)
The last incquality is by 31 < 1.
|(t)||
B2
c (t)b(()) cx(t)
64

---

d
S 12
Σ
ca(( c (()(())
8(t)
S v 3dd (25(t) exp( exp(-13221) + (xp((-n/17t))
B2
< 2 2d,,, (1x()) + + 2dd2d(() (xxxdd(())x((xxxx) exp(-11811)
By multiplying them pairwise, we obtain:
(t) - (2v2d, (xpp exp(-n)11tt) 2
S 8dnd exp(mm@rpp-)))
The last inequality is by exp(-2n S exp((1))))))
(((t)=,,(() (2 (sxp2 22
+ (, +
S ((1) (( exp(-)3322))
The last inequality is by 32d < 233 because B2 ((dd²) d. and cxp(-2n81nt) S
( exp(-2n exp(-n8174),
S exp(-11321) exp(-1((31 + +12))()= S exp(-11821)).
b(O) - (2d22d,,2d(()))ppp(((>) exp(-))2nt) ddd, exp(-n.37tt)
= ( exp( - 21/ - dd exp(- 21131(1)
+ 8dp (t)) exp( (31 + (2)
< 16d,, 8(1) expl-1122 31](2) - dd ((()² (xp(-113171))
The last inequality is by 32 S 83 because B2 2 1, and exp(-2n31nt) exp(-1817t),
exp(-2n1277)
(12(-1))27)),
exp(-n)(6 (B1 (2) S
c (1)b,(0) S |a0 |b,00|||
s (2 2dh sxp(-n3ind) + 2 2x00 -2 dd) exp(m(m1x)))))))
I 8d ((1)2 exp(-n/311 + ( exp(()))))x
The last inequality is by 1ap(-2m3144 S cx exp(-n(3r + 3ehrt) I eex((-n32r)
ब. (+)b(()) - | ((0)
< (2 ex) (//
4d, ((1)² e
+ 832d b(t) 14 43yd
+ B2)
< dd 5(1) exp(-113n1 26)) (xp)66((()xx(2(8817)
The last inequality is by 32d < 23
32 < 82 and exp(-2n8811)
5
- exp(-n(B1 32) S exp(- (1)
exp(-n)171),
exp(-2n3211) S exp( 31221),
6. (1)6(1) < |0:00| |600||
$ 2//2d, d6 exp() (2 (2222d,2.d6(t) + (xd
< df (0(1) exp(mnx2)))))) + 4d, d (xpp
The last inequality is by exp(-2n < exp(- ( + e exp(-
65

---

D Discussion
In this section, we show that orthogonal initialization Mamba can be trained to ICL solution, and compare method with some previous works.
Now we assume that each column of WB and We are initialized with orthogonal columns of unit norm. Then we have
Orthogonal Initialization
C (0)C(0) = B (0) = I, B (T)B(())) = C (0)b(0) = 0.
Consider the following update rule as part of lemma 5.1
1) = ((t) n3,C(() C B(t),
(45)
1) + 3 (())))))))B()())))))) -226)))))))))))2222))))))) C((),,
(46)
b(t + 1) = b(t((( (()))))) b(t).
(47)
By (Eq.
45 ) (Eq. 46 and (Eq.
we have:
B (()
B() C())C(I)T (0))))0))0))))000))0))
( (())))((()))))))))))) COCCTCCCCCCT))) b())
C(t ((+))))6--)))((())))))))))))))) (()))))(c(()))
b(t) B()BB() (O((I)
n82 (82C(()T ((t)6(1)TTT)
C())) +
b0)
((t)b(t)Tc((CCC)) b(t)
Combining B (0)b(0) CT (0)b(0) with induction, can derive that B (1)((())
(+)b((()) 0 fort 2 0. Thus we only need to consider the following dynamics.
B(
+ (00 B(),
(48)
1(++1) = C(t)+ + ()) C(t)
(49)
Denote B()) B()) = D()), C(1)T((() = F(t) and C(I) B(I) = F(t) then by (Eq.
(48) and
(Eq. (49
we have
F(1 +1 = = F((1) n/s (D()) +E()) n F())(()
+1) - (18 - m131 E()F((
(50)
F() F(I)
= - (())+ + F(I)) (F((F))) F(I) + F(L)
- E(t)
(51)
(" B(t)F(t)
E(t+ E + (F(()) F(t) F(I) F())
+ - + D()F(()
D(1)
(52)
Note that D(0) = E(0) I and F(0) 0. By induction
that D(t), E(t) and F(t)
are diagonal matrix for > 0. Because of the symmetry, we have D(t) = E(t). Now we denote
D(1) E(1) = I and F(I) I(I).. Then based on (Eq. (50). (Eq. (51) and (Eq. (52).
have:
g(t 1) = gft - (t) + ( - (t) - 3,h(t))
(53)
h(1 1) h(()
(2) -
(54)
Since g0) 1 and (()) 0 at initialization, (ut) will converge to & (i.c. CT B 8 D.
66

---

Compare with Other Techniques (Eq. 45 (Eq.
(46) and (Eq.
47
can be viewed as the
gradient descent that minimize the following target:
[| W - 1
(55)
where X and Y satisfy:
B1
B3
xxT
IR (d(++))(a++) XY
[d++))((
31
Bg 01xd
B2
To establish convergence for this problem under gaussian initialization, Arora et al. (2019) require the standard deviation to be small enough, while Du and Hu (2019) require larger dimension de because their method relies the condition number of X. Our method balances the requirements initialization and dimension. The fine-grained nature of our analysis (particularly the Vector-coupled Dynamics) enables extension various problem beyong (Eq. (55)).
E
Additional Experimental Results
eupiiees lass Dhoniicl
1.-
1.
06-
04-
19
token length
N
(a) Initialization
(b) Trained parameter
(c) Loss curve
Figure 2: (a) Visualization of matrix product C WB before training: (b) Post-training visualization
of matrix product C WB; (c) Test loss versus token sequence length N. Blue curve: experimental loss; orange dashed line: theoretical loss
31
Experiments Setting We follow Section to generate the dateset and initialize the model. Specifi- cally, we set dimension d = 4, d, = 80, prompt token length N 50, and train the Mamba model on 3000 sequences by gradient descent. Moreover, we vary the length of the prompt token N from to 80 and compare the test loss with the theoretical loss. For each N, we conduct 10 independent experiments and report the averaged results. All experiments are performed on an NVIDIA A800 GPU.
Experiment Result Figure 2a and Figure show that C B can be trained diagonal matrix from random initialization. Figure 2c show that the experimental loss aligns with the theoretical loss C(0) = d (1
noting that the theoretical loss 8 has an upper bound 3d(++) 1) that
2N
decays linearly with N. These experimental results further verified theoretical proof.
Mamba vs Linear Attention Optimal linear attention outperforms Mamba under our construction, and they have ([I/N) error upper bound with different constant factors. We provide a comprarison of loss between optimal Mamba (under our Assumption 4.1) with optimal linear attention as Table 2 with setting d = 10, N 10, 20....... 80.
When smaller than d 6,
We also test the case when N < dl Table 3 with setting = 20,
20.
67

---

Convergence of wa We set = 0 in the assumption. Now we show that random initializd W can converge to 0 experimental. The results is in Table 4
Different d2 Table shows the mean value and standard deviation of the loss for smaller d, (in 10
repeated experiments). We d 4, 30, and the theoretical loss 0.2954.
Table 2: Comparison of Mamba and Linear Attention
10
20
30
40
50
60
70
80
Mamba
2.6671 2.6190
1.8189
1.3800
1.1117
0.9308
0.8005
0.7022
0.6255
Linear Attention
1.7742
1.3415
1.0784
0.9016
0.7746
0.6790
0.6044
Table 3: Experiment for N S d
N
4
6
10
12
14
16
18
20
Experimental Loss Theoretical Loss
8.5911
7.8292
7.7009
6,8235
6.4004
6.0612
5.9689
5.6193
5.1426
8.4484
7.8425
7.3173
6.8579
6.4526
6.0926
5.7706
5,4810
5,2190
Table 4: Convergence of wa
Epoch
0
10
20
30
40
50
60
70
80
w
0.8883
0.7513
0.4821
0.3331
0.2444
0.2026
0.1868
0.1799
0.1773
wa
0.7891
0.5645
0.2324
0.1109
0.0597
0.0410
0.0349
0.0324
0.0314
Table 5: Different da
dp
6
10
12
14
16
18
20
mean(loss)
0.2912
0.2933
0.2899
0.2887
0.2929
0.2951
0.2967
0.2959
std(loss)
0.0075
0.0055
0.0116 16
0.0052
0.0105
0.0097
0.0110
0.0142
68


Reviews:
[{"review_id": "TGe5p7swP6", "reviewer": "Reviewer_8rgY", "summary": "This paper studies the training dynamics of Mamba models on linear regression in-context learning (ICL) tasks and established a connection to online gradient descent among demonstrations. The authors analyze how, under certain assumptions, the learned Mamba can effectively emulate online learning behavior.", "questions": "In Assumption 4.1, the requires hidden dimension is $d_h=O(d^2)$. What is the underlying reason for this high dimension, especially considering that linear typically only needs $O(d)$ embedding dimension? If this assumption is relaxed to use a smaller hidden dimension, are similar results still achievable? If so under what conditions?\n\nMinor typo: line 43.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \nThe work provides an analysis of training dynamics of a state space model in an ICL setting, which is timely given the increasing interest in SSM as alternatives to transformers.\n\nWeaknesses:\n1. The claimed connection between SSMs and online learning is not particularly novel and related ideas have been discussed in prior literature. For example:\n-\tYang, Songlin, et al. \"Gated linear attention transformers with hardware-efficient training.\" arXiv preprint arXiv:2312.06635 (2023).\n-\tLi, Yingcong, et al. \"Gating is weighting: Understanding gated linear attention through in-context learning.\" arXiv preprint arXiv:2504.04308 (2025).\n-\tYang, Songlin, Jan Kautz, and Ali Hatamizadeh. \"Gated Delta Networks: Improving Mamba2 with Delta Rule.\" arXiv preprint arXiv:2412.06464 (2024).\n-\tBehrouz, Ali, et al. \"Atlas: Learning to optimally memorize the context at test time.\" arXiv preprint arXiv:2505.23735 (2025).\n-\tBehrouz, Ali, et al. \"It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization.\" arXiv preprint arXiv:2504.13173 (2025).\n\nThese works explore similar connections between linear attention, gating mechanisms and online optimization. So the novelty should be better clarified. \n\n2. Under Assumption 4.1, Mamba reduces to a weighted linear attention mechanism where the weights (determined by $\\alpha$ and $\\beta$) are fixed. The paper would benefit from more discussion of what makes his particular setup novel compared to prior analyses of linear attention's training dynamics under gradient flow. \n\n3. The construction of $\\alpha$ and $\\beta$ under Assumption 4.1 needs more clarification. In the isotopic Gaussian setting in the paper, the optimal choice of $\\alpha$ and $\\beta$ should ensure that every demonstration contributes equally to the final prediction. Specifically, the condition $(1-\\alpha)\\alpha^{i+1}\\beta=1$ (Theorem 4.1(b)) achieves the same prediction as optimal linear attention with equal weighting. The current construction appears suboptimal and the paper should discuss this limitation.\n4. The experimental setup is too simplistic. For example, the dimensionality is very low ($d=4$), and the important implementation details are missing (e.g., number of layers, whether MLP blocks are used, whether the construction strictly follows Assumption 4.1). \n5. Empirical comparison to linear attention baseline would be valuable. Especially under the construction of Assumption 4.1, linear attention might outperform Mamba.", "was_revised": true, "final_justification": "The authors have addressed most of my concerns, and I have decided to raise my score.", "rebuttal": "Thanks for your constructive feedback! We address your questions and concerns as follows.\n\n---\n\n> **W1.** Some prior works explore similar connections between linear attention, gating mechanisms and online optimization. So the novelty should be better clarified.\n\n**A1.** Thank you for sharing these papers, we provide a brief comparison with these paper as follows.\n\n[Yang, Songlin, et al.] and [Behrouz, Ali, et al.] primarily leverage the online-learning connection to design new architectures. While their works have explored connections between various models and online learning, the specific online learning pattern for Mamba is not straightforward. Specifically, the form of $W_B$ and $W_C$, and how Mamba utilizes the context to generate the final prediction, remain unclear. This deeper understanding is lacking in previous work and our work can fill this gap.\n\n[Li, Yingcong, et al.] theoretically analyze the ICL mechanism for gated linear attention, while we analyze Mamba's. Their model architectures are different.\n\nWe will add the comparison in the revised manuscript.\n\nRefference\n\n- Yang, Songlin, et al. \"Gated linear attention transformers with hardware-efficient training.\" arXiv preprint arXiv:2312.06635 (2023).\n- Li, Yingcong, et al. \"Gating is weighting: Understanding gated linear attention through in-context learning.\" arXiv preprint arXiv:2504.04308 (2025).\n- Yang, Songlin, Jan Kautz, and Ali Hatamizadeh. \"Gated Delta Networks: Improving Mamba2 with Delta Rule.\" arXiv preprint arXiv:2412.06464 (2024).\n- Behrouz, Ali, et al. \"Atlas: Learning to optimally memorize the context at test time.\" arXiv preprint arXiv:2505.23735 (2025).\n- Behrouz, Ali, et al. \"It's All Connected: A Journey Through Test-Time Memorization, Attentional Bias, Retention, and Online Optimization.\" arXiv preprint arXiv:2504.13173 (2025).\n\n---\n\n> **W2.** More discussion of what makes our particular setup novel compared to prior analyses of linear attention.\n\n**A2.** Prior analyses of linear attention's training dynamics usually consider merging key-query (e.g. $W:= W_Q W_K^\\top$), specific initializations (e.g. $W_Q = W_K = I$), and gradient flow to simplify the optimization analysis, some of which have been discussed in lines 151-155.\n\nIn contrast, our work:\n\n1. Considers the dynamics for both $W_B$ and $W_C$, which will introduce non-convexity and complexity.\n\n2. Initializes $W_B$ and $W_C$ with Gaussian distribution, which is more practical, and more difficult to establish the convergence given the non-convexity.\n\n3. Trains the Mamba with gradient descent, which is typically more difficult to analyze than gradient flow.\n\nThese setting are more realistic, and we propose some novel techniques to address the difficulties. The techniques are introduced in Section 5 (Proof Sketch). We will add more discussion in the next version.\n\n---\n\n> **W3.** The current construction (unequal weighting) appears suboptimal vs. optimal linear attention (with equal weighting).\n\n**A3.**\nIt is a limitation of Mamba under this construction. As for $\\alpha$, see Eq.(11):\n\n$$h_l^{d+1} = \\alpha h_{l-1}^{d+1} + (1-\\alpha) y_l B_l$$\n\n$\\alpha$ should be in (0, 1) to ensure that $h_l$ will not explode and can also utilize the current sample ($y_l B_l$). Therefore, the previous samples will be multiplied by a factor $\\alpha < 1$ at each recurrent step, and thus the weighting on each sample is unequal.\n\nWe chose a $\\alpha$ close to 1 to ensure that all the samples share the same order of weights ($\\Theta(1/N)$). Althought the weights are different and the Mamba achieves suboptimal prediction, it share the same order of the error ($O(1/N)$) with optimal linear attention.\n\n---\n\n> **W4.** The experimental setup is simplistic.\n\n**A4.** \nOur experimental setup strictly follows our theoretical setup and assumption (e.g., one layer Mamba, no MLP). We have performed experiments for higher $d$ and $N=40$. The experimental results compare with theoretical results are as follows.\n\n|  d   | 6  | 8 | 10 | 12 | 14 | 16 | 18 | 20 |\n| --- | ---  | --- | --- | --- | --- | --- | --- | --- |\n| **experimental loss**  | 0.4738 | 0.7424 | 1.1706 | 1.4141 | 2.1359 | 2.4401 | 3.1154 | 3.9027 |\n| **theoretical loss** | 0.4618 | 0.7583 | 1.1117 | 1.5154 | 1.9636 | 2.4517 | 2.9753 | 3.5310 |\n\nThe experimental results are closed to the theoretical results.\n\n---\n\n> **W5.** Empirical comparison to linear attention baseline would be valuable.\n\n**A5.** Linear attention indeed outperforms Mamba, and they have $O(1/N)$ error upper bound with different constant factors. We provide a comprarison of loss between optimal Mamba (under our Assumption 4.1) with optimal linear attention as follows.\nsetting: d=10, N=10, 20, ..., 80.\n|  N   | 10  | 20 | 30 | 40 | 50 | 60 | 70 | 80 |\n| --- | ---  | --- | --- | --- | --- | --- | --- | --- |\n| **Mamba**  | 2.6671 | 1.8189 | 1.3800 | 1.1117 | 0.9308 | 0.8005 | 0.7022 | 0.6254 |\n| **Linear Attention** | 2.6190 | 1.7742 | 1.3415 | 1.0784 | 0.9016 | 0.7746 | 0.6790 | 0.6044 |\n\n---\n\n> **Q1.** **(1)** What is the underlying reason for this high dimension $d_h = \\tilde{\\Omega}(d^2)$? \n**(2)** If a smaller hidden dimension $d_h$ is used, are similar results still achievable? If so under what conditions?\n\n**A6.** **(1)** We require that the 2-norm of vectors $b_i, c_i, b$ are enough larger than the inner products between different vectors at initialization with high probability (e.g., $\\| b_i \\|_2^2 >> c_i^\\top b_i$, proved in Lemma A.1, line. 473), so that we can establish the initial conditions $\\mathcal{A}(0), \\mathcal{B}(0)$, and $\\mathcal{C}(0)$ for the induction (line. 306-310).\n\n**(2)** $d_h = O(d)$ is enough under some conditions.\n\nFor example, if each column of $W_B$ and $W_C$ are initialized with orthogonal vectors (O(d) embedding dimension can ensure that), we can easily establish the convergence. We provide the discussion for orthogonal initialization in Appendix D (line. 905).\n\nBesides, we list some works that can address the similar optimization problem with $d_h = O(d)$ dimension as follows.\n\n- [Braun, Lukas, et al.] requires that the network's weight matrices are zero-balance at initialization, i.e., $W_1(0) W_1(0)^\\top = W_2(0)^\\top W_2(0)$\n\n\n- [Dominé, Clémentine CJ, et al.] requires that the network's weight matrices are $\\lambda$-balance at initialization, i.e., $W_2(0) W_2(0)^\\top - W_1(0)^\\top W_1(0) = \\lambda I$\n\n- [Arora, Sanjeev, et al.] requires that the weight matrices are initialized by Gaussian distribution with a very small variation O(1/poly(d)).\n\nMoreover, under standard gaussian distribution initialization $\\mathcal{N}(0, 1)$, [Du, Simon, and Wei Hu.] require larger dimension $d_h$ than ours.\n\nSimilar optimization problems in our paper can be solved through different theoretical techniques under different assumptions. We want to **highlight** that our method can provide some new insights for solving these problems and has the potential to be extended to various problems beyond it.\n\nReference\n- Braun, Lukas, et al. \"Exact learning dynamics of deep linear networks with prior knowledge.\" Advances in Neural Information Processing Systems 35 (2022): 6615-6629.\n- Dominé, Clémentine CJ, et al. \"From lazy to rich: Exact learning dynamics in deep linear networks.\" arXiv preprint arXiv:2409.14623 (2024).\n- Arora, Sanjeev, et al. \"A convergence analysis of gradient descent for deep linear neural networks.\" arXiv preprint arXiv:1810.02281 (2018).\n- Du, Simon, and Wei Hu. \"Width provably matters in optimization for deep linear neural networks.\" International Conference on Machine Learning. PMLR, 2019.\n\n---\n\nOnce again, thanks for your constructive feedback and we hope our response can resolve your concerns!"}, {"review_id": "yk2XtcWZRQ", "reviewer": "Reviewer_Znup", "summary": "The authors considered the question of how state-space models learn linear regression in context. The main result of the paper is that the Mamba-like SSM emulates online gradient descent, in contrast to (batch) gradient descent, which is generally believed to be the underlying mechansm of transformers. The result is established via analyzing the training dynamics of a Mamba-like model. Experimental results support the theoretical analysis.", "questions": "1. The results show the difference between the mechanism of the Mamba and the transformers. Though the order of the error is the same, it appears to suggest that transformers are more efficient in this particular setting. Do experimental results confirm the difference and understanding? \n2. The experiments need some clarification. Are the model trained and tested on the same N? What if it is trained on one N value, but tested on various N values?", "limitations": "Some limitations were briefly discussed at the end of the paper. See my previous weakness comments for additional limitations.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n1. The analysis of the learning dynamics of the Mamba-like model is important and non-trivial, and the authors introduced several new techniques to accomplish this.\n2. The results appear convincing and a natural consequence of the SSM structure, in contrast to the transformers.\n3. The insight is helpful to understand the difference in the capabilities between SSM and transformers.\n\nWeakness:\n1. Theorem 4.1 is somewhat strange, where the value delta is not bounded. Usually, we would need to have either \"for any \\delta\\in (0,1)\", or a delta value that is associated with the approximation error. The current statement does not exclude the possibility that delta is very close to 1, and therefore, the statements (a)-(c) only hold with almost zero probability.\n2. The value of \\alpha determines the exponential decay of the weights of the in-context examples. Its value is important and should be learned, which can reflect a key difference between SSM and transformers. One issue with the assumption is that this parameter is set rather than learned.\n3. For linear regression, the case when the number of samples N>=d is significantly different from the case when N<d. This difference does not manifest in the analysis at all. This appears to be an important missing piece. \n4. The result is only regarding one-layer Mamba. Analysis or experiments on multiple layers of Mamba are not included. \n5. Experimental results appear somewhat weak. Firstly, more experiments to verify the assumptions would be needed. Secondly, given the probabilistic nature of the result, there may be cases where the trained solution does not converge to the predicted pattern. A statistical analysis would be helpful here.", "was_revised": true, "final_justification": "I have taken into account the rebuttal/discussion and decided to keep the score.", "rebuttal": "Thanks for your constructive feedback! We address your questions and concerns as follows.\n\n---\n\n> **W1.** The value $\\delta$ is not bounded (in Theorem 4.1).\n\n**A1.** $\\delta$ determines the requirement on $d_h$, we can choose any $\\delta$ as long as $d_h = \\Omega(d^2 \\log(O(d^2 / \\delta)))$. e.g., if we choose $\\delta=0.01$, then $d_h = \\Omega(d^2 \\log(O(100d^2)))$ is enough for the proof. We provide the exact requirement on $d_h$ in Appendix C (line. 624), and the logarithmic term $\\log(O(d^2 / \\delta))$ is hidden by using $\\widetilde{\\Omega}$ notation. We will explicitly state it in Assumption 4.1 in the next version.\n\nWe also want to mention that if $W_B$ and $W_C$ are orthogonally initialized, the requirement on $d_h$ will be $d_h = O(d)$ and Theorem 4.1 will hold with a probability of 1 (We provide a brief discussion in Appendix D, line 902). \n\n---\n\n> **W2.** The parameter $\\alpha$ is set rather than learned. \n\n**A2.** We set $\\alpha$ to be fixed because the training dynamics are much more complex if $\\alpha$ is trainable. Moreover, a fixed $\\alpha$ is sufficient for Mamba to in-context learn linear regression task. In the experiment, the $\\alpha$ usually converges to a number close to 1, which supports our setting.\n\n---\n\n> **W3.** The case when $N < d$ should be considered.\n\n**A3.** We set $N \\ge d$ for the convenience of establising convergence. For a trained Mamba, Theorem 4.1 will hold for the case when $N < d$ because Theorem 4.1 is depends on the trained Mamba's parameter.\n\nWe have also performed some experiments on the case wehen $N < d$. Specifically, we set $d = 20$ and $N = 4, 6, 8, 10, 12, 14, 16, 18, 20$.\n\n|  N | 4 | 6 | 8 | 10 | 12 | 14 | 16 | 18 | 20 |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| **experimental loss**  | 8.5911 | 7.8292 | 7.7009 | 6.8235 | 6.4004 | 6.0612 | 5.9689 | 5.6193 | 5.1426 |\n| **theoretical loss** | 8.4484 | 7.8425 | 7.3173 | 6.8579 | 6.4526 | 6.0926 | 5.7706 | 5.4810 | 5.2190 |\n\nThe exprimental loss is close to the theoretical loss.\n\n---\n\n> **W4.** The result is only regarding one-layer Mamba. Analysis or experiments on multiple layers of Mamba are not included.\n\n**A4.** \nOur work focuses on the theoretical understanding for Mamba's ICL mechanism. One layer of Mamba is enough to learn linear regression ICL task and demonstrate its mechanisms, consistent with the single-layer assumptions common in Transformer-based ICL theory. \n\n---\n\n> **W5.** A statistical analysis to verify the assumptions and results would be helpful.\n\n**A5.** Our experimental setup strictly follows our theoretical setup and assumption. Because the assumptions for theoretical proofs are typically more stringent than experimental requirements, under our assumptions, the probability of the experiment converging to the predicted pattern is close to 1. \n\nExperimentally, We provide the error bar for the loss curve in Figure 2 (c) (line. 933).\n\nBesides, We have performed experiments beyond our assumption 4.1. (e.g., $N<d$, smaller $d_h$), almost all the experiments will converge to the predicted pattern.\n\nThe experiments for $N<d$ is in the above answer **A3.**. The following is the mean value and standard deviation of the loss for smaller $d_h$ (in 10 repeated experiments).\nWe set $d=4$, $N=30$, and the theoretical loss is 0.2954.\n\n| $d_h$ | 6 | 8 | 10 | 12 | 14 | 16 | 18 | 20 |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| **mean(loss)**  | 0.2912 | 0.2933 | 0.2899 | 0.2887 | 0.2929 | 0.2951 | 0.2967 | 0.2959 |\n| **std(loss)** | 0.0075 | 0.0055 | 0.0116 | 0.0052 | 0.0105 | 0.0097 | 0.0110 | 0.0142 |\n\nThe results shows that with smaller $d_h$, the Mamba can typically converge to the predicted pattern (which can achive the theoretical loss), and the deviation of loss is very small.\nWe will include the statistical analysis in the revised manuscript.\n\n---\n\n> **Q1.** **(1)** The results show the difference between the mechanism of the Mamba and the transformers. **(2)** It appears that transformers are more efficient. Do experimental results confirm the difference and understanding?\n\n**A6.** \n**(1)** As for the mechanism understanding, Transformer simulate gradient descent for ICL, while Mamba simulate online gradient descent. Our experiments in Figure 1 (b) (line. 320) shows that Mamba adjusts is hidden state to align $w$ when processing the prompt tokens, which reflects the process of Mamba's online learning.\n\nMoreover, some related works can reflect the characteristics of this mechanism:\n\n[Park, Jongho, et al] experimentally verify that Transformers can learn vector-valued MQAR tasks in the context which Mamba cannot, while Mamba succeeds in sparse-parity in-context learning tasks where Transformers fail. \n\n- Park, Jongho, et al. \"Can Mamba Learn How To Learn? A Comparative Study on In-Context Learning Tasks.\" At ICML2024.\n\n**(2)**\nOptimal linear attention Transformer has a lower loss than Mamba while they share the same order of error. The following is the result of loss for optimal Mamba and optimal Transformer.\n\nSetting: d=10, N=10, 20, ..., 80.\n|  N   | 10  | 20 | 30 | 40 | 50 | 60 | 70 | 80 |\n| --- | ---  | --- | --- | --- | --- | --- | --- | --- |\n| Mamba  | 2.6671 | 1.8189 | 1.3800 | 1.1117 | 0.9308 | 0.8005 | 0.7022 | 0.6254 |\n| Transformer | 2.6190 | 1.7742 | 1.3415 | 1.0784 | 0.9016 | 0.7746 | 0.6790 | 0.6044 |\n\n---\n\n> **Q2.** **(1)** Are the model trained and tested on the same $N$? **(2)** What if it is trained on one $N$ value, but tested on various N values?\n\n**A7.** **(1)** In our experiments, the model trained and tested on the same $N$.\n\n**(2)** If the Mamba is trained on sequence length $N$, but test on sequence length $M$, then the order of test loss will be $O(1/N + 1/M)$. This difference between training and test comes from the property that the optimal weights are determined by N under MSELoss. We provide a brief explaination as follows:\n\nGiven $x_i \\sim \\mathcal{N}(0, I_d)$, we have $\\mathbb{E}[ \\frac{1}{N}\\sum_{i = 1}^N x_i x_i^\\top] = I$ (independent of $N$), and $\\mathbb{E}[(\\frac{1}{N}\\sum_{i = 1}^N x_i x_i^\\top)^2] = (1 + \\frac{1+d}{N}) I$ (dependent on $N$). Because MSELoss introduces a quadratic term, the optimal weights for Mamba will depend on N during training.\n\nThe similar phenomenon arises in the theoretical analysis of ​linear Transformers:\n\n- Zhang, Ruiqi, Spencer Frei, and Peter L. Bartlett. \"Trained transformers learn linear models in-context.\" Journal of Machine Learning Research 25.49 (2024): 1-55.\n- Wu, Jingfeng, et al. \"How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?.\" ICLR. 2024.\n\n---\n\nOnce again, thanks for your constructive feedback and we hope our response can resolve your concerns!"}, {"review_id": "PkRiAFRtkW", "reviewer": "Reviewer_rRWX", "summary": "This paper investigates how the recently proposed Mamba architecture, a state space model (SSM) based alternative to transformers, behaves when trained on implicit algorithmic reasoning tasks. These tasks are designed such that the optimal solution is a known algorithm, allowing precise evaluation of learned behavior.", "questions": "How difficult is it for you to analyze the ICL of Mamba to simulate the gradient descent of a deep neural network, similar to the results in [Wu25] for the Transformer model?\n\n[Wu25] In-Context Deep Learning via Transformer Models", "limitations": "1. The experimental domain is synthetic, which may not accurately reflect the challenges encountered in real data.\n\n2. Scalability to deep neural networks other than linear regression.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": ">Strengths\n\n1. Highly Insightful Experimental Design: The use of algorithmic tasks with known optimal solutions allows unambiguous interpretation of results.\n\n2. Comparative Analysis with Transformers: Shows that transformers, even when heavily overparameterized, struggle to learn these problems efficiently.\n\n3. Theoretical and Empirical Depth: Offers visualizations of internal states and updates, demonstrating the convergence toward optimal logic.\n\n>Weaknesses\n\n1. Limited Task Scope: Generalization to more complex or noisy real-world tasks is not tested.\n\n2. No Formal Proof of Optimality Emergence: The paper does not offer a theoretical guarantee that Mamba will always converge to the optimal algorithm.\n\n3. Transformer Baseline Might Be Underserved: Transformer models used as baselines may not incorporate all recent efficiency improvements (e.g., linear attention, memory tokens).", "was_revised": true, "final_justification": "Thanks for the response from the authors. All my concerns are solved, and I assign equal weights to my concerns. I appreciate the contribution from the authors, but I feel like the work does not supervise me a lot. Thus, I keep my score as 4.", "rebuttal": "Thanks for your constructive feedback! We address your questions and concerns as follows.\n\n---\n\n> **W1.** Limited Task Scope: Generalization to more complex or noisy real-world tasks is not tested.\n\n**A1.** The primary contribution of our work is to provide the ​theoretical understanding​ of Mamba's training dynamics and its underlying mechanism for (ICL). Achieving this requires focusing on a setting amenable to rigorous analysis. The linear regression ICL task serves as a classic model for ICL research in Transformers, and has not been theoretically studied in Mamba yet. Our work fill this gap. The theoretical foundation established in our paper will provide insight for understanding Mamba's performance on more complex real-world ICL tasks.\n\n---\n\n> **W2.** No Formal Proof of Optimality Emergence.\n\n**A2.** In Theorem 4.1, we garrantee that the results hold with at least $1-\\delta$ probability. We can choose any $\\delta$ as long as $d_h = \\Omega(d^2 \\log(O(d^2 / \\delta)))$. e.g., if we choose $\\delta=0.01$, then $d_h = \\Omega(d^2 \\log(O(100d^2)))$ is enough for the proof. We provide the exact reuirement on $d_h$ in Appendix C (line. 624), and the logarithmic term $\\log(O(d^2 / \\delta))$ is hidden by using $\\widetilde{\\Omega}$ notation. We will explicitly state it in Assumption 4.1 in the next version.\n\nUnder Gaussian initialization, it is imposible to guarantee that Mamba will always converge to the optimal algorithm. For example, if $W_B$ and $W_C$ are 0 at initialization, they will stay at zero because the gradients of them are zero (we cannot rule out this situation under Gaussian initialization).\n\nWe also want to mention that if $W_B$ and $W_C$ are orthogonally initialized, $d_h = O(d)$ is enough to ensure that the Mamba will always converge to the optimal algorithm. (We provide a brief discussion in Appendix D, line 902). This type of optimization problem can be proven through different methods under different assumptions. Our method can serve as a supplement to these methods, while also being more flexible and has the potential to provide insights for the study of other optimization problems.\n\n---\n\n> **W3.** Transformer Baseline Might Be Underserved.\n\n**A3.** As for the comparison between Transformer and Mamba, Our paper focuses on understanding their mechanisms. \n\nWe have performed some experiments on linear Transformers as follows.\n\nSetting: d=10, N=10, 20, ..., 80.\n|  N   | 10  | 20 | 30 | 40 | 50 | 60 | 70 | 80 |\n| --- | ---  | --- | --- | --- | --- | --- | --- | --- |\n| Mamba  | 2.6671 | 1.8189 | 1.3800 | 1.1117 | 0.9308 | 0.8005 | 0.7022 | 0.6254 |\n| Linear Attention | 2.6190 | 1.7742 | 1.3415 | 1.0784 | 0.9016 | 0.7746 | 0.6790 | 0.6044 |\n\n---\n\n> **Q.** How to analyze the ICL of Mamba to simulate the gradient descent of a deep neural network.\n\n**A4.** Based on our reading of [Wu25], they encode the weights, activations, and gradients of a deep neural network into the embedding space of a Transformer. Each layer carries out a specific subtask: forward activation, back propagation of gradients, and weight update. Residual connections pass the resulting states to the next layer, so the whole training procedure for a deep neural network is simulated in one forward pass without changing the Transformer parameters.\n\nWe can use multi-layer Mamba + MLP to simulate these procedure. Now consider the following Mamba layer:\n\n$$ h_l^{(i)} = \\alpha h_{l-1}^{(i)} + (1 - \\alpha) u_l^{(i)} B_l \\tag{1}$$\n$$ o_l^{(i)}  = C_l^\\top h_l^{(i)} \\tag{2}$$\nwhere $B_l = W_B u_l, \\, C_l = W_C u_l$, and $u_l$ is the embedding for the $l$-th token.\n\nWe next show how to simulate one forward step of a neuron with weight $w$ in the deep neural network. Assumming that $u_l$ encode the $l$-th prompt and weight $w$, position indicator $p_l$: $u_l = [x_l, y_l, w, p_l]$. By (1) and (2),we have:\n$$ o_l^{(i)} = \\alpha u_l^\\top W_C^\\top h_{l-1}^{(i)} + (1 - \\alpha) u_l^{(i)} u_l^\\top W_C^\\top W_B u_l$$\n\nCarefully designing $W_C, W_B$ and the position indicator $p_l$ for the $l$-th token (e.g., $u_l^{(p_l)} = 1$ and $u_l^{(p_{l^\\prime})} = 0$ for $l \\ne l^\\prime$), we can ensure that $\\alpha u_l^\\top W_C^\\top h_{l-1}^{(p_l)} = 0$ and $(1 - \\alpha) u_l^{(p_l)} u_l^\\top W_C^\\top W_B u_l = (1 - \\alpha) \\frac{1}{1 - \\alpha} w^\\top x_l^\\top = w^\\top x_l$, and thus $o_l^{(p_l)} = w^\\top x_l$ stores the pre-activation for the $l$-th token. Applying a MLP (e.g. ReLU), we can get the activation for the $l$-th token with weight $w$.\n\nFor multi-layer neural network: By concatenating the weights of multiple neurons into $u_l$, we can simulate the forward prosedure of one layer of neural network. Then we can use residual connections to transfer the activation to the next layer of Mamba, and simulate the activation for the next layer of neural network.\n\nAs for the backpropagation process, MLPs can be used to simulate the gradient calculation, and finally add the gradients descent on the old weight $w$.\n\nTo summarize:\n- $W_B$ and $W_C$ can be used to project the weight $w$ and input $x_l$, and then calulate their inner product (pre-activation) $w^\\top x_l$, and then add an activation layer to calulate its activation, e.g., $ReLU(w^\\top x_l)$.\n- By carefully designing $W_C, W_B$ and the position indicator $p_l$, we can place the simulation results at a certain place in the embedding space.\n- By using residual connection, we can ensure the simulation results can be used by the following layers.\n\n---\n\nOnce again, thanks for your constructive feedback and we hope our response can resolve your concerns!"}, {"review_id": "tqTRUEjBcq", "reviewer": "Reviewer_eGZy", "summary": "This paper theoretically analyzes in-context learning for linear regression with state space model S6, and shows that the model implements online gradient descent for this task. The claim originates from the difference in how Transformer and S6 processes sequences (globally v. ‘online’), leading to a result that the estimator for S6 is a non-uniformly weighted combination of input data. The weights for this combination decrease as the context length increases, so that the later samples in the sequence have more influence on the final estimate for the query, similar to an online gradient descent setup. They propose specific techniques to handle optimization analysis of matrices $B, C$ that have coupled dynamics due to the problem / model structure. Theoretical results are also supported by experimental evidence.", "questions": "*Theory*:\n\n- Line 252: “...with partial a_i = 0 can also enable convergence, which is an undesirable scenario.”: could you elaborate why?\n- Line 315: “This induction answers the question \"How to avoid saddle points?\" …” - could you please elaborate how, beyond “preventing stagnation of partial diagonal entries of $C^TB$ at zero”? (I also couldn't follow the explanation from Line 580 in the Appendix) \n\n*Experiments*: \n\n- Are the Mamba models used in your experiments initialized following Assumption 4.1? \n- Additionally, is it possible to run online gradient descent on the same sequence that is used for Mamba inference, to see how the estimate $\\hat{y}_q$ compares while varying the sequence length?\n\n*Related work*: There exists prior theoretical work on using SSM for ICL in linear regression [1] which was not discussed in this paper. Could the authors compare and highlight how their results differ from this work?\n\n*Small typos*:\n\nLine 171: rewrited->rewritten  \nLine 223: recieving->receiving\n\n[1] Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond. https://arxiv.org/abs/2407.10005", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n- Interesting theoretical result on in-context learning for linear regression with S6 state space model.\n- I found the reasoning for online gradient descent in S6 quite well-explained intuitively comparing it to Transformers.\n- Theoretical findings on loss / weight matrix convergence / hidden state approximating weight vector ‘w’ are well-supported by experiments. \n- Paper is well written and clearly organized. \n\nWeakness: \n\nAssumption 4.1 seems too strict, is it realistic in practice? Specifically, with A = -I, w_\\delta = 0, and b_\\delta = constant: “Given the zero-mean and symmetric distribution of embeddings, w_\\delta can naturally converge to\n0 during gradient descent, and we fix it as 0 for simplicity.” - does this hold in experiments?", "was_revised": true, "final_justification": "The authors have addressed my clarifying questions about \n\n- some theoretical statements, \n- theoretical assumptions and their experimental validity \n- some experimental setup questions,\n- comparison to a related paper\n\nThey also compared their method to online gradient descent as in one of my questions. I've maintained my score of 5: Accept for this submission.", "rebuttal": "Thanks for your constructive feedback! We address your questions and concerns as follows.\n\n---\n\n> **W1.** **(1)** Assumption 4.1 seems too strict, is it realistic in practice? **(2)** Does \"w_\\delta can naturally converge to 0\" hold in experiments?\n\n**A1.** **(1)** The assumptions of theoretical proofs are typically more stringent than experimental requirements. We have tried using smaller $d_h$, smaller $N$ to do some experiments as follows, and the Mamba usually converge to the optimal solution.\n\nSetting: $d = 20$ and $N = 4, 6, 8, 10, 12, 14, 16, 18, 20$.\n\n|  N | 4 | 6 | 8 | 10 | 12 | 14 | 16 | 18 | 20 |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| **experimental loss**  | 8.5911 | 7.8292 | 7.7009 | 6.8235 | 6.4004 | 6.0612 | 5.9689 | 5.6193 | 5.1426 |\n| **theoretical loss** | 8.4484 | 7.8425 | 7.3173 | 6.8579 | 6.4526 | 6.0926 | 5.7706 | 5.4810 | 5.2190 |\n\nThe exprimental loss is close to the theoretical loss.\n\nThe following is the mean value and standard deviation of the loss for smaller $d_h$ (in 10 repeated experiments).\nWe set $d=4$, $N=30$, and the theoretical loss is 0.2954.\n\n| $d_h$ | 6 | 8 | 10 | 12 | 14 | 16 | 18 | 20 |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| **mean(loss)**  | 0.2912 | 0.2933 | 0.2899 | 0.2887 | 0.2929 | 0.2951 | 0.2967 | 0.2959 |\n| **std(loss)** | 0.0075 | 0.0055 | 0.0116 | 0.0052 | 0.0105 | 0.0097 | 0.0110 | 0.0142 |\n\nThe results shows that with smaller $d_h$, the Mamba can typically converge to the predicted pattern (which can achive the theoretical loss), and the deviation of loss is very small.\n\n**(2)**\nAs for $w_\\Delta$:\n\nTheoretically: The gradient for $w_\\Delta$ is $\\nabla_{w_\\Delta} \\mathcal{L}(\\theta) = C w_\\Delta$ for a positive $C>0$, so $w_\\Delta$ will converge to 0 by gradient descent $w_\\Delta(t+1) = w_\\Delta(t) - \\eta \\nabla_{w_\\Delta} \\mathcal{L}(\\theta) = (1 - C\\eta) w_\\Delta(t)$.\n\nExperimentally: We randomly initialize $w_\\Delta$ and perform experiments to track the norm of $w_\\Delta$ during training. The results are as follows.\n\n|  Epoch   | 0  | 10 | 20 | 30 | 40 | 50 | 60 | 70 | 80 |\n| --- | ---  | --- | --- | --- | --- | --- | --- | --- | --- |\n| $\\Vert w_\\Delta \\Vert_2$  | 0.8883 | 0.7513 | 0.4821 | 0.3331 | 0.2444 | 0.2026 | 0.1868 | 0.1799 | 0.1773 |\n| $\\Vert w_\\Delta \\Vert_2^2$ | 0.7891 | 0.5645 | 0.2324 | 0.1109 | 0.0597 | 0.0410 | 0.0349 | 0.0324 | 0.0314 |\n\nThe experimental results show that the norm of $w_\\Delta$ gradually decreases during the training process.\n\n---\n\n> **Q1.** Why \" $C^\\top B = Diag(a_1, \\dots, a_d)$ with partial a_i = 0 can also enable convergence. \" (Line. 252)\n\n**A2.** Considering $B = [b_1 \\dots b_d]$, $C = [c_1 \\dots c_d]$, for a certain i, if $b_i=c_i=0$ occur simultaneously at a certain step during training, then the gradients of $b_i$ and $c_i$ are both zero (see Lemma 5.2, line.258). Therefore, $b_i=c_i=0$ will always hold after that, and so does $a_i = c_i^\\top b_i = 0$. That is why we call it \"also enable convergence (to a saddle point)\".\n\nBecause $C^\\top B$ should be $\\frac{\\beta_3}{\\beta_1} I$ to minimize the loss, the case when some $a_i$ converge to 0 should be avoided.\n\n---\n\n> **Q2.** Why our induction \"prevents stagnation of partial diagonal entries of $C^\\top B$ at zero\" (Line. 315)\n\n**A3.** \nWe restate properties $\\mathcal{A}(t)$ and $\\mathcal{B}(t)$ (line. 307) for your reference:\n\n$ \\mathcal{A}(t): $\n\n$$ d_{h}/2 \\leq b_{i}^{\\top}(t) b_{i}(t), c_{i}^{\\top}(t) c_{i}(t), b^{\\top}(t) b(t) \\leq 2d_{h}. $$\n\n$ \\mathcal{B}(t): $\n\n$$ |\\beta_3 - \\beta_1 c_i^{\\top}(t) b_i(t)|\\leq\\delta(t)\\exp(-\\eta\\beta_1\\gamma t),\\quad | c_i^{\\top}(t) b_j(t)|\\leq 2\\delta(t)\\exp(-\\eta\\beta_1\\gamma t), \\quad | c_i^{\\top}(t) b(t)|\\leq 2\\delta(t)\\exp(-\\eta\\beta_2\\gamma t)+\\frac{\\delta(t)}{\\beta_2}\\exp(-\\eta\\beta_1\\gamma t).$$\n\nAs we discuss in **A2.**, partially $b_i=c_i=0$ are condition for saddle points. Our induction (Line. 307) can prevent this condition.\n\n(1) $b_i=c_i=0$ also means $b_i^\\top b_i = c_i^\\top c_i = 0$. Property $\\mathcal{A}(t)$ provide a positive lower bound for $b_i^\\top b_i$ and $c_i^\\top c_i$. So this case would not happen.\n\n(2) Property $\\mathcal{B}(t)$ guarantee that $c_i^\\top b_i$ will converge to $\\frac{\\beta_3}{\\beta_1}$ (optimal solution). Although $c_i^\\top(t) b_i(t)$ may be zero at some training iteration t, it will eventually converges to $\\frac{\\beta_3}{\\beta_1}$ (Note that $c_i^\\top(t) b_i(t) = 0$ does not mean that $b_i(t)=c_i(t)=0$).\n\n(3) **A brief explanation** for it is that, property $\\mathcal{B}(t)$ guarantee that $C^\\top B$ converges to its optimal solution, so it would not converge to the saddle point. \n\n---\n\n> **Q3.** Are the Mamba models used in the experiments initialized following Assumption 4.1?\n\n**A4.** Yes, our experiments setting strictly follows our theoretical setup. We follow Assumption 4.1 to set the hyperparameters and initialize $W_B, W_C$ with standard gaussian distribution $\\mathcal{N}(0, 1)$.\n\n---\n\n> **Q4.** Run online gradient descent to see how the estimate $\\hat{y}_q$ compares while varying the sequence length.\n\n**A5.** \n\nConsider a task $f(x) = w^\\top x$, we use online gradient descent (Online GD) to train $\\hat{f}(x) = \\hat{w}^\\top x$ with trainable parameter $\\hat{w}$. Specifically, given sample $x_i, y_i$ and loss $\\ell_i = \\frac{1}{2}(\\hat{w}^\\top x_i - y_i)^2$, the gradient of $\\hat{w}$ should be $\\nabla_{\\hat{w}} \\ell_i = (\\hat{w}^\\top x_i - y_i) * x_i$, and the update rule for $\\hat{w}$ is $\\hat{w}(t + 1) = \\hat{w}(t) - \\eta \\nabla_{\\hat{w}(t)} \\ell_i$ with $\\eta = 0.1$. We generate a new query $x_q$ and denote $\\ell_q = \\frac{1}{2} (\\hat{f}(x_q) - w^\\top x_q)^2 = \\frac{1}{2} (\\hat{w}^\\top x_q - w^\\top x_q)^2$ as the test loss. All settings for data generation are the same as Mamba's experiment. The results of $\\ell_q$ are as follows (taking the average of 500 experiments).\n\n|  Iteration   | 0  | 10 | 20 | 30 | 40 | 50 |\n| --- | ---  | --- | --- | --- | --- | --- |\n| Mamba  | 0.1179 | 0.1051 | 0.0936 | 0.0841 | 0.0761 | 0.0696 |\n| Online GD | 0.1137 | 0.0728 | 0.0448 | 0.0275 | 0.0167 | 0.0102 |\n\nNote that Mamba simulate the gradient for $\\min_{\\tilde{h}} \\frac{1}{2} \\Vert \\tilde{h} - \\beta y_i x_i \\Vert^2$ and update $\\tilde{h}$ by $\\tilde{h} _ i = \\tilde{h} _ {i-1} + (1-\\alpha)(\\beta y_i x_i - \\tilde{h} _ {i-1}) $ (Eq.(8), Line. 172), which is different from the gradient for loss $\\ell_i = \\frac{1}{2}(\\hat{w}^\\top x_i - y_i)^2$. So the decreasing rate for their loss is different.\n\n---\n\n> **Q5.** Related work: Compare and highlight how our results differ from [1].\n\n[1] Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond.\n\n**A6.**\nThe key distinctions can be summarized as follows:\n\n**Architecture**: The SSM studied in [1] is H3, while we consider Mamba.\n\n**ICL mechanism**: \n\n$$ H3: \\hat{y}_{q}=x_q^{\\top} WX^{\\top}(\\omega \\odot y) ~ \\mathrm{where} ~ X = [x_1 \\dots x_n] ~ \\mathrm{and} ~ y = [y_1 \\dots y_n]$$\n\n$$ Mamba: \\hat{y} _ {q} = x_{q}^{\\top} \\sum_{i=0}^{N-1} (1-\\alpha) \\alpha^{i+1} \\beta y_{N-i} x_{N-i} $$\n\nWhile H3 and Mamba both use the combination of $y_i x_i$ to estimate $w$, they differ from the ways of processing input tokens. H3 in [1] performs sample-weighted preconditioned gradient descent (WPGD) in a single, parallel pass over all tokens. In contrast, Mamba emulates online gradient descent, processing tokens sequentially and updating its state at every step.\n\n**Theoretical techniques**: [1] pre-specifies model weights and analyzes their expressivity. We consider training dynamics and establish convergence guarantees for a randomly initialized Mamba under gradient-descent training, and then derive its expressive power. To address the challenges of studying training dynamics, we develop new proof techniques.\n\n---\n\nOnce again, thanks for your constructive feedback and we hope our response can resolve your concerns!"}]

Author rebuttal:

Meta-review: This paper provides a rigorous analysis of the in-context learning (ICL) capabilities of Mamba, a popular sequence model. It follows a line of recent work analytically studying various sequence models on in-context linear regression tasks. The analysis focuses on 1-layer, no MLP models. Interestingly, it reveals that Mamba learns to implement a form of online gradient descent, which differs from the full-batch 1-step gradient descent method that (linear) self-attention has been shown to learn in various previous analyses.

This is a technically strong contribution on an actively researched topic. Although there are by now many analyses of the learning dynamics of self-attention on in-context linear regression, much less is known about modern state-space models, a model class of which Mamba is a prominent member. The paper reveals interesting differences to linear self-attention, and the analysis requires mild assumptions only. In particular, it applies to randomly initialized Mamba layers. An obvious limitation is its focus on 1-layer, no-MLP networks, but the reviewers and I agree that this is not sufficient to preclude from publication.
Decision: Accept (poster)