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
Venue: ICLR 2025 Poster (2025)
Keywords: high precision
least squares
algorithm learning
Transformers
gated convolutions
linear regression
in-context learning
Primary area: transfer learning, meta learning, and lifelong learning
TL;DR: This work uncovers limitations of Transformers for high-precision numerical tasks, and makes progress towards training gated convolutional models to machine precision as an alternative.

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.

Published as a conference paper at ICLR 2025
TOWARDS LEARNING HIGH-PRECISION LEAST SQUARES ALGORITHMM WITH SEQUENCE MODELS
Jerry Liu
Jessica Grogan', Atri Rudra Department of Computer Science & Engineering University at Buffalo Buffalo, NY, USA
Institute of Computational Mathematical Engineering Stanford University Stanford, CA, USA
Owen Dugan, Ashish Rao, Simran Arora, Chris Ré Department of Computer Science Stanford University Stanford, CA, USA
ABSTRACT
This paper investigates whether sequence models can learn 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 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, identify alternate class of architectures, comprised en- tirely of polynomials, that can efficiently represent high- 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 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 train near machine precision. Applied iteratively, our models obtain 100, 000 lower MSE than standard Trans- formers trained end-to-end and they incur a 10, 000x smaller generalization gap ofitt--distributtin problems. We make progress towards end-to-end learning of numerical algorithms for least squares.
INTRODUCTION
Least squares is the workhorse of modern numerics: it is well understood theoretically (Boyd & Vandenberghe, 2004; Trefethen & Bau, 2022) and has important downstream applications in science and engineering, including solving regression problems and differential equations (Orszag, 1972; Trefethen, 2000). Thus, least squares has gained interest as natural testbed for investigating how well ML models can learn to implement algorithms (Garg et al., 2022; Von Oswald et al., 2023).
surge of recent work suggests that Transformers (Vaswani al., 2017) can learn solve least squares using optimization algorithms like gradient descent and Newton's 's method (Akyürek et al., 2022; Fu et al., 2023; Ahn al., 2024; Bai et al., 2024: Zhang et al., 2023b). These arguments rest on two observations: (1) simplified Transformer architectures (e.g. non-causal linear attention) can exactly implement such algorithms; (2) standard (softmax attention) Transformers learn solutions with similar properties (e.g. convergence rates) as iterative algorithms. Crucially, these works focus on statistical least squares: they evaluate Transformer solutions in underddetrmiieddnoisy settings and compare to Bayes-optimal estimators. However, scientific applications like climate or fluids
Conresponding author: fww50@stantord.ecr
I Equal contribution.

---

Published as a conference paper at ICLR 2025
Statistical least squares
Applying learned GD iterates
1.0
Transformer (E2E)
100
Least Squares
10-2
0.8
10 4-
Averagg loss
12
Average loss
GD (true)
0.6
Attention
BaseConv
0.4
10
Lincar Attention
Transfomerr (E2E)
10-10
10-12
0.0
10-24
5
20
30
40
60
80
100
4 In no--ontext examples
Iterations
Figure I: Prior work focuses on statistical least squares: Transformers approximate Bayes-optimal estimators (left, adapted from Garg al. (2022)). In this work, focus numerical least squares: Transformers struggle to obtain precise solutions (inset). Using a high-precisionn training recipe, we train two polynomial architectures, BASECONv and linear attention, to perform high-precision gradient descent iterates on least squares (right): applied iteratively, they reach s 10-13 MSE
modeling require numerically precise solutions to least squares, e.g. to accurately model turbulence or to maintain stable temporal rollouts (Frisch, 1995; Wilcox, 2006). Prior works do not engage with the issue of high precision, is still unclear how well Transformers can solve least squares from this perspective.
In this work, we thus study whether existing approaches can solve numerical least squares. Specif- ically, numerical analysis requires that algorithms exhibit (1) machine precision, i.e. they should obtain solutions that are accurate to near floating point error, and (2) numerical generality, i.c. they are computational procedures that should apply broadly across problem instances. (See Section 2.1 for details.) Since traditional least squares algorithms (e.g. gradient descent and conjugate gradients) provably meet these criteria (Trefethen & Bau, 2022), it is crucial to evaluate machine learning methods against these same standards to determine their ability to learn numerical algorithms.
We focus learning the gradient descent (GD) algorithm for least squares. Our study has three parts:
We benchmark standard Transformers for precision/generality and identify an expressivity gap on least squares. When we replicate the standard end-to-end training setup for least squares with Transformers, we find that solutions do not exhibit machine precision and numerical generality (Figure la, 2). We identify high-precision multiplications as a fundamental challenge for softmax Transformers. Empirically, on synthetic element-wi wise multiplication task, we find precision scales poorly with larger Transformers: an 8-layer model trains an MSE that is still 10 million times worse than machine epsilon (Figure 3). Theoretically, we argue that a single layer of softmax attention is unable exactly express element-wise multiplications. Since implementing GD involves high-precision multiplications, this observation suggests standard Transformers are unable even precisely express GD, much less precisely learn the algorithm.
We identify an alternate architecture class which does not suffer from expressivity problems. Motivated by the expressivity limitations of softmax attention, we investigate alternate sequence mixer architectures. Prior work notes that non-causal linear attention is able to exactly implement algorithms like GD) and Newton's method (Von Oswald et al., 2023; Giannou et al., 2024) because it consists entirely of polynomials. We provide unified framework understand existing expressivity results from the lens of arithmetic circuits. In our work, we focus on BASECONV, a gated convolutional architecture, as case study, since is provably equivalent the entire class of polynomial architectures (Arora et al., 2023; 2024). We demonstrate that gated convolutions can express high-precision GD algorithm (=z 10 MSE when implemented in practice, Figure 4).
We identify an optimization precision bottleneck and propose high-precision training recipe. Although polynomial architectures can precisely express the GD algorithm, we find that standard training procedures struggle to find solution with sufficiently high precision (10-1 MSE, Figure 9). Therefore, towards disentangling precision bottlenecks during training, we first focus on the intermediate task of explicitly learning GD iterates. We identify stochastic gradient noise from
2

---

Published as a conference paper at ICLR 2025
minibatching as the main optimization bottleneck, and we find that simple metric, cosine similarity minibatch gradients (Liu et al., 2023b), is diagnostic of precision saturation. Towards reducing stochasticity, propose (1) learning rate (LR) scheduler that adaptively adjusts LR based on the cosine similarity metric, and (2) to apply EMA over optimizer updates to maintain strong gradient signal. Our high-precision training recipe allows
machine precision for the first time. We successfully train two 3-layer models, with BASECONV and linear attention, that learn to perform a single high-precision iteration of GD (Figure 1b). Excitingly, we can also learn multiple GD iterates at once, scaling up to 4 iterations with 10 MSE.
train MI. architectures
Overall, our work makes the following contributions: (1) we specify the desiderata of learning numerical algorithms, machine precision and numerical generality, and we demonstrate that standard Transformers fall short because of expressivity limitations of softmax attention; (2) we provide a unified framework using arithmetic circuits to investigate the expressivity of the class of polynomial architectures; (3) we address additional precision bottlenecks that emerge during training, even when using expressive polynomial architectures. Although we do not achieve end-to-end learning of GD, we make significant headway: we propose high-precision training recipe, which, for the first time, allows us learn iterates of the GD algorithm to near machine precision.
2 LEARNING NUMERICAL ALGORITHMS FOR LEAST SQUARES
In this section, we distinguish between statistical vs. numerical least squares and discuss the two properties we want our models to inherit from numerical algorithms: machine precision and numerical generality. We then briefly discuss prior work and, in doing so, tease apart two increasingly end-to- end notions of performing algorithms with ML: expressing an algorithm in-weights and learning algorithm iterates.
2.1 PROBLEM FORMULATION AND RELATED WORK
In this work, our goal to train a model that solves least squares problems: find æ RD given A E RNXD and b E V such that Ax b. Here, we briefly discuss two different perspectives on least squares: statistical (in the form of in-context learning) and numerical.
Statistical least squares. Originally motivated by applications in language modeling, prior works on solving least squares with Transformers typically take a statistical perspective. Transformers are trained using an in-context learning setup (Garg et al., 2022; Akyürek et al., 2022): problem instances Ax are sampled from a pre-specified distribution Dirain and the model is trained to minimize mean squared eтror (MSE) over Dtrain- Trained models then evaluated instances, both in and out-of-distriuution, and their performance is compared to Bayes-optimal estimators (Garg et al., 2022; Akyürek al., 2022). We define the in-context least squares training setup in Appendix B and leave a more detailed discussion of related in-context learning work to Appendix A.
unseen problem
Numerical least squares. In this work, we instead take a numerical perspective on least squares. prototypical numerical algorithm for least squares is GD. For problem instance Ax b, initialize Co. an estimate of æ, and iteratively improve our estimate via
F;41 2; V)xx),
(1)
where C()) ;= 2 Al - b1l2 is the squared residual error. GD exhibits two properties numerical algorithms that want models to inherit:
Machine precision. Numerical algorithms provably obtain high-precision solutions. For GD, obtaining higher precision simply requires performing more iterations until convergence ma- chine precision (i.e. the smallest achievable error with floating point arithmetic) (see Chapter 11 of Trefethen & Bau (2022)). this work, we use float.3 32 throughout, where machine precision is 2 ~ 1.19 10-7, so we hope for MSEs around 2-10 ~ 1.42 x 10-11.
Numerical generality. Although the convergence rate of GD depends on the spectrum of A (see Chapter of Boyd & Vandenberghe (2004)), the computational procedure comprising GD is general and can be applied broadly to problem instances. This is unlike statistical generalization and notions of in vs. out-oof-distrbutiin In this work, we are interested to study how closely ML models emulate the numerical generality of algorithms despite training on data distribution.
3

---

Published as a conference paper at ICLR 2025
2.2 OUTLINE OF THIS WORK
recent line of work probes the estimators learned by Transformers on in-context least squares, and suggests that Transformers learn to solve least squares by mimicking iterative algorithms like gradient descent and Newton's 's method (Von Oswald et al., 2023; Ahn et al., 2024; Fu al., 2023; Giannou et al., 2024). These works typically analyze simplified models theoretically and extrapolate to standard training regimes, backed by empirical observations:
results for simplified models, e.g. non-causal linear attention can implement GD using a specific choice model weights.
Theoretical
Empirical experiments training standard Transformers, c.g. decoder-only softmax attention Transformers trained end-to-dn and-to-end on in-context least squares display convergence rates reminiscent of iterative algorithms.
Although prior works suggest that trained Transformers learn to solve least squares with algorithms, it still unclear whether statements about learning algorithms in simplified settings transfer to standard Transformers trained end-to-end We note two significant gaps between previously reviously-analyzed settings and standard in-context least squares:
Architectural differences.
Standard Transformers softmax instead of linear attention, causal instead non-causal sequence mixers, and include MLPs and LayerNorms (Ba et al., 2016).
Optimization. Even if a model can express a precise and general algorithm, it is unclear whether the model can learn the algorithm from data.
In this work, we tease apart bottlenecks caused by architecture expressivity limitations (Sections 3.3, 4) and optimization difficulties (Section 5) by investigating two increasingly sophisticated notions of performing GD for least squares with ML: expressing GD in-weights and learning GD iterates.
TRANSFORMERS DO NOT LEARN NUMERICAL ALGORITHMS IN-CONTEXT
In this section, evaluate standard Transformers, trained end-to-end, on the criteria of machine precision and numerical generality. Surprisingly, we demonstrate that existing approaches fail to exhibit these properties: the precision of Transformer solutions (in MSF) saturates million times worse than machine precision (Section 3.1), and their performance further degrades as problem instances deviate from the model's 's training distribution (Section 3.2). These results suggest that Transformers are not learning proper algorithms as numerical analysis defines them.
Towards identifying expressivity bottlenecks, we identify three linear algebra primitives that comprise standard algorithms including GD and Newton's 's method (Section 3.3). We find empirically that Transformers struggle to implement
multiplication, and theoretically we argue that softmax attention faces an expressivity gap when trying to exactly express multiplications.
3.1 TRANSFORMERS STRUGGLE TO REACH MACHINE PRECISION
Recent work (Von Oswald al., 2023; Ahn et al., 2024; Fu et al., 2023; Giannou al., 2024) studying in-context least squares suggests that Transformers learn to mimic iterative algorithms like GD and Newton's method. Note that if Transformers the depth of the model should correspond to the number of iterations performed. We thus focus on the simplest case of fully determined least squares problems with fixed size design matrices and investigate whether precision improves as we scale to larger and deeper models.
able to implement iterative algorithms,
In Figure Ib, following prior work (Ahn al., 2024), Transformers end-to-end on least squares, scaling up to L = 64 layers. We compare their precision the convergence rate of the full-batch gradient descent algorithm least squares. For more details about the training setup, refer Appendix B.3.1.
fix the size of A E 20x5 and train
At first, Transformer precision scaling exceeds the convergence rate of gradient descent: this finding mirrors similar results reported by Fu et al. (2023), who suggest Transformers may instead be learning higher-order algorithms like Newton's 's method. However, we further observe that the precision gains for Transformers rapidly diminish, such that we observe very little difference in precision between L = 32 and = 64 layers. The deepest Transformer models we are able to train achieve an MSE around 10 In contrast, gradient descent converges linearly machine precision, almost 1, 000, 000x better precision. The diminishing returns of the Transformer precision scaling imply that Transformers are not learning standard numerical algorithms like GD.
4

---

Published as a conference paper at ICLR 2025
3.2 TRANSFORMERS DO NOT EXHIBIT THE GENERALITY OF GRADIENT DESCENT
We further investigate whether Transformers learn solutions to least squares that exhibit numerical gen- erality. Recall that models are trained on a predefined distribution of least squares problems, Pirain. If Transformers Iearn to solve least squares using stan- dard numerical algorithm like GD, then we expect the performance of the model should be robust to out-of-distrrbutiin inputs.
Out-af-Distribution Target Vectors
10a
100
10
10
Transfomer
MSE
GD (32 iters)
10.
GD juntil convergence)
BaseConv (GD iteri
For GD specifically, the convergence criterion (0 < 9 < 2/o depends on omars the maxi- mum singular value of A, and the optimal rate of convergence depends on the condition number of
10
10-50
10-12
A. K maz min (Boyd & Vandenberghe, 2004). Thus we specify our training distribution Dtrain over least squares problems (A E R b Ax R)0)
as follows. First, as in prior work (Garg et al., 2022), we sample the entries of A and æ i.i.d from a stan- dard Gaussian (0, 1). We then shift and rescale the singular values of A so that mma K 5. After training a 12-layer Transformer model on the in-context objective, we evaluate our model on out- of-distribution regression targets b.
Figure Transformers generalize poorly to onttof-distribution regression targets. In con- trast, using our training recipe, we train a BASECONV model to perform high-precision GD iterates. Applied iteratively, our BASEC- ONV model incurs 10, 000 less generaliza- tion error on out-of-distribution target vectors than the Transformer.
We define (o) by sampling each entry of x i..d. from [00, or) and computing b = Ax. Al- though the distribution b's and æ's changes with a, because the spectra of the A's is consistent, we know that GD with fixed choice of 77 will provably converge to high precision.
We find that compared to GD, the Transformer solutions are brittle to unseen regression target distributions. Simply scaling the inputs by a factor of 10x, the MSE of the trained Transformer degrades by factor of 108. In contrast, GD is robust: a fixed number of GD iterations consistently converges to the same order of magnitude of precision (Figure 2). The brittleness of the learned Transformer solution compared to GD again suggests that Transformers are not performing standard numerical algorithms.
3.3 IDENTIFYING AN EXPRESSIVITY GAP WITH STANDARD TRANSFORMERS
Toward understanding the limitations of the Transformer architecture, we start with GD and Newton's method, two algorithms used to solve least squares, and look into primitives that comprise them.
Linear algebra primitives. We observe that GD and Newton's 's method can be expressed as compositions of three simple linear algebra operations: sequence-wise read/write (READ), affine transformations (LINEAR), and element-wise multiplications (MULTIPLY). For input Nx
[[kk a:b] k + READ(i,j,a,,)00 [[ii, a:b]
k k=
LINEAR (u) uH, where H 1RD + Rdou is linear, MULTIPLY(a,, b, dout)( = u/: +duul @ ul: dob+d
In Appendix D.2, we define these primitives formally and describe how GD and Newton's method iterates can each be expressed as a composition of these primitives. Intuitively, READ is required to transfer information across the sequence dimension, LINEAR to transfer information across the hidden dimension, and MULTIPLY to compute high-degree interaction terms (like dot products or element-wii squaring).
Empirical analysis: standard Transformers struggle with multiplication. We train Transformers on synthetic formulations of these tasks to investigate how precision scales with model size. Details about training setups are in Appendix B.3.2.
5

---

Published as a conference paper at ICLR 2025
Linear
Multiply
10ml
AbDi, Lt Alln, no LX
Actn. UN Altn LN
Attn. LM Alin,
D5E
10
12
Number of layers
Figure 3: Precision vs. Transformer depth, with and without LayerNorms (LNN), on synthetic tasks. While shallow Transformers are able to learn the READ and LINEAR tasks to high precision 1100 with 2-layer models), precision on the MULTIPLY task scales poorly with depth (only 10 with 8-layer models).
In Figure 3, we show that even 2-layer Transformers are able to achieve 10-8 MSE on the READ and LINEAR tasks. However, we find that Transformers struggle with the MULTIPLY task: precision scales poorly with model depth, such that an 8-layer Transformer is only able to achieve 10 6 MSE. In Appendix C.I, we further show that precision on the MUI.TIPLY task also scales poorly with increased attention dimension, number attention heads, and MI.P upscaling factor.
Theoretical analysis: softmax attention struggles to exactly express multiplication. In Ap- pendix D.2.4, we provide a proof that single layer of softmax attention cannot exactly express the simple element-wiss squaring function SQUARE ()) ji u jú, J1. (intuitively, because softmax cannot implement polynomials). Crucially, we note that element squaring is a special case of element-wise multiply, so softmax attention cannot exactly implement MULTTIPLY either:
Theorem 3.1 (Informal statement of Theorem D.31 and Corollary D.32). One-layer single-headed (causal) softmax atlention cannot exactly represent SQUARE and MULTIPLY for all possible inputs.
Since precisely implementing numerical algorithms like GD hinges performing nigh-precision multiplications, this result suggests that the standard Transformer architecture struggles to precisely implement these algorithms because of fundamental expressivity gap.
We mention briefly that these findings do not conflict with prior results (Yun et al., 2020b) proving universal approximation theorems for Transformers, because they typically require parameter count to scale exponentially with dimension: see Appendix A.
4 ALTERNATE ARCHITECTURES CLOSE THE EXPRESSIVITY GAP
Motivated by the finding that softmax attention struggles to precisely express multiplications, next investigate alternate sequence mixer architectures. We are inspired by prior results (Von Oswald et al., 2023; Giannou et al., 2024) that show non-causal linear attention is able to exactly implement algorithms like GD and Newton's 's method. Thus, we focus on the class of polynomial architectures, ie. sequence mixers comprised entirely of polynomial operations, in order to explicitly bake in multiplications. In this section, we present a unified framework that integrates previous findings through the perspective of arithmetic circuits. Specifically, we focus BASECONV, gated convolutional model that combines element-wise multiplications (gating) with long convolutions. We work with BASECONY for two reasons:
Recent work (Arora et al., 2023; 2024) has shown that BASECONV is equivalent to general arith- metic circuits, including all polynomial architectures. Thus, existing results with other polynomial architectures, e.g. linear attention, transfer directly to BASECONV.
Empirically, gated convolutional models have been shown to perform comparably to attention- based architectures on tasks like language, audio, and DNA modeling (Arora et al., 2024; Nguyen et al., 2024; Zhang al., 2023a).
We emphasize that although we find gated convolutions are convenient to work with theoretically and empirically, we believe that other sequence mixer architectures may also be able to alleviate the expressivity issues we highlight in Section 3.3. In particular, we show promising empirical results for non-causal linear attention in Section 5.2.

---

Published as a conference paper at ICLR 2025
4.1 GATED CONVOLUTIONS ARE EQUIVALENT TO ARITHMETIC CIRCUITS
BASECONV definition. In this work, we focus on variant of the BASECONY operator from Arora al. (2023). Given an input u E RRNDD, BASECONV(u) is defined as:
(uWgate gate bgase)
(h * (uWin nn + bin) + boons)
W out bout
(2)
Linear Projection
Convolution
where the layer is parameterized by learnable filters E &RAND linear projections
Win: W gate: Wout E DxD and bias matrices Deonv Diny Dyute bout RNDD Here, O rep- resents the Hadamard product, and convolution two matrices is computed as convolution of the corresponding columns.
squares: BaseConv
BASECONVS can exactly express linear algebra prim- itives. In Appendix D.2.1, we provide explicit construc- tions of single-layer BASECONV models that exactly im- plement the READ, LINEAR, and MULTIPLY primitives from Section 3.3.
2314-
2314
We note that this result stronger than prior BASEC- ONV expressivity results (e.g. Theorem H.21 from Arora al. (2023)), which imply poly-log- factor increase in parameters (specifically layers) translating from arbitrary arithmetic circuits. Here, we show by construction that these specific primitives, and any circuits that are compo- sitions of them, incur only a constant factor loss.
MSE
12:44
32
BASECONVS can perfectly recover linear algebra prim- itives from data. In Appendix D.5, for SQUARE and LINEAR, we further show the following under mild as- sumptions. which our input distribution satisfies (see de- tails in Assumptions D.45, D.46, D.48, D.49):
Iterations
Figure BASECONv can express high- precision gradient descent: our imple- mentation of the weight construction reaches 10-13 MSE in practice.
Theorem 4.1 (Informal statement of Theo- rems D.62, D.60). BASECONV perfectly recovers SQUARE and LINEAR when it achieves zero population gradient w.r.t. MSE loss.
We note that although results of the form "exact solution implies zero population gradient" exist in the literature (Ahn et al., 2024; Mahankali et al., 2023), to the best of our knowledge. we the first to show the converse ("zero population gradient implies recovery of exact solution") for sequence model architectures. In Appendix C.I, we show that BASECONv models can learn the READ, LINEAR, and MULTIPLY primitives to high precision in practice (Figure 6).
BASECONVS are universal approximators. Finally, we show in Appendix D.4 that BASEC- ONV can efficiently approximate smooth functions by implementing polynomials:
Theorem 4.2 (Informal statement of Theorem D.39). Given a k-times differentiable function J [-1, + R. define f [-1,1]N-D + ND which applies J element-wise to all inputs. Then
Ve > 0, there exists BASECONV model approximates f to within error e, with O ( V -
+ depth
and O(ND) parameters, where (|) T..
We additionally prove universal approximation theorem for general smooth multivariate functions Appendix D.4 (Theorem D.44).
4.2 BASECONV CAN PRECISELY EXPRESS GRADIENT DESCENT FOR LEAST SQUARES
We now focus on the gradient descent algorithm for least squares. Explicitly, given least squares problem instance Ax and an initial iterate Fo. a single iteration of gradient descent computes
0 = xo V L(wu), where V.L. = AT(Azz b).
We provide two explicit (1)-lrre weight constructions to express a GD iterate using BASECONY in Appendix D.3.1. One requires a ODD) state size using non-causal model (i.e. each entry access
7

---

Published as a conference paper at ICLR 2025
any other entry of the sequence) and one requires O state size using causal model (i.e. entries cannot access later entries of the sequence). In Appendix D.3.2, we prove that both constructions are asymptotically optimal with respect state size.
In Figure 4, we implement our non-causal weight construction into a deep BASECONY model as proof of concept. We confirm that gated convolutions empirically implement high-precision gradient descent notably, roundoff errors due to machine precision do not significantly accumulate practice, despite scaling up depth-100 BASECONV model.
5 TOWARDS TRAINING MODELS TO MACHINE PRECISION
Although BASECONVS are expressive enough to solve least squares precisely, we find that simply swapping out softmax attention with BASECONV and training end-to-end is insufficient for high precision: our BASECONV models perform as poorly as standard Transformers (Figure 9). This suggests that additional precision bottlenecks are present during high-precision training. In this section, we thus investigate what it takes to train polynomial architectures machine precision.
Recent works (Rodionov & Prokhorenkova, 2023; 2024) on algorithm learning find that intermediate supervision is crucial for learning long computation trajectories. We hypothesize that end-to-ent least squares faces a similar challenge. Thus, to study high-precision optimization, we first investigate a simplified setting: learning perform explicit GD updates for least squares.
Using this task as a benchmark, we identify a fundamental bottleneck in figh-prccssion regimes, gradient variance from minibatching, and we identify metric based on cosine similarity of successive gradients that is diagnostic of precision saturation during training. We then propose high-precision training recipe, which for the first time allows to train ML models to machine precision. Using our training recipe, we learn to perform explicit GD updates to 10 average MSE (Figure 1b), and we can also learn up iterates of GD once with an MSE of 10-10 (Table 7).
Simplifying the training setup. We first define sequence of k-th iterate tasks, where the goal is explicitly produce the k-th iterate of GD given least squares problem instance (A, b), initial iterate Po. and step size y):
{a1 b (aN: (x)) mo) - Øg where 24 = F; - v)VL() i [k - 1]. We then define the explicit gradient task, where the goal is produce the GD update vector:
{(ai, bi)..... (aN,br), bw), VL(xo).
(5)
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 taking k xo. Thus, the explicit gradient task is a natural simplification of standard in-context least squares, and the k-th iterate task allows us to smoothly interpolate between the two extremes difficulty. Refer to Appendix for more details.
5.1 TOWARDS HIGH-PRECISION TRAINING RECIPE
Our theoretical results in Section imply that 3-layer BASECONV is expressive enough to solve the explicit gradient task, so we use training 3-layer BASECONY on this task our benchmark for studying the challenges of high-precision learning.
Precision saturates with standard training procedures. Motivated by prior work (Garg et al., 2022; Von Oswald et al., 2023; Ahn et al., 2024), we start by investigating two basic optimization procedures: Adam with constant learning rate (LR) and with exponentially decaying LRR
In Appendix C.2 (Figure 10), we sweep initial LR and LR steprate across 2-3 orders of magnitude for constant and decaying LR schedules. We find:
Precision saturation occurs with both constant and decaying LR schedules. After a number of training iterations, the average loss saturates and is unable to improve. We note that this occurs even while gradients magnitudes and LR non-zero.
Slower-decaying LR schedules perform better but require exponentially more training iterations. Figure 8, we further analyze this phenomenon in the simpler case of 1-layer Transform- ers/BAsECONVs on the MULTIPLY synthetic. We observe a power-law relation between precision
8

---

Published as a conference paper at ICLR 2025
Cosine similarity indicates precision saturation
Adaptive LR improves precision
Adaptive LR boosts
gradient signal
Cosine similarty
10
Cosine similerity
10
0.5
Averagge loss
Initial
4
0.0
0.0
0.5
1.0
1.5
1x10-
0.0
0.5
1.0
1.5
1e6
3x10-2
1e6
1x10-2
10-2
10-2-
10-
Averoge
Adaptive
10-5
20-5
10"1
10.4.
10-9.
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
Training Rers
1c6
Training iters
le6
Training iters
1c6
Figure 5: Gradient metric is predictive of precision saturation (left). We propose simple adaptive LR scheduler that alleviates precision saturation (middle). Adaptive LR effectively boosts gradient signal during training (right).
and number of training iterations sweep steprate; although it may be possible to train to high precision in theory, this approach seems infeasible in practice.
With aggressively-decaynn
I.R schedules, higher initial I.R is better. For fixed scheduler step rate, increasing initial LR leads to significant improvements in final MSE, e.g. in Figure 10, an improvement of x000 simply by increasing initial LR from 10-3 to 10 Choosing LR that is too large causes training instability, so in practice we set it to the largest value that trains stably.
Our analysis suggests that Adam with an exponentially decreasing LR scheduler gets only part of the way to machine precision training recipe. We next address the issue of precision saturation.
Stochastic gradients bottleneck precision. We identify minibatch gradient variance as the main source precision saturation. Although our goal is to minimize the expected loss over problem instances from Dtroin. in practice minimize over finite minibatch samples instead. Minibatch training is the standard in ML, but interestingly we find that the variance in minibatch gradients can dominate the population gradient signal in high-precision regimes, causing the loss stagnate.
To demonstrate this, we define a simple metric to assess the strength the gradient signal during training. At given training step, take the current model weights, sample different minibatches of least squares problems, and compute the minibatch model gradients [gI....,9h J. We then compute the average cosine similarity between all pairs, as in Liu et al. (2023b):
2
g
0g = n(n 1)
(6)
99
Tal2219,11
We observe that this cosine similarity metric is predictive of precision saturation across MSE scales and optimizer hyperparameters (Figure 5).
adaptive LR scheduler boosts gradient signal beyond precision saturation. We thus propose an adaptive LR scheduler based on the gradient variance. Our scheduler is motivated by two intuitions:
Whenever the cosine similarity metric is high, gradient signal is strong. In order refine the precison bits of the model weights, we need to slowly decrease the LR.
Whenever the cosine similarity metric is low, the model weights are stuck local region of the loss landscape. To allow the model to escape this region, we need to increase the LR.
The basic scheduler we use in this work simply decreases the LR exponentially while the metric is above threshold oth. and increases the LR instead if the metric is below oth- In Figure 5, we show that this simple approach alleviates the loss saturation phenomenon: we see boost in population loss across our LR settings, and observe the models consistently improve We note that proper choice of LR hyperparameters is still crucial for efficient convergence to machine precision we leave speeding up the convergence rate via better adaptive schedulers to future work.
continue training.

---

Published as a conference paper at ICLR 2025
Exponential Moving Average (EMA) over optimizer updates. Finally, motivated by our obser- vation that gradient variance bottlenecks precision and inspired by recent works (Lee et al., 2024; Pagliardini et al., 2024), we apply an additional EMA over Adam's 's update vectors to help smooth out minibatch noise. Empirically, we find this boosts the final MSE by as much as 100, 000 x on the explicit gradient task: see Appendix C.2 (Figure 11).
Our training recipe for efficient high-precisin preciiseciionn convergence thus involves two techniques: (1) an adaptive LR scheduler that exponentially increases decays LR according to the cosine similarity metric; and (2) applying EMA over optimizer updates.
5.2 LEARNING HIGH-PRECISION GRADIENT DESCENT WITH POLYNOMIAL ARCHITECTURES
Using our training recipe, we successfully train two 3-layer models with polynomial architectures, BASECONV and non-causal linear attention, on the explicit gradient task. For the first time, we are able to train to near machine precision: we achieve an average loss of 10 13 MSE.
In Figure 1b. we slot our trained models into the standard GD algorithm, using their predictions in place of the true least squares gradients VL. Specifically, for a least squares problem Ar and initial iterate Zp, we repeatedly compute 41+1 == 2; m,, where A; == To [A b, .x.) is the prediction of the model. We iteratively apply the model until convergence to a fixed point Coo.
We find that both our models achieve high precision. In this setting, we reach average MSE of 10-12 12 (Figure I, right): this is 100, 000x better MSE than the biggest Transformers we are able to train end.to-end. Moreover, our BASECONv model exhibits better numerical generality than the Transformer, incurring a 10, 000 smaller generalization gap on problems outside its training distribution (Figure 2). Interestingly, we find that our linear attention model exhibits markedly worse generality: its out-of-distribution performance nearly matches the Transformer's, and the model iterates eventually diverge: see Figure 13.
Learning k-iterates of GD for larger k. We find that our training recipe also allows us to learn up to = iterates of GD at once with 10-10 MSE: see Table 7 and Figure 14 for results. We are not able to stably train deeper models without reintroducing non-polynomial normalization techniques like LayerNorms, which causes precision bottlenecks. For small k we observe that LayerNorms worsen precision by over 1, 000 See Appendix C.3 for details.
Experiments with in-context ODE solving. Finally, towards high-precision ML for more realistic tasks, we provide preliminary results on in-context ODE solving. We find that our proposed techniques outperform standard Transformers by up to 1, 000, x00x in MSE (up to s 10 with iterative BASECONVS vs. R 10-4 with 12-layer Transformers). See Appendix C.4 for details.
DISCUSSION AND LIMITATIONS
In this work, we investigate learning to solve least squares from a numerical perspective. We find that Transformers fail to learn solutions that exhibit the properties of machine precision and numerical generality. Disentangling effects from the model architecture and optimizer, we find that standard design choices perform surprisingly poorly from the lens of numeries. We identify expressivity limitations with softmax attention, and find surprisingly that even MLPs and LayerNorms significantly affect precision (up 1, 000, 000 worse MSE on the explicit gradients task). On the optimization front, we find stochastic gradient noise from minibatch training becomes a precision bottleneck in high-precision regimes. We propose an adaptive LR scheduler that alleviates this issue on simplified task, but we suspect that this issue remains fundamental challenge on harder problems. Crucially, although we make progress toward learning to solve numerical least squares end-to-end, our techniques struggle to maintain stable and precise training with deep networks.
We note that the numerical criteria we consider in this work represent fundamentally different type of learning and generalization from statistical notions that are prevalent in MI. We believe these numerical perspectives may be relevant to the wider scientific ML community. For example, existing approaches to solving PDEs have shown promise but are known to be brittle outside their training distributions (Wang& & Lai, 2023; Rathore et al., 2024). This inhibits their usefulness in nigh-impact applications like climate or fluids modeling, where high precision and robustness crucial. We believe learning to implement precise numerical algorithms directly from data is an exciting prospect that has the potential to unlock new capabilities across science and engineering.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
We provide all the code and configuration files necessary reproduce our experiments https: https://githut githhb.. done using synthetic data and tasks. All experiments were conducted using PyTorch on NVIDIA A100/H100 GPUs. Detailed hyperparameters (learning rate, batch size, and optimizer settings) and proofs of all theoretical claims are provided in the supplementary materials.
In this work, all experiments are
ACKNOWLEDGMENTS
We thank Yasa Baig, Mayee Chen, Rajat Dwaraknath, Sabri Eyuboglu, Chris Fifty, Neel Guha, Hermann Kumbong, Benjamin Spector, Aman Timalsina, Alyssa Unell, Ben Viggiano, Michael Zhang, and Dylan Zinsley for their helpful feedback and discussion during this work.
We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize); NSF under Nos. CCF2247015 (Hardware-Awand Awarre.. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), 1937301 (RTML), DGH-22146755 (GRFP), and PHY-2019786 (IAIFI); US DEVCOM ARL under Nos. W991NF-23-2-0188 (Long-context) and W91 1 (Interactive Human-AI Teaming); ONR under Nos. NO001423 12333 (Deep Signal Processing); Stanford HAI under No. 247183; NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Eriesson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI- GCP Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Meta, Google, and VMWare. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, NSF, ONR, or the U.S. Government. JL is supported by the Department of Energy Computational Science Graduate Fellowship under Award Number DE-SC00231 12. JG and AR's ' research supported by NSF grant CCF#2247014. OD supported by the Hertz Foundation Fellowship, the Stanford Enight-Henness, Scholarship, and the NSF GRFP.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Advances in Neural Information Processing Systems, 36, 2024.
Kabir Ahuja, Madhur Panwar, and Navin Goyal. In-context learning through the bayesian prism. arXiv preprint arXiv:2306. 04891, 2023.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algo- rithm in-context learning? investigations with linear models. arXiv preprint arXiv:2211.1566) 2022.
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Arhitec- tures and algorithms. arXiv preprint arXiv:2401.12972 2024.
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and Improving Recall in Efficient Language Models, 2023.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall- thalll-trouughuu tradcoff, 2024.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016. URL https:/ //////avvvvv//////////////e//aaa /ttp//iivrrg/////////../......lt orgg/bb/// 607 04400.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.
11

---

Published as a conference paper at ICLR 2025
Stephen P Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1101, 2020.
Wuyang Chen, Jialin Song. Pu Ren, Shashank Subramanian, Dmitriy Morozov, and Michael W Mahoney. Data-efficient operator learning via unsupervised pretraining and in-context learning. arXiv preprint arXiv:2402. 15734, 2024.
Xiang Cheng, Yuxin Chen, and Suvrit Sra. Transformers implement functional gradient descent to learn non-linear functions in context, 2024. URL /atpiv/.////vvav/a//eeeeeeeiaeve httpa/awwwvvvorraaa/aaaaaaa 2ttp://abbbb11222...... 06528.
David Chiang, Peter Cholak, and Anand Pillay. Tighter bounds on the expressivity of transformer encoders. In International Conference on Machine Learning, pp. 5544-5562. PMLR, 2023.
Liam Collins, Advait Parulekar, Aryan Mokhtari, Sujay Sanghavi, and Sanjay Shakkottai. In- context learning with transformers: Softmax attention adapts to function lipschitzness, 2024. URL
/eetps:/azzee org/ /bs////2/22 11639.
D. Jackson. The theory of approximation. Amer. Math. Soc. Colloq. Publ., vol. 11, Amer. Math. Soc, Providence, R. I., 1930.
Tri Dao, Nimit S Sohoni, Albert Gu, Matthew Eichhorn, Amit Blonder, Megan Leszczynski, Atri Rudra, and Christopher Ré. Kaleidoscope: An efficient, learnable representation for all structured linear maps. arXiv preprint urXiv:2012.14966 2020.
Ishita Dasgupta, Andrew K Lampinen, Stephanie CY Chan, Antonia Creswell, Dharshan Kumaran, James L McClelland, and Felix Hill. Language models show human-like content effects on reasoning. arXiv preprint arXiv:2207.0705: 2022.
Uriel Frisch. Turbulence: the legacy of AN Kolmogorov. Cambridge university press, 1995.
Daniel Y Fu, Tri Dao, Khaled Saab, Armin Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212.1400 14052, 2022.
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint arXiv:2310. 170866,, 2023.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-3053 2022.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. In International Conference on Machine Learning, pp. 11398-11 1442. PMLR, 2023.
Angeliki Giannou, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, and Jason D Lee. How well can transformers emulate in-context newton's 's method? arXiv preprint arXiv:2403.03118,, 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:23/2.00752 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396 2021.
Michael T Heideman and C Sidney Burrus. Multiplicatiie complexity, convolution, and the DFT. Springer, 1988.
Maximilian Herde, Bogdan Raonić, Tobias Rohner, Roger Käppeli, Roberto Molinaro, Emmanuel de Bézenac, and Siddhartha Mishra. Poseidon: Efficient foundation models for pdes, 2024. URL htttps://aret org/ abs////////// /205 19400.9900.... 01.
12

---

Published as a conference paper at ICLR 2025
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023.
Jaerin Lee, Bong Gyun Kang, Kihoon Kim, and Kyoung Mu Lee. Grokfast: Accelerated grokking by amplifying slow gradients, 2024. URL https: wtwpvvarivvvv/va////v//..///..vv .twpe/i.o.o.oooooooo.... abs//405 20233.
Jerry Weihong Liu, N Benjamin Erichson, Kush Bhatia, Michael W Mahoney, and Christopher Re.
Does in-context operator learning generalize to domain-shifted settings? The Symbiosis of Deep Learning and Differential Equations III, 2023a.
Zhuang Liu, Zhiqiu Xu, Joseph Jin, Zhiqiang Shen, and Trevor Darrell. Dropout reduces underfitting. org/ aww.aas///00 2303. 01500.
2023b. URL
Arvind Mahankali, 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.
Nick McGreivy and Ammar Hakim. Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations. Nature Machine Intelligence, (10):1256-1269, September 2024. ISSN 2522-5839. doi: 10. 101100//4255--044000---- URL
http://des doi. owt. 1038 / s42256 024 00897-5.
William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics, 11:531-545, 2023. doi: 10. 10.11622 10111622acc__000555
URL httpthttp////o ac wwwpenhooooy..ooo//00oe
William Merrill and Ashish Sabharwal. A logic for expressing log-precisiog transformers. Advances in Neural Information Processing Systems, 36, 2024.
Eric J. Michaud, Ziming Liu, and Max Tegmark. Precision machine learning. Entropy, 25(1):175, January 2023, ISSN 1099-4300. doi: 10.3390/c25010175. URL http://dda doi org/ 10. 3390 / 22511015.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability, 2023. URL https: 05217.
/2300
Eric Nguyen, Michael Poli, Matthew G. Durrant, Armin W. Thomas, Brian Kang, Jeremy Sul- livan, Madelena Y. Ng. Ashley Lewis, Aman Patel, Aaron Lou, Stefano Ermon, Stephen A. Baccus, Tina Bernaadez-Boussarr, Christopher Ré, Patrick D. Hsu, and Brian L. Hie. Se- quence modeling and design from molecular to genome scale with doi: 10. 10.1101/2224.02.27..555 58822344 URL https::/p/// htwp://wwwwwt...o www.bworw.....wwwwwwwwwwwwwwwwww org/ wontent/eerntttttee/eeeeee /024/02 00/4//02222/2222 02 27. 27.582234.
bioRxiv, 2024.
Steven A Orszag. Comparison of pseudospectral and spectral approximation. Studies in Applied Mathematics, 51(3):233-25, 1972.
Matteo Pagliardini, Pierre Ablin, and David Grangier. The ademamix optimizer: Better, faster, older, 2024. URL httpt:ttt/aaakee wwtp:iiv.rrgaaooo..o..e abs /22/224000 2409. 03137.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al. Rwkv: Reinventing mnn for the transformer era. arXiv preprint arXiv:2305. 30488, 2023.
Peter Bürgisser and Michael Clausen and M. Amin Shokrollah. Algebraic Complexity Theory. Springer, 1997.
Tobias Von Petersdorff. Polynomial approximation and interpolation. 2015. Numerical Analysis Class Notes. https: wwww math. umd. wdw/pptterdd////eeeeeeeeeee /6we///////6////// amsc6 6 pdf..
W. Plešniak. Multivariate jackson inequality. Journal of Computational and Applied Math-
ematics. 233(3):815-88) 820, 2009. ISSN 0377-0427. doi: https//tbo.orr//1.0106j... 2009.02.095. URL https: / / www. scwencedirectt.t com/ hta///cceenne/aticceepppl.l
$03770427090013309 9th OPSFA Conference.
13

---

Published as a conference paper at ICLR 2025
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International Conference Machine Learning, pp. 28043- PMLR, 2023.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: General- ization beyond overfitting on small algorithmic datasets, 2022. URL https://axxii.ccoom abs/2201. 02177.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Pratik Rathore, Weimu Lei, Zachary Frangella, Lu Lu, and Madeleine Udell. Challenges in training pinns: loss landscape perspective, 2024. URL htwps: http///aaavvv/ee//ee/eeeeeeaael/ wttp::iwwwrr/a//////////aaa/ /2002. 01868.
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. Advances in Neural Information Processing Systems, 36, 2024.
Gleb Rodionov and Liudmila Prokhorenkova. Neural algorithmic reasoning without intermediate supervision, 2023. URL htt.t:htt//aattte
/206 13411.
Gleb Rodionov and Liudmila Prokhorenkova. Discrete neural algorithmic reasoning, 2024. URL pettss:/arxxer org/ wtwp//////////oooo//too///a..tt/ wwwp////........o
Günther Schulz. Iterative berechung der reziproken matrix. ZAMM-Journal of Applied Mathematics and Mechanics/Zeitschrif für Angewandte Mathematik und Mechanik, 13(1):57-59, 1933.
Gilbert Strang. Linear algebra and its applications. 2012.
Lloyd N Trefethen. Spectral methods in MATLAB. SIAM, 2000.
Lloyd N. Trefethen and David Bau. Numerical Linear Algebra, Twenty-fifth Anniversary Edition. Society for Industrial and Applied Mathematics, Philadelphia, PA, 2022. doi: 10.113771.9781611177665.. URL https: org. wttg/////as//...1111// 9781611977165.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Petar Veličković, Adrià Puigdomènech Badia, David Budden, Razvan Pascanu, Andrea Banino, Misha Dashevskiy, Raia Hadsell, and Charles Blundell. The clrs algorithmic reasoning benchmark, 2022. URL https: / arxiv. org/ abs /2222225.15555
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference Machine Learning, pp. 35151-35174. PMLR, 2023.
Yongji Wang and Ching Yao Lai. Multi-stage neural networks: Function approximator of machine precision, 2023. URL https: http:pp:////axxx.oo///////////// 230.. 08934.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities large language models. arXiv preprint arXiv:2206.0766 2022.
Sanford Weisberg, Applied linear regression, volume 528. John Wiley Sons, 2005.
D.C. Wilcox. Turbulence Modeling for CFD. Number v. in Turbulence Modeling for CFD. DCW Industries, 2006. ISBN 9781928729082. htttst/p::///// //books google e com/books ? dd LENNPGAACARJ.
Steve Yadlowsky, Lyric Doshi, and Nilesh Tripuraneni. Pretraining data mixtures enable narrow model selection capabilities in transformer models. arXiv preprint arXiv:231 1.00871, 2023.
14

---

Published as a conference paper at ICLR 2025
Liu Yang, Siting Liu, Tingwei Meng, and Stanley J. Osher. In-context operator learning with data prompts for differential equation problems. Proceedings of the National Academy of Sciences, 120(39), September 2023a. ISSN 1091-6490. doi: 10.10773pmas.2311422220. URL http:
//dx. doi. wt./................. 1073 /pwa./nnsss...... 231 01 4142120.
Liu Yang. Siting Liu, Tingwei Meng, and Stanley Osher. In-context operator learning for differential equation problems. arXiv preprint vrXiv:2304.0799/ 2023b.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequenee functions?, 2020a.
Chulhee Yun, Win Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. O (n) connections are expressive enough: Universal approximability of sparse transformers. Advances in Neural Information Processing Systems, 33:13783- 2020b.
Michael Zhang, Khaled K. Saab, Michael Poli, Tri Dao, Karan Goel, and Christopher Ré. Effectively modeling time series with simple discrete state spaces, 2023a. URL
abs/2303. 09489.
Ruiqi Zhang, Spencer Frei, and Peter Bartletr. Trained transformers learn linear models in-context. arXiv preprint Xix:22306.09927, 2023b.
15

---

Published as a conference paper at ICLR 2025
APPENDIX
The appendix is organized as follows:
Appendix provides a more detailed overview of related work.
Appendix provides details about experimental setup.
Appendix C provides additional experiments and ablation studies.
Appendix D provides details about our main theoretical results.
A EXTENDED BACKGROUND
A. LEAST SQUARES
Least squares, Ax b, is well-understood theoretically, and we know of simple numerical algorithms for solving least squares to high precision (Weisberg, 2005; Boyd & Vandenherghe, 2004). We focus on two algorithms: gradient descent and Newton's method.
Gradient descent Given guess for x", minimize the least squares loss
1
((x)
Σ (af x by)²
(7)
2
via gradient descent on æ:
N
VoLN = x T
Σ
a; bi)a
(8)
1=1
0t+1
8r (Lx()
(9)
Ordinary Least Squares and Newton's method In the noiseless, full determined regime, the Bayes-optimal estimator is ordinary least squares (OLS) (Weisberg, 2005):
x OLS
(A A) A
(10)
where
+- a1 +
A =
b
(11)
ON
by
Note that this estimator requires a matrix inverse, which expensive compute exactly. An alternative is to use Newton's method to approximate the matrix inverse term (Schulz, 1933). To estimate (A7 A)-1, we can perform the following iterative algorithm:
M M2 A)
(12)
where M, converges to (ATA)
A.2
RELATED WORK
In this section, we detail prior work on in-context learning, Transformer expressivity, gated convolu- tional architectures, and algorithm learning.
In-context learning. The capability of Transformers to perform in-context learning on language and pattern matching tasks has been well-documented (Brown et al., 2020; Dasgupta et al., 2022: Wei et al., 2022). More recently, flurry of work has investigated in-context learning for regression-style tasks. Garg et al. (2022) first formulated the mathematical framework to analyze the estimators Transformers implement in- in-context focusing linear regression and other least squares problems. A number of works further observed empirically that Transformers seem to approximate Bayes-optimal estimators distributional problems. For example, based the task distribution, the performance of
16

---

Published as a conference paper at ICLR 2025
in-context Transformers mimics optimally-tunet LASSO on sparse linear regression, ridge regression on noisy dense linear regression, and Bayes-optimal priors for task mixtures (Akyürek et al., 2024; Raventós et al., 2024; Yadlowsky et al., 2023; Ahuja et al., 2023; Bai et al., 2024). Beyond standard least squares problems, other works have investigated the ability of Transformers to in-context solve broader problems of scientific interest like differential equations (Yang al., 2023b; Chen et al., 2024; Liu et al., 2023a).
Towards explaining these observations, recent works have focused on understanding the expressivity and optimziation landscapes of Transformer variants (typically non-causal linear attention) on linear regression. Linear attention has been shown be expressive enough implement numerical algorithms for solving linear regression, including gradient descent (Akyürek et al., 2022; Von Oswald et al., 2023) and Newton's 's method (Fu et al., 2023; Giannou et al., 2024). Recent work (Ahn et al., 2024; Mahankali et al., 2023; Zhang et al., 2023b) has also begun to investigate the optimization dynamies for linear attention on least squares. Finally, we highlight that recent work (Bai al., 2024; Huang al., 2023; Collins et al., 2024; Cheng et al., 2024) makes progress on theoretically understanding non-linear attention, e.g. with softmax or ReLU activations.
Unlike prior work, we investigate the capabilities of standard Transformers, focusing exploring their capability to perform high-precision optimization algorithms. Noting a gap between empirical performance and theoretical claims regarding in-context least squares as gradient descent, we further investigate alternative architectures to softmax attention.
Expressivity and approximation ability of Transformers. Although Transformers were initially designed for discrete tasks like language modeling, recent works have investigated the ability of the Transformer architecture express general continuous- sequence-to-sequence maps. We briefly mention three classes of prior work:
Constructive arguments. We highlight Giannou et al. (2023), which proposes looped- Transformer weight construction that implements a basic mathematical instruction set. Using compositions of these instructions, the authors demonstrate that Transformers are expressive enough to implement numerical algorithms, including matrix inversion and SGD on linear models.
Universal approximation results. Several works, such as Yun et al. (2020a;b), provide bounds on the number of parameters and layers required to approximate smooth sequence- to-sequence functions to arbitrary precision using Transformers. However, these results typically require parameters to scale exponentially with respect to problem size, which quickly becomes impractical in practice.
Complexity theory results. Recent works (Chiang et al.. 2023; Merrill & Sabharwal, 2023; 2024) prove that log-precision Transformers lie in TC", limited complexity class of circuits.
Gated convolutions. Gated convolutional models are a class of architectures that serve as an efficient alternative to attention. These models, consisting of gating (element-wise multiplication) and long convolutions (filter size equal to sequence length), stem from earlier work (Gu al., 2021) inspired by the signal processing literature. In this work we focus on the BASECONV model from Arora et al. (2023), but a recent surge of interest in efficient attention replacements has led to a flood of gated convolutional architectures (Poli et al., 2023; Peng et al., 2023; Gu & Dao, 2023).
Recent architectural innovations within the class of gated convolutional models have been largely motivated by language modeling tasks (Fu al., 2022; Arora et al., 2023). Unlike these prior works, which focus on matching attention's 's performance on discrete tasks, we observe that the connection between gated convolutions and arithmetic circuits implies they able exactly express range of important numerical algorithms for cittinuous-valleed tasks. We further investigate their ability to learn these algorithms in-context.
Algorithm learning.
We mention two lines of work related to learning algorithms using ML:
Grokking. Several works (Power et al., 2022; Nanda al., 2023; Lee et al., 2024) have observed the ability of Transformers to learn to perfectly perform small discrete algorithmic tasks, e.g. modular arithmetic.
17

---

Published as a conference paper at ICLR 2025
Neural Algorithmic Reasoning. Recent work (Rodionov & Prokhorenkova, 2023; 2024) investigates the ability of graph neural networks to learn fundamental algorithms like breadth-first search (Veličković et al., 2022).
Crucially, we note that these previous works focus on learning discrete algorithmic tasks, which Transformers excel at. far as we know, we are the first to investigate whether Transformers are able to learn numerical algorithms, which rely addressing key challenges with high-precision floating-por poating-pointtooo arithmetic.
Precision and scientific MI.. The importance and difficulty of high-precision ML for scientific settings well-established: although the scientific MI. 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 (McGreivy & Hakim, 2024). Despite this, we are aware of only few works which directly focus on investigating high precision for ML. We highlight (Michaud et al., 2023; Wang & Lai, 2023), which focus small MLPs for regression tasks and propose alternate training recipes.
As far as we are aware, we are the first to investigate and isolate effects of model architectures and optimizers on precision a controlled setting: in-context least squares. We find that typical training recipes for sequence models (e.g. softmax attention, Adam, and standard LR schedulers) encounter surprising precision barriers when applied to numerical tasks.
18

---

Published as a conference paper at ICLR 2025
B EXPERIMENTAL SETUP
Here, we provide additional details about our experimental setup.
B.1 MODEL ARCHITECTURE
We base our Transformer and BASECONV models off the GPT2 family (Radford et al., 2019). Unless
otherwise specified, the following default settings for Transformers:
Config
Setting
Embedding size
64
Number of layers
12
Number of heads
MLPs
True
MLP hidden size
4x embedding size
MLP activation
ReLU
LayerNorms
True
Input dim
Sequence length
20
Table 1: Standard Transformer architecture details.
and the following settings for BASECONVS:
Config
Setting
Embedding size
64
Number of layers
MLPs
False
LayerNorms
False
Input dim
Sequence length
20
Table 2: BASECONv architecture details.
Finally, we describe the settings use for our linear attention experiment (Figure 2):
Config
Setting
Embedding size
256
Number of layers
Number heads
16
MLPs
False
LayerNorms
False
Input dim
Sequence length
20
Table 3: Linear attention architecture details.
B.2 OPTIMIZER
We describe two sets optimizer settings we use throughout this work.
The first, representative of standard training procedures, is inspired by prior in-context learning setups (Garg et al., 2022; Von Oswald et al., 2023).
The second, training recipe, is for our high-precision experiments, where find more aggressive learning rate scheduler is essential. Note we use the adaptive learning rate scheduler and EMA described in Section 5.
19

---

Published as a conference paper at ICLR 2025
Config
Setting
Batch size
256
Optimizer
Adam
Learning rate
10-3
Scheduler
StepLR
Training iterations
10t
Step rate
10*
Decay rate
0.9
Table 4: Standard optimizer settings.
Config
Setting
Batch size
1024
Optimizer Learning rate
Adam
10-2
Scheduler
AdaptiveLR
Training iterations
2.5 x 106 3 10
Step rate
Decay rate
0.9
EMA decay
0.98
EMA lambda
2
Table 5: High-precision training recipe settings for BASECONV.
Finally, we describe the optimization settings we used for high-precision linear attention, which we found needed slightly different learning rate scheduler.
Config
Setting
Batch size Optimizer Learning rate
1024
Adam
10-2
Scheduler
AdaptiveI.R
Training iterations
2.5 x 106 3 10
Step rate
Decay rate
0.9
EMA decay
0.98 2
EMA lambda
Table 6: High-precision training recipe settings for linear attention.
20

---

Published as a conference paper at ICLR 2025
B.3 TASKS
Each of our in-context learning tasks can be viewed as a sequeence-o--seqeence map
M M + Nout x Dout
In this subsection, we provide details about task implementations. specifying the input/output formats for each of the synthetic tasks and in-context least squares variants implement.
B.3.1 IN-CONTEXT LEAST SQUARES.
We consider MLS [Nx((++) + RP, where above the inputs are formatted as
[a aN by
Win
bi
and the expected output is
(o(um)) -1:]
B.3.2 PRIMITIVES.
For each of the following linear algebra primitives, increase the task size, setting D 20 and N 40.
READ is defined as MRead NRD + R DD
where the inputs are formatted as
Uin E NxD [
ØN)
and the expected outputs To(u E RNN such that
[k, = uum i1,
k =
Uin [k
k +
for task parameters i JEE [N]. LINEAR is defined as M Minearr NND + Nx
where the inputs are formatted as
Win RR NxD = 1
ΦN]
and the expected outputs are
To(uin) = |07
xi,h
where h R is task parameter.
MULTIPLY is defined as Maultiply as
NND R NxD2 where the inputs are formatted
Uin RR NxD_
ΦN]
and the expected outputs
To(u.)) = (ail::: o D/2:
E.l::D/21 ZN D/2 :I).
B.3.3 EXPLICIT GRADIENT UPDATES.
In Section 5, we investigate simple training setting, in which the model is explicitly trained predict the gradient of the least squares loss. We proceed to define the task M Mgradient. ((N=1)x(D*+2D+1) + RDD.
The inputs formatted as
[ai bj
aN
Co
Uin
bn
The expected outputs are
To 1-1:,
:1::;;: :D] V v.noo)))
21

---

Published as a conference paper at ICLR 2025
B.3.4 k-TH GRADIENT DESCENT ITERATE.
Finally, toward end-to-end least squares, we investigate series of increasingly end-to-end tasks in which the model is explicidy trained to predict the k-th gradient descent iterate. We proceed to define the task M iter (DD*+)))( - RDP
The inputs are formatted as
ai
2o aN
Un
b1
bn
0
The expected outputs are
To (00n)]1:::. :D] = a.-
B.4 DATA GENERATION
At each training step, we produce random training prompt e4.. by sampling each variable randomly: from the isotropic Gaussian distribution N(0, I) for continuous-valued parameters, and from the uniform distribution for discrete parameters. Concretely:
For the in-context linear regression tasks, input vectors .i...., aN sampled from NOD 1)) and the unknown linear function determined by w", also drawn from N(OD I.).
For the synthetic tasks READ, LINEAR, MULTIPLY (Section 3.3), each column of the inputs Unn E Nx is sampled from the isotropic Gaussian distribution N NOOP, P). The tasks READ and LINEAR require specifying additional parameters as follows:
For READ, at each iteration, i # E [N] are sampled uniformly.
For LINEAR, at each iteration, the affine transformation h is sampled from ((OP,II")
For the explicit gradient task and the k-th gradient descent iterate task, the random initializa- tion wh is also drawn from N(OP, ID).
The model is trained to minimize mean squared error over the distribution of prompts.
22

---

Published as a conference paper at ICLR 2025
C ADDITIONAL EXPERIMENTAL RESULTS
C.I ABLATIONS: LINEAR ALGEBRA PRIMITIVES
In Figure 6, we train Transformers and BASECONVS, with MLPs, with and without LayerNorms (LN), on the READ, LINEAR, and MULTIPLY primitives from Section B.3.2. We vary the model depth E {1, 2, 4, 8) and investigate how precision scales with number of layers. In these experiments, we use standard exponentially decaying LR schedule for Adam.
We show that Transformers and BASECONVS both achieve high precision (< 0(10-9)) on the READ and LINEAR tasks. However, the Transformers struggle to implement MULTIPLYY to high precision, and performance scales poorly with model depth. We observe that BASECONY without LayerNorm generally performs the best across all three primitives, consistently outperforming BASECONY v with LayerNorm by 2-4 orders of magnitude.
Read
Linear
Element-wise multiply
DaneCeno, OeeeCoTA
DanaCory
Saeceen UM
Jer, Un
ArT LM
kmn
22
10
101
12.
Number of layers
Number of layers
Number of layers
Figure 6: Attention vs. BASECONV, with and without LayerNorms, on synthetic tasks. Precision consistently scales better with depth for BASECONv models than for Transformers. While both models solve READ and LINEAR tasks to at least 10 MSE, the precision of Transformers scales poorly for the MULTIPLY task.
Focusing on 2-layer Transformers and the MULTIPLY task, we additionally find that precision scales poorly with multiple scaling axes, including hidden dimension, number of heads, and MI.P upscaling factor (Figure 7).
Multiply
10-4
MSE
Attention dimension
Attention heads
MLP
Figure 7: Precision of (2-layer) Transformers on MULTIPLY task scales poorly with attention dimension (left), number of heads (middle), and MLP width (right, where MLP hidden dimension width x attention dimension).
23

---

Published as a conference paper at ICLR 2025
Element-wist multiply
BaseConv
Attn
10-2
10-3
MSC
10
10-6
10a
10f
10
10a
Training iters
Figure 8: Scaling number of training iterations for 1-layer Transformer vs. BASECONV on the MULTIPLY task. Both models improve precision by 2-3 orders of magnitude as training duration increases by 3 orders of magnitude.
Finally, we investigate the effect of training duration on precision. In Figure 8, we train 1-layer Transformers and BASECONVs, with MLPs and LayerNorms, the MULTIPLY primitive and vary the number of iterations for which the model is trained. Recall that since new data is sampled at each iteration, also effectively scale the dataset size proportionally. To keep the learning rates consistent across runs, we scale back the scheduler step size accordingly:
num_iters E {10", 10 10?, 10*}
step_size E (102, 10 10 10%}
We observe power law, particularly clearly for BASECONV, as we scale from 10 to 10 iterations. Both models achieve 2-3 order of magnitude improvement in precision, but this requires also increasing training duration by orders of magnitude.
24

---

Published as a conference paper at ICLR 2025
C.2 ABLATIONS: HIGH-PRECISION OPTIMIZATION
In Figure 9, we try directly training on the end-to-end least squares task, simply replacing softmax attention with BASECONV in the standard Transformer architecture. We find are unable reach high precision using this training procedure.
squares, Attn/BBaseonv vs.
10a
Transformer (end-to-end]
BaseConv (end-ta-ende
(baseline)
10
10
10
10
10-10
10-12 -2
10-14
10
30
50
Iterations Depth
Figure 9: Replacing softmax attention with BASECONV in the standard Transformer architecture and training end-to-end least squares is not enough achieve high-precision solutions. BASEC- ONV models trained end-to-end perform as badly as Transformers at small scale, and our largest models perform 100x worse than parameter-matched Transformers.
In Figure 10, we ablate the effects of constant and exponentially decaying LR schedulers with Adam (cutting off training after 10t iterations). We find that neither able efficiently train machine precision on the explicit gradients task. For exponentially decaying LR schedule, we find that the LR steprate crucial parameter: on the explicit gradient task, difference of 10, 000x between precision saturation thresholds using 1 x 103 vs 3 x 10 for example.
Decaying
Dacaying
UR
Intai
UR steprate
10:1
1x30->
10-1
3x10*
10
10
3x10
10
3x10'
Avorage b55
10-3
10
1=110*
1x10'
10-4
10a
10
10-
10-11
10
166
Tranning
Figure 10: Training with Adam on the explicit gradient task, we ablate LR for constant scheduler (left), initial LR (middle) and LR steprate (right) for decaying scheduler.
25

---

Published as a conference paper at ICLR 2025
In Figure 11, ablate the effect of applying EMA the update vectors from the Adam
optimizer. Empirically, we find that this boosts the final MSE by as much as 100, 000 x on the explicit gradient task.
Decaying (ablating EMA)
3x10-1
10-2
1x10
1x10-
10
10.5
Average
10
10
0.0
0.2
0.4
0.6
0.8
1.0
Training
Figure 11: Training the explicit gradient task. applying EMA over Adam's update vectors consistently boosts final MSE, by up to orders magnitude.
In Figure 12, we ablate the effect of restoring the MLPs and LayerNorms to BASECONV models. Surprisingly, we find that even these architectural components worsen the model's 's precision: on the explicit gradient task, by factor of up to 1, 000, 000 MSE. We note that due to training instability with the BASECONV+ML. model, we used a less aggressive LR schedule with initial LR 10-3 and LR steprate 10a
BaseConv (baseline)
10
+LLPs+LNS
+LNs
0.8
+MLPs
10-4_
Cosine similarity
Average looss
0.6
10
10-1
BaseConv (baseline)
10-10
0.2
+MLPs+
+LNs
+MLPs
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Training iters
1e6
Training iters
1e6
Figure 12: Training on the explicit gradient task, adding MLPs and LNs consistently bottlenecks precision: here, by up orders of magnitude.
26

---

Published as a conference paper at ICLR 2025
In Figure 13, evaluate 3-layer BASECONV and linear attention models trained on the explicit gradient task. As in Section 5.2, we apply them iteratively until convergence, We then evaluate on out-of-distributiin regression targets, in Section 3.2.
We surprisingly find that linear attention demonstrates poor numerical generality, despite training to 4, the iterations of linear attention near machine precision the training distribution. Beyond
diverge.
This result suggests that although different polynomial architectures may equally be able to express algorithms, they may learn solutions that exhibit vastly different numerical properties.
Out-of- Distribbtiiin Target Vectors
10
10a
10
Transformer
10
(12 iters)
NSE
GD (until convergence)
10
(GD iter)
Linear Attn (GD star]
10
10-10
10-1
Figure 13: While BASECONV demonstrates improved numerical generality compared to end-to-end trained Transformers, the generalization gap for linear attention is as bad as the Transformer.
27

---

Published as a conference paper at ICLR 2025
C.3 K-TH ITERATE GD
In this section, we investigate how well our proposed techniques can learn the k-th GD iterate tasks defined in Section 5:
{(a1, ...,.... (aN.bN).
æg, where 2;+1 =; -VV[V i [k 1].
(13)
Recall that k = 1 is equivalent to the explicit gradient task, while taking k - 20 is equivalent to the standard in-context least squares task. Here, we are interested in understanding how well our techniques extend to larger k. towards learning end-to-end least squares. See Appendix B for more detailed description of the training setup.
Our theoretical results in Section 4 imply that k + 2-layer BASECONY is expressive enough to solve the k-th iterate task to machine precision. Thus we train lt 2-layer BASECONv models on the k-th iterate task for k 2 1 using our training recipe.
k MSE
2
3
5.0 x 10 13
2.5 x 10
2.5 x 10"
3.1 x 10
Table 7: We can learn up to 4 iterations of GD at once with our current training techniques. Model stability becomes bottleneck with harder tasks.
Training on the k-iter GD task, we find our training recipe scales to before training instability occurs. Adding LayerNorms, least 1, x00 for small k: see Figure 14.
able to train deeper models, but we find MSE worsens by at
k-iter GD
10-4
Without LNN
With LNs
10-5
10.5
10-7
Average
10-8
10-9
10
11
10-12
Number GD iters (K)
Figure 14: BASECONV with LayerNorms are able to stably scale to deeper models, but LayerNorms present precision bottleneck: even on small k, MSE degrades by over 1, x00x.
28

---

Published as a conference paper at ICLR 2025
C.4 IN-CONTEXT ODE SOLVING
In this section, we demonstrate the generality of our insights on the more practical setting of in-context ODE solving. We note that solving differential equations in-context with Transformers is framework that has been explored in recent papers (Yang et al., 2023a; Herde et al., 2024; Liu et al., 2023a), and thus represents natural first step towards extending our techniques realistic scientific ML problems.
Experimental setup.
follow setup from Liu al. (2023a):
We train distribution of ID ODEs over E [-1, 1], defined by
u' u' ajd(t) + aqu(t) + ag.
(14)
For each operator, we provide 25 in-context examples of foreing functions, initial conditions, and their corresponding solution values at a fixed time Equery [-1,1]. We then give the model query forcing function and initial condition, and the goal is to predict the
corresponding solution Lguery-
We sample our parameters 01 On.5 1.5)). 02 Unif([-1,1]), ag Unif([-1,
Initial conditions sampled from u(0) Unif([-1,
Forcing functions c(t) are sampled from Gaussian process with RBF kernel K(xxx')
(2-2²) with length-scale parameter & = 1. We sample each forcing function on 28
exp
equispaced points over [-1,
ODEs are solved pseudospectrally on N nodes: we find this is sufficient for machine- precision solutions with FLOAT32 datatype.
We find that our observations from least squares transfer to the setting of in-context ODBs:
Transformers struggle learn precise solutions. We find that 12-layer, 9M parameter Trans- former model only achieves ~ 10-4 MSE, almost 10 worse than the threshold FL.ATT32 machine epsilon implies. Furthermore, as with least squares, we observe precision saturation with model size. In Figure 15, we find that scaling the depth of the model by up to 2x does not improve precision. We further note that precision saturation already seems to occur with 4-layer Transformers. We hypothesize that the depth at which precision saturation begins is dependent the task difficulty.
In-context ODEs with Transformers
10-3
MSE
10
5.0
10.0 12.5
17.5 20.0 Number of layers
25.0
Figure 15: Transformers fail to learn precise algorithms for solving ODEs in-context. As with least squares, precision saturates with deeper models: in our experiments, we observe no significant performance boost between 4-layer and 24-layer Transformers.
29

---

Published as a conference paper at ICLR 2025
COP OUL Perameters
0O0 Forcing Funstions
OOD intill Conditionss
10
107%
1011
10 GP Longth Parameter
ODE Parometor Scale
Inttial Conditions Scale
Figure 16: Transformers fail to learn numerically general solutions: performance brittle to out- of-distribution ODE parameters (left), forcing function smoothness (middle), and initial condition distribution (right).
Transformers exhibit brittle generalization. We observe that Transformers robust to changes to the distributions of ODE parameters, forcing functions, and initial conditions. We describe our experimental setup below, mirroring Liu al. (2023a):
Out-of-distribution ODE parameters. We parameterize out-of of-diistritttion ODEs via scale parameter opp, where a Unif([1- Tp. 1 opl) and a22 dr ~ Unif([-os Topl). As we increase Cop. we sample from a wider distribution of ODE solution operators, including those with larger operator norms and worse-conditioned design matrices.
aut-of-distribution forcing functions. We vary the length parameter of the Gaussian process from which we sample our forcing functions, which effectively controls their smoothness.
Out-of-distribution initial conditions. We sample out-of-distribution initial conditions as u(0) ~ Unif([-ov -ac, (rc)) As we vary arc, we widen the distribution of the solution values at 0, which increases the overall magnitudes of the solutions.
We note that in all out-of- experiments, the Transformer's MSE explodes to near O(1): refer to Figure 16.
Our proposed techniques obtain precise and general solutions. Liu al. (2023a) shows that in-context ODEs can reduced to solving least squares problems. Thus, we train a 3-layer BASEC- ONY architecture the explicit gradient task for the equivalent least squares problem, and apply our model iteratively, as in Section 5. We compare the performance of our iterative model with end-to-end Transformers, least squares solvers, and standard gradient descent applied to the equivalent least squares problem.
We note that our ODEs reduce to least squares problems that are ill-conditioned. In this set of experiments, we find the condition numbers of our design matrices are O(10*). Since the theo- retical convergence rate of gradient descent on least squares depends inversely on the condition number (Boyd & Vandenberghe, 2004), we expect our iterative models and standard gradient descent will require orders of magnitude more iterations than in the least squares problems of Section As such, we limit the number of iterations for our BASECONV model and standard gradient descent to 10, 000. Nonetheless, we find that our BASECONV model learns to high enough precision that we are able to maintain the stability of the iterative algorithm for up to O(10a) steps. In our experiments, iteratively apply our BASECONV model until convergence to a fixed point and report final MSEs.
We find that BASECONY learns precise and general algorithm for in-context ODEs:
Precision. In Figure 17, we show that our BASECONY model, applied iteratively, converges about 10 10 MSE, 000, 000x higher precision than our best Transformers.
Generality, In Figure 18, we find that our BASECONV model exhibits more robust gen- eralization than the Transformer model: in all the dfitt--disttiuution settings we test, our BASECONv model achieves higher precision than Transformers in-distribution. Like above,
30

---

Published as a conference paper at ICLR 2025
evaluate distribbtitiitto ODE parameters, forcing functions, and initial condi- tions. particular, we note that the performance of our BASECONV model almost exactly matches proper gradient descent, even in -ittr--dittriuttion settings. Additionally, the generalization behavior of our BASECONV model matches the generalization of a proper least squares solver with preconditioning, except for out-of-distribution initial conditions, where we note that the iterative procedure suffers from slow convergence and times out at 10, 000 iterations.
find
We believe these preliminary results show the promise of our techniques towards learning numerical algorithms for more complex tasks, such as solving PDEs, directly from data.
Error Comparison
Least squares
10-4
BaseConv
GD
Transformer
10-6
MSE
10
10-
10 10-12
Least squares
BaseConv
Transformer
Figure 17: In-distribution error comparison between Transformer, BASECONY, gradient descent, and least squares.
COP OUT Parameters
0OD Foroing Functions
COD inttlal Condittons
16a
32-9
1175
0
32-1
1071
38
00'11
JCP
ODE Parometor Scale
GF Length Parameter
Inttial Conditions Scale
Figure 18: Out-of-distribution error comparison between Transformer (orange), BASECONY (blue), gradient descent (gray), and least squares (green): we evaluate out-of-distribution ODE parameters (left), forcing function smoothness (middle), and initial condition distribution (right). BASEC- ONV learns numerically general algorithm that closely matches proper gradient descent and least squares.
31

---

Published as a conference paper at ICLR 2025
D THEORETICAL RESULTS
This section is organized follows:
We detail notation and definitions in Appendix D.1.
In Appendix D.2, include theoretical results regarding the primitives from Section 3.3: expressivity results with BASECONV and attention, and iterative algorithms as compositions of primitíves.
Appendix D.3, we discuss upper and lower bounds for implementing gradient descent on least squares using BASECONV, supplementing Section 4.1.
In Appendix D.4, we provide theoretical details regarding the universal function approxima- tion properties of BASECONV.
Finally, in Appendix D.5, we provide technical details about the claims from Section 4.1 that BASECONV can perfectly recover SQUARE and LINEAR.
D.1 NOTATION
We heavily borrow notation from Appendix H of Arora et al. (2023), which we recollect below. We denote the all row vector of size k, given by [1 1
1 1]. and the all 0 row vector of size 0), as 1 and 0*, respectively. We also construe the standard basis vector as column vector in this appendix, and adhere to the following matrix indexing convention: M[i,j] is the entry in the row and the jth column, M[i,:] Fix= denotes the ith row, and [[:j]] E Fuax1 denotes the jth column of M E |mxx, where F is a field (the reader can assume that IF the field real numbers i.e. F R). We then use mxxn 0m E |mmxn denote matrix of all is and Os, respectively. We note that some notation differs from those used in earlier sections.
given by [0 0 ...
Next, we denote the Hadamard product of vectors to matrices by applying the Hadamard product column-wise across the matrices. This is commonly referred to as (element-wise) gating. For vectors u, E F", also denote their linear (or acyclic) convolution as u v and cyclic convolution as u @ V.
F" u the operation can be extended
Polynomial Notation. Since convolution is equivalent to operations on polynomials, it is convenient to use them discuss the inputs and outputs of gated convolution models. Let us define maps poly: F" - F[X]/(X") such that
poly(u) Σ x
This allows us to map between vectors and polynomial. Accordingly, we also define coeff F[X]/(X*+) F the map converting polynomials back vectors: coeff cocff(u(()) (u(X)) with
ulli] defined as the coefficient in uXX) at degree i.
These operations allow to interpret the convolution of vectors terms of polynomial multiplica- tion (Heideman& & Burrus, 1988). More specifically, we have
coeff (u(X) x(X) mod X")
The following notation for a polynomial will be used in this section: Definition D.I. A polynomial P(X) with degree d and some coefficients c E [Rd+1 1 is defined
P(X) = Σ e.
==0
Further, the degree of P(X) will be denoted deg( P).
32

---

Published as a conference paper at ICLR 2025
Function Approximation. In this part, collect notation and known results about function approximation. We will reference some definitions from Plešniak (2009); Petersdorff (2015).
The following notation is to denote the kth derivative of function:
Definition D.2. For some function f R + I. f(h) == Z f(x) is the kth derivative f.
Define set of univariate functions with notion of continuity:
..... Definition D.3. We denote C* (a, b) for k = 1, the space of univariate functions f [a, bj + R,
which have derivatives f(.),... f(k) that are continuous on the closed interval [a, 6].
Next define of multivariate functions with a notion of continuity: Definition D.4. A function f [a, b)
TR is in C* (a. b)th for h = 1,2... if all partial derivatives
Oa
drn f() 42 ................................ ,......9.) vn)
0r #1 Ou 22
exist and are continuous, for every a1, ,2,,..., an E Zzo. such that a1 + 02 + ..... + an S k and every (y1,---9a)) Un) E [a, b).
We use the following notation for the set of all univariate polynomials: Definition D.5. For any integer d 2 0, we define
Pa(X) {co e1 + caX E E}.
In other words, Pa(X) the space of univariate polynomials of degree less or equal to d.
We use the following notation for multivariate polynomials: Definition D.6. For any integers n, 2 0, we define
PT (X1. X,)
Σ
X
X
X
E RR Σ d
t=0
Then ( X₂) is the space of n-variate polynomials of degree less or equal to d.
The following notation is for considering the pointwise absolute value of matrix: Definition D.7. For M E NxD define,
M
OCTNN 03
BEaK
[M
Now lets define the corresponding oo norm for functions:
Definition D.8. For g [-1, N NxD define
|g x 1E--11 NXD gg(x))
We will use the following version of Jackson's theorem for univariate inputs:
Theorem D.9 (D. Jackson (1930) Jackson's Theorem for C [-1, 1-1,1},) 1].). Let d, k integers with d + 1 k 2 0 and f E C* 1}. Then
1 (d 1)d .. (d + 2)
fll Plas. S PEP,I
inf
(
(15)
We will use the following version of Jackson's 's theorem for multivariate inputs: Theorem D.10 (Plešniak (2009) Jackson's Theorem for [*[-1,1]" Let d, k be integers with d + 1 2 k 2 0 and f E 1"" Then
gb+1
inf PEPS |f
Pll.o 5
Ck
Σ
f(x)
(16)
O
jml
where Cg. is positive constant.
33

---

Published as a conference paper at ICLR 2025
We will use the following definition of univariate smooth functions: Definition D.11. We call & times differentiable function f 11 + R to be (k, L)-smooth if [| S.
Next, we observe that given univariate smooth function, there's a univariate bounded degree polynomial that approximates it some error, E:
Corollary D.12. For some (k, -moothhh univariate functtin f (as in Definition D.11), then there exists a polynomial (z(x) with
(11 (12
deg 0 00
+k
such that for all [-1,1]
- ()))) c.
Proof. We will be bit more specific on upper bound of deg((s)) We pick:
deg( Pf) =
+ k
(17)
Let 5 where Py the polynomial that achieves the left hand side Equation (15). Then we have error at most
k
(())
(d l)d (d + 2)
Using the definition a (As, -mooth univariate function in Definition D.1I we get the error at most
L
k L
(d 1)d (d + 2)
2 (d k-)a
where the inequality follows since each d 1, 1,d, ......,ee k-k + 2 > (d k). Plugging in Equation (17) for d we get the error is at most:
L
2
desired.
We will use the following definition of multivariate smooth functions that map to single value:
Definition D.13. We call a times differentiable f [_1,1]" - to be (ks, Lm-shh L for S m I n.
f(0 ||=
Now we show the corresponding observation for multivariate functions and polynomials: Corollary D.14. Let deg( Pÿ) d. For some (k, I.) -mooohh multivariate function f (as in Defini- tion D313), then there exists polynomial Py(x) with
dog(P;)) Ok
such that for all x [-1, 1)"
- Se.
Proof. Let Py be the polynomial we get from the left hand side of Equation (16). We want to upper bound the error
g
C₁ t
fxx
Se.
dz k+T
0
34

---

Published as a conference paper at ICLR 2025
which follows if
Ck
==1
since f is (k, s.ssmohth. The above is the same
C dc
S
or equivalently
S d.
Picking d
V
eanL
suffices.
Arithmetic Circuit Notation. We briefly recall arithmetic circuits Peter Bürgisser and Michael Clausen and M. Amin Shokrollah (1997). An arithmetic circuit C with variables X & {11, TTn,.,,.,,TA over a field is interpreted as a directed acyclic graph, where the input nodes lahelled by either the variables from X constants from F and the internal nodes are labelled by + or x with the output being the polynomial computed at the output node.
We shall also refer to the size' of the circuit C the number of wires (or edges in C),, the depth the circuit as the length of the longest path between an input node and the output node, and the width of the cireuit as the number of wires that will be intersected by a horizontal 'cut' through the eircuit. Moreover, the degree of a circuit is defined as the degree of the polynomial computed by the circuit. We summarize this with the following definition:
Definition D.15. An arithmetic circuit C is (n, s, A, wircrcuit if C is n-variate arithmetic circuit of size s, depth at most A, and width w.
BASECONY Architecture. In the following definitions we formally define the BASECONY model Arora et al. (2023). To formally define BASECONV, we will need the Kaleidoscope hierarchy Dao al. (2020) as well.
To start, we define butterfly factors:
Definition D.16. A butterfly factor of size k 2 2 (denoted as Bg) is a matrix of the form B
D1 D2 where each D; a 5/2 x 5/z diagonal matrix. We restrict be a power of 2.
D3 Daj
The following definition is for butterfly factor matrix, which is made up of the above butterfly factors:
Definition D.17. A butterfly factor matrix of size n with block size k (denoted as B is block
diagonal matrix of z (possibly different) butterfly factors of size
(n)
diag [Ba],
B
B
Now lets define a butterfly matrix:
Definition D.18. A butterfly matrix of size (denoted B is matrix that can be expressed as
product of butterfly factor matrices: B (*) B
n B B Equivalently, we may define B (1)
(n)
recursively matrix that can be expressed in the following form:
B (n)
(Note that B and B
may be different.)
NNote that if all the gates of arithmetic circuit have bounded arity then the number of wires and gates asymptotically the but in this appendix will consider gates with unbounded arity.
35

---

Published as a conference paper at ICLR 2025
Using these butterfly matrices, lets define the Kaleidoscope Hierarchy:
Definition D.19 (The Kaleidoscope Hierarchy (Dao et al., 2020)).
Define B as the set of all matrices that can be expressed in the form B (") (for some n.)
Define (BB*) as the set of matrices M of the form M M,M for some M1. M2 E
Define (BB*) w as the set of matrices M that can be expressed as M = M M,M1,
with each M, (BB*) (1 S Kis w). (The notation represents width.)
Define (BB")² as the set of matrices M that can be expressed as M = SES for some x en matrix E E (BB*)" where S Fnxen = [L, upper-left corner of E). (The notation represents expansion relative to n.)
0] (i.e. M is the
Here now formally define BASECONV layer:
Definition D.20 (BASECONV (Arora et al., 2023)). Given input sequence RNxDD where N is the sequence length and D is the model dimension, learned weight matrix W E DxD and biases B1, Bg N and a matrix of convolution filters H E RXDD, BASECONY layer computes the following:
H BASECONV = (uW + O B2) NND
(18)
where the jth column of H u E RNDD is defined as H [[j]]] ulr.j].
The corresponding pseudocode for BASECONV layer is follows:
Algorithm 1 BASECONV(u, W B B2)
Require:
Input sequence u E R N linear map W E DxD, convolution filter H E NxD and bias matrices B1, B2 NND
1: In parallel for () < n < N afn.:
u n W
2: In parallel for 0 < { D 2[:() H
u [144
In parallel for f D yl: - (*/.. + B B. @ (2|,,) Bzli,)). See eq. (18)
4: return
Remark D.21. The definition of BASECONV layer in Equation (19) has the input go through linear layer before the convolution operation. For this section we will assume the linear layer the identity matrix, it is not needed for the results in this section.
Assumption D.22. Moving forward assume the weight matrix W E R DxD in Definition D.20 poly- log D Consequently, each matrix W has Ó(D) parameters also has the property W (BB* poly- and runtime for matrix vector multiplication Dao al. (2020).
In this section, we will establish some additional basic primitives that expect need implement via a BASECONV layer: shift and remember. We specify them below:
Definition D.23. shift (..r, f) Shift an sequential input of length N up down by entries:
INPUT: y E XD 2 O.
OUTPUT: T NxD where z shi ft dow_downff (y, 8) and =- shi (t__p())
36

---

Published as a conference paper at ICLR 2025
yo
ys
:
11-1
+ -
YN
y =
=+
2-
yi
yo
UN-1
yN
0
The following proposition is defining the
convolution Kernel that computes the
shift_down (s (. 2
N 1) primitive:
Proposition D.24. Define H R2NND as
D =
ifk
H
1
otherwise
For any u E 22NXD, u will result in
0 NxD
H *
u1
+
ug
u1
where u1 U2 RND
Proof. The convolution operation: H
111
where each column of H is convolved with each
U2
column of be restated polynomial multiplication. For column i, 0 S i 2N,
111
[.,
c(X ul::i((X) mod x²N).
12
= Note that the columns of H are all eN basis vectors and =
When multiply the term through the input polynomial we get,
coeff (x** N [u[o][i] + u[]][X
+ u|2N 1|11| mod X X2N)
N+1 .... + u [2N 1| mod x 22) cocff (c + u[1][4xx
With the lower order terms all becoming zeros, the above is same as
coeff ((0 0X +x
x u[]][i]xN+1
u/2N 1 mod X2N).
After we take the mod 2NN we get
coeff(( + OX +... +0X u +... u(N 111)
which implies that H
u1
Ug
NxD
u1
as desired.
We also define the following primitive:
37

---

Published as a conference paper at ICLR 2025
Definition D.25, remember(y,,r,t INPUT: y E R xd' , Z, 1 E 2, RR-
R on E R R2, R where y is defined as
below.
OUTPUT: IR''''' xd' which is defined follows:
11
t- 11 +
T +
t f((x) 7
0 0xxd
22
U2
0
We will need the following BASECONv implementation of remember:
Proposition D.26 (Arora et al. (2024), The Remembering Primitive). For any E Ruxdd Uz E t r contained in some y E RN' xd' such that is in the first Brxd' U2 mM where
rows, 2: is in the next n rows, Os fill up the next 8 rows, and V2 are in the next m r rows, for some 3n + 3m + 2s + 2t < N' so that for h Raxd and W E Idd xd' with x E [n++s) and v * E [m(++))xd" where v E mRd is defined as +2+ Ft_down (11, m r), there exists N", d BASECONV that computes remember(y, J), where be implemented I layer of BASECONv through the parameters W Rdd h RN N'xx xd'
88 (N",8,
RNN xa' by
defined below:
(uW (axdd
+ ( bi (sxdd ((BRce)
)) + u h
b2 Osxx' ))
f(u)
We will also need the following generalization of the above result:
Corollary D.27 (Arora et al. (2023)). Let be in Proposition D.26 but now let f be implemented with BASECONY(N, L, D. D). Then remember(y, r, r,t, f) where t can be implemented with BASECONV via (N, ((L), D. N, D) BASECONV.
The rest of Appendix will use this 5 tuple notation for BASECONV: Definition D.28. Lets define 5-tuple notation for BASECONY layer (, BASECONV with layers such that:
1. Input and output are N D matrices.
2. Each layer is defined by Definition D.20 where and D are replaced by N' and D'. Le. each layer takes in N' x D' matrices and output N' x D' matrices. We refer to the tuple (N', D') as the inner dimension of the model.
3. The matrices are projected from (N, D) - (N', D') (and vice-versa) via a linear projection.
We state the following bounds parameters and runtime for single BASECONV layer: Proposition D.29 (Arora et al. (2023)). An (N,1, D) BASECONv requires Ö(ND) param- eters and runtime.
We state the following result that says arithmetic circuit can be represented as BASECONV model: Theorem D.30 (Arora et al. (2023), Theorem H.21). For any (ND, A, w) arrithmeet circuit C, there exists an equivalent (N, D, D') BASECONY with A' = O(Alog w), N' (w)), D' D that simulates C.
38

---

Published as a conference paper at ICLR 2025
D.2 PRIMITIVES
In this section, we provide theoretical results about primitives.
In Appendix D.2.1, we implement the three primitives (READ, LINEAR, and MULTIPLY) from Section 3.3 using BASECONV, each using a single layer.
Next, in Appendix D.2.2 and D.2.3, we briefly sketch how the three primitives READ, LINEAR, and MULTIPLY can be used in composition to exactly express gradient descent and Newton's ss method iterations least squares (see Appendix A).
Finally, in Appendix D.2.4, we provide proof that a single layer of causal softmax attention cannot exactly represent the entry-wise squaring function. As corollary, since entry-wise square is a special case of MULTIPLY, this implies that attention cannot exactly express the MULTIPLY task for all arguments.
BASECONY parameterization
We recount the parameterization of BASECONv from Equation
y (u W gate b ggate) C (h * (u Win + bin) bomu)
W Wout + bout
(19)
Linear Projection
Convolution
Wont (W gate (u) @ Conw(W,,, (u)))
where Win W gates Wout are linear projections RD + D.
D.2.1 1-LAYER BASECONV CAN IMPLEMENT LINEAR ALGEBRA PRIMITIVES
Below, we formally define the linear algebra primitives discuss in Section 3.3, and we describe our BASECONY weight constructions.
Read
The READ operator, which maps inputs R to outputs y E Rxd is:
READ(i,j.a,b)(u) = [ulk, b] [uli k ==j
k ##1
(20)
Our implementation requires the use of the positional encodings and residual connections within the BASECONv architecture. Concretely, consider the input
C1
62
EN (N.:
Uin
[11,:]
u[2.
where the basis vector eμ represents the positional encoding for the k-th entry of the sequence. Define the output the BASECONV layer with residual connection:
y W out (W gate (u) @ Conv (u)))) u).
Then the following weight construction is equivalent READ(i, j, b):
Wgate( (kk :)) == u D
. Coon(WWnn [k, == ualk + - i-j.j u[(k:=]
Wont = proj
In particular, Wyate is defined such that
DD k -=j
W gate (ulk;:) =
100
k: j
Thus
uu/k lk + j.:] ulk: :] = ua(i,:] =
k = k + j
W Wgaze (u) o Conw(W..
0
Finally,
Wgate Wgate(u) O Conw(Wnn (u)) + == [uli::] k = j
uulk; k + j
the final output of this layer will be exactly equivalent READ b).
39

---

Published as a conference paper at ICLR 2025
Linear transformation R XXdout is:
The LINEAR operator, which maps inputs I Wxxd outputs y
LINEAR (u) uH
(21)
where H : Rdin + Rdout is linear map.
Define Cono (Wior(Win())) 1D, gate I, and out H. Then
W gate (u) o (W
so
Wont out ( gate(u) @ Con (n))) uH.
Thus the output of this layer is exactly equivalent to LINEAR( H)..
Element-wise multiply The MULTIPLY operator, which maps inputs E R Nxdin to outputs
y IR Xdout, is:
MULTIPLY(a, b, (u) uli:a: dowl @ ul: b: b + dout]
(22)
Define Cone = Identity, W., proj(a: a+ dow)), W gate proj(b: b+ dowt). ) and Wout = I.
Then
W (ga @ Cone ( in (u)) = ali,a: a + down] C ul;:b: b.bbb.. b + doutl.
Since Wont I, the output of this layer will equivalent to MULTIPLY (a, b, dout).
D.2.2 GRADIENT DESCENT We assume our input is the form
a1 bi
aN
1o 0
bn
Our goal is to compute the gradient update
200
(ang a, b:)a,
(23)
Intuitively, our argument proceeds similarly to the causal gradient descent construction from Ap- pendix D.3.I:
First, repeatedly apply READ and LINEAR move the information {a,, b; Vi into e.g. the final entry of the sequence. Without loss of generality, we omit the rest of the sequence, and assume we have access of arbitrary amounts of memory.
large enough embedding dimension that we can make use
After this phase, our u is of the form
(xo
a1
aN
bj
bn ...)².
Next, we use MULTIPLY and LINEAR to compute and store (af a: for all We will end up with
=.... (xr 0 {a:hi {bile (afa; ...).
We use LINEAR to compute and store {w) a, - 6.) for all
u = x {ails (bi)s (xfath {xfaa - bib. ...).
( We use MULTIPLY and LINEAR to compute and store (( - b.ja,} for all
u = (xo 0 {as}: (bi}:
- by b:)a ...). {oa a:); ((xTa. (
Finally, LINEAR compute the gradient update:
(I0 - 1 EN
(of a; b,)a, O (ath. (be), (afa;), ((afa b.)ar)s ...).
40

---

Published as a conference paper at ICLR 2025
D.2.3 NEWTON'S S METHOD
We assume our input is of the form
a1
AN
Mo[L.:]
My D.
bj
by
Our goal is to compute the Newton's ss iterate:
M1 M2 (a Mo),
(24)
where
b1
a
b =
(25)
ON
bN
For any matrix M E RRxp, let fll denote the flatten operation, so that (MI(M) represent vectorized version of M: fli(M) E RRP.
We proceed similarly to the argument from Appendix D.2.2.
First, we repeatedly apply READ and LINEAR to move all information {a:}: Vi and ())(M) e.g, the final entry of the sequence. We omit the rest of the sequence for notational ease, and we assume we have access to a large enough embedding dimension that we can make use of arbitrary amounts of memory. After this phase, we have
(f.(()))
{ai}: ...).
Using LINEAR, we can copy and rearrange the a:'s to construct copies of fit(a) and fit(a"):
(Ju(M)) {ai), sli(aa) su(() ...).
Now, that we can represent the matrix multiplication a a as linear combination of the entries of the element-wise multiplication flt(a")) @ flt(a). This means that we can obtain fit( a a) using a single application MULTIPLY and LINEAR:
= .... (o)) {aihi fit(a")) flt(a) flt(a"a)...).
By the same argument, we can obtain fli((a" using another application MULTI- PLY and LINEAR:
u (JU(MO) {as): fll(a") flt(a) (lt((aaa)Moo....
Finally, we have that fit(Mi) == 2flt(Mo) flt((a fft(((Iaa)M)) M)) can be obtained using LINEAR once more:
... (f...()((MMi) {art flt(a") fit(a) (
41

---

Published as a conference paper at ICLR 2025
D.2.4 SOFTMAX ATTENTION CAN'T TT IMPLEMENT ELEMENT- WISE SQUARING.
In this section, we consider the following parameterization of softmax attention:
Attn(w) softmax ((uW))uWWk) M) (uWv + B),
(26)
where u E RND Wo. WK. Wv RXD, BERNND and M NNN is the causal attention mask:
1-0 i<j
(27)
0
otherwise
Theorem D.31. One-layer single-headed causal softmax attention cannot exactly represent the entry-wise squaring function SQUARE: RNX + +NND s.t.
SQUARE(u)) =
for all RND.
Proof. We proceed by contradiction. Let's assume there exists WQ. WK. Wv, B R and B R NxD such that Vu RVD.
softmax (wW))) + M) (uWy + B) SQUARE(U).
(28)
Consider the set of inputs u E NDD with at most one non-zero entry, defined as
Uab (i,j) (a, b)
u|i, = 10
(29)
else
for an arbitrary choice a E [N]. b E [D]. Then:
0
ON
Q = uWQ
(30)
Uab W [blb;:]
ON
where Q's rows are all N except for the a-th, which und Wolk [b.
Similarly:
0
: uWK
(31)
uab W Wxll [b.c]
N
42

---

Published as a conference paper at ICLR 2025
and
0
:
0
uWy
(32)
Uab [bbb,
0
Then the pre-softmax attention matrix, A' = QKT, satisfies
A/5 = 2 (W (W 6j
= =
(33)
otherwise
Define
C= u (WQ W6bbo
(34)
Now consider what happens after we apply the softmax operator. Recall that the softmax operator is defined as
exp(=|i))
softmax [t]
(35)
Σー cxp(x(j))
for E P. Then A = softmax (A' + M) satisfies
i i
Aij =
i == a $ (i.j) (a,a)
(36)
Exp(C
exp(C)
exp C
Now let's consider the output softmax attention:
- A + B)
(37)
such that O SQUARE(u).
Note that for + a:
O
[ (V [k]
(38)
and this must also be equal to ON SQUARE u) j :]. We consider three cases:
First, consider < a in order from i = l,..., verify that (V + :] must equal 0 for all i <
Since this equality is true for all <
Next, looking at = a + 1, we have
((V B)aa (V + BB + 1,:]) NN
(39)
must have a
(V + B)aa -V + Ba 1,:]
(40)
Finally, from 2 a + 1, we can again conclude that (V + Bi. must equal ON for all
> 1.
43

---

Published as a conference paper at ICLR 2025
This means the only rows of V + B that might not be zero are (V B)(a, and (V + B) 1, Thus looking at the a-th row:
exp(C)
exp +a-1 (V + B) = SQUARE(u)(a..
[0
u
0]
Recall that from above,
Via uub Wvlb Then analyzing entry-wise, we have:
(41)
exp(C)
(Uhbb (uab Wv B(a,
(42)
exp(C) +
for all j + b, and
exp(C)
exp(C) + (uab (6, + Bla. b]) = uab
(43)
We now plug back in our expression for C' and simplifying the latter equation. For case of notation,
denote A = (Wo WK kb V == Wv lb, b), and B = Bla, b). Then the expression simplifies to:
V eexp(Aud (xpp(Aug)) uab + B exp( (xux()uat) exp( (Au uug + (a 1) 1)
This must hold for all non-zero values of uab We can take V B 0, but still left with
exp( Au ) (a 1)u
exp((Aua)
However, there is no choice of A such that this statement holds. This completes the proof by contradiction.
As corollary, we have
Corollary D.32. One-layer single-headed causal softmax attention cannot exactly represent the entry-wise multiply function MULTIPLY RNN - +R x s.t.
MULTIPLY b, (u) ul::a: (::a:aa a done] @ bl::b: b: dawe]
(44)
for all u N and all choices of a, b, dout.
Proof. Note that for b 0, and dout D,
SQUARE(U) MULTIPLY(a, b, dw)
Since softmax attention cannot exactly represent SQUARE for all u. it also cannot represent MULTIPLY for all
44

---

Published as a conference paper at ICLR 2025
D.3 UPPER AND LOWER BOUNDS WITH BASECONV FOR GRADIENT DESCENT
In this section, we detail upper and lower bounds for implementing gradient descent using BASEC- ONV, discussed in Section 4.1.
Upper bounds. We provide two explicit constructions for implementing iterations gradient descent linear regression: for non-causal BASECONV requiring 0(1) layers and ((D) state size, and one for causal BASECONV requiring ((1) layers and ( state size.
Lower bounds. In Appendix D.3.2, we prove that our constructions are asymptotically optimal with respect to layers and state size.
D.3.1 UPPER BOUNDS: BASECONV CAN IMPLEMENT GRADIENT DESCENT FOR LINEAR REGRESSION
In this section, we provide weight constructions for exactly implementing gradient descent on linear regression, Recall:
N
nN
x a; - b.)²
(45)
2N
So
1
V VpLN
x a; b.)a;
(46)
N
N
b;a;
Σ a:a
x
(47)
N
Non-causal BASECONV This weight construction uses Equation 46 to compute the gradient descent update.
We note that non-causal constructions for in-context linear regression are standard in the literature: e.g. Von Oswald et al. (2023); Ahn et al. (2024).
We start with input:
a
ON
0gg by
b
b
0
We define the initial embedding:
a1
...
ON
bx
0
1
0
0
0
0
a
0
We drop the bottom two rows of the block matrix representation for now and show how to perform the gradient descent update with the rest of the embedding.
45

---

Published as a conference paper at ICLR 2025
Layer 1:
1
<- a; +
bj
6 1
b;
20
f D
30 7
a
t 20
a O Zo
0
0
f
0
conv(in_praj()
gate_prost
t- a4
b; T
b,
< Fi 4
t +D +
out_proj(()
t a; ② DD
f a. ① DD +
t 0 +
(x) a, - bb.11
Layer 2:
f- 04
(- 1
f- -
+ bi +
f -
b,
ED +
t D
Zg +
t a, ① Zo +
1
f a, ② D0 +
(xf a; - bb:11
(xf a; bi)a; -
a
conu(in_proj(t)
gate proj(])
f- a. 7
f- a; -
f- b; +
<- bi -
f- xi0 -
f- Bo -
out_projt = Edentity
a; ZD
f- a. Bg
f (2) a; - -b,ja,
+ (x a. b,)a; -
Layer 3:
f- ai -
f- a; -
<- bi ->
<- b; ->
+- Zo -
x- Zg +
cono((n_prof())
+ a, Bg
a; @ Py +
+ (a) a; bi)a, +
si (2g a; - -b:ja: +
46

---

Published as a conference paper at ICLR 2025
-- a;
+ b; +
t +
b, 7
-- 20 7
20
1/ Vo((wo)
t ④ 20 -
out out_proj()
f- - +
Σ a a; b,)a; 7
1
-Va((00)
Causal BASECONV
This weight construction uses Equation 47 to compute the gradient descent
update.
We start with input:
0
a1
aN
b
by
0
0
We use two BASECONY layers construct an initial embedding, after which each gradient descent update step will only require single BASECONY layer.
In the following construction, we use flt to denote the flatten operation, which maps an M N matrix to M. N-entry vector with the same elements.
Layer 1:
0
1
a1
AN
b
by
0
1
D
0
(- 1 +
0
b,1
by1 0D
aN
(1())(1P)F)
(M(x((1PFF) (oP))))P))))
(u(1Paf) [ (ODP)
mate_projt)
0
0D
01
ON
a1 AN
b
bn
0
b bn
0
0
D
0D DD 87
out out_proj= dentity
b1aj biay 0D
biaj bia 0"
f(t(aaaa) flt((aaa))
(
flt(aia?) flt(ana() f(((0P(OP)z))
47

---

Published as a conference paper at ICLR 2025
Layer 2:
01
ON
0DD
f 1 D
a1
ON
00
b
by
0
+ 1
by
0
00
0
a
@
D
0
0
N
b:a,
t 1
Σ
b,a,
E flt((,a))
- 10a
Ei fit(a,a?)
conv(in_projj
gate_prodt
0
0
aN
a1
ON
by 0
6,
bn
0 D
Co
0D
0
Zo
oud_projildnntity
> 1==1
b:a;
+
+
Eim b:a:
-
E Em flt(a,a))) +/
N flt(a,a?)
Now, we use single BASECONV layer to implement a gradient descent update.
ON
D
+ 1"
a1
aN
0
by
0
t
1
+
6
by
0
0D
0D
+ 1D
0D
0D
x
0
0
D
f 12
0
0
10
D b:a; +
f 1
Σx
+
EN flt(aaa)
1
Σ N
fitifttaaaa?
Di b:a;
0
0D
D
0D
0D
Σn b,a,
En flt(a,a?)) +
00a
OD3
f1((1
0
00º
Σ ((t1a;(;00x,))))
conv(in_prao(())
gate_pros(-)
Note that the gradient
N Σ a
vaL(2n)
Σ b.a:
wo
can be written as a linear combination of the vector
MN b.a.
ΣN flt(a,(a, o xz)
48

---

Published as a conference paper at ICLR 2025
write weight construction for out_proj that updates wy - wo 1/ a C((x)):
a aN 0 D
aN 0
a1
b
by
b
by 0
0
0
0
0
Do
V
0r
0" P
0" r
0" 1
Σ b,a;
Σ b,a; 7
out_proj
D flllaaa?)) 7
2i=1 sllaaaa) 7
N
0
0
E b,a: Eiml
00
OD
Σ b:a:
O
N
0
O
ΣN flt(aila, O oz))
D.3.2 LOWER BOUNDS: BASECONV CONSTRUCTIONS ARE ASYMPTOTICALLY OPTIMAL
Note that the non-causal weight construction in Appendix D.3.1 requires O(1) layers and O(D) state size, while the causal weight construction in Appendix D.3.1 requires 0(1) layers and O state size. Clearly the ODD) state size requirement for non-causal models is tight, since one needs to store the gradient V E IR D. In this section, we prove that the O(D²) state size requirement for causal models is also asymptotically tight.
Theorem D.33. Any single-pass (causal) algorithm computing the gradient
N VaL
N
Σ b,aj
Σ
a,a
T
x
given inputs ((a,,b)...... (aN, by); æ. with (a,, bi) E ((++))N) and x E R.,, requires Q( D?) state size in the worst case, where bj I and aj. a E RRD.
Proof. For simplicity, we pick N = D for large enough D.
Since we can compute Σド= b,a in O(D) space, we focus on computing the expen- term. Assume there exists a single-pass algorithm A that computes (Σm= æ exactly for all choices of a a R., Now consider the following two claims:
sive
Define sD be the state of the algorithm after seeing d1...., ap. Then we claim that sp must have enough information to exactly reconstruct Mp = Σ3= a.a
This follows since the algorithm must be correct for any value E RD takes on. In partieular, setting & e; for i [D), we observe that the algorithm must be able exactly recover Mpe, Mp i), [D].
2. The space of matrices
Σ a
over all choices of aj RD, j [d] contains the set of all real symmetric matrices in TR DxD
49

---

Published as a conference paper at ICLR 2025
This holds since for any real symmetric matrix A, obtain set of possible aj 's via its eigendecomposition Strang (2012):
D A = QAQ"=
a
j=1
where a₁ = [:
From the first claim, conclude that sp must contain enough information to be able recover MD for any possible value MD can take on (over all choices of 0,,..., ap E PR From the second claim, we have that the space of possible values of MD includes the set of all possible real symmetric matrices. Since we know that this set requires (D)D+1) parameters to represent, we can conclude
that [ 2 (D)(++1) S2D22).
50

---

Published as a conference paper at ICLR 2025
D.4 BASECONY AND JACKSON'S THEOREM
In this section we prove BASECONV's ability to approximate arbitrary univariate and multivariate smooth functions.
D.4.1 UNIVARIATE FUNCTION APPROXIMATION
We start with a special case of smooth functions that apply entry-wise univariate smooth functions:
Definition D.34. Let [-1, + R be a (k: L)-smooth univariate function. Then define
I: 1-1, 1] Nx - IR. N
as follows. For all < N.O S < D, and u |-1,11xxx
[i = [(uli,
Now we will state a simple observation on BASECONV's ability to approximate these functions. Lemma D.35. For any smooth function J as defined in Desinition D.34, let g(x) = Py(x (x(x with Pj being the polynomial from Corollary D.12. Then for all [_-1, 1JxD,,
llg(x) - f(x) S.
Proof. Follows from Definitions D.7 and D.34 and Corollary D.12.
Next we will state construction of an arithmetic circuit for a function that applies a univariate polynomial to all entries in [-1, 1|xD
Lemma D.36. Let P(X) be a degree d univariate polynomial. Then there is D, OND)), ((d)) -iiiiruut compute P(u) where P(u) defined as follows. For an imput u E [|1,1]
P(u)[i. P(uli,]))
Proof. Let the univariate polynomial be
d
P(X) Σ xx
1=0
where coefficients E R.
Next we state the natural arithmetic circuit to compute P(x) for Algorithm 2:
Algorithm 2 circuit Cp(x):
8o + Co
2: mo +-1
3: for 1,2,...... do
4:
m, +- m
> Multiplication gate
5:
1j 4 Gj mj
B Multiplication gate & Addition gate
6: 8j 4- 85-1
7: return 8d
D 8 is the output gate
Next we apply the above circuit in parallel to form the circuit that computes Puu) in Algorithm
Algorithm 3 Circuit for P(u):
1: for = 0, 1....,,N .N do
2: for 0, 1,..,, D do
3: 2|i,j] = (p(((i,))
Do this in parallel is the output matrix
4: return
51

---

Published as a conference paper at ICLR 2025
Looking at Algorithm 2, the depth of the circuit is 3d, O(d), since that is the bound on iter- ations of the for loop, and each iteration we compute sequential operations. Therefore it's a (1, O(d), cirreci.
For Algorithm 3. The width is O(ND), since we have our input of size x D, which goes through the circuit in parallel, stated in Algorithm 3. Therefore have an OND,)(N)) O(DD)) circuit that computes P(u).
Since BASECONV has the ability to represent any arithmetic circuit, we get the following:
Corollary D.37. We can implement P(u) (where P(u) is as defined in Lemma D.36 when deg (P) with a (N, O(dlog(ND)) D, (ND), D) BASECONV.
Proof. Follows from Lemma D.36 giving the (ND, O(ND), (d), -ciruut for an arbi- trary polynomial and Theorem D.30 gives us the BASECONY model to implement the circuit.
We will prove a tighter bound showing we represent P(u) using constant number of BASECONY layers (for constant (eg(
Theorem D.38. We can implement P(u) when deg( () d with (O(N), D, BASECONV model.
Proof. We will convert the steps done in Algorithm 2 to layers of BASECONV. Since Algorithm is essentially running Algorithm 2 in parallel over all entries of input u E [-1, 1|NxD, the latter happens automatically in our BASECONY implementation.
For this proof, define
P2(X) X
and let C2 be the matrix of size N D and all the entries are Go.
We expand the input to our BASECONY layers follows,
u'
O3NxD)
This means that the size of the internal dimension of our BASECONY layers will be (4N, D). To begin iterations of the for loop we need to store initial values into the extra space in Taking us from
u NDD
11 1NDD NND C.
NND OND
Up
ONXD IND NXD Co
We do this via BASECONV(ui,
DDDD I
ANDD 1NDD) which computes
O NXD N N Co
u
11 NXD NXD NND
N
N N
DXD
+
OLNDD
4NDD
The above simplifies to
OND 1 INDD 1 N Co
NDD
(1 (1NDD).
NND 0 NxD
+
which gives us
u N
1 N
Uo.
C.
52

---

Published as a conference paper at ICLR 2025
desired
This was done with a (4N,1, D, D,,,,,NN,DD BASECONY layer.
Our goal is, the end of iteration to compute uy RANxD such that,
P P C, o P, 1 +... C2 O Pj(u)
4
We will view the above matrix in terms of the variables in the Algorithm 2 as follows
m
C Co Ci O ((( Cj
t
Sj
The for loop for values of - d which the remainder of this proof will replicate. There are three lines in the for loop in Algorithm 2 which we will cover how these operations happen in constant number of BASECONV layers.
In line 4, the first line in the for loop computes
11
11
m
m
(1)
Ug
fj-1
1
8j-1
88-11
Note that my = my-1 @ u.
We the remember primitive compute u follows
(1) from uy-1- Define R2NxD R2NN as
m3-1) m m Ou
If we can compute with BASECONY layers then can compute u.") (1) for u,-1 by calling rememberru, 0, 2N - f)
((n)
We show BASECONV
DDD 02NxD H, (NNx.)
maps
m,
m
m
where H is defined as in Proposition D.24. follows:
We plug the matrices into the BASECONV layer as
u (m--1)
DXD + 2
11
(INXD)
c
(ONXD)
mj-1
We know from Proposition D.24 that this convolution operation is a shift down by N rows. Therefore the above simplifies to
((OYDD)
11 (m)1))
DXD 02NxD)
(IM*D)
(INxD)
which simplifies
N
u
u
m
m o
m
as desired. Therefore by Proposition D.26, line 4 can be computed by (4N, D, 4N, D) BASECONV.
53

---

This results in Published as a conference paper at ICLR 2025
For line 5 of the for loop we need to compute
m
m
(2)
tj
S
5
Note that t, = C2 O my.
To do this we will use three BASECONY layers. We use the remember primitive to compute u from u (1), Define g I22xD > 2 as follows,
m
my
C
we can compute with BASECONV layers then we can compute u for u by calling rememberruy (1) .N,NN - 1,9).
Indeed,
show the g
be (
computed by first computing
BASECONV
m, (c)-))
DDD O2NxD ONNxD
DxD + 02NDD]
2 (2NDD
m
1 (ONX.)
tj--
t
+
which simplifies to
1 1XXD)
(NNDD
t
m, ONDD
We pass into the next layer, BASECONV (OoKD)) DDD (ONxD) defined as in Proposition D.24:
H. ()
where H is
(ONRD)
my
DXD (ONXD) 1NxD
m, H
(1xD))
ONDD
(NND))
Since the kernel H is as in Proposition D.24, this simplifies to
(1
(ONxD)
1 N N
m (1NRD)
+
m,
The above simplifies further
1 NxD
my 1 N
m,
which results in:
m m;
DDD ONND O2NxD (INXD)
We pass the above to BASECONV
(m:)
m
DDD + 2NDx
(2NxD
m m,
NxDD
+
mj
Cj
which simplifies to The above results in
m
INDD
m
C2
m
m
C
=
t
54

---

Published as a conference paper at ICLR 2025
desired.
Therefore by Corollary D.27, line was computed by (4N, D) BASECONV. For line 6, the final line of the for loop, want
u1
(2)
m
U
f
{j
$j
Note that 8; Sj
Define function h: R2NND + 22NxD as follows.
t,
h
8j-1
8j
If we can compute h with BASECONV layers then we can compute uy for uz-1 by calling remember(u)) (2) 2N, - 1,h).
Indeed
show
that
h
be
computed by computing
((1))
tj 0 DxD 1 x xD H. 02NxD)
) where kernel H E 2 is defined
BASECONV
55-1
as:
D if {O,
H
0
otherwise.
This layer computes
ti (s)-))
02NXD
12NxD
H
tj 8j-1
+ 022xD
This simplifies to
(1 22NxD o
H *
t 8j-1
= H
tj 8/-1
Now we compute this convolution for column < i
2N. For notational convenience, let
tj
be noted as matrix V. Then we have:
$--1
H
Vi. coeff + ( mod x²N),
where (1 (1+ +N is the polynomial representation of the columns of H (since there's one in the 0th index and one in the Nth index of each column).
The expression simplifies
coc + (X)x mod X
which can be broken down to coeff ([V[0 v[[]i]X +.... V/2N 2 mod X2N)
+ coeff ((V]))X + V[I[i] +... V/2N 1[ mod X2N)
with the lower order terms in the second coefficient vector being zeros,
cooff ([V[0][++ + V +... V/2N 1[[:]x2-11) mod X2N)
coeff ((0+ + OX + 1O V( +...+ V(2N - x3N1)x mod X2N)
After taking mod X2N get
coeff (V[o + V + + V2N - 1111]22N-1
coeff (0 + 0.X ....0x N N- v +... V[N 2N
55

---

Published as a conference paper at ICLR 2025
The first set of coefficients is the input matrix as is. And the second one is the input matrix shifted down as seen Proposition D.24. Therefore when we add these vectors we are doing
1 8j-1 +
(ONXD) (0,,)) 1
8j
as desired. Therefore by Proposition D.26, line 6 is computed with by 1, D) BASECONV.
The Sy matrix gives us Co + C1 O m + ... + Ca 3 m. Recalling that
Co C1
Ca m Σ
Coon's P(u),
==
and hence Sa is our desired output.
We have d layers, each consisting of 0(1) BASECONY layers. Giving us O(d) many layers to implement Algorithm 2.
Therefore, via the ability to stack BASECONV layers to do function composition, the for loop was computed by (4N, Odd), D,
D,d,, 4N,
D) BASECONV, as desired.
The following states BASECONV'S ability approximate univariate smooth function:
Proposition D.39. Let f be the (k,, -smooth function defined in Definition D.34. Then there is (N.O vVP) + D, D) BASECONV model that approximates within error
Proof. Follows from Corollary D.12, Lemma D.35, and Theorem D.38.
D.4.2 MULTIVARIATE FUNCTION APPROXIMATION
We consider the following multivariate functions:
Definition D.40. For < < N, S j D, let Figs [-1,1]8xD E be a L)moothh
multivariate function. Then define
f(x) 1-1, 1 NND
as follows. For all S N, S < D, u NXD define
f
Lemma D.41. For any smooth function f as defined in Definition D.40, let g, XNXD)
(y........ XNxD) be the polynomial from Corollary D.14. Then for all x E [-1,1]]xD,
llg(x) - s(x]l Se.
Proof. Follows from Definitions D.7 and D.40 and Corollary D.14.
Next will state construction for arithmetic circuit for function that takes [-1, N variable input:
Lemma D.42. Let P(X) be a degree multivariate polynomial. Then there is (n, Old n (ddog(()) ((n)))-eircuil to compute P(u) on any input u E 1_.
Proof. Let the multivariate polynomial be as defined Definition D.6. We build the circuit to compute this in Algorithm 4,
56

---

Published as a conference paper at ICLR 2025
Algorithm 4 circuit Cp(x):
1: for a (a1,..., an) 220 such that Ei=1 a4 S do
2:
ma
3: for i =
n do
b Done in parallel
4: if ai then
5:
Ma f ma 24
6: la Ca Max
7: for a ( an)) 8g0 such that Ei=1 ai S d do
8:
Σta
Done in parallel
9:
return S
We compute the for loop starting line 3 by making multiplications in parallel. Therefore obtaining depth O(og(()) We also have the for loop starting on line 7, making pairwise addition operations, resulting in depth of O(dlogg())
We again use the result that BASECONV can represent any arithmetic circuit to get:
Corollary D.43. We can implement P(u) (where P(u) is as defined in Lemma D.42) when XND)) d with (N, O(dlog(N D. ((ND)))) BASECONY where u 11.111
Proof. Lemma D.42 gives the arithetmic circuit that computes this polynomial. Then via Theo- rem D.30 we get a (N, D,), D) BASECONY model to implement the circuit.
Finally we state BASECONV's ability to approximate multivariate smooth functions:
Proposition D.44. Let be the function defined in Definition D.40. Then there O(d D)), D, D) BASECONV model that approximates J to within error with Ox( // NDL...
Proof. We get the existence of polynomial that approximates for some from Corol- lary D.14. Then via Corollary D.43 we get that we can represent any polynomial, implying (N, Oddlog(ND)) D. O((ND)"), D) BASECONV represents any polynomial that approximates the multivariate smooth function
57

---

Published as a conference paper at ICLR 2025
D.5 ZERO POPULATION GRADIENT BASECONV ON PRIMITIVES RECOVERS EXACT SOLUTION
In this section, we prove we can recover the functions LINEAR and MULTIPLY exactly given the expected gradients of their respective loss functions being 0 along with some necessary assumptions.
D.5.1 NOTATION
We start by defining additional notation for this subsection.
For readability, we will redefine how we index entry of 2 dimensional matrix note that we using O indexing [M] = (0, ...... N - 1]. For an entry matrix A where the row number and is the column number, define the parameters of the layer as follows. We have the weight matrix, W
denote it A.j. Now recall our BASECONV layer in 18, will
{W Rdd 1, E R and
the kernel matrix, K = {Ki}} E RNxd, the first bias matrix B(") the second bias matrix, B = (B E TR- We denote the array of these parameters as 0 (W. B(), B. Therefore have the BASECONv layer operation,
Z BASECONV (0, u, c, doua) def [u + B()) @ (K* x B0O) l:c:ct c+ dow - 1] (48)
for some integers dout [d] (with dout # 0) and [d dout] that we use truncate columns of the output layer to match function input and output size as stated below.
Recall that u is the imput to BASECONY layer, {u.,} Rxd.
Moving onto our target function:
f Nxd + R Nxdout
Naturally, for i E [N] and j E [doa]]. we'll denote (S(u)) [.,]) by (())
Next, we define the training input distribution.
D.5.2 TRAINING INPUT DISTRIBUTION
1. Let A be the training distribution on R.Nxd such that:
Assumption D.45, Given a monomial, II, (Uirde) "h if my is odd for some & then E [[ me
0. Otherwise, E [T, (u. [m] > 0.
Assumption D.46. Assume that the training data is generated as
u A as input
Output is y fu) + & where & {E,,} ) Nxdost is the random error matrix such that
The distributions & and A are independent. (Call the distribution on be Ag)
][E,,] = for all (i.j) E [N] x [doue]
Loss function
Define for E [N] and j E [dout]
(u, 0, (Zi, - - = (Ze. - f(u),s - -Ens)
(49)
L(u) = E 1=
dout -
Li (u, E)
Training loss, (0) [(t = A~ [L[u]]
E~AR
VoT.0) (0) EE, DN'1 Eラ=0
V (
58

---

Published as a conference paper at ICLR 2025
The Goal Given target function f. what can infer for 0 (W, B(!) BB2) from VoL (0) 0?
1. Ideally, would like assume that can be represented exactly by I-layer BASBCONV.
2.
For now, let's assume that ( only depends on u,;
This includes special cases:
f(u) dout U. bob for some integers a, b E [dow]]
f(u) W for W E Rdxd
We want to prove that when the gradients of the expected loss function are O, then the set of parameters that satisfy the condition perform exactly these functions.
D.5.3 A GENERIC PARTIAL DERIVATIVE
Let's try and reason as much as we can for a generic partial derivative. Let x E
(W, K, B B())) Then from Equation (49), we have that for any (i,j) E [ND] x [dowr]:
cTig =2 Zi.j f(u)(,) Eij)
azij
dr
Tx(
-T(?3
-T(33
=2
bid
where
aZis
((() Zii Oz
aZi.j
T( = f(u) i.J dx
ÖZij
T = Et,) dx
Proposition D.47. Eu, (T (3)
0.
Proof. Follows from the facts that and Ag are independent, and B[E,,] =
From now on, we will ignore the term ( because of Proposition D.47 we can (in expectation)
assume that (3) 0.
D.5.4 SETTING THE GRADIENTS TO 0
In this section we will prove the gradients of the loss function are O, under some given assumptions on the input data and parameters, when the functions we're learning are MUI.TIPLY and LINEAR. In order to do so we need to have another restriction on the target function f which is that f must be defined with a linear map, W E Rdxd that has non-zero columns. We make further assumptions on W, W,B (1), K, B(2), which we will justify later. We note that this assumption is satisfied for both the MULTIPLY and LINEAR functions.
Assumption D.48. For j E [doux] and & E [d - dont). (i) either (K.jta + 0) or (B.2) # 0)
and (i Wo * O. Further, we have (11) = O.
The target function f Nxd + N Nodus is
1. Implementable with 1-layer BC W K, W,K,B") ()) B(O) such that for all j [dovs] and (d dout), Waite #
59

---

Published as a conference paper at ICLR 2025
2. f only depends u
We make another assumption to assist with the following theorems, Assumption D.49. For all j E (c....,c+ c dout - 1}. (W W ..) + 0 and W is 7 0.
The main results follows. First for the MULTIPLY function we have Theorem D.50. Given Assumptions D.45, D.46, D.48, and a function
f(u, b, dow) = -oo::toout-- 2 U, bibbbbuustt1
where a, b E [d - doal and with a S and c = a. Let Og be such that, EVo. 0 then BASECONV(u, Oo) [=; (o:cc.. dout] f(u).
We prove a similar result for LINEAR function:
Theorem D.51. Given Assumptions D.45. D.46, D.48, D.49, and function
f(u) uW.
Let 0g be such that, EVoLo0-0 0 with and doat = d. Then BASECONv(u, 6,, 0, f(u).
The following is to provide information about each entry in the output of BASECONV layer.
Lemma D.52. For all (i.j) [N] x [dout], the entries of resulting layer of BASECONY, Z. are:
/d-1
(()
Zi, Σ
u.e Wes+e + B(ste)
Kk.tte us + Bijto
B (2)
(50)
Proof. To begin, from Equation (48), we know that a layer of BASECONY yields matrix Z as,
u W + B o K (K*u * u B(2) B(3) [;,c c dout - 1].
Looking the W operation, know that for row i [N] and column j [dout], the vector dot product is computed as
T W. Wwe+e)
Σ u W Weyee-
Meaning each entry in the resulting matrix is defined as such
(u Dijto
Σ
u
Wii+e-
E=
To sum the matrix (()) to this operation, we simply add the corresponding index giving us
dd1
uu W B(L)
Σ
(51)
Then, for the convolution operation between K and u, that's computed column by column, we have for all j [doue]:
(K * u). = Kiste * U.ite
i.c for any i [N].
(K * u (.
KR.tt Ui-kjtt-
k=
Finally, to sum (2) we add the corresponding entry giving us
K (K B())
Σ Keifte uitkitte + B (2)
(52)
iift'
k=
Combining Equations (51) and (52), gives us Equation (50) as expected.
"These assumptions without loss of generality.
60

---

aB (2) Published as a conference paper at ICLR 2025
D.5.5 SOME PARTIAL DERIVATIVES ARE ALWAYS ZERO
To simplity future computations in this section, we will state simple lemma on some partial derivatives that always go to O.
Lemma D.53. Fix E [ E [dout] . E [d] CE [d dour)- Then for any j' +j+cOss<N, and 0 - k < N have,
(Z 0Zis OW e.j' kK
= 0.
Further, any (i,j - c) + (V, we have,
ÖZ 0.
aB (1)
V
Proof. Follows from Equation (50) and definition of partial derivatives.
D.5.6 GENERIC FORM OF PARTIAL DERIVATIVES PLUS A CONSEQUENCE
Given Lemma D.52 we can conclude the following.
Lemma D.54. For 0 S < N, 0 S <
dont and 0 S c < douts any entry
{ iite B
(1)
aZ
0 (KKott u..+o) [i] + (22) iifte) dx ((uf., W Wiste) + B
dx
then for any entry x E {Kistor B.itch.
aZ.j
à (KK u.ste) [1] B(2)
((uT.o Waste) Bijte) Or
dz
1.j+c)
consequence of Lemma D.54 is the following.
Corollary D.55. Let 0 (W. K, (1), B =
Then for all parameter variables x, we have
0Z,,
Or
Specifically,
VOI (0) l0 =
Corollary D.55 implies that initializing 0 0 a good choice for initializing parameters since is a local minima.
We can exactly figure out the partial derivatives Lemma D.54 by the following.
Lemma D.56. Fix i [N], [d]. Then for any 0 l N we have,
a
((uT) W.y) + B u.e
oWly
and
(1uT. W.s)) B 1.
aB(1) (1)
Also,
0
(K u [0]+ 1.
aB (2)
61

---

Published as a conference paper at ICLR 2025
Proof. Let's begin by looking at Next, for any 0 H I we have
0
(Kiy (Kins+ u + (1(++BBB))
u,
for all l i.
(Kag * u / O.
0K..'
((u) W..) + BB
oWey
Expanding this out gives
Σ
ung We + B
aw i
P=D
When we take the partial derivative of this with respect to We,s the B term
term goes to 0. And the
Σ u We's =iet
oW
P==
as desired, since We,5 only shows up in the summation when E' = l.
Next, let us look at
(uuT,
W..) + (
s.
Since B. doesn't show up in the dot product of the vectors, know that piece goes zero, giving us
a
(u.
W W.+)+
B
0B
11) aB
aB
as desired.
Next, for any 0 S k S i we have
oKg,'' (Kog)
u. [i]
(
The ((2) '. term goes 0 as we're taking the partial derivative with respect to Kkyy. So we have
OK (KK J" u [i])
K uI-k ui-kj'
OKr
===0
as desired, since Kay only shows up in the summation when ke k. Next, for we have
i
i
((K u
K u.
(53)
oKr
OKE
as desired, since Key will never show up in the summation as ke Finally, let us look at the fourth piece.
(K. u.
+
B(9)
0 '.j'
62

---

Published as a conference paper at ICLR 2025
The term By doesn't i appear in the result of the convolution operation, therefore that piece goes 0, giving us
①
0B L 1,
(K. +u +
(2) (
0B 1.
6if'
desired.
Definition D.57. For the rest of the section, we will redefine (W, K, B()). Note that we are just removing (1) since is all zeros as per Assumption D.48.
Lemma D.58, Given Assumption D.45 and recall that B O. Fix i E [N]. E [dout] c (d doa) Then we have
E
Zi,i
UB
(2)
=
B
s.f+c
(2)
E
W
Pifte
i.jtc_
#1=0
Next, for any $ S we have
aZ,
d-1
Zi,5 OK Kk.tte]
K Kiitte
W e+o [ui,e ul-k.stel
='=0
For k >
i ÖZ
E
Zo.j
0K Kk.tto
Finally, for any S 1,
OZi. Zi.i aW, (Wott
(B 2
E
W Wijto
Σ
K
[u
W We [uf
k==0
Proof. Given Lemma D.54 and Lemma D.56 (along with the fact that B()) = 0) we have
0Zi
E [Kijo s+ol]+ Bl3to) (u. Wiste)
E
Zi
0B kote
(2)
E uu uu We We WecK jtc
p= 00=0"" ==
B 2
E Uge WerteWWerite Σ
P==
In the above, the first summation goes to since for all l', l" k. by Assumption D.45, the expected value the product three u's will be since there's 's an odd number of them. Again, by Assump- tion D.45, the second summation will be non-zero if and only if l 8". Therefore get the following,
oZ,,
d
E
Zi.
BBC
[
W2
0
(2)
iitt
1.
='=0
desired.
63

---

Published as a conference paper at ICLR 2025
Moving onto the next piece, using Lemma D.54 and Lemma D.56 (along with the fact that B(!) 0) we have for 0 < k
I i
E
iZ. Zi.i K Rkj+o_
= 13
[K (Kitte u.j+
[
B (u Wis+o) u uisko
d-1 d-1
L Mi We
6' # pr
d-1
d-1
+ B
Σ
E
We
P= P
By Assumption D.45, only expected values of terms with square monomials are non-zero. Specifically, the first summation has the u1-A. k.j term, therefore, we need K' the same reasoning for f" = L" Therefore, the first summation is non-zero if and only if k' = k and = t". The second summation will be 0 since for all f,, L" the expected value the u's 0 since there's s odd number of them, there will always be odd exponent. So we get
get an even exponent. This
OZ. Zu.j OK kKtt]
E
Kajto
Σ =0=0
W Je
[uge uitt.
desired.
When & > i,
ÖZ.
E
Zi.i 0K k.joc]
0
since index the convolution piece at 0K for le >> i will never be in the piece we're re taking the derivative of. Moving onto the final piece, given Lemma D.54 and Lemma D.56 and ( we have
8 0Z.
= E ((K
(u W Wastta Ung of u. [ (B
E
Zi.i aw ejjt_
E Ho stou
0=
+ 2B
(2)
E u; Mt KesteWer
ijtc
K=0
(B(3) (isto
E uud We
We again Assumption D.45 simplify the summations. The first summation has the u,y term, therefore to get an even exponent on it we need l' = This is the same reasoning for k k" Therefore the first summation will be non-zero if and only if l l and k' k. The second summation will be for all K', E since we're taking the expected value of an odd number of u products, there will always be an odd exponent. The third term will be non-zero if and only if l E to get an even exponent on u's entry. Therefore we have,
E Zi. OW [Witte]
oZi,
Weste Σ KZ B
[ stc u + (B()t)) W Wes+oB [u
M=0
as desired.
D.5.7 LINEAR
The following lemma will be for when the function considering is linear map.
64

---

Published as a conference paper at ICLR 2025
Lemma D.59. With 0 and dout d. i E [ E [d]. and W E Rdxd
Then we have
aZ.
d-1
E
W. (.))
Σ
Wt, W es [
0B
For all k,
OZ
E
u W. ..))
oKkj.
For all E
aZ.
u
B
W
E
[u
WWej
1.
Proof. Given Lemma D.54 and Lemma D.56 (and the fact that B 0) we have
iZ..j
u
E
(u.
W [
0B
d-1d-1
ΣΣ
Wt
W
[ue
8=0==0
From Assumption D.45 we get that the summation will always be non-zero if and only if so that the variable has even exponent. Therefore get,
dZ. W.
WayB [
YB (2) 1.j
as desired.
Next, for all k by Lemma D.54 and Lemma D.56 (and the fact that (1)) 0)
aZ,i
E
U W. (3)
EE Ku
War)
u (u W w
0Kki
d-1d1
ΣΣ 2=0==0
We, W
u, u
We simplify the above using Assumption D.45. The summation goes to since there are an odd number of terms, there will always be odd exponent. Note that this is true for i. Recall from Equation (53) that for K > i.
aZ, 0.
OK
Therefore, for all k,
ÖZ
W
..)
K
as desired.
Next, for all E, by Lemma D.54 and Lemma D.56 (and the fact that ((1) 0)
oz,
U W.s)
E [(K...
u [i]
B
u (u. Wis)
K W e
d-1
+ B (2)
W e We [u.euee]
6=0
65

---

Published as a conference paper at ICLR 2025
We simplify the above using Assumption D.45. The first summation will always be 0 due odd exponent on the u's. The second summation piece will always be non-zero if and only if = E. giving us the even exponent the variable. Therefore get,
öZ. OWe..
E
B 12 W
E u
as desired.
Next, we restate Theorem D.51 and prove it:
Theorem D.60 (Theorem D.51, restated). Given Assumptions D.45, D.46, D.48, D.49, and function
((u) uW.
Let 0, be such that, EV 0+00 0 with 0 and dout d. Then BASECONV(u, 00 dout) f(u).
Proof. From Lemma D.53 we get that
aL
N N-18
daue -1
ÖLe
E
(B
OB!2) b
(2)
5
dLij
E
ÖB
(
Recall loss function from Equation (49). Then given Proposition D.47, Lemma D.58, and Lemma D.59 we know that
OL
az: Rj
E
21E
Z
(uW)
12)
aB 2)
1.j
d-1
d-1 Σ
= 2
W f
We, W We,B [uid
t.
t=
t=
Setting this to 0 and solving for (
1.j gives us,
B
D ΣE
W
Wes
Given Assumption D.49 know that the numerator will be non-zero and given assumption Assump- tion D.48 we know the denominator will always be non-zero as well. Therefore we get that for all j
W 5 def bj.
B
(54)
(W W
Next, via Lemma D.53 we have for all 2 0:
aL
dout
OLv3
E
E
KK
==0 j'=0
N
E
[Kaij
Z=
From Equation (49) and Proposition D.47 we get
OL
N Σ
Zeri (uW)..)
E
2E
OK
R.j.
ok
='=0
66

---

Published as a conference paper at ICLR 2025
Then by Lemma D.58 and Lemma D.59 get
N-
aZ
iZ 0KuJ
N-1
d-1
21E
Zvd
uu
2
Σ
Kkj
Σ
W
up
t..
===0
P==
P=0
d-1
N-1
=K
Wes
E Σ P==
-
P=0
Since we know that the summation piece over always non-zero due to the even exponents on the terms and at least one of We' + 0. Therefore, when setting this to 0 and solve for KEg gives us,
0T.
E
=
kK.. W
KK.
implying for all k. KEj 0. In other words,
K =
Finally, via Lemma D.53 we have for all f:
aL
N-1
doul
OL.Y
E
E
aWei
WWej
P=
j'=00
N
E
From Equation (49) and Proposition D.47 we get
aL
N Σ
Zed - (wW)xa)
2EE
OW
='=0
Then from Lemma D.58 and Lemma D.59 we get
N-1
az
aZ OW
N
21E
Zea
(WW)
=
Σ Wis Σ
K
aWei
='=0
y=0
k'=0
+ 2 (B/?))
(2)
We
W
[
B
W e,
Which simplifies to
N-1 We
N-1
2
[u?
E.
Σ
B
B
===0
='=0
We can drop the summation with K in it as we know K O. Recall that from (54), B bj. Then rewrite the above
N-1 26,
E
aL aW
E
(We,by-Wa
P=0
know from Assumption D.45 that the first summation will always be non-zero since there's an even exponent on the u variable and we know that by is non-zero. Therefore setting
aL
E
0
ow
P.j
tells us that We,hhy W Ł.j or,
We E
W =
(55)
b
67

---

Published as a conference paper at ICLR 2025
Given the above value for We, and recall have K 0 and B (bobi...ba-z) where each by a column vector comprised of all by values. Therefore, when we take BASECONV(u) we get
BASECONv(u) (uW)
[x u B
(uW)
(
We can rewrite (2) as
bo 0
0 b1
0
B (1xx
ba-1
let us call this diagonal matrix on the right, D. Then note that by Equation (55)
W WD
Therefore, we have
BASECONV(U) wWD o N
=uW @ 1 Nxd
=uW,
desired. In the above the second inequality follows since D diagonal matrix.
D.5.8 MULTIPLY
Note that for i [N],j E [dout], the j, entry of MULTIPLY b, dow)) is
MULTIPLY (a, b, dout)i.s = Uij+a Uij++b
Lemma D.61. Fix ix b, where i [N] and a, b. E [d dowl and j E [doat]. Then we have
E
(uijta Ui.j+b)
(56)
0B 1.j+0_
When is 0:
aZ 0.
E
(uista u(,+b)
(57)
0K KKitto]
When a = b then for all c:
8Zi.j
E
(uista Ub.j+b)
R. [u ubstel W Witasto
KKojte
When then for all b:
aZ.,j
B
(uista We.+t))
[u] soo m Withite (58)
[Kojte]
When b = whennfoor all
dZ.,
[uis+o wi,+
= E [uite utjttl Wstastt
OKojte
For all other values of a, b, (i.e. a, b, are all distinct).:
0Z.
E
(uista (i.++b)
0.
OKujtel
Next, when E (= and b =c
aZi,j
E
(uista Uistt)
= [u ifte u² uiigtal Kojte
aW [e
68

---

Published as a conference paper at ICLR 2025
When 8 = b and a c,
iZ
E
(uigta- Uig+b)
I [ua lott u² 1g40)) Koito
(59)
@Wr.to.
When ll = j + and a b,
iZ.
E (uajto Uig+b)
E [u² ijto u² udite Kojte
OW e.
For all other values of t,a, b,
OZ. OW @Wejta. e.jte
E (uij+a U(j4b)
0.
Proof. Let begin with
aZ
E
Uizta
Wig+b
(2)
From Lemma D.54 and Lemma D.56 we can simplify this to the following (recall that B(1) 0):
E Uijta Wijtb (u W Wog+e))
(Kifte + Unite [ + B (2)
(
(2)
istt
[u " U (u W Wo))+))]]
This can be rewritten the following
E lu.jfaa Mufouuel We
From Assumption D.45 we know this is always as there's an odd number of u's being multiplied together each iteration of the summation. Therefore get
0Z.
E
Uij+a
dij+b
0.
(
0B (2)
s.jtc_
Next, let consider for lt < i,
0Zij
E
(ui.iaa u.gj-b)
[Kijto
From Lemma D.54 and Lemma D.56 we can simplify this to the following (recall that B()) 0):
E Uijta Ung+b (u Was+e)) 0Kk,ttc
(Kijte* unjto [1] + B b.jte)
=B [uata Uig+b (LuT Waste)) ui-kjto]
This can be rewritten the following
Σ E (u.f+au.J+oul u.jkjte Weette
We have the following cases about the expected value of the above:
1. When k > 0 for any l' a, b, or e we get the expected value is 0.
2. When b, we get the expected value is E [uā ufytel W Witegie if and only if i and l'
69

---

3. When c, we get the expected value is I u j ui,40] Ko.ito and only if k' = and f j b. Published as a conference paper at ICLR 2025
3. When a = c, we get the expected value is E uf u u Wo if and only if =
and =j --j+0
4. When b = c, we get the expected value is E [ui uaa+ and = + a.
W Witagte if and only if k = 0
5. For all other values of f", a, b, and get the expected value is 0.
Via Assumption D.45, the reasoning for and is that there will always be an odd exponent on the u's. Then the reasoning for Items 2 to 4 that there will always be an even exponent the u's. Next, let us consider,
aZ
E
(uiita Ugg
aW (WejteJ
From Lemma D.54 and Lemma D.56 we can simplify this the following:
a WW ei
E Uijta
Uig+b K (Ko
* u. [i]
B (22)
i.jte)
(u Wisto))
=E uiyta uni+b
Kijto u.oto + 1,ite) u
(2)
This can be rewritten
EE Uu Hig+b" K Uil K (+e) + IE (u.j+a Hog+bbl B
This second term goes to zero via Assumption D.45 there will always be an odd exponent on the term we take the expected value of. The first summation will be zero or non-zero given specific cases, just as we previously saw. Here they are:
1. When k' > 0, for any l, a, or b get that the expected value is O.
2. When b, we get the expected value is EE [uā ur and L = j
Koste if and only if k' =
4. When = c, we get the expected value is E [ue vaa and L = j
Koste if and only if k' =
5. For all other values of f, a, b, and e, we get that the expected value is
Via Assumption D.45, the reasoning for and is that there will always be odd exponent the u's. Then the reasoning for Items 2 to is that there will always be an even exponent on the u's. Each of these scenarios, covers the pieces in the lemma statement.
Next, we restate Theorem D.50 and prove it:
Theorem D.62 (Theorem D.50, restated). Given Assumptions D.45, D.46, D.48, and a function
f',, b, dow) = Uu W.
where a, b [d - dout] and without loss of generality assume a < b then take Let Oo be such that, E. TV then BASECONV(u, (:) 0: dout 1] f(u).
70

---

Published as a conference paper at ICLR 2025
Proof. For io E [N] and j.j E [dout] and a, b, defined in theorem statement, let us consider
àL
N 1
doue-1
E
=
E
0B!2 Ljte
Ljtc
v==0
0
TLis
E
where the second equality follows from Lemma D.53.
Recall our loss function from Equation (49). Then given Proposition D.47 we have
aL
iZ (2)
2E
(Z
(Uigia Uij+b))
( ijte
Simplifying and plugging in values from Lemma D.58 and Equation (56) from Lemma D.61 we get
öZ.
OZz (2)
2E
(Z.,) (uifta mi5+b))
= =B
Z
(uij+a
u.g+b)
(2)
0B
biftc
i.ftc
1.f+c
d-1
2
2
[ut
W
10+0
Pite
8==0
From Assumption D.46 we know the expected values of the squared input terms will be positive. Then from Assumption D.48 we know that least one entry in Woj+e is non-zero. Therefore the above summation will be non-zero. Further implying, when we set
OL
E
0,
(2)
b.ft
we can conclude that ( Bejtt for all i,j. Explicitly,
B (2)
pettdou--1 Nxd-1
(60)
Recall from Lemma D.52 that we only consider values of the parameters in column range (c....,c- dous - 1}. Therefore, moving forward, the (22) terms will be dropped from equations to simplify them. Next let us consider for all 0,
OL
N-1 dout
E
0K kKks+e
0Kiitol
1==0
j'=0
ÖLv
E
the second equality follows from Lemma D.53. Then from Equation (49) we get
Ol
N-1
E
2F
(Zr
(urjta my)+)))
OK
KKjtte]
ÖKkjte
y==
Simplifying and plugging in values from Lemma D.58 and Equation (57) from Lemma D.61 we get
N-1
aZr
N-1
d-1
Σ
2E
OKkjjto
(Zy - (uajta urrste)
= 2
Σ Kiitte
c
jite
[u lr u
k.je]]
P=0
p=0
#=00
N-1
2 K KNoje
Σ
W
Σ ==
E [u e u
-kteel
P=0
71

---

Published as a conference paper at ICLR 2025
From Assumption D.46 know the expected values of the squared input terms will be positive. Then from Assumption D 3.48 we know that at least one entry in Waj+e is non-zero. Therefore the summation piece will always be non-zero. Therefore, when setting
OT.
E
0,
we can conclude that Kk.tte = for all j. > This implies that we have for all j E [doue]: Kit+o # + Ko,+e + O.
(61)
Next let us consider,
OL
N N-1 doue
E
OW ejjt
Weijtel
N ai
E
WWejite
r'=0
The above follows from Lemma D.53. Then from Equation (49) we have
aT. ÖWe.jte
N-1
aZ
E
Σ 2
(Ze
(uritta up ue
WWijoo
After simplifying and plugging in values from Lemma D.58 and Lemma D.61 (and recall that ( [; dout 1] O and c) get the following for & + bj
N-1
JZr (Zr's - (urg+a (rr)+)))
2E
=
Weite
Σ K
[ut Mife
AW.jto
P=0
F'=0
k/=0
1'
2 Wejte Ejtt
ΣΣ
K
[už
- Mite
u
k'=0
From Assumption D.45 we know the expected values of squared input terms will be non-zero. Then from Equation (61) and Assumption D.48 we know that the 0-th entry each column of is non-zero. Therefore, this summation is always non-zero. Further
say, when setting
E
0,
aw eijtoJ
can conclude that We, = 0 for l # + b. Explicitly,
W.f+c + 0 + Wi+bs+e # 0.
(62)
Now let us consider the following for = j + b. by Lemma D.58 and Lemma D.61 (and recall that 2 [ c
dmm
1] and c)
N
0 (Zys - (Uij+a up u+)+)))
Ow
stbif+o
-
N-1
=
W
Wjtbjte
Ke Be iu u
goo - E [uf ub Kosto
H=0
k==0
From Assumption D.48 and Equation (61) let us simplify the above the following:
N-1
Σ
W Witbs+e K2
E. Ju stc
[u
ste
Ko
N
N-1
K2 2 WythsteK Ogite
Ju Jtc
Kojtt
Σ
E
[uu
3==0
R=0
N
2 (W Ko
0.4+c
Ka,tte)
E u J---" u uz.
===
72

---

Published as a conference paper at ICLR 2025
From Assumption D.45 we know the summation pieces are both always non-zero. Therefore when setting
aL
=0 = 2 W j+ K
Ko,j+o)
OW
j+bjtte
0.j+c
WobbjteK 0.jte 0,jt Koofte
Witbs+eK Koito Ko.tt
WitbiteKo.st Ko
In the above the last equality follows form the fact that Ko.ie O. Thus, the above gives us
Ko,f+c Ws+bi+t
(63)
Note that this is valid assignment since Equation (62) and Assumption D.48 implies W j+b.i++ o O. Therefore, given Assumption D.48 and the above values; W W of for all j E [dout] and W 0 for all other (t','") Note that multiplication
the right of u with this lower left shift matrix will shift the input to the left by c. And we have cecdoue
Wi+b.++" -latdddo--ttttdduue- DV-Ixdout Ixdous
(2
O Axxouee
Let us use these pieces to show 2 U. bibt (without loss of generality,
0 that BASECONV(u\[ + dout 1] u where a).
Recall that (I) 0. Then indeed,
BASECONv(u)]:::: doot 1] (u. W), weccdmi ① (Keoddou (scotdau--)) 1):
Plugging in values get
0
0
0
0
0
0
0
O
0
W
(b.c)
0
1 dont 1]
0
0
W (b-t -oodoot--)
0
0
0
0
0
0
0
W(b.e)
Wbbtoou -.c--ut---c-ou-t
e: c dout - 1] + O Nxdoue
*
N -1xdout
Let's define
W (b.o)
0
0
D def
0
0
0
W (B++oou Loctdoue-))
73

---

Published as a conference paper at ICLR 2025
Then say
0
0
0
0
W (b.c)
0
be: (.c:cdou-11 U.b
0
0
0
W(b Wo-doue l.cddou-
0
:
0
0
0
Also note that
0
0
W Wib.e)
K :ootdouu--1 # U actdour- U, dootduru
Whbdoout
-.eetdout-1 D-1
which gives us
(uW & * uKuueeee c+ dont - 1] = (U D) O (U D-1)
--bbddau UI ce-
dob U. a:a
Where in the last equality we used the fact that a = c. Therefore, we have shown that the gradients of the expected loss function is O,
BASECONV(u)]:.::c daut] do -1 ④ bbbbb-ooo
as desired. For the case of с b, the proof remains the same, just values for and are swapped where necessary.
A corollary of the above is that MULTIPPLY implements the SQUARE function
Corollary D.63. Given Assumptions D.45, D.46, D.48, and a function
f(u) ②
Let O, be such that, EVoLee-0 0 with and dout d. Then BASECONv(u, 0, fu))
Proof. The proof follows when we have values e and dout d for the Theorem D.62.
We now revisit the importance of Assumption D.48. Specifically, the following definition is stronger version of the complement Assumption D.48. The following essentially states that there are many ways get the expected gradients of the loss function be though this doesn't 't imply that we have learned the exact solution, as we recover in Corollary D.63.
Definition D.64. Define (assumption) to be
(2) = 0
For all j E [d], either
(i) W. W 0 and Kn
74

---

Published as a conference paper at ICLR 2025
( K and W j.
The following theorem is to emphasize, there are many ways to get expected value of the gradients of the loss function to be O.
Theorem D.65. Let satisfy Definition D.64. Then, EVoL o-0- 0 when f(u) = u O u (where
0 and dout = d).
Proof. This proof considers values of j E [d]. Via Lemma D.58 and Lemma D.61 when k > 0 we have
aT.
N-1d- d-d1
E
KEJ
Σ
Σ
W W
u3,e us
[K..
P=0
This expected value goes to zero since every column j either k, or W., is O. When k we have
aL oKojj
N [utt
( W - 1) W
This expected value goes to zero since in (i) and (ii) from Definition D.64, Ko.; = Wis Next have for all L.j where L j,
Wes Σ E [už , (K [už BB
E
OW
This expected value goes zero since column either K.j W.j is O. Then when f j we have
N-1 = Σ
W
ju K
[u Ko, [ub
OW j.j
===0
This expected value goes to zero since in (i) and (ii) from Definition D.64, Ko.; Wjj =0
And we know that since (2) is all zeros, we don't t need to consider the gradient of the loss function it.
What the above proves is that there are infinite instantiations of parameters such that the expected gradient loss is However note that in Definition D.64, for all j either Kx.; for & + 0 or We, for é e j are unconstrained. In other words, we can set these values arbitrarily, which means get L. - inf but we still have the expected gradient loss to be O. This shows that some form of Assumption D.48 is necessary prove Corollary D.63.
75


Reviews:
[{"review_id": "cMKt7T3dGC", "reviewer": "Reviewer_frTT", "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.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n\n2. 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.\n\n3. It is unclear whether the proposed setup can generalize to practical scenarios beyond simple linear regression.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gKVMDMG0W9", "reviewer": "Reviewer_TtTT", "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).", "questions": "n/a", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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).\n\n- Theoretical and experimental claims are well-substantiated. The experiments are thorough and well-executed within the scope of the problem.\n\n- 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:\n\n- 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.\n\n- 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.\n\n-  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.\n\n(Minor) Missing related works:\n\n[1] Zhang et al., Trained Transformers Learn Linear Models In-Context, JMLR 2024 -- Like Ahn et al. 2024, shows that linear transformers learn GD.\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1Nswc7ho3I", "reviewer": "Reviewer_XPf9", "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.", "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?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)