Linear Transformers are Versatile In-Context Learners

Summary

This paper proves that linear transformer layers maintain a weight vector for implicit linear regressions, including a more challenging scenario where data is corrupted with different levels of noise. In the theoretical analysis, this paper shows the intrinsic mechanism of the gradient scent in linear attention, where implicit variants play various roles in parameter update. Besides, this paper shows that simplified linear attention with diagonal transformation also maintains a powerful GD behavior. The experiments demonstrate the flexibility of linear transformers which outperforms GD++ as well as other linear regression solutions. This research promotes the understanding of Transformer weights and implicit learning capabilities of attention-based models.

Strengths

1. The theoretical analysis in this paper is sound and insightful for understanding the behavior of linear attention. 2. This paper broadens the analysis framework from fixed noise linear regression to flexible noise, making a step forward to more complex problems. 3. The experiment settings can prove the mathematical formulation in this paper.

Weaknesses

The analysis framework is still on linear regression and linear attention with simplified modifications, which is similar to previous references. The theory can not be generalized to other architectures and context settings.

Questions

1. How do you think about different linear attention variants in theory? For example, the most trivial one is $QK^TV$, but there are many improved versions, such as RetNet, SSM, and GLA. Do these architectures bring stronger modeling capability in theory? 2. With linear regression layer-wisely, does your claim that linear Transformer with FFN only maintains linear regression still work? For example, in Theorem 4.2, it is easy to build a quadratic FFN to make more complex computations.

Rating

6

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

The limitation is discussed in the conclusion section.

Summary

In this paper, the authors study linear transformers trained on linear regression problems and prove that each layer of every linear transformer maintains a weight vector for an underlying linear regression problem. Furthermore, the authors consider the mixed linear regression problem with varying noise levels and empirically demonstrate the flexibility of linear transformers for this problem.

Strengths

**Overall well-presented** For someone who is familiar with the field, the presentation is certainly very good. For others, it might not be super obvious that the linear transformer is trained on a bunch of other linear datasets with different noise levels. Finally, I didn't find the abbreviations before Section 2.2 particularly useful. They don't save a lot of space, but the reader has to go back to it. Not clear why the authors chose to use those. **Relatively Interesting Insight** I found it somewhat interesting that the transformer is doing so well in the varying noise levels scenario. But then again, might be not surprising given it had a lot of training data.

Weaknesses

**Incremental Nature** The authors build on top of a lot of existing literature. While this is not necessarily bad, the delta compared to previous work might be too small. Since I'm not an expert in this domain, I'll have to rely on the opinion of the other reviewers to confirm my understanding. **Unclear Baselines** I'm not sure what the baselines are for. I understand that it is useful to give an idea how well the transformer does in the task. However, the author make it look as if it was important to beat them. Furthermore, the authors claim that these methods have an advantage given that they can use matrix inversion. In my understanding, the transformer is trained on much more data which could technically used to improve the noise-level estimation for the linear models and improve their performance as well. To me it is not really clear why the authors are so defensive in this experimentation. The authors highlight that the transformer does so much better in the varying noise levels scenario compared to the linear models. However, there was little to no effort to make the linear models good in this scenario in the first place.

Questions

What are potential practical uses of this insight?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

n/a

Summary

This paper demonstrates that linear transformers maintain a weight vector for an implicit linear regression problem. The authors provide theoretical analysis showing that (1) linear transformers maintain a linear regression model at every block and (2) a diagonal parameterization of attention heads does not compromise the expressive power of the model. The authors support their theoretical findings with experimental validation.

Strengths

1. This paper extends von Oswald et al. (2023) by further investigating that each linear transformer layer maintains a weight vector that can be used for regression problems. Moreover, they show that GD++ is a second-order optimization alg for the least square problem. 2. Experiments on three different parameterization matrices are consistent with the analysis.

Weaknesses

1. Limited empirical validation of Section 5.3: The results in Section 5.3 suggest that the update of y occurs every two steps. However, the paper lacks empirical studies to illustrate this phenomenon, particularly in a diagonal parameterization scenario. Can the authors provide: a) Concrete experiments demonstrating this two-step update pattern? b) Visualizations or quantitative analyses that show how this behavior manifests in practice?

Questions

1. Comparison of attention matrices in full and diagonal parameterizations: Given the similar loss patterns shown in Figures 2 and 3 for full and diagonal parameterizations, it raises questions about the final structure of the attention matrices. Specifically, is it possible that the full matrices converge to a near-diagonal structure during training? If so, what implications does this have for the choice between full and diagonal parameterizations in practice? If not, how do the differences in attention matrices reconcile with the similar performance observed?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

The paper tries to understand the reasons of the strong performance of Transformers. The authors study linear Transformers trained on linear regression problems. Moreover, the authors explores the problem of regression where the labels have variable noise levels.

Strengths

- The problems studied are important, especially with the wide adoption of Transformers. - The case of regression with variable noise is interesting - Theoretical analysis is presented

Weaknesses

- Although the focus is on analyzing Transformers, the paper could have benefited from adding more empirical analysis.

Questions

See above

Rating

4

Confidence

2

Soundness

2

Presentation

2

Contribution

2

Limitations

NA

please respond to rebuttal

Dear reviewers, The author-reviewer discussion period will end soon. Please make sure you read the authors' rebuttal and respond to it. If you have additional questions after reading the rebuttal please discuss with the authors. For those who have done so, thank you! AC