Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression

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.

Matching between Transformer layers and Iterative Newton steps

We 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.

Generalization to other problems

We 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.

Not matching GD is expected

Garg 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. We 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.

Comparison with Precondition GD

We 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.

Weakness 2.2: mismatch between theoretical construction and empirical Transformers

There is some degree of misunderstanding from the reviewer about our main claim and we would like to reiterate our claim and evidence.

• 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.

• 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.

• Finally, we show expressivity results that Transformers can indeed implement Iterative Newton’s method, a representative higher-order method.

Covariates across layers

We 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.

ReLU activation

As 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.

Purpose of computing $w^{\mathrm{Newton}}_{k}$

It 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 $w^{\mathrm{Newton}}_{k}$.

Moreover, $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.

[1]Shivam Garg, Dimitris Tsipras, Percy Liang, Gregory Valiant. What Can Transformers Learn In-Context? A Case Study of Simple Function Classes. In NeurIPS, 2022

[2] Johannes von Oswald, Eyvind Niklasson, E. Randazzo, Joao Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, 2022.

[3] Ekin Akyurek, Dale Schuurmans, Jacob Andreas, Tengyu ¨ Ma, and Denny Zhou. What learning algorithm is incontext learning? investigations with linear models. The Eleventh International Conference on Learning Representations, 2024

[4] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent. Advances in Neural Information Processing Systems 36. 2024

[5] Yu Bai, Fan Chen, Haiquan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36. 2024

Thank you for your thoughtful suggestions and comments.

We have conducted an additional experiment on a non-linear function class: 2-layer neural network with ReLU activations, a setting also studied by pioneering work, Garg et. al. (2022) in Fig. 5(c) in their NeurIPS 2022 paper version. Please refer to https://anonymous.4open.science/api/repo/transformer_icl_neurips_rebuttal-2E55/file/Transformer_ICL_rebuttal_additional.pdf for experimental results. TL;DR: Transformers still converge superlinearly, and exponentially faster than GD -- indicating higher-order optimization in-context.

In the same prompt setup $(x_1, y_1, \cdots, x_t, y_t)$, each $y_k$ is generated by $y_k = a^\top \mathrm{ReLU}(Wx_k + v) + b$. There are 100 neurons, aka, the hidden size, and the problem dimension $d$ is still 20. As shown in Figure 3a, even on 2-layer neural network tasks, Transformer shows superlinear convergence rates. As shown in Figure 3b, Transformer shows an exponentially faster convergence rate than GD’s, because GD’s steps are shown in log scale and the trend is linear – similar to Figure 9 in the main paper.

Although our title indicates the study mainly focuses on linear functions, we find the extension to non-linear functions quite meaningful. We will make the experiments extending to non-linear function classes more thorough and include more discussions in the camera-ready version.

In analogy to many optimization algorithms, what's updated along the iterations are the estimators $\hat{w}$ than the covariates $x_i, y_i$. We can understand by analogy to any other iterative method, such as GD. The final $L-1$ steps of GD solve some problem whose input is the solution to the first step, $\hat{w}_1$. However, the $x_i, y_i$ for the linear regression instance would not change, it’s just that the optimization algorithm starts with a better estimate. In analogy to Newton, the goal is to compute better and better estimates of the inverse covariance iteratively. Please let us know if you have further questions.

Dear Reviewer, as the discussion period is about to end, we would like to know if our newest experiments on two-layer MLP with ReLU activation (thread here: https://openreview.net/forum?id=L8h6cozcbn&noteId=7Gu8rWpfSn) have cleared out any of your concerns. We are also happy to answer any further questions. Thanks!

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.

Direct Impact

We 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.

Impact of Hidden Dimension Size

Empirically, 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. Theoretically, 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)$.

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.

Generalization to noisy problems

We 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.

Additionally, 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.

Choice of step size for GD

Notice 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.

Empirical, 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.

Generalization to logistic regression and classification, or even non-convex problems

We 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.

We thank the reviewer for detailed comments and suggestions. We are pleased that the reviewer thinks our ideas intriguing and empirical evidence convincing.

Evidence for higher-order is indirect

Our 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.

contradiction in ill-conditioned case

There 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,

Mismatch between the empirical and the theoretical results

Surely, 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)$.

Our 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.

There 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.

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.

In some cases, Transformers can behave differently from Newton's method

Admittedly, 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.

Clarify the difference between this paper and Garg et. al.

Garg 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.

Whether Transformers evidently learn a particular algorithm

Again, 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.

LSTM only matches 5-step Newton’s Method

Yes. 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.