Official Review · NEURIPS 2024

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. The 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): 1. 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. 2. 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). 3. 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$. The 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.) - The experiments are thorough and successfully show that transformers can learn optimizatio