How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression

Summary

This paper empirically studies how the different heads works across different layers of the transformer: while the first layer uses all heads, the later layers mainly relies on a single head. In addition, the authors also propose a preprocess-then-optimize algorithm.

Strengths

This paper is easy to understand and provides interesting theoretical results.

Weaknesses

The main flow of Section 4 and 5 is that, it (1) constructs a first layer which can perform the exact preprocessing considered in the proposed algorith, (2) studies the generalization of the preprocessing, (3) constructs a the later layers which can perform gradient descent, and finally (4) studies the generalization of the optimization. The paper uses a linear attention-only transformer. Given the above, there are several main weaknesses regarding the contributions in this paper: (1) Missing theory or theoretical intuition to explain why multi-head attention works in the way observed in this paper. Given the series of in-context learning research which focuses on aligning multi-layer transformers with gradient descent, the contribution is limited. (2) The linear attention-only transformer can be improved to better align with the practical LLMs. In the recent year, various studies consider single/multi-head attention with softmax attention [1,2,3]. The authors are encouraged to add softmax attention in the paper. In particular, compared to the linear attention result in [4], the softmax result in [1,2,3] shows that softmax attention works differently from linear attention. The authors are also encouraged to study how MLP works in the system. These elements in the transformer may result in different behavior compared to the linear attention-only transformer. (3) The paper constructs a transformer which aligns with the preprocess-then-optimize algorithm, rather than showing that a trained multi-layer transformer indeed works in the way designed in this paper. The authors need to provide more evidences for this, e.g., [5]. The authors may also want to highlight their different to [5]. A convergence analysis similar to [1] and [2] is even better. (4) The generalization properties are built upon a linear model, whose analysis is supposed to be routine. The authors may also highlight the challenges when performing the analysis. Given the above limitation, the theoretical contributions are not sufficient enough. There is no critical difference between existing theories in literature and the study in this paper. References: [1] Huang, Yu, Yuan Cheng, and Yingbin Liang. "In-context convergence of transformers." arXiv preprint arXiv:2310.05249 (2023). [2] Chen, Siyu, et al. "Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality." arXiv preprint arXiv:2402.19442 (2024). [3] Cui, Yingqian, et al. "Superiority of multi-head attention in in-context linear regression." arXiv preprint arXiv:2401.17426 (2024). [4] Zhang, Ruiqi, Spencer Frei, and Peter L. Bartlett. "Trained transformers learn linear models in-context." arXiv preprint arXiv:2306.09927 (2023). [5] Ahn, Kwangjun, et al. "Transformers learn to implement preconditioned gradient descent for in-context learning." Advances in Neural Information Processing Systems 36 (2024).

Questions

I would like to raise my score if the authors could provide an extra detailed theoretical result for any one of the above weaknesses. The more the better. Thanks.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

NA

Summary

This work presents an analysis of how Transformers perform in-context learning by experimenting with a sparse linear regression problem setup. The author’s analysis combines both empirical and theoretical analysis. They first examine the properties of real models using pruning and probing methods. From this, they identify a key property of the Transformers, that multiple heads are needed in the first layer but in the remaining layers only a single attention head is necessary. From this, they present theoretical analysis and propose a preprocess-then-optimize learning approach to explain the observed behavior in transformers. The authors also present results showing the effectiveness of this learning approach.

Strengths

While this work is not in my area of expertise, I believe the work is well organized and clearly presented. I appreciate the authors combination of empirical and theoretical analysis, to demonstrate that the proposed learning method may actually be employed by real transformers in real tasks and are not only achievable through hand-crafted weights. The authors results are mathematically rigorous as well in supporting the proposed preprocess-then-optimize method.

Weaknesses

The main weakness I see in the work pertains to the simplifications made for the analysis. The authors also acknowledge these limitations in section 7. In summary, the analysis focuses on a simplified problem (linear regression) and uses a simplified version of a transformer, which is also done by other similar works. As a result, it is hard to say if a similar learning approach is learned by transformers in more practical real-world problems. It is also impossible to say how the removed layers may alter the learning process in a full transformer. Overall, I believe this work warrants an accept rating, though I will defer to the other reviewers if they have significant issues with the work.

Questions

Line 108: should this instead read "to perform a tractable theoretical investigation"? Line 209: typo "proprocessing" instead of "preprocessing"

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors have discussed the limitations I mentioned in the section above. I do not see any clear potential negative societal impact for this work.

Summary

This paper studies the mechanism of transformers under the in-context sparse linear regression problem. The authors reveal that the transformer pre-trained for this task has the first layer preprocessing the data, and the remaining layers implement gradient descent. More intriguingly, only one head in the second to last layers is dominantly utilized.

Strengths

1. The paper is a nice extension of the recent line of work on ICL formulated by regression problems. The study of multi-head also provides an important direction. 2. The theory is solid and clear. 3. The authors have found many interesting phenomena in their setting. I especially find their study on multi-head very interesting and beneficial for LLM interpretability.

Weaknesses

1. The P-probing lacks a controlled experiment. Should you also try regressing on the hidden states before the first layer? I understand that $h=1$ might be a controlled experiment, but it may still have some preprocessing on the $x$. 2. Is Theorems 5.1 and 5.2 a fair comparison? Should the $\tilde{w}^t_{gd}$ in Theorem 5.1 also include the parameters in the first preprocessing layer? My understanding is that we want to compare the excess risk of "preprocessing+gd" v.s. " gd". Now we are comparing "gd with preprocessed data" and "gd". 3. The orthogonal design is an over-simplified setting.

Questions

1. I find the observation that only one head dominates in subsequent layers very interesting. Does this also occur in other settings of ICL tasks?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

The orthogonal design is an over-simplified setting, which the authors do not properly address. 1. It is necessary to include ablation studies with non-orthogonal design. It's fine to get different results without the orthogonality, but it would be important to report the findings. 2. If the MLP layer is added, an alternative lasso method should be available that uses the closed-form solution under orthogonal design. Can the authors present the experiment results? Overall, I think there should be more ablation studies in different settings. Although a simplified setting is fine for theory, I want to see experimental results in more realistic settings.

Summary

This work seeks to provide a deeper exploration of the use of multi-heads, at different layers in a Transformer, to perform in-context learning tasks. More specifically, the goal of the paper is to experimentally discover additional insights on the interactions of multi-headed attention across layers. Subsequently, the authors find that multiple heads are primarily used in the first layer, whilst the remaining layers of the transformer typically leverage a single head. Furthermore, the authors provide a hypothesis for their observation before building a preprocess-then-optimize algorithm.

Strengths

Strengths: - This work provides a timely analysis of the underlying mechanism of multi-headed attention in Transformer models. - The authors provide strong empirical and theoretical evidence to justify the insights derived from on in-context learning. - Additionally, the authors provide a novel Preprocess-then-optimize Algorithm for training Transformers.

Weaknesses

Weaknesses: - As discussed by the authors, the reviewer's primary concern with this work is the lack of account for the remainder of the Transformer architecture. In particular, the role of feed-forward layers is sidesteped when considering the effects of the Preprocess-then-optimize Algorithm.

Questions

The reviewer would like to better understand how the authors think about feed-forward layers with respect to the observations made by the authors regarding the necessity of multi-headed attention in the first few layers. In particular, would alterations to the structure of the first few feedforward layers provide a beneficial impact for transformers?

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have clearly stated any limitations with their current work. Additionally, the reviewer does not foresee any potential negative societal impact from this paper.