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

Thanks for your constructive feedback! We address your questions and concerns as follows.

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?

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.

Setting: $d = 20$ and $N = 4, 6, 8, 10, 12, 14, 16, 18, 20$.

The exprimental loss is close to the theoretical loss.

The following is the mean value and standard deviation of the loss for smaller $d_h$ (in 10 repeated experiments). We set $d=4$, $N=30$, and the theoretical loss is 0.2954.

The 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.

(2) As for $w_\Delta$:

Theoretically: 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)$.

Experimentally: We randomly initialize $w_\Delta$ and perform experiments to track the norm of $w_\Delta$ during training. The results are as follows.

The experimental results show that the norm of $w_\Delta$ gradually decreases during the training process.

Q1. Why " $C^\top B = Diag(a_1, \dots, a_d)$ with partial a_i = 0 can also enable convergence. " (Line. 252)

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)".

Because $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.

Q2. Why our induction "prevents stagnation of partial diagonal entries of $C^\top B$ at zero" (Line. 315)

A3. We restate properties $\mathcal{A}(t)$ and $\mathcal{B}(t)$ (line. 307) for your reference:

$\mathcal{A}(t):$

$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}.$

$\mathcal{B}(t):$

$|\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).$

As we discuss in A2., partially $b_i=c_i=0$ are condition for saddle points. Our induction (Line. 307) can prevent this condition.

(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.

(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$).

(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.

Q3. Are the Mamba models used in the experiments initialized following Assumption 4.1?

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)$.

Q4. Run online gradient descent to see how the estimate $\hat{y}_q$ compares while varying the sequence length.

A5.

Consider 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).

Note 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.

Q5. Related work: Compare and highlight how our results differ from [1].

[1] Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond.

A6. The key distinctions can be summarized as follows:

Architecture: The SSM studied in [1] is H3, while we consider Mamba.

ICL mechanism:

$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]$

$Mamba: \hat{y} _ {q} = x_{q}^{\top} \sum_{i=0}^{N-1} (1-\alpha) \alpha^{i+1} \beta y_{N-i} x_{N-i}$

While 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.

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.

Once again, thanks for your constructive feedback and we hope our response can resolve your concerns!

Thank you for your response and the additional experimental results; these have clarified my questions about the submission, and I have maintained my initial rating of accept. I would encourage the authors to incorporate the results / clarifications in this rebuttal in their main paper as well.

Thanks for your constructive feedback! We address your questions and concerns as follows.

W1. Limited Task Scope: Generalization to more complex or noisy real-world tasks is not tested.

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.

W2. No Formal Proof of Optimality Emergence.

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.

Under 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).

We 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.

W3. Transformer Baseline Might Be Underserved.

A3. As for the comparison between Transformer and Mamba, Our paper focuses on understanding their mechanisms.

We have performed some experiments on linear Transformers as follows.

Setting: d=10, N=10, 20, ..., 80.

Q. How to analyze the ICL of Mamba to simulate the gradient descent of a deep neural network.

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.

We can use multi-layer Mamba + MLP to simulate these procedure. Now consider the following Mamba layer:

h_l^{(i)} = \alpha h_{l-1}^{(i)} + (1 - \alpha) u_l^{(i)} B_l \tag{1} o_l^{(i)} = C_l^\top h_l^{(i)} \tag{2} where $B_l = W_B u_l, \, C_l = W_C u_l$, and $u_l$ is the embedding for the $l$-th token.

We 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: $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$

Carefully 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$.

For 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.

As for the backpropagation process, MLPs can be used to simulate the gradient calculation, and finally add the gradients descent on the old weight $w$.

To summarize:

• $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)$.
• 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.
• By using residual connection, we can ensure the simulation results can be used by the following layers.

Once again, thanks for your constructive feedback and we hope our response can resolve your concerns!

Thanks for your constructive feedback! We address your questions and concerns as follows.

W1. Some prior works explore similar connections between linear attention, gating mechanisms and online optimization. So the novelty should be better clarified.

A1. Thank you for sharing these papers, we provide a brief comparison with these paper as follows.

[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.

[Li, Yingcong, et al.] theoretically analyze the ICL mechanism for gated linear attention, while we analyze Mamba's. Their model architectures are different.

We will add the comparison in the revised manuscript.

Refference

• Yang, Songlin, et al. "Gated linear attention transformers with hardware-efficient training." arXiv preprint arXiv:2312.06635 (2023).
• Li, Yingcong, et al. "Gating is weighting: Understanding gated linear attention through in-context learning." arXiv preprint arXiv:2504.04308 (2025).
• Yang, Songlin, Jan Kautz, and Ali Hatamizadeh. "Gated Delta Networks: Improving Mamba2 with Delta Rule." arXiv preprint arXiv:2412.06464 (2024).
• Behrouz, Ali, et al. "Atlas: Learning to optimally memorize the context at test time." arXiv preprint arXiv:2505.23735 (2025).
• 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).

W2. More discussion of what makes our particular setup novel compared to prior analyses of linear attention.

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.

In contrast, our work:

1. Considers the dynamics for both $W_B$ and $W_C$, which will introduce non-convexity and complexity.

2. Initializes $W_B$ and $W_C$ with Gaussian distribution, which is more practical, and more difficult to establish the convergence given the non-convexity.

3. Trains the Mamba with gradient descent, which is typically more difficult to analyze than gradient flow.

These 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.

W3. The current construction (unequal weighting) appears suboptimal vs. optimal linear attention (with equal weighting).

A3. It is a limitation of Mamba under this construction. As for $\alpha$, see Eq.(11):

$h_l^{d+1} = \alpha h_{l-1}^{d+1} + (1-\alpha) y_l B_l$

$\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.

We 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.

W4. The experimental setup is simplistic.

A4. Our 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.

The experimental results are closed to the theoretical results.

W5. Empirical comparison to linear attention baseline would be valuable.

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. setting: d=10, N=10, 20, ..., 80.

Q1. (1) What is the underlying reason for this high dimension $d_h = \tilde{\Omega}(d^2)$? (2) If a smaller hidden dimension $d_h$ is used, are similar results still achievable? If so under what conditions?

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).

(2) $d_h = O(d)$ is enough under some conditions.

For 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).

Besides, we list some works that can address the similar optimization problem with $d_h = O(d)$ dimension as follows.

• [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)$

• [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$

• [Arora, Sanjeev, et al.] requires that the weight matrices are initialized by Gaussian distribution with a very small variation O(1/poly(d)).

Moreover, under standard gaussian distribution initialization $\mathcal{N}(0, 1)$, [Du, Simon, and Wei Hu.] require larger dimension $d_h$ than ours.

Similar 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.

Reference

• Braun, Lukas, et al. "Exact learning dynamics of deep linear networks with prior knowledge." Advances in Neural Information Processing Systems 35 (2022): 6615-6629.
• Dominé, Clémentine CJ, et al. "From lazy to rich: Exact learning dynamics in deep linear networks." arXiv preprint arXiv:2409.14623 (2024).
• Arora, Sanjeev, et al. "A convergence analysis of gradient descent for deep linear neural networks." arXiv preprint arXiv:1810.02281 (2018).
• Du, Simon, and Wei Hu. "Width provably matters in optimization for deep linear neural networks." International Conference on Machine Learning. PMLR, 2019.

Once again, thanks for your constructive feedback and we hope our response can resolve your concerns!

Thanks for your constructive feedback! We address your questions and concerns as follows.

W1. The value $\delta$ is not bounded (in Theorem 4.1).

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.

We 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).

W2. The parameter $\alpha$ is set rather than learned.

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.

W3. The case when $N < d$ should be considered.

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.

We 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$.

The exprimental loss is close to the theoretical loss.

W4. The result is only regarding one-layer Mamba. Analysis or experiments on multiple layers of Mamba are not included.

A4. Our 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.

W5. A statistical analysis to verify the assumptions and results would be helpful.

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.

Experimentally, We provide the error bar for the loss curve in Figure 2 (c) (line. 933).

Besides, 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.

The 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). We set $d=4$, $N=30$, and the theoretical loss is 0.2954.

The 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. We will include the statistical analysis in the revised manuscript.

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?

A6. (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.

Moreover, some related works can reflect the characteristics of this mechanism:

[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.

• Park, Jongho, et al. "Can Mamba Learn How To Learn? A Comparative Study on In-Context Learning Tasks." At ICML2024.

(2) Optimal 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.

Setting: d=10, N=10, 20, ..., 80.

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?

A7. (1) In our experiments, the model trained and tested on the same $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:

Given $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.

The similar phenomenon arises in the theoretical analysis of ​linear Transformers:

• 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.
• Wu, Jingfeng, et al. "How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?." ICLR. 2024.

Once again, thanks for your constructive feedback and we hope our response can resolve your concerns!