Paper: In-context Convergence of Transformers (ICLR-2024)

Reviewer comment (Experiments):
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.

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

Reviewer comment (Experiments):
It will be interesting to compare with linear transformers such as Zhang et al. 2023 where in certain task scenarios, linear transformer fails.

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

Reviewer comment (Experiments):
Another concern is the nonlinearity of transformers are examined on features from linear functions. Would it be possible to develop some insights on features from nonlinear functions?

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

Reviewer comment (Writing):
The paper does not introduce/define Θ(K) in the notations.

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

Reviewer comment (Theory):
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, 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)?

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

Reviewer comment (Experiments):
Did the authors investigate other tasks with softmax transformer as discussed by the Garg et al. paper?

Author response:
It is a very interesting future direction to consider tasks of a broader class of functions including nonlinear functions. We are getting started to explore some simple nonlinear cases, for example, quadratic functions or sparse linear regressions. Even those problems seem to be very challenging to characterize analytically.

We thank the reviewer again for the helpful comments and suggestions for our work. Given that we have provided the requested experiments (which seems to be the main concern of the reviewer), if our response also resolves your other concerns to a satisfactory level, we kindly ask the reviewer to consider raising the rating of our work. Certainly, we are more than happy to answer any further questions that you may have.

Reviewer comment (Theory):
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.

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

Reviewer comment (Experiments):
The paper also consider only the attention mechanism and no residual.

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

Reviewer comment (Experiments):
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?

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

Reviewer comment (Theory):
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.

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

Reviewer comment (Presentation):
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.

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

Reviewer comment (Theory):
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?

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

Reviewer comment (Writing):
Could you make the dimensions of the main terms clear (e.g. in Definition 3.1) otherwise the readers have to work them out.

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

Reviewer comment (Theory):
The bounds on $T^*$ have quadratic dependence on $K$ (the number of distinct feature vectors). Is this order of dependence necessary? What would happen if the features are drawn from a continuous distribution, say multi-variate gaussian?

Author response:
**A6:** Great question! Our hypothesis is that the bound will depend on the underlying dimension of the space from which features are sampled, rather than the actual number of distinct feature vectors in the space. For example, if features are drawn from a $d$-dimensional multi-variate Gaussian distribution with rank-$K$ covariance matrix, then the bound will depend on $K$. More generally, for the continuous feature space with $K$ basis vectors (as we discuss in A4), the bound will depend on $K$, which is the dimension of the space, not the actual number of features, which is infinite. Of course, this hypothesis will require rigorous mathematical proof, for which our current proof techniques seem to support. We leave the detailed proof as future work.