Linear Transformers are Versatile In-Context Learners

We thank the reviewer for raising these insightful questions, which highlight the importance of generalization of our findings. While a complete exploration of these aspects is beyond the scope of our current work, we would like to provide some initial thoughts and discuss potential future directions.

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?

While we focused on simple attention for clarity and ease of analysis, these more advanced architectures could indeed offer stronger modeling capabilities in theory. They can implicitly capture higher-order feature interactions. These architectures often employ techniques like kernel functions or randomization, which could allow them to implicitly capture higher-order interactions and learn richer function classes. More specifically, we believe our Theorem 4.1/4.2 would generalize to the linear version of RetNet, but would not generalize to SSM and GLA (because SSM is time-varying and GLA has nonlinearities in the form of gates). Even in the setting of RetNet, the specific algorithms discovered could be more sophisticated. Exploring how our findings extend to these richer attention mechanisms is a promising direction for future work.

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.

Indeed, our current analysis considers linear transformers without FFNs. This simplification allows us to isolate the core behavior of the attention mechanism. With nonlinear FFNs, Theorem 4.2 no longer holds. A quadratic FFN can indeed introduce higher-order relationships, however the output of each attention layer is still a linear combination of its inputs (though the weights might be non-linear functions of the input). Exploring the complex interaction between attention and FFNs in the context of implicit optimization is crucial for future research and may require more advanced tools and techniques.

Thanks for your response! I'd like to see your future work to explore the sophisticated modeling nature of the whole linear models, which is closer to real data and more valuable for the interpretation of LLMs. I will keep my score.

We thank the reviewer to take their time to evaluate the paper and for valuable suggestions on improving the presentation. Indeed, the linear transformer is trained on various generated sequences of noisy linear regression. We will make it more obvious in the introduction and preliminary section. We will also move the notation from the end of Section 2.1 closer to the Theorems where they are actually being used.

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.

When training the model, the amount of the training data is just one factor. Another big factor to consider is the capacity of the model. Not every model can fit any given function. The original motivation for this paper came from the surprising observation that Transformers have the capacity to solve certain problems provided in-context. Previous research has demonstrated that for a single noise level this ability might be due to the transformers implementing a form of gradient-based optimization on an implicitly defined objective function. For multiple noise levels previous work (Bai et al.) relied on complicated constructions with large networks and nonlinearities. In this paper we demonstrate that even linear transformers with digaonal attention metrics have enough capacity to solve quite complex and non-obvious problems or noisy linear regression.

..the transformer is trained on much more data which could technically be 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.

Thank you for this observation. What we were trying to convey is not defensiveness, but a genuine surprise by the results that we have observed! The problem of linear regression with variable noise level, while simply stated, is practical and complex. The fact that a simple linear transformer can learn quite a sophisticated algorithm that works on par or better than many baselines that us, humans, can come up with, is quite surprising.

What are potential practical uses of this insight?

We think that there are several implications of our work. One, the algorithm discovered by the linear transformer, with its adaptive rescaling and momentum-like behavior, could be directly applied to real-world noisy linear regression problems in domains such as robust control, time series analysis, or finance. Second, our work strengthens the evidence for transformers' ability to implicitly learn sophisticated algorithms, opening up exciting possibilities for automated algorithm discovery in other machine learning tasks.

It would be interesting to see how variations in problem complexity and structure affect both the performance of the transformer and its ability to discover novel algorithms. This could involve exploring different data distributions, underlying function classes, and noise structures, ultimately leading to a deeper understanding of the factors influencing algorithm discovery in transformers and potentially uncovering a broader class of implicit algorithms with practical applications in various domains.

We thank the reviewer for their insightful comments and valuable feedback. We address their specific points below:

Empirical validation of Section 5.3 is lacking that the update of y occurs every two steps.

Indeed, since $w_{xy}$ controls the how the current $y_t^l$ prediction affects the prediction of $x_t^{l+1}$ for the next layer, it doesn't change the prediction $y_t^{l+1}$. It takes at least two layers for this effect to reach the prediction $y_t^{l+2}$. However, it is quite challenging to empirically isolate the effect of every component, since all of the elements work simultaneously at every layer. All four terms per layer are helping to improve the loss and it is not-trivial to show the exact effect of every term empirically.

Comparison of attention matrices in full and diagonal parameterizations.

We were quite surprised to see that the diagonal parametrization performs almost identical to the full one. As the reviewer has predicted, the learned attention matrices for the full case (see example in the PDF attached above) converge to the near-diagonal matrix, but with some additional structure as well. Our preliminary experiments with diagonal plus low-rank parameterizations yielded similar results to the full attention mechanism. However, given the comparable performance and interpretability of the diagonal approach, we chose to focus on this simpler diagonal model.

In practice, the choice between full and diagonal parametrizations involves several factors. The diagonal approximation is easier to interpret (only 4 terms per layer!), is faster and works just as well as the full within the settings we consider. However, it is quite possible that for different distributions of $x$, $w$ or $\sigma$ beyond the ones that we consider in our paper, the difference between full and diagonal models would be more pronounced. We plan to investigate this further in future work.

We will make sure to include the attention weights matrices as well as the consideration above to the final version of the paper if accepted.

We thank the reviewer for taking time to look at our paper. We believe that the empirical evaluation provided in the paper is thorough and is appropriate to cover the claims and contributions presented in the paper. We would love to know what kind of empirical evaluation the reviewer has in mind that we should add?