Benign or Not-Benign Overfitting in Token Selection of Attention Mechanism

Authors: Keitaro Sakamoto, Issei Sato
Primary area: learning theory
Keywords: benign overfitting, attention mechanism, prompt tuning, gradient descent

Abstract:
Modern over-parameterized neural networks can be trained to fit the training data perfectly while still maintaining a high generalization performance. This “benign overfitting” phenomenon has been studied in a surge of recent theoretical work; however, most of these studies have been limited to linear models or two-layer neural networks. 
In this work, we analyze benign overfitting in the token selection mechanism of the attention architecture, which characterizes the success of transformer models. We first show the existence of a benign overfitting solution and explain its mechanism in the attention architecture. Next, we discuss whether the model converges to such a solution, raising the difficulties specific to the attention architecture. We then present benign overfitting cases and not-benign overfitting cases by conditioning different scenarios based on the behavior of attention probabilities during training. To the best of our knowledge, this is the first study to characterize benign overfitting for the attention mechanism.

Decision: reject
Reviewer scores: 6, 5, 5, 8, 5

Reviews and discussion:
<thread>
creation date:
1730399743686

thread_type:
official_review

<thread_object>
creator: reviewer_vumc

<content>
summary:
This paper provides a theoretical understanding of the benign overfitting phenomenon within the framework of training Transformers. It demonstrates the existence of benign overfitting in the Transformer architecture and characterizes the conditions under which a gradient descent-based approach can converge to this benign overfitting regime, supported by theoretical guarantees.

strengths:
Overall, this paper is well-written. (1) The motivation and problem formulations are clearly articulated. (2) The discussion of related works is thorough and supports the claim that this paper is the first to study benign overfitting in the context of the attention mechanism. (3) The theoretical results are reasonable and intuitive, with sufficient explanations and a detailed description of the assumptions.

weaknesses:
The primary weakness lies in the assumption of a favorable initialization of \( v \) (as indicated in Lemma 3.1). I assume that the scenario in Lemma 3.1 simplifies to a fully connected neural network when \( p=0 \). Consequently, the results of Lemma 3.1 could be seen as a straightforward extension of existing findings. However, it is essential to clarify the conditions under which Lemma 3.1 holds. 

Additionally, if we assume such an effective initialization, it appears that this model performs sufficiently well even without the attention mechanism, which raises questions about the necessity of analyzing the attention mechanism.

Furthermore, I did not find any results concerning \( W \), representing \( W_Q W_k \) in the attention layer.

Finally, the training approach differs from standard Transformer training practices, as this paper requires \( W_v \) to be fixed from the start. The author should verify whether this approach is effective in real applications or if the insights are transferrable to practical scenarios.

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1730456432179

thread_type:
official_review

<thread_object>
creator: reviewer_hy7u

<content>
summary:
This paper investigates the phenomenon of benign overfitting within the attention mechanism, marking the first study to do so. Specifically, the authors demonstrate the existence of benign overfitting solutions under certain conditions. They establish that both signal learning and noise memorization are essential for benign overfitting in token selection. Furthermore, they illustrate that benign overfitting can be achieved through gradient descent, given specific conditions. Theoretical findings are supported by experimental results.

strengths:
1) The first work to perform theoretical analysis of benign overfitting in the attention mechanism.
2) The paper is well written.
3) Theoretical findings are supported empirically.

weaknesses:
1) N/A

questions:
1) Any practical takeaways from the analysis of the benign overfitting phenomenon for transformer-based models (especially about their training)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700279270

thread_type:
official_review

<thread_object>
creator: reviewer_wjup

<content>
summary:
This work studies benign overfitting and harmful overfitting in the token selection mechanism of the attention architecture. The theoretical analysis first shows the existence of the attention mechanism and then provides the convergence analysis. Conditions on benign and harmful overfitting are characterized theoretically.

strengths:
1. The problem to solve is interesting and significant.

2. The provided theoretical analysis makes sense and seems solid.

weaknesses:
1. The presentation is not satisfactory. 
i) In Introduction (line 52), it says Theorem 4.1 explains the mechanism in the token selection of attention architecture. However, I cannot find any related discussion around Theorem 4.1. ii) Assumptions A3-A6 can be presented as conditions needed in the theory instead of assumptions. iii) Section 4.3 mentions some challenges, but does not mention how this paper addresses them.

2. Equations 11 and 13 are indirect, which makes the results weak. Meanwhile, it seems it cannot be associated with Figures 3 (a) and (b) about $d$ and $||\mu||$. One way to improve is to further interpret Equations 11 and 13 in terms of $||\mu||$ and $d$.

3. No experiments on real-world datasets are conducted. 

4. Only $p$ is optimized, while other parameters are not changed, which makes the problem to solve less challenging.

questions:
The legend in Figure 3 is confusing. There are five different lines but with the same name "irrelevant". Figure 3(a) shows that $W\_{-Y^*}$ is the largest in the attention weights, which seems contradictory to the sentence in lines 483-484.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703786617

thread_type:
official_review

<thread_object>
creator: reviewer_auqu

<content>
summary:
This paper presents theoretical analysis of benign overfitting in transformer attention mechanisms, examining how these models can perfectly fit noisy training data while maintaining generalization capability. Authors have maintained that it is the first work to analyze benign overfitting in the attention mechanism of transformer models. Based on a set of assumptions, the work provides theoretical guarantee for token selection behavior and characterizes conditions for benign versus harmful overfitting. While the paper makes a novel theoretical contribution to understanding attention mechanisms, several critical limitations make it difficult to comprehend the full significance of this work. The highly simplified setting, strong assumptions, and lack of real-world validation raise concerns about the practical relevance of the findings. Here’s what this reviewer feels about the proposed method

strengths:
Originality:
First theoretical analysis of benign overfitting in attention mechanisms and unique analysis of token selection behavior and how attention handles noisy labels. Several setting build on existing benign overfitting analysis frameworks

Significance: 
Good theoretical insights into transformer attention behavior but limited practical impact due to simplified setting and need for strong assumptions. Results may not translate to real-world applications and missing key transformer applications (e.g., next token prediction, multi-head attention)

Summary of Strengths:
1.	Novel theoretical contribution to understanding attention mechanisms 
2.	Clear mathematical analysis and proofs 
3.	Well-structured validation of theoretical findings 
4.	Potential implications for prompt-tuning applications based classification

weaknesses:
1.	The proposed method has been implemented in a oversimplified setting with single-layer attention and binary classification through single prompt-tuning. Hence, it is difficult to give verdict that the method can be a general tool for understanding benign overfitting in transformers’ attention mechanism setting.
2.	Some assumptions in the paper seems very strong or unrealistic, for example, in the data distribution, a constant ratio of relevant tokens and weakly-relevant tokens in an input sequence is assumed, with experiment done on single relevant and irrelevant token in a sequence assumed. This may certainly differ in a real data distribution
3.	Limited empirical validation on the proposed theories with only small synthetic data. The applicability of the proposed method is also unclear and there is not enough discussion on the scalability.
4.	This paper is a first step of understanding benign overfitting on attention mechanism, however, it is not clear from reading the paper the real motivation, applicability or significance of the analysis and understanding of benign overfitting in transformer attention mechanism. This should be addressed in the writing.

questions:
Please refer to the weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730718210431

thread_type:
official_review

<thread_object>
creator: reviewer_xrje

<content>
summary:
Authors investigate "benign overfitting" (where model generalizes as well as overfit the training data) for attention architecture. Benign overfitting has not yet been defined for transformers, so the authors take a first step by studying it in the context of token selection. They theoretically prove its existence and studies different training trajectories when it can happen.

strengths:
- Novelty: Authors are first to study it for transformer models.
- Contributions seems to be rigorously proved.

weaknesses:
- Paper is challenging to follow.

Scope/Practicality of the paper seems limited as:
- Papers made various assumptions both while formulating the problem (like tokens being splitting nicely into three groups) and while defining the data. 
- Only synthetic data has been used.

questions:
1. Please justify the use of synthetic data.
2. Does real data also adhere to the various assumptions made?
</content>
</thread_object>

</thread>


Metareview:
metareview:
This submission received mixed reviews. During the discussion phase, Reviewer HY7u, who gave the highest rating (8), did not participate in the discussion, and the associated comments were quite brief, providing insufficient grounds to champion the paper for acceptance.

Meanwhile, Reviewers vumC and wjUP raised key concerns about the paper’s contributions, summarized as follows:

(1) The model is perceived as overly simplistic and does not sufficiently differentiate itself from existing works, even if it addresses a different problem.

(2) The real-world experiments presented at the end of the Appendix focus primarily on noisy training and partially verify the theory. However, more practical insights that could guide real-world improvements are preferred.

Given these concerns and the current state of the submission, I regretfully recommend rejection.