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

Summary

This paper aims at deriving a simplified model for theoretical study of Transformers' optimization. The authors study a pure-attention without softmax Transformer variant (which is called linear attention) on a synthetic regression task. By conducting extensive empirical study, the authors claim that the proposed linear attention Transformer exhibits similar optimization properties that have been observed in the training of vanilla Transformers. Thus, this paper concludes that the proposed linear attention Transformer can be a realistic proxy for understanding Transformer optimization.

Strengths

- The paper is well-written and easy to follow. - Sec. 2 offers a concise and informative overview of existing study on Transformers' optimization. It's helpful for readers to digest the context of this study. - The empirical study and comparisons look solid and convincing. I'm looking forward to optimization theories based on the proposed simplified model.

Weaknesses

- The proposed simplified model does not have feed forward network (FFN), which is yet another difference from standard Transformers. The authors should explicitly point this out (e.g., in the introduction and the formulation of model architecture) in the paper. - I strongly recommend the authors to summarize the detailed training recipes (e.g., hyperparameters) in a table and release the code for calculating relevant quantities (e.g., gradient error, robust condition numbers) to facilitate reproducibility and follow-up research. - In the beginning of Sec. 3.1, $x^{(i)}$ is defined as a row vector. But in the input matrix it seems that one has to interpret $x^{(i)}$ as a column matrix. Moreover, the matrix shape here is very confusing. The input of $F$ is $(n+1)\times (d+1)$, while $Z_0$ is $(d+1)\times (n+1)$. I suggest following the "_Attention Is All You Need_" paper and using shape $(n+1)\times (d+1)$ consistently. The authors should fix the notation and formulas here. **Minor issues.** - Sec. 2.1: _The features defined above are typically of ..._ $\to$ _The concepts defined above are typically of ..._ - I recommend the authors to capitalize the first letter of "Transformers".

Questions

The question here is just for clarification. - What is the box in the upper right of the left 3 plots in Fig. 3?

Rating

6: marginally above the acceptance threshold

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

3 good

Presentation

3 good

Contribution

3 good

Summary

This paper offers empirical evidence highlighting parallels between shallow linearized transformer models and the dynamics observed during practical Transformer training. These properties span various well-known aspects: the superiority of adaptive methods over SGD, the presence of heavy-tailed stochastic gradient noise, a robust condition number characterizing the optimization landscape, and the directional smoothness gap. The findings advocate for the potential of simple linear transformers as an insightful and accessible model to elucidate the mystery of Transformer optimization trajectories.

Strengths

The presented work is overall clear and of notable quality. Its innovative aspect is the unexpected empirical evidence indicating that, under certain conditions, shallow linear transformers can closely resemble conventional attention models, revealing a significant similarity between the simple model and its real-world analogs. Given the historical challenges associated with analyzing attention mechanisms involving softmax or other nonlinearities, it is commendable to identify such marked resemblances under broad observations using a simplified framework. This study illuminates the theoretical analysis for optimization on linear transformers, which makes a valuable contribution to the literature.

Weaknesses

The paper is somewhat limited in its data distribution and optimization objective setting to support its conclusion. The data distribution discussed in this setting focuses on in-context linear regression tasks—a largely statistical setting, rather than one that closely resembles language. This might be attributed to the limited expressivity of linear transformers. I fully understand that the authors are trying to *find the simplest abstraction* as a representative of the transformer’s landscape. However, since the landscape is likely heavily dependent on the data distribution, it would be more convincing to see a variety of data models that exhibit the same phenomenon in simple models. Furthermore, it's difficult to assert that the linearized model is sufficient to understand Transformer optimization, especially since the properties verified in this paper represent only a portion of the optimization narrative.

Questions

- Are there other possible tasks that can be expressed by linear transformers? Do those tasks on simple linear transformers also exhibit those behaviors? - If we add other linearities (such as softmax) to the shallow neural networks, can you also show those properties along the training trajectory?

Rating

8: accept, good paper

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.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Summary

This paper identifies "use shallow linear Transformers for regression task" as a simple setup that serves as an effective mathematical abstraction for studying Transformer optimization. - The linear Transformer has attention (inner product) and residual connection, without softmax and MLP. For attention, a single $Q$ matrix is used to replace the product $W_Q^\top W_K$. - The aim is to find the _simplest_ setup as a surrogate for studying optimization: this paper uses a 3-layer network with 5-dim covariate, on the linear regression task.

Strengths

- The simple setup is able to reproduce several characteristics of the training dynamics, which are not captured by canonical optimization theory. These include heavy-tailed gradient noise, robust condition number (larger for SGD than for Adam), as well as better directional and generalized smoothness from Adam. - The synthetic setup allows for more precise control. - The paper varies the number of layers, data distribution, and context length. - It's affordable to grid search over 10 learning rates.

Weaknesses

- The main claim is that the simplified setup can be used as a proxy for larger scale systems. However, some connections are drawn by having "qualitatively similar" results. - For example it's known that higher-order moments of the gradient noises matter, which cannot be captured by examining plots visually. - e.g. How to judge whether the noises are similar based on Fig 3, or that 8-layer is qualitatively similar to 3-layer based on Fig 8? - The current empirical results are for confirming existing findings, but it's unclear what new discoveries can be enabled by this simplified setup.

Questions

- What are some phenomena that are not able to be captured in this simplified setup? - e.g. using the $Q$ parameterization rather than $W_Q^\top W_K$ might lead to different implicit regularization effects. - There might be effects about architectural choices, e.g. consequences and limitations from the softmax. - In general, it would be helpful to better clarify what is consider in-scope and what is not. - Could you give some examples of new discovery or actionable advice from the simple setup? - The current setup is entirely linear; if the data is non-linear, would there be similar optimization characteristics?

Rating

6: marginally above 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.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Summary

This paper studies the property of the shallow linear Transformer model. The evaluated phenomenon includes comparing Adam and SGD, heavy-tailed gradient noise, condition number, smoothness, etc. The experimental results show that the linear Transformer model reproduces the phenomena that have been observed for full Transformers. The results also show that more heavy-tailed data distribution and more layers can enhance the conclusion. -------------------------------------------------------- **After rebuttal**: Thank you for your response. I am sorry for the late reply. I am satisfied with the feedback and revisions. I will keep my rating as 6 and support you. A minor point and suggestion: Please clarify what the four figures in Fig. 15 are to be compared within the main body. Later, I realize you are mentioning some figures on page 3. It is not a big issue anyway.

Strengths

This work is novel in terms of new experiments on linear transformers from many aspects of evaluation. The paper is well-written and clear to follow. The studied problem is interesting to the community, from my understanding. The designed experiments are concise to support the conclusion.

Weaknesses

1. This work lacks theoretical understanding or explanation after the experiments in Section 3.2, 4.1, 4.2. 2. The limitation is not discussed. One thing that should be emphasized in the introduction or the abstract is that the experiments are linear regression, which is a good fit for linear Transformers.

Questions

1. One conclusion may not be obvious for linear Transformers compared with softmax Transformers. That is, the attention weights are more concentrated (even sparsely) on some key features/tokens after training, which is observed in several existing works [Li et al., 2023a, Li et al., 2023b, Oymak et al., 2023] for softmax Transformers. Can the authors provide a comparison empirically on this? Also, I think it is better to cover such a discussion in the revision. Oymak et al., 2023, "On the Role of Attention in Prompt-tuning. " Li et al., 2023a, "A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity." Li et al., 2023b, "How do transformers learn topic structure: Towards a mechanistic understanding."

Rating

6: marginally above the acceptance threshold

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

3 good

Presentation

3 good

Contribution

3 good