One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention

We thank the reviewer for their feedback, and respond to the comments below:

Practical Insight: Our work makes progress towards understanding the mechanism by which transformers perform in-context learning. Theoretical analyses usually need to start with simple cases to build intuition. Our setting with i.i.d. examples $(x_i, y_i)$ is also motivated by the few-shot in-context learning setting seen in practice.

Experiments: Please see Appendix D (newly added during discussion period), where we give experimental results with the goal of showing that the linear self-attention layer implements 1 step of GD, with learning rate matching our theoretical prediction. We observe that the learning rate decreases as sigma increases.

Fully-connected neural networks (NNs) and Assumption 1: We can show that fully-connected NNs satisfy this assumption as follows. For simplicity, consider a two-layer neural network $f(x) = W_2 \sigma(W_1 x)$ where $\sigma$ is a nonlinearity, and $W_1$, $W_2$ have i.i.d. Gaussian entries. Let $R$ be a fixed rotation matrix. Then, $(f \circ R)(x) = W_2 \sigma (W_1 R x)$. Since $R$ is a fixed rotation matrix and $W_1$ is a Gaussian matrix, $W_1R$ has the same distribution as $W_1$. Thus, the random function $f \circ R$ has the same distribution as $f$.

Contribution relative to Zhang, et al. 2023: We note that Zhang, et al. 2023 is a concurrent/independent work, which was posted on arxiv around the same time as ours.

We thank the reviewer for their positive review, and for mentioning that the results are noteworthy.

Clarification on global minimizers: We note that all global minimizers must implement the same linear predictor, by Lemma 2. Thus, while (Wk, Wq, Wv, h) is not unique, the function which the resulting transformer implements is unique. We will include this clarification in the final version.

We thank the reviewer for their comments. We emphasize that all the previous works on this topic, including [1], are empirical. von Oswald et al. empirically show that a 1-layer transformer with linear self-attention trained with GD will implement 1 step of gradient descent, while our contribution is to give a theoretical proof.

We note that in the case of a nonlinear target function, while it is clear that the 1-layer transformer will always implement a linear model, it is not necessarily clear that it should implement 1 step of gradient descent.

We thank the reviewer for their detailed review. Below we respond to each of the comments.

The dataset assumption is simple: Theoretical analyses usually need to start with simple cases. Compared to the bigram language model setting of [1], our setting with i.i.d. examples $(x_i, y_i)$ is closer to the few-shot in-context learning setting seen in practice. Also, many previous theoretical works analyze deep linear neural networks to obtain intuition, such as Gunasekar, et al. (2018).

Consider the minimizer of the empirical loss, how the minimizer of GD generalizes: We believe that our analysis of the global minimizer of the population loss is already an interesting finding. We also note that understanding the population loss is, in general, a big step towards understanding the empirical loss, as standard concentration bounds can be applied to bound the difference between the population and empirical losses, given enough samples. There are also several existing works which study the implicit regularization of gradient descent on linear or non-linear neural networks, such Gunasekar, et al. (2018) or Damian, et al. (2021). We consider the minimizer of the empirical loss or the generalization ability of GD to be orthogonal questions for our work.

Damian, et al. (2021) Label Noise SGD Provably Prefers Flat Global Minimizers Gunasekar, et al. (2018) Implicit Bias of Gradient Descent on Linear Convolutional Networks

Experiments: Please see appendix D (newly added during discussion phase), where we include some preliminary experiments with the goal of confirming that a trained linear self-attention layer learns to implement 1 step of GD with our theoretically predicted learning rate.

Is the global minimizer unique?: The global minimizer is essentially unique (up to rescaling of $w, M$, etc.) since as we show in Lemma 2, any global minimizer of the population loss must implement the linear predictor whose weight vector is $\eta X^\top y$. However, this minimizer will not attain zero population loss, due to the output noise. Furthermore, the linear predictor which minimizes the Bayes risk is given by ridge regression, but finding the optimal weight vector for ridge regression requires performing matrix inversion, which cannot be represented by a single linear self-attention layer.

Learning rate v.s. Hessian singular value: The learning rate $\eta$ is roughly the same size as $2/\lambda$ - we will include a proof in the final version. We can simplify the numerator and denominator of $\eta$ using the equality $E[yy^T | X] = XX^T + \sigma^2 I$, as well as the observation from this link https://stats.stackexchange.com/questions/589669/gaussian-fourth-moment-formulas, to find that the learning rate is $\approx 1/(N + d)$, assuming that the output noise variance $\sigma^2$ is very small. Meanwhile, the Hessian of the least-squares problem is $X^T X$, which has a maximum singular value which is $O((\sqrt{d} + \sqrt{N})^2) = O(d + N)$.

Different Population Covariance: We believe our analysis in Section 4 can extend to the case where the weight vector w has the identity covariance.

Minimizer of $g(u)$: Here we are using the fact mentioned in e.g. Section 4.3 of Akyurek, et al. (2022) that the ridge regression weight vector is the linear predictor which minimizes the Bayes risk in the noisy linear regression setting.

Akyurek, et al. (2022). What learning algorithm is in-context learning? Investigations with linear models

Comparison with related works: The work of Takakura and Suzuki theoretically shows that transformers can achieve low approximation error when the target function is shift-equivariant on sequences of infinite length, subject to certain regularity conditions. Bai, et al. propose various constructions through which transformers can solve problems such as Bayesian linear regression and generalized linear models, achieving low test error - they also show that transformers can represent a form of model selection. Guo, et al. similarly provide constructions of transformers which can represent certain algorithms. Compared to these works which give approximation-theoretic constructions, we show that one step of gradient descent is the best predictor that can be implemented by 1 layer of linear self-attention. Interestingly, it follows from our Lemma 2 that this is the unique linear predictor that can be implemented by 1 layer of linear self-attention which minimizes the population loss.

Tarzanagh, et al. [4] analyze the global minima and training dynamics of transformer layers with general data. They show that the global minimum corresponds with a type of max-margin solution, and give conditions under which the training dynamics converge to this global minimum. We note however that this work was released on arXiv at around the same time as ours, and was published in NeurIPS only shortly before the ICLR submission deadline.