Transformers Learn Higher-Order Optimization Methods for In-Context Learning: A Study with Linear Models

We thank the reviewer for providing constructive feedback.

The reviewer asked why we use cosine similarity on errors. Cosine similarity is better for cases where values are small. Usually errors (also known as residuals) have quite small magnitudes. Directly measuring the prediction will usually yield high similarities. We also ran comparisons using the L2 norm and didn’t find a difference in convergence rates. Our alternative metric (similarity of induced weights) also agrees with our cosine similarity metrics.

The reviewer questioned the validity of using Induced Weights. Samples used to induce these weights are coming from the same distribution as our in-context examples. As Transformers can make good predictions on $x$ sampled in-distribution, and the true $x \rightarrow y$ mapping is linear, it’s reasonable to approximate Transformers with a linear function too. The induced weights are calculated through solving a deterministic linear system. We wouldn't think of out-of-distribution samples (for example scaling) at this stage. This is also done by [2] in Eqn. 16 in section 4.1.

[2] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models.

The reviewer questioned using best-matching hyperparameters to compare algorithms. The hyperparameters we match for comparisons are the number of steps of GD/Newton v.s. the number of layers of Transformers. The intuition behind is: how many more steps of Newton or GD are needed to match the performance of an additional layer of Transformers. Afterwards, we can use this to compare the rate of convergence between algorithms, and we show that Transformers and Iterative Newton can have the same convergence rate and both are exponentially faster than GD.

The reviewer asked whether we can use Causal Attention for theory. It can be relaxed if we only care about cases when the number of in-context examples are greater than the data dimension. Under such assumption, in Eqn. 32, $\frac{1}{n} \sum_{j=1}^n x_j x_j^\top$ will become $\frac{1}{t} \sum_{j=1}^t x_j x_j^\top$, when $t \geq d$, this is a valid estimator of the covariance matrix $\hat{\Sigma}$ and the rest of the proof remain similar. However, for $t < d$, this is not a valid estimator any more.

Thank the authors for the detailed response. My concerns regarding the performance metric and causal attention are well addressed. It is encouraged to also include the simulation results about $l _2$ loss in the main paper or at least mention them.

My concern about the best-matching parameters still remains. The overfitting may appear due to the "best-matching". For example, a naive algorithm is an ergodic random walk on the subset that the results of GD and transformers belong to. Then we can always find a step of this naive algorithm that is close to the results of transformers. Thus, I will maintain my scores.

We thank the reviewer for discussions.

In the experiments, there are two sets of hyper-parameters:

1. Hyper-parameters specific to algorithms (learning rate of GD and initialization scale for Newton): these hyper-parameters are determined before running the algorithms. As long as they are chosen to guarantee convergence, they won’t affect the rate of convergence.
2. Hyper-parameters that determine the rate of convergence: number of steps of GD, number of steps of Iterative Newton, and number of layers of Transformers.

In comparing the rate of convergence, we compare the second set of hyper-parameters (steps/iterations). In order to compare which algorithm convergences faster, we need do the following: Fixing error rate $\epsilon$, if Transformers can achieve $\epsilon$ error with $\ell$ layers, how many steps of Iterative Newton or GD are needed to achieve the same $\epsilon$ error. We find that Iterative Newton needs $\mathcal O(\ell)$ steps and GD needs $\mathcal O(\exp(\ell))$ steps. Our experiments are designed accordingly to compare this.

We acknowledge the confusion of using the same name “hyper-parameters” to address two sets of parameters and we’ll modify the terminology of the second to be “best-matching steps for comparing convergence”.

The reviewer also mentioned that “linearly combining GD and Newton” solutions can produce a better algorithm than Newton. This is not true. Ignoring the $\kappa$ term in well-conditioned cases, Newton converges to $\epsilon$ error with $\mathcal O(\log \log (1/\epsilon))$ steps, while GD is exponentially slower with $\mathcal O(\log (1/\epsilon))$ steps. Combining the two may produce faster algorithms than GD but any interpolation of a fast algorithm (Newton) and a slow algorithm (GD) cannot produce a faster algorithm than the fast one (Newton).

We thank the reviewer for providing constructive feedback.

The reviewer asked for details about hyperparameters. Iterative Newton has no learning rate. We choose the best learning rate for GD via grid search to ensure convergence, and here it’s 0.01. As long as GD converges, it won’t affect our claim on the rate of convergence. The best matching hyperparameters (number of steps in GD/Newton v.s. number of layers in Transformers) are selected by taking the argmax of similarity.

The reviewer pointed out GD should converge to the OLS. This is indeed the case, however it converges much more slowly. We have thus updated our claims to be about convergence rates. we revised our experiments so that all three methods of interest converge to the OLS solution and the major difference is the rate of convergence.

We thank the reviewer for providing constructive feedback.

The reviewer asked whether another first-order method can be more similar to Newton's method. In terms of rate of convergence, the lower bound for gradient-based methods is $\Omega(\sqrt{\kappa(S)} \log(1/\epsilon))$ [1]. It is still exponentially slower than Newton’s method.

[1] A.S. Nemirovski and D.B Yudin. Problem complexity and method efficiency in optimization.

The reviewer pointed out a missing reference on Looped Transformers that can implement matrix inverse by Iterative Newton’s method. We thank the reviewer for pointing it out. We believe their construction could be used as a subroutine in an alternative proof of our claim, however one major difference is that our construction explicitly considers regression tasks whereas their construction concerns matrix inversion. Our construction never explicitly calculates the Newton inverse $M_k$ (at step $k$): we always consider $M_k x_j$ for some input token $x_j$. Replacing our construction with looped transformers also requires additional circuits to compute $X^\top y$ and make predictions using $\langle M_k X^\top y, x_\mathrm{test}\rangle$. We will discuss more about this in our next revision.

Question 1: The reviewer asked about the step size of GD. We choose the best learning rate for GD via grid search to ensure convergence, and here it’s 0.01. As long as GD converges, it won’t affect our claim on the rate of convergence.

Question 2: The reviewer pointed out GD should converge to the OLS. This is indeed the case, however it converges much more slowly. We have thus updated our claims to be about convergence rates. we revised our experiments so that all three methods of interest converge to the OLS solution and the major difference is the rate of convergence.

We thank the reviewer for providing constructive feedback.

The reviewer asked whether trained Transformers can do something smarter than Newton. We agree this is possible. As discussed in Section 5 and Appendix B.3, we showed that Iterative Newton is a specific version of the “sum of moments” method. Trained Transformer is also such a method, but could have better coefficients. Nevertheless, this won’t change the asymptotic rate of convergence (measured as big-O), since Newton’s method matches the lower bound for the problem.

Question 1: The reviewer pointed out GD should converge to the OLS. This is indeed the case, however it converges much more slowly. We have thus updated our claims to be about convergence rates. In our revised manuscript, we fix this issue: GD, Iterative Newton, and Transformers will all converge to the OLS solution and the major difference is the rate of convergence. GD converges exponentially slower than the other two. The reviewer asked for learning rates used for GD. We choose the best learning rate for GD to ensure convergence, and here it’s 0.01 in well-conditioned cases. As long as GD converges, it won’t affect our claim on the rate of convergence.

Question 2: The reviewer asked whether LSTM is also more similar to Newton with 5 steps. We also updated the OGD experiments by initializing w to zero (previously it was initialized randomly). In the new Fig. 5, LSTM matches well with both OGD and 5-step Newton in terms of in-context performance. However, LSTM and OGD have the same level of forgetting issues but Newton does not. In the updated table, even compared vertically, LSTM matches OGD the most.

Question 3: The reviewer asked how Transformers can decide $\alpha$ which can depend on $S S^\top$ in Eqn. 8. For Iterative Newton’s method, $\alpha$ can be any value between 0 and $\frac{2}{||S S^\top||_2}$ to ensure convergence. We pick $\frac{2}{||S S^\top||_2}$ in experiments in order to reduce randomness. Transformers can learn any fixed $\alpha$ as long as it is sufficiently small.

Question 4.1: The reviewer suggested we should describe $\Sigma$ as fixed in problem setup. In most of the experiments we fix $\Sigma$ as identity (see the ending sentence at the first paragraph of Sec. 3). We also need to sample $\Sigma$ for ill-conditioned problems (so that the eigenbasis is randomly sampled to prevent Transformers from being lazy by memorizing a fixed $\Sigma$), where it further verifies that both Transformers and Newton won’t suffer from large condition number but gradient-based methods do.

Question 4.2: The reviewer questioned about cosine similarity as the distance measure. Cosine similarity is better for cases where values are small. Usually errors (also known as residuals) have quite small magnitudes. We also ran comparisons using the L2 norm and didn’t find a difference in convergence rates. Our alternative metric (similarity of induced weights) also agrees with our cosine similarity metrics.