Transformers learn to implement preconditioned gradient descent for in-context learning

Summary

The paper tries to tackle an important and relevant problem in LLM training: whether in-context learning properties can be explained by a form of gradient descent? For the setting, the authors have chosen to focus specifically on the linear transformer with a single or multiple layers with an input that is given by the data generated using a linear model. Within this setting, the authors were able to show a series of results for various specially constructed transformers. Specifically, they look at a (1) single-layer transformer, (2) multi-layer with sparse parameters and (3) multi-layer with more general parameters. For (1), they have provided a characterization of the unique global optimum of a single-layer transformer with linear attentions. They showed that foer the optimal parameters found, the transformer implements a single step of preconditioned gradient descent. For (2), they have shown that global minimizers of the training objective of a two-layer linear transformer over isotropic regression instances. For the most general scenario (3), they found that certain critical points implement preconditioning using inverse-data-covariance of the input. Experimentally, the authors have confirmed using a toy-example that optimizing a three layer model indeed results in a minimizer that converges to an inverse of a covariance matrix.

Strengths

The paper provides an interesting analysis that can be useful to understand the intriguing properties of in-context learning. It studies a series of transformer models, ranging from a simple single layer to a more general multi-layer ones and shows that for all the settings the in-context learning implements the preconditioned gradient descent. While by no means comprehensive, this paper can be useful as a stepping stone for any researcher that wants to further study intriguing properties of in-context learning. The paper is easy to read and easy to understand. The assumptions and the math mainly seem to check out.

Weaknesses

Generally, the paper restricts itself to a very narrow setting of a specific data that is centered at zero, linear model generating the labels, linear transformer, specific sparsity of the parameter matrix. The transformers that they study are linear with a single head only. While I do appreciate the value of the analysis and I understand that this is just a step of many, it would be great to include a short paragraph outlining what the authors think about generalization of the proposed approach when the assumptions are removed. Also, given the series of the assumptions, I would make sure that they are stated early on in the abstract and introduction. I also think that the experimental validation can be slightly improved by considering different choice of layers, dimension of x, conditioning factor of the input etc. It would also be great to have an empirical plot demonstrating how close A is to the Gram matrix as the number of points increase (the results from the end of section 3)

Questions

The equations for the gradient do not immediately pop out from the eq 7. I would encourage the authors to expose the equations for the gradient to the reader (maybe after eq 5) so that the connection between the gradient and the preconditioner would become more apparent. What are the properties of A in sections 3 and 4? Can a transformer find a positive definite matrix through training? For the evaluation of the Theorem 4 and 5, it is a bit hard to interpret the numbers, since the used distance metric is close to zero, but not exactly zero. The obtained value of ~0.1-0.2 is quite hard to interpret, apart from the fact that it is smaller that the distance with respect to the identity. Also, the metric that the authors are using involves the minima over the space of scalars $\alpha$. It would be interesting if the authors can actually provide the scalars $a_i$ from Theorem 4 that they have found empirically. Do these scalars remain the same over multiple restarts of the algorithm? How come the same $\Sigma$ matrix is used for both data and the weights (theorem 4 and 5)? It doesn't sound very realistic that the data and the weights use the same covariance matrix. Would be great to add more information to Section 4.1 about adaptive coordinate-wise step sizes. If I understood correctly, Theorem 3 talks about the existence of the global minimizer under certain conditions. How exactly does adaptive coordinate-wise step sizes help find this solution? I would encourage the authors to add a simple table at the end of the introduction clearly stating the assumptions that each section has. As written, It is quite hard to understand, especially when the notation change from section to section. Small comments: - Typo l.5: “Can transformers can” - L.83: dimension of Z doesn’t match the dimension of W_k, and W_q in eq (1). - In Lemma 2 z^(1) etc are mentioned, but not properly defined nor used. - Gamma_k in l.155 plays a similar role to a_i from l.184. Maybe call both of these variables gamma? - In fig.1 (d) the log(loss) converges to -1.0, which is not a zero value, as the text suggests in l.213. - How come the authors choose L-BFGS as an optimizer? It is a fairly specialized algorithm that is not often used these days. What properties of L-BFGS are desirable in the settings that authors used it for? Can the same results be achieved with SGD or Adam?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

It would be important to highlight limitations of the analysis. E.g. same covariance matrix, singe-head linear attention...

Summary

This paper studies linear transformer and validates its ability in solving linear problems via applying conditioned gradient descent. The authors firstly state that there exists a unique global minimum for single attention layer. Then they explore multi-layer linear transformer and demonstrate that stationary points can be found that perform preconditioned gradient descent algorithm.

Strengths

This paper theoretically studies that a $k$-layer linear transformer can perform as a $k$-step preconditioned gradient descent algorithm when forwarding a linear ICL sequence.

Weaknesses

1. Related work is absent in this paper. 2. This paper simplifies the problem a lot by treating all components as linear. However, for a standard transformer, the softmax and activation functions are two key components and introduce nonlinearity. A discussion on the relationship between linear and standard transformers is missing. 3. Compared to the previous work (von Oswald et al., 2022), the novelty is insufficient except the introduction of non-isotropic data distribution. 4. Although simulations are provided in this paper, it is not sufficient. It would be interesting to know: How performance varies with the number of layers? How a $k$-layer linear transformer performs compared with a direct gradient descent approach with different step sizes or a standard transformer (during training or after sufficiently trained)? ...

Questions

1. Although $w^\star$ is well-defined in the paper, it is not clear how $w$ relates to the model parameters or data. minor: 2. Presuming you are studying least-squares loss, Eq.(5) may miss a square symbol. 3. There seem to be notation typos: $k\to i$ in Eq. (9) and $A_k\to A_i$ in Lemma 2.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

While the authors did not explicitly discuss limitations, they provide potential future directions in the discussion section that are related to their work and have not been explored.

Summary

LLMs trained with transformers have shown remarkable abilities that have led to a lot of recent theoretical research towards understanding the origins of these abilities. In context learning is one such ability that refers to a model's ability to adapt in-context to new tasks without weight updates. In the work of Garg et al., the authors proposed a simple and powerful data generation process to understand the mechansims behind in-context learning. In Garg et al., it was empirically shown that transformers can learn to imitate ordinary least squares algorithm. Since that work several works have taken steps to understand why and how transformers implement these algorithms. In recent work from Oswald et al., it was empirically shown that linear transformers implement one step of (modified) gradient descent per self-attention layer. In this work, the authors take an important step to explaining why do transformers implement gradient descent. The authors show that for shallow transformers pre-conditioned gradient descent is the global minima of the in-context training loss. For multi-layer transformers, the authors establish that pre-conditioned gradient descent form the critical points of the loss landscape.

Strengths

This is a very good paper. I was very happy to read it as it solves an important unresolved problem that has to come to fore in the recent investigations on in-context learning. On the different axes, I would rate the paper as follows. Originality: I really liked the approach that the authors have taken. It is quite a neat idea to consider sparsity driven constraints that permit tractability. Quality: This is a very good paper. The results resolve some very important questions towards understanding ICL. Clarity: The paper is very well written and a joy to read. Significance: . This paper serves as a crucial step towards the goal of understanding ICL is one of the central problems towards demistifying LLMs

Weaknesses

I have some concerns that I would highly appreciate if authors can throw light on in the paper. i) The authors have studied linear attention mechanisms, which I very well understand are themselves non-trivial to understand. If we were to consider softmax attention and carry out the same numerical experiments with sparsity constraints of the form of equation 8, then do we continue to achieve global minima. If authors can add such an experiment, then this can serve as very solid motivation for someone to study ICL with softmax attention (and sparsity constraint) and develop theory for the same. ii) The authors reparametrize the query and key matrix as $Q= W_k^{\top} W_q$. This is not standard in transformers. What would happen if we were to consider the standard parametrization. Do we still obtain stationary points in the general case? Does the update dynamics still converge to the global minima (the verification step of Theorem 4 and 5)? iii) The authors use LBFGS to learn the parameters. Do the results change under standard gradient descent or SGD or Adam? iv) In the most general case of linear transformers, when $P$ and $Q$ are not constrained to be sparse at all, do the authors conjecture that similar results as the ones shown in the paper continue to hold.

Questions

Please see the weakness section.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The authors discuss future directions, which stem from the limitations of the current study. The authors can elaborate further on limitations of the theory further and highlight some other potential gaps that can be addressed in future work.

Summary

The paper presents an in-depth analysis of in-context learning (ICL) in transformers, particularly focusing on linear transformers. The authors propose a novel interpretation of ICL as an implicit form of '**preconditioned**' gradient descent. They further extend this analysis to multi-layer transformers and provide insights into the optimal parameters for the in-context loss under certain constraints. The paper also includes empirical results on synthetic datasets to support their findings.

Strengths

1. The paper provides a novel perspective -- the preconditioning/adaptive perspective -- on in-context learning in transformers, which could potentially lead to new ways of understanding and improving these models. Given the growing interest in understanding ICL from a theoretical perspective, the `*preconditioning gradient descent*' can be served as an important building block. This preconditioning perspective is also new in the literature (as far as I know). 2. The authors provide a rigorous mathematical analysis to support their claims, including single and multi-layer cases. The analysis tools developed in this paper could be helpful for future research on understanding ICL. 3. The paper is well-structured and includes visualizations to help understand the concepts and results.

Weaknesses

1. (minor) In the experiments, I did not find the ICL experiments --- x-axis represents the number of ICL samples, and the y-axis represents the test loss (similar to the experiments in Figure 2 of Garg et al. (2022)). It would be interesting to check those results besides the current ones on training dynamics.

Questions

1. In the current setup, I find there is no MLP block in the architecture, empirically that seems to be an important component in the transformer and ICL, could the authors provide more discussions/explanations on the roll of MLP (could also be listed as future directions)? 2. Typo: missing '.' at the end of line189.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

Limitations

NA

Thanks

I have read the response from the authors. The response has clarified some of the confusions I had. I would encourage authors to make changes to the paper in light of the above discussions. I am happy to maintain my score.