Linear attention is (maybe) all you need (to understand Transformer optimization)

Thank you for your fruitful comments! We revised our paper based on your comments and all the edits are highlighted in orange.

• We added in our revision that "our setting doesn't have FFN" in the places that you suggested.

• For the dimensions in the beginning of Section 3.1, those were typos. Thank you for catching them. We now corrected them.

• We also corrected the "minor issues" raised by you; we also capitalize "Transformers" throughout the manuscript.

• For the upper right box in Fig. 3 are the quantile-quantile plots (QQ-plots) to see how heavy-tailed the distributions are. We added the QQ plots for our experiments as well as descriptions in our paper.

• We added the summary of hyperparameter choice in Appendix A. Also, we attached the code in the supplementary material.

Thank you for appreciating the values of our work and your constructive comments!

Admittedly, we do not provide theoretical understanding of the phenomenon, but what we claim as our main contribution is the discovery of a simple abstract setting where the key characteristics of Transformer optimization persists, and we believe that our findings could serve as a useful testbed for the future theoretical investigations.

In order to address your concern regarding the limitation, we conducted an additional set of experiments for a nonlinear regression. In particular, we consider the case where the covariates $x^{(i)}$ are distorted by a ReLU network $MLP$, i.e., the responses are generated as $y^{(i)} = \langle w_\star , MLP(x^{(i)}) \rangle$. In order to cope with the nonlinearity we also added a MLP in our Transformer architecture. As you can see from Appendix B, the results for the nonlinear regression are largely similar to the linear regression. We hope that our additional set of experiments addresses some of your concern.

Moreover, thank you for point us to the additional relevant works regarding the softmax Transformers. Although it's not directly related to the main message of our paper, we answer your question regarding the attention scores as below. As a summary, for the linear Transformer for linear regression, we need to combine appropriate covariates $x^{(i)}$, $i=1,2,...,n$ to figure out the missing prediction $y^{(n+1)}$ for $x^{(n+1)}$. So, the sparsity of the attention weights will depend on instances. For instance, when the covariates $x^{(i)}$ are orthogonal, we empirically observe that the attention scores are sparse both for linear and softmax Transformers, which is consistent with the results from the previous papers. As per your suggestion, we added a discussion regarding the softmax Transformer on Page 6.

Thank you for appreciating the values of our work and constructive comments! We also believe that having a simplified framework is valuable for the future rigorous study on Transformer optimization.

We also appreciate your comment "I fully understand that the authors are trying to find the simplest abstraction as a representative of the transformer’s landscape." That is precisely the main scope of this work.

To address your comments about "the landscape being likely dependent on the data distribution," in Appendix B, we conducted an additional set of experiments for a nonlinear regression. In particular, we consider the case where the covariates $x^{(i)}$ are distorted by a ReLU network $MLP$, i.e., the responses are generated as $y^{(i)} = \langle w_\star , MLP(x^{(i)}) \rangle$. In order to cope with the nonlinearity we also added a MLP in our Transformer architecture. As you can see from Appendix B, the results for the nonlinear regression are largely similar to the linear regression. Although we agree that nonlinear regression also can't fully explain the practical language modeling, we hope that our additional set of experiments addresses your concern.

Thank you for your constructive comments. We agree with your comment that our simple setting lets us do more precise controlled experiments.

Let us address your points one by one. We revised our paper based on your comments and all the edits are highlighted in orange.

-- More quantitative way of checking noise distribution? Following [Kunstner et al., 2023], we added the quantile-quantile plot (QQ-plot) in our revision. A QQ-plot is a plot of the quantiles of two distributions against each other. In particular, we add QQ-plots that compare the quantiles of the stochastic gradient noise (y-axis) against those of its best-fit Gaussian distribution. One can see that QQ-plots go beyond $y=x$ line toward the right, suggesting that the stochastic noises are indeed heavy-tailed. Also, for the heavy-tailed covariates experiments and the more layers experiments, we now can also see clearly for which cases the noises become more heavy-tailed.

-- Regarding the $W_Q W_K$ parametrization. Thank you for your question. In the left plot of Figure 8, we added a result for the standard Transformer parameterization of $W_Q W_K$, and observed that the result looks similar to our simplified parameterization.

-- Regarding softmax activation. We mostly considered linear Transformer since it's better suited for the task of solving linear regression in-context as shown in Figure 7.

-- Clarifying the scope of this work As per your suggestion, we further clarified the main scope of this work in our introduction, by saying "....development of efficient optimization methods for Transformers. However, such directions are out-of-scope of this work, and left for future work."

-- Our setting entirely linear? We would like to clarify that despite removal of nonlinearity, the linear Transformer is still highly nonlinear function due to the attention mechanism. In particular, even the output of single layer linear Transformer is a multivariate 3rd order polynomial function of the input.

[Updated Response]: In order to address your concerns about the linearity in the data distribution, in our newest revision, we conducted an additional set of experiments for the nonlinear regression, and it's presented in Appendix B.