In-context Learning of Linear Dynamical Systems with Transformers: Approximation Bounds and Depth-separation

We thank the reviewer for their kind and helpful feedback and appreciate the positive assessment of our work. First, since several reviewers cited similar weaknesses, we have a general response:

$\**Experiments:**$ We ran some experiments to validate our theoretical results. Due to the current rebuttal guidelines, we cannot show the plots from our experiments. However, we plan to add our experimental results to the revised version. Our results are as follows. We trained transformers with both linear and softmax attention to learn random linear dynamical systems in-context. For computational reasons, the loss we use is averaged over the task distribution, as opposed to the supremum loss we use for theory. For one-dimensional systems with noise variance $\sigma^2 = 0.05$ and context length $T = 500$, we present the minimum test loss as as a function of depth in the following tables:

We also have numerical experiments for multi-dimensional systems (up to $d=5$), but the results are most illustrative for one-dimensional examples. Additionally, while we are unable to produce results for depth greater than $L=5$ due to the time constraint, we plan to include experimental results for deeper transformers in the revised version. Nonetheless, we hope that our inclusion of these experimental results strengthens the overall contribution of the paper. To summarize, our plots confirm the following points:

1. There is a clear separation in the error achieved by single-layer and multi-layer transformers. This provides evidence for Theorem 2 in our paper.
2. Increasing the depth beyond two layers marginally improves the error, but the most significant difference is between one and two-layer transformers.
3. For each choice of the number of layers $L$, the minimum loss achieved by the softmax attention transformer is comparable to the minimum loss achieved by the linear attention transformer. This suggests that the non-vanishing lower bound persists when using softmax attention. The minimum loss values are shown in the table below.

$\**Choice of architecture for the lower bound:**$ Broadly, the goal of our paper is to provably demonstrate the importance of depth for in-context learning when the data is not iid. We think that the linear transformer is a natural architecture through which to study this phenomenon, since several prior works have studied the in-context learning ability of this architecture on iid data. In particular, the recent works [1,2,3] which prove upper bounds on the loss for single-layer linear transformers on iid data, provide a direct baseline to compare with our lower bound. Currently, we are investigating whether our upper bounds can be extended to linear transformers with finite depth, as well as the approximation error of nonlinear transformers, such as those with with softmax attention modules; however, these directions are beyond the scope of the current work.

Concerning your specific questions:

1. (Extension to long-context problems): It would be very interesting to extend the positive results of our paper to linear dynamical systems which depend on a context length $p > 1$. Since our proof for the case $p=1$ constructs a transformer which implements the least-squares estimator, the statistical guarantees derived in [4] of constrained least-squares estimators could potentially be used to adapt the proof of our result to the long-context case (although the transformer construction would have to be modified, since the current proof uses the specific form of the least-squares predictor in the $p=1$ case). We plan to add a discussion of this problem in the conclusion of the paper and cite the reference [4].

2. (Optimization dynamics): In general, the optimization dynamics of multi-layer transformers for this learning task is highly nontrivial and beyond the scope of this work. However, recent work [5] has demonstrated that transformers linear transformers can provably learn to perform multi-step gradient descent for in-context linear regression using chain-of-thought (CoT) training. Since the ICL mechanism in [5] is similar to ours (multi-step gradient descent vs multi-step Richardson iteration), a natural extension of our work is to study the dynamics of CoT training to enforce the construction of Theorem 1. We plan to investigate this direction more carefully in the future.

3. ($\sup$ loss versus average loss): The reviewer asks an interesting question. Our experiments, which use the average loss instead of sup loss for computational reasons, suggest that a depth-separation is still present in this case. The heuristic behind the lower bound is that the covariance of a trajectory depends on the task $W$, and thus the expressions in the least-squares predictor which involve the covariance cannot simply be implemented directly by the attention weights. The $\sup$ loss naturally captures this limitation, but we plan to investigate whether a non-vanishing lower bound also holds under the $W$-averaged loss in the future.

References:

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

[2] Ahn, Kwangjun, et al. "Transformers learn to implement preconditioned gradient descent for in-context learning." Advances in Neural Information Processing Systems 36 (2023): 45614-45650.

[3] Mahankali, Arvind, Tatsunori B. Hashimoto, and Tengyu Ma. "One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention." arXiv preprint arXiv:2307.03576 (2023)

[4] Yüksel, Oğuz Kaan, Mathieu Even, and Nicolas Flammarion. "Long-Context Linear System Identification." arXiv preprint arXiv:2410.05690 (2024).

[5] Huang, Jianhao, Zixuan Wang, and Jason D. Lee. "Transformers learn to implement multi-step gradient descent with chain of thought." arXiv preprint arXiv:2502.21212 (2025).

We thank the reviewer for their kind and helpful feedback. Since several reviewers cited the same weaknesses, we have the following general response:

$\**Experiments:**$ We ran some experiments to validate our theoretical results. Due to the current rebuttal guidelines, we cannot show the plots from our experiments. However, we plan to add our experimental results to the revised version. Our results are as follows. We trained transformers with both linear and softmax attention to learn random linear dynamical systems in-context. For computational reasons, the loss we use is averaged over the task distribution, as opposed to the supremum loss we use for theory. For one-dimensional systems with noise variance $\sigma^2 = 0.05$ and context length $T = 500$, we present the minimum test loss as as a function of depth in the following tables:

We also have numerical experiments for multi-dimensional systems (up to $d=5$), but the results are most illustrative for one-dimensional examples. Additionally, while we are unable to produce results for depth greater than $L=5$ due to the time constraint, we plan to include experimental results for deeper transformers in the revised version. Nonetheless, we hope that our inclusion of these experimental results strengthens the overall contribution of the paper. To summarize, our plots confirm the following points:

1. There is a clear separation in the error achieved by single-layer and multi-layer transformers. This provides evidence for Theorem 2 in our paper.
2. Increasing the depth beyond two layers marginally improves the error, but the most significant difference is between one and two-layer transformers.
3. For each choice of the number of layers $L$, the minimum loss achieved by the softmax attention transformer is comparable to the minimum loss achieved by the linear attention transformer. This suggests that the non-vanishing lower bound persists when using softmax attention. The minimum loss values are shown in the table below.

$\**Choice of architecture for the lower bound:**$ Broadly, the goal of our paper is to provably demonstrate the importance of depth for in-context learning when the data is not iid. We think that the linear transformer is a natural architecture through which to study this phenomenon, since several prior works have studied the in-context learning ability of this architecture on iid data. In particular, the recent works [1,2,3] which prove upper bounds on the loss for single-layer linear transformers on iid data, provide a direct baseline to compare with our lower bound. Currently, we are investigating whether our upper bounds can be extended to linear transformers with finite depth, as well as the approximation error of nonlinear transformers, such as those with with softmax attention modules; however, these directions are beyond the scope of the current work.

Concerning your specific questions:

1. (Existing works): We had not encountered the work [4], but we agree that it also studies transformers' ability to learn from non-IID data. Compared to our work, the setting is quite different, since [4] focuses on sentences generated from a discrete vocabulary, whereas we study linear dynamical systems on the continuous state space $\mathbb{R}^d$. We have updated the literature review of our paper to include a discussion of [4] and other relevant papers we missed.

2. (Optimization): In general, the optimization dynamics of multi-layer transformers for this learning task is highly nontrivial and beyond the scope of this work. However, recent work [5] has demonstrated that transformers linear transformers can provably learn to perform multi-step gradient descent for in-context linear regression using chain-of-thought (CoT) training. Since the ICL mechanism in [5] is similar to ours (multi-step gradient descent vs multi-step Richardson iteration), a natural extension of our work is to study the dynamics of CoT training to enforce the construction of Theorem 1. We leave the investigation of this direction to future work.

3. (Using full attention parameterization and multi-head attention): We suspect the same proof strategy can be used to prove a lower bound for the full parameterization and for multi-head attention, as it just amounts to more additive terms in the limiting loss function. For the single-head case, we argue that the parameters we chose to zero out are not necessary to predict the dynamical system. Indeed, we know that the next token $x_{t+1}$ takes the form $x_{t+1} = Wx_t + \textrm{noise}$, and the loss function minimizes the difference between the transformer prediction $y_{\theta}$ and the conditional mean of $x_{t+1}$ given the history $x_1, \dots, x_t$. This motivates us to parameterize the transformer to output an estimator of the form $y_{\theta} =W_{\theta} x_t,$ where $W_{\theta}$ is a matrix depending on the transformer parameters $\theta$ and the history $x_1, \dots, x_t$. Under the full parameterization, the transformer output is of the form $y_{\theta} =W_{\theta}^{(1)} x_t +W_{\theta}^{(2)} x_{t-1}.$ This motivates us to choose a parameterization of the transformer which corresponds to setting $W_{\theta}^{(2)} = 0$.

References:

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

[2] Ahn, Kwangjun, et al. "Transformers learn to implement preconditioned gradient descent for in-context learning." Advances in Neural Information Processing Systems 36 (2023): 45614-45650.

[3] Mahankali, Arvind, Tatsunori B. Hashimoto, and Tengyu Ma. "One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention." arXiv preprint arXiv:2307.03576 (2023)

[4] Huang, Ruiquan, Yingbin Liang, and Jing Yang. "Non-asymptotic convergence of training transformers for next-token prediction." Advances in Neural Information Processing Systems 37 (2024): 80634-80673.

[5] Huang, Jianhao, Zixuan Wang, and Jason D. Lee. "Transformers learn to implement multi-step gradient descent with chain of thought." arXiv preprint arXiv:2502.21212 (2025).

We thank the reviewer for their kind and helpful feedback. Since several reviewers cited the same weaknesses, we have the following general response:

$\**Experiments:**$ We ran some experiments to validate our theoretical results. Due to the current rebuttal guidelines, we cannot show the plots from our experiments. However, we plan to add our experimental results to the revised version. Our results are as follows. We trained transformers with both linear and softmax attention to learn random linear dynamical systems in-context. For computational reasons, the loss we use is averaged over the task distribution, as opposed to the supremum loss we use for theory. For one-dimensional systems with noise variance $\sigma^2 = 0.05$ and context length $T = 500$, we present the minimum test loss as as a function of depth in the following tables:

We also have numerical experiments for multi-dimensional systems (up to $d=5$), but the results are most illustrative for one-dimensional examples. Additionally, while we are unable to produce results for depth greater than $L=5$ due to the time constraint, we plan to include experimental results for deeper transformers in the revised version. Nonetheless, we hope that our inclusion of these experimental results strengthens the overall contribution of the paper. To summarize, our plots confirm the following points:

1. There is a clear separation in the error achieved by single-layer and multi-layer transformers. This provides evidence for Theorem 2 in our paper.
2. Increasing the depth beyond two layers marginally improves the error, but the most significant difference is between one and two-layer transformers.
3. For each choice of the number of layers $L$, the minimum loss achieved by the softmax attention transformer is comparable to the minimum loss achieved by the linear attention transformer. This suggests that the non-vanishing lower bound persists when using softmax attention. The minimum loss values are shown in the table below.

$\**Choice of architecture for the lower bound:**$ Broadly, the goal of our paper is to provably demonstrate the importance of depth for in-context learning when the data is not iid. We think that the linear transformer is a natural architecture through which to study this phenomenon, since several prior works have studied the in-context learning ability of this architecture on iid data. In particular, the recent works [1,2,3] which prove upper bounds on the loss for single-layer linear transformers on iid data, provide a direct baseline to compare with our lower bound. Currently, we are investigating whether our upper bounds can be extended to linear transformers with finite depth, as well as the approximation error of nonlinear transformers, such as those with with softmax attention modules; however, these directions are beyond the scope of the current work.

Concerning your specific questions:

The main message of Theorem 1 of our paper is not only the construction of a specific transformer, but that our constructed transformer achieves a next-token prediction error bound of $O(\log(T)/T)$. This is where much of the work of proving Theorem 1 comes from, as it requires us to synthesize our particular construction with statistical guarantees of the least-squares estimator (LSE) on noisy data, and to control the error incurred on the event where the LSE bounds fail to hold. In contrast, Theorem 1 in [4] only presents the construction of a transformer which approximates the ridge regression algorithm for dynamical systems and does not bound the next-token prediction error of their construction. In addition, the two constructions are not directly comparable because their construction applies to a wider class of nonlinear dynamical systems, but our construction requires a much simpler transformer. In particular, while [4] constructs a transformer with $L + O(\log(1/\epsilon))$ layers, 5 heads, and nonlinear MLP layer a hidden dimension of $D_{hid} = 2D + d + 10$ to implement $L$ steps of gradient descent for ridge regression to accuracy $\epsilon$, our construction requires $L + O(1)$ layers, 1 head, and a linear MLP layer with a hidden dimension of $d$ to implement the least-squares predictor with the inverse covariance approximated by $L$ steps of Richardson's iteration. This demonstrates the benefit of our algorithm unrolling technique compared to the traditional GD approach (which is a different form of algorithm unrolling): it allows us to exploit the structure of the dynamical system considered herein, thus reducing the model capacity. References:

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

[2] Ahn, Kwangjun, et al. "Transformers learn to implement preconditioned gradient descent for in-context learning." Advances in Neural Information Processing Systems 36 (2023): 45614-45650.

[3] Mahankali, Arvind, Tatsunori B. Hashimoto, and Tengyu Ma. "One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention." arXiv preprint arXiv:2307.03576 (2023)

[4] Guo, Tianyu, et al. "How do transformers learn in-context beyond simple functions? a case study on learning with representations." arXiv preprint arXiv:2310.10616 (2023).

Thank you for the detailed response. Most of my questions are solved, especially regarding the comparison with existing work and the benefit of the algorithm unrolling technique, which could be a contribution itself, and I willl raise my score to 4.

I do have one more question regarding the experiment though, Theorem 2 states that the loss can't be smaller than a constant, does your experiment observe a platue for the one layer attention at the end of the training epochs compared with deeper layers?

We thank the reviewer for their kind and helpful feedback. Since several reviewers cited the same weaknesses, we have the following general response:

$\**Experiments:**$ We ran some experiments to validate our theoretical results. Due to the current rebuttal guidelines, we cannot show the plots from our experiments. However, we plan to add our experimental results to the revised version. Our results are as follows. We trained transformers with both linear and softmax attention to learn random linear dynamical systems in-context. For computational reasons, the loss we use is averaged over the task distribution, as opposed to the supremum loss we use for theory. For one-dimensional systems with noise variance $\sigma^2 = 0.05$ and context length $T = 500$, we present the minimum test loss as as a function of depth in the following tables:

We also have numerical experiments for multi-dimensional systems (up to $d=5$), but the results are most illustrative for one-dimensional examples. Additionally, while we are unable to produce results for depth greater than $L=5$ due to the time constraint, we plan to include experimental results for deeper transformers in the revised version. Nonetheless, we hope that our inclusion of these experimental results strengthens the overall contribution of the paper. To summarize, our plots confirm the following points:

1. There is a clear separation in the error achieved by single-layer and multi-layer transformers. This provides evidence for Theorem 2 in our paper.
2. Increasing the depth beyond two layers marginally improves the error, but the most significant difference is between one and two-layer transformers.
3. For each choice of the number of layers $L$, the minimum loss achieved by the softmax attention transformer is comparable to the minimum loss achieved by the linear attention transformer. This suggests that the non-vanishing lower bound persists when using softmax attention. The minimum loss values are shown in the table below.

$\**Choice of architecture for the lower bound:**$ Broadly, the goal of our paper is to provably demonstrate the importance of depth for in-context learning when the data is not iid. We think that the linear transformer is a natural architecture through which to study this phenomenon, since several prior works have studied the in-context learning ability of this architecture on iid data. In particular, the recent works [1,2,3] which prove upper bounds on the loss for single-layer linear transformers on iid data, provide a direct baseline to compare with our lower bound. Currently, we are investigating whether our upper bounds can be extended to linear transformers with finite depth, as well as the approximation error of nonlinear transformers, such as those with with softmax attention modules; however, these directions are beyond the scope of the current work.

Concerning your specific questions, we have the following comments:

1. Extending our results to nonlinear transformers is an ongoing direction. Our recent experiments suggest that the depth separation observed in Theorem 2 also holds for transformers with softmax attention. However, it is technically challenging to compute the limiting loss as we did for the linear case, because it involves computing moments softmax vectors of the linear dynamical system. It therefore seems that the softmax case may require a different analytical approach to the problem.
2. While our experiments do confirm a depth separation, it is difficult to validate these exact rates because they are asymptotic results. In addition, the optimal rate for in-context learning linear dynamical systems is unknown. The $\frac{\log(T)}{T}$ rate is optimal (up to the log factor) for estimating the matrix $W$ from the dynamical system trajectory, but it may not be optimal for our in-context learning problem, namely, predicting the next token from the history $x_1, \dots, x_T$.

References: [1] 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. [2] Ahn, Kwangjun, et al. "Transformers learn to implement preconditioned gradient descent for in-context learning." Advances in Neural Information Processing Systems 36 (2023): 45614-45650. [3] Mahankali, Arvind, Tatsunori B. Hashimoto, and Tengyu Ma. "One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention." arXiv preprint arXiv:2307.03576 (2023).