Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression
Authors: Deqing Fu
Tian-qi Chen
Robin Jia
Vatsal Sharan
Venue: NeurIPS 2024 poster (2024)
Keywords: transformers
in-context learning
linear regression
Primary area: interpretability_and_explainability

Abstract: Transformers excel at *in-context learning* (ICL)---learning from demonstrations without parameter updates---but how they do so remains a mystery. Recent work suggests that Transformers may internally run Gradient Descent (GD), a first-order optimization method, to perform ICL. In this paper, we instead demonstrate that Transformers learn to approximate second-order optimization methods for ICL. For in-context linear regression, Transformers share a similar convergence rate as *Iterative Newton's Method*, both *exponentially* faster than GD. Empirically, predictions from successive Transformer layers closely match different iterations of Newton’s Method linearly, with each middle layer roughly computing 3 iterations; thus, Transformers and Newton’s method converge at roughly the same rate. In contrast, Gradient Descent converges exponentially more slowly. We also show that Transformers can learn in-context on ill-conditioned data, a setting where Gradient Descent struggles but Iterative Newton succeeds. Finally, to corroborate our empirical findings, we prove that Transformers can implement $k$ iterations of Newton's method with $k + \mathcal O(1)$ layers.

Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression
Deqing Fu
Tian-Qi Chen Department of Computer Science University of Southern California
Robin Jia Vatsal Sharan
{deqingfu, tchen939, bus.ijiaassaaaaa/u.s.c.... edu
Abstract
Transformers excel in-context learning (ICL)- from demonstrations without parameter updates- how they do so remains a mystery. Recent work suggests that Transformers may internally run Gradient Descent (GD), first-order optimization method, to perform ICL. In this paper, we instead demonstrate that Transformers learn to approximate second-ordes optimization methods for ICL. For in-context linear regression, Transformers share similar convergence rate as Iterative Newton's Method; both are exponentially faster than GD. Empirically, predictions from successive Transformer layers closely match different iterations Newton's Method linearly, with each middle layer roughly computing 3 iterations; thus, Transformers and Newton's s method converge at roughly the same rate. In contrast, Gradient Descent converges exponentially more slowly. We also show that Transformers can learn in-context on ill-conditioned data, a setting where Gradient Descent struggles but Iterative Newton succeeds. Finally, to corroborate our empirical findings, we prove that Transformers can implement l iterations of Newton's method with & + O(1) layers.
Introduction
Transformer neural networks [Vaswani al., 2017] have become the default architecture for natural language processing [Devlin et al., 2019, Brown et al., 2020, OpenAl, 2023]. As first demonstrated by GPT-3 [Brown al., 2020], Transformers excel at in-context learning (ICL) from prompts consisting of onput-output pairs, without updating model parameters. Through in-context learning, Transformer-best based Large Language Models (LLMs) can achieve ofate-of-the-ar few-shot performance across variety of downstream tasks [Rae et al., 2022, Smith et al., 2022, Thoppilan al., 2022, Chowdhery et al., 2022].
Given the importance of Transformers and ICL, many prior efforts have attempted to understand how Transformers perform in-context learning. Prior work suggests Transformers can approximate various linear functions well in-context [Garg et al., 2022]. Specifically to linear regression tasks, prior work has tried understand the ICL mechanism, and the dominant hypothesis is that Transformers learn in-context by running optimization internally through gradient-based algorithms [von Oswald et al., 2022, 2023, Ahn et al., 2023, Dai al., 2023, Mahankali et al.. 2024].
This paper presents strong evidence for a competing hypothesis: Transformers trained to perform in-context linear regression learn strategy much than a first-order method like Gradient Descent (GD). In particular, Transformers approximately implement second-order method with convergence rate very similar to Newton-Schule Method, also known as the Iterative Newton's 's Method, which iteratively improves an estimate of the inverse of
similar to second-order optimization method
Our codes are available at http:://tttttbboooo//o/gggggataa
com
-icl- secl-secon---e order.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

the data matrix to compute the optimal weight vector. Across many Transformer layers, subsequent layers approximately compute more and more iterations of Newton's 's Method, with increasingly better predictions: both eventually converge to the optimal minimum-nove -orm solution found by ordinary least squares (OLS). Interestingly, this mechanism is specific to Transformers: LSTMs do not learn these same second-order methods, as their predictions do not even improve across layers.
We present both empirical and theoretical evidence for our claims. Empirically, Transformer layers demonstrate similar rate of convergence to the OLS solution as second-order methods such as Iterative Newton, which is substantially faster than the rate of convergence of GD (Figure 2). The predictions made by the Transformer at successive layers closely match the predictions made by Iterative Newton after a proportional number of iterations, showing that they progress in similar ways the same rate. In contrast, to match the Transformer's predictions after A layers, GI would have to run for exponential in many steps (Figure 3). Some individual Transformer layers make progress equivalent hundreds of GD steps: these layers must be doing something more sophisticated than GD. Furthermore, a crucial aspect of second-order methods is that they can handle ill-condttioned problems by correcting the curvature. We find that the convergence rate of Transformers is not significantly affected by ill-conditioning, which again matches Iterative Newton but not GD. To provide theoretical grounding to our empirical results, we show that Transformer circuits can efficiently implement Iterative Newton: one transformer layer can compute one Newton iteration (given (1)
processii layers), and requires hidden states of dimension (dd) for a d-dimensional linear regression problem. Overall, our work provides a mechanistic account of how Transformers perform ICI. that explains model behavior better than previous hypotheses, and hints at why Transformers are well-suited for ICL relative to other architectures.
2 Related Work
In-context learning by large language models. GPT-3 [Brown et al., 2020] first showed that Transformer-basee large language models can "learn" to perform new tasks from in-context demon- strations (i.e., input-output pairs). Since then, a large body of work NLP has studied in-context learning, for instance by understanding how the choice and order of demonstrations affects results [Lu et al., 2022, Liu et al., 2022, Rubin et al., 2022, Su et al., 2023, Chang and Jia, 2023, Nguyen and Wong, 2023], studying the effect of label noise [Min et al., 2022c, Yoo et al., 2022, Wei et al., 2023], and proposing methods to improve ICL accuracy [Zhao et al., 2021, Min et al., 2022a.b].
In-context learning beyond natural language. Inspired by the phenomenon of ICL by large language models, subsequent work has studied how Transformers learn in-context beyond NLP tasks. Garg et al. [2022] first investigated Transformers' ICL abilities for various classical machine learning problems, including linear regression. We largely adopt their linear regression setup in this work. Li et al. [2023] formalize in-context learning as an algorithm learning problem. Han et al. [2023] suggests that Transformers learn in-context by performing Bayesian inference on prompts, which can be asymptotically interpreted as kernel regression. Other work has analyzed how Transformers do in-context classificatinn [Tarzanagh et al., 2023a.b, Zhang et al., 2023], the role of pertaining data [Raventós et al., 2023], and the relationship between model architecture and ICL [Lee et al., 2023].
Do Transformers implement Gradient Descent? A growing body of work has suggested that Transformers learn in-context by implementing gradient descent within their internal representations. Akyürek al. [2022] summarize operations that Transformers can implement, such multiplication and affine transformations, and show that Transformers can implement gradient descent for linear regression using these operations. Concurrently, von Oswald et al. [2022] argue that Transformers learn in-context via gradient descent, where one layer performs one gradient update. In subsequent work, von Oswald et al. [2023] further argue that Transformers are strongly biased towards learning to implement gradient- based optimization routines. Ahn et al. [2023] extend the work of von Oswald al. [2022] by showing Transformers can learn to implement preconditioned Gradient Descent, where the pre-conditioner can adapt to the data. Bai et al. [2023] provide detailed constructions for how Transformers can implement a range of learning algorithms via gradient descent. Finally, Dai et al. [2023] conduct experiments on NLP tasks and conclude that Transformer- based language models performing ICL behave similarly to models fine-tuned via gradient descent; however, concurrent work [Shen et al., 2023b] argues that real- wolld LLMSs do not perform ICL. via gradient descent. Mahankali et al. [2024] showed that implementing gradient descent is a global minima for single layer linear self-attention. However, we study deeper models in this work, which can behave differently from

---

single-layer models. In this paper, argue that Transformers actually learn to perform in-context learning by implementing a second-order optimization method, not gradient descent!.
Mechanistic interpretability for Transformers. Our work attempts to understand the mechanism through which Transformers perform in-context learning. Prior work has studied other aspects of Transformers' internal mechanisms, including reverse-engineering language models [Wang et al., 2022], the grokking phenomenon [Power et al., 2022, Nanda et al., 2023], manipulating attention maps [Hassid et al., 2022], and circuit finding [Conmy al., 2023].
Theoretical Expressivity of Transformers. Giannou et al. [2023] provide construction of looped transformers to implement Iterative Newton's method for solving pseudo-inverse, and each Newton iteration can be implemented by 13 looped Transformer layers. In contrast, our construction needs only one Transformer layer to compute one Newton iteration.
3 Problem Setup
In this paper, we focus on the following linear regres- sion task. The task involves n examples {æ;, y: ) where 2; E Rd and y/ E R. The examples are generated from the following data generating distri- bution Pp. parameterized by distribution D over (d x d) positive semi-definite matrices. For each se- quence of n in-context examples, we first sample a
9:
Yr
Transformers
In-Context
ground-truth weight vector w* itd. I), E IR and a matrix Σ Lid. D. For i E [m], we sample each 2; i.d. (0, Σ). The label y! for each æ, is given by Y: w" H;. Note that for much of our experiments D is only supported on the identity matrix I and hence Σ = I, but we also consider some distribu- tions ill-conditioned matrices, which give rise ill-conditioned regression problems. Most of our results are on this noiseless setup and results with the noisy setup are in Appendix A.3.2.
Examples
Figure 1: Illustration of how Transformers are trained to do in-context linear regression.
3.1 Standard Methods for Solving Linear Regression
Our central research question is:
What convergence rate does the algorithm Transformers learn for linear regression achieve?
To investigate this question, first discuss various known algorithms for linear regression. We then compare them with Transformers empirically in $4 and theoretically in $5, to evaluate if Transformers more similar to first-order second-order methods. We particularly about algorithms' convergence rates (the number of steps required to reach an error).
For any time step f, let X [ the labels for all the datapoints seen so far. Note that since can be smaller than the data dimension (1) can be singular. We now consider various algorithms for making predictions for 2441 based on x and (t) When it is clear from context, drop the superscript and refer to (t) and (1) as X and y, where X and y correspond to all the datapoints seen so far.
xe] be the data matrix and y = [y]
y/ be
Ordinary Least Squares. This method finds the minimum- norm solution to the objective:
1
L(w X, y)
X
2n
The Ordinary Least Squares (OLS) solution has closed form given by the Normal Equations:
i OLS
( X)
After an initial version this paper, Vladymyrov et al. [2024] found that a variant of Gradient Descent can mimic Iterative Newton by approximating the inverse implicitly and getting second-orde rates, which also supports our claim.
3

---

where S == ===TXx and St is the pseudo-inverse [Moore, 1920] of S.
Gradient Descent. Gradient descent (GD) is first-order method which finds the weight vector uGD with initialization vi 0 using the iterative update rule:
GD 1-1
W
GD V ( X y).
It is known that GD requires O (s((S log(1/e)) steps converge an error where k(S) the condition number. Thus, when kSS) is large, GD converges slowly [Boyd and Vandenberghe, 2004].
( Arain (S)
Online Gradient Descent. While GD computes the gradient with respect the full data matrix X at each iteration, Online Gradient Descent (OGD) is an online algorithm that only computes gradients on the newly received data point (#k,, yk at step k:
OOD ensures that the new weight vector v 1k+1 makes zero error on (æk, 3k
u GOGD
Th V C(w C
OGD
Tk yk)).
Picking 7/k
Iterative Newton's Method. This is a second-order method which finds the weight vector v Newton by iteratively apply Newton's method to finding the pseudo inverse of X X [Schulz, 1933, Ben-Israel, 1965].
Mg = aS, where a = SS 11
w
Newton
MoX
(5)
M 2 MM
W W+ll
Newton
M
This computes an approximation of the psuedo inverse using the moments of S. contrast to GD, the Iterative Newton's 's method only requires Olog x(S) log log(1/e)) steps to converge to an e error [Soderstrom and Stewart, 1974, Pan and Schreiber, 1991]. Note that this exponentially faster than the convergence rate of GD. We discuss additional algorithms such as Conjugate Gradient, BFGS, and L-BFGS in the Appendix A.2.3.
3.2 Solving Linear Regression with Transformers
We will use neural network models such Transformers to solve this linear regression task. As shown in Figure 1, at time step t + 1, the model sees the first in-context examples {z,, y/ V- and then makes predictions for 2t+11 whose label y/++1 is not observed by the Transformers model.
We randomly initialize our models and then train them on the linear regression task to make predictions for every number of in-context examples /, where / E [n)]. Training and test data are both drawn from Pp. To make the input prompts contain both m; and yı, we follow same the setup as Garg et al. [2022]'s to zero-pad y/ 's and use the same GPT-2 model [Radford et al., 2019] with softmax activation and causal attention mask (discussed later in Definition 3.1).
We now present the key mathematical details for the Transformer architecture, and how they can be used for in-context learning. First, the causal attention mask enforces that attention heads can only attend to hidden states of previous time steps, and is defined as follows.
Definition 3.1 (Causal Attention Layer). A causal attention layer with M heads and activation function a is denoted as Alln on any input sequence H = [h..-.., hy] E R RDDN where the dimension of hidden states and N is the sequence length. In the vector form,
M Σ XVL 1=1
h. = [Attn (H)) h, +
a ((Qmhe, Kmitz)) Vnneg.
(6)
Vaswani et al. [2017] originally proposed the Transformer architecture with the Softmax activation function for the attention layers. Later works have found that replacing Softmax(-) with ! ReL.U(-) does not hurt model performance [Cai et al., 2022, Shen al., 2023a, Wortsman et al., 2023]. The Transformers architecture is defined by putting together attention layers with feed forward layers:

---

Transformer Errors vs, e Layers
Iterative Newton Errors v.s. @ Steps
Gradient Errors
201
12
Dne
20-1
10-1
30-2
bat pnd
Be
23 Coitee cur saaata
F Coaratee 31
Toaefomer Lager indea
traite Newter Stape
Gostere Boskert Sheas
(a) Transformers
(b) Iterative Newton's 's Method
(c) Gradient Descent
Figure 2: Convergence of Algorithms. Similar to Iterative Newton and GD, Transformer's perfor- mance improve over the layer index E. When > the Transformer model, from layers 3 to demonstrates a superlinear convergence rate, similar to Iterative Newton, while GD, with fixed step size, is sublinear. Later layers of Transformers show a slower convergence rate, and we hypothesize they have little incentive to implement the algorithm precisely since the error is already very small. A 24-layer Transformer model exhibits the same superlinear convergence (Figure in $A.4.2).
Definition 3.2 (Transformers). An L-layer decoder-based transformer with Causal Attention Layers denoted as TFa and is a composition of MLP Layer (with skip connection) and Causal Attention Layers. For input sequence ( the transformers f-th hidden layer is given by
TH (0))
MLP (Attngca (HH(-11)). attn
0(4)
where 0 (0
( 12 mlp att and (0) (Q (0)
K!
10 m==
has M heads layer L.
m
In particular for the linear regression task, Transformers perform in-context learning as follows
Definition 3.3 (Transformers for Linear Regression). Given in-context examples { (æ ( Vi...... Bt }. Transformers make predictions query example 2 through a readout layer parameterized as Preadout = uu, v). and the prediction TTF is given by
TF Di+1
ReadOut, T ((21, Vr," Tts Dr, 8c+ +1))
uuu H (L) .2+1 +
To compare the rate of convergence of iterative algorithms to that of Transformers, we treat the layer index É of Transformers as analogous the iterative step k of algorithms discussed in 1..1. Note that for Transformers, we need to re-train the ReadOut layer for every layer index l so that they can improve progressively (see $4.1 and for experimental details) for linear regression tasks.
3.3 Measuring Algorithmic Similarity
We propose two metrics to measure the similarity between linear regression algorithms.
Similarity of Errors. This metric aims to measure similarity of algorithms through comparing prediction errors. For a linear regression algorithm A, let A(#x+1 {æx, I1 3i) denote prediction on the (t + 1)-th in-context example 24+1 after observing the first t examples (see Figure 1). We write A(#+41) = A(x (x,, ys ) for brevity. Errors (i.e., residuals) on the sequence are:?
E(A (#, Vi ) = A(#22) y2, A(#n+1) yn+t)
The similarity of errors for two algorithms A. and As is the expected cosine similarity of their errors on a randomly sampled data sequence:
SimE,An, Ab)
E ( Pp
the indices start from to because evaluate all cases where i can choose from 1,
5

---

Here C(u, v) PD the data generation process discussed previously.
the cosine similarity, is the total number of in-context examples, and
Similarity of Induced Weights. All standard algorithms for linear regression estimate a weight vector τò. While neural ICL models like Transformers do not explicitly learn such weight vector, similar to Akyürek et al. [2022], we can induce an implicit weight vector tù learned by any algorithm A by fitting weight vector to its predictions. We can then measure similarity of algorithms by comparing the induced w. To do this, for any fixed sequence of in-context examples (æt, Mi -=) iid. we sample > query examples Øs. ~ (O, E), where k E [T]. For this fixed sequence of in-context examples {æz, Me we create 7' in-context prediction tasks and use the algorithm A to make predictions A( { Pi We define the induced data matrix and labels as
A((11
X
Y
TT
A(Tr (2o,, Me MiHl
The induced weight vector for A and these t examples is:
き (A) (AA {z,, v.lii)) = (
(8)
The similarity induced weights between two algorithms A,, and A. is the expected average cosine similarity" of induced weights ui; (Aa) and wü (AAb) over all possible < l S on randomly sampled data sequence:
Sim SimW(AR ((a,, An)
E
c (ai Di w. (Abl (1)1))))))))
Matching steps between algorithms. Each algorithm converges to its predictions after several steps for example the number of iterations for Iterative Newton and GD, and the number of layers for Transformers (see Section 4.1). When comparing two algorithms, given a choice of steps for the first algorithm, we match it with the steps for the second algorithm that maximize similarity.
Definition 3.4 (Best-matching Steps). Let M be the metric for evaluating similarities between two algorithms Aa and Ab. which have steps Pa E [0, Ta] and E [0, To], respectively. For given choice of Pa. define the best-matching number of steps of algorithm As for Aa as:
Pb (pa) == arg max M(Aa (: Pp)..A(( pt)). PoE(O.TT]
In our experiments, we chose Ta. Th be large enough integers so the algorithms converge. The matching processes can be visualized as heatmaps as shown in Figure 3, where best-matching steps highlighted. This enables two algorithms converge at the same rate, the best matching steps between the two algorithms should follow linear trend. We will discuss these results in $4. See Figure 26 on how best-matching steps help compare the convergence rates.
to compare the rate of convergence of algorithms. In particular, if
Experimental Evidence
We primarily study the Transformers-bases GPT-2 model with 12 layers and 8 heads per layer. Alternative configurations with fewer heads per layer, or with more layers, also support our findings; I and later
we defer them $A.4.1 and 4...22 We initially focus on isotropic cases where consider ill-conditioned Σ in $4.3. Our training setup is exactly the same Garg et al. [2022]] models are trained with most n 40 in-context examples for 20 (with the same learning rate, batch size etc.).
We claim that Transformers learn high-order optimization methods in-context. We provide evidence that Transformers improve themselves with more layers in $4.1; Transformers share the same rate of convergence Iterative Newton, exponentially faster than that of GD, in $4.2; and they also perform well on ill illconditioned problems in $4.3. Finally, we contrast Transformers with LSTMs in $4.5.
Alternttiv metrics such as &a distance gives the same observation. Here cosine similarity is better since emors usually have small magnitudes, and directions of induced weights meaningful.

---

Similarity Errors
Transformet Layur Index
Similarity Errors (Tranfformers Gradlent Descent)
Tracstomer Layor 18,00x
.000
912
392
1.50
200
250
.214
3100
850
Wetiin Stoge
610
30
054
11
32
14
.580
750
3922
8000
15
.401
850
000
15
3411
20
2133
395
22
363
903
003
Figure 3: Heatmaps of Similarity. The best matching steps are highlighted in yellow. Transformers layers show linear trend with Iterative Newton steps but exponential trend with GD. This suggests Transformers and Iterative Newton have the same convergence rate that exponentially faster than GD. See Figure 10 for additional heatmap where GD's steps are shown in log scale: that plot there is linear correspondence between Transformers and GD's steps. This further strengthens the claim that Transformers have exponentially faster rate of convergence than GD.
4.1 Transformers improve progressively over layers
Many known algorithms for linear regression, including GD, OGD, and Iterative Newton, are iterative: their performance progressively improves as they perform more iterations, eventually converging to a final solution. How can a Transformer implement such an iterative algorithm? von Oswald et al. [2022] propose that deeper layers of the Transformer may correspond to more iterations; in particular, they show that there exist Transformer parameters such that each attention layer performs one step of GD.
Following this intuition, first investigate whether the predictions of trained Transformer improve the layer index l increases. For each layer of hidden states H (see Definition 3.2), re-train
the ReadOut predict y/2 for each 1; the new predictions are given by ReadOut (5) HH))) Thus for each input prompt, there are Transformer predictions parameterized by layer index &. All parameters besides the ReadOut layer parameters are kept frozen.
As shown Figure 2(a) (and Figure 7(a) in the Appendix), as we increase the layer index E, the prediction performance improves progressively. Hence, Transformers progressively improve their predictions over layers f, similar to how iterative algorithms improve over steps. Such observations consistent with language tasks where Transformers-hatt language models also improve their predictions along with layer progressions [Geva et al., 2022, Chuang et al., 2023].
4.2 Transformers are more similar to second-order methods, such as Iterative Newton
We test the specific hypothesis that the iterative updates performed across Transformer layers are similar to the iterative updates for known iterative algorithms. First, Figure 2 shows that the middle layers of Transformers converge rate similar to Iterative Newton, and faster than GD. In particular, the Transformer and Iterative Newton both converge at a superlinear rate, while GD converges a sublinear rate.
Next, we analyze whether each layer the Transformer corresponds to performing i steps of some iterative algorithm, for some & depending on f. We focus here GD and Iterative Newton's Method; we will discuss online algorithms Section 4.5, and additional optimization methods in Appendix A.2.3. We will discuss results noisy linear regression tasks in Appendix A.3.2.
For each layer & of the Transformer, we measure the best-matching similarity (see Def. 3.4) with candidate iterative algorithms with the optimal choice of the number of steps k. As shown in Figure 3,

---

107
Convarponce Ill-Conditooed
Ynanforee
15TM
Tranaformene
strpe
Mottoed
Jarabee Nowttn
Meeg Sapangg Sree
10
10--
10.
10 107
104
14 15
CoineConnett Caurgiir
Timee Stamp
Figure 4: Transformers perfor- mance on ill-conditioned data. Given 40 in-context examples, Transformers and Iterative New- ton converge similarly and they both can converge to the OLS so- lution quickly whereas GD suf- fers.
Figure 5: In the left figure, we measure model predictions with normalized MSE. Though LSTM is seemingly most similar to Newton's 's Method with only 5 steps, neither algorithm converges yet. OGD also has similar trend as LSTM. In the right figure, we measure the model's error rate on example Zn-g after seeing n examples, for different values of the time stamp gap y (see Appendix A.6), and find both Transformers and not-converged Newton have better memorization than LSTM and OGD.
the Transformer has very high error similarity with Iterative Newton's method all layers. Moreover, we see a clear linear trend between layer 3 and layer 9 of the Transformer, where each layer appears to compute roughly 3 additional iterations of Iterative Newton's method. This trend only stops the last few layers because both algorithms converge to the OLS solution; Newton is known to converge to OLS (see 1), and we verify in Appendix A.2 that the last few layers of the Transformer also basically compute OLS (see Figure 14 in the Appendix). We observe the same trends when using similarity of induced weights as our similarity metric (see Figure the Appendix). Figure 11 in the Appendix shows that there is similar linear trend between Transformer and BFGS, an alternative quasi-New method. This is perhaps not surprising, given that BFGS also gets a superlinear convergence rate for linear regression Nocedal and Wright [1999]. Thus, we do not claim that Transformers specifically implement Iterative Newton, only that they (approximately) implement some second-order method.
In contrast, even though GD has comparable similarity with the Transformers at later layers, their best matching follows an exponential trend. As discussed in the Section 3.1, for well-conditioned problems where 1, to achieve error, the rate of convergence of GD is (log((//e) while the rate of convergence of Iterative Newton is Olog log(1/))). Therefore the rate of convergence of Iterative Newton is exponentially faster than GD. Transformer's 's linear correspondence with Iterative Newton and its exponential correspondence with GD provides strong evidence that the rate of convergence of Transformers is similar to Iterative Newton, i.c., ([l([ee)(L/e)) We also note that is not possible to significantly improve GD's convergence rate without using second-order methods: Nemirovski and Yudin [1983] showed a 52 log((/c)) lower bound on the convergence rate of gradient- based methods for smooth and strongly convex problems, and Arjevani et al. [2016] shows a similar lower bound specifically for quadratic problems. In the Appendix, we show that limited-memory BFGS Liu and Nocedal [1989] and conjugate gradient (see Figure 12), which do not use full-second order information, also converge slower than Transformers. This provides further evidence for the usage of second-order information by Transformers. We also show more evidence by investigating alternative function classes such as linear regression with noises in Appendix A.3.2 and 2-layer neural network with ReLU or Tanh activation function in Appendix A.3.3.
Overall,
we conclude that a Transformer trained to perform in-context linear regression learns to implement an algorithm that is very similar to second-order methods, such as Iterative Newton's 's method, not GD. Starting at layer 3, subsequent layers of the Transformer compute more and more iterations of Iterative Newton's 's method. This algorithm successfully solves the linear regression problem, as it converges to the optimal OLS solution the final layers.
4.3 Transformers perform well on ill-conditioned data
We repeat the same experiments with data a;; (O Σ) sampled from an ill-condition covariance matrix Σ with condition number x(E)= = 100. and eigenbasis chosen uniformly at random. The first
lid
8

---

Transformers Hidden
Transrormers (Hidden Sze=10)
Transforers (Hldddnn Ste-122
Transforees (Hidden Suwm64)
0.8
Luast Squares
0.6
90.4
6.2
0.0
10
15
20
35
In-contoxt examples
Figure 6: Ablation on Transformer's Hidden Size, For linear regression problems with d =20, Transformers need (d) hidden dimension to mimic OLS solutions.
d/2 eigenvalues of Σ are 100, and the last d/2 are 1. Note that choosing the eigenbasis uniformly at random for each sequence ensures that there is a different covariance matrix Σ for each sequence of datapoints.
As shown in Figure 4, the Transformer model's s performance still closely matches Iterative Newton's Method with 21 iterations, same when Σ I (see layer 10-12 in Figure 3). The convergence of second-order methods has a mild logarithmic dependence on the condition number since they correct for the curvature. On the other hand, GD\'s convergence is affected polynomially by conditioning. As H² E) increase from to 100, the number steps required for GD's convergence increases signifi- cantly (see Fig. 4 where GD requires 2,000 steps to converge), making it impossible for 12-layer Transformers to implement these many gradient updates. We also note that preconditioning the data by (X X)1 can make the data well-conditioned, but since the eigenbasis is chosen uniformly at random, with high probability there is no sparse pre-conditioner or any fixed pre-conditioner which works across the data distribution. Computing (X Xx1 appears to be as hard as computing the OLS solution (Eq. 1)-in in fact Sharan et al. [2019] conjecture that first-order methods such as gradient descent and variants cannot avoid polynomial dependencies in condition number in the ill-conditioned case." See Appendix A.3.1 for detailed experiments on ill-conditioned problems. These experiments further strengthen our thesis that Transformers learn perform second-order optimization methods in-context, not first-order methods such as GD.
4.4 Transformers Require (dd) Hidden Dimension
We ablate 12-layer 1-head Transformers with various hidden sizes on 20 problems. As shown in Figure 6, we observe that Transformers can mimic OLS solution when the hidden size is 32 or 64, but fail with smaller sizes. This resonates with our theoretical results on (dd) hidden dimension in Theorem 5. and in this case, the theorem ensures construction of transformers to implement Iterative Newton's 's method.
4.5 LSTM is more similar to OGD than Transformers
As discussed in $A.1, LSTM is an alternative auto-regressive model widely used before the introduc- tion of Transformers. Thus, natural research question is: If Transformers can learn in-context, can LSTMs do so as well? If so, do they learn the same algorithms? To answer this question, we train a LSTM model in an identical manner to the Transformers studied in the previous sections.
Figure plots the error of Transformers, LSTMs, and other standard methods number of in-context (i.e., training) examples provided. While LSTMs can also learn linear regression in-context, they have much higher mean-squared error than Transformers. Their error rate is similar to Iterative Newton's 's Method after only 5 iterations, point where it is far from converging to the OLS solution. Finally, we show that LSTMs behave more like an online learning algorithm than Transformers. In particular, its predictions are biased towards getting more recent training examples correct, as opposed to earlier examples, as shown in Figure 5. This property makes LSTMs similar to
function of the
'Regarding preconditioning, we also note that even for collicondt instances preconditioned GD still gets linear rate of convergence, whereas Transformers and Iterative Newton get superlinear rates.

---

online GD. In contrast, five steps of Newton's method has the early examples, showing that the LSTM implements a very different algorithm from a few iterations of Newton. We hypothesize that since LSTMs have limited memory, they must learn in a roughly online fashion; in contrast, Transformer's attention heads can access the entire sequence of past examples, enabling it to learn more complex algorithms. See $A.1 for more discussions.
average for recent and
5 Theoretical Justification
Our empirical evidence demonstrates that Transformers behave much more similarly to Iterative Newton's than GD. Iterative Newton is second-order optimization method, and is algorithmically more involved than GD. We begin by first examining this difference in complexity. As discussed in Section 3, the updates for Iterative Newton are of the form,
W k++
Newton M X y
where M 2 M,SM
(10)
and My = aS for some a > 0. We can express Mx in terms of powers of S by expanding iteratively, for example M1- 2a5 4a 5 M2 4a
2 + 2 S 16a²5", and in general
M2- = - 1=1 11 3, for some 3. E R (see Appendix B.3 for detailed calculations). Note that k steps of Iterative Newton's requires computing S2(2*) moments of S. Let us contrast this with GD. GD updates for linear regression take the form,
GD W 1++1
Wk GD
( X y).
(11)
Like Iterative Newton, we can express CD terms of powers of and X y. However, after steps of GD, the highest power of S is only (k)). This exponential separation is consistent with the exponential gap in terms of the parameter dependence in the convergence rate O (s(SS log(1/c)) for GD vs. Olog (S) + log (o))) for Iterative Newton. Therefore, natural question is whether Transformers can actually as complicated of a method such as Iterative Newton with only polynomially many layers? Theorem 5.1 shows that this is indeed possible.
Theorem 5.1. For any there exist Transformer weights such that on any set of in-context examples {æi, Vi Ji and test point Prest, the Transformer predicts on Ptost using D 20, Newton Here Newton are the Iterative Newlon updates given by i Newton = M, y where Mj is updated as i
Mj -2 -
S i Mo =o,
for some a > and = X The dimensionality of the hidden layers (d), and the number of layers is + 8. One transformer layer computes one Newton iteration. 3 initial transformer layers are needed for initializing Mo and 5 layers the end are needed to read out predictions from the computed pseudo-inverse M
We note that our proof uses full attention instead of causal attention and ReLU activations for the self-attention layers. The definitions of these and the full proof appear in Appendix B.
Conclusion and Discussion
In this work, we studied how Transformers perform in-context learning for linear regression. In contrast with the hypothesis that Transformers learn in-context by implementing gradient descent, our experimental results show that different Transformer layers match iterations of Iterative Newton linearly and Gradient Descent exponentially. This suggests that Transformers share a similar rate of convergence to Iterative Newton but not to Gradient Descent. Moreover, Transformers can perform well empirically on ill-conditioned linear regression, whereas first-order methods such as Gradient Descent struggle. This empirical evidence - when combined with existing lower bounds in optimization - suggests that Transformers use second-order information for solving linear regression, and we also prove that Transformers can indeed represent second-order methods.
An interesting direction is to explore wider range of second-order methods that Transformers can implement. It also seems promising to extend our analysis to classification problems, especially given recent work showing that Transformers resemble SVMs in classification tasks [Li al., 2023, Tarzanagh et al., 2023a]. Finally, a natural question is to understand the differences in the model architecture that make Transformers better in-context learners than LSTMs. Based on our investigations with LSTMs, we hypothesize that Transformers can implement more powerful algorithms because of having access to a longer history of examples. Investigating the role of this additional memory in learning appears to be intriguing direction.
10

---

Acknowledgement
We would like to thank the USC NLP Group and Center for AT Safety for providing compute resources. DF would like thank Oliver Liu and Ameya Godbole for their extensive discussions. DF and RJ were supported by a Google Research Scholar Award. RJ was also supported by an Open Philanthropy research grant. VS supported by NSF CAREER Award CCF--2332265 and an Amazon Research Award.
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. ArXiv, abs/2306.00297,, 2023. 1,
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. ArXiv, abs/2211.156611 2022. 2, 3.3, B.1, B.4, B.2, B.5
Yossi Arjevani, Shai Shalev-Shwarte and Ohad Shamir. On lower and upper bounds in smooth and strongly convex optimization. Journal of Machine Learning Research, 17(126):1-51, 2016. URL http://jam jttpr./mmmrreeer//// wrg. wtgp/papers/v1715-11115
htmml. 4.2
Yu Bai, Fan Chen, Haiquan Wang. Caiming Xiong, and Song Mei. Transformers statisticians: Provable in-context learning with in-context algorithm selection. ArXiv, abs/2306.04637, 2023.
Adi Ben-Israel. An iterative method for computing the generalized inverse of an arbitrary matrix.
Mathematics of Computation, :9452) 455, 1965. ISSN 00255718, 10886842. URL http: //www. jstor org/ htt//saabee22003//// 3.1
Stephen P Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004. 3.1
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jef- frey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Ben- jamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc., 2020. URL https: // http:///pooce/nneeeeeeeee// neurips iierippssco cc/ paper paper_filest files httpa///pppppp/po///oo/////////a 2
2a--222444PPPPPea paaper.pde .dd. 1,
Han Cai, Chuang Gan, and Song Han. Efficientvit: Enhanced linear attention for high-resolution
wwwcoomptttiion visual recognition. ArXiv, abs/2205. 2022. 3.2
Ting- Yun Chang and Robin Jia. Data curation alone can stabilize in-context learning. In Proceedings of the 61 Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 8123-8144, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 12.28 accloo/o0-... URL https: https////aaaaatog/////o org/
org/2023. 2022 acl-long.
452. 2
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi. Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev. Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowyez, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan

---

Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways, 2022.
Yung-Sung Chuang, Yujia Xie, Hongyin Luo, Yoon Kim, James Glass, and Pengcheng He. Dola: Decoding by contrasting layers improves factuality in large language models, 2023. 4.1
Arthur Commy, Augustine N. Mavor- Parker, Aengus Lynch, Stefan Heimersheim, and Adrià Garriga- Alonso. Towards automated circuit discovery for mechanistic interpretability. 2023. 2
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models secredly perform gradient descent meta-optimizers. ArXiv, abs/2212.10559, 2023. 1,2
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies. Volume (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/ URL https: hatps////caaathhlggg// www.ollloooooo.g/i..... 14233
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes. ArXiv, abs/2208.01066, 2022. 1, 2, 3.2, A.3.3, D
Mor Geva, Avi Caciularu, Kevin Wang, and Yoav Goldberg. Transformer feed-forward layers build predictions by promoting concepts the vocabulary space. In Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, pages 30-45, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 1/.23/2211/2022. emnlp-main.3. URL https: http:///claaahhooooo org/202/ 2022. comlp-main.. 3. 4.1
Angeliki Giannou, Shashank Rajput, Jy-Yong Sohn, Kangwook Lee, Jason D. Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Pro- ceedings of Machine Learning Research, pages 11398-11442. PMLR, 23-29 Jul 2023. URL https:/ http///poceeeeeeee/ mlr, press/v202/ http:///22//g/an222/// giannou23a. html. 2, B.5
Chi Han, Ziqi Wang, Han Zhao, and Heng Ji. In-context learning of large language models explained as kernel regression, 2023.
Michael Hassid, Hao Peng, Daniel Rotem, Jungo Kasai, Ivan Montero, Noah A. Smith, and Roy Schwartz. How much does attention actually attend? questioning the importance of attention in pretrained transformers, 2022. 2
Sepp Hochreiter and Jürgen Schmidhuber. Long Short-Term Memory. Neural Computation, 9 (8):1735-1780 11 1997. ISSN 0899-7661 7667. doi: wwtp://nccc.199..988. URL https:
//doi. org/o0/ 10. 11w1//eeo...oee /ww.. 1997 9. 1735. A.1
Ivan Lee, Nan Jiang, and Taylor Berg- Kirkpatrick. Exploring the relationship between model architecture and in-context learning ability, 2023.
Yingcong Li, Muhammed Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, 2023. 2, 6
Dong C. Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical Programming, 45:503-528, 1989. URL https: : /appi. semant i cscholar. org/ CorpusID: 5681609. 4.2
12

---

Jiachang Liu, Dinghan Shen, Yizhe Zhang, Bill Dolan, Lawrence Carin, and Weizhu Chen. What makes good in-context examples for GPT-3? In Proceedings of Deep Learning In- side Out (DeeLIO 2022): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, pages 100-114, Dublin, Ireland and Online, May 2022. As- sociation for Computational Linguisties. doi: 10.1865///202..deiioll.l00 URL https: //aclanthologe. org/ 202./22222ca2 deelie-1.10.0...11 2
Yao Lu, Max Bartolo, Alastair Moore, Sebastian Riedel, and Pontus Stenetorp. Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 8086-8098, Dublin, Ireland, May 2022. Association for Computational Linguisties. doi: 10. 1116666//w1/2222allllloggg
URL https::ppp. htw.s//hoolaoalooooooooooo oxg/2022.oo. acl-lcng 1on...56......
Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In The Twelfth International Conference on Learning Representations, 2024. URL https: h http///pppeeeeeeeeeeeeeeeeeeeeet net/ wwte/foummmeeeeemeeeeee forum?ide 8p3fu561Kc. 1,2
Sewon Min, Mike Lewis. Hannaneh Hajishirzi, and Luke Zettlemoyer. Noisy channel language model prompting for few-shot text classification. In Proceedings of the 60th Anmual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pages 5316-5330, Dublin, Ireland, May 2022a. Association for Computational Linguistics. doi: 10. 10.18653/11/2222.acl-00 365. URL https: https////caaaathooooooo org/2022..c. sel-long 365. 2
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. MetalCL: Learning to learn in context. Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages States, July 2022b. Association for Computational Linguistics. doi: 10.18653/v 1/263///12022.naac main. 201. URL https: wtttttss//aaaaatoo//o///ooo
Seattle, United
2222//222.cr naacl main. 201. 2
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zeulemoyer. Rethinking the role of demonstrations: What makes in-context learning work? In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 11048-11064, Abu Dhabi, United Arab Emirates, December 2022c. Association for Computational Linguistics. doi: 1111666/11222222mmmmmaaaa emnlp-man maillpmmai.... 759.
https::: /awaaaaaacaaaaca.caaaaaaaaoo.ala or2/2022.
E.H Moore. On the reciprocal of the general algebraic matrix. Bulletin of American Mathematical Society, 26:394-395, 1920. 3.1
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability, 2023.
A.S. Nemirovski and D.B Yudin. Problem complexity and method efficiency in optimization. 1983. 4.2
Tai Nguyen and Eric Wong. In-context example selection with influences. arXiv preprint arXiv:2302. 11042, 2023. 2
Jorge Nocedal and Stephen J Wright. Numerical optimization. Springer, 1999. 4.2, A.2.3
OpenAl. Gpt-4 technical report, 2023. URL URLLittp://aaxdv..or/// abs http://aabb223300777777/ 1
Victor Y. Pan and Robert S. Schreiber. An improved newton iteration for the generalized inverse of a matrix, with applications. SIAM J. Sci. Comput., 12: 109-1130, 1991. 3.1
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, perforaefrr deep learning library, 2019. D
13

---

Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Mísra. Grokking: General- ization beyond overfitting on small algorithmic datasets, 2022. 2
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. models are unsupervised multitask learners. 2019. 3.2
Language
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican. Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring. Susannah Young, Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George van den Driessche. Lisa Anne Hendricks, Maribeth Rauh, Po-Sen Huang, Amelia Glaese, Johannes Welbl, Sumanth Dathathri, Saffron Huang. Jonathan Uesato, John Mellor, Irina Higgins, Antonia Creswell, Nat McAleese, Amy Wu, Erich Elsen, Siddhant Jayakumar, Elena Buchatskaya, David Budden, Esme Sutherland, Karen Simonyan, Michela Paganini, Laurent Sifre, Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro, Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki Lazaridou, Arthur Mensch, Jean-Baptiste Lespiau, Maria Tsimpoukelli, Nikolai Grigorev, Doug Fritz, Thibault Sottiaux, Mantas Pajarskas, Toby Pohlen, Zhitao Gong, Daniel Toyama, Cyprien de Masson d' Autume, Yujia Li, Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy. Chris Jones, James Bradbury, Matthew Johnson, Blake Hechtman, Laura Weidinger, lason Gabriel, William Isaac, Ed Lockhart, Simon Osindero, Laura Rimell, Chris Dyer, Oriol Vinyals, Kareem Ayoub, Jeff Stanway, Lorrayne Bennett, Demis Hassabis, Koray Kavukcuoglu lu, and Geoffrey Irving. Scaling language models: Methods, analysis & insights from training gopher, 2022.
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression, 2023. 2
Ohad Rubin, Jonathan Herzig, and Jonathan Berant. Learning to retrieve prompts for in-context learning. Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 2655-2671, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022.nnaa http:///222222.mmmaiaa 191. URL https: http//cccaannoooggg.gooooo org maacl-main, 191. 2
Günther Schulz. Iterative berechung der reziproken matrix. Zeitschrift fürr Angewandte Mathematik und Mechanik (Journal of Applied Mathematics and Mechanics), 13:57-59, 1933. 3.1
Vatsal Sharan, Aaron Sidford, and Gregory Valiant. Memory-sample tradeoffs for linear regression with small error. In Proceedings of the 5 Annual ACM SIGACT Symposium on Theory of Computing, pages 890-901, 2019, 4.3
Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, and Jiang Bian. A study on relu and softmax in transformer, 2023a. 3.2
Lingfeng Shen, Aayush Mishra, and Daniel Khashabi. Do pretrained transformers really learn in-context by gradient descent?, 2023b. 2
Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, Elton Zhang, Rewon Child, Reza Yazdani Aminabadi, Julie Bernauer, Xia Song, Mohammad Shoeybi, Yuxiong He, Michael Houston, Saurabh Tiwary, and Bryan Catanzaro. Using deepspeed and megatron train mieatrron-tring nlg 530b, large-scale generative language model, 2022. 1
Torsten Soderstrom and G. W. Stewart. On the numerical properties of an iterative method for computing the moore- penrose generalized inverse. SIAM Journal on Numerical Analysis, 11(1): 61-74. 1974. ISSN 00361429. URL http: h/wwww http://ww.j.ttt...ooooo jstor . org/ sttsaa/eeeeee533///// 3.1
Hongjin Su, Jungo Kasai, Chen Henry Wu, Weijia Shi, Tianlu Wang, Jiayi Xin, Rui Zhang, Mari Ostendorf, Luke Zettlemoyer, Noah A. Smith, and Tao Yu. Selective annotation makes lan- guage models better few-shot learners. In The Eleventh International Conference on Learning Representations, 2023. URL ptatttss:/opprreeee
net weet/forumiddggIhlll
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak, Transformers as support vector machines. ArXiv, abs/2308. 168999 2023a. 2, 6
14

---

Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism, 2023b.
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng- Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, YaGuang Li, Hongrae Lee, Huaixiu Steven Zheng. Amin Ghafouri, Marcelo Menegali, Yanping Huang, Maxim Krikun, Dmitry Lepikhin, James Qin, Dehao Chen, Yuanzhong Xu, Zhifeng Chen, Adam Roberts, Maarten Bosma, Vincent Zhao, Yanqi Zhou, Chung-Ching Chang. Igor Krivokon, Will Rusch, Mare Pickett, Pranesh Srinivasan, Laichee Man, Kathleen Meier-Hellstern, Meredith Ringel Morris, Tulsee Doshi, Renelito Delos Santos, Toju Duke, Johnny Soraker, Ben Zevenbergen, Vinodkumar Prabhakaran, Mark Diaz. Ben Hutchinson, Kristen Olson, Alejandra Molina, Erin HoffmannJohn. Josh Lee, Lora Aroyo, Ravi Rajakumar, Alena Butryna, Matthew Lamm, Viktoriya Kuzmina, Joe Fenton, Aaron Cohen, Rachel Bernstein, Ray Kurzweil, Blaise Aguera- Arcas, Claire Cui, Marian Croak, Ed Chi, and Quoc Le. Lamda: Language models for dialog applications, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, edi- tors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, neurips.. cc/ cw/paper.ppp..//.ee/e wwt.ess/pppee///////iiii..
Inc., 2017. URL https: : httts::/proceeediii 4f5e22454444e99.f000333c.d.a.a Paperr pdf. 1, 3.2
Max Vladymyrov, Johannes von Oswald, Mark Sandler, and Rong Ge. Linear transformers are versatile in-context learners, 2024.
Johannes von Oswald, Eyvind Niklasson, E. Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, 2022. 1, 2, 4.1
Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Blaise Agüera Arcas, Max Vladymyrov, Razvan Pascanu, and Joao Sacramento. Uncovering mesa-optimization algorithms transformers. ArXiv, abs/2309.05885, 2023. 1,2
Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Inter- pretability in the wild: circuit for indirect object identification in gpt-2 small, 2022. 2
Jerry Wei, Jason Wei, Yi Tay, Dustin Tran, Albert Webson, Yifeng Lu, Xinyun Chen, Hanxiao Liu, Da Huang, Denny Zhou, and Tengyu Ma. Larger language models do in-context learning differently, 2023. 2
Mitchell Wortsman, Jaehoon Lee, Justin Gilmer, and Simon Kornblith. Replacing softmax with relu in vision transformers, 2023. 3.2
Kang Min Yoo, Junyeob Kim, Hyuhng Joon Kim, Hyunsoo Cho, Hwiyeol Jo, Sang- Woo Lee, Sang-goo Lee, and Taeuk Kim. Ground-truth labels matter: A deeper look into input-label Conference on Empirical Methods in Natural Language Processing, pages 2422-2437, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10.18653/v 12555//112222..lnmmmaall attps://acllanthology wrg.ey..rrr/eee www/ emm.pmmmmlmmmmm..mmm 155. 2
demonstrations. In Proceedings of the 2022
URL
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context.
ArXiv, abs/2306.09927 2023. 2
Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving few-shot performance of language models. In International Conference on Machine Learning, pages 12697-12706. PMLR, 2021. 2
15

---

Appendix
Additional Experimental Results
16
A.1
Contrast with LSTMs
16
A.2
Additional Results on Isotropic Data without Noise
17
A.2.1
Progression Algorithms
17
A.2.2
Heatmaps
17
A.2.3
Comparison with Other Second-Order Methods
20
A.2.4
Additional Results on Comparison over Transformer Layers
22
A.2.5
Additional Results on Similarity of Induced Weights
22
A.3
Varying Data Distribution or Function Class
23
A.3.1
Experiments on Ill-Conditioned Problems
23
A.3.2
Experiments with Noisy Linear Regression
25
A.3.3
Experiments with Non- Function Class (2-Layer MLP)
25
A.4
Varying Transformer Architecture
27
A.4.1
Experiments Transformers of Fewer Heads
27
A.4.2 Experiments on Transformers with More Layers
28
A.5 Heatmaps with Best-Matching Steps Help Compare Convergence Rates
29
A.6
Definitions for Evaluating Forgetting
29
B Detailed Proofs for Section 5
30
B.1 Helper Results
30
B.2 Proof of Theorem 5.1
B.3 Iterative Newton Sum of Moments Method
34
B.4 Estimated weight vectors lie in the span of previous examples
................................
35
C Computes
36
D License
36
Limitations
36
Broader Impacts
36
A
Additional Experimental Results
A.1 Contrast with LSTMs
While our primary goal is to analyze Transformers, we also consider LSTMs [Hochreiter and Schmidhuber, 1997] to understand whether Transformers learn different algorithms than other neural sequence models trained to do linear regression. In particular, we train unidirectional L-layer LSTM, which generates a sequence of hidden states ( for each layer E, similarly to an L-layer Transformer. As with Transformers,
add readout layer that predicts the 2!-TTM from the final 224+1
hidden state at the final layer, H (L)
16

---

Transformers
LSTM
Newton
0.991
0.920
GD
0.957
0.916
OGD
0.806
0.954
Table 1: Similarity of errors between algorithms. Transformers are more similar to full-obser oollevbbevvtnoon methods such as Newton and GD; and LSTMs are more similar to online methods such as OGD.
We train 10-layer LSTM model, with 5.3M parameters, in identical manner the Transformers (with 9.5M parameters) studied in the previous sections.
LSTMs' inferior performance to Transformers can be explained by the inability of LSTMs to use deeper layers to improve their predictions. Figure shows that LSTM performance does not improve across layers-a readout head fine-tuned for the first layer makes equally good predictions as the full 10-layer model. Thus, LSTMs seem poorly equipped to fully implement iterative algorithms. Similarly, Table shows that LSTMs more similar to OGD than Transformers are, whereas Transformers are more similar to Newton and GD than LSTMs.
A.2 Additional Results on Isotropic Data without Noise
A.2.1 Progression of Algorithms
Errors vs. In-Context Examples
Errors v.s. In-Contert Examples
Ervors vss In-Context Examples
Desikmmmeee
becatos Frotoo end 0
10TV
Temkmmlee
Nintes eit
1575
maskkmriiee
Roota
Tucebhe Kromo add 1eee
150
Traskmmm
Tumehe Nonton e10
Tranirmmi Lone
(app))
L0TV Zenre
(a) Transformers
(b) Iterative Newton's 's Method
(c) LSTM
Figure 7: Progression of Algorithms. (a) Transformer's ' performance improves over the layer index f. (b) Iterative Newton's performance improves the number of iterations k, in way that closely resembles the Transformer. We plot the best-matching to Transformer's l following Definition 3.4.
(c) In contrast, LSTM's performance does not improve from layer to layer.
A.2.2 Heatmaps
We present heatmaps with all values of similarities.
'While the LSTM has fewer parameters than the Transformer, we found in preliminary experiments that increasing the size of the LSTM would not substantively change results.
17

---

Stodlarity
Newtoo)
Trgaformer Logur
Index
Sundiviit
Gradient Descent)
Trinformer Lague
Index
1000
9.20
.16
716
.692
.534
.233
.618
.913
.620
.916
nia
12
270
.792
.522
.60
855
.551
.558
577
341
822
.228
.33
205
.745
.713
.000
.694
.230
.895
.922
100
5931
2.43
583
4131
961
2881
.789
.75
241
2735
732
235
733
1.0
531
531
6779
781
871
835
873
974
972
873
6970
9.09
733
239
101
206
7833
777
774
277
774
240
535
524
.072
.722
.553
.93
-24
.777
J376
274
.974
.274
591
4991
15.0
.580
1224
817
.514
250
521
521
6887
7677
830
932
924
.224
.778
8777
.977
.977
R32
450
11
1952
1155
RS2
854
6832
300
350
538
018
538
518
.664
703
850
832
922
925
.77
073
3101
390
.951
-90
J50
979
.977
.500
.979
oo Stegie
615
877
835
911
885
866
B85
8B6
888
850
936
972
36.3
A81
.98
3991
700
310
1211
9115
9131
91111
.982
.58
508
723
826
907
856
8:12
6936
5134
1935
911-
552
703
.988
.953
.550
02
.539
528
688
792
882
969
976
800
1965
.922
1962
1472
511
512
545
844
.923
.023
3544
355
354
.284
533
.77
.973
7141
511
521
1522
594
922
202
213
3114
385
24
.524
.528
969
.771
265
JB0
SRS
3160
.9788
.975
.9/13
750
511
510
1653
702
843
925
960
302
965
885
965
885
810
-510
510
.852
.751
.442
.019
.065
J122
-355
35-4
.985
.985
15
.547
-517
.969
.760
452
.227
S428
J1BS
JA67
.965
.982
115
900
510
.100
510
510
.652
6552
.700
750
8911
.411
915
915
5668
960
.62
9177
-966
.01.5
963
.966
915
945
.366
17
M50
.399
3.9
.852
.750
.581
925
S122
-258
515
18
1533
.313
055
754
846
924
973
R18
S92
361
.93
1991
5:58
917
967
966
886
-96
.986
20
511
511
053
751
342
820
970
RI7
6R3
563
9R3
9921
21
531
511
902
528
.39
.215
916
J8L
J86
D55
2200
200
22
510
510
098
748
A3B
817
967
RM4
8R1
563
993
043
1250
530
530
6650
700
R33
916
300
474
23
.58
508
.60
1300
538
508
.630
7668
838
915
560
281
355
055
0100
05/
Figure 8: Similarity of Errors. The best matching steps are highlighted in yellow.
Similunity Induced Weight (Euundrruurs va. Berative Newioul
Traedornue
Luqur Indes
Trastoreer Lages 2ndCX
.000
.001
800
611
742
.199
714
711
711
711
713
712
9
12
.000
.00L
.700
260
.760
.760
004
0772
671
354
.001
J00L
65/0
8190
.616
3026
595
205
1928
1.00
.030
895
3710
371
.71
971
.00L
.877
340
845
845
845
845
150
636
835
712
861
2:1
973
471
478
873
893
J00
1016
8F/7
2177
200
820
805
740
858
921
.057
075
076
.75
073
.770
.075
.005
D04
-000
000
9226
3222
3123
S11
J921
901
850
250
300
J020
.020
818
105
005
005
746
740
747
106
854
655
.17
.215
910
.254
355
075
300
.75
975
377
.76
576
374
575
874
3576
.077
876
3776
.077
876
D05
0001
780
231
9-11
RE77
9715
010
1210
D05
917
954
975
377
375
3972
.00G
-000
936
.4.7
2447
805
853
975
377
977
DOW
.001
760
81
431
95A
R55
4555
955
155.5
Terre rebleee Ten tm
.000
10000
706
961
.961
.50
.006
.000
750
862
923
.53
.066
.955
JO65
.964
.965
.945
J220
805
835
J744
.62
.15
992
.953
975
.77
377
377
3/8
700
636
A25
74h
R61
+15
PP3
677
876
870
878
.006
.000
764
658
.000
J00
.060
.063
.963
.370
.970
750
.020
405
.745
.61
.314
.S52
.97
3'5
.77
276
.778
.78
J220
D05
745
151
.214
JS3
.977
.F77
3776
370
18
.006
.000
791
851
414
S4R
267
SP1
.172
.01
.772
.072
960
850
.020
820
055
800
.745
743
861
A51
.314
314
.U52
802
.97
.77
.75
.77
.77
576
876
.978
.778
879
950
.020
095
.743
851
.314
.022
.0P1
.75
.77
576
.078
.78
18
.007
008
739
J948
912
SEC
966
9XX0
47.
9072
.975
.974
4230
035
745
214
37%
.77
978
970
19
.000
973
895
977
378
673
23
.007
.082
736
J948
911
.45
965
800
373
903
9774
974
21
.000
973
095
51.
22
.007
.000
738
.848
310
.044
.65
.70
.733
373
974
574
1250
-230
895
.B.0
213
993
975
3/3
23
1300
818
825
742
913
651
475
A76
Figure 9: Similarity of Induced Weight Vectors. The best matching steps are highlighted in yellow.
18

---

Similarity of Errors (Transformers v.s. Gradlent Descent)
Trasformer" Layer Index
3
6
10
11
12
953
953
.800
.771
.692
.645
.620
.609
605
599
606
.602
.910
.910
.903
.826
.750
.703
.676
.665
660
655
.661
.677
.728
.759
.759
.886
.903
.876
.846
.220
807
801
.788
.801
.799
Gradient Steps
16
.678
.677
.831
.895
.910
.903
.886
873
867
865
867
865
32
.610
.610
.768
.858
.914
1938
.934
.924
.918
.16
.917
.916
64
.563
.563
.717
.817
.897
.917
.961
.954
.950
.918
.948
.917
256
.521
.666
.766
.858
.932
.973
.979
.977
.977
.977
512
.513
.513
.856
.755
.847
.923
.971
.982
.984
.922
.883
.983
1024
.509
.509
.522
.749
.810
.917
.967
.982
.986
.985
.86
986
2048
-507
.507
.648
.745
.836
.913
.964
.980
.966
.986
.987
.977
.666
Figure 10: Similarity Errors of Gradient Descent in Log Scale. The best matching steps are highlighted in yellow. Putting the number of steps of Gradient Descent in log scale further verifies the claim that Transformer' rate of covergence exponentially faster than that of Gradient Descent.
19

---

A.2.3 Comparison with Other Second-Order Methods
In this section, we ablate with alternative second-order methods, such Conjugate Gradient, BFGS, and its limited memory variant, L-BFGS.
Conjugate Gradient Method. For linear regression problems, the Conjugate Gradient (CG) method solves the linear system
X w
X y
CG finds the weight vector D with initialization we by maintain set of conjugate gradient (Aw),-'' J. ). It follows the iterative update rule
d2 V(((wk)
k-1
d SAw, Aw,
Aw = d2
Aw. SAw;
(12)
==0
0g arg min L(ws + adm)))
wk++1 wk + ag Awk
The conjugate Gradient method requires O (V log(1 1/c) steps to converge to an error on quadratic objectives such as linear regression.
BFGS. Broyden- Fletcher-Gooldfrrb--hiee approximate the inverse Hessian B₃ = (²((w))) The BEGS updates are given by
(BFGS) is Quasi-Newton method, designed to
Wx+1 = wh or Bg VL(ww))
(13)
where
8% Wh+1 WE
y% vL(Wk+1) v[(wk)
BRYKY B
8k.8k
B&+1 B
yk Bayk
Uk 8
When & is large, Bj: approximates the inverse Hessian well.
L-BFGS is menited-memory version of BFGS. Instead of the inverse Hessian BR, L-BFGS maintains a history of past m updates (where m is usually small). Recall the iterative update rule of Bş in BFGS
Kiimited-memory)-Bree
Buyky Bu
8k
B B
(14)
U BuYK
y 8
Unlike BFGS, which recursively unroll to an initialization Br, L-BEGS only unroll to Ba-m but replacing Bt-m with Binit In this regard, running steps of L-BFGS only requires (mm) memory, which more memory-efficied
than BFGS who requires (n²²) memory. The trade-off is that L.- BFGS won't 't have good estimate of the inverse Hessian when < d, where d is the dimensionality of the quadratic problem. In this regard, will converge slower than full BFGS.
In Figure 11 and Figure 12, we compare Transformers with BFGS, L-BFGS, and Conjugate Gradient method on the metric of similarity of errors. We find that Transformers have similar linear correspondence with BFGS. This is perhaps not surprising, given that BFGS also gets superlinear convergence rate for linear regression Nocedal and Wright [1999]. Meanwhile, Transformers show substantially faster convergence rate than L-BFGS and CG.
20

---

Stndleitty Droen Crandereen Mlpo
Stalarry (TrNe Remarr LaMgn
Laqur
156
F75
126
623
177
-886
679
-16
1812
136
-27
2747
743
-264
795
149
1.23
103
4.35
8.38
30.9
8333
924
462
294
833
101
-037
1312
A22
720
120
676
2776
813
672
372
372
515
-779
-223
841
021
001
223
0.33
333
J0L
167
111
224
1325
436
171
2.56
281
.224
232
781
255
311
276
576
J74
16
110
491
182
500
114
16
333
1311
262
4.5
264
-384
16
567
-31
467
183
160
-300
J82
-48
231
18
264
433
133
-881
314
J34
402
30:
233
397
492
161
824
P4
111
161
10
282
253
748
846
330
371
516
-36
423
133
850
-30
317
26
201
25
100
300
519
340
315
349
J84
331
PV
15
511
386
467
137
R51
13
373
283
0131
55:
27
289
467
146
012
414
.811
PSC
21
221
4664
540
40
411
556
463
789
Km
612
16
114
10
637
151
545
173
36
146
-20
243
745
830
412
36
J84
.322
-66
-214
14
516
-64
464
133
845
-04
472
J4B
231
34
340
457
185
J84
PRC
-4
14
546
4641
546
10
163
2141
-50
345
264
.36
266
26
-83
230
235
433
395
722
145
39
286
-26
463
785
8:30
816
363
J84
.B62
244
-6
-01
30
535
-64
462
361
846
AG
21:
815
(a) Transformer v.s. BFGS
(b) Transformer V.s. L-BFGS
Figure 11: Similarity of Errors between Transformers and BEGS or L-BFGS. The best matching steps highlighted in yellow. We find that Transformer, from layers 6 to 11, has linear correspon- dence with BFGS. For L-BFGS, due to its limited memory, it approximates second-order information more slowly and results in a slower convergence rate than Transformers.
Stullarity Emore (Tnaadmmsss Conjegata Gradieed
(ransloemaa Layer Indars
30
11
12
.441
.41
E56
707
71
<1p
-775
270
-665
E70
-657
.033
11250
[1:
-772
.118
E19
-107
885
-11
927
024
188
181
173
160
477
-51
565
91.4
925
932
033
430
501
F31
502
352
-32
-711
816
1314
-1.4
-250
-147
J564
243
344
100
01:1
1900
5400
562
1952
10
334
136
185
105
255
1122
-315
-211
.250
.207
200
322
180
134
575
1460
352
.200
.297
.008
575
2012
424
1771
700
1774
483
941
67H
.877
676
3320
1572
366
142
122
-26
.961
262
1992
116
218
1773
172
1925
2-11
312
365
.811
10
515
3355
1957
1020
1992
67W
482
3m
462
162
2114
-31
1995
270
481
5m
300
.907
21
512
1313
1921
1524
R76
461
308
332
1992
350
.200
.289
520
244
310
-330
1983
100
D33
.315
.250
.221
.203
.25
.00
310
1993
319
.291
.297
680
27
340
230
1035
130
1833
D31
.218
.250
.252
.091
50:
308
153
1812
530
375
-200
.252
-02
10
345
.308
1057
1833
515
.217
-550
.258
.192
.12
.01
307
158
1850
512
337
.255
.292
-12
360
155
186
1810
R7B
416
360
363
15
505
106
655
154
1945
R72
876
389
360
691
561
36
366
186
151
154
1966
R37
489
891
18
385
186
66
163
8M1
R76
915
389
641
A91
345
4044
347
928
315
.100
.26
104
204
Figure 12: Similarity of Errors between Transformers and Conjugate Gradient. Transformer's convergence rate is still faster than conjugate gradient methods.
21

---

A.2.4 Additional Results Comparison Transformer Layers
Coxine Similarity
Cosine Similarity
Steff meaeeemrrn
MokMhaateemr StoMrTaartemmn 6E51
Stef Cnartmmann 6651
0.
Index
Layer Index
(a) Similarity of Errors
(b) Similarity of Induced Weights
Figure 13:
Similarities between Transformer and candidate algorithms. Transformers resemble Iterative Newton's ' Method the most.
A.2.5 Additional Results Similarity of Induced Weights
We present more details line plots for how the similarity of weights changes the models more in-context observations {ær, ) } i.e., as n increases. We fix the number of Transformers layers and compare with other algorithms with their best-match steps to É in Figure 14.
Layer
Layer
Layer
Golll dnci
Coeetro Stimn wntty
Chadto Steer untty
Contine Stmslaltt
0.4
00
02
13 Number Contert tommples
Number Context Examples
Number in Content Enamples
Figure 14: Similarity of induced weights over varying number in-context examples, on three layer indices of Transformers, indexed as 2, 3 and 12. We find that initially layer 2, the Transformers model hasn't 't learned so it has zero similarity to all candidate algorithms. As we progress to the next layer number 3, we find that Transformers start learn, and when provided few examples, Transformers are more similar to OLS but soon become most similar to the Iterative Newton's Method. Layer 12 shows that Transformers in the later layers converge to the OLS solution when provided more than example. We also find there is dip around n = for similarity between Transformers and OLS but not for Transformers and Newton, and this probably because OL,S has more prominent double-descent phenomenon than Transformers and Newton.
22

---

A.3 Varying Data Distribution Function Class
A.3.1
Experiments on Ill-Conditioned Problems
In this section, we repeat the same experiments as we did on isotropic data in the main text and in Appendix A.2, and change the covariance matrix to be ill-conditioned such that Ai [) 100.
Errors v.s. In-Context Examples
Errors v.s. # In-Context Examples
Herza 1G zaps
Taraferrorr eop
Larafomarr Ligar
heethe Newre 203 vom)
Dusth Nerza #641 exgo
Tarafferars Ligar a00
baath Nersa #66113 depo
Taraferaar Liger eon
Nersa
Tarafeeaars Ligar eop
baath Nersa ecs dept
Ianafeeaar Lager
913
Hoothe Newvo Meptl
In-Contest frrempiss
(a) Transformers
(b) Iterative Newton's Method
Figure 15: Progression of Algorithms on Ill-Conditioned Data. Transformer's performance still improves over the layer index &; Iterative Newton's 's Method's performance improves over the number of iterations and we plot the best-matching / to Transformer's é following Definition 3.4.
We also present the heatmaps find the best-matching steps and conclude that Transformers are similar to Newton's 's method than GD in ill-conditioned data.
Sindlurity
Trueoomer Loper Indas
Sumillantt Errres Cirusstomers Gradicat DesceoD
Tranfurmar Laque index
.E00
322
.733
.208
337
.335
519
325
532
.332
.522
750
401
5
462
413
416
15
4i0
13
413
12
1.00
-200
553
-016
D61
200
100
100
317
IEC
.735
-736
2.4
A18
.733
479
.055
830
.342
433
352
4030
200
461
4.1
835
164
522
327
327
129
433
433
051
472
343
.443
593
180
745
A13
A:8
783
777
775
775
777
775
668
600
415
413
307
075
266
.262
.262
493
672
414
603
B16
16-5
945
1911
700
410
A15
A38
354
721
201
601
401
375
207
670
.20
6G1
3E8
.266
.285
Met Strges
21-2
464
445
615
820
561
455
603
8:55
11.3
5.14
111
.574
274
-225
12:00
.011
A33
533
771
501
.072
400
578
.E75
.355
375
.272
422
425
230
820
.963
540
.22-
2133
1300
110
436
572
774
H13
805
597
FON
336
22
415
.14-
420
211
.274
278
.277
.775
23
411
412
276
784
337
.003
526
-224
2114
615
60N
S76
14
-07
.944
406
400
257
770
J77
.71
570
D13
.35-
D4S5
S422
1860
1760
.157
-377
318
337
396
357
170
513
513
813
.372
.075
.23
.82
.311
462
.273
.200
461
.21
.211
.26
.262
18
.000
400
233
755
.00
D75
274
S74
17
207
382
5.58
750
ST4
S7M
22141
-34
354
201
J055
J502
103
18
.944
394
232
755
S02
D13
2E2
202
23.0
-164
392
554
755
201
.005
206
J602
J82
383
.244
19
382
2400
1160
233
543
785
807
483
106
193
JFS
485
J61
29
.30
.30
34
J448
200
9:4
.99
D25
JB5
267
100
.355
334
207
.053
380
.204
.355
285
246
167
22
357
188
345
J48
883
947
.92
DRS
.163
254
-80
105
2000
2800
.252
332
393
351
543
782
N61
203
601
J96
.83
463
280
465
K01
.24
444
225
222
46C
3000
334
392
801
401
Figure 16: Similarity of Errors on Ill-Conditioned Data. The best matching steps are highlighted in yellow.
23

---

Simolanity Induccd
Beratioo
Tegalomnar
7
Index
Sinilarity of Induced Welight (rannfecuees U.S. Trunformrr Laqur
Gradieot Descont
.003
.033
.023
0224
609
778
723
7747
.721
.550
.026
.044
.017
.054
.515
.633
.612
.908
1.00
1000
-039
.052
4030
-925
.222
-317
.213
.346
.533
294
207
201
250
5
264
-79
10
-200
2//
13
-24
.287
12
000
4228
600
834
.805
.726
.019
.685
.987
.883
.979
200
638
011
611
624
F06
836
830
0710
-038
001
.088
-350
.943
.441
.02
422
664
65
495
.864
.779
.751
.740
.728
.754
.729
500
1330
002
589
1:54
102
A10
.775
(850
600
1838
833
555
7018
8923
197
J50
35/4
J877
J170
D02
022
15
915
882
505
777
7155
764
760
715
700
8141
.018
1338
10.
013
.577
1.7
-25
772
.386
5300
.532
-230
.00
391
200
.884
-225
863
.886
.85
J822
netetnn Stage
D02
022
R10
R27
420
R06
820
782
71B
F78
77%
7710
11:00
0037
016
330
654
265
.224
895
.08
913
.908
905
307
.01
000
7BB
.187
.783
1210
1837
493
-559
1024
200
.924
.9.0
.113
.211
212
.90
929
.02
0221
R21
415
917
8317
81.0
795
732
7154
1:300
037
014
120
100
630
219
910
920
474
812
A12
AE1
12
.033
509
.133
812
.092
.023
502
800
310
.930
868
.946
.854
.532
.823
.525
1650
0032
004
546
150
8:15
913
925
918
800
.022
579
.556
830
1710
1937
015
:544
123
8:22
9T1
421
823
R00
500
.002
.021
562
768
873
.07
.22
.878
.977
.876
.673
1800
1814
-037
1038
.815
015
1582
-500
.744
-211
.339
1:31
.509
.207
.11
.927
.928
.92
-211
921
3rd
923
.211
.02
M29
16
0992
.021
.544
.44
249
914
-1P3
.534
.195
17
1902
1120
828
722
826
905
SI09
SG0
911
2200
0037
005
015
1544
656
724
820
831
.902
81.
.31
.20
926
.25
823
836
927
18
.092
-020
.155
.886
JM22
J225
9115
.P18
3124
.121
2310
1037
.535
/22
.227
.A01
.009
J2L
.927
328
19
.921
.018
505
699
S2N
975
922
9014
2400
0377
015
131
720
825
600
300
930
828
R27
828
979
20
.00
018
488
680
781
A63
861
831
626
52B
993
813
21
9021
0018
492
77%
SQN
935
915
3700
003
035
520
818
959
X11
22
.000
015
488
.006
766
B48
880
826
926
525
835
616
2000
2000
.037
1837
485
815
.525
-27
712
713
810
817
.855
FM
.07
S00
924
R31
950
916
824
023
051
R112
002
23
.003
.85
.25
3000
032
015
1524
210
814
191
331
Figure 17: Similarity of Induced Weights on Ill-Conditioned Data. The best matching steps are highlighted in yellow.
Similaatty Errors& (Trandda rooer v.a Blgs)
Similasity of Erom ('andommn v.S. L-Bigs)
Lay
Layer Index
1.0
13
12
7.44
227
825
825
21.2
876
874
650
705
80
768
743
133
732
731
725
729
728
122
347
200
407
822
817
S1G
814
513
813
913
536
528
833
619
815
R0B
806
R06
484
471
886
855
KM
HBS
BR1
MRD
473
548
E7S
613
883
BTB
876
878
870
100
B17
565
430
382
/33
10.4
225
948
425
523
743
975
1276
577
475
.775
376
433
4027
269
1975
295
265
D65
268
416
585
730
AAA
RED
BR7
9417
887
SAT
430
400
761
850
3400
8860
576
475
872
R71
877
BFC5 Stupp
413
382
J25
2011
427
200
735
M12
./18
V/2
41.
078
731
BEL
593
1954
394
423
750
882
962
1885
384
R33
861
360
280
341
MER
872
.055
.03
409
526
291
.2%
294
408
075
728
857
835
S78
880
862
853
363
634
387
417
437
745
873
585
667
367
887
M47
408
325
2/7
390
.91
295
418
385
280
100
202
200
260
305
408
575
721
550
055
277
.251
283
393
.03
204
414
736
055
201
291
200
.289
265
230
450
408
575
737
856
336
891
9/2
483
821
584
413
340
A65
591
8M1
2A0
291
408
322
222
800
036
51G
983
350
853
362
853
384
412
661
862
347
681
561
003
682
38
862
563
408
324
225
155
836
SMP
sen
950
382
850
284
412
380
234
B83
436
201
892
962
891
488
374
225
555
057
476
D85
.099
292
302
453
204
431
579
1983
345
063
201
2013
203
J02
408
074
736
855
837
976
888
850
363
851
384
411
478
730
B12
944
662
992
(a) Transformer V.S. BFGS
(b) Transformer V.s. L-BFGS
Figure 18: Similarity of Errors on IllConditioned Data with Quasi-Newton Methods. The best matching steps are highlighted in yellow. Transformer also matches BFGS linearly, from layers 4 to 11. L-BFGS still suffers due to its limited memory but still better than Gradient Descenthecause L-BFGS also attempts to approximate second-order information.
24

---

A.3.2 Experiments with Noisy Linear Regression
We repeat the same experiments on noisy linear regression tasks with y wo x + where e NO0, with noise level a = 0.1. As shown in Figure 19, Transformers still show superlinear convergence on noisy linear regression tasks. Since the predictor is u = (xT AT) xTx y for some A, the iterative newton's ' method is applied to = X X + AI. Iterative Newton's 's method still keeps the same superlinear convergence rates. As it's also shown in Figure 19, Transformers and Iternative Newton's ' rates match linearly, as in the noiseless linear regression tasks.
Transtormer Errors v.s. Layers
10-1
10
InCantest Exampee
10-2
I-Cantest Exmmpee
I-Cantest Esmm
Co-ott Examp
Conoet Examp
CCnoet Examp
Cinoett Fammpes
10
2
6
8
10
12
Transformer Layer Index
Sinillaitt Emnes (Tranfemoss Vos Toeatioe Newton)
6
10
11
12
Stullarit Eepee in Trrrrrnddden doemers Uo Gradient Descunt) Layar
7206
.72
.15
385
578
320
.585
1556
.594
626
518
865
281
454
658
6655
638
657
812
812
736
202
701
A92
721
.006
.595
.008
73H
601
82R
7.83
772
7011
7112
A85
922
822
465
804
818
824
.221
8:23
.822
Lammiie SREMOR Starpe
441
A15
(192
1922
582
722
.223
32.1
31.7
920
31.
5:20
3523
696
451
876
789
012
J495
335
.567
2:37
527
.652
775
.67
J778
200
357
315
2008
.258
.203
4612
363
4112
9776
9776
677
645
:6
507
.640
740
.875
455
.266
.587
.586
.567
17
072
R52
341
.00
FB55
437
472
6GD
18
19
501
.831
755
.983
.45
.383
220
321
-02
.002
2020
.02
20
499
419
736
882
844
.HI
10:
563
2016
HF1
22
490
637
704
651
478
491
592
28
497
826
783
867
446
360
910
9-1
293
5-92
.453
404G
.626
730
410
6G0
461
Figure 19: Experiment results on Noisy Linear Regression. (Top) Transformers have superlinear convergence rate. (Bottom) Transformers match Iterative Newton's 's rate and are exponentially faster than Gradient Descent.
A.3.3 Experiments with a Non-Linear Function Class (2-Layer MLP)
To extend experiments to non-linear cases, adopt the same 2-layer ReLU neural network studied by Garg et al. [2022]: see Fig. 5(c) in their paper. For any prompt (x1. V/. ... De, Yt.. instead of generating labels y/k = w x as mainly studied in the paper, we study 2-layer neural network function class parameterized by W E R Rddiddenn xd so that
v R Rdaidder,
R Rddidden, and b R.
V/ fw. = a. RRLU (Wax + +0) +b
(15)
Then we repeat the same probing experiments as in the main paper. As shown in Figure 20, even on 2-layer neural network tasks with ReLU activation, Transformer shows superlinear convergence rates. Transformer shows an exponentially faster convergence rate than Gradient Descent's, 's because
25

---

Gradient Descent's steps are shown in log scale and the trend is linear similar to Figure 9 in the main paper.
Transformer Errors # Layers
Staalurtty Exquu Classicann Guadeal DINOMD lides
Triniterrer
312
244
.71
125
Stags
2B03
182
A91
181
184
M47
141
461
Errors
.63
5-13
112
204
777
770
752
Emtuiam doveel
H41
11.4
2013
508
424
188
183
.334
.13
254
733
7754
774
7768
073
131
143
M2
945
141
Tansfomer Layer Index
Figure 20: Empirical Results on 2-Layer Neural Network Regression with ReLU activation function. Transformers have superlinear convergence rates and match Gradient Descent's convergence rate exponentially
would be interesting to ablate the activation funetion used in Equation (16). We further consider the case when it's using the Tanh activation instead of ReLU, i.e.
y/ fw (xr))) a Tanh Wxx x + u) +
(16)
Repeating the same experiments before, as shown in Figure 21, we find that Transformers use entire first 5 layers to pre-process and then only the next few layers show exponentially faster convergence rate compared to Gradient Descent. We further note that in both Figure 20 and Figure 21, the cosine similarities between Transformers and Gradient Descent are significantly lower than the experiments with linear regression tasks. This might due to the paramcterization of the function class and Transformers and Gradient Descent may arrive at different optima.
Transformer Errors v.s. # Layers
10+4.
Erros
Dog
J73
373
.20
339
262
322
-72
722
.43
.762
7758
2009
930
886
831
P04
X18
830
K14
K17
333
2007
.377
8.06
855
280
.777
875
2075
313
870
854
243
542
.717
7777 717
7717
738
351
885
2224
352
800
830
243
11.44
J62
440 664 6662
20-4
316
1
10
11 32
2382
618 615
2002
405
Tansforer Layer Index
3754
A2A
614
610
2072
500
Figure 21: Empirical Results on 2-Layer Neural Network Regression with Tanh activation function. Transformers have superlinear convergence rates and match Gradient Descent's convergence rate exponentially
would be interesting for future research explore further this function class of 2-layer MLP to understand fully how Transformer solve the regression problem in-context and whether it achieves a different optimum compared alternative algorithms such as (Stochastic) Gradient Descent.
26

---

A.4 Varying Transformer Architecture
A.4.1 Experiments on Transformers of Fewer Heads
In this section, we present experimental results from an alternative model configurations than the main text. We show in the main text that Transformers learn second-order optimization methods in-context where the experiments are using a GPT-2 model with 12 layers and 8 heads per layer. In this section, we present experiments with a GPT-2 model with 12 layers but only I head per layer.
Simdlurit
Tradformer Lager
Index
Sumillrit Eresey Cirusffrmnta Gradicat Desceal
Traafomer Lager Intex
9908
261
785
J077
871
667
331
625
.514
M02c
[540)
5
630
8
A15
10
13
.b44
12
-128
190
307
0858
.741
750
7225
-026
.702
.894
1.50
539
539
.053
713
.226
341
223
.973
.911
.73
7.33
793
844
872
806
867
602
.110
788
7784
776
200
532
532
598
1555
705
619
928
455
9777
.775
250
528
528
1583
.600
700
813
802
335
954
874
.970
.078
.054
.054
715
777
A20
910
31.
800
864
B61
BGO
892
310
300
525
532
525
522
.551
1570
.337
635
3027
188
510
R07
280
FR7
3236
930
-3:3
310
275
875
.900
900
.5/9
980
Hommtte Beneton Stoge
.224
.924
900
327
924
912
594
.003
600
020
520
.57
632
602
894
685
931
J51
576
.981
.962
.590
5148
711
780
9210
9154
1520
M12
M422
630
515
538
391
803
184
933
35:
377
5776
757
867
320
947
947
.941
.442
950
876
9833
.558
5.50
1.76
7316
S430
12
3:00
602
725
917
45.B
516
79P
948
.537
5315
1951
712
9025
315F
SF/3
.878
1//22
14
530
5200
1417
1644
704
A17
910
907
9779
514
516
.020
.90
.277
.578
922
.47
.51
.985
.525
.125
.082
R06
807
380
.984
.984
516
516
185
797
177
925
247
975
985
16
17
5519
502
.318
522
.176
176
.634
806
800
.833
R55
S81
987
.999
810
9141
850
-314
513
515
515
538
514
.559
.570
550
.226
8225
675
.586
686
.785
796
785
.576
877
R76
J22
924
920
J145
.47
-46
275
076
876
.82
.62
961
3915
.085
986
18
5588
517
.573
.929
.686
.759
.880
.928
.553
.981
067
.91
1000
.914
514
.569
.224
.683
.75
.676
.23
346
576
.522
.386
19
536
.516
.772
826
.969
360
.966
992
1050
1131
533
-338
.725
9221
S46
378
380
20
515
515
.571
.928
.88
.756
.276
.024
J948
.579
.966
.222
794
675
875
886
22
513
513
569
.924
.83
.783
.722
.22
J445
.576
963
.991
513
682
794
945
1300
513
533
-229
582
921
244
981
Figure 22: Similarity of Errors on an alternative Transformers Configuration. The best matching steps are highlighted in yellow.
Stuailactty of Induced Weight IT (Crunada ETOSTS V.a Layer
Berathvv Newion)
10
Trasfomer Layer Index
1000
0003
1.56
7:50
734
717
213
710
211
712
.011
-003
100
1767
770
802
782
706
762
758
760
761
.001
0331
-377
8755
(758)
703
2015
351
475
3863
SFS1
.002
566
.311
.07
505
-206
1.00
032
.508
.02
7.66
875
153
300
953
868
.070
002
-003
580
751
787
8G8
864
850
5:44
B45
8/66
150
.032
039
-505
.637
.781
571
.800
9:9
370
.071
.313
.003
.67
796
.90
881
.878
.B76
200
833
638
500
655
738
669
90h
951
871
877
8770
.003
-002
555
716
780
802
.1.1
.915
.103
.901
.102
.902
250
.033
029
.511
.03
737
867
.27
.50
359
571
.972
.075
.003
.923
300
053
1052
029
550
4.59
86b
665
922
R06
950
12:8
97.8
372
873
11/3
.003
-003
531
4856
770
80
920
J024
3555
.S26
.957
.922
Herntton Deovin Staps
20
.002
-000
515
.966
753
880
J32
445
.42
JS54
.955
.985
019
.033
819
.800
724
.244
961
J57
576
.00
0000
106
655
741
928
917
954
963
RG4
650
600
.002
033
409
018
.46
496
.000
1550
773
700
8853
661
.944
904
J68
.57
417
072
A73
072
9771
.77
876
714)
JU33
.977
1009
1863
JS24
367
3772
311
3172
14
.000
.001
301
6550
736
R6
924
915
954
966
RGB
750
.33
618
457
608
733
603
948
357
872
4771
877
.003
.001
.500
.22
.444
.044
.967
.9G9
800
.013
.099
.457
.60
733
.982
.023
J57
572
.073
.077
16
.003
.001
409
.647
.722
.861
.221
.43
.553
.976
.972
850
1033
029
039
8952
865
923
357
357
972
9F2
3972
3772
17
.003
.001
498
842
.903
.972
850
662
957
18
.003
.001
406
.645
.731
.859
.319
J422
.952
.68
.70
.77
19
[000
.001
455
.645
-230
J858
.219
341
952
377
19/3
20
.003
.001
407
.644
.730
.858
.318
J411
.55
JE7
.77
.77
019
947
357
21
0001
.001
497
.648
-224
.188
941
951
A62
379
972
1150
1200
.03
833
019
018
496
.609
6A8
732
732
861
661
J23
903
947
847
J54
476
972
873
./3
577
472
22
000
000
697
6668
72R
910
851
SE7
970
977
1250
.03
.477
.668
.32
861
.02
347
356
072
.7%
23
0001
1001
457
.643
722
940
951
369
1500)
.033
039
.46
868
/322
863
802
347
324
272
4//2
3172
Figure 23: Similarity of Induced Weights on an alternative Transformers Configuration. The best matching steps are highlighted in yellow.
We conclude that our experimental results not restricted to specific model configurations, smaller models such as GPT-2 with 12 layers and head each layer also suffice in implementing the Iterative Newton's 's method, and more similar than gradient descents, in terms of rate of convergence.
27

---

A.4.2 Experiments Transformers with More Layers
In this section, investigate whether deeper models would behave similarly or differently. We work on Transformers with 24 layers and 8 heads each.
Simdltrity of Revses Chrandeemera yoe Beruthre Newrooa
10 11 13 15 30 18 20 21 22 24
Stantlarity Ercupe Cirmnalormer Gradient Desceab Cyaddommet
541
1.00
1.50
Staps
10
11
550
12
6000
600
7000
14
750
15
16
17
19
1100
21
1150
1200
1250
23
1300
Figure 24: Similarity of Errors on 24-layer Transformers Configuration. The best matching steps are highlighted in yellow.
Transformer Errors v.s. # Layers
10
10
Enors
10
Context Examples
In-Gontext Examples
10
In-onneet Eaamples =
In ContoxtExxoppis -
In-Context Examples
Context Examples
In Context Examples
10-4
3
11 13
15
17
19
21
23
Transformer Layer Index
Figure 25: Transformers with 24 layers also converge superlinearly, similar Iterative Newton.
28

---

A.5 Heatmaps with Best-Matching Steps Help Compare Convergence Rates
In this section, we show the heatmaps with best-matching steps among known algorithms.
Blrrative Nowtoon Stape
2- 10 12 13 14 15 30
30
120
43.0
450
550
530
620
630
12
230
210
15
17
16
1050
1130
1130
1230
21
1230
1830
(a) Iterative Newton v.s. Gradient Descent
(b) Iterative Newton v.s. Iterative Newton
Terative Nowton Stops
10
11
122
13
17
15
19
20
25
(c) Iterative Newton v.s. BFGS
Figure 26: Best-Matching Steps on Similarity of Residuals Help Compare Convergence Rates. (a: top-left) When comparing Iterative Newton and Gradient Descent, there is an exponential trend showing Iterative Newton converges exponentially faster than Gradient Descent. (b: top-right) When Iterative Newton is compared with itself in sub-figure, there is a linear trend showing they have the same convergence rate. (c: bottom) When Iterative Newton is compared BFGS in sub-figure, there linear trend after there are enough steps for BFGS to approximate second-order information showing Iterative Newton and BFGS share a similar convergence rate after sufficient BFGS steps.
A.6 Definitions for Evaluating Forgetting
We measure the phenomenon of model forgetting by reusing an in-context example within {x,, Vi 1=1 as the test example Ptest In experiments of Figure 5, we fix 20 and reuse Etest 2,, We denote the "Time Stamp Gap" the distance the reused example index from the current time stamp 20. We measure the forgetting of index as
Forgelting(A,3
E
MSE A(xx; (xt, (1).))))1))
(17)
( PD
Note: the further away i is from the more possible algorithm A forgets.
29

---

B Detailed Proofs for Section 5
In this section, we work on full attention layers with normalized ReLU activation o(-) = n2 given examples.
R.LLU.)
Definition B.1. A full attention layer with M heads and ReLU activation is also denoted as Attn on any imput sequence -.,...... hyl E DxN, where the dimension of hidden states and N is the sequence length. In the vector form,
M E
h, = (Attn(H)) h4
ReLU ((Qmh,, kmhj)) Vnhj
(18)
Remark B.2. This is slightly different from the causal attention layer (see Definition 3.1) in that at each time stamp i, the attention layer in Definition B. has full information of all hidden states j E [n], unlike causal attention layer which requires j [t].
B.1 Helper Results
We begin by constructing useful component for our proof, and state some existing constructions from Akyürek et al. [2022].
Lemma B.3. Given hidden states (h1.... h there exists query, key and value matrices Q K, V respectively such that atfention layer compute Σ}=1 hj-
h+ 1
Proof. We can pad each hidden state by 1 and O's such that h/
g2d+1
We con-
Odx
Oax1
Oaxd
struct two heads where Qi
K
Q2
O1xd Odxd
O1xd
and K₂ = -K,. Then
Odx1
O Oxxd_
O dxd
Oax1
Oaxd
[
Oxxd
1
Oxxd
11
dxd
Odx1
Odx
0g
h i
O (d41)xx
0(+1)x(d+1)) Oax(d+1)
0
Os+1
Let V.
V22
so that VT,
nlaxd
nh
O22
We apply attention layer these 1-padded hidden states and have
h. h4
ReLU ((Qmh'y Kmh;))
V
j=1
0a+1
h1j
ReLU(L) ReLU(-1)
nh;
(19)
T1
[2
ht 1 Σラ=1 hej
Os+1
=
+
hj
02
n
Proposition B.4 (Akyürek et al., 2022). Each of mov, aff, mul, div can be implemented by single transformer layer. These four operations are mappings Dx RDN expressed as follows,
mov(Hr H: s,t, i, ".j,"',"" selects the entries of the s-th column of И between rows i and j. and copies them into the t-th column (t Z s) of H between rows i' and j'.
mul (H; o. b, c, (i, ('))), (i",, j")): in each column h of H. interprets the entries between and j as b matrix A1, and the entries between i' and j' as a x matrix A2, multiplies these matrices together, and stores the result between rows i" and j". yielding a matrix in which each
column has the form hh 1. A, A2, hzy".] This allows the layer implement inner products.
30

---

div H: (1, (8", ")))): in each column h H, divides the entries between and j by the absolute value of the entry at i', and stores the result between rows i" and j", yielding a matrix in which every column has the form [h, h.hh.//
aff (H; (t.5), (C' (t.J J')), (t".3)) 3"". W1. W2. b): in each column of H, applies an affine transforma- tion to the entries between and j and and j', then stores the result between rows i" and j" (t",
yielding matrix in which every column has the form [h. W i hag + W2h + b, hgr.] This allows the layer to implement subtraction by setting W1 = I and W = --I.
B.2 Proof of Theorem 5.1
Theorem 5.1. For any k, there exist Transformer weights such that on any set of in-context examples {ar, Ut Vi and test point Frest, the Transformer predicts on Drest using Frest u Newton Here w Newton are the Iterative Newton updates given by u Newton M y where Mj is updated as
Mj = 2M 1,
M =a.
for some a > and T The dimensionality of the hidden layers is O(d), and the number of layers is k 8. One transformer layer computes Newton iteration. 3 initial transformer layers are needed for initializing Mo and layers at the end are needed to read out predictions from the computed pseudo-inved mindo-invere Mi.
Proof. We break the proof into parts.
Transformers Implement Initialization ( aS. Given input sequence H := (...... I h. with E IR we first apply the operations given by Proposition B.4 (similar to Akyürek et al. [2022], we show only non-zero rows when applying these operations):
æ1
mov
21
Zn
(20)
21
We call each column after mov as hj. With an full attention layer, one construct two heads with query and value matrices of the form Q Ki
Iaxd
Oaxa
-Q K₂
such that for any
Oaxd
Odxd
t E [n], we have
Σ ReLU ((Qmhy Kmh,)) RcLU(xx x,) ReLU(-xx, x;)) (It, 2,)
(21)
Laxd Odxd
O Odxd
Let all value matrices Vn
for R. Combining the skip connections,
O Ouxd
we have
h,
1 n
Σ
(or,
2zz) nax
x 0
x
2.
xc
I; +
axS
(22)
72
1=1
Now
the aff operator make subtractions and then
xt + aSx
aff
(x: + a.Sx.)) 2c
Taxx
(23)
2t
We call this transformed hidden states as (0) and denote (0) aS:
H
(0)
h (0)
h
(0)
wn
(24)
Notice that S is symmetric and thereafter ( is also symmetric.
Transformers implement Newton Iteration. Let the input prompt be the same as Equation (24),
(T()) x
(T)) xn
H (0) = h
(0)
h
(0)
(25)
31

---

We claim that the E's hidden states be of the similar form
T x1
T( mm
H
h
h
(E)
(26)
æ1
We prove by induction that assuming our claim is true for &, we work on l + 1:
Oa
la
Ia
Os
Let Qm
Qm
Km
Km
where Q, K1 I, K2 -I and
Oa
Oa
Oa
Ia
Od
V = V22
2-head self-attention layer, with ReLU attentions, can written has
Oa
O8
h (f+1) [Attn H ()) h1
ReLU ((emth kkmbt")) Vinhg
(27)
m=1
where
ReLU (mmmlh K Kmn/ V h
m=1
ReLU (Gh!o) (Gh!)) Q K, (Jh ReLU (G Q K2 (Jh (Jh99)
(28)
[RLLU((Gh, T (Jh, )) + ReLU(-(Gn,)) (Jh,)
(Gh(") (Jh
(Jn))) (Gh,
Plug in our assumptions that h
(9)
(x)))
we have Jh
TT(0x,
and Gh,
(0)
3 xz
0a
0.4
we have
h (++1)
(T()) cr
1
Tr(() zz
[r(e)22 xr
+
O2
Oa
1=1
T()xx: x0
(T (t) x ()) (T xx,) x
(29)
T( mz
1T(())
| (T)) T
xce
(T)) (T((S)()))
ST())
xc
Now we pass over an MLP layer with
h
(e+1)
- h
(f+1)
+
Ia
Os ht
((+1)
-
T() TT(ST
(30)
Oa
Bt
Now we denote the iteration
(f++)) (T( (1) ST ST(OT
(31)
We find that T (e++1) since ( and S are both symmetric. reduces to ((++)) (T) (1) ST(())
(32)
This is exactly the same as the Newton iteration.
Transformers can implement u TF ( x y. Going back to the empirical prompt format
{01, V1,-.. En, }. We can let parameters be zero for positions of y and only rely the skip
32

---

T(0x
connection up to layer E, and the H is then
2j
We again apply operations from
0
Proposition B.4:
T æj 3j
T "
T
T
T(
0
moy
0
0 Vj
(33)
Vj
j=1
U1
T yy
j=1
Now we apply Lemma B.3 over all even columns in Equation (33) and we have
T
Output =
Σ
0
= T
E Σj=1 Uy
=
& T X
(34)
yJ
y
1=1
T
where & denotes irrelevant quantities. Note that the resulting ( X is also the same as Iterative Newton's predictor wy
M y after k iterations. We denote w T X y.
Transformers can make predictions on Trest by (i Feesst/
Now can make predictions
query Prosti
E i TF 0
Dtest
Etest
E
TF
Prest
mov
TF W 0
Etest TF W
mul
Ptest i (w
(35)
ii
Frest
0
TF
(Iestt
Finally, we can have an readout layer BReadOut = {u, 2} applied (see Definition 3.3) with
(0sd 1]" and v= 0 to extract the prediction (nb Beestt) at the last location, given by Buest This exactly how Iterative Newton makes predictions.
TF
To Perform l: steps of Newton's iterations, Transformers need O(k) layers.
Let's count the layers:
Initialization: needs (1) layer; gathering aS needs ((1) layer; and aff needs ((1) layer. In total, Transformers need ((1) layers for initialization.
Newton Iteration: each exact Newton's 's iteration requires O(1) layer. Implementing iterations requires ((k) layers.
Implementing u.. We need one operation of mov and mul each, requiring O(1) layer each. Apply Lemma B.3 for summation also requires ((1) layer.
Making prediction on test query: We need one operation of mov and mul each, requiring ((1) layer each.
Hence, total, Transformers can implement k-step Iterative Newton and make predictions accord- ingly using (kk) layers.
Remark B.5. We note that Giannou et al. [2023] used 13 layers to compute one Newton Iteration, and in construction, need only one Transformer layer (with one attention layer and one MLP layer) to compute one Newton Iteration. At the same time, we didn't use Akyürek et al. [2022] for constructing Newton Iterations. Akyürek et al. [2022] is applied to initialize Newton and for reading out the prediction.
In our construction, only the initialization and read-out prediction components use causal attention and softmax because Akyürek al. [2022]'s construction applied. To more specific, those are the first 3 layers in initializing Iterative Newton and the last 5 layers in reading out the predictions from the computed pseudo-inverse. All the layers corresponding to the Iterative Newton updates are using full attention and normalized ReLU activations.
33

---

Remark B.6. We note that proof he extended to causal attention for n. sufficiently larger than d. Under causal attention (see Definition 3.1) with normalized ReI.U activation, Equation (29) can be
Oa
Ia
rewritten as follows, given > we first choose
where the coefficient on the
Oa
upper right block instead of -2 originally. Then
h
(e-1)
( x
+
T z
T x
T
2
1
Wt
0a
02
0g
==1
T(
(T()) xz (T((() 23)) Tax
3t
(36)
T( 1T(1)) (t (/E)= x T(1)
(T()) (TTO T
where Σ = ビ 3 is the estimate the covariance matrix given seen in-context examples {ay, 3/ 1 so far. Since d. Σ an unbiased estimate for Σ 1
S if n is sufficiently large. The
x
rest of the proof follows similarly, up to the perturbation introduced by the enror in the estimate of Σ.
We also note when & < d, the estimate Σ = : 2 is no longer a valid covariance matrix since it's singular. Then this gives different ( for different time stamp / < and such error may propagate in our proof. Hence, a formal extension to causal models requires extensive analysis of the error bounds and is beyond the scope of this work. Nonetheless, we provide a plausible direction of such extension.
B.3 Iterative Newton as Sum of Moments Method
Recall that Iterative Newton's 's method finds St as follows
Mo
ST.
M, = 2M M 4k 2 1.
(37)
SS
12
We can expand the iterative equation moments of S follows.
Mi 2M M, 2aST TT SS 2a8 S3
(38)
Let's 's do this one more time for M2.
M2
2M, M,SM, 2(2nS 4a25" (2x.5
4a 4a2 4025a)
4a.S 2 5 16a 5 16a 57
(39)
4a.S 12a2 53 16a 5 16a4 5"
We can see that M, are summations of moments of S, with respect to some pre-defined coefficients from the Newton's 's algorithm. Hence Iterative Newton special of algorithm which computes an approximation of the inverse using second-order moments of the matrix,
24+1
M
Σ
8,S
(40)
with coefficients 8, E RR.
We note that Transformer circuits can represent other sum of moments other than Newton's method. We can introduce different coefficients 83 than the proof of Theorem by scaling the value matrices through the MLP layers.
34

---

B.4
Estimated weight vectors lie in the span of previous examples
What properties can we infer and verify for the weight vectors which arise from Newton's method? A straightforward one arises from interpreting any sum of moments method as kernel method.
We expand S follows
S"
Σ i=
x
T
Σ i=1
Σ f1
II v=l
(I Biu+s)
(41)
Then we have
2
8 M.X y
Σ
3,S*X X
==1
2
2m
Σ
83
Σ
(as Bij
II
F B joti)
2
Σ yo
i=1
Σ
Σ
8.
Σ
(2 and
a I Joti
xi
(42)
j1. de
21+1-1
Σ
Σ
Ballia (mi F.j.
II (zj jjuts)
J1.
deli 1, x
where X is the data matrix, coefficients of moments given by the sum of moments method and dx(+) is some function which assigns some weight to the i-th datapoint, based on all other datapoints. Therefore if the Transformer implements then its induced weight vector tū (Trrnsforrerss (2,,):)i after seeing in-context examples (a,, e should lie in the span of the examples 1t
of moments method (such Newton's method),
n. (Transformers (a,, v:) _ Span( ,
Σ aæ,
for coefficients a,- (43)
We test this hypothesis. Given sequence of in-context examples {o, yiM we fit coefficients (a in Equation (43) to minimize MSE loss:
{a.)i==
arg min 01.0g RR
w. w, (Transforrmrss
(x,, (it
Σ a,æ,
(44)
We then measure the quality of this fit across different number of in-context examples 4, and visualize the residual enror in Figure 27. We find that even when < d, Transformers' induced weights still lie close the span of the observed examples x,'s This provides an additional validation of our proposed mechanism.
35

---

Linearity Error # In-Context Examples
Transformers
0.00012
M
O1.5
0.00010
Linearity Error (MSTE)
0.00008
0.00006
0.00004
0.00002
0.00000
10
15
20
25
30
35
40
lo-Cnntett Examples
Figure 27: Verification of hypothesis that the Transformers induced weight vector ao lies in the span of observed examples {x:}.
C Computes
All experiments involving fine-tuning GPT2 models to learn in-context linear regressions are trained on one NVIDIA A6000. Linear probing experiments also used one NVIDIA A6000.
D License
We used PyTorch Paszke al. [2019] as code framework and used PyTorch implementation of LSTMs. PyToorc is licensed under the Modified BSD license.
We used GPT-2 Model as our backbone, and it's released under MIT License. We used trained GPT-2 checkpoints linear regression by Garg et al. [2022] and it's released under MIT License.
E Limitations
In this work, analyses of Transformers mostly based only simple task: linear regression. might not able to extrapolate to any arbitrary algorithmic tasks. would be interesting for future work to extend such analysis to an extensive class of problems.
F Broader Impacts
This paper presents work whose goal is to advance the field of Machine Learning. Through mechanistic understanding Transformers, the backbone of modern large language models (LLMs). this work can help advance building safe and trustworthy models.
36

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitations is discussed in Section E.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We provide theorem in Section 5 with its proof in Appendix B.
37

---

Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification:
We provide the details of experimental settings in $4.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer:[Yes]
38

---

Justification:
reproducible, they We will release codes and data generation processes.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
swiiiissisoiiiiioio for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ /nips. wtt.cc/cpppcccpcc.eeeeeeooooclet gww.//gw/ee///eeg /Codee/mmoeeubiiicca setssubbmisioniiie asionnolicyy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: Yes]
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
39

---

For asymmetric distributions, the authors should be careful show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The detail of the computing resourse provided at $C
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https http://eeerupppeeeee http:/pc...cpp...caccccccctaaaaa
Answer: [Yes]
Justification: The authors have read the NeurlPS Code of Ethics and made sure the paper follows the NeurIPS Code of Ethics in every aspect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
The potential societal impact is discussed in $F.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinfornation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
40

---

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification: Our paper works on simple linear regression tasks. We believe there is no such risk.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
See $D for details.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- e-paakaaed, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
41

---

Justification:
We did not introduce any new assets in this paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such as the institution conducting the review.
42


Reviews:
[{"review_id": "2WYmhMM5Ja", "reviewer": "Reviewer_vdXr", "summary": "This paper studies in-context learning in transformers, and find that gradient descent converges more slowly than transformers/Newton's method for linear regression. Previous works proposed that transformers do in-context learning using gradient-based algorithms, based on experiments with linear regression tasks. This paper suggests that transformers use a higher-order optimization method rather than a first-order method, for linear regression. The evidence for this is as follows: (1) transformers converge to the OLS solution at a similar rate as iterative newton. The predictions after each layer match those of Newton's method after a proportional number of iterations. (2) GD requires a much larger number of steps (exponentially many) to match the transformer's errors.\n\nThe problem setup/experimental setup is similar to prior works in this area - in each example sequence, they sample a weight vector $w$, and a matrix $\\Sigma$ which is from some distribution over PSD matrices (in some of the experiments, $\\Sigma$ is an identity matrix, and in others, $\\Sigma$ is ill-conditioned). The examples $x_i$ within the sequence are sampled from a Gaussian with covariance $\\Sigma$. They use the following metrics to measure similarity between transformers/other algorithms (namely Newton's method and GD):\n1. Similarity of errors - cosine similarity of the errors that two algorithms achieve. On a given data sequence of length $n$, they take the vector of errors that each algorithm achieves on $y_2, \\ldots, y_{n + 1}$, where for $y_i$, the algorithm is given $x_1, y_1, ..., x_i$. Then, they take the cosine similarity between these two error vectors.\n2. Similarity of induced weights - fit a weight vector to the predictions of this model, and compare two algorithms by comparing their induced weight vectors (on the same sequence of in-context examples).\n3. Matching steps between algorithms - Let $p_a$ be a particular number of steps taken by algorithm A. Then, the best matching number of steps for algorithm B is the one that maximizes the cosine similarity with algorithm A at step $p_a$.\n\nThe results in Figure 2 show that the middle layers of transformers converge at a superlinear rate, similarly to Newton's method. Figure 2 also shows that gradient descent has a sub-exponential convergence rate, and gets slower as the number of steps increases. It requires a much larger number of iterations. They also match the steps between transformers and Newton's method and GD (in other words, for each layer L of the transformer, find the step of Newton's method/GD whose predictions are the most similar to the transformer at layer L). Figure 3a shows that there is a linear relationship between the transformer layer index and the corresponding step index of Newton's method. (Figure 8, which matches steps based on weight vectors instead of errors, seems to show a stronger linear trend.)", "questions": "- Figure 5, left, seems to be showing a drastic difference between the performance of transformers and Newton's method. This may be because the red curve only uses 5 steps of Newton's method. Is the result different if more steps of Newton's method are used?", "limitations": "Yes.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- The experiments are thorough and successfully show that transformers can learn optimization algorithms much faster than gradient descent.\n- The paper is very well-written.", "weaknesses": "- As the authors mention, the experiments ultimately show that transformers behave more similarly to some higher-order method, rather than Newton's method in particular. In some cases, transformers can behave differently from Newton's method as well.\n    - For instance, in Figure 20 (left), it could be argued that the transformer converges even faster than iterative Newton's method, i.e. the relationship between the iteration of Newton's method and the transformer layer index is not exactly linear. For instance, layer 14 of the transformer seems to correspond to about 4 iterations of Newton's method.\n\n- It would be useful to clarify the difference between the takeaway of this paper and e.g. Garg, et al. (2022), which finds that transformers can match ordinary least-squares very closely.\n    - If this submission shows that transformers learn Newton's method in particular, then the difference would be clear. This work gives evidence that transformers learn faster than GD, but somewhat less strong evidence that transformers learn a particular algorithm.\n\nGarg, et al. (2024) \"What Can Transformers Learn In-Context? A Case Study of Simple Function Classes\"\n\n==================================\n\nThe authors clarified my questions in the rebuttal - I am increasing my score from 6 to 7.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for detailed comments and suggestions. We are pleased that the reviewer thinks our experiments are thorough and successfully support our main claim. We are also happy to see the reviewer thinks our paper is well-written.\n\n## **In some cases, Transformers can behave differently from Newton's method** \nAdmittedly, Transformers could behave slightly differently compared to Newton’s method as the reviewer pointed out. This could come from optimization noises during training or probing. Additionally, we are not claiming trained Transformers models are exactly implementing Iterative Newton’s method and this is why our title says “higher-order method” rather than “Newton’s Method”. Our main claim is that Transformers resembles a class of optimization algorithms, and a slight deviation from Newton’s method wouldn’t hurt such claims.\n\n## **Clarify the difference between this paper and Garg et. al.**\nGarg et. al. empirically showed that Transformers, when performing in-context linear regression, match the ordinary least-squares (OLS) solution. However, what remains unclear is **how** Transformers converge to the OLS solution. One common hypothesis is that Transformers converge to the OLS solution via Gradient Descent (GD), while the main contribution of this work is to show that (1) empirically, Transformers can converge exponentially faster than GD, which suggests that they emulate higher-order optimization methods; (2) theoretically, there is a construction of Transformers to implement a particular higher-order optimization method – the Iterative Newton’s method.\n\n## **Whether Transformers evidently learn a particular algorithm**\nAgain, we are not claiming trained Transformers models are exactly implementing Iterative Newton’s method. We also compared Transformers with many alternative higher-order optimization methods such as Conjugate Gradient and (L-)BFGS, and the conclusion is that Transformers are learning a particular class of algorithms, that are utilizing higher-order information. \n\n## **LSTM only matches 5-step Newton’s Method**\nYes. As shown in Fig. 6(c), LSTMs couldn’t improve with more layers and the errors are still quite large. This implies that LSTMs could only converge to an estimator $\\hat{w}_{LSTM}$ that achieves the same performance as Iterative Newton’s method after only 5 steps, which is far from converging to the least squares solution.", "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": "Puvwhb5atX", "reviewer": "Reviewer_F4MD", "summary": "In the paper the authors study the nature of in-context learning in transformers. Starting from the hypothesis of previous work on the fact that transformers may internally implement gradient descent algorithms to correctly perform linear regression on test data, the authors put forward the theory that transformers actually implement a different, higher-order algorithm. The authors back their claim with empirical evidence on the convergence rate of the algorithm across layers, and they provide a theoretical construction for the implementation of Newton's iterative algorithm using the transformer architecture.", "questions": "What do you think may be causing the mismatch between the theoretical and the empirical results, i.e., what do you think is the limit of your construction? Do you think that it may be improved, e.g, by some additional architecture components, or by the substitution of softmax instead of ReLU?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I believe that the idea that transformers implement a higher-order algorithm for linear regression is intriguing, and I think that the empirical evidence provided by the authors definitely rule out the possibility that transformers simply implement gradient descent and not a more sophisticated algorithm.", "weaknesses": "The limit of the work is two-fold: (1) the authors only provide circumstantial -- and not direct -- evidence that the network is learning a higher-order algorithm, and the provided evidence on ill-conditioned data actually seem to contradict the theory that the network is actually implementing Newton's iterative method and not something ever more sophisticated; (2) there is quite a bit of mismatch between the empirical and the theoretical results: there is a 8-layer baseline required for the proposed construction to work that doesn't seem to appear in the experiments, and 8 is quite a large number for the setting presented in the paper (e.g., in Figure 2 the transformer seem to achieve small error already after 8 layers).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for detailed comments and suggestions. We are pleased that the reviewer thinks our ideas intriguing and empirical evidence convincing.\n\n## **Evidence for higher-order is indirect**\nOur main evidence is the convergence rate. We do believe that convergence rate is a solid categorization of optimization algorithms where higher-order methods will have $\\log\\log(1/\\epsilon)$ rates whereas first-order could only achieve $\\log(1/\\epsilon)$ rates. We also note the $\\Omega(\\log(1/\\epsilon))$ lower bound of gradient-based methods and it is not possible to improve Gradient Descent’s convergence rate *without using second-order methods*.  \n\n## **contradiction in ill-conditioned case**\nThere might be some degree of misunderstanding about our ill-conditioned experiments. Our experiments show that both Transformers and Newton’s method are *not* susceptible to ill-conditioned data – a common trait for higher-order optimization methods. In both well-conditioned and ill-conditioned experiments, our heatmaps (Fig 3 and Fig 15) show that one Transformer layer could implement three Newton iterations, which indicates that Transformers could implement some algorithms even more complicated and stronger than Iteratvie Newton’s method. However, they are still higher-order optimizations with a similar convergence rate to Newton’s method,\n\n## **Mismatch between the empirical and the theoretical results**\nSurely, there’s a mismatch between empirical and theoretical results for the initial constant number of layers needed, but they are both  $\\mathcal O(1)$. Empirically, we observe there’s an  $\\mathcal O(1)$ number of layers needed for a **warmup**: the initial layers where transformers cannot show improvements in prediction errors. Please see Figure 2(a) for 12-layer transformers, 2 layers at the beginning are reserved for warmup; in Figure 21 for 24-layer transformers with iterative newton, 4 layers are reserved at the beginning. Similarly, our theoretical construction needs 8 layers, but it still lands in the category of $\\mathcal O(1)$. \n\nOur main focus for theoretical construction is to show that, to implement $k$ Newton steps, we need $k + \\mathcal O(1)$ Transformers layers, with a particular focus on the 1-to-1 correspondence between **intermediate layers** and Newton steps. Admittedly, we didn’t optimize our construction for the $\\mathcal O(1)$ constant number of layers needed to initial layers. There could be alternative constructions with a smaller amount of warmup layers than 8, but that’s not the main focus of our theoretical analysis.\n\nThere could be many causes for the mismatch of the constant. For example, the number of layers of the Transformers, the hidden dimension size, etc. At the same time, optimization algorithms used to train these Transformers could also be another factor on the empirical side. On the theoretical side, as the reviewer pointed out, the substitution of softmax activation by ReLU could be an important factor. Nonetheless, matching the exact number between empirical and theoretical results would be tedious but matching the number within the same complexity of  $\\mathcal O(1)$ is simply what’s presented in this paper.", "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": "lGiZMzEwBg", "reviewer": "Reviewer_dusD", "summary": "This work investigated the ability of transformers to implement higher-order optimization methods for in-context learning of linear regression tasks. The authors considered a noiseless linear regression setting and compared the output of each layer of TF with few steps on gradient descent (GD), online gradient descent (OGD), iterative Newton's method, and empirically showed that the output has a linear trend with the iterations of the Newton's method but an exponential trend with GD's iterations. Therefore, they concluded that TFs are more likely to implement high-order methods instead of first-order method for ICL tasks. They also performed experiments on ill-conditioned tasks and observed similar trends. In addition to the experiments, they  theoretically proved that there exists a polynomial size TF that can implement  iterative Newton's method.", "questions": "What is the choice of stepsize for GD? Is it optimal? Would different choices of stepsize for GD affect the exponential trend  shown in Figure 3 (right)?\n\nDo similar observations appear in linear regression tasks with non-zero noise? For example, in this case, would TF perform high-order optimization methods to learn the optimal ridge predictor?", "limitations": "This work compared the outputs of TFs and some first-order/higher-order algorithms and concluded that TFs are more likely performing higher-order optimization based on the similarity of their outputs. However, this does not imply that TFs  internally implement Newton's method or any other specific higher-order method. \n\nThis work focused on ICL with linear regression tasks. It is unclear if TFs still behave like higher-order methods in other convex ICL tasks (e.g., logistic regression for classification) and what TFs approximate if the learning task is non-convex.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-written. The experimental results and their implications are well-discussed. The similarity metrics used to compare different algorithms are reasonable. Novel theoretical results on the approximation ability of TFs are proved.", "weaknesses": "Assumption on the linear regression tasks. This work considered a noiseless setting in which the optimal predictor in the OLS. It is unclear to what extent the observations in this work can be generalized the setting where the noise exists.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for detailed comments and suggestions. We are pleased that the reviewer thinks our experimental results and their implications are well-discussed. We are happy to see that the reviewer regards our theoretical results novel, our comparison metrics reasonable, and our paper well-written.\n\n## **Generalization to noisy problems**\nWe thank the reviewer for this suggestion. We run experiments on noisy linear regression with a noise standard deviation of $\\sigma = 0.1$. Please refer to Figure 1 in the rebuttal PDF for detailed results. We observe that in noisy setup, transformers and iterative Newton’s method still share the same convergence rate, and are both exponentially faster than GD.\n\nAdditionally, under the noisy linear regression setup, the closed form solution would be $\\hat{w} = (X^\\top X + \\lambda I)^\\dagger X^\\top y$. Our theoretical construction stills holds with slight modifications: replacing $S = X^\\top X$ by $S = X^\\top X + \\lambda I$. We will discuss this more clearly in our next revision.\n \n\n## **Choice of step size for GD**\nNotice that the optimal step size for GD is $\\eta = \\frac{2}{\\beta + \\gamma}$ where $\\beta$ and $\\gamma$ are the smallest and largest eigenvalues for $S = X^\\top X$ respectively. However, each sequence is sampled randomly and will have different optimal $\\eta$, and it would be an unfair comparison if extra computations are allowed for eigenvalues of $S$ – which involves computing higher moments of $S$ if using numerical methods, for example, power methods. To circumvent this, we sampled 10,000 sequences and computed the average optimal step size $\\bar{\\eta} = \\mathbb E[\\eta]$ and used the same step size $\\bar{\\eta}$ for GD. \n \nEmpirical, we showed that transformers match Iterative Newton’s convergence rate. Theoretically, Iterative Newton’s convergence rate is exponentially faster than GD with optimal step sizes (see lines 121 and 130). Combining the two, we can deduce that transformers will also be exponentially faster than GD with optimal step sizes.\n \n\n## **Generalization to logistic regression and classification, or even non-convex problems**\nWe believe that linear regression is the right starting point to understand how transformers perform in-conext learning. It would be interesting to see the extensions to logistic regression, classification, and even non-convex problems. However, this is beyond the scope of this paper.", "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": "wG6Xq3IfcC", "reviewer": "Reviewer_p5Dq", "summary": "This paper demonstrates that transformers learn to approximately perform higher order algorithms, building on the previous works that demonstrate that transformers can approximate gradient descent which is a first-order algorithm. For the problem of linear-regression, it is empirically shown that that prediction errors of transformer are similar to that of iterative Newton's method, and so are their corresponding rates of convergence. Further, it is shown by explicit construction of weights how a $k+\\mathcal{O}(1)$ layers of transformers can implement $k$ steps of the Newton's method.", "questions": "(1) The effect of the number of layers is that the prediction after each layer progressively improves. However, one might also be interested in the effect of embedding dimension on the rate of convergence, which is not discussed. I am curious if you have any comments on this?", "limitations": "(1) Within the scope of the problem, there is no specific limitation that requires attention (pun intended).", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "(1) The work provides a very good understanding of how transformers can efficiently solve (even ill-conditioned) linear regression problems, which is of interest to the research community.\n\n(2) This paper establishes that transformers are better algorithmic engines than other deep learning models (for example, LSTMs, in Figure 6 of Appendix A.2.1).\n\n(3) The performance of the trained transformer is very close to that of a higher-order algorithm, even though transformer is not trained for different number of layers $\\ell$ separately (only the ReadOut layer is retrained). This demonstration strongly supports the crucial claim that the representations after each layers learned by the transformers in fact mimic the updates of some higher-order algorithm.\n\n(4) The experiments are well-executed and their presentation is commendable.\n\n(5) The proof of the main theoretical contribution (Theorem 5.1) is very clean.", "weaknesses": "(1) Although the result is very interesting that leads to a better understanding of the abilities of a transformer, there is little direct impact in terms of application.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed comments and suggestions. We are pleased they find our work valuable in demonstrating how transformers efficiently solve linear regression problems, including ill-conditioned ones. They appreciate our comparison showing transformers as superior to models like LSTMs, and acknowledge our experimental results show strong evidence that transformers resemble higher-order algorithms. We thank the reviewer for commending our well-executed experiments, clear presentation, and clean proof of our main theorem.\n\n\n## **Direct Impact**\nWe acknowledge that this is a more interpretability and theoretical work whose direct impact to applications is unknown. However, we believe that interpretability could be useful in general, for example, it facilitates our understanding of architectures, and motivates us to make reliable changes in the future.\n\n## **Impact of Hidden Dimension Size**\nEmpirically, we ablate 12-layer 1-head transformers with different hidden dimensions: 64, 32, 16, and 8. Please see results in the Figure 2 in the rebuttal PDF. We find that transformers with hidden dimensions 64 and 32 are able to converge to the OLS solution. On the contrary, transformers with hidden dimensions 16, and 8 could not.\nTheoretically, as stated in Theorem 5.1, the hidden dimensions in our theoretical construction require a size of $\\mathcal O(d)$, where $d$ is the dimension of the linear regression problem. That being said, as long as the hidden dimensions are of order $\\mathcal O(d)$, transformers are able to converge to the OLS solution, and the higher-order optimization methods are learned in-context. This coincides with our empirical results, where $d = 20$ and the hidden dimensions need to be $\\mathcal O(d)$.", "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": "eIOdzZmK7t", "reviewer": "Reviewer_J6bG", "summary": "The authors show that when trained on ICL tasks encoding a regression problem, the layers of a transformer implement an implicit higher order optimization method. They contrast this with prior work on this subject that suggest that the transformer in this setting implicitly implements a preconditioned gradient method. They also construct a transformer capable of implementing this proposed algorithm.", "questions": "1. What happens to the covariates across layers? To implement these algorithms, the transformer needs to implicitly compute some second order moment of the original data, but at each layer can only work with the “modified covariates” at that layer. It would be interesting if these quantities somehow turned out to be invariant across layers.\n2. Is it fair to say that ReLU activations don’t hurt performance? Softmax seems to be pretty crucial, or why would people go to the trouble of implementing something that has a larger run-time? Even most of the works cited only say it isn’t much worse, sometimes. Also, it seems softmax is used for the experiments in this work anyway, and ReLU is only used for some of the layers in the construction. imo this is not a big issue, just pointing it out.\n3. I dont understand the description of the Iterative Newtons Method (Page 3). It seems $S$ is an empirical covariance, and then $M_k$ is just a function of $S$. What is the purpose of computing the intermediate $\\hat{w}_k^{\\text{Newton}}$ (other than to show that they match the Transformer)?", "limitations": "Yes", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "There appears to be a better match between the layers of the transformer and iterations of iterative newtons method as compared to the iterations of the gradient methods they consider.", "weaknesses": "1. In the comparisons to Iterative Newton: If the Transformer had 40 layers do you expect that each layer would implement a single iteration or that it would match more layers of IN overall? In other words, is there something fundamental about taking 3 layers together to compute one iteration?\n\n2. Does the transformer also match the higher order methods when solving other problems? That is, if the objective was not OLS for linear regression? The problem with studying only linear regression is that subsequent layers interpolating between 0 and the OLS solution is somewhat what one would “expect”. Not matching with GD fits into this interpretation because GD exhibits an oscillatory behavior for ill-conditioned problems that an “interpolator” would have no reason to match. \n    1. What happens to Appendix A.3 when we compare the transformer to the appropriately preconditioned GD? \n    2. Figure 15 is only an illustration of the best matching iterations. I have two concerns with this. The construction given in the paper is theoretical and not necessarily the one used by actual transformers (it is also deeper than 3 layers). I think this reduces the impact of saying that iterative newton matches every third layer of the transformer, since even GD is matching (though with a varying number of iterations). And the match quality of GD after those varying number of iterations is not much worse.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for detailed comments and suggestions. We are pleased that the reviewer agrees that our experiments show a better match between the layers of the transformer and iterations of iterative newtons than when compared with gradient descents.\n\n## **Matching between Transformer layers and Iterative Newton steps**\nWe believe this is some degree of misunderstanding by the reviewer. As shown in the heatmap in Figure 3 (left), one Transformer layer is implementing 3 iterations of Newton steps, not the other way around as the reviewer suggested.  The reviewer can also take a look at Figure 21 on a 24-layer transformer model, where each layer is approximately implementing one Newton step until convergence. \n\n## **Generalization to other problems**\nWe believe that linear regression is the right starting point to understand how transformers perform in-conext learning. It would be interesting to see the extensions to logistic regression, classification, and even non-convex problems. However, this is beyond the scope of this paper.\n\n## **Not matching GD is expected**\nGarg et. al. [1] empirically showed that Transformers, when performing in-context linear regression, match the ordinary least-squares (OLS) solution. However, what remains unclear and what is the main focus of this paper, is **how** Transformers converge to the OLS solution. \nWe wouldn’t say not matching GD is expected because many existing work claims that Transformers converge to the OLS solution via Gradient Descent (GD) [2,3,4,5]. The main contribution of this work is to show that (1) empirically, Transformers can converge exponentially faster than GD, which suggests that they emulate higher-order optimization methods; (2) theoretically, there is a construction of Transformers to implement a particular higher-order optimization method – the Iterative Newton’s method.\n\n## **Comparison with Precondition GD**\nWe would like to emphasize that even with well-conditioned data, our experiments show Transformers and Newton converge exponentially faster than Gradient Descent. Under the well-conditioned case, preconditioning would not do anything because the preconditioner (inverse of the Hessian matrix) is identity. Even under the ill-conditioned case, where the preconditioner is not identity, one needs to compute the inverse of the Hessian matrix first, where such inverse computation is already computationally heavy and involves second-order information, and to compute the inverse efficiently, one needs to use the Iterative Newton’s Method. In our setup, the eigenbasis of the covariance matrix $\\Sigma$ is sampled at random, so there is no way the Transformer stores the preconditioner during training so they need to compute the inverse at ICL inference time. \n\n## **Weakness 2.2: mismatch between theoretical construction and empirical Transformers**\nThere is some degree of misunderstanding from the reviewer about our main claim and we would like to reiterate our claim and evidence. \n\n- First, we don’t claim that Iterative Newton is the algorithm that a trained Transformer model is implementing. What we really claim is that Transformers learn some algorithm that convergence exponentially faster than the gradient descent method, and algorithms with such log log convergence rate fall into the same category of “higher-order” optimization methods, where Iterative Newton is one of them. We do believe that convergence rate is a solid categorization of optimization algorithms. We also note the $\\Omega(\\log(1/\\epsilon))$ lower bound of gradient-based methods [1] and it is not possible to improve Gradient Descent’s convergence rate *without using second-order methods*.  \n\n- Second, our empirical results show that Transformers share the same convergence rate as Iterative Newton and such results can categorize Transformers as a higher-order method, and they are all exponentially faster than Gradient Descent algorithms. Please refer to Fig 3 for the well-conditioned case, Fig 15 for the ill-conditioned case, and Fig 21 for deeper transformers. The intuition why we don’t believe Transformers are doing GD is that, for example in Fig 3 (right), there is no way for one Transformer layer (from layer 7 to layer 8) to implement 500 gradient descent steps. \n\n- Finally, we show expressivity results that Transformers can indeed implement Iterative Newton’s method, a representative higher-order method.\n\n## **Covariates across layers**\nWe would like to humbly ask the reviewer to clarify the question. We also would like to point to our theoretical proof that the intermediates $M_k$ do not need to be stored in each layer. In this case, it will be difficult to extract the exact intermediate covariates.\n\n## **ReLU activation**\nAs the reviewer pointed out, this is not a bit issue as many existing work studies either linear attention (no activation applied to attention layers) or with ReLU activations.\n\n## **Purpose of computing $w^{\\mathrm{Newton}}_{k}$**\nIt allows us to compute the convergence rate of algorithms, for which one needs to show the relationship between $||w_{k+1} - w_{\\star}||$ and $||w_{k} - w_{\\star}||$ for the optimal solution $w_{\\star}$. In this case, we need to compute the intermediate \n$w^{\\mathrm{Newton}}_{k}$. \n\nMoreover, $w^{\\mathrm{Newton}}_{k}$ is needed to make predictions on any test sample, and we can measure how good each iteration has been by looking at the errors on test samples.", "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:
We sincerely thank the reviewers for their time and effort in reviewing our work. We are encouraged that the reviewers praise our work for providing a very good understanding of how transformers can efficiently solve linear regression (p5Dq) and find our ideas of higher-order optimization intriguing (F4MD). We are happy to see that reviewers find our experiments thorough (vdXr) and well-executed (p5Dq), our results reasonable (J6bG) and convincing (vdXr, F4MD), and their implications are well-discussed (dusD). We are also pleased to see the reviewer commend our theoretical results (dusD, p5Dq). Finally, we are delighted to see reviewers find our paper well-written (vdXr, dusD) and our proof clean (p5Dq).

We conducted two additional experiments per the reviewers’ comments and the results are attached in the **PDF**:

- **Noisy Linear Regression**. As reviewer dusD suggested, we conduct additional experiments on noisy linear regression problems, and find our claim that transformers learn high-order optimization methods in-context still holds in noisy linear regression setups. 

- **Varying Hidden Dimension Sizes**. As reviewer p5Dq suggested, we ablate on the hidden dimensions of transformers, and observe the resonance between empirical results and our theoretical claim that a hidden dimension of $\mathcal O(d)$ is necessary for Transformers to converge to the OLS solution and thus learn higher-order optimization methods in-context.

We have addressed the comments in the responses to each reviewer and will incorporate all feedback into the paper.

Meta-review: The paper analyzes in-context learning of linear models, showing that parametric transformers can achieve a super-linear convergence rate. To explain this rapid convergence, this paper demonstrates that transformers are expressive enough to implement the "iterative Newton's method" which enjoys a super-linear convergence rate.

Based on the discussions and reviews, I recommend to accept a poster presentation with the following revisions in the final draft:

- Narrowing the focus in the title, abstract, and introduction to linear regression. The implementation of second-order optimization for in-context learning, as claimed in the abstract, is beyond the paper's scope.
- Discussing the gap between theoretical analyses and experiments, including the differences between the architectures used in theory and those used in experiments.
- Including results for linear attention mechanisms, and then extending the analysis to attention layers with ReLU.
Decision: Accept (poster)