Towards Learning High-Precision Least Squares Algorithms with Sequence Models

Authors: Jerry Weihong Liu, Jessica Grogan, Owen M Dugan, Ashish Rao, Simran Arora, Atri Rudra, Christopher Re
Primary area: transfer learning, meta learning, and lifelong learning
Keywords: high precision, least squares, algorithm learning, Transformers, gated convolutions, linear regression, in-context learning

Abstract:
This paper investigates whether sequence models can learn to perform numerical algorithms, e.g. gradient descent, on the fundamental problem of least squares. Our goal is to inherit two properties of standard algorithms from numerical analysis: (1) machine precision, i.e. we want to obtain solutions that are accurate to near floating point error, and (2) numerical generality, i.e. we want them to apply broadly across problem instances. We find that prior approaches using Transformers fail to meet these criteria, and identify limitations present in existing architectures and training procedures. First, we show that softmax Transformers struggle to perform high-precision multiplications, which prevents them from precisely learning numerical algorithms. Second, we identify an alternate class of architectures, comprised entirely of polynomials, that can efficiently represent high-precision gradient descent iterates. Finally, we investigate precision bottlenecks during training and address them via a high-precision training recipe that reduces stochastic gradient noise. Our recipe enables us to train two polynomial architectures, gated convolutions and linear attention, to perform gradient descent iterates on least squares problems. For the first time, we demonstrate the ability to train to near machine precision. Applied iteratively, our models obtain $100,000\times$ lower MSE than standard Transformers trained end-to-end and they incur a $10,000\times$ smaller generalization gap on out-of-distribution problems. We make progress towards end-to-end learning of numerical algorithms for least squares.

Decision: accept (poster)
Reviewer scores: 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729086512602

thread_type:
official_review

<thread_object>
creator: reviewer_xpf9

<content>
summary:
This paper investigates whether sequence models can learn to perform numerical algorithms. This paper finds that Transformer fails to meet criteria, such as machine precision and numerical generality, and instead propose alternative architectures that overcomes these limitations.

strengths:
The strength of this paper is not only showing the limitation of Transformer, but also proposing alternative architecture that experimentally makes progress toward learning to solve numerical algorithms. Furthermore, there are several theoretical analyses to support experiments. For instance, it is interesting to show that, in Theorem D.34, single-head causal softmax attention cannot represent a square function.

weaknesses:
As mentioned in Discussion and Limitations, proposed methods seem to be challenging in maintaining stable and precise training with deep networks.

questions:
Out of interest, could you explore some theoretical analysis of proposed methods such as universal approximation or convergence of training? For instance, similarly in [Yun et al 2020] that shows the universality of Transformer, can we (expect to) show the universality of proposed architecture?
</content>
</thread_object>

</thread>
<thread>
creation date:
1729980299340

thread_type:
official_review

<thread_object>
creator: reviewer_tttt

<content>
summary:
This paper observes empirically that standard (softmax-attention-based) transformers cannot achieve machine precision or numerical generality for solving linear regression and proposes an alternative polynomial architecture based on the BaseConv module that achieves nearly machine precision and numerical generality close to that of gradient descent (GD).

strengths:
- The observations that standard transformers cannot solve linear regression to machine precision and are brittle to changes in distribution are interesting as they discredits a hypothesis supported by a plethora of recent literature that simplified and standard transformers learn to perform GD on linear regression problems. This is especially the case for the machine precision observation, as it is novel as far as I know (brittleness to changes in distribution has been observed elsewhere, e.g. [1] below).

- Theoretical and experimental claims are well-substantiated. The experiments are thorough and well-executed within the scope of the problem.

- The paper is very well-written.

weaknesses:
While I think that observation that standard transformers are not learning to perform GD or Newton's method on linear regression is an important and relevant contribution, I'm not sure that coming up with a new architecture and developing a training procedure for solving linear regression to machine precision is of interest to the ICLR community, for the following reasons:

- Outside of scientific ML, solving linear regression to machine precision is not of interest. The paper does not discuss whether the proposed model and training procedure can shed insights on other problems besides solving linear regression to MP.

- Within scientific ML, it is not clear why practitioners would want to solve linear regression with a large and costly model that requires 50-80 layers to solve even small-scale problems to machine precision (as suggested by the experiments) rather than simply running GD.

-  I can understand why it is interesting from a theoretical perspective, however to cater to a theoretical audience, I think the paper can use more theoretical development. In particular, I think the paper should move their proof of the inability of softmax attention to express multiplication from the appendix to the main body, and further develop this theory to show that standard transformers cannot express GD. Of the two theoretical results currently in the paper, one is borrowed from another paper, and the other is an informal statement that does not fully show the claimed expressivity of their proposed BaseConv model, since it does not show that BaseConv can achieve zero population gradient.

(Minor) Missing related works:

[1] Zhang et al., Trained Transformers Learn Linear Models In-Context, JMLR 2024 -- Like Ahn et al. 2024, shows that linear transformers learn GD.

[2] Collins et al., In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness, arxiv 2024 -- (and Related Works therein) -- also argues that standard transformers are not solving regression with GD.

questions:
n/a
</content>
</thread_object>

</thread>
<thread>
creation date:
1730059276673

thread_type:
official_review

<thread_object>
creator: reviewer_frtt

<content>
summary:
The paper shows that the in-context learning setup for transformers cannot learn linear regression solutions with machine precision or numerical generality, and identify architectural and optimization factors that contribute to this phenomenon. It then proposes a new recipe for training to machine precision, including using BaseConv architecture, explicitly learning gradients, and reducing stochastic gradient noise.

strengths:
The paper addresses a novel idea, namely learning high-precision algorithms with transformers. Their proposed training setup is well motivated by empirical observations.

weaknesses:
1. Not all the notation needed to understand the main text is defined in the main text itself, e.g. What are the arguments to MULTIPLY? This makes some parts hard to understand without jumping to the appendix.

2. As part of their training setup, the authors propose to learn GD iterates explicitly, which requires knowing the value of the gradient. It is not clear to me why this setup is comparable to the standard in-context setup, which never explicitly computes gradients of the regression instance, or how this scenario could arise in practice.

3. It is unclear whether the proposed setup can generalize to practical scenarios beyond simple linear regression.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732417006509

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for their time and effort reviewing our work and for their constructive comments! Reviewers appreciated the novelty of our insights about the precision limitations of Transformers [frTT, TtTT, XPf9]; in particular, Reviewer TtTT noted our findings are particularly surprising as they challenge “a hypothesis supported by a plethora of recent literature that simplified and standard transformers learn to perform [gradient descent] on linear regression”. Reviewers found our experiments “thorough and well-executed” [frTT, TtTT] and our theoretical results “interesting and well-substantiated” [TtTT, XPf9].

In this common response, we:
- contextualize the relevance of our contributions within the in-context learning and SciML literatures;
- demonstrate the generality of our insights by providing new experiments on the more practical setting of _in-context ODE solving_; and
- summarize changes to our manuscript.

Please find our comments for individual reviewers in their respective threads.
# Contextualizing our contributions
## In-context learning
As noted by Reviewer TtTT, **our observations about the precision limitations of Transformers directly challenge the prevailing hypothesis that Transformers implement in-context learning via gradient descent.** This hypothesis is argued by a long line of recent work [1,2,3,4,5], which is focused on the setting of in-context linear regression. In our work, we provide three pieces of evidence that suggest this story is _incomplete_:
1. **Transformer precision saturates with depth.** For the simplest least squares problems (small, noiseless, fully determined systems), we scale Transformer depth and show that MSE saturates $1,000,000 \times$ worse than we expect based on float32’s machine epsilon. In contrast, gradient descent provably converges to machine precision on least squares.
2. **Transformers are brittle to out-of-distribution inputs.** In contrast to theoretical convergence guarantees for gradient descent on least squares $Ax=b$, we observe that Transformers are not robust to changes to the distribution of $b$. We show that a simple rescaling of inputs by a factor of $10 \times$ worsens MSE by $100,000,000 \times$.
3. **Transformers struggle to implement basic linear algebra primitives precisely.** Specifically, we show Transformers struggle to precisely implement simple floating-point multiplications both empirically (scaling Transformers in depth, width, and training iterations) and theoretically (proving that 1-layer causal softmax Transformers cannot exactly implement multiplication). Since implementing gradient descent on least squares requires precise multiplications, our findings suggest that it is unlikely the “algorithm” Transformers learn to perform is gradient descent.

Our findings highlight that Transformers learn qualitatively different algorithms from gradient descent that are limited in both precision and generality. Our proposed techniques represent a first step towards developing model training recipes that can learn proper generalizable algorithms, with potential implications for both theoretical understanding and practical capabilities of sequence models.
## Scientific ML
Reviewers frTT and TtTT raised important questions about the relevance of learning gradient descent (GD) iterates with ML models in practice and about the generalizability of our insights to more realistic settings, e.g. for scientific tasks. We address these points here.

We first highlight that our work builds upon a broad effort in the scientific ML literature towards tackling problems (e.g. PDEs) traditionally solved with numerical methods [6,7,8]. 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.

We point out that the importance of high-precision ML is well-established. Specifically in scientific settings, **high-precision ML presents a unique challenge** that is crucial for:
- Accurately combining equational and data-driven modeling. This includes equation discovery [9], e.g. recovering equational knowledge from data, and physics-informed learning [10], e.g. enforcing equational knowledge within a model.
- Maintaining stable temporal rollouts in simulations of nonlinear phenomena, e.g. climate or turbulence modeling, where small errors can accumulate exponentially when making long time-horizon forecasts [11].

Although the SciML community has made exciting progress in recent years, numerical methods are still known to outperform existing ML methods in precision even on simple PDE benchmarks (e.g. refer to [12], a survey of 76 works in this space).
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper examines the in-context learning ability of transformers and shows that it cannot learn in-context linear regression solutions with machine precision or numerical generality. Additionally, the paper proposes a new approach for training to machine precision. The paper also provides empirical evidence to verify these claims.

The reviews for the paper were borderline to leaning positive. Most reviewers acknowledge that the paper solves an interesting problem and highlight the novelty. While the practical applicability of this approach is somewhat arguable, I think it is worth to let the community know about this issues and broadly understand/address them. I recommend acceptance.