Benign Overfitting in Token Selection of Attention Mechanism

Thank you for your detailed review. We appreciate that you highly evaluate our work.

---

1: It would be beneficial for the authors to conduct experiments with varying levels of noise in real-world datasets to further validate their findings. Additionally, a heatmap visualization on real data would provide valuable insights into the model’s behavior under different noise conditions.

In response to your comment, we are making the following additions to Section G.2:

1. Table similar to Table 2, presenting results for different label noise levels $\eta = 0$ and $0.2$ in addition to the $\eta = 0.1$ results in the main text
2. Heatmap illustrating how accuracies change as the label noise $\eta$ varies continuously within the range $[0, 0.2]$ for each dataset. While these results do not directly serve as an experimental validation of the SNR boundary in the theorem, we believe that they provide additional insights into the behavior of real-world datasets.

---

2: While the authors aim to introduce a brief proof outline in Section 4.1, the presentation feels more like a summary of multiple results. A more detailed explanation would enhance clarity. I recommend adding a section in the Appendix to provide an overview of the proof, including key derivations. For instance, is it possible to give a short discussion on how is the function $g(x)=2x+2\sinh⁡(x−\log⁡T)$ developed?

To address your concern and enhance the motivation, we have made the following improvements in Section 4.1 and Appendix:

• To prevent the abrupt start of Section 4.1, we have added the following overview after left column, line 308: "In the analysis of benign overfitting, it is necessary to track the model behavior on the training data while also its generalization ability. We first present the result of training dynamics, and the generalization result is shown at the end of this section in Lemma 4.8."
• We have clarified the reason for introducing Definition 4.4 by adding the following sentence after left column, line 310: "This quantity is useful for evaluating the softmax probability at each time step $\tau$ for a given training example."
• We have added the motivation of the function $g(x)$ at line 329 as follows: "This function naturally arises when expressing the evolution of the attention gap using the weight updates in Equations 4 and 5 and evaluating the dynamics via the quadrature method. Please refer to the Appendix for further details on the derivation."
• We have added a brief explanation of how the dynamics of the attention gap are derived after Lemma 4.5: "This lemma is shown by tracking the gradient descent dynamics and conducting induction argument with several desirable properties."
• We have added a new section titled "Proof Sketch" before Section C to provide an overview of the proof structure and the motivation behind our analytical approach.

The motivation for using the function $g(x) = 2x + 2\sinh(x - \log T)$, where $T$ is the sequence length of the inputs, naturally arises in the analysis of token selection dynamics. Specifically, the parameter updates are governed by the softmax probability term $s(\tau)(1 - s(\tau))$, which can be expressed in terms of the "attention gap" in Definition 4.4 and the equation $(2 + 2\cosh(x - \log T))^{-1}$ by Lemma E.1 and E.2. The function $g(x)$ appears naturally when applying the quadrature method to analyze the training dynamics (please refer to the discussion around Equation 96).

---

3: I notice that in this paper, the authors assume that the number of iterations $T$ can grow exponentially. I am curious about the potential impact of this assumption on the conclusions. Since $\log ⁡T$ can be polynomial, are the results dependent on the exponential property in the loss function? Additionally, do the proofs establish that certain desirable properties hold for all t, allowing $\log T$ to grow polynomially instead?

Thank you for your question. We answer the questions as follows.

1. As you correctly pointed out, the main theorem states that the required time steps for benign overfitting is in an exponential order. In our paper, please note that $T$ is used to denote sequence length rather than time steps.
2. This exponential time steps arises not from the shape of the loss function but from the softmax function in attention mechanism. We provide this intuition around line 376 in the main text.
3. As you noted, Lemmas D.11, D.12, and D.13 in the Appendix use an inductive argument to establish that some desired properties hold for all time steps. In particular, Proposition $C(\tau)$ implies that the evolution of the attention gap, as defined in Definition 4.4, follows a logarithmic order. Consequently, the number of time steps required for generalization scales exponentially.

Thank you for your detailed review. We hope that our answer has addressed all your concerns and pray this reply will change your decision.

---

1: The comparison among not, benign, and harmful overfitting lacks in real-world experiments. It needs more explanation and experiments for harmful overfitting.

The primary objective of real-world experiments is varying the training size $n$ to support the transition between the not overfitting and benign overfitting in Theorem 4.1. This is because we cannot control the values of $\\|\mu\\|_2$ and $d$ unlike synthetic experiments. However, Table 2 also provide insights into harmful overfitting. For instance, in the AG-news dataset, while the model perfectly fits the noisy samples, increasing noisy samples leads to a decline in test accuracy, which is a harmful overfitting scenario.

In response to the review, we are adding following experiments to Section G.2: softmax dynamics in real-world settings similar to Figure 3 and new tables similar to Table 2 for different label noise $\eta$, to improve the generality of experimental results.

---

2, 3: What if there is FFN and nonlinearity on top of the attention layer? Once $\nu$ is also trained, what will it influence the attention mechanism?

If the FFN and nonlinearity are fixed, the problem remains a token selection problem given the token scores, and similar results hold. As explained below, fixing the head $\nu$ is a necessary problem setting rather than an assumption or limitation of analysis.

Since our focus is analyzing benign overfitting in the token selection mechanism, jointly optimizing the head would unnecessarily expand the components enabling noise memorization and benign overfitting. Specifically, when we show benign overfitting, it is unclear whether this is because of the token selection inside softmax or simply a consequence of the linear model $\nu$, which has already been extensively studied (e.g., [Bartlett+, 2020]).
Assumption 3.3 is necessary to formulate the problem of token selection given the assigned token scores. To clarify this point, we have moved this part from Assumption to the Problem Setting. Furthermore, please note that this setting is practically plausible in the area of parameter-efficient fine-tuning, such as prompt tuning and LoRA.

Our study focuses on aligning with a more realistic data model instead of joint head optimization. We newly analyze the token selection in the presence of label noise (see Table 1), which highly complicates softmax dynamics by introducing different training directions within the same training run. Since the parameter updates depend on the softmax probabilities (Equations 8 and 11), addressing this difficulty requires carefully tracking and evaluating softmax dynamics, and a significant portion of our proof is dedicated to resolving this new challenge.

---

4: The assumptions (A1-A8) seem very strong.

In Section 3.5, before (A1-A8), we explicitly list and compare the assumptions used in prior works on benign overfitting. This comparison shows that our assumptions align with commonly used ones in the literature. While the assumptions may seem complicated, the essential part is the relationships among $d$, $\\|\mu\\|_2$, and $n$. It is common to express parameter assumptions using big-O notation ignoring logarithmic dependencies [Cao+, 2022; Jiang+, 2024], but we present them explicitly.

Could you please share which specific assumption you find particularly strong compared to the existing studies? Based on your feedback, we can further elaborate on why it is necessary for our analysis.

---

5: Three distinct token groups and a consistent scale $\rho$ are implausible.

Building on previous research, we have carefully designed the analysis setting to better align with real-world scenarios.

The existing benign overfitting works commonly assume that the input image is split to two parts: signal and noise [Cao+, 2022; Kou+, 2023]. Similarly, in the prior studies on the attention dynamics, it is often assumed that each input is made of a single optimal token and the other irrelevant noise vectors that are orthogonal to signal [Tarzanagh+, 2023b; Jiang+, 2024], as summarized in Table 1. Our study introduces several realistic elements to the analysis, including middle states termed "weakly relevant tokens" and non-orthogonality between signal and noise. To show the correspondence with real-world datasets and strengthen the plausibility of data model, we provided an example of medical imaging in the paragraph from line 4105.

Finally, our analysis also holds if the scale $\rho$ varies across examples. We have added this clarification after Definition 3.1.

---

6: What is the motivation for using $g(x)$?

Due to the character limit of OpenReview, please refer to the second comment in Reviewer BNsU at the bottom. We have made several updates to clarify the motivation throughout Section 4.1, not just for $g(x)$.

Thank you for your detailed review. We hope that our answer has addressed all your concerns and pray this reply will change your decision.

---

1: The proof techniques lack motivation. Section 4.1 lacks sufficient explanation or intuition in the current draft.

We have added the following improvements to enhance motivation. Due to the character limit of OpenReview, please refer to the second comment from Reviewer BNsU at the bottom for details on the updates.

• The overview of Section 4.1 after left column, line 308
• The reason for introducing Definition 4.4 after left column, line 310
• The motivation of the function $g(x)$ at line 329 (see also Comment 9 below)
• The brief explanation of how the dynamics of the attention gap are derived after Lemma 4.5
• The new section titled "Proof Sketch" before Section C

---

2: The real-setting experiments are somewhat trivial and expected—showing that more data leads to better loss. The authors should include additional experimental results that better highlight the intuition behind their findings.

Increasing the number of data containing label noise does not trivially lead to better loss. Depending on the dataset, fitting to label noise can result in worse generalization; for example, the AG-news dataset in Table 2 exhibits a harmful overfitting scenario. To further improve the generality of experimental results, we are adding the following new experiments to Section G.2: softmax dynamics in real-world settings similar to Figure 3 and new tables similar to Table 2 for different label noise $\eta$.

---

3: Yun et al. (2020) is an incorrect reference.

You are correct, and we have already fixed it. Additionally, we have reviewed the references again to ensure there are no errors.

---

4: Assumption 3.3, (A1)–(A8), and the first two lines in Theorem 4.1 seem overly restrictive, and their significance is unclear.

Assumption 3.3
Our study is the analysis of benign overfitting in the token selection mechanism; thus, training the head $\nu$ is inappropriate as it obscures which part enables noise memorization and benign overfitting. Specifically, when we show benign overfitting, it is unclear whether this is because of the token selection inside softmax or simply a consequence of the linear model $\nu$, which has already been extensively studied (e.g., [Bartlett+, 2020]). Assumption 3.3 is necessary to formulate the problem of token selection given the assigned token scores. To clarify this point, we have moved this part from Assumption to the Problem Setting. Finally, we have added a new proposition after Appendix E showing that a one-step optimization from a randomly initialized $\nu$ satisfies Assumption 3.3. This result supports the validity of the setup to some extent.

(A1-A8)
In Section 3.5, we explicitly list and compare the assumptions used in prior works on benign overfitting. This comparison shows that our assumptions align with commonly used ones in the literature. Several papers express parameter assumptions using big-O notation ignoring logarithmic dependencies [Cao+, 2022; Jiang+, 2024], but we present them explicitly.
Could you please share which specific assumption you find strong compared to the prior studies? Based on your feedback, we can further elaborate on why it is necessary for our analysis.

Assumption in Theorem 4.1
The assumption $\\|\nu\\|_2 = O(1 / \\|\mu\\|_2)$ is made to ensure that the token scores remain at most of constant order. This corresponds to appropriately scaling down the network output, and satisfying the assumption is straightforward.

---

5: Running head should be fixed.

Thank you for your comment. We've already fixed it.

---

6: The authors should provide the gradient update equations for $W$ and $p$.

For gradient updates, they are provided in Equations 8 and 11 in Appendix B due to space limitations in the main text. We have added the following sentence after the gradient update equations in Section 3.4: "The weight updates with specifically calculated loss gradients are provided in Appendix B."

---

7: What happens in “Not Overfitting” case after the time step $\tau$?

In the not-overfitting case, signal learning is dominant, and the model fits only the clean samples as stated in the theorem. This is discussed in detail in Figure 2, the paragraph from line 291, and Appendix D.2.

---

8: Time step $\tau$ should depend on uncertainty $\delta$ but it seems there is no such dependency.

Uncertainty $\delta$ indirectly affects the time step $\tau$ through the parameter assumptions (A1)-(A8). A similar order notation for the time step is also found in existing benign overfitting studies [Cao+, 2022; Jiang+, 2024].

---

9 : Why does $g(x)=2x+2\sinh(x− \log⁡T)$ appear in the proof techniques?

Due to character limit of OpenReview, please refer to the second comment from Reviewer BNsU at the bottom. We have added an explanation in the main text after line 329.

Thank you for your detailed review. We hope that our answer has addressed all your concerns and pray this reply will change your decision.

We first answer your question because it is an important point relating to our contribution.

---

1: What is the specific difficulty of label noise setting?

The presence of label noise introduces a significant challenge due to the existence of competing training directions within the same training run. This results in two key difficulties:

1. Signal learning vs memorization in each token selection (Figure 2).
2. Clean samples vs noisy samples in the learning direction of class signals​.

These challenges become even more difficult because the weight updates depend on the softmax probability $s(\tau)$ (Eqs. 8 and 11) and are NOT statically determined by pre-training quantities such as SNR or label noise $\eta$. This is a fundamental difficulty that does not appear in previous analyses of benign overfitting (e.g., two-layer NN [Frei+, 2022]). For instance, depending on the convergence speed—how quickly $s(\tau)$ converges to $0$ or $1$—it is possible that even when label noise $\eta$ is very small, the actual contribution to weight updates at some time step can be dominated by noisy samples. Thus, it is crucial to track the whole training dynamics of each token to analyze 1 and 2, making the analysis inherently difficult. A large portion of this paper is dedicated to addressing and resolving this issue. In addition to the existing explanation from line 86, we have added a new section titled "Difficulty of Label Noise Setting" in Appendix A, incorporating an explanation similar to this reply.

---

2: The novelty and impact are limited compared to [Jiang+, 2024].

The above question highlights a challenge that is specific to the attention mechanism but is entirely absent in [Jiang+, 2024]. As discussed above, addressing this challenge is a major focus of our work and not a straightforward extension. Additionally, we discuss finer differences starting from line 627 in the Appendix.

---

3: There is no sufficient comparison with [Magen+, 2024].

We discuss the differences from line 143 and a more detailed difference from line 650 in the Appendix, which clarifies the novelty of our work. In response to the review, we have added a row for [Magen+, 2024] in Table 1. Furthermore, since the Appendix previously only provided a theoretical comparison, we have also newly included a comparison of experimental validation, as suggested.

---

4: There is no definition of “good run”.

We have already defined "good run" precisely in Definition C.2 and have confirmed that this term is not used prior to this definition, including the main text.

---

5: There are some concerns about the threshold between benign and harmful overfitting.

Assumption (A2) determines the lower bound on SNR for benign overfitting as $\Omega(d^{-1/4})$. We have added a clarification regarding this point below Theorem 4.1. Although this is not a continuous boundary, we discuss in Remark 4.2 that when $\text{SNR}^2 = o(d^{-1/2})$, the model exhibits harmful overfitting. This boundary also appears as a minimax generalization bound in [Xu and Gu, 2023]. Furthermore, the boundary is demonstrated through synthetic experiments (Figure 4, right). In response to your comment, we have explicitly annotated Figure 1 with the specific values instead of "large" and "small".

---

6: Fixing head $\nu$ as in Assumption 3.3 limits the generality of the results.

Our study is the analysis of benign overfitting in the token selection mechanism; thus, enlarging the trainable components, including head $\nu$, unnecessarily would obscure which part enables noise memorization and benign overfitting. Specifically, when we show benign overfitting, it is unclear whether this is because of the token selection inside softmax or simply a consequence of the linear model $\nu$, which has already been extensively studied (e.g., [Bartlett+, 2020]). Our work establishes a novel result that benign overfitting can occur solely through the attention mechanism.

Assumption 3.3 is necessary to formulate the problem of token selection given the assigned token scores, rather than being a limitation. To clarify this point, we have moved this part from Assumption to the Problem Setting. In this setting, our work is directed toward a more general token-selection problem instead of head training. We emphasize that analyzing token selection with label noise involves handling nontrivial softmax dynamics, dedicating a significant portion of our proofs to address this difficulty. The novelty and difficulty are explained in Comment 1 above and in Table 1 of the main text.

Finally, in response to your comment, we have added a new proposition after Appendix E showing that a one-step optimization from a randomly initialized $\nu$ satisfies Assumption 3.3, which supports the validity of this setting to some extent.