Paper: One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention (ICLR-2024)

Reviewer comment (Theory):
I am not sure about the practical insight from the theoretical analysis of this work.

Author response:
**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.

Reviewer comment (Experiments):
Can you verify that the $\eta$ in Theorem 1,2 are as predicted by experiments? Specifically, can you show how $\sigma^2$ in Theorem 2 affects $\eta$ by experiments?

Author response:
**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.

Reviewer comment (Theory):
I don't know why fully-connected neural networks satisfy Assumption 1 (1). Can you provide a proof for this claim?

Author response:
**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$.

Reviewer comment (Theory):
The usage of the statement "(Wk, Wq, Wv, h) is a global minimizer" in Theorem 1 raises questions about the specifics of this minimization process. Further clarification or details regarding this construction might be beneficial for readers.

Author response:
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.

Reviewer comment (Theory):
The dataset assumption is simple and the authors only considered i.i.d. data sequence with linear teacher model. This setting helps the analysis but may not be able to fully capture the properties of the self-attention layer. Besides, the proofs rely on the rotational invariance of Gaussian distribution. It would be interesting to generalize the results in non-Gaussian datasets or consider more dependent structures in the data sequence, like the Bigram language model in [1].

Author response:
**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).

Reviewer comment (Theory):
This paper focuses on the global minimizer of the population square loss of the self-attention layer which simplifies the analysis. It would be natural to consider the minimizer of the empirical loss during the pre-training process and how the minimizer of the GD or stochastic GD with finite step sizes generalizes in the test point.

Author response:
**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

Reviewer comment (Theory):
How large the learning rate $\eta$ is? Following the remark after Theorem 1, we know the global minimizer is equivalent to a step gradient descent on the empirical loss of the least squares problem with zero initialization and learning rate $\eta$ defined in Theorem 1. How large $\eta$ is, compared with the largest eigenvalue of the Hessian matrix $H$ of this least squares problem? Is it just close to or larger than the maximal learning rate $2/\lambda_{\max}(H)$?

Author response:
**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)$.

Reviewer comment (Theory):
In Section 4, when we consider data covariance in $\Sigma$, why do we renormalize back by $w\sim\mathcal{N}(0,\Sigma^{-1})$? Can we consider $w$ has another different population covariance like [2]?

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

Reviewer comment (Theory):
In the proof of Lemma 1, after Eq. (29), why is the minimizer of $g(u)$ given by $\hat{w}_{\tilde{D}}$? Here, for $g(u)$, you only have one data point. Can you explain more?

Author response:
**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

Reviewer comment (Novelty):
It may be worthy to mention or compare with some of the references among [1] and [3-7].

Author response:
**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.