Title: In-context Learning of Linear Dynamical Systems with Transformers: Approximation Bounds and Depth-separation
Venue: NeurIPS 2025
Accepted: yes
Keywords: In-context learning

linear dynamical systems

depth-separation
Abstract: This paper investigates approximation-theoretic aspects of the in-context learning capability of the transformers in representing a family of noisy linear dynamical systems. Our first theoretical result establishes an upper bound on the approximation error of multi-layer transformers with respect to an $L^2$-testing loss uniformly defined across tasks. This result  demonstrates that transformers with logarithmic depth can achieve error bounds comparable with those of the least-squares estimator. In contrast, our second result establishes a non-diminishing lower bound on the approximation error for a class of single-layer linear transformers, which suggests  a depth-separation phenomenon for transformers in the in-context learning of dynamical systems. Moreover, this second result uncovers a critical distinction in the approximation power of single-layer linear transformers when learning from IID versus non-IID data.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the in-context learning capabilities of linear transformers. The authors propose the following problem: learning, using a model parametrized by a multi-layer linear transformer, on sequences of next-token prediction trajectory that take the form of noisy linear dynamical systems, with context length equal to one, and design matrices of each LDS being sampled according to some distribution.\nAfter presenting this problem, the results of the paper are two-folds. First, the authors prove that provided the depth is larger than logarithmic in $T$ (the length of each trajectory), there exists a linear transformer that predicts well the last iterate of such LDSs (i.e., performs in context learning of these tasks), with a test error that scales as $\\log(T)/T$. Then, the authors prove that depth 1 is not enough to have a test error that decreases to 0.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\nS1) In-context learning capacities of transformers is a very timely subject, and the simplified problem studied in this paper is a very nice and elegant proxy to study this problem. It of course does not capture the complexity of real LLMs, but that's completely out of reach for now.\nS2) The depth-separation result is quite strong in my opinion: theory of deep learning and LLMs hardly capture such results, as far as I know. Such a result in the context of next-token prediction and ICL really is a strength.\nS3) The paper is overall very well written: intuitions, assumptions, results are clearly stated, in a way that is reader-friendly. I did not read the proofs fully, but skimming through them and the help of the proof sketch, makes them overall clear.\n\n\n**Weaknesses**\nW1) The proposed results characterize the fact that if a certain depth is satisfied, there exists some linear transformer that satifies desired properties. However, this does not takes into account optimization perspectives of the problem. \n\n\nOverall, I think the contributions and the clarity of this submission are quite strong, and definitely make it above the acceptance threshold.\n\nQUESTIONS:\nI have several questions and remark:\nQ1) How about long-context problems such as long-context linear dynamical systems (such as e.g. *Long-Context Linear System Identification, Yuksel et al., ICLR 2025*)?\nQ2) What would be the authors take on the optimization perspectives of the problem?\nQ3) In Eq. (4), I agree that putting a $sup_W$ in the test loss captures a form of worst-case robustness. However, would the depth separation result appear also if it were an average over the sampled matrices ?\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper theoretically studies ICL in the non-IID setting, within the linear dynamica system frameword. A upper bound for error of a multi-layer linear transformer and the error lower bounded for a singla layer linear transformer is provided.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1). This work provides a theoretical study of the ICL capabilities of transformers for linear dynamical systems, which is a novel contribution.\n\n2). The lower bound result demonstrating that a single-layer linear transformer lacks sufficient expressive power for in-context learning of linear dynamical systems, to me, it is particularly interesting.\n\nWeakness:\n\n1). Some existing works have already studied ICL for non-IID data theoretically, even if not explicitly framed as such. For example, \"Non-asymptotic Convergence of Training Transformers for Next-Token Prediction\".\n\n2). Theorem 1 focuses on the expressiveness of transformers, showing that they can efficiently optimize training loss for linear dynamical systems by mimicking an iterative algorithm via construction. However, this result does not guarantee that a transformer will actually converge to such a solution during pre-training.\n\n3). Theorems 1 and 2 rely on a linear transformer architecture, which is quite restrictive. While this is not a major criticism, introducing non-linear activations in the attention block would strengthen the theoretical contribution, especially for the lower bound.\n\n4). The lower bound result is based on constrained attention and mlp layers. Although the authors discuss a few sentences about the fully parameterized case, no theoretical results are provided. Could the constraints on the linear attention blocks be relaxed, even if the constraints for mlp layer are kept?\n\n5).  Experimental validation of Theorem 2 would strengthen the paper, since it would show the incapability of a linear transformer structure with one layer.\n\nQUESTIONS:\nSee Weakness section. Besides, is it possible to extend the lower bound to a multi-head transformer with linear attention?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the in-context learning of linear dynamical systems for transformers and presents two main results. The first is the approximation error for multi-layer transformers which implements the Richardson iteration. The second is a lower bound for a single-layer transformer.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n1. The structure of the paper is clear and is generally well-written.\n\n2. The lower bound which showcases the inability of a single-layer transformer to perform ICL on dynamical systems is interesting.\n\n## Weaknesses\n\nPlease see questions.\n\nQUESTIONS:\nQ1: My main question is about the significance of Theorem 1. [1] already proved that  transformers can implement ridge regression on nonlinear dynamical systems, then what new message does Theorem 1 bring compared with [1]? \n\nQ2: Theorem 1 used the algorithm unrolling technique, what's the benefit of this approach compared with the conventional gradient descent upadate in [1] [2] [3] [4] etc.?\n\nQ3: While Theorem 2 suggests that a single-layer Transformer is theoretically insufficient for ICL of linear dynamical systems, could this limitation be empirically verified? Additionally, although determining the exact number of layers required remains challenging from a theoretical perspective, an empirical investigation could provide valuable insights.\n\n\n[1] Guo, T., Hu, W., Mei, S., Wang, H., Xiong, C., Savarese, S., & Bai, Y. (2023). How do transformers learn in-context beyond simple functions? a case study on learning with representations. arXiv preprint arXiv:2310.10616.\n\n[2] Von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., & Vladymyrov, M. (2023, July). Transformers learn in-context by gradient descent. In International Conference on Machine Learning (pp. 35151-35174). PMLR.\n\n[3] Ahn, K., Cheng, X., Daneshmand, H., & Sra, S. (2023). Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 45614-45650.\n\n[4] Bai, Y., Chen, F., Wang, H., Xiong, C., & Mei, S. (2023). Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 57125-57211.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the in-context learning capabilities of transformers on non-IID data generated by noisy linear dynamical systems. The authors establish a depth-separation phenomenon. Their primary results show that while deep transformers with logarithmic depth can achieve an error rate of O(log(T)/T), which is comparable to the least-squares estimator, a class of single-layer linear transformers is fundamentally limited, suffering from a non-vanishing approximation error of Ω(1). This highlights that for certain classes of correlated, task-dependent data, network depth is a crucial architectural requirement for effective in-context learning, a finding that contrasts with results from IID settings.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n- The paper provides a significant depth-separation result for in-context learning in a non-IID setting, which is more realistic than prior work.\n- The analysis is sound and well-written, offering insights into why depth may be crucial for learning from correlated data sequences.\n\n**Weakness:**\n- The key lower bound result, which establishes the failure of shallow models, is proven for a significantly simplified single-layer transformer (one-dimensional system, linear attention, and a fixed identity MLP matrix). While the authors briefly discuss potential extensions, this simplification limits the immediate generality of the claim. It remains an open question how these results would change with more complex components like non-linear MLPs or softmax attention.\n- The work is purely theoretical, and it would be valuable to include numerical experiments to empirically validate the predicted performance gap between shallow and deep models, particularly within a more realistic Transformer architecture. In this context, the phrase “We also provided numerical results that confirmed…” at line 298 is unclear: could the authors clarify which results are being referred to and how they relate to the theoretical findings?\n\nQUESTIONS:\n- The lower bound is established for a strongly simplified single-layer model. While the authors discuss several possible generalizations (such as to multi-dimensional settings or non-identity MLPs) it remains unclear how they view the role of non-linearity in standard softmax attention. Do they expect the depth separation phenomenon to persist in the presence of such non-linearities?\n- The results identify a task that a one-layer linear transformer cannot learn. Have numerical experiments been conducted (using standard architectures and training) to empirically validate the predicted Ω(1) vs O(log(T)/T) error gap between shallow and deep models, or to explore potential improvements in these bounds?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}