Thank you for your thoughtful review! We appreciated that you found our experimental results with polynomial-based architectures promising and our theoretical results interesting.

We’d like to highlight our common response, where we contextualize the contributions of our work within the literatures of the in-context learning and scientific ML communities. Briefly:

• Our observations that Transformers fail to learn precise and general solutions directly challenge the prevailing hypothesis about in-context learning in the literature. This suggests that Transformers must be learning qualitatively different algorithms than gradient descent, with important implications for how we understand and potentially improve in-context learning.
• As far as we are aware, we are the first to investigate and isolate effects of the model architecture and optimizer on precision in a controlled setting. Our analysis of least squares serves as an essential testbed for understanding fundamental precision limitations that impact more complex numerical tasks, e.g. solving differential equations.

Please refer to the common response for a more detailed discussion. Below, we address your questions and comments.

W1: Stable training with deep networks remains a bottleneck

Thanks for this point! We agree that scaling to deeper networks is the key bottleneck to learning more complex numerical algorithms.

We want to highlight that disentangling the effects of model depth, architecture, and optimizer is non-trivial and comprises an important contribution of our work. Although precision has been a known challenge in the SciML literature (e.g. refer to [1], a survey of 76 works in this space), our work is, as far as we can tell, the first that investigates the separate effects of architecture and optimizer on precision in a simple controlled setting: least squares. We think the observation that scaling can detrimentally impact precision, even with a theoretically expressive architecture, is a surprising and important takeaway from our work, which the wider ML community may find interesting.

Q1: Universal approximation and training dynamics for BaseConv

Thanks for this thoughtful question! We address these two points below:

Universal approximation:

Universal approximation of our architecture in fact follows from Proposition D.43 (now D.44 in the revision), which proves that the BaseConv architecture can efficiently approximate arbitrary smooth multivariate functions using the universal approximation property of polynomials. Previously, this result was relegated to Appendix D.4. We have revised the manuscript to include a clearer discussion of this result in Section 4.1.

Training dynamics:

Thank you for this interesting question! We are also very interested in obtaining theoretical results about the training dynamics of BaseConv. However, we point out that the equivalent result for Transformers is rare except for simplified models -- the only such result we are aware of examines 1-layer non-causal linear attention under gradient flow [2] -- and represents an entire effort on its own. As such, we leave this question to future work.

As an aside, we want to emphasize that although results of the form “exact solution implies zero population gradient” exist in the literature [3,4], the converse (“zero population gradient implies recovery of exact solution”) is rare. Our existing results roughly say that BaseConv perfectly recovers the Square and Linear primitives if its population gradient is zero. To the best of our knowledge, ours is the first such result for sequence model architectures (such as Transformers) that have been used to create large language models. We emphasize this point in the revision of Section 4.1.

We thank you again for your insightful questions, which we believe have improved the clarity of our work!

References:

1. Nick McGreivy, Ammar Hakim. Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations. Nature Machine Intelligence, 2024.
2. Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. JMLR 2023.
3. Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.
4. Arvind Mahankali, Tatsunori B. Hashimoto, Tengyu Ma. One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention. ICLR 2024.

Proving softmax attention cannot express gradient descent on least squares:

Thanks for this interesting question! We believe the inexpressibility of Square and Multiply tasks is the “right” level of abstraction to anchor on, in that we believe it is the most fundamental observation that accounts for why softmax attention Transformers struggle to precisely learn numerical algorithms like GD in practice:

• Our proofs for Square and Multiply capture the main theoretical observation: softmax cannot express polynomials. At a high level, our proofs about the expressibility of Square and Multiply with softmax attention reduce to the simple but important observation that softmax can’t exactly express polynomials (e.g. $u_1^2$ or $u_1 u_2$ for real-valued inputs $u_1$, $u_2$). We believe that a proof for the inexpressibility of e.g. GD for least squares would reduce to the same basic observation. In our revision, we now include a brief discussion of this intuition in Section 3.3.
• Our primitives are fundamental. We show in Appendix D.2 that basic numerical algorithms, GD and Newton’s method, can be expressed as compositions of three linear algebra primitives: Read, Linear, and Multiply. The fact that Transformers struggle to learn the simple Multiply primitive provides strong evidence that Transformers would also struggle to learn numerical algorithms that require compositions of primitives, including Multiply.
• Our theory results are supported by empirical results. We show experimentally that Transformers struggle to precisely learn the Multiply synthetic, even when scaling in depth, width, and training iterations (Figures 7 and 8 in Appendix C.1).

We agree that it would be interesting to work towards a result that “softmax Transformers cannot implement GD for least squares”. However, we believe such a proof would only hold for a specific formatting/prompting of the least squares task, and the proof idea would be to show that GD for least squares is equivalent to computing a polynomial of inputs, which softmax cannot express exactly. One result we believe should be possible to prove with our techniques is that “1-layer causal softmax attention cannot implement our specific task formulation of GD for least squares”. Due to the limited discussion period, we chose to spend our time obtaining additional experimental results, but we would be happy to commit to producing the detailed proof of such a result if the reviewer believes it would be interesting to the community.

Proving BaseConv can train to zero population gradient:

This is also a great point -- thank you for highlighting this important detail! Our existing results roughly say that BaseConv perfectly recovers the Square and Linear primitives if its population gradient is zero. We want to emphasize that although results of the form “exact solution implies zero population gradient” exist in the literature [2,3], the converse (“zero population gradient implies recovery of exact solution”) is rare. To the best of our knowledge, ours is the first such result for sequence model architectures (such as Transformers) that have been used to create large language models. We emphasize this point in the revision of Section 4.1.

We are also very interested in obtaining theoretical results about the training dynamics of BaseConv. However, we point out that the equivalent result for Transformers is rare except for simplified models -- the only such result we are aware of examines 1-layer non-causal linear attention under gradient flow [4] -- and represents an entire effort on its own. As such, we leave this question to future work.

In response to your detailed and insightful comments, we believe the presentation and clarity of Section 4 has been greatly strengthened -- thanks again!

W4: Related work

Thank you for pointing out these missing related works! We now include them in the extended related work section of Appendix A.2.

References:

1. Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, Christopher Ré. Zoology: Measuring and Improving Recall in Efficient Language Models. ICLR 2024.
2. Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.
3. Arvind Mahankali, Tatsunori B. Hashimoto, Tengyu Ma. One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention. ICLR 2024.
4. Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. JMLR 2023.

Thank you for your thoughtful and detailed review! We appreciated that you found our observations insightful and our experimental and theoretical results thorough and well-substantiated. Below, we address your questions and comments.

W1: Transferring our insights beyond linear regression

This is an important point -- thank you for highlighting it! We focus on least squares as a first step for two key reasons:

• High precision is crucial for scientific ML, where small errors can compound catastrophically (e.g., in climate modeling and turbulence simulations). Least squares provides a controlled setting to examine fundamental precision limitations. (We further discuss this point in our response to W2 below.)
• Least squares serves as a useful primitive for differential equations. As a first step towards exploring these connections, we include new experimental results on solving ODEs in-context in Appendix C.4. Our insights about limited precision and generality transfer directly to the ODE setting, whereas our proposed techniques achieve high precision. Please refer to the common response for more details.

W2: Practicality of our least squares techniques

Thanks for this comment! We agree that learning to perform gradient descent with ML on least squares may be of limited practicality for its own sake. Rather, we think our results may be interesting to the broader scientific ML community as a fundamental analysis of precision bottlenecks within typical training setups. In our common response, we discuss the importance of high precision for SciML and contextualize our contributions within the existing SciML literature. Briefly, as far as we are aware, we are the first to investigate and isolate effects of the model architecture and optimizer on precision in a controlled setting. Please refer to the common response for a more detailed discussion.

W3: Theoretical results are not comprehensive (e.g. proving BaseConv will train to zero population gradient)

Thank you for this thoughtful and detailed comment! We first make one clarifying comment, then address the points you raise:

Key theoretical result (Theorem 4.1) is borrowed from another paper (Theorem H.21 of [1]):

This is an important point that deserves a clearer discussion in the manuscript -- thank you for bringing it up! It’s true that we use the statement of Theorem H.21 from [1], showing the equivalence of arithmetic circuits and the BaseConv architecture, as key theoretical grounding for our method. However, we want to highlight that for the class of numerical algorithms we consider in this work, our theoretical results are stronger. Specifically, while Theorem H.21 implies a poly-log-factor increase in parameters (specifically the number of layers) translating from general arithmetic circuits to BaseConv, in our work we show by construction that the primitives Read, Linear, and Multiply incur only a constant factor loss -- this improvement is especially important in the context of number of layers that we need. Since the numerical algorithms we consider are compositions of these primitives, this poly-log savings also holds for GD and Newton’s method. We emphasize this point in the revision of Section 4.1, and move the result we borrow from [1] to Appendix D.

Highlighting that softmax attention cannot express multiply:

Due to space limitations, we are unable to move the proof of this key result to the main text. However, in the revision, we now include a proper theorem statement for this result in Section 3.3, with a pointer to the full proof in Appendix D.2.4. We hope this improves the clarity of the main text!

Thank you for your thoughtful review! We appreciate that you found the high-precision algorithm learning setting novel and found our proposed techniques well-motivated.

We’d like to highlight our common response, where we contextualize the contributions of our work within the literatures of the in-context learning and scientific ML communities. Briefly:

• Our observations that Transformers fail to learn precise and general solutions directly challenge the prevailing hypothesis about in-context learning in the literature. This suggests that Transformers must be learning qualitatively different algorithms than gradient descent, with important implications for how we understand and potentially improve in-context learning.
• As far as we are aware, we are the first to investigate and isolate effects of the model architecture and optimizer on precision in a controlled setting. Our analysis of least squares serves as an essential testbed for understanding fundamental precision limitations that impact more complex numerical tasks, e.g. solving differential equations.

Please refer to the common response for a more detailed discussion. Below, we address your questions and comments.

W1: Clarity of notation in the main text (e.g. definition of the Multiply primitive)

Thank you for this comment! Towards making the main text more easily understandable on its own, we make the following changes to the manuscript:

• We move the definitions of the Read, Linear, and Multiply primitives to Section 3.3.
• We now introduce and discuss the connections between the simplified training setups we propose, explicit gradient descent (GD) and $k$-step GD, at the beginning of Section 5. We hope these changes improve the clarity of the work!

W2.1: Relation between the explicit gradient task and standard in-context least squares

This is an important point that deserves a clearer discussion in the manuscript – thank you for bringing it up!

The explicit gradient task and standard in-context least squares are connected via the multistep GD task. Specifically, given the standard GD iterates

$\mathbf{x}\_{i+1} := \mathbf{x}\_{i} - \eta \nabla \mathcal{L}(\mathbf{x}\_i),$

the goal of the $k$-step GD task is to predict $\mathbf{x}_k$ given inputs $(\mathbf{A}, \mathbf{b}, \mathbf{x}_0)$. Note that (up to a residual connection), the explicit gradient task is equivalent to $1$-step GD, and standard in-context least squares is equivalent to taking $k \to \infty$ (though in practice, Figure 4 suggests $k \approx 80$ is equivalent, up to machine precision). Thus, we view the explicit gradient task as a natural simplification of standard in-context least squares, where the $k$ parameter allows us to smoothly interpolate between the two extremes of difficulty. Note that at the end of Section 5, we show that our proposed techniques allow us to scale to $k=4$, and using our controlled setting, we identify model depth as the remaining bottleneck to learning more complex algorithms.

Previously, we briefly mentioned the $k$-step GD task at the end of Section 5, relegating a detailed definition to Appendix B.3.4 and discussing experimental results in Appendix C.3. We have added a brief discussion of these connections to the start of Section 5. Thanks again for raising this question -- we hope this discussion more clearly motivates the simplified training setups we investigate!

As far as we are aware, we are the first to investigate and isolate effects of the model architecture and optimizer on precision in a controlled setting: least squares. We make the following contributions:

• We find that typical training recipes for sequence models (e.g. Transformers, Adam, and StepLR schedulers) encounter surprising precision barriers when applied to numerical tasks. Our results about the limitations of Transformers are timely. Recent works [13,14] have investigated the use of Transformers for foundation models for differential equations, so we believe it is important to highlight precision bottlenecks that are fundamental to the Transformer architecture.
• In our controlled setting of least squares, we show that properly learning a simple numerical algorithm (GD) requires a major reworking of standard architecture and optimizer design choices. In doing so, we identify an important gap between expressivity and learnability. We believe this observation is particularly relevant to SciML, where theoretical backing for standard methods (e.g. PINNs [6], neural operators [7]) focuses on existence results like universal approximation theorems (e.g. [6,15,16]). Our work shows that expressibility is not enough -- better training procedures are also required to learn high-precision models, and our proposed techniques represent a first step towards developing them.

Additional experimental results on in-context ODEs

Beyond connections to prior in-context learning work, we focus on least squares because it serves as a useful primitive for differential equations. To demonstrate this, we include new experimental results on in-context ODE solving, addressing comments from Reviewers frTT and TtTT. We note that solving differential equations in-context with Transformers is a framework that has been explored in recent papers [13,14,17]. Here, we follow the setup from [17] and include a comprehensive discussion of experimental details in Appendix C.4.

We find that our observations from least squares transfer to the setting of in-context ODEs:

• Transformers struggle to learn precise solutions. We find that a 12-layer, 9M parameter Transformer model only achieves $O(10^{-4})$ MSE, almost $10^{10} \times$ worse than float32 machine precision would imply. Furthermore, as in least squares, we observe precision saturation with model size -- we find scaling the depth of the model by $2\times$ does not improve precision.
• Transformers exhibit brittle generalization. We observe that Transformers are not robust to changes to the distributions of ODE parameters, forcing functions, and boundary terms. For example, we show that for initial conditions, a simple rescaling by a factor of $10\times$ worsens MSE by $100,000\times$. Please refer to Appendix C.4 for more detailed results.
• Our proposed techniques obtain precise and general solutions. Prior work [17] shows that in-context ODEs can be solved by reducing to an equivalent least squares problem. We train a BaseConv architecture on explicit gradients for the equivalent least squares problems, and obtain $O(10^{-10})$ MSE ($1,000,000\times$ better than our best end-to-end Transformers) when applied iteratively.

We hope these preliminary results on differential equations demonstrate that our insights generalize to broader tasks of interest to the SciML community.

Summary of changes

• Move definitions of primitives to Section 3.3 (Reviewer frTT)
• Move theorem about inability of softmax attention to exactly express Square and Multiply to Section 3.3 (Reviewer TtTT)
• Restructure Section 4.1 to highlight our theoretical contributions: Theorem 4.1 (zero population loss implies perfect recovery of primitives) and Theorem 4.2 (universal approximation) (Reviewer frTT, TtTT, XPf9)
• Discuss connections between $k$-step GD, explicit GD, and standard in-context least squares tasks in Section 5 (Reviewer frTT)
• Incorporate additional related work (Reviewer TtTT), including relevant SciML papers, in Appendix A.2
• Restructure Appendix C:
  • Move discussion of LayerNorms/MLPs ablation to Appendix C.2
  • Include preliminary results on in-context ODEs in Appendix C.4 and Section 5 (Reviewers frTT, TtTT)
• Add a formal proof of the inability of softmax attention to exactly express Multiply in Appendix D (Corollary D.32) (Reviewer TtTT)