In-context Convergence of Transformers

We thank the reviewer for providing the helpful review! We have addressed the reviewer’s helpful comments and revised the paper accordingly. The changes are marked in bright blue color in our revision. Please feel free to check and comment.

Q1: The major weakness is lack of empirical experiments. The linear regression setting according to task distribution and data distribution should not be too challenging to do synthetic experiments for balanced and nonbalanced features similar to Garg et al. or Zhang et al., 2023a. It will be interesting to compared with linear transformers such as Zhang et al. 2023 where in certain task scenarios, linear transformer fails. Another concern is the nonlinearity of transformers are examined on features from linear functions. Would be possible to develop some insights on features from nonlinear functions?

A1: Many thanks for your suggestions! In Appendix A of our revised paper, we have provided experiment results to validate our theoretical results. We summarize our key results as follows: (1) Figure 2 depicts the prediction error for each feature versus the training epochs, and highlights the stage-wise convergence process between the dominant and under-represented features for imbalanced case. (2) Figure 3 illustrates the evolution of attention score heatmap during the training, and verifies our theoretical results on attention score concentration and the multi-phase characterization of training dynamics. Please feel free to check and comment. Our current experiment code has adopted the existing well-packed transformer package, which does not allow change to linear transformer easily. Given this, our handcrafted code on linear transformer may not provide a fair comparison between the two. Note that although Zhang et al. (2023) studied linear attention layers in their theory, their experiments are also on more complex, large, and nonlinear transformers, which do not provide a code on linear transformers either.

Regarding features from nonlinear functions, this is definitely an interesting question to ask. It also appears to be a challenging problem to analyze from theory side, at least not something straightforward. Our preliminary thoughts along this direction are as follows. To consider such nonlinear tasks/functions, (1) some additional layers may be required to learn the non-linear functions; and (2) the softmax weight matrix structure will change and may not admit any simplified zero entries, which is different from the current structure of linear tasks as in Eq. (2).

Q2: The paper does not introduce/define $\Theta(K)$ in the notations.

A2: Thanks. In our revision, we have formally defined $\Theta(K)$ in the notation section as suggested.

Q3: The paper adopts $W^V$ and $W^{KQ}$ from linear SA by Zhang et al. for example by setting $v =1$. Although the authors justify the use of $v=1$ with two reasons. However, the objective from Zhang et al. involves a residual connection term. In this paper, there is no residual connection term involved. Could the authors justify either the use of v =1 as a scaling factor in their objective as in Equation (1) or not including the residual term in Equation (1) and subsequently in equation (3) and (4)?

A3: The residual term in Zhang et al. can also be removed. Both Zhang et al. and our paper choose the loss function to only measure the distance between $y_{query}$ and $\hat{y}_{query}$. The residual term, which is associated with the prediction output $\hat{y}$, equals 0. (See Eq. (3.3) and Eq (3.4)). As a result, the residual term can be removed without any influence on the loss and training dynamics. Furthermore, as we discuss in Remark 1 in our paper, our current specific parameterization (including setting $v=1$) does not lose optimality up to an exponentially small error.

Dear Reviewer xLBH,

As the author-reviewer discussion period ends soon, we will appreciate it if you could check our response to your review comments. This way, if you have further questions and comments, we can still reply before the author-reviewer discussion period ends. If our response resolves your concerns, we kindly ask you to consider raising the rating of our work. Thank you very much for your time and efforts!

We thank the reviewer for providing the helpful review! We have addressed the reviewer’s helpful comments and revised the paper accordingly. The changes are marked in bright blue color in our revision. Please feel free to check our revised paper and comment.

Q1: That all training points are one of $K$ possible vectors. This is not compatible with the training process for linear regression. In general, in which setting we only have $K$ vectors that we sample them in that way.

A1: Our data assumption, where each data point is randomly selected from a predefined set of features, is standard and well accepted in many recent theoretical work, for example, [3] on theoretical deep learning, [4] on transformer theory, and [5] on vision transformer theory. Such a model captures the case where data are sampled from a subspace where those $K$ features serve as basis vectors. Our analysis can be extended to study such a case by allowing the data sample to be correlated. Thus, our assumption indeed aligns with the typical linear regression setting, where the feature vectors are representative of a broader spectrum of data points typically encountered in linear regression scenarios.

Q2: The weight matrices are very constrained and sparsified to match the ones proposed in the construction of [2]. This construction also refers to linear transformers. The authors are multiply the matrices $W_K^\top W_Q$ and train them as one matrix (this also impacts the training dynamics).

A2: We respectfully disagree with the above comments.

(i) Although our weight matrices adopt the construction in [2], our model is NOT linear transformer, but a softmax nonlinear transformer. This is the key difference of our study from that in [2]. Our main contribution lies in characterizing the training dynamics for the more practical softmax based transformer, which is more challenging to analyze than linear transformer.

(ii) In our paper, we have shown that specializing our weight matrix to the construction in [2] does not lose optimality up to an exponential-delay error. Hence, such a choice of weight matrix should not be considered as "very constrained". Further, only one column and one row of our weight matrices, as detailed in Eq.(2), have zero entries, which are not sparse either.

(iii) Training $W_K^\top W_Q$ as one matrix has been adopted in several recent theoretical studies of transformers, for example in [1,2].

Q3: The paper also consider only the attention mechanism and no residual.

A3: The residual term can be removed without any influence on the training dynamics. Please see more detailed discussions in A3 in our response to Reviewer xLBH.

Q4: Could the authors clarify the motivation of this setting? Fixed data points and constraint sparse matrices while considering the task of linear regression.

A4: Our motivation is to understand the learning dynamics of non-linear, softmax-based attention structures trained via gradient descent under in-context learning, which is very different from previous work studying linear attention. The softmax-based attention plays a crucial role in the success of transformers and has not been investigated. To this end, we have introduced some reparameterization and simplifications, which are standard in theoretical work and do not hurt the essence of the problem. We remark here, (1) as explained in A2, our weight matrix is not sparse and does not lose optimality; (2) as explained in A1, the distinct feature assumption is standard to simplify analysis and can be generalized; (3) [2] also adopts the task of linear regression and similar reparameterization, where the only difference is that they study linear attention model, whereas we focus on practically more popular softmax attention model.

Dear Reviewer kkxb,

As the author-reviewer discussion period ends soon, we will appreciate it if you could check our response to your review comments. This way, if you have further questions and comments, we can still reply before the author-reviewer discussion period ends. If our response resolves your concerns, we kindly ask you to consider raising the rating of our work. Thank you very much for your time and efforts!

We thank the reviewer for providing the helpful review! We have addressed the reviewer’s helpful comments and revised the paper accordingly. The changes are marked in bright blue color in our revision. Please feel free to check and comment.

Q1: Are the phases artifacts of the analysis or is it really the case? Could you provide some simulations that demonstrate this multi-phase convergence of gradient descent on in-context learning with transformers?

A1: The division of different phases captures actual transitions of attention score when applying gradient descent to optimize the one-layer transformer. In Appendix A of our revision, we have provided numerical results that validate the multiple phases and other aspects of our theoretical results. The key findings are summarized as follows:

(1) The prediction error for each feature versus the training epochs is shown in Figure 2, which highlights that the training process has stage-wise convergence between the dominant and under-represented features for imbalanced case.

(2) The evolution of attention score heatmap during the training is presented in Figure 3, which demonstrates that the concentration of attention score indeed follows multiple phases during the training by gradient descent.

Please feel free to check and comment.

Q2: The results and analysis seem to abstract away the task diversity (different $w$ vectors). In-context learning is about generalizing to unseen tasks and it looks like the analysis is based on doing gradient descent on $L(\theta)$ which is the expectation over the task vectors $w$ and the features vectors $x$. The setup for in-context learning in Garg et al. 2022 (Figure 1) is more pragmatic and I was expecting an analysis on this. This analysis studies the convergence of gradient descent in minimizing $L(\theta)$ with transformer/self-attention as an underlying model. It's a nice study, given the additional complexity introduced due to softmax, but I think it is not exactly studying in-context learning instead it's a study of convergence of gradient descent on regression with one-layer transformers. The title also seems confusing, it suggests that the paper is about studying the convergence of transformers to the right output during inference time.

A2: We clarify that we indeed study an in-context learning setup. (1) Our transformer setup and loss functions described in Sections 2.2 and 2.3 are exactly the same as the in-context learning setup in the recent studies on in-context learning [1-3]. In particular, the transformer is set up to learn the attention of the query to the tokens in the prompt of the same task. (2) In the last paragraph of Section 3.1 on Page 5, we illustrate the in-context learning ability of unseen tasks. We show that even if the weight $w$ is selected from a different set from the training task space (and is hence an unseen task), the in-context learned model that we characterize in Theorem 3.1 can still well predict the function value for the query in the test prompt.

Q3: Could you simplify the presentation and expand on the main intuition behind the multi-phase convergence? Especially after page 7, instead of the lemmas some simulations/illustrations could be very helpful in understanding the key ideas behind the multi-phase convergence in balanced and imbalanced settings.

A3: Many thanks for your suggestions! In Appendix A in our revision, we have provided experiment results to illustrate and validate the multi-phase convergence. We will move this part to the main body of the paper in our final version.

Q4: Why do you need the features to be orthonormal vectors and be drawn from a finite set? Does the analysis depend heavily on this assumption?

A4: Many thanks for the comments! These assumptions are made to simplify the analysis, so that our characterization of the convergence and the dynamics of the attentions scores won't be over-complicated by non-essential aspects. Following the reviewer's comment, we made an initial exploration of the case where the features are drawn from a subspace with $K$ features serving as basis vectors. This setting allows the features to be generally correlated. Our techniques can be extended to such a setting with further characterization of how correlation among features affects attention coefficients. Also note that in such a setting, we don't assume the set of features is finite; features can be sampled from an infinite set, although the set contains $K$ basis vectors.

Q5: Could you make the dimensions of the main terms clear (e.g. in Definition 3.1) o.w. the readers have to work them out. Minor, the upper case $K$ is used for $W^K$ and also for the number of feature vectors.

A5: Thanks for your suggestions. We have revised them and made them clear accordingly in our revision.

Dear Reviewer VMhr,

As the author-reviewer discussion period ends soon, we will appreciate it if you could check our response to your review comments. This way, if you have further questions and comments, we can still reply before the author-reviewer discussion period ends. If our response resolves your concerns, we kindly ask you to consider raising the rating of our work. Thank you very much for your time and efforts!

The original review by Reviewer AAX6 is copied below to serve as a reference, in case the reviewer will update the review later on.

Summary:

The paper delves into the behavior of transformer architectures, particularly focusing on their convergence properties in various contexts. Transformers, known for their self-attention mechanism, have gained immense popularity in NLP and other domains. The paper's primary objective is to investigate the factors affecting the convergence speed and stability of transformers when trained in different contexts.

Key contributions and discussions of the paper include:

1. An exploration of how the context (e.g., input data distribution, task complexity) impacts the convergence properties of transformers.

2. The introduction of a novel metric to quantify and measure the convergence speed and stability of transformers.

3. A series of experiments that highlight the varying convergence behaviors of transformers across different tasks and datasets.

Soundness: 2 fair

Presentation: 3 good

Contribution: 2 fair

Strengths:

1. Given the widespread use of transformers in various domains, understanding their convergence properties is of paramount importance. This paper addresses this gap by exploring the role of context in transformer convergence.

2. The paper provides a series of well-designed experiments that shed light on how transformers behave across tasks, datasets, and other contextual factors. This empirical evidence strengthens the paper's claims and findings.

3. The paper is well-structured, with clear explanations and visualizations, making it accessible to both experts and those less familiar with transformer architectures.

Weaknesses:

1. While the paper does a good job of examining transformers across various tasks and datasets, there's a limitation in the breadth of architectures studied. Delving into different variants of transformers or comparing with other architectures could have provided a more comprehensive picture.

2. While the introduced metric is innovative, there could be more rigorous validation or comparison against other potential metrics. This would strengthen the metric's claim as a standard measure for convergence properties.

3. The paper could benefit from a discussion on the practical implications of the findings. For instance, how can the insights on convergence be used to improve training methodologies or model selection in real-world applications?

4. A deeper exploration into the individual factors affecting convergence (e.g., model size, training data size, task complexity) could provide a granular understanding and more actionable insights.

5. This paper contains several typographical errors, which detract from the reading experience. For example, there is a missing space before "We" in the 9th line of the summary.

Questions: See Weaknesses.

Flag For Ethics Review: No ethics review needed.

Rating: 5: marginally below the acceptance threshold

Confidence: 3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Code Of Conduct: Yes

I support your behavior in this matter, perhaps we can contact the SAC or PC?

Your actions are quite reasonable for maintaining a good community!