Benign Overfitting in Token Selection of Attention Mechanism
Authors: Keitaro Sakamoto
Issei Sato
Venue: ICML 2025 poster (2025)
Keywords: Attention mechanism
token selection
benign overfitting
gradient descent
over-parameterization
Primary area: deep_learning->attention_mechanisms

Abstract: Attention mechanism is a fundamental component of the transformer model and plays a significant role in its success.
However, the theoretical understanding of how attention learns to select tokens is still an emerging area of research.
In this work, we study the training dynamics and generalization ability of the attention mechanism, under classification problems with label noise.
We show that, with the characterization of signal-to-noise ratio (SNR), the token selection of attention mechanism achieves ``benign overfitting'', i.e., maintaining high generalization performance despite fitting label noise.
Our work also demonstrates an interesting delayed acquisition of generalization after an initial phase of overfitting.
Finally, we provide experiments to support our theoretical analysis using both synthetic and real-world datasets.

Benign Overfitting in Token Selection of Attention Mechanism
Keitaro Sakamoto Issei Sato
Abstract
it remains unclear whether such max-margin solutions ac- tually generalize well. In contrast, theoretical studies such as (Jelassi et al. 2022 Li et al. 2023a; Jiang et al. 2024) analyzed the generalization ability of attention mechanisms, but in the data models assumed by these works, token selec- tion proceeds in the same direction for all training samples. Under annotation errors or adversarial label flips, the behav- ior of token selection in both clean and noisy samples and its impact on generalization remain unclear. These existing studies raise the following research questions: (Q1) How do the training dynamics of token selection evolve under label noise? (Q2) Does the obtained solutions generalize well?
Attention mechanism is fundamental component of the transformer model and plays significant role in its success. However, the theoretical under- standing of how attention learns to select tokens is still an emerging area of research. In this work. we study the training dynamies and generalization ability of the attention mechanism under classifi- cation problems with label noise. We show that, with the characterization of signal-to-noise ratio (SNR), the token selection of attention mechanism achieves "benign overfitting". i.e., maintaining high generalization performance despite fitting label noise. Our work also demonstrates an inter- esting delayed acquisition of generalization after initial phase of overfitting. Finally, provide experiments to support our theoretical analysis using both synthetic and real-world datasets.
To this end, we analyze the token selection of attention mechanisms in the context of benign overfitting studies, which allows us to study the training dynamics of gradient descent and the generalization performance. Modern over- parameterized neural networks achieve a high generaliza- tion while perfectly fitting to the training data (Zhang et al., 202I). This "benign overfitting" phenomenon has attracted attention over the past few years because it contrasts with the conventional wisdom that achieving better generaliza- tion requires balancing training error and model complexity.
1. Introduction
The transformer models (Vaswani et al. 2017) have greatly succeeded across a wide range of fields, including natural language (Devlin et al. 2018; Brown et al. 2020), vision (Dosovitskiy et al. 2021 Touvron et al. 2021). and have be- come a foundational architecture in modern machine learn- ing. The key component that characterizes the transformer model is the attention mechanism, which was originally introduced in recurrent neural network (RNN) and long short-term memory (LSTM) to capture the long-range struc- ture of sequence (Bahdanau 2014 Xu 2015). The attention architecture can process lertthle-length input sequences and flexibly select important tokens based on the input. The seminal work the training dynamics of this token selec- tion mechanism (Tazzanagh et al., 22033a b Vasudeva al.) 2024| Sheen et al., 2024) studied the implicit bias of gradi- ent descent to the max-margin token separator. However,
There are lines of studies analyzing the benign overfitting phenomenon in various settings, including linear classifi- cation and two-layer neural networks, but the analysis of benign overfitting is mostly limited these types of architec- ture. In this paper, following the common setup in existing attention work (Tarzanagh et al.. 2023a b Oymak et al. 2023 Sheen et al. 2024), we analyze a one-layer attention network f(X; W. = A x S(WW p) E R, where S(-) is the softmax function, X (x1,...., xT) E TRTx is the sequence of input tokens, W Rddd is the trainable key-query matrix, p Rd is tunable token, and Rd is a pretrained linear head, on a binary classification task. Here, p corresponds to the trainable [CLS] token (Devlin et al.
2018 Dosovitskiy et al.) (021) or prompt tuning (Li &
Liang 2021 Lester et al. 2021) in the application of trans- formers, and represents the output at that token position. To isolate the role of benign overfitting in token selection, we fix the linear classifier v, and focus our analysis on the training dynamics and generalization behavior that arise solely from optimizing the attention mechanism.
'Department of Computer Science, The University of Tokyo, Tokyo. Japan. Conrespondence to: Keitaro Sakamoto -sakakei-19999@ w.w.u.o..c...ooo....o jp>, Issci Sato <sato@@ g. ccc.u- to...
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Benign Overfitting in Token Selection of Attention Mechanism
o(d 1/4) small
large 1/2)
SNR (
Harmful Overfitting
Benign Overfitting
Not Overfitting
Train: Y +1
Train: Y +1
Train: Y +1
1411 direction
20-
Train: Y
Train: Y
Train: Y
Train
10
Test: Test:
Test:
Test: Y +1 Test: V
11+1 direction
20
Test
Select one token for each data X
Test
10-
10
A4-1 -direction
14_ direction
11 -direction
Figure 1: Projection of one selected token per sequence. Each point indicates x selected by attention from each input in the direction of class signals μس1 and μi_ for the three scenarios of harmful overfitting, benign overfitting, and not overfitting. Top: Training data with label noise. Bottom: Test data. The decision boundary common because the head is fixed, but the model can select an appropriate token that belongs to the desired output region. Here, ||ll| denotes the strength of the class signal, d is the dimension of the data and parameters, of is the variance of the input noise, and is the size of the training set.
( = (x]] ....,,................. x7 )T
Contribution. We show that, under the condition based on SNR, benign overfitting occurs in the token selection mechanism when optimizing W and p with gradient de- scent from random initialization. Specifically, when SNR is high, the model disregards fitting to noisy samples and fits only clean samples, and the generalization ability is high (not overfitting case). As SNR decreases, the model per- forms distinct token selection for clean and noisy training samples, fitting both while still generalizing well even in the presence of label noise (benign overfitting case). This reflects a balance where sufficient class signal strength is required for generalization, while noise memorization is necessary for fitting noisy samples. Figure S illustrates the benign overfitting in token selection.
setting without fundamentally modifying the argument.
2. Related Work
Token-selection in Attention. The theoretical analyses the attention mechanism from the perspectives of gen- eralization ability (Jelassi et al. 2022 Li et al. 2023a b 2024a) and training dynamics (Tian et al. 2023 2024) have been progressing in recent years. The most related line of work to ours deals with the implicit bias of gradient descent for a one-layer attention model (Tarzanagh et al.) 2023a b Li et al. 2024b Vasudeva et al. 2024 Sheen et al.) 2024 While our work follows these studies terms of model settings, their primary focus is the training dynamics with- out generalization analysis. Furthermore, our work is also influenced by the problem setting of (Oymak et al.) (023). However, they analyze the initial few steps of training under the data model without label noise. As for the technical dif- ferences in proof, we show in Appendix [A that our results cannot simply be obtained by adapting their analysis the context of benign overfitting. Table [ provides a concise comparison of the analysis focus with the existing studies of the attention mechanism for classification.
Technically, we provide method to evaluate the evolution of the softmax probabilities for each training example, both from above and below. We need it because the learning direction of the class signal not determined solely by the number ratio of clean noisy data but requires careful eval- uation of the softmax probability ratio at each time step. Moreover, our results demonstrate that grokking (Power) et al. 2022. Nanda et al.) 2023) occurs in the benign over- fitting case. While fitting the training data is achieved at a relatively early stage of training, further training in exponen- tial order required to learn the class signals sufficiently and reduce the generalization error. In this paper, we con- sider binary classification for simplicity of discussion, but as shown in Appendix E.I it can be extended to the multi-class
Benign Overfitting. The success of modern over- parameterized models has led to numerous studies attempt- ing to understand why and when benign overfitting occurs. The analysis is interesting because the standard generaliza-
2

---

Benign Overfitting in Token Selection of Attention Mechanism
Table 1:
Comparison of the theoretical works of the attention mechanism on the classification task. "Label Noise" indicates
whether distant token selections occur depending on the training example.
Paper
Dynamics Generalization
Label Noise
Token Type
Tarzanagh et al.
2023a
One optimal, the others non-optimal One optimal, the others non-optimal One optimal, the others non-optimal
Vasudeva et a.
2024
(Sheen et al. 2024)
Oymak et al. 2023
Relevant and irrelevant tokens
(Li ct al.
2023a)
Relevant and irrelevant tokens One relevant, the others irrelevant One relevant, the others irrelevant One relevant, the others weakly or irrelevant
Jiang et al. 2024)
Magen et al.
2024
Ours
tion bound based on uniform convergence does not explain this phenomenon well (Nagarajan & Kolter, 2019). For comprehensive surveys of the literature on this topic, please see work such as Bartlett et al. (202T); Belkin (2021).
solely through optimization within the attention mechanism. Please refer to Appendix A for further details.
3. Problem Setting
Benign overfitting in regression has been studied in the lin- 2022) and ker- Tsigler& & Bartlett
In this section, we introduce the notation and the problem settings in the rest of the paper.
ear model (Bartlett al. 2020 Hastie al.
nel regression (Liagg & Rakhli 2020
2023). It is complicated to analyze the classification task because the explicit formula for the min-norm separator is not obtained. One approach is to track the training dynamics with gradient descent in the linear classifier and two-layer
3.1. Notations
Let [n]] be a shorthand for the set (1.....n). n2.. We de- note multivariate Gaussian distribution with mean vec- tor n and covariance matrix Σ by NV,,, E). Denote by S R + R S(v); = exp(v.)/ Eree exp(ve) the softmax function. The standard Big-O notations 0(.), (-), 22-), o(.), and w(-) are used to hide absolute constants, and we denote inequality ignoring constant factors by Z, S.
neural network (Chatterji & Long 2021| Frei et al. 2022
Xu & Gu 2023 Cao al. 2022; Kou et al. 2023 George
let al.
2023] Meng et al.
2024| Xu al.) 2024). Another
line of work is built on the results of implicit bias to the max- margin solution (Cao et al.) 2021] Wang et al. 2021 Wang & Thrampoulidis/20222 Frei et al.. 2023a). These studies base their discussion of convergence on existing research on implicit bias (Soudry et al. 2018 Ji & Telgarsky 2019 Lyu Li 2020 Frei et al. 2023b). Specifically, they analyze the properties of the solution through the KKT conditions of the max- margin problem in order to examine whether the solution at convergence shows benign overfitting or not. Furthermore, there is a research direction in investigating the degree of benignity of overfitting (Mallinar et al. 2022 Wen et 2023 Kornowski et al. 2024).
3.2. Data Model
In the analysis of benign overfitting, we typically need to consider the specific shape of the data distribution to evalu- ate the generalization error without using a uniform conver- gence argument. We consider the following data distribution P defined over (X, Y) E Rxxd / x {=1]. In this paper, we consider binary classification for simplicity of discussion, but the same argument applies the multi-class case. Please refer to Section E. in the appendix for more details.
Recently, Jiang et al. 2024) studied benign overfitting in simplified vision transformer model. A major difference is that our analysis is under label noise, where the direction of token selection differs between clean and noisy data in the same training run, and signal learning for generaliza- tion competes between clean and noisy data. Concurrently, Magen et al. (2024) has also analyzed benign overfitting in similar model setting; however, our analysis is beyond the initial few training steps and is conducted on more general data model. Furthermore, our research put more focus on the token selection mechanism than these stud- ies, demonstrating that benign overfitting can be achieved
Definition 3.1. Let 14+1 4-1 E Rd be fixed class sig- nal vectors satisfying |All = || = || and (M p_1) = 0. The input X [x..... xr R Txd has T tokens that are split into three groups: relevant token R {1} containing the strong signal for true class, weakly relevant token W S [T] \ R containing weak class signals, and irrelevant token I = [T] \ (R UW) containing only noise. Let clean distribution P be the distribution over T x {+1} such that (X, Y*) is sampled as follows:
1. The clean label Y is sampled from Unif( (=1))
2. The noise vectors (e. (e()) are sampled independently
3

---

Benign Overfitting in Token Selection of Attention Mechanism
from ((0,o2L))
with trainable weights W o. W K, Wy Rdxd and
W. Rdxm Here, the softmax function S(1) is applied row-wise with the abuse of notation.
3. The relevant token is given by X1 = My- + 61.
In practice, an additional tunable token E R is concate- nated to the input, and this position is used for the model prediction. This setup is widely used in, for example, the classification token [CLS] in BERT (Devlin al. 2018) and
4. The weakly relevant tokens x, for u E W are given by x PM4,, + Eu. where < 1 is small scale parameter and W/ E {+1} for u E W. For simplicity, we assume single confusing token X2 = pp_y- y-- + 62 and other tokens x,, = ppy- + Eu for u E W \ (2).
VIT (Dosovitskiy al. 2021), and prompt-tuning technique
[Li & Liang 2021 Lester et al. 2021). Let the concate- nated input be X p P. x
E [[(++)xdd then the cross-attention feature between Xp and X is given by
5. The irrelevant tokens are given by x, = Ev for e I.
We use the notation W+1 {u E W wn = +1} and W_1 = (u E W wi, = -1.. The data distribution Pi is defined as the label-corrupted version of P with the level of label noise 1 > 0. The data point (X, Y) from P is generated by first sampling (X, Y*) from clean distribution Pt and then setting Y = Y with probability 1) and Y Y* with probability 1 17. We denote signal-to-noise ratio by SNR- = [1llla/(a, vd). For simplicity of notation, we assumed a fixed scale p for weakly relevant tokens, but the analysis holds even if the scale varies across examples.
[[x)]T [sx(X)]
= =X WX ( JXWW
SSp wx ((XWX XWrW...
where we use W to denote a key-query weight matrix W e W and the output corresponding the position of P is denoted by S(X) = w o W X S(XW p) E RRm.. In this work, we use the model output for binary classifica- tion, leading to the output dimension being m = 1, and we denote the value prediction head by WvW. E R Therefore, the model under our analysis is of the form
Training data S = (x(,,, Y (1) ) j==1 are sampled i..dd from P. We denote the clean data (i E [m] y = y*o) and noisy data {i E [n]] y # y by C and N, respectively. The set of data pi E C| y = are denoted as C+. and (i C | = -11 }} are denoted as C_. The same notation is applied to NC The superseript (i) denotes that the variable corresponds to the training data i E [n].
f(x) S(XW p).
(1)
The output can be regarded as an affine combination of the token scores {7h := -'' xx xteretr using the learned softmax probabilities. Let E H² be a shorthand for the softmax vector S(XW p).
(4) for t E [T]. w0 and I(o)
For instance, we write
Data models based on signal and noise widely appear in the existing benign overfitting studies (Chateerji & Long 2021 Frei et al., 2022| Cao et al., 2022, Jiang et al. 2024]] Meng et al., 2024). Such data models based on signal and noise not limited to benign overfitting works but also com- monly observed in other analyses of attention architecture (Jelassi et al. 2022 Li et al. 2023a Oymak et al. 2023).
To clarify the role of token selection in the attention mech- anism with respect to benign overfitting, we fix the linear classifier v, for which benign overfitting has already been extensively studied (Bartlett et al. 2020 Chatterji & Long 202T), and analyze the training dynamics and generaliza- tion behavior arising solely from the optimization of the token selection mechanism inside the softmax. To formu- late the token selection problem based on given token scores {1h hee()), we assume a pretrained head v that assigns ap- propriate scores each token. Specifically, we consider such that
Remark 3.2 (Weakly relevant token and label noise). Weakly relevant tokens represent tokens with weak signal strength and confusing class information, reflecting a more realistic scenario than clean separation into relevant and ir- relevant tokens. Furthermore, such weak class information, including that which is confusing, is likely to lead to lower annotation quality, making it more plausible to consider the presence of label noise.
cos cos Og > O(),
(2)
for k E (+1), where 9g denotes the angle between z and Ug. Jointly training v is itself an intriguing setting, but it would obscure whether benign overfitting occurs in the softmax weights or in the linear classifier. The analytical setup in this paper enables a stronger claim: that noise memorization and benign overfitting can occur solely within the token selection mechanism. We further discuss the rationale behind Equation ( in Appendix E.2
3.3. Attention Model
Given a sequential input X = (X1,... X1) TT E Rxx_ a single-head self-attention layer SSA : Rxx RTxm is
(X) X
4

---

Benign Overfitting in Token Selection of Attention Mechanism
Remark 3.3 (Relevance to practical scenarios). The ana- lytical setup is relevant to practical scenarios, especially in the context of parameter-efficient fine-tuning, For example, prompt-tune tuuiig (Li & Liang] 2021 Lester et al.] 2021) trains only the tunable input tokens, and LoRA (Hu et al. 2022) fo-
in the existing studies, including Frei al. (2022) Cao et al. ) Jiang al. (2024). Assumptions A5 and A6 the common assumptions to evaluate the class balance in the training data and the amount of noisy data. For exam- ple, n = ((polylog(d) is assumed in (Cao et al. 2022 Jiang et al. 2024). Assumption A7 is put to focus on the dependencies on d. || and n in the asymptotic notation, following (Jiang et al., 2024). Finally, Assumption A en- sures that the attention probabilities initial weights reasonably uniform, and this type of condition is observed in existing studies, including other architectures (Cao et al.
cuses on training only the attention weights. The results on the training dynamics and the generalization performance provide guarantee regarding the required training time and the generalization when applying the model to low-quality downstream tasks containing label noise.
3.4. Gradient-descent Training
2022 Kou et al. 2023 Meng et al. 2024 Jiang et al. 2024).
The learnable parameters (W, p) are trained minimize the empirical risk objective:
Now, we state our assumptions in the following. Let ô_ max(a 1/ Given a small failure probability
8 > 0 and a large enough universal constant C, we make
L(W, p) 1 e (r(.f((x )(x())),
the assumptions for each parameter follows:
(3)
1=1
d z Ca 4/3 log (Tn/0),
(A1)
{(2) = log(1 1 + exp(- -2));
(4)
|| 2 0/ log (og(TT/O),
(A2)
where :: + R a binary cross-entropy loss.
C' C||اا| a. log(T 6) SPS
(A3)
Let each element of W and be initialized as (
a S maxxlllllaa n 2 C'oo(d/o),
(A4)
N (0, o2) and p()); ~ N(0, o,), respectively. The param-
(A5)
eters are optimized by gradient descent with a step size a > 0:
7 7 1/C, T T=
(A6)
(A7)
W(T + -V WH--))-)))--WWWWW...).... p(T))
(5)
ow (mxx( 1a (Tnn/))
p(r 11= = -(1
(6)
(A8)
Let fr be the f after the r gradient descent step and [(r) be shorthand for ((W((). The weight updates with specifically calculated loss gradients are provided in Appendix B.I.
4. Main Results
In this section, we provide the main results regarding the training dynamics and generalization of attention mecha- nisms. The key techniques used in the proof are presented in Section 4.1 By dividing ranges for the signal-to-noiss ratio, we obtain the following results for the convergence.
3.5. Assumption on Parameters
In this section, we first discuss the necessity of the assump- tions on parameters and compare them with existing studies, followed by list of the assumptions used in our work.
Theorem 4.1. Suppose that the norm of the linear head scales as ||/|| = 0(1/ [lll)) Under the parameter as- sumptions in Section 3.5 we have
The condition on d in Assumption training in an wer-paramettrrize setting, and similar terms
A1 is necessary for
1. (Not Overfitting) If SNR? = w(n')), then with probability at least 1 - 5, there exists time step such that the weights (W )) fi only the clean data, and the test loss is sufficiently low:
2||m|| can be found in Chatterji & Long 2021 Frei et al.
2022, Xu & Gu 2023 Kou et al. (023). Assumption A2 is required for generalization The lower bound with d/4 appears in (Xu & Gu 2023), and 2/444/1/4 and rddt/4 are found in (Xu et al. 20244] Meng et al. 2024), respectively. While these conditions may seem intricate, the key aspect in the analysis is the relationship among d, ||w|||2 and n. Assumption
T= alleelaala ligd maxi
WEEC 1=(x0)==11)) =
Vj (x(())) # ((),,
A3 is natural setting to represent weak class information. The lower bound is set such that the weak signal strength is on a larger scale than the inner product of the class signal and a random noise vector. Assumption A4 for a sufficiently small learning rate, which is widely set
Pr
Isign (X)) # Y <
2. (Benign Overfittin)) If SNR? = onn ) then with probability at least - 5, there exists time step
5

---

Benign Overfitting in Token Selection of Attention Mechanism
== +
Iallllal ad max(og. such that the weights (W(T), P(T)) overfit the training data, and the test loss is sufficiently low:
expín 1NR ")
the strength of SNR, which leads to the different cases in Theorem 4.1 The middle in Figure 2
corresponds to the not-overfitiing case, where signal learning dominates and the model does not fit noisy data. The right figure illustrates the benign overfitting case, where noise memorization be- comes dominant. The vertical axis in Figure 2 represents the value of 81 (1 - 81), which determines the amount of parameter updates. As attention probabilities become more concentrated around 0 or 1, this value decreases, reducing the influence of the example on parameter updates. This makes analyzing the token selection dynamics inherently challenging. Our analysis under the label noise setting must account for two competing training directions within the same training run: 1) between signal learning and noise memorization (Figure 2 and 2) between clean and noisy samples for learning the class signals. These balances de- pend on softmax probabilities and are not determined by pre-training quantities such as SNR or label noise n. This is specific difficulty with the attention mechanism, which is absent in existing benign overfitting studies. For instance, depending on the convergence speed-howg quickly s(r) ap- proaches 0 or 1-it is possible that even when label noise n is quite small, the actual contribution the weight updates at some time step can be dominated by noisy samples. This motivates us to carefully analyze the dynamics of softmax probabilities to evaluate the direction of these competing relationships.
[ (x((0))= = y(o) [sign (X)( Y < 6.
Pr
(X.
The assumption ||r||2 = PPL/||p||| (| in the theorem ensures that the token scores remain bounded by constant. This condition can be easily satisfied by appropriately scaling down the model output.
This theorem demonstrates that under the parameter assump- tions in Section 3.3 the model does not fit noisy data and achieves high generalization accuracy with high SNR, and the model shows benign overfitting with low SNR. In both cases, high generalization is achieved despite the presence of label noise. Since both the signal required for generaliza- tion and the noise memorization for fitting label noise are significant, their balance is essential for benign overfitting. An upper bound on this balance is provided in the state- ment of the theorem, while the lower bound determined by the parameter assumption A2 Specifically, by extract- ing the relationship among fz | d, and n, we see that SNR² = 1/4-1/4) holds in this benign overfitting case.
Remark 4.2 (Harmful Overfitting). Theorem 4.1 presents cases of not overfitting and benign overfitting, where the test loss is sufficiently low. With a lower SNR that violates Assumption A2, specifically SNR² o(d-1/2 1/2), even inner products among the random noises in the input dominate the selection of class signal (see Lemma B.I in the appendix for details). It becomes inherently challenging to select class- relevant tokens and generalize effectively. The low SNR case corresponds the left column of Figure [
4.1. Key Techniques
In this section, we present the key techniques used in the proof of Theorem 4.1 The whole proof is provided in Appendix C The proof for the not-overfitting case in the theorem is more straightforward and is presented first in the complete proof. In this section, to capture the essence of the proof in the main text, we focus the more non-trivial benign overfitting case, i.c., SNR² o(n 1)
Remark 4.3 (Implication for Grokking). The time step demonstrated in Theorem is also an important result. In the case of benign overfitting, fitting the training set requires a similar time step order to the not-overfitting case; how- ever, to sufficiently learn the signal components and achieve high generalization, the exponential term of the SNR in the numerator is necessary. This delayed generalization abil- ity after fitting the training set implies a connection to the phenomenon of grokking (Power et al. 2022. Nanda al. 2023). For the necessity of this exponential term, please refer to the end of Section 4.
In the analysis of benign overfitting, it is necessary to track the model behavior on the training data while also its gen- cralization ability. We begin by presenting the results on training dynamics, and the generalization result is shown at the end of this section in Lemma 4.9 To this end, we intro- duce the following values that represent the relative strength of token selection of the learned weights. This quantity is useful for evaluating the softmax probability at each time step T for a given training example.
Definition 4.4 (Attention gap). We define the attention gap between a significant token and other tokens as follows:
We illustrate the differencee between the two scenarios of the main theorem using Figure Selecting the class-relevant token X1, i.e., increasing 81 towards 1, leads to fitting clean data but fails to adapt to noisy labels. For clean data, both signal learning and noise memorization contribute to in- creasing $1, whereas for noisy data, these two processes are in competition, The learning direction is determined by
Ar,(T)=
(xx² (xx
p(F),
[T, == =xx]) (x()
W()T
for E [m], E [T] \ [1), and a E [T] \ (2).
6

---

Benign Overfitting in Token Sclection of Attention Mechanism
do (T))
Noise Signal
5 (r) of ()))
< Noise Signal +
SNR w(nnt)
( (T)
Noise Signal +
<
SNR o(
+
(r)
== 0
t
( 1- 7
T
Clean Data 81 (T)
(4)
Noisy Data $ (T)
(5) Noisy Data $1 (T)
Figure 2: Illustration for the training dynamics of the probability assigned to relevant tokens $1 in clean data i E C and noisy data E N.. The y-axis shows $1 (1) 81 (T)), which determines the magnitude of the gradient descent update. This value converges to 0 as (T) approaches 0 or 1, and consequently, the contribution of this training example to the gradient update diminishes. The middle figure corresponds to the not-overfitting case in Theorem 4.1 and the right figure represents the benign overfitting case.
In this paper, we first show that the training dynamics of the attention gap can be described by the monotonically increasing function g(x) = 2x + sinh(x - T). This function naturally arises when expressing the evolution of the attention gap using the weight updates in Equations ( and (6) and evaluating the dynamics via the quadrature method. Please refer to Appendix for further details on the derivation. By using this function, the token selection of relevant token in clean data i C evolves as follows:
Stage 2, learning progresses to select the confusing token x that can fit the label noise.
The next lemma shows that Stage I, the value of g de- creases linearly over time, leading to a decrease in the atten- tion gap A, in contrast to Lemma 4.5
Lemma 4.6 (Attention gap dynamics of noisy data in Stage
1). For some Ti θ an-ta; Mlla Mllad max o
with probability at least 1 - 6, we have for all time step T E [0, Ti] that
Lemma 4.5 (Attention gap dynamics of clean data). For any exp(n SNN 2)
Ts = + an-1oo Mf alllad max و.ل ). with probability at least 5, we have for all time step T E [0, that
(Aga(T)) 99 (A
maxfo² -T o max(o o
g (Ai,((T) g (0))
+ 2 (an "fllelaalll| maxfo of)),
for any E (7] (1).
for any t E [T] (1).
that is, to Subsequently, the training proceeds to select x2 (i) increase the attention gap I j,a.
This lemma is shown by tracking the gradient descent dy- namics and conducting induction argument with several desirable properties. Due to the shape of the hyperbolic sine function, the function is close to linear around x = log I", and gradually becomes exponential as & moves away from To This lemma states that while the value of evolves linearly over time, the actual evolution of the attention gap becomes increasingly logarithmic factor. This finding cor- responds to the left in Figure where (7) approaches 1, (8) (T) $1 (7)). which determines the
Lemma 4.7 (Attention gap dynamics of noisy data in
Stage 2). Let Ti be the time step in Lemma 4.6 For any exp(n "oppn--SSIB SSNR-2)
(ogg)) ). with probability at least - 8, we have for all time step T E [Ti,Ty) that
Tr = -n-
9 (T (T) log 1) 9 (0), log p 11
+ θ (pam - - max o;1)
g (T) = y (T (0)
and the value of 81
scale of gradient updates, diminishes accordingly. This part
+ θ (pom | max oz))
is also illustrated in Figure i in the appendix.
for any u E (T) (1,2).
In contrast, the behavior of noisy data j E N is described across two training stages. During the initial Stage 1, the softmax probability assigned to the relevant token x(5) is suppressed. In this stage, as illustrated in the right column of Figure the influence of noise memorization dominates, and the training progresses in a way that avoids selecting the relevant token, unlike the clean data case. In the subsequent
Similarly to Lemma 4.5 these lemmas are proved by using the parameter updates to track the evolution of the attention gap and the softmax probabilities assigned to each token. At this point, is necessary to carefully evaluate several factors: the softmax probabilities, the norms and inner prod- uets between the model weights and both signal and noise

---

Benign Overfitting in Token Selection of Attention Mechanism
vectors, as well as the relative contributions of other training examples. In particular, the presence of label noise com- plicates the analysis involving softmax probabiliitess in the attention mechanism. Since the gradient updates depend on the softmax probabilities of training examples, the learning direction of the class signal cannot be determined solely by the ratio of clean to noisy data sample sizes, which is different from the existing benign overfitting work on linear or two-layer neural networks. Therefore, it is essential to evaluate the ratio of softmax probabilities across different training examples at each time step. This technique is useful for the analysis when dealing with distinct token selection behaviors among training samples, not limited to the label noise setting. The whole proofs based on mathematical induction are provided in Appendix C
Synthetic experiments. We train the same model / as defined in Equation ① with gradient descent W and p.
using the same data model as Definition 3.1 with o, = 1. Specifically, we consider the setting with n = 20, T = 8, 7 = 0.2, = 0.1 and a =5e-3, -3, changing the value of the dimension d and the signal size I l2. For simplicity, we set [++| = [W_1 = 1 and v x p - 1111 which aligns
with our problem setup.
Figure 3 shows the dynamics of softmax probabilities for clean and noisy training samples from the initial weights. The left column represents the harmful overfitting case, where the dimension is larger compared to || The bottom figure shows that confusing weakly relevant token X2 is not always selected for noisy data; instead, irrelevant tokens that fit the label noise can be picked. This model can fit the training data with noise components, which hinders signal learning and reduces generalization ability. The mid- dle figure shows the case where the balance between signal and noise is achieved, and benign overfitting is observed. In this figure, selecting the weakly relevant token 2 E W-y-G) for the noisy data j E / aligns with our analysis in The- orem 4.1 In the right figure, where signal norm || is large, fitting the noisy data does not happen. The model is trained to select relevant token x; for both clean and noisy samples, supporting the not-overfitting case in Theorem 4.1 For additional synthetic experiments, please refer to Ap- pendix F.I We provide the loss heat-map when varying d and || 2l. which supports the SNR boundary established in the main theorem. Furthermore, we conduct similar experi- ments using a more general one-layer transformer encoder, providing results beyond our analytical setting.
Remark 4.8 (Two-stage Analysis). In the two-stage analy- sis in Lemmas 4.6 and 4.7 it is implicitly considered that cxp(n 1 SNR ²) is of larger order than p= meaning that the log/ 1 > 0 scale is neglected on the benign overfitting condition SNR² (n)). However, even if this is not the case, the benign overfitting result can still be obtained by replacing the numerator of T2 with p p-
Finally, we present results related to generalization. The attention value of the class signal at time step r is bounded from lower as follows:
Lemma 4.9 (Attention value of class signals). Let T2 be the time step in Lemmas 4.5) 4.7 With probability at least 1 - 8, we have for all time step T E [0, T2] and l E {+1} that
T (0) (0) My (() p(T()) Z #e
+ nSNR2 log (7-an - -02 || aallll 2d max of)).
Real-world experiments We further conducted real- world experiments on image and natural language datasets for classification. For each task, we used the pre-trained VIT (Dosovitskly et al.) 202T) and BERT (Devlin et al.) 2018) models. To align as closely as possible with our analysis setup, we initialized the weights within the final attention layer and trained only these weights using a dataset with label noise. This setup corresponds treating the fixed backbone model up to the last layer as feature extractor and its output as the input to one-layer attention model. We used datasets from various types: 10-class image classi- fication with MNIST (LeCun et al., 2010) and CIFAR-10 10 Krizhevsky et al. 2009) anomaly detection in medical im- age with PneumoniaMNIST and BreastMNIST (Yang et al., 2023 topic classification of text with AG-news (Zhang et al. 2015). and question type classification with TREC (Li & Roth 2002). For detailed descriptions of these datasets, please refer to Appendix F2
Please note that under the current condition SNR² o(n-1) the value of nSNR2 is of small order within the range that satisfies the parameter assumptions. To select the class relevant token for unseen examples and achieve a high generalization performance, p4 (T) T P(T) is required to sufficiently exceed e' T (F) P(T), where E is a random noise vector. Lemma 4.9 states that this scenario is satisfied if the training time r becomes exponentially large, thereby compensating for the smallness of nSNRR.. This is why we need a time step of exponential order in Theorem 4. Such an exponential time scale is not necessary merely for fitting training examples, which implies that a significantly longer training period is required to achieve generalization.
5. Experiments
In this section, we provide experimental results to support our analysis. The code is available on GitHub
Table 2 presents the training loss and test accuracy when varying the training size n. Since the SNR cannot be con- trolled due to its dependency on the dataset and the pre-
https:// http:///...ooo.eetttsttstooo attentioo
8

---

Benign Overfitting in Token Sclection of Attention Mechanism
Clean sample
Clean sample
Clean sample
Softmax probability
1.011
Releoet
Softmax probability
1.00
Relecser
Softmax probability
1,00
s.. Reieaet
0.75
My. Wishly Rdccani W-y-
0.75
Sy. Wishly Recoaant W.
0.75
iy. Vically Rdceani W.
Weaky Reievnn My-
Weaky Reteont Py-
Werily Relevant W...
3.50
Kt, Imiievnt
0.50
24. Imlevart
0.50
st. Iire eart
xy. Irtiesant
22. Incleaart 2
2; lize evant
3.25
2% Irtlovant.
Ke.
0:25
2). Inclevant
Se. Indevartt
0.25
dy.
a:: lise ovart
lirelsont
Ireeevont
lire rreettt?
3.30
0.00
0.00
Noisy sample
Noisy sample
Noisy sample
Softmax probability
1.00
Softmax probability
1.00
Softmax probability
1,00
3.75
0.75
0.75-
3.50
0.50
0.50
0.25
0:25
0.25
3.30
0.00
0.00
200
100
800
1000
200
-03
600
1000
200
100
800
1000
Time step
Time step
Time step
(a) Large noise setting: d = 5000, |||| Final training accuracy is 1.0 and test 20. Final training accuracy is 1.0 and test 100. Final training accuracy 0.8 and test accuracy is 0.87 (Harmful overfitting).
(b) Balanced setting: d 2000, |||2||
(c) Large signal setting: d 1000, ||
accuracy is 1.0 (Benign overfitting)
accuracy is 1.0 (Not overfittin))
Figure 3:
Dynamics of softmax probability. The top represents a clean sample, while the bottom represents a noisy sample. From left to right, the figures correspond to the cases of harmful overfitting, benign overfitting, and not overfitting
Table 2: Training loss and test accuracy when only the final attention weights are trained on a label-noisy sub-dataset (1) = 0.1) for 2000 epochs. The results show the average over three different runs with the standard deviation.
overfitting. In contrast, BreastMNIST does not overfit the training data, and AG-news demonstrates overfitting to the training set but with test accuracy negatively affected by noisy samples, suggesting a harmful overfitting case.
Training Size n
Dataset
Eval
6. Conclusion
20
200
1000
MNIST
train
0.00
0.01
0.03
In this paper, we analyze the training dynamics and the generalization performance of the token selection of the attention architecture. Specifically, we showed that benign overfitting occurs in the token selection mechanism con- ditioning on SNR, and the model sufficiently learns class- relevant signals after overfitting. Furthermore, we supported the analysis with experiments. general, this study help- ful in analyzing the balance of convergence speeds when the model selects different tokens for different training ex- amples, including the label noise setting, As natural next step, extending the analysis setup to next-token prediction or a more general data model can be considered.
test
88.5=3.2
90.020.4
91.4220.2
train 0.00++..01
0.07
0.12=0.01
CIFAR-10
test 95.920.4
95. 1
1
Pncumonia
train 0.00++..0
0.00
0.02
MNIST
test
76.2
78.5a3.2
81.0#1.
Breast
train 0.07++.00
0.144v.01
MNIST
test
74.1=2.0
77.1117
train 0.00++.00
0.00
0.00
AG-news
test
83.6=1.9
12.211.1
7..
train 0.00
0.02
0.05
TREC
test 80.1
78.3
78.5.0.22
trained model, we varied the training size n to observe the scenario transitions shown in Theorem 4.1 Please recall that under our parameter and data model setup, Theorem 4.1 states that redueing n for fixed SNR causes a transition from not-overfitinng benign overfitting. In the range of n shown in Table 2 datasets such as CIFAR-10 10 and TREC exhibit such an increased fitting to the training set as decreases. MNIST and PneumoniaMNIST maintain very low train- ing losses across all n, and the increase in noisy samples with larger n does not negatively impact test accuracy. This scenario is closer to benign overfitting rather than harmful
A Acknovledgements
This work was supported by JSPS KAKENHI Grant Num- ber JP24H00709.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

---

Benign Overfitting in Token Sclection of Attention Mechanism
References
Frei, S., Vardi, G., Bartlett, P., and Srebro, N. Benign overfitting in linear classifiers and leaky relu networks from kkt conditions for margin maximization. In The Thirty Sixth Anmual Conference on Learning Theory, PP. 3173-3228. PMLR, 2023a.
Bahdanau, D. Neural machine translation by jointly learning to align and translate. arXiv preprint A4Kvv:44000445 2014.
Bartlett, P. L., Long, P. M., Lugosi, G., and Tsigler, A. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 117449):0033330070 2020.
Frei, S., Vardi, G., Bartlett, P., Srebro, N., and Hu, W. Implicit bias in leaky reLU networks trained on high- dimensional data. In The Eleventh International Confer- ence Learning Representations, 2023b. URL https //oeenrvieew net forum? http:iiddppppppp.... SEWW
Bartlett, P. I.., Montanari, A., and Rakhlin, A. Deep learning: statistical viewpoint. Acta numerica, 30:87-201, 2021.
George, E., Murray, M., Swartworth, W. J., and Needell, D. Training shallow reLU networks on noisy data using hinge loss: when do we overfit and it benign? In Thirty- seventh Conference Neural Information Processing Systems, 2023,
Belkin, M. Fit without fear: remarkable mathematical phe- nomena of deep learning through the prism of interpola- tion. Acta Numerica, 30:203-248, 2021.
Boucheron, S., Lugosi, G., and Bousquet, O. Concentration inequalities. In Summer school on machine learning. pp. 208 240. Springer, 2003.
net/
Forum?id=s 1ROXEK jh
Hastie, T., Montanari, A., Rosset, S., and Tibshirani, R. J. Surprises in high-dimensional ridgeless least squares in- terpolation. Annals of statistics, 50(2):949, 2022.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
Hu, E. J., yelong shen, Wallis, P., Allen-Zhu, Z., Li, Y.,
Wang, S., Wang, L.,, and Chen, W. LoRA: Low-rank adap- tation of large language models. In International Confer- ence on Learning Representations, 2022. URL https //openrevient wtt.://tooooommmdddddeeeeeee
Cao, Y., Gu, Q., and Belkin, M. Risk bounds for over- parameterized maximum margin classification on sub- gaussian mixtures. Advances in Neural Information Pro- cessing Systems, 34:8407-8418, 2021.
Jelassi, S., Sander, M., and Li, Y. Vision transformers provably learn spatial structure. Advances in Neural In- formation Processing Systems, 35:37822-37833 2022.
Cao, Y., Chen, Z., Belkin, M., and Gu, Q. Benign overfitting in two-layer convolutional neural networks. Advances in neural information processing systems, 35:25237-25250, 2022.
Ji, Z. and Telgarsky, M. The implicit bias of gradient descent on nonseparable data. In Beygelzimer, A. and Hsu, D. (eds.), Proceedings of the Thirty-Seconte Conference on Learning Theory, volume 99 of Proceedings of Machine Learning Research, pp. 1772- PMLR, 25- 28 Jun 2019. R v99 j119a 9a. html
Chatterji, N. S. and Long, P. M. Finite-sample analysis of interpolating linear classifiers in the overparameterized regime. Journal of Machine Learning Research, 22(129): 1-30, 2021.
ml press/
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding. arXiv preprint arXiv: 180000408 2018.
Jiang, J., Huang, W., Zhang. M., Suzuki, T., and Nie, L. Unveil benign overfitting for transformer in vision: Train- ing dynamics, convergence, and generalization. arXiv preprint arXiv:2409.199: 19345, 2024.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image worth 16x16 words: Transformers for image recognition at scale. In International Conference Learning Representaatins, 2021. URL https:// wdb..cdddiiyyyyy...ccc
Kornowski, G., Ychudai, G., and Shamir, O. From tempered to benign overfitting in relu neural networks. Advances in Neural Information Processing Systems, 36, 2024.
Kou, Y., Chen, Z., Chen, Y., and Gu, Q. Benign over- fitting in two-layer relu convolutional neural networks. In International Conference on Machine Learning, pp. 77615-17659 PMLR, 2023.
openreview
Frei, S., Chatterji, N. S., and Bartlett, P. Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. In Conference on Learning Theory, pp. 2668-2703 PMLR, 2022.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
10

---

Benign Overfitting in Token Sclection of Attention Mechanism
LeCun, Y., Cortes, C., and Burges, C. Mnist hand- written digit database. ATT Labs [Online]. Available: 2, 2010.
Mallinar, N., Simon, J., Abedsoltan, A., Pandit, P., Belkin, M., and Nakkiran, P. Benign, tempered, or catastrophic: Toward refined taxonomy of overfitting. Advances in Neural Information Processing Systems, 15:1182-1195, 2022.
Lester, B., Al-Rfou, R., and Constant, N. The power of scale for parameter-efficient prompt tuning. arXiv preprint 10Xiv22004006611 2021.
Meng, X., Zou, D., and Cao, Y. Benign overfitting in two-layer reLU convolutional neural networks for XOR data. In Forty-first International Conference on Machine Learning, 2024. httthtttt///peettt
Li, H., Wang, M., Liu, S., and Chen, P.-. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In The Eleventh International Conference on Learning Representations,
net./
2023a.
Nagarajan, V. and Kolter, J. Z. Uniform convergence may be unable explain generalization in deep learning. Ad- vances in Neural Information Processing Systems, 32, 2019.
id= [d=jCIGv3
Li, H., Wang, M., Ma, T., Liu, S., ZHANG, Z., and Chen, P.-Y. What improves the generalization of graph trans- formers? theoretical dive into the self-attention and positional encoding. In Forty-first International Con- ference on Machine Learning, 2024a. URL https: V/openreview met/ffoumm iidimiddddmii ddddddmnhllssssm sZzzE
Nanda, N., Chan, L.,, Lieberum, T., Smith, J., and Stein- hardt, J. Progress measures for grokking via mechanistic interpretability. In The Eleventh International Confer- ence on Learning Representations, 2023. URL Btthttps //openreview.o ftrttt/ooummiie 9XF SbDPmdw
Li, X. and Roth, D. Learning question classifiers. In COLING 2002: The 19th International Conference
Oymak, S., Rawat, A. S., Soltanolkotabi, M., and Thram- poulidis, C. On the role of attention in prompt-tuning International Conference on Machine Learning, pp. 16724-26760 PMLR, 2023.
Computational Linguistics, 2002. URL https: //www
ac selweb.orgs www.annnhhhaaoooeooooooooo CC0 1150
X. L. and Liang, P. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint arXiv:2/01.00190, 2021.
Papyan, V., Han, X., and Donoho, D. L. Prevalence of neural collapse during the terminal phase of deep learn- ing training. Proceedings of the National Academy of Sciences, (17(40))2655222466 2020.
Li, Y., Ildiz, M. E., Papailiopoulos, D., and Oymak, S. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference Machine Learning, pp. 19565-19594. PMLR, 2023b.
Power, A., Burda, Y., Edwards, H., Babuschkin, I., and Misra, V. Grokking: Generalization beyond overfit- ting on small algorithmic datasets. arXiv preprint arXiv:2201.022 227.00277, 2022.
Li, Y., Huang, Y., Ildiz, M. E., Rawat, A. S., and Oymak, S. Mechanics of next token prediction with self-attention. In International Conference on Artificial Intelligence and Statistics, pp. 685-693. PMLR, 2024b.
Sheen, H., Chen, S., Wang. T., and Zhou, H. H. Implicit regularization gradient flow on one-layer softmax at- tention. arXiv preprint arXiv:2403. 00099, 2024.
Liang, T. and Rakhlin, Just interpolate: Kernel "ridge- less" regression can generalize. The Annals of Statisties, 1333):132- 2020.
Soudry, D., Hoffer, E.,, Naeson, M. S., Gunasekar, S., and Srebro, N. The implicit bias of gradient descent on sep- arable data. Journal of Machine Learning Research, 19 (70):1-57, 2018.
Loshchilov, and Hutter, F. Decoupled weight decay reg- ularization. In International Conference on Learning Representations, 2019. URL net/ wtt://orrmmmmm wwwp:mmiiickkkoccxg...
Tarzanagh, D. A., Li, Y., Thrampoulidis, C., and Oymak, S. Transformers as support vector machines. arXiv preprint arXiv:2308. (an::2308.18998, 2023a.
Lyu, K. and Li, J. Gradient descent maximizes the margin of homogeneous neural networks. International Confer- ence on Learning Representations, 2020. URL https: //openreview.o wttpiiiwwweeeeeeemmm ott.://eooormmm111sseeeeeee
Tarzanagh, D. A., Li, Y., Zhang, X., and Oymak, S. Max- margin token selection in attention mechanism. In Thirty- seventh Conference on Neural Information Processing Systems, 2023b. URL http:h///:////tt /ptp://////ereveieee/mm /ttus nnt/
Magen, R., Shang. S., Xu, Z., Frei, S., Hu, W., and Vardi, G. Benign overfitting in single-head attention. arXiv preprint arXiv:2410.077 2024.

---

Benign Overfitting in Token Sclection of Attention Mechanism
Tian, Y., Wang, Y., Chen, B., and Du, S. S. Scan and snap: Understanding training dynamics and token composition in 1-layer transforerr. Advances in Neural Information Processing Systems, 36:7191 (6:71911-71944 2023.
M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jer- nite, Y., Plu, J., Xu, C., Scao, T. L.., Gugger, S., Drame, M., Lhoest, Q. and Rush, A. M. Transformers: State- of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natu- ral Language Processing: System Demonstrations, pp. 38-45, Online, October 2020. Association for Compu-
Tian, Y., Wang, Y.. Zhang, Z., Chen, B., and Du S. S. JoMA: Demystifying multilayer transformers via joint dynamics of MLP and attention. In The Twelfth In- ternational Conference on Learning Representations,
tational Linguistics. URL https:: aclweb org/ http/atthlogyyy200 2020 emn www. p-d.pos.eee.
2024. URL
id=LbUgRGNYCF
Xu, K. Show, attend and tell: Neural image cap- tion generation with visual attention. arXiv preprint arXiv:1502.03044 2015.
Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and Jégou, H. Training data-efficient image transform- ers & distillation through attention. In International con- ference on machine learning, pp. 10347-10357. PMLR, 2021.
Xu, X. and Gu, Y. Benign overfitting of non-smooth neural networks beyond lazy training. In International Confer- ence on Artificial Intelligence and Statistics, pp. 11094- 11117. PMLR, 2023.
Tsigler, A. and Bartlett, P. L.. Benign overfitting in ridge regression. Journal of Machine Learning Research, 24 (123):1-76, 2023.
Xu, Z., Wang, Y., Frei, S., Vardi, G., and Hu, W. Be- nign overfitting and grokking in reLU networks for XOR cluster data. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https //openreview. htt.:t/e/...otmmmooooooxooopooop
Vasudeva, B., Deora, P., and Thrampoulidis, C. Implicit bias and fast convergence rates for self-attention. arXiv preprint arXiv:2402.057338 2024.
Yang, J., Shi, R., Wei, D., Liu, Z., Zhao, L.,, Ke, B., Pfister, H., and Ni, B. Medmnist v2-a large-scale lightweight benchmark for 2d and 3d biomedical image classification. Scientific Data, 10(1):41, 2023.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, E., and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):10-- 115, 2021.
Vershynin, R. High-dimensione probability: An introduc- fion with applications in data science, volume 47. Cam- bridge university press, 2018.
Wainwright, M. J. Asymptotic Viewpoint. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2019.
Statistics: A Non-
Zhang, X., Zhao, J., and LeCun, Y. Character-level convolu- tional networks for text classification. Advances in neural information processing systems, 28, 2015.
Wang, K. and Thrampoulidis, C. Binary classification of gaussian mixtures: Abundance of support vectors, be- nign overfitting, and regularization. SIAM Journal on Mathematics of Data Science, 4(1):260-284, 2022.
Wang, K., Muthukumar, V., and Thrampoulidis, C. Benign overfitting in multiclass classification: All roads lead to interpolation. Advances in Neural Information Processing Systems, 34:241644222 24179, 2021.
Wen, K., Teng, J., and Zhang, J. Benign overfitting in classi- fication: Provably counter label noise with larger models. In The Eleventh International Conference on Learning Representations, 2023. URL httphhtttt///erreeeehe net/ net/forung wtw.:ummiiiiuUUueeee......
Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz,
12

---

Benign Overfitting in Token Sclection of Attention Mechanism
Appendix
Table of Contents
Comparison with Existing Work
13
A.1
Difficulty of Our Label Noise Setting
13
A.2
Comparison with Implicit Bias Studies
A.3
Comparison with Benign Overlitingg Studies
A.4
Comparison with Other Studies
B
Preliminaries
17
Notation
17
Proof Overview
B.3
High-probability Events
Proof of Main Theorem
24
Preliminary Lemmas
24
Analysis of Not Overlitting Case
28
Analysis of Benign Overlitingg Case
41
D Technical Calculations
58
D.1 Softmax Probability
58
D.2 Attention Updates for Not Overlitinng Case
63
D.3 Attention Updates for Benign Overfitting Case
71
E Further Discussion
E.1 Multi-class Setting
77
Head Optimization
F
Additional Experimental Results
80
F.
Additional Synthetic Experiments
F2 Additional Information for Real- World Experiments
82
Additional Real-World Experiments
A. Comparison with Existing Work
In this section, we highlight the novelty of our analysis and clarify the position of this study through a comparison with existing studies. Before presenting the existing work, we first describe the analytical challenges inherent in our setting, which are addressed in this study. This serves to clarify the novelty of our contribution.
A.1. Difficulty of Our Label Noise Setting
The analytical setting that incorporates label noise distinctively characterizes our study, as shown in Table | By introducing variability in token selection across training examples, the presence of label noise enables a more general analysis of the token selection mechanism. Specifically, label noise introduces significant challenges due to the existence of competing training directions within the same training run. This results in two key difficulties:
Signal learning vs memorization in token selection of each example (see Figure
13

---

Benign Overfitting in Token Selection of Attention Mechanism
Clean samples vs noisy samples in the learning direction of class signals, i.e., M41 and 1-11-
These challenges are further complicated by the fact that the weight updates depend nontrivially on the softmax probability, as will be shown later in Equations ( and [2). Therefore, the training direction is not statically determined by pre- pre-training quantities such as SNR or label noise ratio 7). This is a fundamental difficulty that does not appear in previous analyses of benign overfitting, including two-layer NNs. For instance, depending on the convergence speed-how quickly the softmax probability s(l) converges 1 or 0-it is possible that, even when label noise is very small and the number of noisy samples is negligible, noisy samples can still dominate the weight updates at certain time step. Thus, it is crucial to track the whole training dynamies of each token to analyze the above two relationships, making the analysis inherently difficult. A large portion of this paper is dedicated to addressing and resolving this issue.
In the following, we present specific related studies and clarify the differences and novelty of our work.
A.2. Comparison with Implicit Bias Studies
We compare with the previous papers on the implicit bias of the attention mechanism (Tazzanagh et al. 2023a b
Li et al.
2024b Vasudeva et al. 2024 Sheen et al. 2024). The two major differences from ours are:
(a) Analysis of a fixed training set S, rather than considering the underlying distribution.
(b) Single optimal token, with all other tokens having identical token scores.
In our setting. (b) corresponds to considering a single relevant token, irrelevant tokens x, N /411I1//lla
μ 1-1M-1/ / || ()). and the linear head x 1 14 - In this setting, the noise vectors are orthogonal to the class signals and have an identical token score "n- (141 1-1) x; O, which aligns with (b) However, we consider more general data model, including an intermediate state termed weakly relevant tokens and non con-orthogonaliyy of noise and signals. Furthermore, for data with label noise, selecting any irrelevant token results in the model with output O, and the model cannot fit the training data. Therefore, this data model is too simple to analyze the distinct behavior of the token selection, which is the focus of our study. While we follow the model setting, our analytical methods differ entirely.
A.3. Comparison with Benign Overfitting Studies
Comparison with Jiang et al. (2024). The major differences from ours are:
(a) Absence of label noise.
(b) Differences in the data model.
(c) Focus on vision transformer (ViT) model.
Their setup assumes no label noise in the target Y, meaning all training data is clean. In contrast, we focus the benign overfitting concerning label noise, in addition to the input noises {Er).. As explained in Appendix A. the ability to fit label noise of the attention mechanism has not been analyzed before, and even when possible, the training direction is not straightforwadly determined due to the intrinsic softmax difficulties. These challenges are specific to the attention mechanism and are newly addressed in our work. For (b), they assume orthogonality between class and noise vectors -_ - This assumption overlooks the
because the noise covariance matrix is given by = I "4+1/T///lel interactions between learned parameters and training samples, and it becomes more unrealistic assumption in a multi-class setting in Appendix I Additionally, our data model newly considers intermediate states as weakly relevant tokens, which reflects the real setting more closely. Finally, for they explored the advantage of ViT compared to linear models and CNNs from the perspective of benign overfitting. Their analysis includes the optimization of self-attention and value matrix, which involve challenging interactions during training. In contrast, our focus is on the training dynamics and generalization ability of token selection mechanisms. By focusing on this aspect, we enable analyses in the above general settings. Additionally, our work demonstrates that benign overfitting can be achieved solely through the token selection mechanism within the softmax. This is demonstrated for the first time in our analytical setting. Training the linear head would obscure whether noise memorization and benign overfitting originate from the attention mechanism or from the linear classifier, the latter of which has already been extensively studied in the literature.
Comparison with Magen et al. (024). This work also studies the benign overfitting of the attention mechanism with the same type of architecture, under label noise setting, but it largely differs from our work terms of the following points:
14

---

Benign Overfitting in Token Selection of Attention Mechanism
Differences in the data model.
(b) Two-step gradient descent of p and
For a they simplified the data model, such as signal-noise orthogonality and the clean separation of signal and noise vectors. Our analytical setup fills this gap by introducing noise into the relevant tokens, and by incorporating intermediate states between signal and noise vectors, which we refer to weakly relevant tokens. We also do not assume the signal-noise orthogonality. For [b) present general gradient descent results for p and key-query matrix W. This allows us to discuss differences in the convergence of gradient descent conditioning SNR, as well as the difference in the time scale. Instead, they provide the characterization of SNR for the max- max-margin solution.
We also comment on the differences in experimental validation. Their work conducted the empirical validation under the data model their theoretical setting, not only for the model architecture analyzed but also for other settings such as self-attention and multiple layers. In contrast, we validate our theoretical analysis using synthetic data for both the analytical model architecture and one-layer transformer encoder. Furthermore, we present results on real-world data, highlighting the connection to the analytical setting in paper. We also provide visualization of the heat map for the synthetic data.
A.4. Comparison with Other Studies
Comparison with Oymak et al. (2023). We emphasize that while we followed the analysis setups from this work, our results cannot be obtained by simply applying their results to the context of benign overfitting. The key differences from study are as follows:
(a) Absence of label noise.
(b) Not focusing on overfitting.
(c) Three optimization steps p()), In particular, the learning within the softmax is single-step optimization
of p.
(d) Minor differences in input distribution.
For (8) and (b) we must analyze the competing training directions due to the presence of label noise, explained in Appendix A.I This largely complicates the analysis of softmax probabilities because it is not a simple selection of the
relevant token. [cr the analysis is in the very early stage of training, so it is difficult to fit noisy data containing label noise. Therefore, their setting cannot be applied in our study. Additionally, the optimization within the softmax a single step. Only at the first step, suffices to analyze the gradient descent concerning the average the tokens because of the weight initialization, and the complex softmax term does not need to be considered in the gradient calculation. Optimizing token selection beyond a single step introduces the softmax probabilities into the gradient calculation, resulting complex dependency on the current parameters. To address this, needed to handle the relationships among attention probabilities across all training steps using inductive approach.
15

---

Benign Overfitting in Token Sclection of Attention Mechanism
Table 3: Notations used in this work.
X
Sequence of imput tokens, i.c., X (x1,...,xT))
Y
Training label.
Y*
True label.
n
Number of training data.
T
Length of input tokens.
d
Dimension of token embeddings and model weights.
M+1 1111
Signal vectors for class and 1,, respectively.
Noise component in x; for E [n]] and E [T].
o?
Variance noise vector, i.e., c ~~ (N0,,
o; o
Variances for weight initialization, i.e., p(0); N(0, W(O),) N(0, for i, [d].
SNR
Signal-to- -erial-o-nooie noise ratio given by SNR == [Ialla((o.vd)
(T)
Probability vector for (() at r-th step, defined as S(x (1) W(T)
Token score, defined as x(")
R
Set of relevant tokens; R. {1} and X1 py + 61.
W
Set of weakly relevant tokens; PM + + W,w, E {=1}.
W-1.W-1
Set of weakly relevant tokens with specific label; {u E W wy = +1} and {u e W Wi 1}. We assume that the confusing token consists W-y- = (2).
I
Set of irrelevant tokens; x,, Ev,
c
Set of clean data; [n] (() Y-))).
C+,C_
Set of clean data with label and -1, respectively.
N
Set of noisy data; {i E [n| y * Y.
N4,N_
Set of noisy data with training label and 1, respectively. Learning rate.
a
n
Label noise ratio.
X41(7),A-11)
Attention scores for class signal in Definition C.I A+1 (W())4))P(7) and
x-1 = (W(r)/_),P( P(T)).
P(e(T)
Attention scores for noise vectors in Definition C.1
[,4(T), I,g((T))
Weighted inner-produce terms defined in Definition C.2 C.2
(T), TW (T), TW T((T)
Pa.(T)
Attention gaps defined in Definition C.3 Aa,(T) (x)) xx²") W(r)T
(T, = (x2 -x. W(r) P(T). WP))).(T).
16

---

Benign Overfitting in Token Selection of Attention Mechanism
B. Preliminaries
B.1. Notation
We first list the notations used in this work in Table 3
Furthermore, the basic computations are presented here for convenience. The gradient Vw p) and Vpc(W, p) used in the training be explicitly computed as follows. Since the derivative of the softmax function is given by
VvS(v) = diag(S(v)) - S(v)S(v)',
where diag(v) Txx denotes the diagonal matrix whose (i. t) entry equals to e,. Using the denominator layout, have
WWTL(W.P)
V PwTf(x())
(7)
E y()) x (+)T (diag(S(x 11 W p)) - S(X p) D
p)') x xx up
(8)
T
4 Y(O.
Σ
(1)
7t
(
( (0al
x
(1) x2"07
(9)
where E is abbreviation for (( n x TSX w p)). From a similar calculation, we have
V VpŻ(W, p)
y(i) p))(x())
(10)
Σ
& y(i)
( diag(S(X (6) WT p)) - "(x("T p) WTDSss" (Tp))
p) x )
(11)
" i=l
T
1
4
((i)
Σ 1=1
(6)
(6) 7t
( (6) Wxx!) (6)
(12)
71
u=l
B.2, Proof Overview
This section provides roadmap for the proofs in the appendix, aiming improve readability. For the precise statements and proofs of individual lemmas and propositions, please refer to the corresponding sections.
We begin in Appendix B.3 by establishing high-probability events. Under parameter assumptions in Section 3.5 we derive bounds on norms, inner products, and class balance among samples. Lemma B.I states the collections of these events, and their proofs are provided via standard concentration inequalities, divided across Lemmas
B.3 B.4 and [B6 to B.11 As
preparation for the main proof, Lemma B.12 presents evaluation of the token score 7t- The following Appendix builds upon the high-probability events established in this section.
Appendix C presents the proof of the main theorem. In Appendix C.I we first introduce technical notations (Definitions C. to C..). simple calculations regarding training with gradient descent (Lemmas C.4 to (.6)), and several initial properties at the start of training (Lemmas C.7 and C.8 We also provide a result regarding the balance of loss gradients across
examples in Lemma C.9 Appendices C.2 and C.3 correspond to the not-overfitting and benign overfitting cases, respectively.
In the not-overfitiing case, Lemma C.II properties such as norm and inner product evaluations and balance in softmax probabilities. The most important proposition C(r) characterizes the dynamics of the attention gap A(T) introduced in Definition 4.4 (main text) and Definition C.3 (appendix), governed by the function s(x) = 2r + 2 sinh( log 7). This inductive argument, which holds for all time steps 0 < T < Ti, leads to the proof of the main theorem in Appendix C.2.2 establishing both the not-overfitting and generalization guarantees. The benign overfitting case follows a similar structure, but as described in Section 4.1 the analysis is divided into two stages. The inductive argument is split into Lemmas C.12 and C.13 and Appendix C.3.2 concludes with the proof of the main theorem of this case, including both the overfitting and good generalization.
is the central result. The proof proceeds inductively, incorporating desirable
17

---

Benign Overfitting in Token Sclection of Attention Mechanism
g 2x + 2 sinh( log T)
g(Ai,((T))
Evolve linearly over time ((Arx(r)) (0),,(00) (coefficlent)
(Lemma C.11, C.12 and C.13)
$ (r)
Fast
Slow
( $1 (7)
(T)
based on the function g. The right figure similar to Figure 2 in the main text and illustrates that 81 (r) approaches 1, the term (T) (T)), which determines the magnitude of the one-step update, becomes smaller corresponding slower training of the weights. For the relationship between (()( (T) and Ar,(T), please see Lemma D.1 which shows that 81 (7) sI())) decreases as A(r) moves away from log'T.
Figure Illustration of the evolution of the attention gap defined in Definition C.3
To help illustrate the development of the attention gap driven by the function g. we provide a visualization in Figure 4 The function g(z) behaves approximately linearly near x = log T but grows exponentially as moves away from this region. The inductive lemmas in Appendix C show that ((A()) and ([[(()) evolve linearly, which implies that the growth of A(r)
and (()) themselves gradually slows down.
Although the main theorem is fully proved in Appendix C Appendix SerruddxxDsuuee supplements the proof by presenting several auxiliary calculations. Appendix D.I provides results linking the weight parameters to the softmax probabilities. Appendices D.2 and D.3 correspond to the not-overfitting and benign overfitting cases, respectively. These sections give one-step gradient calculations under the favorable weight conditions, showing how signal learning and noise memorization progress and supporting the induction arguments in Appendix
B.3. High-probabinn probability Events
We first show that the following events occur simultaneously with high probability under the assumptions in Section 3.5
Lemma B.1 brobb-probabblityy events). Suppose that the parameter assumptions in Section 3.5 hold. There exists some constant Cı, C2 > 0 such that the following hold simultaneously with probability at least over the realization of training data S and model parameters (O)), (0):
(i) (Norm concentration) For all E [n], [T]. we have
(1 v s
(13)
(14)
(1 - o(1)
(15)
(== (1 <
(16)
[|] < <1 (1+o(1))o,Vd
(17)
18

---

Benign Overfitting in Token Selection of Attention Mechanism
(ii) (Inner-product concentration) For any i,j E [n], [T] such that (i,1) + (G. u), we have
Ke (e <occ
(18)
(WW())+,, N01
(19)
(W(0))+1 "O
(20)
(W(O)),, N,N
(21)
" (MM00A...000)00ccccoc000)2111110
(22)
(23)
(24)
(0 [100000)0
(25)
(M41 <20,1l ///g///
(26)
(M_ E < <200__0llll|1100([111
(27)
(R, c, )1 < <320 </2oo//lll//////g//////
(28)
(iii) Regarding the clean and noisy samples, we have
3n
2-n
n S -1+1,11115
(29)
4
4
7n S W1+1,
31 n2
(30)
4
Definition B.2 (Good run). We denote by "good run" the trial that the events from ii to (iii
oceur.
In the proof of the main theorem, we will proceed by conditioning on these events. Lemma B.I implies that a good run occurs with the probability at least 5 over the realization of training data S and the weights initialization.
In the rest of this section. we will prove each high-probability event in Lemma B.1 Specifically, can show it by combining Lemmas B.6 B.
B.9 and B.11 using union bound argument. In the following proofs, suppose that the parameter assumptions in Section 3.5 hold. First, we show the norm concentration of the Gaussian noise vectors. The next lemma gives the lower bound for the expectation of the L2 norm.
Lemma B.3, For a Gaussian vector x ~ (0, Σ), we have
V (r(E)) - -15 BESEEI [||
Proof of Lemma B.3 We use the Gaussian Poincaré Inequality ( (Boucheron et al. 2003). Theorem 3.20):
Var (f(x)) < E [| fxx)
(31)
where f Rdd_ + R is any continuously differentiable function. By taking f as f(x) = |x 22, since we have E [
=
E ||xx|||>|] = 1,
1 22 Var(f(x)) E [||x|]] (E (1) Tr(E) - (E 1|}
(32)
Rearranging the terms, we get
V (TrE)) - < [|||||]]
(33)
which concludes the proof.
The following lemma is about the concentration Lipschitz functions and used to prove the norm concentration.
19

---

Benign Overfitting in Token Sclection of Attention Mechanism
Lemma B.4 (Rephrased from (Wainwright 2019. Theorem 3.16)). For any L-Lipschiie,functin / Rd - RR, we have
(- 40 1
w²
Pr (If(e) [[f(e)] 2 w) S 2exp e~N((,
(34)
Note that the coefficient 2 could be removed in the case of one-sided inequality.
Remark B.5. Generally, this holds for strongly log-concave distributions, i.c., distributions with density p(z) exp(-1b - (x)), where 1o R + I is strongly convex function. Here, we used the fact that the Gaussian distribution
(0, Σ) is strongly log-concave distribution with parameter min m = 1/ Jman (2).
We are now ready to prove the norm concentration as follows.
Lemma B.6 (Norm concentration).
With probability least 1- - 8/4
(1 (1=
(1 o ( (12000/00MMBM///X000000003344040
(1 o10_::/0111402000000022333444110
(1 (1000//00///////////////0/11/111
< (1+
for all i [n],te e [T]
Proof of Lemma B.6 From the definition of noise distribution, et NO. La for all i E [n], E (T]. To begin with,
(
we show the norm concentration of the Gaussian vector. For x. y E Bd since we have
|x||2
|yyl
|| |x||2 - Ilyl|2
x y1l2 s ||x y112.
(35)
|x y l22
f(x) = ||||| [|2 1-Lipschitz function. Using Lemma B.4 we get
(-w² w²
Pr
Pe(l1e] c 8 Ila - E [|el|aa]
w s 2exp
'Ao?
(36)
for some E [n], E (T). Taking union-bound gives
w
Vi it] Vt [T], Pr
{ ||e| 1l2 -E [[]e[[a] > w < 222n
(37)
4o
Lemma B.3 and Jensen inequality lead the following bound the expectation of the Gaussian norm:
Vold-15 1 < a. Vd.
(38)
Using this and Equation (37 we have with the probability at least 5/20,,
40Tn
(40Tn
oodd - 1 1
21 o2
log
< |le S so.va 2/ o² log 8
(39)
for all [m), (T). By using d w(log( (oolog(Tn/))) (8) in the parameter assumptions, the second term on both sides becomes dd..vd). Combining this with (a_ (o (1-1/(o2d) < Vo I. we have the desired result. Regarding other inequalities, by Gaussian initialization of parameters, W(0) is normally distributed with mean 0 and covariance of 1 and we repeat the same discussion. For (0)c (0) I|,. considering the probability for parameter initialization under the realization of the training set, we can use the same argument and union bound. The inequality for ((0) is derived similarly.
20

---

Benign Overfitting in Token Sclection of Attention Mechanism
Next, we move to the concentration inequality for Gaussian random variables. Lemma B.7 (Gaussian tail bound, ((Vershynin) 2018), Prop 2.1.2)). For Gaussian variable x N(0, the tail bound is given by
0
0² w
1
2²
1 o w) S
w
exp
S Pr
f2
exp
2m
2r²
V 2
2n²
Using this, we can show the following result for the inner products of the noise vectors.
Lemma B.8 (Inner-product of two Gaussian). There exists some constant ei such that with the probability at least 1 6/4,
(W()))4 = (VO04
(W(O) <oot:_l"""llccc//////0ooooo/
W(0)c 9
(W(0))41 (0)) < <aowoo/111])///log((111
(W(0) (0) c10 <clooo//l/|a/(0lg((///
(W(0))) (() c1 llt.oollllllllllllll(/1oo/ooopo
for all i,j je [n],t,u E [1]] such that (i, (..).
Proof of Lemma B.8
Before delving into the main part of the proof, we first show
w²
Pr (le]le 2 w/ S 2xxp
(40)
20
for the Gaussian vector e ~ o0 Unlike Lemma B.6 handles the norm itself rather than the deviation around the mean of the norm, which more useful to prove this lemma. We have for any X that
Pr fllell2 2 w) S Pr {lel]] 2 w)
(41)
I exp( (>x) exp (Alea))
(42)
A=1
Eexp(-Xux
II E exp (Ae)
(43)
22 (
= 2exp
of w1
(44)
where the second inequality follows from Markov inequality and the last follows from the moment- geneenttggnertingg function of Gaussian distribution. Minimizing the upper bound over À gives the desired inequality Equation 440
Fix i [n] and E [T'] such that (i,t) u). For any > O, we have
Pr [e (8) 2 o) {
Pr ke 1 P.))))
oellc w/ + Pr
> w].
(45)
where we used the inequality Pr(A)= Pr( + Pr(BO) Pr & Pr(B) Pr(A BC) for the event A, B, which gives tighter bound when outlier event A and event B share large common parts. Under the condition ( is fixed, (6))
since (c.9) ca follows NO0, o2(|c(()) [12),
Lemma B.7 gives
2 0elic ll2 11 o
222
Pr Ke (
exp
(46)
v
V2T
20 1
21

---

Benign Overfitting in Token Selection of Attention Mechanism
Thus, the conditional probability is bounded
Pr [(e)) (0
< > oelle 1|2 S w
2 w
exp
v
(47)
v
2
2w
Combining Equations
40 and (47.. then applying union bound on Equation
(45
we obtain
Pr (ii.j tn [T], (i,t) * (i. u), s.t. (e, .e. >}
Pr [i, E [n],t, (T), (t # (1 u), s.t. (i c v Vj [n], [T). = w]
+ +P Pje (T], s.t. T. ||اb 112 > w)
(48)
122
w²
$
V2TV exp
2w?
+ cxp
200dd
(49)
Let = 4 (log(1 (log(T'2//0)) n//))
1/2 -1/2
for some constant of > 0. By Equation 49). have
3i, Inl (T), (i,t) (i, s.t. ((e))) (0) < w}
Pr
n2
ei (log(Tn/8)) -1/2
8 T'
2/(2())
+ exp
GFUZ
(50)
2
20
Further, let v= c vd log (T (TT/o where c1 > 0 is some constant. We have
Pr J E [n]. t, u E [T], (1, + (i, s.t.
16,), -, > of
2
1/2 5
1/(26))
8 6/t//2
& (log(Tn/6))
+ 2I'
(51)
2m
Tn
T'n
8
8
(52)
64
64
32'
where the last inequality is satisfied with the appropriate choice of c1, c > 0. This completes the proof for the first inequality. The other inequalities are also the inner product of two Gaussian vectors and can be proved with the same argument.
Lemma B.9 (Inner-product of signal and noise). There exists some constant 02 > 0 such that with probability at least 1 8/4,
1M41: < c20 <goo://lllllll/ooooTTrr
|p_1; c (0)) < C2a, c20. (T//)
[|, (01)) < c2 log
for all i [n],t [T].
Proof of Lemma B.9 We will show that the inequality for u411 holds with probability at least 8812.. The same discussion
applies to μ_1 and v. For the fixed i E [m)) [T], since (upne, (.)) follows the Gaussian distribution N(0, o² 12)
Lemma B. gives
20. ||al|12 cxp
Pr (p & ( > w
(53)
20
||
Let w 020, 1/11 (og())/)) for some constant C2 > 0, then applying union bound on Equation 53
gives
Pr Pr{36e (n],t e [T]] s.t. KM4 { > w
2Tnoz
w
exp
(54)
v 2TT
20 14
2//
2Tn
8
(55)
V2TC2 (Tn/o)
Tn
12
where the last inequality is satisfied with the appropriate choice of 02 >
22

---

Benign Overfitting in Token Selection of Attention Mechanism
Lemma B.10 (Hoeffding Inequality). Let X1..... ...... Xn be i.i.d. random variables such that 0 s X S 1 almost surely. Then for all O, we have
n X; - -
Pr
w) S 2exp (-2nw²).
n
Lemma B.11 (Number of Samples), For all d' > 0, the following hold with probability at least 1 - 8/4:
1-1 ạ
n S |c+1s
7
c'
n,
2
1- 11
n S 1c-1s
1
c
n1
2
2 n W+
11
1
11
2
( - n 2
+c)
2
72
Proof of Lemma B.17
We show the first equation holds with probability at least - 5/16. The proof of remaining cases follows similarly, and the desired result is achieved by using union bound.
The training data i [n] belongs C+ when its true label (+) 1 and label flip does not occur. This event occurs independently, and its probability is calculated (1 n2/2. Since |C+ Ei= 1y00=y+6==1- applying Lemma 1.10 to X; = =y=0)=00=1 leads
8
Pr [1C+| (1- n// cn} S 2 exp (-2nd²)
(56)
16
where the last inequality follows from n > C Colog11/)) /6) in parameter assumptions in Section
3.5
At the end of this section, we prepare an evaluation the token scores for the training data on a good run. Lemma B.12 (Token score). Suppose that the linear head satisfies Equation constants (c+)))()1)) such that vk E {+1}, Cs > 0 and we have for the clean data i E C that
Then, good run, there exist
Cyc (1 - o(1)) ||l1lll11|1 S y()) 71 S cr (1 + o()) |||||||1||
(()) 71 e (p(l112221111)) 00,
(() 7/ = = (al1111211) < <0,
h0 = 00 (o.llvllay/ log(TTn0))
where & E w(i) yigguu W(( -Y v I00 Similarly. for the noisy data j N, we have
-on (1 o(1)) |1 11lllllll22 y-) 21t0 -cr (1+ + o(1)) ||2||||11
Vτ) 7 -0 (p(laelaaaaaaal) <0,
ytl =0)= (pplutlaalla))l > 00,
h = 0 Dallllllev/o(((//5)
where t E W&) E W
-Ym,, E ZO),
Proof of Lemma B.12
Using Lemma B.I and Equation 2 we have for i E C+ that
(()) 710
(#_1 + { {
(57)
= =I|M|lalle||2 cos 0+1 + O(a.lll/l//og log(Tn/6))
(58)
E + = [lllllllll
(59)
where recall that cos 0+1 is defined as the angle between v and μ41- In the last line, we used a Fey/oogTT/// = ([|| in the parameter assumption and replaced cos f41 with constant C41. The other equations are derived as well.
23

---

Benign Overfitting in Token Sclection of Attention Mechanism
C. Proof of Theorem 4.1
In this section, we present lemmas concerning the gradient descent dynamics of token selection, and provide the proof of Theorem 4.1 For clarity, the essential lemmas are included in this section, while minor and more technical lemmas are deferred Appendix D Appendices C.2 and C.3 provide the main lemmas and the direct proof for the not-overfitinng and benign overfitting cases in Theorem 4.1 respectively.
C.1. Preliminary Lemmas
In order to analyze the complicated dynamics of the token mechanism, we introduce the following notations. Definition C.1 (Attention to signal and noise). Let ((r) and p(r) denote the parameters at the r-th gradient descent step.
Then, we define the attention to the signal and the noise as follows:
(+1(7) =
X-((T) = (W(F) P()(T) = (W(T)4', p(T))
Definition C.2 (Weighted interaction terms). and token scores.
We define the following interaction terms weighted by softmax probabilities
(T)
(T)
1t
i
(T)
(x -41):
(T) =
4 (r) (
(x12 p-11)
88
72
(T) =
(0 (T)
(x( -,(
W (T) (T)= =
(T) 1
(W())x W(T)mi)
(r)= Σ 81 (7)
Σ
10 (T) (
(w()) 1 W(T)L-))
TW (T) = Σ (7)
(0)
0 (T) (0)
(W(T)x 9
(WT))") W(T)
7
((()) =
(T)
19 (T)) (W(r)x 1 ,((T))
P(T))
The next definition is the restatement of Definition 4.4 in the main text.
Definition C.3 (Attention gap between significant token and other tokens). We define the following values representing the attention gap between relevant token and other tokens:
Ar,(T)== = (xx]=xxx) (x -x(") p()),
i [n], [T] \ (1). Additionally, we define the attention gap between confusing weakly relevant token and other tokens:
T [T,((T) (x (4) - W(T p((),
for E [n], u E [T] (2).
For clarity, we provide below the results of applying the data setup Definition 3.1 to these values. For E C4 U._
24

---

Benign Overfitting in Token Selection of Attention Mechanism
{i E [n] yx(
(i [n] y -1}. have
(1- (1) + (T, - p(,(T)
for E W
Age(T)
+ (i,((T
for E W = (2):
for E I()
( (p,a(( - -pja(T)
for / +1
(2),
Agr(T)=
(1 (1-p))-1(7)+1 + (7,1(7)- - -Pj,(TT) 1-((T) + PT)((T) - -),((T))
for / E =
for t I
fort (T) (1), and we have
pİ-1(T) - \+1(T) P12(T) P(,((T
for u E R {1},
(T)
pA-1(T) - p)+1(1) + PI.((T) PL,u(T)
for u E W
p)-1(T) + (T2(T) - (T,((T)
for u E Z(1,
P)+1(T) - A-1(T) + PT)((T) -Pj,((T)
for u E R = (1},
(T)
- (4+1(t)-p1-(1)) + P),2(T) - PT)
for u E W
p)41(T)+ pz,(()) - (T)
for I(),
for u E [T] \ (2).
In the following, we calculate the one-step updates for various quantities that will be frequently used subsequent proofs. Lemma C.4 (Updates of signal and noise attention). The updates of (+1(1), 1-) and Pz,(()) for [nil, E [T). which are defined in Definition C.1 are given by
+
(((T) y(. (T) + Vaw
V
A-1(++ + ---((T) = 0
(-)( y(. (Ww.--HH1ooiBBBB.-----eee////T/-
i=1
n
Pj + ---(() = a
Σ
(T) y (1())))))))))))) + + ( (()))([(T)
n
Proof of Lemma
C.4
From Equations 9 and 12 have
W(r+1) p(r +1) - -W))
(60)
-W(T) (-avpī(r)) + (-avw((T)) p(T) + (Vw2(T))) (Vpî(r)
(61)
(((T)) 1 Σ (7)
- (T)( (6) (
V WT V()))
(62)
Following the definitions of À and p. taking inner products with 1411 11 and ( yields the desired update equations.
The notations defined in Definition C.2 are applied here.
Lemma C.5 (Updates of p). The update of [[p(T]] is given by
2x
|p(T + 1)11 - |
Σ
(--((T)) . y() I[(T) + a² (
1=1
25

---

Benign Overfitting in Token Selection of Attention Mechanism
Proof of Lemma
C.5
From Equation
12
we have
|p(T + 1)ll2 |Ip()) = -(p() +|
(63)
2a
(-((()) y [P(T)) + a² |
(64)
where in the last line, we used the notation defined in Definition
Lemma C.6 (Updates of W). The updates of the terms related to (()) are given by
20
[W + 11 (T
(()) y(i) I, (
==1
20
[W(r + -1
(-((T)) y(. (V
n
20
[W(r I
(-((()) y(0) 1
for any j E [n], E (T]. Additionally, we have
(W(T + ( W(T W(++1))m)) - W(T)
a n (-M(T)) (.00::/0......000000040000000000 (W
(W(r 1) W(+1)))))))) (W
(-M(7)) y(0 (1) (V (wwrr))
n
(W(T + 1)
- -WW.
(-M(()) y (,0
n
(W 1)e W(T -(Te))) - (W) W(
(T)) 0 (L. + (1)())))))) a (VW((()))) Vw
for any j, k E [n], u, [T], uj + v)
Proof of Lemma
C.6 From Equation we have
[[W(T 1)141 12 W ((T) 1
= - ((T))) + VlV (VV
(65)
20
((()) y()
(
7t
(5)
Σ
8
) (ITIix")) p(T) ((())e)))) +
(66)
20 Σ (T(()) y(.) 100..14400 (7) +
(67)
where in the last line, we used the notations introduced in Definitions
C.1 and C.2
The other two equations are shown in
26

---

Benign Overfitting in Token Selection of Attention Mechanism
similar manner. In the same way, have
(()) (W + 1) W(T 1)p (W())
= -W(r)) - - - aV VV (W
(V
(68)
) (III*") p(()) ( W(T)HH))
(-(((T)) y(0.
Σ
(
(0
?1
Su
n
1=1
"
T
T
(-((T)) y()
Σ
(1)
(1)
Su 1s
(8)
(uTIx(") p(7) (W(T)-)))
==1
+ a V ((T)()) M41 Vw Vw((T)mL
(69)
((() y(0 (4.-(()A++17)) + (7) ( Vwā(r)m-
(70)
The other equations are derived as well.
In the following, we prepare the lemmas concerning the weight initialization. Lemma C.7 (Softmax probabilities at initialization). Under the parameter assumptions in Section 3.5 and on a good run, we have that for any constant e >> 1,
1 I
(0)
1
T
T
for all i E [n]] and t [2].
Proof of Lemma C.7
By definition, we have
exp
Xi
T
W(O)
P(0)
1
max exp ((x)) (xx[-xx{²) NE[T]
"w()) T(O)) }
(71)
((O)TP(())
T
exp
X
From Equations
25 have
max ME [T]
{
x(
"(O0) [(O)]
< 3e1 T max [lullavd,, (1 +o(l))add) log(Tn/0)
(72)
< log (1+c-)).
(73)
where the last inequality follows from Assumption A8 The same argument applied to the lower bound, which completes the proof.
Using similar argument to that in Lemma
we confirm that the following lemma holds.
Lemma C.8 (Attention gap at initialization). Under the parameter assumptions in Section 3.5 and on a good run, the
attention gaps introduced Definition C.3 at time step T O are bounded as follows:
[A,+(0]== = o (0)] o(1)
for all [m], E (T) (1), and (2).
Proof of Lemma We show the case of E I(s)
Since p < 1/C from the parameter assumptions, we have the same
result for E a and E wi( By Definition C.3 we have
(0)) < 3 max (0=1
max (T)))
1Pp.:
(74)
27

---

Benign Overfitting in Token Selection of Attention Mechanism
From Equations
23
25
have
max (A+1 (0),
(og((T/)) = o(1), (75) (0)3 - oi max (1 log(T//)
max sE(A],EE[T]
where the last inequality follows from Assumption A8 proof.
The exact same argument is applied to (0). This concludes the
Next, show that the ratio of the loss derivative can be bounded by constant. As byproduct, we obtain that the loss derivative itself is bounded by constant from both sides.
Lemma C.9 (Ratio of loss derivative). Suppose that the assumptions in Theorem 4.1 are satisfied. There exists an absolute constant Ce > 0 such that on a good run, we have that for all time step T 2 0,
(T)
< ce.
A (T(T)
More strongly, (T) is bounded both above and below by positive constants, uniformly over all i E [M] and 2 0.
Proof of Lemma
C.9 Recall that the derivative of the loss function is given by
(() + exp (-y(0) ΣT. 5t (T)("))
(76)
for E [n], 2 0. On a good run, Combining Lemma B.12 and ||2l22 = =0(1/| gives the following token scores;
In (0) = 0(1), hr = 0(p). h(4) = o(1),
(77)
for all i E [n]. / E W(a and a E Z(, Thus, there exists some constant e > 0 such that 1n6²0 < e for any e [n] and e TT. Since 1/(1 exp(-z)) is monotonically decreasing, we have
+ exp(c) (T) < + exp(-c)
(78)
This leads to the conclusion with the constant c = 1 exp(c)/ (1 + exp(-c)).
Remark C.10 (Ratio of loss derivative). This lemma, which shows that the gradients of loss function for clean data and noisy data remain within a constant factor of each other at every time step, is a critical component of the proof in the existing
analyses of linear classifiers and two-layer neural networks (Chatterji & Long 2021 Frei et al. 2022 Xu & Gu 2023 However, in the learning token selection, the output always an affine combination the token scores, and the output scale is not changed. Therefore, the training dynamics need not be considered as long as the balance of the loss derivatives in the token scores is maintained. To ensure that the derivative of the loss function for each token remains within constant factor, small linear head scale, described in Lemma gradient will be generated for clean data even at the initial weights, and learning the signal vectors will not progress.
CV is required. If the scale of the linear head is too large, little
C.2. Analysis of Not Overfitting Case (SNR²
In this section, provide the proof for the not-overfitiing case in the main theorem. We first present the main lemma using mathematical induction and then proceed to prove the main claim.
C.2.1. MAIN LEMMA
Lemma C.11. Suppose that the signal-to-noie ratio satisfies SNR?
-w(nn) For any time step Ti = M max on a good run, the following propositions hold for all time step T E [0, Ti]:
A(T): There exists constant C such that
(1 - 1/Ci) Vn ( .
28

---

Benign Overfitting in Token Selection of Attention Mechanism
B(T): There exists constant C2 > 1 such that
(1 - (C0 ( aullallvdd
(1  (1-V0))))llllllkkd0000000)3211)) ( amllull2Vd,
(1-1 -wo l1 owOed
for all i [n] and t E [T].
(W(T) (oll:
(W(F))LI W(T) (0) [ (n))))
(W(F)&) ( W(T)e < 0 0 1og
for all i,j E [m] and I, u E [T] such that (i,1) # u)
C(T): Let g be a monotonically increasing function g(x) 2x 2sinh (x log T). Then, there exist constants
C3, C4 O with cg C4 such that
g (At,(7)) 2 g (A), + c3a caal:llll
og),
9 (Acu(T)) S g (A),(0)) + c max 771
for any i [n] and E [T] \ {1}.
D(T):
s (T) O(1,
for all [n].
E(T): There exists constant C5 >> 1 such that
exp (A -A,())) cs.
for all i.j E [n] and t, u E [T] (1). F(T):
There exist constants cs. C7 0 such that
(0 (7) < <6.
JEM
max
0)
(T)
and
(T) (1) (T)
1 (7)
5 (0 (T) < <
$
max A.ETT(()
< C6
( (T)
for any i [n] and E (1).
G(T)):
D+1(+)]= 0(1), (--1 0(1), |
for any i E [n] and t. E [T].
We will prove the above propositions by induction argument for E [0, Ti).
29

---

Benign Overfitting in Token Sclection of Attention Mechanism
Base case. It is obvious that C(0) holds. From Lemmas E(D), F(D) and G(O) hold.
B.1
C.7
and C.8
the all other base cases A(O), B(0), D(0),
((T)) = 1() We only show the case of y = 1. i... i E C, U._. The same argument can be applied to the case of E C_ U.N. The conditions of Lemmas D.6 and D.7 are satisfied from A(7) A B(7) D(F) A F(r). From these lemmas, we have that for any i C+ U.W,
Proof of ArrST (A(7") A B(T') C(T") A D(7)
A+I(T + - (+) 2 of
of).
(79)
and for any clean data i E C+ and t E [T] {1}, we have
PT,(T + 1) Pt,(()) 2 (7) -( -famm o221lllllllllll max of:),
(80)
Pt,e(T + 1) p(.(TT) -ss (r)s (T)
caan I max of
(81)
and for any noisy data j E EN and t E {7]] \ [1), we have
PT + 1) - pT,((T) 2 -s (T) (p (F)
d m og)
(82)
PT + -pj((T) S ( (-)) -han '0 ||2l| max og).
(83)
From the definition of Ai, in Definition
C,3 we have for i E C+ and E I that
A + 1) Ar,(T)) (1+1(+ -+1(7)) (p,,(T + 1) pa.(T)) (pt,0(++1)) (s,((T)
(84)
2 - (T)
(ci
11111 1 + cnn oo?d) max(o of)}
+ (T)s (T) (T) Gan of 111ld max o/)
(85)
2 (r) - (T))
11d max o7)
(86)
Here, we bound with the signal update because it is dominant from the condition SNR² = " i.e. For noisy data> j E N, the signal update becomes dominant under the current SNR setting as follows:
|| = w(n
Aga(T + 10-1)--A0 - AAja()) 2 ( (r) =f - |||ا| 2d (cillиa1² -
max o7)
- 5 (T)s
(3) (r) caan 2 212) d2 max o
(87)
2 2 ( (-) - (T) G0|l max a0))
(88)
for some constant c4 > 0. Since the same lower bound is obtained for clean and noisy data, we proceed with the discussion
for i E C., U._. We note that we have the same result for weakly relevant tokens L E w U w( When L E wi()) the 1/C, the problem is reduced to the case of by substituting
update of X41 is multiplied by factor of (1 11-p) p). Since constants. For t e w(() we also have to consider the update of A_1. However, by similarly applying Lemmas D.6 and D.7
along with the result of the softmax probability ratio in F(T), the case for t ( reduces to that of e w Therefore, we obtain the following upper
the above inequalities hold for any & (T) \ (1). Similarly, from Lemmas D.6 and D.7 bound:
As, + 1) (),)) s s| (((()) ca | ||
e ma (o og)
(89)
for i E C_ UN_ and [T)
We will evaluate the softmax probabilities in the coefficients of the above inequality using Lemma D.1
As we will see later, the increase or decrease of the bounds changes at (T) = log T, so we divide the analysis into cases before and after this point. At the beginning of training T = O. we have |A,, (0)l = o(1) from Lemma C.8 Since (r, is monotonically increasing from Ay's,C(r") there exists a time step changing to ATT(T) > log T. We proceed with the analysis by considering cases before and after this transition.
30

---

Benign Overfitting in Token Selection of Attention Mechanism
Case (As,((T) < log T): From Lemma
D.I have
c-
(T) (7)) 2 2 2cosh (A1,(T)) - log
(90)
Substituting this to Equations (86) and
88 provides us
(2 2 cosh (A) - log T)) (T +1 AA,((T)
of of)
(91)
for some constant c6 > 0. To apply quadrature method to Equation 91p.. will replace the coefficient with A + 1). Using (T)(T) < A1, (T + 1) and the convexity of cosh, we have
cosh (A 1) log T) 2 cosh (Ar,(T)) log T) + sinh (AL - log (Aci(T + 1))((++))),((77))
(92)
2 cosh (ALa(T) - log T) (1 (A + -
(93)
where the second inequality follows from sinh(x) 2 cosh(ar) for all æ E R. Using [A + Aco(T) from
Lemma D.4 we have
(1 +1+00 (()) cosh (A + 1) - -o0) 2 cosh (Ai,(()) -log T).
(94)
By substituting this to Equation 911 we have
(2+ cosh (Ai,((T + 1) - log T)) (A + +1) Ar,((T) 2 Gallalllll max o/h
(95)
where e/ = (1 -())) Since this inequality holds for any 0 S T' S from the induction hypothesis, summing both sides from 0 to 7 gives
Σ (2 + 2 cosh (A,, 1) llgg T)) (A,,((T' +1) (,)(())) (T+ 11+11))aalll r'=0
Galle ( max o (96)
Recall that g(x) = 2x + 2sinh(x - T). Since AA,(TT) is monotonously increasing in 7' E [0, + 1). and the function
2 + 2 cosh(r - log T') is monotonously decreasing in this case, we have the following lower bound by the quadrature method:
(
g (A), +1) - (Ac,(0))
(2+ 2csh (x - log T)) dr
(97)
SA..(0)
2 Σ (2 + 2 cosh (An,(T'+ + - log (An,(T' +1) - Ar,((r))
(98)
+|=0
2 (T + 1) G |11 max ogh
(99)
which gives the desired result. Therefore, we have
g (Ace(T= + 1)) 2 g (0), + +T Galla/lal|l1 max (o o.)
(100)
We also provide the upper bound. By applying Lemma D.I to Equation 89 there exists a constant de > 0 such that
(2 + 2 cosh (Ac+(T) - -log lo) (Ar,(7 + 1) - (r,(T)) s c [|lalll1||@ ( max(o? -
(101)
Since this inequality holds for any 0 S T, summing both sides from 0 to and using the quadrature method, we have
(re(T+1)
(Ar,(( + 1)) - 10g-gg/9 (A),
2 2 cosh (x log T) dx
(102)
AA.(())
<
(2 + 2 cosh (Ace(T') - log 7)) (Ace(T' +1) - Ai,(T'))
(103)
+'=0
< (7 + 1) 6 M 1elllaalllll@dxx((l og.
(104)
Therefore, have
g (M + 1)) S (A,, + (T+ ( GallelballelBddaxx oz)
(105)
31

---

Benign Overfitting in Token Sclection of Attention Mechanism
Case 2 (log? T (T) Let To be first-time step falling into Case 2. Since + 2 cosh(a log T') is monotonically increasing in this case, by summing both sides of Equation (91 from To and using the quadrature method, we have
As,((+11)
9 (A + 1)) (Ao)
2+ 2 (a log T)) dr
(106)
(To)
Σ
(2 + 2 cosh (A,,(T') - log (A + + -1,-)(((77)
(107)
7'
2 (T To + 1) 6 Gallllllllll o/).
(108)
Combining this with Equation (100p, we have
9 (Ace(T (T + 1)) 2 , (A; (0)) + (T+ Gallee/22l pell max oz
(109)
which follows from & (1 - ((1)) < c.
Next, we provide the upper bound. In Case 2, to obtain an upper bound using the quadrature method, it is necessary to replace the coefficient Equation (101) using (T, + 1). Since cosh(x - log T) is convex and monotonically increasing in this case, we have
cosh (Ai,e(T + 1) - log 7) I cosh (Ar,(r) - log 7) + sinh (Ai,e(T + 1) - log T) (Nea(T + -1)-)((+))) (110)
Using |A + 1) ((,((T)) = (1) from Lemma D.4 and cosh(r) < sinh(x) for all R, we have
cosh (A,,(( + 1) - log T) (1 o(1) cosh (A,,(T + log 7') (1 (Ar,(T + Ar,(T)))
(111)
cosh (At,(T) - loo T).
(112)
By applying this to Equation (101). we have
(2 + 2 cosh (Aia(T + 1) log (A (Ar, +
Ai,,(T)) aolpe//llllllllllll/xxxxoeoo o/
(113)
where bg = (1 + +(1)ccs Again, using the quadrature method, we have
(1+(++1)
9 (A + 1)) (Tr) =
2 2 cosh (I - log T)) dr
(114)
AT,(
<
+2 2 cosh (Ar,(T' + + - log T)) (Ar, +1) - Are(T)))
(115)
7¹ == To
$ (T
To 1) call 1l22 max oz)
(116)
Combining this with Equation (055), we have
g (ALi(T + 1)) S 9 (0), + (r 1) Gallaelllll o7),
(117)
which follows from cg (1 + o()) cg.. This concludes the proof.
Proof of C(T) D(T). For any E [m]. we have
(T) 1 Ei= exp (-A,,((T) 1 El exp (-A,,,00) s (0) 2 0(1),
(118)
where the incquality follows from ((T)) which states that (T) is greater than (0). The last line is derived from D()).
Proof of C(T) + (T()). We proceed with the proof by contradiction. Suppose that E(()) does not hold, i.e., for any constant e 1, there exist i,j E [n]] and t, u [T]] (1) such that
exp (A,,(T) - Aj)(T)) 2
(119)
We continue the analysis, dividing into two cases.
32

---

Benign Overfitting in Token Selection of Attention Mechanism
Case (Ag,((T) log 7): In this case, from C(7) and the monotonic increase of g. we have
9 (A (0)) + 03al|ull2 I max ottt S (A و (2 < 4og + T T Since 9 (0), = -(1 (1
(120)
(l))T from the definition of g and JA (0) o(1) in Lemma C.8 have
37
(121)
Again from C(T), have g (Acx(r)) - g (T)) S (A)) 9 (A),a + (cr [ og))
(122)
3(c4 c3)
3e4
<
T
T.
(123)
C3
C3
On the other hand, under Equation (19 we have
g (Ar,(T)) (T) 2 2 2oog e + 2 (sinh (Aj,u(T) (T) ) + log c - -log T" -sinh (Aj.u (T) - -o)
(124)
c-1
> 2 c
T exp (ap(Aj((F)) + 1
exp ( (T))
(125)
(Vo- Vo
3c4
> 2 2lo +
T.
(126)
C3
The last line is the result of AM-GM inequality, the parameter assumption T = O(),, and an appropriate choice of c. This leads to a contradiction with Equation (23).
Case 2 (2 log T < Agu(T)): From C(T), we have
(Ar,(T)) (0)
<
(127)
(T) (A (0))
for T 2 1, Additionally, from the definition of and Lemma C.8 we have that
(0) (0) + + exp (At, T exp ( (-))
< 11+0(1) <2,
(128)
9 (O) a. (0) + exp (A,, T exp (-Aj. (0))
where the inequality comes from the fact that the third term is dominant and (0) - A0, = o(1)
Before delving into the part of contradiction, first show that the following inequality holds. By simple calculation, obtain the following for > 2 lo
c
2 sinh (xr + log c - log T) = exp(x) - exp(-=2 (x-
(129)
T
2c
(c (c_ 1
2 sinh (x log T) + exp(z) +
T Texp(-2)
(130)
3T
3 c)
2sinh (a log T') +
c
+
T
c
(131)
3
3
where the last line holds for constant satisfying c/3 1/c, and used exp(x) 2 (2) Tr for 2 2og By using Equations (199. (128) and [131. we have
y (Au)(T) - y (0)
2Acx(T) + 2 sinh (Ai,(T) - log 9 (A))
(132)
1
2 (A) + log c) + 2sinh (A (T) + log - log T) 7 (A)
(133)
2 (2 + c/3) A Aj.
+ c/3 - 2inh (A),u(T) - log T) + ((1 o(1))c/3 + 1/2) (2122((00))) (-g(As,(00))
(134)
> ca/c3 (g (A - (A,,a (0)))
(135)
where the second inequality. we also used - (0) = (1 o(1) T. The last line holds sufficiently large >> satisfying + c/6 > c4/C3, c/3 > ca/03 and (1 o(1) c/3 1/2 > 24/03- Therefore, this contradicts Equation 127
33

---

Benign Overfitting in Token Selection of Attention Mechanism
Thus, since a contradiction arises in both Case and Case 2, it follows from Equation 119 that there exists constant 1 such that
exp (Ari(T) - AAj((T)
(136)
for any i,j [n]] and i,u u E (T) (1).
Proof of E(r) = F(r). We have
Σi= exp (-Ar,(T))
(1) + exp exp (-Aee(F)
(137)
Ei=g exp exp(-Ai, (-Ai,(7))
Then, we have
o (T) (()
2 exp (r,((T))
T 1+ 2t exp (app-AAjri (-Ajie(T)) 2
(138)
(T)
Eiz exp
(-A)((T))
Σr Σi=y exp (-,,((T))
1 (-)r,(T))
1. 1- (-)))(T))
-mxx
(139)
21=2 exp
(-))
1= exp (-),,(TT)
1= exp (-A(((T))
E 1=2 exp
AT,(T))
max
(140)
2i=2 exp
(-A),((T)
2 exp (A),((T)
exp
(-)))(TT)
exp (-A),(())
S max
max 251ST
(141)
exp
(-Ai,(()))
2SEST
exp (-Arx(T))
where the inequalities are the result of mediant inequality, i.e., (E,a:/ (. < max:(as/b;) for a; > 0, Vi. Therefore, using E(T), we have
s (9 (1- S Sc.
(142)
(JE s (0) (T) (()
Additionally, we have
X ( W(T) [ ((T))
(T)
cxp
exp (-Ar,(T))
max ()) (T)
max ())
(.E(T)(()) exp (-Ai,((T) I Co.
max
(143)
exp
((T))
for any i E [n]. Finally, since Aia(T) is monotonically increasing for any 1 E [T] \ (1), Lemma
D.3 concludes that
(1) (7) 81 (T)) is dominated by 81 (6) (7) - of (T)) ignoring constants.
Proof of ArrT (A(7') A B(7') A C(T))
F(r)) G(r+1). The conditions of Lemma D.6 are satisfied from
A(T) A B(T) D(T) A F(7). From Lemma D.6 we have
A41(T -4
(0) <
Σ s 7'=0
callullall [ax o
(144)
We use Lemma D.5 to derive the upper bound of the summation of softmax probability terms. Since Lemma D.5 follows from Ar'gr ( E(7') and the definition of Ti, we have
(1)11
(T))
(145)
1ll2 p mdx oz}
Substituting this to Equation
144
we have
A+1(T + 1) (0) O(1)
(146)
34

---

Benign Overfitting in Token Sclection of Attention Mechanism
Since (0) o(1) from Lemma C.8 Equation (146) leads \41(T + 1) O(1). In the same way, have 1-1(7 1) & (()). For noise memorization terms, from Lemma D.7 have that for any [n].
+ 1) (a) - Σ sf can of oallalaalle max of].
(147)
We proceed with the same argument as in the signal update case discussed earlier. From the current SNR condition SNR (( and (0) o(1), we have (T,( I o(1). The bound of pT,((T) is obtained well.
Proof of Ar'ST (B(T)(((T)) A C(T) A F(F) A G(T)) = A(T + +1) It follows from Lemma C.5 that
||p(T + 1>1 ||]0)]
Σ (IIP(T + 1) 1IP(7)
(148)
2a
(-((T)) y 17(7'))))))))))(())) + a² Σ vp((tr)
(149)
1=1
r'=0
By definition, we have
T
[((T) Σ
81 (+))
71
(1)
-
0 (+)) (W(F)x)). 4
(150)
==1
11 (r) -1-)((
max { 176 -
,(7'))
(151)
"E[T]
S (t)( - (79))
max
[1:411] (E[2]
max (444))())))))))
(152)
where in the last line, we used the results in F(T') and T (1) in the parameter assumptions. By substituting this into the first term of Equation (149), we obtain
2a
((()) y (0) ((l))alulll)l))))
Σ
(()))))((((()))))
-1 max (((=((())))
1=
(153)
2 a
a | 12111121
1
(154)
max of]
(155)
max ob)
where the inequalities follow from Lemmas B.12 and C.9 the results in G(7'), and Lemma D.5 Equation (15) becomes (o).d). From the parameter assumptions in Section 3.5 we have
We confirm that
| o?
0² maxfo oz) 11111/d
2 min
d.
(156)
log (Tgn/T//))
C2/3/1
2 min
C2 1/3
3/18
(157)
"log" (Tn/8)
=w(1),
(158)
where the first inequality follows from Assumption A8 Equation 155) becomes d))
while the next one relies on Assumptions Al and A2 Thus,
Finally, we confirm that the second term in Equation 149 can be ignored. From Equation 12 and the results in B (T), we
35

---

Benign Overfitting in Token Sclection of Attention Mechanism
have
T
T
2 o (T'' ((+))x)))
V pp((T)] a² Σ
(( r(.
Σ 3t
1 (7')
(6)
Σ
(159)
=0=0
2
(5/ (
(he (1x2(1)
(T)
[12]
max dEim],ue[2)
max se|mi,t([T]
max se(n],te[T]]
(160)
|||اا|||| Σ of (1) (allu)lallalll0 max{
|aax|||lolle pl o²d
(161)
where the last inequality follows from the dominance of o (1) - (T)) established by F(r'). Using ||]| ((1/[]a)]2) and the parameter assumptions a, R. the latter part of the last line becomes o(1). Consequently, this
quadratic term is absorbed into Equation (153). Combining these results with Equation (149). we have |p(T + 1) e
(1 + Dod, for some constant C1 > 1, which concludes the proof.
Proof of Ar'ST (A(T') A C(F) A G(T)) = B(T + 1). From Lemma
C.6
we have
[[W(T + 1)21)1111 - ||
(I + 1) 112 - (1)())4+1l1))
(162)
+'=0
2a
R
(()))) (F)AA())))))) + Σ WW
n +==
y'=0
(163)
We analyze the two terms separately. Recall the definition of Is, in Definition C.2 and using the same discussion as Equation (122) in the previous proof for A(T + 1), we have
1,+(7') 5 (7) - (r)) max
1h iE(T).rE(T (1))(1=1(())) |1
max rE(o).rE(T)
(1(x()
14111
(164)
(165)
which follows from Lemmas
B.1 and B.12 Therefore, for the first term, we have
2a
((7')) y (7)444()))) & (1)) a:/1e/2211141
(166)
1=
d max (o., oz) = ol m 111
(167)
where the inequality follows from G(r'), Lemma D.5 and the same evaluation in Equation (18 The analysis of the second term in Equation
(63) follows from similar argument that of Equation [61 so omit the
proof. Therefore, Equation (63) leads
WW 1) 12 = W(0 + (m([d = (l11)l
(168)
The same result holds for [[ + 1)m,, ( H (d). Similarly, we have
= | 1)e; ( 112 = (()) +
(169)
Additionally, we will show the case of the inner products as well. From Lemma C.6 we have
(W(T + 1) W(++ (W(0)44 W(0)
(())) y0 (J.- + (.,4())))))) d² Σ (V V (170)
36

---

Benign Overfitting in Token Sclection of Attention Mechanism
For the first term, from similar argument to Equation 167
have
((T)) y ( (T)-(()))+1()) + (1,(0)))177))) dmax fo::
O Olooolll3
n i==
o ob))
(171)
where the last line is derived from the parameter assumptions; specifically, we have
σ m o2} -11 (oo.onnn////)o Z min { 11/2 1
(172)
o²d
log (Tn/6)
d1/2
Ca
2 min
(173)
(Tn/8)
(og(Tn/8)
2(1).
(174)
The second term can ignored similar manner, and we have
(W(r 1) W( + 1)p_)) = =0 (0 /
(175)
We will show that the other equations are shown in the same way. Since the quadratic term can be ignored by the above discussion, we denote this by O(a?) ) in the following. Similarly, it follows from Lemmas B.1 and C.6 and SNR wn
that
(W(T + 1) W(T + 1) - (W(0))41 W(0)
(T')) y()) (1 + La,t (F') +
(176)
n
odd
max
n
+ +0(a²)
(177)
dd fox )
max oo]
(aa)) O log
(178)
d dmax o
which follows from the parameter assumptions; specifically, we have
{ d
||ll|2
1
t max -, 2 log(Tn/6) 2 min
(179)
1llll o²
(Tn//Tn/0)
log
l|ul|2
2 min
Comillull2
1/3
(180)
o? (Tn/6)
= 2(1)-
(181)
The last line can be derived by combining Assumptions Finally, have
and
A2
The same result holds for (W(++1)m_,,
((++))cd))))
(W + 1)c(f) W(r + +1c(" (W(0)c?) (())c!
(((')) y(. (T)),(('))))((1)) + (,j.a("))),, O(a²)
(182)
n
o2d max o2)
+ 0 log (n()))))),
(183)
||
which follows from the same discussion as in Equation 172 This completes the proof.
Proof of Lemma C.77 At time step = 0, all propositions hold from the proof for the base case. For the next time step, A(T + 1), B(T + 1), C(r C(++1) 1) and G(r + 1) are proved based on the propositions up to time step т. As for D(T + 1), F(T + 1), and F(r 1), they are derived from C(T + 1). Thus, the proof is completed by induction.
37

---

Benign Overfitting in Token Sclection of Attention Mechanism
C.2.2. PROOF OF NOT OVERFITTING CASE IN THEOREM 4.1
In this section, we provide the proof of the not-overfitting case in the main theorem. We divide the proof into two parts: behavior on the training data and generalization performance.
Training data. The condition of Lemma C.11 is satisfied with the current SNR condition SNR w(n The
proposition C(T) in Lemma C.11 states that there exists time step Ti O oilpllalllllddxxo..... such that
((Aa,(Ti)) 2 (O),, + @(1) dj.
(184)
for arbitrary constant 24 > 0, for any E [n] and E [T] (1), Here, we used (()))(0)) -1 o - from Lemma C.8 Recall that g(x) 2r + 2 sinh(r log T), and there exists a constant ca such that bg exp(x) > g(x).
Therefore, for all i [n]. have
(T)) 1+ 2i exp (-Ar,
Σ exp (-Ac,(Ti)) -T- 1)
> -6,
(185)
cj
for sufficiently small constant >> O. By using Equation [85 and Lemma B.12 we have that for any clean data C,
y(, fr (x y(i) X 8 X W (Ti) (T:))((T:)) p(Ti))
(186)
21 81 (T)) + Σ y(i) 7 (Ti) (0)
(187)
@ - 11tpb/lllll11-00000000lll/l//// (111111211111))
(188)
0.
(189)
Similarly, for any noise data j E N, we have
y fr, (x()) < [ (1 c) +
O.
(190)
Equations 189) and (190) hold deterministically on a good run. Therefore, at time step r = Ti, we have that with probability least 1 8,
Vi E C, fx( Y EjEN, f+X()) YY))
(191)
Generalization. Let (X. P be the unseen data on which we investigate generalization performance. We first evaluate the attention values of signal vectors at time step Ti. From Lemmas D.5 and D.6. have
-n1(1))A=1(0) N
Σ (T) -()(1))((()))) calle1121 d of)
(192)
> Ti al|illa max of
(193)
2 1
(194)
which follows from the definition of Ti. From Lemma Similarly, we have (T,) 2 1.
C.8 we have [4+1 = o(1), which implies (41 2 1.
Next, we show that the attention scores of the noise vectors (Etherry in the unseen data, i.e., e, (T p(T), become sufficiently small on a good run. While it natural to evaluate (T, pTT;)]l2 and use concentration inequality, is challenging to track the evolution of [W()) (22) Therefore, following induction proof in Lemma C.11 we apply a concentration inequality at time step and show that the result does not change at time step Tj. Let us define & as the event
38

---

Benign Overfitting in Token Selection of Attention Mechanism
that the following inequalities are satisfied:
E [T), (1 - 5| |2 2 (1+ o o.v,
(195)
A (7), vA: [+1}, (x), <oolltlla///ll
(196)
WE [n], vt, u E [T], [(,, (
(197)
14 (7). 44 (+1), (W(0)e, <
(198)
3A [m] 14 [2], (W((),, ["o+=""))==c"llllllee""///////""
(199)
4t E [T), p (ow.o///////////////////o///b//
(200)
14 E (7), ((,, < <20.ll/la//og((/m0))
(201)
where the constants C1 02 the same ones appeared in Lemma B.1 Applying union bound on the modified versions of
Lemmas B.6 B.8 and [.99 the probability of the occurrence of & can be evaluated. Since there is no need to apply the union
bound over the additional n training data points, the outlier probability can be reduced by 1 1/n compared to the original lemma. Therefore, we have
Pr Pr( > 8/n > - 6.
(202)
In the following, using the results of Lemma condition E:
C.11 we will prove that the next proposition holds for all r E [0, Ti] under the
H(r):
(WWoeeepphi (W(T),, < (TliadloggT////0 (Tn/o)). (W(()); ( (Tn)0))
(W(T)), ((T)c o 3/2 log(Tn/8))
for all i E [n]] and l,u E (T].
Proof of HT)). We proceed with the proof by induction. The base case holds from the condition In the following, we suppose that (7') holds for any 7' [0, 7]. By calculation similar that the proof of Lemma C.4 have
(W(r + 1).. 1)) - (W(O)e,,(())
Σ (W(r' +1),, +()
+ 1)) (W(r)),,
(203)
(100)
2 (T)))) ) (IW(')e, (((())) ( + [ x(1))
(T)))) y0 Σ
(
(7')
8
n
x=1
+ Σ
2 E
V wT ((7)))
(204)
7'=0
39

---

Benign Overfitting in Token Selection of Attention Mechanism
Using the induction hypothesis, the condition E, and A(7') in Lemma C.11 the first term is bounded follows:
((())) y(
(8)
0 (+') )
(IW()), (WW(")s w(F'x{"> ( + [I+(+))B(e) )
aa
(7) (7))
max
(oe o o.d (l log(Tn/0)
ie
+ a
Σ -==0
$1
(0) (7') 59 (T))
max
(Te { } ofd o. max
(205)
[E
|| max oz vox log(T/
(206)
a 1a|p | dd max (o og]
M max | oe d|| ² (Tg(Tn/6)
(207)
= 0 po, log ())
(208)
where the first line follows from ((r'), in the same manner as Equation 152p 152p. and the second line follows from Lemma D.5 The last line is derived from the parameter assumptions in Section 3.5 specifically, we have
||ll
Tn 1appl1l1/222 Z
min
Vd,
92(1,
(209)
log (Tn/6)
| a. log (Tn/o) C2dl/4 22(1).
Tw 2 11 Z
min
min
(210)
ae
o? Vd
log" (Tn/o)
For the quadratic term in Equation
204
we have
Σ a'c, W ((T')v(((T))
7==0
2 V (w
(211)
-==0
Σ
$1 (7))
iE(m),e1(1] (17
max
(1="1))
max
(0) (c] x H|p(T)ll2
max (T),
(12) (212)
iEjm),tE(m]
r'=0
(213)
a a|||||l|||1 3a o2]
max (onn/// Vd Vdlog(TT/6) opvd ow max(lpllavd
(214)
(max(]lull2, a, log log(Tn/0)) (mmxx(1eell2.
(215)
o(o_oyoodlog[(n/())
(216)
where the first line is the result of the Mouchy-Schwarr inequality, and the second one follows from the gradient updates provided in Equations 9 and (1)). The third inequality follows from the same discussion as in Equation (206), the condition
E. and A(7') and B(T') in Lemma C.11 The last inequality follows from ||M||2 = ((1|||m||]) Using the parameter
assumption on a, Equation (215) is absorbed in Equation (207). and the discussion is reduced the first term analysis.
Therefore, substituting Equations (208) and (216) Equation 2044 leads
(W(T + L)e,, + 1)) (W())), p(0)) + O (owo,oddloo((nn/)) O (ogga
(217)
The proof for the three inequalities below in the proposition H(T) is omitted because they can be shown in the same manner B(T 1) in Lemma C.I1. under the induction hypothesis and the results of Lemma C.11 Since H(T + 1) holds under the condition (T) is valid. from induction argument, the proposition (T) holds for E 0, Ti].
From H(Ti) and Assumption
A8 we have
d/3 log(Tn/6)
((T)), & o
(218)
40

---

Benign Overfitting in Token Sclection of Attention Mechanism
for some constant 2g Using Equation (944) and taking sufficiently large Ti, have (41(TT) > 21 and A-1
2--- From Equation 218
have that for E [7] (1),
x1) W(T)
-(1 p) max (++1 (Ti) ( (E: E1) w(TT)' (W(T) p(T)) ((T))
(219)
2/3
(1 1/C) 20g
(220)
log (T
<
(221)
where the second inequality holds by < 1/C. Then, the softmax probability of the relevant token is as:
1
T
(Ti)
>1
1
(222)
1+ E I exp (x, x1) W(T:)
expley
for sufficiently small 6 > 0. Consequently, we have
. .T, = Y (T)) Σ Y (Ti
(223)
(1Ie12l11))))) (1 c) - o (pIlellalll|l)) e 0.
(224)
Under the conditioning on &, the output fr, (X) deterministically takes the same sign the true label y*. Thus, the generalization error is bounded as:
Pr (XP- (sig (X)) Y Pr
[sign(fr: (X)) + Y
[8]+
Pr (X,Y*)~p
(89 189< <
(225)
(xP
where we used Pr(A) & Pr(A|B9) + Pr(B) and the result of Equation (202) This concludes the proof.
C.3. Analysis of Benign Overfitting Case (SNR² = o(n ())
Next, we proceed to the analysis of benign overfitting. Although the proof by induction follows similar structure to Lemma C.11 the proof is a two-stage analysis divided into Ti] and [T] address the behavior of noisy data. In Stage 1 (T E [0, Ti), we demonstrate that the probability of selecting the relevant token x(") for noisy data j E W becomes sufficiently small. In Stage 2 (r E [Ti, T2)), we show that the model selects confusing weakly relevant tokens that can the noisy labels. While the coefficients differ from Lemma C.IT many arguments are shared between the not-overfitting case and benign overfitting case. To avoid redundancy and unnecessary page length, the proof repetitions are referred and omitted.
C.3.1. MAIN LEMMA
Lemma C.12 (Benign Overfitting, Stage 1). Suppose that the signal-to-noise ratio satisfies SNR² = o(n For some time step Ti = on con-toz[|e|aall I22d max(o2.03] ) on a good run, the following propositions hold for all time step t E Ti]:
A(T): There exists a constant C₁ > such that
(1 (1-1/C1) lprt v,yd.
B(T): There exists a constant C2 > 1 such that
(1 - 1-1CC)) o
ow oulldll
1/C²) ( ( a=lld|/2Vd
(1 (1-1/C4 1/CC) a.a.d < [[(7) (| < (1+1/Cc) awO.d,

---

Benign Overfitting in Token Selection of Attention Mechanism
for all i [n] and i (T].
(W Wt
(W(T))4. =
(W(T))-1 W(T) =
(W(T)& ( (W(T)) < 0 ( 00 (Tn
for all i,j E [n] and i, u E [T] such that (i,t) + uj.
C(T)): Let g be monotonically increasing function g(x) = 2x + 2sinh (x log T". Then, there exist constants
C3, > 0 with ez < such that
9 (A),((T)) 2 y (0), + T cz0m - 12 max (o m ogh 9 (A), 9 (O)) + cian 22 max op)
for any clean data i E C and t E [T] \ {1}. Addirionally, we have that for some constants es < Co.
9 (A) - g (A),(00) csam - 0 max max (o o 9 (A)a(T)) 2 9 (A),((0) caon 0221||1121|111| max fox oz)
for any noisy data j E N and t E [2] \ {1}.
D(T):
(7) (f) 2 O(p),
for all clean data i and noisy data j N. E(T): There exists a constant 07 > 1 such that exp (Ar, Ag,((T)) < e7, exp (ARa(T) Ar,((T) < cT.
for all i,j E. k, E and u E [T] {1). Additionally, we have
exp (As,(T) + (T) < er, exp (-Ar,((T) - An.a(T)) < CT.
for all i E C, N and uu E [T] (1}.
F(T): There exists constant Ca > 1 such that
(7)
(4) - (())
max JEE (T) s() <88.
Additionally, there exist constants C10
C10 > 1 such that
(T)
(T) - 5f ())
max (T(()
C9.
< <10
(T)
(7) - (T)
for any [n], i' e C, and E (1).
G(T);
[A+1(7] ologg 1) [A-1(T) ollogg 1), (p.,((T) O(logp ))
for any i [n] and e [T].
We will prove the above propositions by induction argument for E [0, Ti).
42

---

Benign Overfitting in Token Sclection of Attention Mechanism
Base case. It is obvious that C(O) holds. By Lemmas F(0) and G(0) hold.
B.1
C.7
and
C.8
the all other base cases A(O), B(0), D(0), E(O),
+ 1). We only show the case of Y 1, i.c., E C+ U.N The same argument can be applied to the case of E C_ U. In the following, we bound the attention gap introduced in Definition C.3 from both above and below, in the proof of Lemma C. for the not-overfiting case. The conditions of Lemmas D.8 and D.9 are satisfied from A(r) A B(T) A ((r) A F(T). From these lemmas and the definition of Ai,, we have that for any i C+ and & I(1,
Proof of Ar'ST (A(7') A ((T)) A D(7) ((T') = C(T
A + 1) 1 AA,(7) (A+1(T + -+(())))))+()))----) + (P:,(T + (., - (pi,((T + +1) -t,((T))
(226)
- (f
(c)d((1)))) (c + en - o²d) max(o" op))
22 fo og)
(227)
28 (T) - 1 (T)) San o max (o op)]
(228)
Here, we bound with the noise memorization term because it is dominant from the condition SNR? = o(n-1), i.e.
| = o(n-to]d) For noisy data j E N_,, using 81 (T) - (T) 2 of (7) +()(7)) from F(T), we have
(r)) for any clean data i
Aga(T + 1) -f (all))aall
- (T) daam b0ll maxfou,
caam 0' + (f
(229)
$ - [|llaaldl (c en o²dd) max (o of)
+ (r)) (f) (r).cam -f max of
(230)
s si (T) - 2/ (r)) of
max oz)
(231)
We note that we have the same result for weakly relevant tokens t w U W When E W( the update of X41 is multiplied by factor of (1 p). Since the noise update is dominant in this case, it does not affect the above bounds. For I E w(( we also have to consider the update of A-1. However, by similarly applying Lemmas D.8 and D.9 along with the result for the softmax probability ratio F(r), the case for E w( reduces to that of E w Therefore, the above inequalities hold for any E [T] \ (1}. Similarly, from Lemmas D.8 and D.9 we obtain the following upper bound:
AG,(T + 1) -n,(T)) < s (0) (T) 1/ 3/0 (T)) Gan | max of),, AgA(T + +1) 1)(T+1)-AMp) - ( s (T) cgan -2 | 22 max oz)
(232)
(233)
for b E C_,j N_ and t [T] {1}.
The evaluation of the evolution of A is identical to the proof of the not overfitting case in Lemma C.11 differing only in coefficients; therefore, we omit the discussion. For noisy data, since (T) is monotonically decreasing, the evaluation in Case (A) < log T) suffices complete the argument.
Proof of C(r) = D(7). The lower bound of clean data is derived from C(T) and D(0). To derive the lower bound for noisy data j E N,, we use a similar technique to Equation (419) in the proof of Lemma D.5, From C(r) and Lemma C.8 we have (T) < o11) holds for any E (T] \ (1). Since we can confirm that the inequality + 2coh(( log 7) <
43

---

Benign Overfitting in Token Selection of Attention Mechanism
-2x- 2sinh(r log 7) 2 -(()) + holds for d(1) by rearranging terms, Lemma D.I and C(T) give
c²
(T) > (T)) > 2 + 2 cosh (Aga(T)) log 7)
(234)
(235)
9 (Aj,e(T')) +2
(236)
-9 (A,, (0)) + Cgam 1020m112210 ||21| 2 p1/2dd max & 02) +2 Op))
(237)
where the last line follows from the definition of T2 and g(),, = (1 + o(1)) This completes the proof.
Proof of C(r) E(r). The proof for the clean data is the same as the not-benign overfitting case in Lemma C. therefore, we omit the discussion. We first provide the analysis of the noisy data with proof by contradiction. Suppose that ((() does not hold, i.e., for any constant c 0, there exist k I /N and i, u E [7] {1} such that
exp (Ax,e(T) (T) 26 c.
(238)
From C(T), we have for any k, E N and L, u E [T] \ {1} that
g (A)) (Ax, (0)) - C6
9 (AI,a(T)) - g (0)
(239)
for 2 1. Additionally, we have (A... /g (Aa, < 2 by the same argument as in Equation 128 (T) < (1) and (T) < o(1) are implied by Lemma C.8 and the fact that A(T) is monotonically decreasing as a result
Also,
of C(T). By a simple calculation, we obtain the following for x < o(1) - logc: c:
T
sinh (I + log log T) exp(x)
exp(-r)
(240)
T
c
3
2sinh (x - log T) + exp(-x) + (c-3) 3 exp(#)
2T
(241)
c
c
T
3
2sinh (I - log - 2TTz
(242)
c
for e satisfying e 3/c. In the last line, we used exp(- -cr for x < a(1) log c. By using Equations 238 and 242
we have
g (AAR,(7)) - g (0), = 27x, + 2sinh (A),)(T) log T) - g (O),
(243)
1
2 2 (A(,(T)) log c) 2sinh (A),a(T) + log c log T)
2 (Ar,,(0)
(244)
> (2 - 27) Ar, (T) + 3 2sinh (T. - log 7) + 2
1
(-9 (At,
(245)
minfcs/co. 1/2} (g (A), (At,)
(246)
In the second inequality, please note that (T) (T) - log o(1) loge e < 0, sinh - log T' <
and -g(Ai,. (0)) (1 + o(1))T O. Thus, the last line follows from sufficiently large c 0 satisfying 3/c
min(cs/co. /c. 1/2). Since the left-hand side of Equation 246 is negative by C(T). Equation (246) leads to
g (T) 9 ()) max
C6 C5
C6
2
(247)
9 (A), -g (0)
C5
which contradicts Equation (239)
The remaining is to show exp(Ar. + AR,,(T)) < e7 and manner as the proof of Lemma [CIII for the not-overfitting case. Specifically, we divide the analysis into the initial learning phase (Case 1) and the subsequent phase (Case 2). In Case 1, the difference between (T)),)()) and (())..(T) is shown to be of constant order in a similar way. In Case 2, since they have the same order of time evolution with opposite signs by C(r), a similar discussion can be applied. This completes the proof.
< c7. This can be shown in the same
44

---

Benign Overfitting in Token Selection of Attention Mechanism
Proof of E(T) = (T). The ratios of (1 -11(1-spp -11))))))))) within the clean data and within the noisy data respectively bounded by constant, by the same argument in Equation (14T) for the not-overfiting case. We discuss the ratio between clean data and noisy data. From the definition of Ai we have
T
1+ Di= exp (-Ar,((T) 1+
1 +
Σ
exp (-Multt))
(248)
10 (F) -17)(1=8'())
(7))
Σiz exp (-A),(TT)
1=2
T
1
= 2 + Σ
exp
(-Ai,(T)) +
(249)
21 exp (-A),((T)
n=2
Thus, for any i and j E N, there exists a constant 2 such that
(T) - x4)
(7)
(250)
$1 (7) (T)
0 (T) - s1 10) (T)
2+ 1- exp (ia(r)
1 1 " exp (-) +
(251)
(1 En' ex=eep exp (-Age(t))) En= exp (-Ar,(7))
2 (T' (T- 1))(1)
$ 1 (T 1)o(1] + exp (Aiu(T) + Aj,((T)) < <t
(252)
for any u, E [2]] \ (1}. In the last line, we used that Ar,(T) > -0(1), Ar,(T) < o(1) and the result of E(T). Similarly,
we have
-1 (T) 1 (T)
1 5-
(253)
(T) (0) (T))
Ex= exp (- (+)
T
1
2 Σ (xp((AAi((T)
(254)
1+ Et= exp (-))(T))
El= exp (-Aji((T))
1=2
exp exp (-A,,))) (r) + (T- - 1)0(1)
< < 2++7" 72
c
(255)
(T 1)o(1)
which follows from (r,(T) -o(),, (T,(T) < o(1) and the parameter assumption T= (1)).
Finally, E(r) implies
St
(T)
(4)
exp
X
W(T)
exp (- (A))(T)) C7.
max (E
max
max (LE exp (-Ai,u(T))
(256)
(0 (7)
M.
(T
(i)
((T)
exp
X
for any i E [n). The dominance of (1) (r) -1 (r) argument as in Lemma C.II
(()())) for i E C, relative to all tokens t E [T] \ (1), is shown by the same
Proof of Argr (A(7') B(7)) A())) G(r+1).
The conditions of Lemma D.8 are satisfied from
A(T) A
A F(7). From Lemma D.8. we have
X41(T + (0)
Σ
@ ( max of)
(257)
r0=0
45

---

Benign Overfitting in Token Selection of Attention Mechanism
for any clean data E C. To derive the upper bound of the summation of softmax probability terms, Lemma D.5
Since Lemma D.5 follows from A('))((T") and the definition of Ti, we have
of (0
log (T aan" an.aaoooo
- 0.? (21l|/ pelleاlllddd max fo og))
max
on -o | 14 max oo og)
an an-to2 |/ | 2d max o/²]
(258)
log (Ti an -000000 1 22 max op1)
< max
(259)
on 2 1 1|2 12ll max m oo)
on 0o / 2 14 2d max (o og))
log p-
(260)
an
p oo? 1 22
max oz)}
Substituting this to Equation 257 and using the current SNR condition SNR² o(n 1), we have
(+ 141 (0) o(log p 1
(261)
Since (0) = o(1) from Lemma C.8 Equation 261 p lead to \41(7 + 1) S ollogp-1).
In the same way, have 1- 1) o(log p ). For noise memorization terms, from Lemma D.9 have for any i [n] that
1)) (0,1 -
10 (())(((()))))))))) can -2 max o2)
(262)
r'=0
We proceed with the same argument the previously discussed signal update case and obtain the upper bound [T (T+ pllogp-'), The bound of Pi.e is obtained as well.
C( F(F') A G(T') = A(T 1). The proof is basically the same the not-overfitting case, but we newly have to care about the behavior of noisy samples. It follows from Lemma C.5 that
Proof of Ar'ST (B(r') C(r') A
2a n
|I((T + 1)11 [|(())
(-M(()) y() 12((') 2 Σ V C(
(263)
0==
By definition, we have
T
T
1(()) = Σ
(7')
7t
(W(F)x>" ,P('))
(264)
t=1
u=1
) max (14x((((""") (ett)))))))) (p))())))
<
81
71
(1)
(7')7
(265)
For clean data i C, the argument similar Equation 152 2 which follows from F(7'), and combine it with G(r'),
Lemmas B. 12 and D.5
Then, we have
Σ
Σ
1 of
max
{ (1T [11 max (1A+1((')) (Pt))
(266)
TE(T]
r'=0
log 2 p-1
(267)
10? max (0z 02 o.]
46

---

Benign Overfitting in Token Selection of Attention Mechanism
For noisy data>> j E N, the softmax probability (7)(() of (+') is not dominant over the other tokens. Instead, we have
T
)
81
(
(7))
71
(1)
Σ
8
(r') (0)
T
T
(70)) (0
T 160-
(7)
(100)
Σ
2 (+)o)))
1 (F)s 18
(+)
(3)
Σ
5 (7)
Σ 0
(7)) (
+
+
88
2=2
k==
(268)
0
13
81
max
(269)
()
r'=0
log p
P
log p-
(270)
on-lo? o? max & og)
on l2d max (o o8]
on on- o? max (ol oz}
where the last line follows from Lemma D.5 we have
( ((llualalll)))) and the definition of T1. Therefore,
log P
Σ
G7)) on to- m w: o
(271)
+'=0
Now we have Equations (667) and (71)) and the remainder of the proof follows from the same discussion in the not-overfitting case.
Proof of Arrgr (A(T') A C(7') A G(T')) + 11). As with the proof of A(T + 1), the contribution of the noisy data can be evaluated in the same order as the clean data, similarly to Equation (271) The proof is essentially the same as the not-overfitting case, so we omit the proof.
Proof of Lemma C.12 At time step r = O, all propositions hold from the proof for the base case. For the next time step, A(T + 1), Bi + 1); C(T + 1) and G(T + 1) are proved based on the propositions up to time step T. As for D(T + 1), E(T + 1), and F 1). they are derived from C((T + 1). Thus, the proof is completed by induction.
Lemma C.13 (Benign Overfitting, Stage 2). Let Ti be the time step in Lemma .3)). on a good run, the following propositions hold for all time step 5 E [2]:
C.12 For any time step T2
exp(n -SSNR θ an o (allllllll all max(o?.o3))
A(T): There exists constant C1 > 1 such that - 1/C1) V,
B(T): There exists a constant C2 > 1 such that
- : ow |112 (
(1 1/C2) (,00))_.olw111221 11 (0th/CC))oolll/lpppp l0lldl2Vd,
1001/2)).0o..d ( |a FwO.d
(1-1/C2) owo
for all i E [n] and t E [T].
(W()),, Www.
(W()) <10 (nn0))
(W(F) (
(W " w(T)e < ( lo
for all i.j E [n] and t, u E [T] such that (i,t) # Uj,
47

---

Benign Overfitting in Token Selection of Attention Mechanism
C(T):: Let g be a monotonically increasing function g(x) 2x + 2sihh (x log To. Then, there exist constants C3, C4 O with cg < Ca such that
g (As 2 9 (0)) + r ezon -1 o 12 max m op) 9 (T), S 9 (A)) + - o ulll max ogh.
for any clean data E C. and t E [T] \ (1). Additionally, we have that for some constants cs < Co.
9 (T) log p )) 2 9 (0) log 1) T. p c5om - -02|||||allll| max op), 9 (T - -log))-oopppp p 1) g (F),(0)) - log p 1 + T p. cgam | max wi op)
for any noisy data j E N. and
9 (T), 2 g (T + r. p c5am coanmlom 022||11221|00||220 max (ou op), (T s 9 (O)) + p- caan -2, "121|||2||u||2d max o/)) E [T] \ (1.2). for any noisy data j E N and t
D(T):
0(1), 52 (T) 2 O(1),
for all clean data i E C and noisy data j E E.
E(T): There exists a constant > 1 such that
exp (A) - (T,((T) < e7, exp (T) - [T)(T)) < c7,
exp (Tx, (T) - [T)( < < l0, exp (T) - [(,(TT) < pOT:
exp (Ai,(T) - [k)((T)) < exp (T) - -Air(T)) < pCT,
exp (Ase(T) - (T)) < C7, exp (T - Aee(T)) < <.7,
for all i,j E C, C.k. t, u E [T] {1} and v, w E [T] {1,2}.
F(T):
There exists a constant og > 1 such that
0 (T) 51
(T) (8) (T)
max (JEÍM)
81 (T)(1 81 (T)
< <0s,
max. KIEN
" (T) - 82
(3) (T))
<(8,
82
(7) -1 (T))
82 (T) 89 (4) (T))
max (R) (T) (9-1) (1)) (x) REC,REN
< pos LEC.KEN
max
s ((6))(15((()) <
<p- lcs
Additionally, there exist constants co C10 > 1 such that
(T)
5 (T) < C9.
(f) < peg.
(P)
()
max
(T)
< CD.
max v.
( (T)
(E(T]M(2 6.6 (T)
max
( (T)
max
< <p co.
(7) 8 (1)
8. ..'( - 80 (7)) <10
81 ( (1)
< C10:
$2 (1)( - (T))
for any clean data E C, noisy data j EN, E [T] \ (1), and E [T] \ (2).
G(T):
[4+1())= - 0(1,, [A-1(F) 0(1), (T) = ('n'R'SNR" ?),
for any i E [n] and e [T].
We will prove the above propositions by induction argument for r E [Ti, T2).
48

---

Benign Overfitting in Token Selection of Attention Mechanism
Base case. The propositions A(Ti) and B(T)) hold from the results in Lemma C.12
We show that (TT)) holds. It is obvious from Lemma C.12 that the incqualities for clean data i E C hold. For noisy data, since T1 is given by θ an-1o3]pp 2ll oa it is sufficient to show g(T
p0g = (1)) and = (1) for E [7] (1,2). From the results in Lemma C.12 we have g (A,.2(Ti)) O(p-1). Thus, we
(T)) J(T
have (Ti) = (Zi) E (1 + o(1) log p leading to 5( (7i) log = @(1). Additionally, from [T) in
Lemma C.12 we have
(Ti
9 X2 x )
W(T)) p(T)) (PT)) (Ti) (Ti) < loge7, cr,
(272)
where e7 is a constant in Lemma C.12 Thus, we have ((A),(TTi) = @(1) which is the desired result.
Regarding G(Ti), the result Lemma C.12 and the current SNR condition n- --SNNR-2 2 = (1) conclude the base case.
Proof of Ar'gr (A(r') A B(7') C(r)) A D(7') A F(T') = C(T + 1). We only show the case of y 1, i... E C. U.W. The same argument can be applied to the case of C_ U.4. The conditions of Lemmas D.10 and D.11 satisfied from A(7) B(T) A D(r) A F(r). For any clean data E C, since we have the same bounds for signal and noise updates as in Lemma C.12. the same argument leads to
A + 1) A 2 (7) - of) 4am 0.
maxx(oa o2h
(273)
Ax + A.,(7) $1 (7) - s (T) Gam | ||||||||112dd max o?h
(274)
for any t E (T) \ (1}. Applying the same integral analysis, we obtain the desired inequalities for ((A,,
For noisy data j E __ from these lemmas and the definition of Tf Tj we have that for any 1 E I(t and E C++
(T +1) 1) - (,)(()) =
-11-1111(1)) + (p52(T) - (z2(TT) - 2 (-) - (al))) Gallullllll max (=..0)) o2)
+ -1 -1 pga(T))
(275)
(276)
+ + - 3/ (7) pan | 1l2llllll22 max o/])
+ (T)) (T) dpon - o | 2f? max of)
(277)
2 82 (T) -s (T)) cppam 22 maxfo oz).
(1)
82
(278)
In the last line, 8 ( (T) (1)))1(())())))))) (7)) S p82 (T) ( (T)), which is derived from F(T). and the current SNR condition
(
SNR = o(n-1) allow us to ignore this signal update. The same result holds for the relevant token E R = {1} and weakly relevant tokens & WU), as the signal update can be neglected by the same argument. In the same manner, the upper bound of (T + 1) - (T) is derived. Using the evaluations on (z) (7) - 8 (r)) in Lemma D.2 and applying the same integral analysis as As, which appeared in the proof of Lemma C.11 we conclude the desired results of T,.-
Proof of C(T) = D(T). For any clean data i E C, we have
(T) (7)
1
(0) (0) 2 0(1)
(279)
I'= T' 1+
ex (-A,,(77)) 1 - (xp(
where the inequality follows from C(T), which states that (T)) is monotonically increasing. Similarly, since (T,(()) for j E N is monotonically increasing, we have 82
(1) (7) (2) (Ti) 2 O(1).
Proof of C(r) + E(T). By the same argument as in Lemma C.11 and using the inequalities of C(T), we can show that exp (ALe(T) (T) is bounded by constant for any i,j E C and u E [T] \ (1). Similarly to clean data, it can be shown that exp (I (T)
(T)) is bounded by a constant for any k, E N and v, w E [T] \ {1,2) as well. At this time, since we have the result for (r) - -logp- instead of T (T), the second line in the statement follows by factoring
log out. Moreover, from C(r), the growth order of (A) and gTT) differs by p. and this ratio also appears in exp
following the same proof as Lemma C.11 The last line in the statement follows similarly by factoring log p out for J
49

---

Benign Overfitting in Token Sclection of Attention Mechanism
Proof of ((r) = F. We first consider the ratio of ( (r) - (r) across all pairs i, E [n]. Using the same argument as Lemma C.II
we can show that this ratio bounded by constant among clean data and noisy data, respectively. For the ratio between clean data and noisy data, from E(T), we have
exp (T) 2x, = exp (A,,(T)) exp
l'x.(TT)) cT.
(280)
exp (xp,-Ax,(TT) A,,((T) exp (Tx, - Ar,((T) CT.
(281)
for any i E C, E (T) X (1), and & E. Additionally, from E(T), have that for any [T] \ (1,2),
(T, = (T) [x,1 < log(pe7) < 0, (r. < (T) + 2 (T, < <,
(282)
(283)
where the first line follows from the parameter assumption p < 1/C for a sufficiently large constant C, and the second line by (T) > Thus, we have l'..= 0 for any E [7] \ {1}. Combining this with Equations (280) and (8 1 and
proceeding similarly to Equation 252 we obtain that there exists constant ej such that
(k) (1))1 (T))
2 + Σ 2= exp (xp(-Aa (-)),(TT)
1
(284)
(T) - s (T)) (T)
1 1 Σ 1= exp ( (-Ax,(TT)
(1+ Ei exp exp(<Ax.(F)) Exp
2 (T -
1 1)0(1)
1 -T- - 1))(1) + exp (T) + Ax,2(T)) < <.
(285)
for any a E [T] {1}. Similarly, by the same argument as in Equation 255 we have that there exists a constant de such that
(1) - ( Eize exp (-
T
2+
Σ t=2
exp exp (xp((Ax)(T)) +
(286)
"1 (1 1 - exp (-Ac+(7)
Et Etzexp expl (-Aue(T))
+ (T (T-1))((1) <g
< + T
exp (-,, An.((T))
(287)
(T (T-1)00 1)o(1)
where the last line follows from Equation 283 and the parameter assumption T = @(1) Equations (285) and (287)
completes the proof. Similarly, the ratio of 2 (T) - 82 (k) (7) across E N is shown to be bounded by a constant.
(k) (T) 82 (1) (T)) for i E C and k E N. For any i E C and We now consider the ratio between of (T) - (8) (T) and 82
k E N, we have
Ere exp (-)),((T)
E exp (-I',,(TT)
(T)
82 (7) -/t
(T)) -1
(1+ exp exp (-A),((T))
(1 E EeenTM22 exp exp((11)))) (-Tx, 2'
(288)
Thus, there exists a constant d' > such that
1 (-)
Ele exp exp (AA)(T))
[1+exp(-Taa(7 + eT=ppp
exp exp (-Ix.c(T)
2
(8) -0 (7)) exp (-Ix,(TT) + Eis exp (-Ir,((T))
1+ ΣT= exp (-Ar,((T)
$2
(289)
max (exp (xu(()))) (1+ -L)))))))))) < pcc. 2E
(290)
where in the first inequality, we use the mediant inequality excluding (T) along with the bounds (T)) (x..)T > -01). The last inequality follows from ETr) and the parameter assumption T = O()).
Similarly, the desired result for the reciprocal is obtained by F(T))
Finally, from E(r), we have
(1) (T)
(1) exp
x
((T)
exp (-Ai,eTT) S C7, exp (-Ai,u(())
max (8) (7)
max (.)()
(291)
(T
(. exp
X
(3)
p(F)
50

---

Proof of Ar'TT Benign Overfitting in Token Selection of Attention Mechanism
for E C. Similarly, for noisy data j E N,, the ratio of 00 (r) across e E [T] {2) can be evaluated using the result of ((7). The last two inequalities are derived by combining the fact that Ai, and I 3, are monotonically increasing from (T(). and Lemma D.3
Proof of Ar'ST (A(T') A A
C(') A D(7)) A F(T))) G( + 1). The conditions of Lemma D.10 are satisfied
from A(T) A B(T) ((T) A F(r). Combining Lemmas D.8 and D.10 we have
(+ + (0) S
of
dallullll 12 max og)
(292)
r|=0
for any clean data i E C. By using Lemma D.5 which is derived from Ar'++(TT)). and the definition of T2, we have
(1)
(293)
log (r. an -omn- 21 maxfo o211)
S max
(294)
on an-lo? || 122111112d max (o 03]
on on | 2 | 22 max (ou- ob)
log (T2 am - o 2|| 1112|11112210 max fo og.)
max
an o || 11 max o. o}}
an an 0² 2 |1112|1114 || p1|2d max oB))
(295)
n SNR
(296)
|122||1122d max fo w. o2)
Substituting this to Equation 2922 we have
X+I(T - -
nn nSNRR O(1).
(297)
Since [0)) = o(1 from Lemma C.8 1- + 1) v 0(1).
Equation
297
leads to A41(T + 1) < 0(1). In the same way, we have
Next, we move on to the analysis of noise memorization terms. For clean data i C, from Lemmas D.9 and D.II similarly we have
(1) Pr, (0)
Σ 81
- - d'n
-oddddmno' -0 "22lee1aallul|d"" max o7) = Onn-ISnR SNR 2).
(298)
The same upper bound obtained for Pt. with E [T] [1}. For noisy data j E N, the upper bound for PJ,1 is derived in the same manner, based Lemmas D.9 and D.II For other tokens, have
(p + 1) -)))(00) < (T)) ())(()) 1p + 1)-1)))00 p3.2(Ti))
(299)
< Ollog p -) Σ
82 (T)(() =66 (T')) c' d'pam p'poam'o??lllll "022|ulaallllll
max of))
(300)
r'=T:
$ Olog p ) + log (Th pam -
maxto?. of)))
(301)
29 Ollog -) + log + n SNRR2
(302)
~SNRR²,,
(303)
which leads the same upper bound in clean data. The upper bounds for / [T] \ (1,2) are derived in the same manner.
AC(T)
+ It follows from Lemma that
|I[(T + 1)I1 - | = 2a
(((T)) v(. 1p((') + a²
IVp((T)1.
(304)
P'= i=1
7'=
51

---

Benign Overfitting in Token Sclection of Attention Mechanism
By definition, we have
17(6))
Σ
1
7t
(1)
Σ
(W(T)x)),P()
(305)
)|
88
(+')
78
Σ
(7')
| max (141(7)))1A (e.(())))
(306)
For the clean data i E C, we have the following upper bound through the same argument as in Equation 267
)
Σ r'=h t=
8t
1 (7')
78
(0)
Σ
00
S (F)( -f
max (E(M)
(7(
(307)
log (T2 . an to? 10 1211 max (ow. o,1) on an-to?21 | 1|2d max 03)
Y2
(308)
n SNN
(309)
on o max (ox oz)
For the noisy data j E N,, we evaluate separately for the token score of the relevant token 71
We have
Σ 1 (7')
(7))
(
Σ
(+')(
78
W=1
/
(f)) )
(7)
71
(1)
Σ
(F))
+ Σ
St (7)) )
0
(7') 6)
+Σ
+
St
( (r')
7t
(
(310)
t=2
1=2
(1) (1) + (T 1) 82
(5)
(T')
82
( (T)))
max ie (m()e
(176111)
(311)
max S (r)(1 - (r)) (ph1 (E
the
(312)
where the last two inequalities follow from the ratio of softmax terms in F(r'), together with the same manipulation as in Equation 152p. By summing up both sides and applying Lemmas B. 12 and D.5
we have
T
T
)
log (T2 p an S
|
2d max(o? of))
8t
(7') 7
(T')71)) -
Σ
1o0 am
on---o oo? 1/ 12 22 maxx(2? 03]
e||1||2/1a11a (313)
T'=
n -SSRR
(314)
an a0- max max o3}²
where the first line follows from max.e(n). the small order term log p in the numerator. By substituting Equations
(7 = O(ple|lel|a|)) from Lemma B.12 In the last line, we ignore
306 (309 and (314) into the first term of Equation 304 we obtain
2a
((-(())) y(. [p((')
n
S a an dd max o3}
n SNR
max (+)) (1-1)) (e,
(315)
n
7'
1=11
SNR-
(316)
no?d? max m: o?2
0² n n||ll|| max o o?}
(317)
52

---

Benign Overfitting in Token Selection of Attention Mechanism
which follows from Lemma assumptions in Section 3.5
C.9
and G(r'). We confirm that Equation
317 becomes o(o3d). From the parameter
have
of max fo og) 2 12d d 2 min
|lull
1 2 S o? (Tn/0)
n
(318)
ogd
C'ndd/4 (Tg/(Tn/0) C'nva
Z min
(319)
((1,
(320)
where the first inequality is by Assumption A8 and the second line follows from Assumption A2 Thus, Equation (317 becomes df
Finally, confirm that the second term in Equation (304) can be ignored. From the parameter assumption the step
size a, it can be shown in the same way Equation [66 in the proof of the not-overfitinng case. Therefore, combining
Equation (304) and the results in Lemma C.12 we have [P(T + 1)ll E (1 = 1/ for some constant C₁ > 1, which concludes the proof.
Proof of Ar's- (A(T') A C(T)) A F(7') A G(T)) = B(T 1). In the proof of the not-overfitting case, we have already shown that the second-orde term with respect to the step size in Lemma C.6 can be ignored under the small step size assumption. Therefore, we write it as O(a²) in the remainder of the proof. From Lemma C.6 we have
WW(r + 1) [W
=
(IW( + 112-1112 [W 13)
(321)
T'=T
2a
((')) y (+)
(322)
t=1
Recalling the definition of I,+ in Definition C.2 proof of A(T + 1), we have
and applying the same technique used in Equations (309 and 314 in the
T
Σ (+(7)== T/=T1
St
(T')
(
Σ
(
(T')
(323)
T'=Ti
SNR
SNR 1| ao?d max f fo? o2
(1 (1)
(324)
an dl max pn of)
Since have (7') (1) from G(7'), we have
2a
n
SNR
2 1|111
(((()) y() 1. (+)(4))+1(
(325)
dd max max(o o3]
Y
d max oß] olo 1112d1)
(326)
where the last inequality follows from the parameter assumptions. Combining this result with the base case r 0, we have [[(r + 1) 1/C2)a [ for some constant C2 1, which completes the proof. The same result holds for [[ + 1)pa l12. Similarly, it follows from Lemma C.6 and G(r') that
W[( + 1)el [WT))e
2a
(((7)) Y())
(327)
n
1SNR -2 o
an od max ( 02 0,2 nn 'S-R 2+ O(a2)
n
(328)
old
+ +(a2?) = o
(329)
n? n |1 11 max o7)
53

---

Benign Overfitting in Token Sclection of Attention Mechanism
which follows from exactly the same analysis as in Equation (317
Additionally, we will show the case of the inner products as well. From Lemma C.6 and the above discussion, we have
(W(r + 1) W(T + 1)) - (W(T))u4 W(Ti)m=))
((7)) y()
(T)A_ 0(a2)
(330)
SNR 1
0(1) 0(a2)
(331)
on o2 de max (o o
V4
+ - 0l
(332)
maxx oo o3)
where the last line follows from the parameter assumptions; specifically, we have
w m (ou 0,) || ||g13/d//ogg(////)))) log(Tn/0) 2 min Vd. 1
(333)
o² VV (Tn/))
vd
C2
2 min
(334)
logg(TT/) (Tn/8) log(I'n/0)
= -11)
(335)
Similarly, we have
(W(r + 1)p W(T W(T 1)c(") (W(()) ((T)))
a
((T)) y (Iis, + I,+ (+):)))))) +
(336)
1=1
n -SNN n1 - ond
n - SNN 2, | n -'SNRR2 + +0(a²)
a.
0(1)+ +
(337)
am l0.d max 2 o2)
on lo?d max f0 [, o2}
V2
+ + 0 a. (11
(338)
max fo o2)
1 max (o
which follows from the parameter assumptions; specifically, we have
d maxto of) o. "nn|llli 11) Z min
|L||
n log (Tn/5)
(339)
a.dd
Oe
n2||a|122
min
(340)
o. log (Tn/8)
-(1).
(341)
The same result holds for (W(T + 1) W(r 11 (8)) Finally, we have
[W (1) W(r+1)c!") (k)) - (W(T)) ((
Ti (T')) ( (II... (T) + (T)) +
+
(342)
n SNNR -2
022
a an o2d2 max ( m: o2)
n -SNR -2 O
(343)
Ad
2/ 2 max
+ +(a2) = 2 log(Tn/))).
(344)
02 m 22
54

---

Benign Overfitting in Token Selection of Attention Mechanism
which follows from the parameter assumptions; specifically, have
n²
d max o..) o. n=3nn² "nn"lullVVVVol((TT 2 min
(345)
oevdd
of
(Tn/0)
C²n2di/4
2 min
C log(Tn/0))
(346)
log(Tn/o))
= -((1,,
(347)
which completes the proof.
Proof of Lemma C.13 At time step 7 T1, A(T)), B(T1). C(T1), and G(T;) hold from the proof for the base case. As for [,), (T)), and (r)). they are derived from C(r). For the next time step, 11, B(T + 1), C(r 1), and G(T are proved based on the propositions up to time step Thus, the proof is completed by induction.
C.3.2. PROOF OF BENIGN OVERFITTING CASE IN THEOREM 4.1
In this section, provide proof of the benign overfitting case in the main theorem. We divide the proof into parts: behavior on the training data and generalization performance.
Training data. The conditions of Lemmas C.12 and C.13 are satisfied with the current SNR condition SNR = o(n
The proposition C(T) in Lemma C.13 states that Tg = an an Mla pll2dd max
exp(n SSNR=))
satisfies
((Aca 2 ((A,,(00)) + + (exp (n 3)) w(1)
(348)
for any clean data E C and / E [7] \ (1). Here, we used g (0)) -(1 + o(1)) (1)) from Lemma C.8 Recall that g(x) = 2x 2 sinh(x log T), and there exists a constant of such that of exp(x) > g(x). Therefore, for all E C, we have
T cxp (-A,, > - (T 1)4 o(1) >
(T2)
>
(349)
1 D exp
(-As,e(T)) (
for sufficiently small constant > O. By using Equation 349 and Lemma B.12 we have for any clean data E C that
y0 (x(1) = (()) v g
w (T2) (TTT)TT(TT))
(350)
=
71
+
y(i)
7t (Tz)
(351)
t=2
0 12
- 0 (111112211111))
(352)
00.
(353)
For noisy data, again from C(T) in Lemma
C.13 we have
g(T),(T))) > g(T, (T) log (0) - log pogp~~) + θ (pexp (n-'RR 2)) = w(p-))
(354)
9(T (TT))(T)) 2 ((T (0)) + θ (pexp (n -n-ISNR-- ?)) = =w(p-
(355)
for E [T) \ (1,2), which follows from -SNR (1) and Remark 4.8
Similarly, we have
1
8y (Tb) I+ Σ, (7]](2 exp (T2)
>1
exp (-I,, > 1 - (T 1)g o(p) > 1ppe. (356)
(2)
55

---

Benign Overfitting in Token Selection of Attention Mechanism
Thus, have
yūi, (r,((()))= _ LI
W(T2) (T(TT))P(l (P(Ta)
(357)
=Y 77 (1) (T2) + Σ yūl 7/ 2 (T2)
(358)
(2)
2 (e(iellaal)))) (1 pe) O 0 (|||||||||11)) pe
(359)
>0
(360)
Equations 353 and (300) hold deterministically on a good run. Therefore, at time step T = T2, we have that with probability at least 1 6,
E E c, fx(x0) Y V.EN, fx(x = YO).
(361)
Generalization. Let (X, p be the unseen data which investigate generalization performance. We first evaluate the attention values of signal vectors at time step T2. From Lemmas D.5 D.8 and D.10 we have
Th-1
x41(T2) (72)))))2(0)>
Σ
8 (T)
(1- carl 2 a
[d maxto
(362)
2 log ((T2 1) an 02 o
ma nSNRR*
(363)
N1, 1
(364)
which follows from the definition of T2. From Lemma C.8 we have (0) o(1), which implies (T2) 2 1. The same result holds for (T2)..
Next, we show that the attention scores of the noise vectors {elrrmy in the unseen data, i.e., o ( (T2), become sufficiently small on good run. While it is natural to evaluate (TW( p(T2)ll2 and use concentration inequality, it is challenging to track the evolution [[T) 122) Therefore, following induction proof in Lemmas C. 12 and C.13 we apply a concentration inequality at time step 0 and show that the result does not change at time step Tz. Let us define & as the event that the following inequalities are satisfied:
E [7], (1 - -(1) | |6.2 < (1)) We [7] Vk [=1), ((E,, < (ooe:lla))l
(365)
(366)
WE [n], Vt, u E [T], (c,, <=
(367)
14 (T). AA (+1), (W(() (co))
(368)
vi [m]] vt, [T], (W(()), w(O)e(")) (o0e))"))ll
(369)
14 E [7], (W (()) < <ww.://llllooloooo//////.../ Vt [T. (v, ei) c20.ll//l///og((1M
(370)
(371)
where the constants C1, 0g are the same ones appeared in Lemma B.1 Applying union bound on the modified versions of
Lemmas B.6 B.8 and B.9 the probability of the occurrence of can be evaluated. Since there is no need to apply the union bound over the additional n training data points, the outlier probability can be reduced by 1/n compared the original lemma. Therefore, we have
Pr Pr
8/n 8.
(372)
In the following, using the results of Lemmas C.12 and C.13 we will prove that the next proposition holds for all 7 E [0, T2] under the condition E:
56

---

Benign Overfitting in Token Selection of Attention Mechanism
H()):
(Wwoeeess (W()),, W(T) ( (1laall (Tn/))) (W(T)) W(T) <0 (o (0000l0100000000000000000000 (Tn/0))))
(((()), ((
σ o. 2 Lo
for all E [n] and 6, u,u. [2].
Proof of H (T). We proceed with the proof by induction. The base case holds from the condition E. In the following, suppose that (T') holds for any r' [0, 7]. By calculation similar to that in the proof of Lemma C.4 have
(W + 1),,))((TTT))).... + 1)) - (W(O)),,
((W(7' + 11E,, p(T'= + 1)) ((())),,((7))
(373)
(())) y(0 Σ c (r') (1(0) -0-'0(')))
(wW()),, w(F'x(") + (I(((")) x(1))
/=0ml
1=
+ Σ a W (()))))(()) 7'=0
(374)
We bound the first term using the induction hypothesis, the condition E, and A(T') Lemmas C.12 and C.13 We first evaluate the softmax probability part. For Ti, using a similar argument Equations (267) and (270). we have that for any E [n]]
(T') 100 - Σ
(+')))
log p-
(375)
on o22 max max o2)
t=1
For 2 Tí, follows from Lemma C.13 and the same technique used in Equations
309 and 314) that
SNR 2
(r'
(376)
on o max fo a3}
for any i [n]. Substituting Equations 335 and (376) to Equation 374.. we have
- (((T)) y(o.
Σ
(7')
Σ 1 (r') ) (WW())e, W(T')x()) ((()))))((r)))))) + II(())I((en
T'=0 1=1
n SNN
- a d [1lll12, o. vd lng()n/0)
11
oo2 d max 10 oo}
SNN
od o. max (
(377)
n dd maxfo2 mi o27
SNN
max (o m
(378)
n odd max 20 oo}
max || o. Vd 2 a. Loog(Tn/// =00ooppp.dlo(((//))))
(379)
where the last line follows from the parameter assumptions, using the same argument as in Equation 208 for the not- overfitting case. For the quadratic term in Equation (374), we can show that this term is small enough to ignore, similarly to Equation 216). Therefore, substituting Equation (399 to Equation 374) leads to
(W(T + 1) p(T + 1)) (W()),,
(o
(380)
57

---

Benign Overfitting in Token Selection of Attention Mechanism
The proof for the three inequalities below in the proposition (r) is omitted because they can be shown in the same manner [T 1) in Lemmas C.12 and C.13 under the induction hypothesis and the results of Lemmas C.12 and C.13 Since
[T +1 holds under the condition (T) is valid, from induction argument, the proposition (T) holds for E 0, T2].
From (T2) and Assumption A8
we have
(W(T)e,,pTT)) & ouonoddoog((/// -
(381)
log(Tn/0)
for some constant e2 > Using Equation (364) and taking sufficiently large T2, we have (72 > 20g and 1 > 2dz. Thus, have that for (
(x -1) " ((1)) & -11 - p) (2)) + (E: -1) ((T)) p(zb)
(382)
2ch
s - - 111111111
1/C) 202- +
(383)
log(Tn/8)
< <
(384)
where the second inequality holds by p < 1/C. Then, the softmax probability of the relevant token is lower-bounded as:
$1 T2
> 1 -
(385)
1 21 exp (xx+ - x1) T
for sufficiently small
0. Consequently, we have
Y (X) = Y (T2) Σ Y (T,
(386)
t=2
ス + (12l21ll1)))))
(1 c) - -0 (elle|la||a|1)) e 0.
(387)
Under the conditioning on E, the output frg (X) deterministically takes the same sign the true label y*, Thus, the generalization error is bounded as:
Pr (XP-)>>- (Xgn YY
Pr (XP-+)p-
(sign(fr; (X)) + Y
[8+
Pr (X,
[39 <
(388)
where we used Pr(A)= & Br) + Pr(B) and the result of Equation 372 This concludes the proof.
D. Technical Calculations
In this section, we provide the small lemmas that are necessary for the proof of the main theorem. The lemmas concerning the softmax probabilities are given in Appendix D.I In Appendices D.2 and D.3 provide lemmas for the attention updates in the not-overfitting case and the benign overfitting case, respectively.
D.I. Softmax Probability
In the analysis of attention dynamics, the values of s1 (1))()))) and (2(r) s2(T)) play significant role. We have the following lemma on the evaluation of this value.
Lemma D.1 (Bounds for relevant token probability). Fix arbitrary i [ni), and suppose that there exists a constant d' > 1 such that for all u E (T) (1), we have exp (Ni, Ai,u) < d. Then, there exists constant such that for any t BE T (1}.
c-
(7) < 2 + cosh (A,,(T) - log T)
2 + 2 cosh (Ar,(T) - log T)
58

---

Benign Overfitting in Token Sclection of Attention Mechanism
Proof of Lemma
D.
Using notations introduced in Definition
we have
(PT)
Σ exp
( x
W(T) p(T)
exp
(1) - (T)
(389)
E 1- exp x
W(T)
Σ exp
X
(1)
W(T)
El= expl (-Ar,(TT)
(390)
1+ Expp
1+ Ei= 1= exp( (A,())
Using the condition of the lemma, we have
d -I exp())xxp(AA1111
of (1) - 1+c (1)--((- 1)exp(--Aa(t))
1 1+--11(T 1) (,,((),,,(TT))
T
1
T Texpl ((,(--),((T))
(391)
T-1
(CT)/(T - 1) T Texpl- exp(-A,,(77)
Texp (-Ar,(TT)
<
(392)
+ 4 (-)))
= 2 2 cosh (A,, To
(393)
where the constant CT/((T - 1) is replaced with e > This gives the desired result. The lower bound shown similar way.
Next, we show a similar lemma that is used in the proof of the benign overfitting case.
Lemma D.2 (Bounds for confusing token probability). Fix arbitrary j E .N, and suppose that there exists constant df >> 1 such that for all E [2] 21, we have
exp (T) 1 (T)( <,, exp (T) (T,(T)) < c exp -C (T)) pd.
Then, there exists a constant > such that for any 1 [T) 21,
(1) (T) $2 (T)) < 2 + 2 cosh (T - log 7')
c
+ 2 cosh (T) log
c-
(T) 52
( (T) <
2 2 cosh - l"
92
2 + cosh (T),1(T) log p log 7"
Proof of Lemma D.2
From the definition of J,, in Definition
we have
Enery exp(-1(,(())
(T) I+ (2) exp(-I2a(7)) I+ >xETT exp((I';,c())
1
(394)
For E (7] \ (1,2), since 1=-ppepp--T.((11 S exp(-T,,177) S cpexp(-T;3<7)) and d' exp(-Ig..()) S exp(-Ir,,u(T)) S. c' expp(-I',e(T)) for any u E [2] (1, 2). using the same discussion as in Lemma D.I and the parameter assumption p 1/C, we have the desired result. For the relevant token, since c' pー (xx((-r,,(7) exp(-I,,(TT) d'p-1
for any / E [T] (1,2), similarly to Equation (393), we have
82 (T) 82 (T)) <
(395)
2 cosh (T))(T) log log T)
The lower bound is derived in the same way, which completes the proof.
In the next lemma, we confirm that the significant term in tracking the gradient descent dynamics, s(1 - 8), is dominated by the token with the highest assigned probability.
59

---

Benign Overfitting in Token Selection of Attention Mechanism
Lemma D.3 (Inequality for s(1 8)). Let R² be probability vector, and let E [T] be such that St max Se-
Then, we have
Su (1 (1-8u) Su) S 8t), Vu [T] (t).
Proof of Lemma D.3 The function f())
[0,1] is monotonically increasing in 00, 1/2]. and f(x) = f(1 holds by the symmetry at 1/2. When 0 < S 1/2, the claim holds from the monotonicity over
x,2 defined
[0, 1/2]. For the remaining case 1/2 v S 1, since f(st) = f(1 - 8t) and
Su S
Σ
Sv = 1 5t
(396)
oE (E(T]R())
hold, the claim follows from the monotonicity over [0, 1/2] again.
In the next lemma, we will see that with the assumption of a sufficiently small step size, the attention values do change significantly in single step of gradient descent.
Lemma D.4 (One-step update of attention). Suppose that the assumptions in Theorem 4.1 are satisfied. Then, under the condition A(r) and BT)). which appear in Lemmas C.I to C.13 we have
D1+1(T + 1))+)++(+)) =+1(())==(())) o(1), |A-(T 1 )-1(r)=== o(1), lp,,(T + -=,
for any i [n] and [T].
Proof of Lemma D.4 From Lemmas C.4 and C.9
we have
A4 + 1) (+1(7)) & max (7 (r), (r) + t N..TII/VW-- V (w-V)r)V
(397)
For the first term, from Definition C.2
we have
max (T. (7), (I
(WVr)))).W))))))) [Ip(T)]13 ke (E
(x( (
Sa- max (E(T)
max (EE))).E(T]]
max
(398)
I <alll|la wll2 max 1ul[2d
(399)
=1],
(400)
where the third inequality follows from the propositions A(T) and B(r) in the conditions. The last line is by |||||2
((Ilul)2) and the parameter assumptions in Section 3.5 specifically, a S max[]] o.d)-1/C and max(o o
θ max (ld o.d] - log 2(Tn//)).
For the quadratic term, we have from Equations and 12) that
a² 2² V (wTL(r)Vpp(t
S a²² ( max
(ha
11)")))²
1 . |Ip(T)]l2 max (e(m)
{ (TT())("'l)
1b
(401)
(iEjn),tE(7)
sE(n),tE[T])
Ilall2 op aw max(llel]avddd S a² ا
(402)
= (1),
(403)
which similarly follows from parameter assumptions. Substituting Equations 400 and (403) to Equation (37 leads to the desired result. The other inequalities for A and Pi. are shown in similar discussion.
Finally, we provide the lemma regarding the bounds for the summation of probability terms s(7) - a(7)). This result necessary to evaluate the attention updates.
Lemma D.5 (Bounds for summation of probability terms).
60

---

Benign Overfitting in Token Selection of Attention Mechanism
(Not Overfitting Case) Let T2 be the time step defined in Lemma C.H. and consider T S Ti. Suppose that the propositions C(r') and ((T') in Lemma C.I hold for any 7' [0, 7). Then, have that for any E [n].
Σ (()))
(Benign Overfitting Case) Let T2 be the time step defined in Lemma and consider T T2. Suppose that the propositions ((T) and E(r') in Lemmas C.12 and C.13 hold for any T 0, T. Then, we have that for any clean data C,
(T))
if = 0
om -o- nيد max
(-)(1 - = T'=0
Σ
log log(r-an r.an ا lldd maxe m (()))
if = Ω on 1 dd max .z}
@ellllllll 2d max o3) 03)
For any noisy data E N, we have
log (T
annn"" an 0.2 '022| max og
Σ - 7==0
ifr s Ti,
am am-to?1 || |||||||||| 2| oax o2)
and
(log (T p. an aom-to? | p || max (o ( p. a. -,? 2d max 0,21
Σ T'
82 0 -1-)(-8800 (9))
ifr 2 Tı,
where Tj is the time step defined in Lemma C.12
Proof of Lemma D.5 We proceed with the proof separately for the not-overfitting and benign overfitting cases.
Fix single t from (T) [1). Since (T,((7) is monotonically increasing and (A,,(()) o(1) by C(r') and Lemma C.& it follows from Lemma D.1 and T = (11) that s (0) si (7)) = O(1) for any r' [0, 7].
Not overfitting case:
This immediately leads to the conclusion.
Benign overfitting case: We first discuss the case of clean data i C. Fix t E [T] {1}, and let To be first-time step where (T,(T) exceeds logT. From C(T) and the parameter assumption T = O(1), we have To
lnn--agllellal #lladd max .rt)/ ). We first consider the case of < To. In this case, the conclusion follows from the same argument as in the overoverii case. This leads to the first line in the statement.
Next, we consider the case of To and first derive the lower bound. From Lemma D.I have
c²
c-1
1 O0 Σ 7==0
(404)
2 + 2 cosh (Ai,(T') - log
2 cosh (Ai,((T') log T)
===0
7' r==0
for any i E [n]. Combining the inequality 2 + 2 cosh (x - log T) < 2r + 2 sinh (x - log for x > log T and the
proposition C(7) in Lemmas C.12 and C.13 we have
1
1
(405)
2 + 2 cosh (Ar,(7') - log 2Ar,(7') + sinh (Ar,(T')) - logg T")
(406)
9 (Ar,((T)))
1
(407)
9 (0), + 7' cian 1oan-lo211 l 1221111/1 pel 22 max max oā]
61

---

Benign Overfitting in Token Selection of Attention Mechanism
Since the function 1/(b ax)) for is convex and monotonically decreasing for b + az > 0, we have
1 2 + cosh (A1,(TT)) - l)
7+1
dr'
(408)
9 (A,,((0) + caa to2 v 11 p 12 max fo oz)
+ +T 9 (O), + (T
caom caam 2 max fo oz)
o3 log
(409)
4am-1 | 2d2 max m fo
9 (Ar, +To c4am 1on 2
2l m o7)
log (T an 'o-o0'? o. e l ll 22 max fo o2))
(410)
an-lo? 2 | 20 max max o2)
where the last line follows from the fact that the denominator inside the log is ( This is ensured by the definition of To. which guarantees that the denominator exceeds g(oog T) = 2 log T = ((1), and by the fact that the one-step update of attention score is (1)), as shown Lemma D.4 Consequently, by substituting this result into Equation (404), we obtain the desired lower bound.
Similarly, provide upper bound using the function g and the result of C(r). From Lemma D.1 have
Σ s (
<
Σ
s
10
(1)
Σ
(411)
2 + 2 cosh (Av,(T') loo T*)
r'=0
r' TT=
To +
(2logT+1)) I (As,((T))
(412)
4
+o+
To
d
(413)
4
con--ooa +'= 9 (Ai, (0)) -nnm lo? 2112 |alllll 2d max (o of)
where the second inequality comes from s) < 1/4, sinh(r) < cosh(r), and the inequality 2 log T cosh(x 7) > for all v, which can be verified through a straightforward evaluation of the minimum. In the last line, we used the result of C(r), and the coefficient replaced with e > from the parameter assumption T @(1). Similarly to the lower bound, we have
7/ To+1 (At,((0)) + + cgan -gan-1o?l1 [ 112221111 (2 12d max oo? of}
dr'
(414)
13an-ooo lo? 1Ul2 2 max max oz] To 9 (Ai,((0)) + +
g (0), + - caan- 1 max fo o2)
Gan caa- 10000o) zd max o2} log
(415)
(0), + +To c-am / 21l1 max oo o2]
am - to? w1l2l|m1|ed max fo o21) log (r am
(416)
oo |ll |M de max o3)
Combining this result with To (11/mm -00 [ max 2o2. of)) from C(r), Equation (413) completes the proof.
For noisy data j E A, we first derive the lower bound for 7 < Ti. From Lemma and the proposition C(r) in Lemma C.12 we have (7) < o(1) for any E (T). Since we can confirm that the inequality 2+2 cosh(x--oog Tog <
62

---

Benign Overfitting in Token Selection of Attention Mechanism
2x 2sinh(x 7) +2 -3)) 2 holds for < (11) by rearranging terms, Lemma D.1 and Lemma C.12 gives us
Σ 0 (f)
(417)
2 + cosh (Az,((T') log T)
c-1
(418)
-9 (A),)(TT)) +
(419)
9 (0)) + coon looamooo?? [ | 2d max m 02) +
Applying the same discussion as in Equation (410), we have
Σ s1 (( +'=0
log (T- an - 12 max oz))
(420)
an an-1o??1 v | 2dL max (o 0%)
As for the upper bound, since we obtain the inequality 2 + 2 cosh(r - log > -I sinh(æ log = - for x < o(1) by rearranging terms, similarly we have
(1)())(1- Σ (A))(0)) + csan -o? -
2c
(421)
p4|2d max max o3]
Using a similar evaluation with integral, we have
log (T. om -- - max max o2,)) an on-10??1 Allly 2d max (o. o3}}
st (()
(422)
which completes the proof for S Ti.
Finally, for the case of T 2 Tí, the desired result follows by repeating the same arguments as in the clean data case, using C(T) in Lemma C.13 and Lemma D.2
D.2. Attention Updates for Not Overfitting Case (SNR² = w)n
Lemma D.6 (Signal updates in Lemma C.I1 Let Tí be the time step defined in Lemma C.11 and let 7 E [0, Ti}. Suppose that the conditions in Theorem 4. and A(T), B(T), D(T), and F(T) in Lemma C.1l are satisfied. Then, good run, there exists some constant
0 such thut
2 (1) ca /a max
1 81 (T) - - 3d max oz),
for any i [m]. Additionally, for some constant d' >> G, we have
141(7 + 1) (41(7) s
1 (()
op.
0 1= x + -7) - daallllllll oz),
for any i E [m].
Proof of Lemma D.
From Lemma
C.4 the update of X+1 given by
X+1(T+1), (+1(T)=
(7)) y(. (7) [Ip(7) a Vwr (423)
n
Using B(T), together with the parameter assumption || = w(a log(Tn/8)), we have that for E C. UN = E [m]
(+) = 1].
(1-o(1) lE (W())x1) W(T) S (1+) "ll
(424)
63

---

Benign Overfitting in Token Selection of Attention Mechanism
and other terms bounded as follows:
(W())x?' W(T) = O (poz (11
(425)
( @ W(r) = O max { (llull [lllldddd] o log(Tn//)) - O (po ll
(426)
(W())x) W(T) O Gelllllladlogg(/n// - (po
(427)
fort E W u E W and I which follows from the lower bound of the weak signal strength p. Additionally, for any E C+ U.W,, we see that $ (T) 2 (11) holds from D(r). Then, we have
T
T
(T) ) (WIT))
(T)
Σ
8t
(T)
Σ
w(T)m)) (WTDD..WWO..))
(428)
1=1
w=1
Σ
St
(7)
Σ
(T) (76 -
-10 (W())x() W(T)m4)
(429)
-f (1) (W(T)x) W(T))4)
(P()s) (W()))) Of W(F)A)
Σ (T)
(1t ( (0) (W()) W(T)m))
-10) +Σ (F) E (T -61 - (W())) W(T)
(430)
2 3 (7) (1)
0 lo
(T)( - 00
0 (po [cd
o(1) (000)
(0o || (2d)
(431)
2 2 s (0
1
(432)
(433)
for a constant of > 0. At the last term in the first inequality, we used (()) (211))11)11) (r) from
the condition D(r). The second inequality follows from Lemma 12 12 and p < 1/C for sufficiently large C. Applying
Lemma 1.12 provides the last equation. Using the same reasoning, we also have the following upper bound for C4 UA=
(T) < s e (1)
So |ل|a||
(434)
for a constant & > such that dg > ch. similar way, we will provide the evaluation for j C_ C_C_U. = {i E [n]
y+(1) -1]. The condition BT)) gives us that
(W(F)x(' W(T) = max [1lulB/v a. allllladd of log(Tn/6)) = =0
001mm mmxx { (p. Vd
log(Tn/6)
0 11 1el12dd
(435)
(W(F)x)) () W(T) = (po (1121111
(436)
(W())x)) 1 W(T) = O (max (pllml13 d, a. 11llladd ou log log(Tn/8)) = O (po 11111//d))
(437)
(W())x?) W(T) 220a:lull12doog((11 O (po 14 ((dd)
(438)
64

---

Benign Overfitting in Token Selection of Attention Mechanism
for E W() E WLy and e E IG), Then, we have
T
(+)
Σ
b
(T)
18
()
Σ
(T)) (W(+))p, w(r)m
(439)
==1
==1
<
86 (7) (T)
max
{>l - - (W())) W(T)
(440)
WE(T]
111
(T)) O [max { P4
log(Tn/6)
s T max 1t 1Ef2
(1) (T) - (T))
1 12 12
(441)
vd
- f (1 (T) (max {p, o(1) o 2122
(442)
where the first inequality is the result of the triangle inequality, and the second line follows from Lemma B.12 and the ovaluations just before. the last line, we used the dominance (T) o (T)), which follows from F(T). and the parameter assumptions 7 (1) and vd = w(lo((Tn/8))
Using these evaluations, Lemma C.9
Equations (29) and (30), have
[((()) y() 1 (T)
Σ
(-((()) y (T)
(443)
a. (2 - 3n)/4, (2-3n)/44 min ((-((T))) min { - -
iEC+
iEc. ((-((T))))
a ce min iEC+
max
{ (T) -1 c (
JEN-
aa. n1/2 ce
min {( (T-)-
-1 (f) max jEC uN,
(444) (max (p,) (al)
ieC,
2 2 (7) - 3/ (T)) Goo u 1 1|
(445)
for any E [n] and some constant &g > 0. The last inequality follows from the balance of the softmax probabilities over i [m). as stated in (T)). We also used 1 1/C and p < 1/C in the parameter assumptions. Similarly, have
(()) YO. (T) S (00 G0o [|l2llal))
(446)
n
for any i [n]] and some constant c4 such that c4 cz.
We now turn to the analysis of It, (T). Lemma B.1 states that, for i CN = fie Y = 1}, we have
(1 o(1)) |1
S (1+ |1
(447)
and
Kx (1) 1141) < + to = O (pllm112)
(448)
(xx(9) (411 eogellllllllog((T = O (p](4112)
(449)
(xx()) 14 < |2 (Tn/5) = (pllm(
(450)
((()) fort E +1+ u E ((()) and Z(0, Here, we used (#+1 1 0 in our data setup.
Then, for any i E C4 UN,, using calculation similar that for IW there exists constant c/ such that
T
T
(,)
Σ
(8) $t
(4 (7)
Σ
(x (xx" m 2 10 (7)( -s() of 11111121
(451)
65

---

Benign Overfitting in Token Selection of Attention Mechanism
Similarly, have the following upper bound:
5 (T) (1) (T) % 1
(452)
for some constant de > 0 such that dg > cz.
Additionally, using similar argument in Equation 442) we give upper bound for data with different label, i.e., for j j UN+ = {i E [n] y-() = =11. follows:
[3,+ (1) (T) (plvlla2)))
(453)
Note that the term log(Tn/6)/ does not appear this case, due to the orthogonality μ411 and p-11- Thus, combining these bounds and (1 1/C1)
o, o, & [l(()]l2 s (1 + 1/C1) o, A(T), we have
(-((() y(0 [IP(T)
Σ
(-((T)) y()) (T)
(454)
UEC,UN-U(C_UNG)
2 aa|p(T>)l (2 3n)/4
min ((-((T)))
min {st 2919 of) d|PlÍ2)
(+))
iEC+
seC+
-<lp(r))1l (31)/4 /4 o min
((-((T)))
max JEN- {st - (x)) (
KEC+
max jEC_UN, (T) - (+)) (111)122)))))
- ==][(())1l (1 / c min
(T)
(455)
sEC+
2 (1 - (T) | 11
(456)
for any i E [n]] and some constant & > We again used the results of F(r) and 17, 1/C in the parameter assumptions. Similarly, have the following upper bound:
n Σ
(-((())) Y, [I(()) ( (aa/)lllll
(457)
for any i E [n] and some constant d3 O such that de > cz. By substituting Equations 445 and 456 to Equation we have that for any i E [n].
423
A+1 (T + (+1(T)
2 (7)(1 min Ja Ille
o a²HT, w ((T)V)((T))
(458)
Finally, we show that the second term can be ignored. We have
a² F V WT [(7)Vp[(7)
V
(459)
aa² max
(st ( - (7)
- (()()))
max
(hr 1n/"
max [MT
H (2) [p(()]l2
max
(T)x,)
(1)
[2]
VE(N],T([T]
ie (E
sE|m],cE|m]
1E()),EE[T]
(460)
$ ( (0 0 3 1 1 op
ap dd an max (ax a.d)
(461)
(T) (() - a d u. of) (alle mal
(462)
= = (F) 5( (al|vil 1dd 1.))
(463)
where the first inequality is the result of the Cauchy-Sehin Schwarz inequality, and the second one follows from Equations and 12 Here, the probability term appears the bound from a similar argument Equation (442. In the third line, we
66

---

Benign Overfitting in Token Selection of Attention Mechanism
used Lemma B.12 the concentration inequalities in A(r), B(r), and the balance of softmax probabilities in F(7). Using |||| O(1 (2) and the learning rate assumption a O [11] (mmax|||all22d, add)=") the latter part of Equation (462) becomes (1)) Therefore, the quadratic term in Equation (458) can be absorbed in the first term.
Thus, there exists constant c >> such that
cool |2) max (o og): (T 1) - (+) 2 19 (T) (()))())((()))))) (7))
(464)
for any i E [n]. In the same way, from Equations
(446) and (457 there exists a constant d' such that
+ - -)) x+1(7) < s ( (T) of (T) c'a 1 alll max ogh,
(465)
for any i E [n]]. By applying the same argument to A-1, we obtain the desired result.
Next, we analyze noise memorization in the attention update. Many equations in the proof are similar those in Lemma D.6 so we provide the proof avoiding the redundant repetition.
Lemma D.7 (Noise updates in Lemma C. Let Ti be the time step defined in Lemma C.11 and let T E [0, Ti}. Suppose that the conditions in Theorem 4.1 and AT)), B(T), D(T), and F()) in Lemma C.1l satisfied. Then, good run, there exists some constant >> 0 such that
Pt,1(T 1) pr,(F)) 81 (T) - s() can - - maxfo
(T 1)),),((T))) -s
(0) (T)s (r) (T) can '02 eella max og
for any clean data E and [T] \ (1), and
- - con of max oz) PJA(T + -5, s (-)
PJa(T + (1 Pj,((T) 2 (7)s (7) -on of llllllllllll max for oph
for any noisy data j E a,d / E [T] \ (1). Additionally, for some constant è > c, we have
-0 (T) dan 22 maxfo" 12)) Pt.1(T + 1) - - s (T) -
Pt,(( + 1) - -t,((T) 2 -s (F) c'an - 0o? 211111221 4 12d max o?)
for any clean data E C and t E [7] \ (1). and
PJ.(T + (z,(()) 2 ()) d'n d'an-'o ofllelllllll max og)
11 -
(1) ( S s (T)s (T) -can -1 0² ||212|211111 lellllll max of]
for any noisy data j E N, and L [T] \ (1).
Proof of Lemma D.7 We first analyze the noise learning in the relevant token of clean data E C. From Lemma C.4 we
have
+ -1),at))
(((T)) (k) (1) (IP())1BAm)) + aee (O (wT((r)V (466)
n
A and A2 have d Cô (Ca, (ogT/))))(( 4/3 log" (Tn/8), from which it follows that
Combining Assumptions
14/3, d 2 14/,))4/8,,²
log (T = wn? log (Tn/6))
(467)
Additionally, from Assumption A1
we have
d w (o_ nnlpal122 (og(((/())))))
(468)
67

---

Benign Overfitting in Token Selection of Attention Mechanism
Therefore, applying Equations (67 and 4668
the results in (T)) have
[ 0
(469)
- o(n on 0220)) (W(()) ((T) 0 ( (Tn/6) = onn / -n
(470)
(W(T),-;, W(r) ) 0 21
(471)
(W(T))) W(r) (o o 0 12 log / = o(n - o
(472)
for any j,j [n] and i, u e (T] such that (i,t) # (1 u). In the case of k = i, using similar argument to Equation have
433
T
T 71- - - (7)) (W(+)x)) W(F)e())
IW (T) = 1W (T) =
Σ % (T) (T)
"1
(473)
t=1
uml
1
si (1) 71t) (o
2 2
(474)
s (r) (f o Go o L 112
1
(475)
for constant ei O. This follows from Lemma B.12 and the dominance of Equation (469). Using the same argument, we also have the following upper bound for some constant c2 > 0:
- (T) = - (T) - -70000000
(476)
In contrast, we have that for k & i,
T
T
kx,, =
W
Σ (7)
7t
(k)
- ->x 0h (T)
(W(+)(")) W(+))())
(477)
u=1
(1)
(T)
81
(k) (T))
max {
? (tt - (W())x (K) W(r)
(k)
(478)
nE[T]
t=1
maxxfss (st (7) - (2) (r)) (n -2=00000 | (1lal
(479)
1E[T]
(2) (T) - i1 (7)) oln - 2 ||ll||||l2d
(480)
where the first inequality the result of the triangle inequality, and the second one follows from Lemma B.12 12 and Equations 470 to 472 In the last line, we used the dominance of (k) (T)(1 (r)) from F(r) and the parameter assumption T (1). Thus, we have
(((T)) (M)) W (
a (-(((T) -8 (P -10011110110 ()) a 2|
72
n
(-(((T) of ( s (T) (n 20 |1llllaaa)
(481)
n
kti
(7) - (T) daam - 1 o² v d
(482)
for some constant &g > 0, which follows from the balance of loss derivative and softmax probabilities over training samples, as established in Lemma C.9 and F(r). We also have the upper bound for da as follows:
((()) (1) I. o9
(T)) 6aan - "2
(483)
68

---

Benign Overfitting in Token Sclection of Attention Mechanism
Similarly, evaluate the terms related (T). Lemma B.1 and Equations (467 and 468 lead
||= (4) ()),d)
(484)
(M+1: 1:: ) = 0fa.llallav/oo((m//00 = o(n onn
(485)
1M-1, 6 ) = Ofo:llllla
o(n o
(486)
(e( |("i = O(o'Valog(Tn/8)) = o(n -o,d))
(487)
for any i,j E [xx] and é, u E [2] such that (i, * (i, u). Similarly to Equations (475) and 476). we have that for k = i,
(F) = [1(1(T) 2 s1 (1 G
(488)
(T) (T) s s| (7) (())(1s))()) (()
(489)
for some constants c,, de > 0. In contrast, by the same argument as in Equation (480). we have that for ls + i
< ()) o(n -ta-to -
(490)
and similarly to Equations 482 and 483
we have
(())) y(b) TR. 2 (T) $ (()
dan" 10om-t0?11//
2 ||| pllad,
(491)
7
(T)) y(8) (R. S 2 (7) (((()) caom 1l 02111112|1101|2d
(492)
1=1
for some constants 4, co > 0. Therefore, using Equations 482) and 491 and [|p(T)]l2 ((a,vd) in A(T), we have
((T++) -)( 10 (T) -s (T)) min(cy, - "1lllllalll max ogt + aaee (1) V V (())V()))(())))
(493)
similar argument to that in Equation 4633 shows that it can be neglected under the small learning rate assumption. Thus, by appropriately redefining the constants, we obtain the desired lower bound for Ps,1. The upper bound is derived in the same way.
For the quadratic term in Equation 493
now turn to the noise memorization term for E [T] {1} of clean data E C. Since the proof is essentially the same as for pi,1 we show only the different parts avoid repetition. From Lemma C.4 we have
Pie(T +1 (I,(T)
(((()) ((). - + [1(())))) + +aze!" Vw: (494)
n
69

---

Benign Overfitting in Token Sclection of Attention Mechanism
When & = i, have
T
T
Iw (T) = 1:
(T) Σ
( (W())) 200 - Σ a (T)
(495)
1
=
(T) (T) 19 (W(T))) W(T)
T
Σ
(+)
10 (0 (W(()) @ W(T)
(7)8) (T)
( (109 -10) (((()x 0 (((T)
Σ (T) E
(T) (077)6600011 -1) (W(())( ((()))))
(496)
NE(TTV(L.))
1 (T) - - (7) 11
o(n on 222
+8 86 (mm)
(-10()))(()) -
+ (1) - - (()00)))))0000())10--00)))) on on
o(n + +0 (7)( (T)) 0( o2
(497)
< <7) (r) ?
(498)
for a constant df > 0. The first inequality follows from Equations 466 to 472 In the second last line, we used
(11) s (T) 21 12 -s (T) from the condition D(T). The last line follows from the comparison between (8) (r) and
1 - 81 (T). as shown in F(r), Lemma B.12 and the parameter assumption < 1/C for sufficiently large C.
Using the same argument, we also have the following lower bound for some constant c/o > 0:
K. 2 -s ( (r)s(
(499)
From similar argument, we have that for 1
TR,A((T) Te, (3) ( (T) Go || 12ll1111dd
(1) (7) 1
(500)
In,i((T) = l(se(T) 2 -51 (T))
(501)
for constants e/1: c/2 > 0. For k # i, (T)) and ((),,,(T)] are bounded as discussed in Equations (480) and (490). Thus, using ||p(T)
Vd) in A(7), we have
Pi.r + 1) - (T)
< < (7) (7) min(ch, Jan - ||lldd max of] + Vw-
(502)
We can show that the quadratic term can be ignored as in the case of (T)) Consequently, there exists constants c' > 0
such that
PT,(TT + -c.(T) S -1))())s (8) (7) can 12112|122111|12220 max(oe, o/) (4) (T) c'an d'an--oo - max oz). PT,(( + 1) pt.(T)) (r)s
(503)
(504)
70

---

Benign Overfitting in Token Sclection of Attention Mechanism
So far, have discussed the updates for clean data. For noisy data j E N, the update equations differ only in the sign due to the flipping of Y. Thus, we conclude the proof.
D.3. Attention Updates for Benign Overfitting Case (SNR² o(n
D.3.1. ANALYSIS FOR STAGE
We first present lemma on the signal update under the current SNR setting. Since most of the discussion overlaps with Lemma D.6 for the not-overfitingg case, particularly in basic concentration inequalities and equality evaluations, we only present the different parts. The main difference from Lemma Dielie lies in the behavior of noisy data, which shows monotonic decrease in 81 (T) rather than a monotonic increase as in the clean data.
Lemma D.8 (Signal updates in Lemma C.12 Let T, be the time step defined in Lemma C.12 and let T E [0, Ti]. Suppose and there exists some constant > 0 such that
that the conditions in Theorem 4.1 and A(T), B(T), D(T), and F(T) in Lemma C.12 are satisfied. Then, on good run,
X+1 +1) (1)++)))++(7)) 81 (T) (1) (T) colle//lll max ( op);
- 11(T+1)-A-1
00 (T) - s1 (T) call//2l 2 p ,²);
for any i E C. Additionally, for some constant > G, we have
A+1(T + - -+1 +
o -0 dalllllallllllddaaaa of)
for any i E C.
Proof of Lemmal D.8 The full proof is omitted because it follows from the reasoning similar to that in Lemma D.6 In the
analysis of Equation ) for clean data, the conditions A(T). B()), and D(T) yield the same results as in Equations (433
and (451). For noisy data, extra care is required because ( (r) > ((1) does not hold; however, noting that s(²) (7) > (p)
holds from D(7), the same evaluation in Equation (433) is obtained follows. Let j E EV,, and we have
T
(n ) (W(()))) W(T)
(T)
Σ
1 st (T)
1:
100 Σ -
(505)
=f) (T()) 70 (W())() w(T)))))
Σ
(T (T) (w()))) W(T)A))
(f) (7)
(1 9 - (0)) (W(()) (9, W(T)mL)
Σ (7) Σ (T) (77)))111 - ~0)) (W())x?), W(T)m))
(506)
"E(TMML.
2 ()) (T)
(
(T( (1)))) (2) 0 1
- -P())(1=(((7)))) (1+ (1) O (po u (2d)
- - -(000 O O (po 111d
(507)
2
(508)
where the fourth term of Equation (577) is applied the result of DT)), and since both the token score and the inner product term have a small order, this term is negligible compared to the first term. Therefore, the same evaluation is obtained as in the clean data case. The same argument is applied to T (7). Regarding data from the different class, i.e., j E [rx] such that
71

---

Benign Overfitting in Token Sclection of Attention Mechanism
y+ 1, the influence can be bounded similarly to Equations (442) and (552) Here, the influence of noisy data E NN+ be shown be in the same order as that of clean data, using a similar argument to Equation (69).. For example, the part corresponding Equation (442) becomes the following:
Σ 5t (T)
(1t"
(
Σ a (n) (W(rXP.W(rr)))))
(509)
81 - (T))
max
(1( O (max ( o ll)ldd
(510)
<E(T].11
10 (f) O mmax {p, o(1)) oz | l1
(511)
(1) from D(T), and that the balance of 2 (T) (+ (r)) holds over all examples & E [n]] as provided in F(r). Consequently, the desired result follows from the balance between the number of clean and noisy data.
for any i E C. Here, we used Lemma B.12 the fact that
p- p - (T)) (T)
(T)
The following lemma on noise memorization is basically the same as in the not-overfitting case, but the behavior of the noisy samples j E A differs. In partieular, (j,((T) for / E [7] \ 2) initially evolves in the same manner as (2,((7) but gradually decreases, changing sign around Ti. As a result, we obtain the following evaluation that spans both positive and negative values.
Lemma D.9 (Noise updates in Lemma C. 12... Let T2 be the time step defined in Lemma C.12 and let r E [0, T.]- Suppose that the conditions in Theorem 4.1 and A(T), B(r), D(T), and F(T) in Lemma C.12 are satisfied. Then, on a good run,
there exists some constant c, d > 0 such that d G, and
PG,(T + 1) (t,1(T) 2 a (7) - of (T)) con vam-'o' "2llllaalllll max (ux Pi,((T + 1) Pt,((T) s -s (r) ( (r)
can 02 maxfo onh Pt,1(T + 1) - (r,((T) S (T) - 31 (T) d'n - || max fou o?), (1) (T)
c'n - -o w l2d max 07) for any clean data E and t E [7] \\ (1). For any noisy data j E., we have
Pi,(( + 1) - (t,((T) 2 (T)s
P5.(T + 1) (T) S - - x1.) (7)) can 1 022|||||2||u|||a max oz), Pai(T +1) - Pp.((T) 2 (F) - ((0(())
c'an" -1 0 ا|||ا|||||ااا|| maxfo of).
and
P),2(T + 1) - -,,((T) 2 (T(()) (7) con - 0 max max (o on P)2(T - pp,((T) S 8 (T)
(1) (r) c'an
22 ام 12d max f0 o -0?2 -aan--oo00
and
Ps.(T +1) - - (7)s))) (r) can - | max oon og)) Pzr(T +1) - (s)(T)) (T) (T) d'om - 2021| max o7))
for any t E [T] (1,2).
Proof of Lemma D.9 For clean data E C, the same reasoning as Lemma D.7 is applied using the conditions of the lemma. Here, the influence of noisy samples in the summation can be treated in the same way as in the not-overfitingg case, by applying the same argument as in Equation (5IT) in the proof of Lemma C.12
For (T,) E NV, we can reduce the analysis to the case of clean data, as in Equation 5088. using s (()(()) O(p). In the rest of the proof, we analyze the update of (T.)(T) for E [T] (1). From Lemma C.4 we have
Pj.(( + 1) - -))1)((())
((()) (k) (T) (P( aa2e (1) (T) Vw (())) (512)
n
72

---

Benign Overfitting in Token Selection of Attention Mechanism
Without loss of generality, consider N4 {i E [m] y*) y = 13. Then, from the data model defined in Definition 3.1 and Lemma B.12 we have
70 (((211221)))))) 72 = -(pllellallal1))
(513)
v -(olle1122lll))))) h6 O (lggTT/o)))((T/0
(514)
for f E _ and u E It). When = j. using Equations
469 to 4722.
we have
T
T
(F))(T) = 1W
= Σ (T)
700) - Σ
su
(Thi)))
(W(T)))) (
(515)
==1
==1
= (T) 81 (7)
7 (W(F)xf w(T))
10
Σ (T) (W(())) ( W
( (T)
(T) (1 11) (W(())x W(T)(()
+ Σ
(T) Σ
(T) -10)) (W())x)) W(T)
(516)
wE/TT\/1.0)
2 - (()(1))((r))) (1e/222l1112 -200
- - o (plllllall n n 20000
+ +T) 56 (T) (0 0(11llallalll -
(517)
(T) ( 2t +7) (p|Pa|2lll|1) n o'oooo
8 (1)
Σ
s (T)(( --)(""
-1)) 0
(518)
uE(TiML.
- -1- (7) o (7)) o
n o0)
(519)
For 2, Equation (518) becomes positive from the evaluation of token scores, and combining this with Equation (517)
and (5) (r) 2 (p) leads to
W. IW. (7) 2 o (T)) (
(520)
The upper bound and the evaluation for (T) follow similarly. For t E [T] \ {1,2}. Equation 518 is bounded from below by -8 (7) 1111--spp0 o ( (7) - o (7)) (7)) O [ |||2|||||111 of o²d"d" Since (1) (7) 2 O(p) from D(T), we choose Ti such
that Equation 518 does not exceed Equation (517) too much. Specifically, we choose that we have
(T) 2 - (0) ( (T)s (T) ollllllllll
(521)
The upper bound is derived trivially. The discussion for W (T) and (T) for k F proceeded a similar way to kJ.
Equation [1I]p, and we conclude the proof by repeating the argument in Lemma D.7
D.3.2, ANALYSIS FOR STAGE 2
Next, we provide the results for Stage 2, i.e., 7 E Ty] in the proof of Lemma C.13 In Stage 2, the signal updates are dominated by the (r) g(0)))) for i C, as well as Stage 1 (Lemma D. D.10). The different part from Stage 1 is the behavior of noise memorization of noisy data, and the learning progresses in such way that the confusing weakly relevant (i) token, i.e., x² would be selected (Lemma D.11
73

---

Benign Overfitting in Token Selection of Attention Mechanism
Lemma D.10 (Signal updates in Lemma C.13). Let T2 and T2 be the time steps in Lemma C.I3 and let 7 E [ T2.
Suppose that the conditions in Theorem 4.1 and A(T), B(r), D(r), and (T) in Lemma C.13 are satisfied. Then, on good run, there exists some constant 0 such that
(+1 + - A+1(1) 19 (T) - (T) ( ca max of)
10 - (7) ca max op), s1
for any i E C. Additionally. for some constant è > G, we have
+ 111(T+1)- - (
op)) - (r( (0) - dalleilal (lllllllllllldddx of)
for any i E C.
Proof of Lemma D.10 We only provide the different part from Lemma D.6 Without loss of generality, consider the case of (T). ). In the summation over i E [n] in Equation (423). we divide the terms into three groups: clean data i E C++ noisy data j E W_, and examples from different class.
For clean data i C+, the conditions A(T). B(T). and D(r) give the same results as Equations (33 and 451 p. For noisy
data j E N_, the ratios of softmax probabilities in F(T) are significant. To derive the lower bound of signal updates, we use the results in F(T); specifically, for any clean data i E C, / E [T] (1), and [T] (2), ( (T)
(7) (T) x
-- (7) 81 (f)). o (1) ..22) (7)) S ( (7) - ( (7)). and the balance between 81 (T) 51 (7) and 2 (7) - (9)
(7)). Under the conditions A(T) and B(T), we have the same lower bound of E N_ as in Equation 433
as follows:
T
(T) Σ
1 (T)
11
( -E (Pn))) (W())x W(T)m)
(522)
(T)
11 (T( ( (W())() ((()m)))
=
Σ
(T)) (W(r)x) W(T)A)
Σ (T))) - 19)) (W()x W(T)L
Σ ( Σ ( - ~ (W(+)x?) W(T)m4)
(523)
(E(T)((L.)
2 (f -1 (
- -0001-s(001 s))) 0 (o
- - (())))0)((()))(0)))))))))
(0o 22)d
- (T- (0)1)) 0 (1) (0o 1 (d)
(524)
2 2f (T) (0( |
(525)
where the second last line follows from F(r) and p < 1/C in the parameter assumptions. The effects of data with different
74

---

Benign Overfitting in Token Selection of Attention Mechanism
labels & C_ UN+ bounded similarly Equation (422. From B(7), we have
(W(()x) (8) W(F) O (max [1 allllllllddd o logg(T/66)) o (max (o: log(Tn/0)
s a 1e
vd
(526)
(W(())x 88) W()) O (1 (e
(527)
(W(()x W(r) (max { { 13 a celllllldd o log(Tn/6))
(528)
(T) RR) W(T) = O oe (oglodlog((T//))))) - O (po (1)
(529)
for E (()) W
+11.0 -1 and I(a). Similarly to Equation 522.. we have
4.- (-)( (T) S 8/ (T) 51 0/th (T) 1m (k) (max ta o(1) o2ll13))))
W
(8) (T) 1 (7) (()))((ph))))
0phhi (0 O (max (p) (
11 (4) (T) si (1) (())))((+00)))) (1+0()) hi (po 1 (d)
(T- max
(T) (k) (T)) O D O (1o (l
(530)
(E(TTM())
( (T) max {p, o(1)} oth b ||
(531)
for any clean data i E C, which follows from F(T). The same discussion is applied to 14 (T),, and the desired result is obtained from the balance between the numbers of clean and noisy data.
Lemma D.11 (Noise updates of clean data in Lemma C.13) J. Let Ti and T2 be the time steps in Lemma C.13 and let E Te). Suppose that the conditions in Theorem 4.1 |and A(T), B(T), D(r), and F(T) in Lemma C.13 are satisfied.
Then, on a good run, there exists some constant e' > 0 such that d' > c, and
Pt,(T + 1) - (t, (7 s( can - "0lll/lallull max of)
Pi,((T 1) - -i.(T)) s (r) (T) con 1 0. 1|| 2d max of Pt,((T + 1) ((.(()) s (r) - -i c'n 0. | ||2 22 max of),
Pi,(T + 1) (T,((T) 2 - -0 (T) (T) d'nn- 1 2d max max og)
for any clean data i E and [7] \ (1). For any noisy data E. we have
PT.(T + (., < (( can -o max oz), Pai(T + P5) -()) can - max og)
and
PT.2(T + 1) p)2(() 2 8 (T)(( .()())) opan =pll)lllll max o of)) P):2(T + 1) pT,((T s (1) (r)(1 (1) (T)) cppon o
and
P)a(T + 1) --0l-))(TT (1) (T) -82 ( (T) cpam 222101/1211m011100 max max o7), PT +1) -))a()())2 2 -
(1)
(T) pan 1 2d max o7)
for any t E [T] (1.2).
Proof of Lemma D.17 The proof for clean data i e essentially the same as the preceding proofs in Lemmas D.7 1,,1 (7) and (T) is derived in the same way as Equations 475 and 488
and D.9 The lower and upper bound of 7W
75

---

Benign Overfitting in Token Selection of Attention Mechanism
As for Th. (7) and (T) for k + i, the difficulty arises because 1 (7) - (f)())) does not dominate for noisy data
k / like Equation (480p. However, it is resolved using the closer analysis as in Equation (312.. Specifically, using the
22 (7)(1 - 2/ (r) - p - (1) - of (T)) and the balance between 81 (T) - s1 (8) (T)) and of (T) -1 (7)), which are derived from F(r), we have that for l *i,
(T))
(T)) of (
(532)
for any i E C. Therefore, we obtain the desired inequalities for Pt, (T) following the same discussion as in Lemma D.7 The
updates (T) are also derived in a similar way to Lemmas D.7 and D.9
For noisy data j E N the inequalities for 7W (T) Therefore, the analysis is reduced to that of clean data, and the desired results for (). are obtained by flipping the sign of updates. We proceed with the analysis of Pj.c(T), for t E [T] \ {1}. From Lemma C.4 we have
(T) and Iر5ي (r) are derived using a similar argument to Equation 525
10 n
(()) (()). -(())))()))))))) (I + (I((7))11 c ()T W (w(Z())p(((r) (533)
Without loss of generality, we consider j E N = {i [n] y-u = -1,Y = 1}. From the data model defined in Definition 3.1 and Lemma B.12 note that we have
T1 -((111112lllll)) 720 @(pllt1lallall))
(534)
6) = -(1lee1aaalll))) hn
(535)
for t E WL] and E IG), Additionally, using B(r), the current SNR condition SNR² o(n-1) and the parameter
assumptions, we have
[W(T)c 0
(536)
(W(T))+) W(r) 211 (lod:oolllllaloog( = o(n pa po 2d))
(537)
(W()))) W(r) (f))
= o(n po 2 0² (2 log (T (Tn/6) o(n -1 po n o. .)))
(538)
(W()) ((T)e("))
(539)
for any l E [n]] and u E [T] such that (ks u) # (j.t). Therefore, when k = j. we have
T
I (T) Σ (T)
(T)
7000 Σ ( (W())X,,W(r)e))
(540)
u=1
( (T)
(T))
( (W(T)x2)
() ((
-22
(T)
$ (T)) 1 + E (Th@))
(W(T)x)), ((F)e(
2
Σ (-))
(1(5) - 72 (1)) (W(T)x (W)))))WW)))))))
(T) Σ
O (T) (-00)=10) (W()))))
--())-
(541)
2/ (T) - (0(11-9(71 (T) (21e||al|и 2d
(T))) (-
0 + (r)) (tcl)))0ccllll///ee((7//111
(T)()(1
(542)
2 (+)((1 - (1)0))))p0030)001)1))))) 1
(543)
76

---

Benign Overfitting in Token Selection of Attention Mechanism
for some constant 2 > In the first inequality, we used o1 (F) S p(1 =²² (r)), which is the result in F(7), and (
$2 (7) > ( from (r)). From the parameter assumptions, the first line and the first term of the second line in
Equation (541) become dominant, leading to the result. Since s5" (T) 52 (7) dominates o p=
()(()(7)( (7) . ( (7)) and
(T) St () (T)) for any E [T] {1, which follows from FT)), similarly to Equation 511.. we have that for k * j. ( < $2 (T) 22 (T)) o(n - po ||212||11122
(544)
The arguments for (T)) and the upper bounds follow similarly, and substituting them into Equation (533) yields the desired result for (T))
Finally, we analyze the update of PT.(T) for j E EN and E [7] \ (1,2). Similarly, using Equations (536) to 539 there
exists a constant 2g such that
(T = Σ 0 (T)
(00 -0 ( (W(+)x?) (( 1 (T))
(545)
9 (T)
(. (W())x2) W(T)
82
72
Σ
(T)s
(
( (W(T)x?) ( w(T))(")
x=
Σ
(T)s (T) (1 ) (W())x)) W(T)e())
NEITM22
Σ (T) Σ (T) (T0)- 11)) (W(()x))) W(T)
(546)
NE(TIME)
xE(T]
22 (T) ( (T) - (pllelaalllll . n po
+ (r)( (0))01 n
2
+ (Ph0))/((0(((e0)))aaaaaaaaaa....
(547)
++ ptpp::/llllllll--ppp_pp0_2 o2)
+ (00009(1))0(00bbaaaa)))0ddd
(548)
+ +
(
(549)
+ (7)
|| n po
< < (F) 2 o² | 11
(550)
In the first inequality, Equations (547) to (549) become dominant. By the choice of Tj in Equation (521) and Equation (548) is bounded by other two terms. This leads to the last line. The same argument applied to the upper bound and inequalities for (T)) we obtain the desired result.
(T) S psz (T)
(
E. Further Discussion
E.1. Multi-class Setting
In this section, we will see that the analysis in the multi-class setting is a straightforward extension of the binary setting studied in the main paper.
Let K be the number of classes and Wv (P (MI...... vK) Rd K be the weight matrix. The model output is given by
f(x) = W. s (WTP) ER
(551)
Let (ejek be signal vectors corresponding to each class, and we consider the data distribution P. which is modified for the multi-class setting from Definition [.1 Here, we assume the orthogonality and norm equality among signal vectors.
Furthermore, we assume the pretrained linear head satisfies cos 0g > O(1), which is modified version of Equation (2) for Va and the class signal P4, for all [K].
77

---

Benign Overfitting in Token Selection of Attention Mechanism
The concentration inequalities in Lemma B.1 derived by taking a union bound for class numbers and appropriately updating the parameter assumptions depend K. The proof of Theorem 4.1 based the analysis of the empirical loss function the training data S sampled i..d. from P. The empirical loss function when using cross-entropy loss with softmax output is:
exp (f(x()))))
L(W, p) = n
Σ
log
(552)
EnE[K] exp (1(x(1))))
1=1
Σ exp (wux - Mytn)" X (XL p))
log
(553)
The derivatives of the empirical loss function are given by
WT ((W.))
ER=1 exp (vx uy())) x (x(WW)) <(OTs(X(W') W (ux - -y()) X ( Σis exp ( (wI- (wI-uy()) X (X p)) (W p))
(x(OWT)
(554)
K x (
P)S(X pp')) x (ux - My(o)p
(555)
71
ii exp (w) vR) T x ( (W p))
and
V Vp((W.p) p)
Σk k= exp (Ux - uyco) ( ( ('"0(x P)) Vp (vx - -yw)) X (x()WT))
n
(556)
n
E exp (w) vyw) x (T ( (x D)
WX (diag(S(x (1) W p)) - -((X W
p)T) x - (uk-Uy(o)
(557)
Σ exp (vi (wI- x ( x ((O pp)
Since we have
-1 EK exp (IV- vn) T x (0 (TTwT p))
-1
(558)
1+ Eie ( exp (nn) - Uz) x x (x(WTT)) p))
this term corresponds to & in Equations constant order discussed in Lemma C.9 Additionally, using Σ (Myt vx) K (uy K ΣkejK] vk)). we can confirm that Equations (553) and (557) correspond to Equations (9) and [1]. The remaining proof is based on the equations of gradient updates and the results of concentration inequalities; therefore, the statement essentially does not change the multi-class case.
9 and 12 Under the scale condition of the linear head, this term becomes
E.2. Head Optimization
In this section, we provide further discussion on Equation 2 in our problem setting. For the sake of generality, we consider the K-class classification setting following Appendix E.1
The problem setting in Equation specifies the alignment between the classification head and the class signals. It is necessary to appropriately assign token scores, which represent the desirability of each token for solving the task, and to formulate the token selection problem. Since the analysis in this paper can generally be interpreted as an analysis of token selection dynamics under fixed token scores, our result may provide broader insights into attention mechanism across different problem settings, data and model architecture settings. However, it is an interesting question to what extent the problem setting in Equation (2) is justified from the perspective of practical scenarios.
78

---

Benign Overfitting in Token Selection of Attention Mechanism
In the following proposition, we show that the gradient direction of the expected risk zero initialization of the weights forms equiangular tight frame (ETF) with class vectors (Papyan et al. 2020). This corresponds the most extreme case of the alignment condition between the linear head and the class vectors described in Equation (2). Note that the
signal vectors have equal lengths and orthogonality from the definition of our data model. In Definition 3.1 we assumed
that among the weakly relevant tokens, there was only one token for each class different from the true class, namely the
confusing tokens. Here, we assume a more general setting that each weakly relevant token X aligns with a class uniformly sampled from ]K]. that is, Xa - p + tu for E W.
Proposition E.I. Suppose that the weights P. W, and Wv are initialized to zero. We also consider a generalized version of the data model in Definition 3.1 where each weakly relevant token aligns with class k uniformly sampled from [K]. Then, the gradient descent direction of the expected risk for Wv forms an ETF geometry consisting of the signal vectors, i.e, we have
- -WV, = =)) x [Ix
1 - 1x1k) K
(559)
where [(Wv)= = EX [(Ww)].
Proof. Taking the gradient of Equation 553
the gradient of expected loss at Wy 0 is given by:
Vu. (J(X) - [(X)r)
CWV = 0) = -BXYY
(560)
MEEK] EeE[K] exp (Ve
(Vc-MKc) -
X XWTp)) O.
B x (1 - (1-kkirzz)
(561)
K
K (1 (1-K) EE Σ X Y = k
+
- R Σ xt (xx1+
(562)
KT
K
K 1
E Σ x, Y
-E
Σ xx 1
(563)
KºT
where we changed the order of gradient and integral in the first line, and in the second equality, we denote by 1 the indicator function which returns 1 if the event A is satisfied and otherwise returns O. The third line follows that Y is sampled from a uniform distribution over [K], and label noise is added to different labels uniformly.
Let is be the mean of class signals ΣnE(K] / K. We have
E
Σ X1 Y
(1 1)
Σ X Y
=k
+
7
E
Σ x Y
(564)
ME(K]((A
(1 m!)
12 +
+
11
Σ ME(K]N(A
Ms +
|
T PFL
(565)
T
K
(1
1) + K-1 (K
14)) + +
(566)
K
KK11" 7 - 1911)) p.
K
K x1 11
b +
(567)
79

---

Benign Overfitting in Token Sclection of Attention Mechanism
Train Loss
Test Loss
5300
0.44
5300
620
0000
400
4300
505
Dimonsion
-000
Immonsoon
3000
2:48
-200
2300
1000
1000
5:42
100
100
39
Signal
Norm Mella
Signal
Norm
Ialla
Figure 5: Heat-map of train loss and test loss under different signal norm || and the dimension d, after 1000 iterations. The yellow color indicates a higher loss value. The model and data follow the setup in our analysis.
where we used the definition of our data model. Therefore, Equation 563 leads to
CWW =
K K-1
(1- 1----1") Ms +
K
K
7 (W/l)) f
K2T
K
K K-1
Σ
(11 K --1- He +
K
K 1- + mml1)) 24) }
K
(568)
K-
WE(K]((K)
K-1
(1- K
(1 K - 11)
1k K-1 (K - m1)
1
(569)
K2T
K-1
K
K
1 K K-1 1
K K-1 (M μ))
(570)
K2T
Consequently, we have Cw = 0) x 1k
p. which concludes the proof.
This proposition is the result for the case of zero initialization. However, by taking the variance of the initialization sufficiently small, as Assumption A8, similar alignment with the class signal can achieved under random initialization.
While we have considered the gradient of the idcalized expected loss, it natural ask about the gradient of the empirical loss. When using the same training set S as in the main theorem, v contains the noise cf" from the training examples, and thus the arguments in Lemma B.12 cannot be applied, meaning that the same result does not hold. This is because noise memorization occurs even in the head, which outside the focus of our study on token selection. On the other hand, if we use different training set S", it is possible appropriately modify the argument and obtain results in the main theorem. Regarding the class signal, can show that v 14 Z || as in this proposition. However, when the input noise is large, specifically when a. /d 2 ||//|/2. the condition k cos 0g > @(1) in Equation (2) may no longer hold. Nevertheless, in the discussion of token scores Lemma 12 suffices to show that v 14 dominates v 6 which indeed holds under the parameter assumptions in Section 3.5
F. Additional Experimental Results
F.I. Additional Synthetic Experiments
Heat-map experiments. We conducted another synthetic experiment when varying d and ||||||2. Figure shows the train and test loss when varying the dimension d and the signal norm || under the same setting as in Section 5 in the main text. This figure shows that the balance between the dimension and the signal norm is significant for achieving low train and test loss. Additionally, while our theoretical results in Theorem 4.1 propose the boundary between not-overfitting and benign
80

---

Benign Overfitting in Token Sclection of Attention Mechanism
Clean sample
Clean sample
Clean sample
Softmax probability s
Softmax probability
1.5
4y. Reeerant R
Softmax probability
Relecser
ss. Feieant R
Weikly Révront 1.
0.75
sy. Wiahly Recoaan W-y-
2t. Weskly Rebvvntt W..-
3.5-
Ty. Weekly Relecont Wy.
Weaky Rebecat Pay
My. Vically Vedcrant Wy-
344
Rv, Inttesantt
St.
0.50
24. Imbevant
Incleaart 2
0.4
sny. Incio
line ocurt
3.21
Nv, Irtiovant.
lireesont
0.25
2). Inclevant Ireeevont
Sv. Indevartt
0/2
sr. Iime mont
lice
Sy, Insievant
my. live event
3.3
0.00
0.0
Noisy sample
Noisy sample
Noisy sample
Softmax probability
Softmax probability
0.5
Softmax probability
0.75
3.3
0.6-
0.30
3.1
3.2
0.2
0.25
1.7
0.0
0.00
600
500
1000
-03
600
1000
200
800
1000
Time step
Time step
Time step
(a) Large noise setting: d = 5000, |||| Final training accuracy is 1.0 and test 20. Final training accuracy is 1.0 and test 100. Final training accuracy 1.0 and test accuracy is 0.49 (Harmful overfitting).
(b) Balanced setting: d 2000, ||
(c) Large signal setting: d 1000, ||
accuracy is 0.99 (Benign overfitting).
accuracy is 1.0 (Benign overfittin).
Figure 6: Dynamics of softmax probability for one-layer transformer encoder. The top represents a clean sample, while the bottom represents a noisy sample. From left to right, the configurations of d and ||||| follow those used in Figure 3 While the right setting in Figure 3 corresponds to the not-overfitting case, it now shows benign overfitting.
overfitting as SNR" O)n-1), the training loss boundary in the left figure approximately forms a quadratic curve. This corresponds to that the ratio between d and || remains constant at the boundary, with a fixed training size n.
One-layer transformer encoder. We also conducted experiments under more practical model setting using the same synthetic data. Specifically, we used one-layer transformer encoder with a single-head attention. Compared to the model in the analysis, this model additionally incorporates non-linear feedforward layers, normalization, and skip-connections. The additional experiments here aim to provide further insights into the behavioral differences arising from joint optimization, rather than to support our analytical results. We first examined the token selection dynamics in the same way as in Figure 3 and the results are presented in Figure 6 The main findings are as follows:
Similar benign overfitting was observed depending on the relationship between d and || In low-SNR setting, harmful overfitting was also observed in Figure 6aa However, Figure bc shows that the model exhibited benign overfitting instead of not overfitting. This aligns with the intuition that increasing model capacity facilitates fitting to the training data.
In the benign overfitting case, token selection for noisy samples progresses more rapidly. The bottom in Figure 6b shows that the softmax probability assigned to token X2, which most aligns to label noise, increases more rapidly than that of Figure 3b The ability to learn token scores dynamically enables cooperative interaction between the token selection mechanism and the feedforward layers, which could lead to faster token selection, as shown in the figure. Note also that, unlike the setting analyzed in the main text, fitting the label noise does not necessarily require selecting X2, and a different token is selected depending on the initialization.
Furthermore, Figure 7 shows the results of heatmap- baased experiment analogous to Figure 5 For the training loss, observed very small values within the plotted range of ||wl|/ and d, which aligns with the observation of benign overfitting instead of not-overfitting in Figure 6c. As for the test loss, we obtained a similar boundary structure, but with clearer separation between the well-generalizing and sporly-generalizingg regions. We suppose that this difference due to the two main factors: i) the ability to learn output scaling and ii) the ability to distribute the effect of noise memorization not only within the attention mechanism but also to the feedforward layer.
81

---

Benign Overfitting in Token Sclection of Attention Mechanism
Train Loss
Test Loss
435
117
4S4
207
Dimension
415
2 Signal Norm ||A4
Signal Norm ||A
Figure 7: Heat-map of train loss and test loss under different signal norm || and the dimension d, after 1000 iterations. The yellow color indicates a higher loss value. The data follows the analytical setup in this paper, but the model used is a one-layer transformer encoder. For comparison with Figure the range of the heat map aligned with that of Figure
Table 4: Details of image datasets.
Dataset
Number Class
Input Size
Training Size for Classifier Head
MNIST
28 x x
10000
CIFAR-
10
32 x 32 x 3
10000
PNeUmoniaMNIST
28 28
3000
BreastMNIST
2 28 x 28 x
200
Table 5: Details of natural language datasets.
Dataset
Number of Class
Average Word Count
Training Size for Classifier Head
AG-news
4
10
10000
TREC
6
10
3000
F.2. Additional Information for Real-World Experiments
In this section, we provide the additional information for the experimental setup. We prepared the pre-trained ViT (Dosovitskiy et al. 2021) and BERT (Devlin al. 2018) models using huggingface transformer library (Wolf et al. 2020). These models use the default configuration; specifically, they consist of 12 attention layers, the embedding dimension is set to 768, and the hidden dimension of the feed-forward layers is set to 3072. Dropout was not performed in any layers during the following training. Since the classifier head is initialized, we first train only the classifier head on a subset of the training data without label noise to align with Equation ( The sizes these sub-datasets vary due to the original dataset size, as summarized Tables and [5] The experiments in Section 5 and Appendix F.3 use training data that a different split from the one used in the classifier head pertaining, and training is performed by initializing and updating only the attention weights in the final layer. As explained in Section this setup is designed to align with our analytical setting closely as possible, treating the pretrained model up to the final layer as feature extractor, and training single-layer attention weights on top of these features. It should be noted, however, that the data model of course does not follow our analytical setups, and differences such as normalization, skip connections, and multi-head attention remain. During the experiments, the AdamW W optimizer (Loshchilov & Hutter (201)) without weight decay was used with learning rate of 5e-5, along with linear warmup and learning rate decay. The used datasets are described as follows.
Image Dataset We conducted experiments on image classification on the following four datasets. The MNIST (LeCun et al. 2010) dataset consists of scale 28 x 28 images with 10 classes. Each image is copied to form a 3-channel input and fed the common image processor for the pre-trained ViT model. The CIFAR-10 (Krizhevsky et al. 2009 the
82

---

Benign Overfitting in Token Selection of Attention Mechanism
Table 6: Training loss and test accuracy when only the final attention query-key weights are trained on sub-dataset for 2000 epochs, sufficiently long time. The results show the average over three different runs with the standard deviation.
Label noise f) 0.0.
Label noise 7) 0.2.
Training Size
Training Size
Dataset
Eval
Dataset
Eval
20
200
1000
20
200
1000
train
0.00
0.00
0.00
train
0.00=0.00
0.04=0.01
0.14=0.02
MNIST
MNIST
test
9..
91.820.1
93.9+0.1
test
84.0+1.2
87.3+1.2
88.4
train 0.00
0.00
0.0000.00
train 0.18=0.07
0.18=0.01
0.33
CIFAR-10
CIFAR-10
test 96.0=o.1
95.9+0.1
96.3+0.1
test 95.5+0.4
94.3+0.3
93.2+0.1
Pneumonia
train 0.00+E..0
0.00
0.00
Pneumonia MNIST
train 0.03=0.04
0.02=z.00
0.07=0.00
MNIST
test
80.423.1
79.022.0
81.520.9
test 79.221..0
81.212.2
82.520.7
Breast
train 0.0840.02
0.122n.01
Breast MNIST
train 0.1220.03
0.17=0.02
MNIST
test 14.4:1.4
78.4
test 74.8122.0
76.312.3
train 0.0010000
0.00
0.00=0.00
train 0.00=0.00
0.00=0.00
0.00 00
AG-news
AG-news
test 84.0=0.3
85.9+0.5
85.5+0.3
test 82.5+0.7
77.8+0.6
69.6+1.1
train 0.00+0.00
0.00
0.00
train 0.24g0.33
0.11
0. 0.12
TREC
TREC
test 81.420.3
7
82.8
test 79.2:2.4
73.3+0.8
69.1
dataset composed of 32 x 32 color images in 10 classes. These classes are mainly made of vehicles and animals. Finally, we focused on the MedMNIST (Yang et al. (2023) dataset, specifically using PNeImoniaMNISTT and BreastMNIST, which are tasks for disease detection. Both datasets consist of 28 x gray-scale images similar to MNIST, and they are binary classification settings based the presence or absence of disease. Table summarizes the details of these datasets.
Furthermore, provide comments how our data model, based on relevant, weakly relevant, and irrelevant tokens defined in Definition 3.1 corresponds to the input images. For example, consider the case of cancer detection using medical images such as MedMNIST. In this case, the patches directly indicating cancer, such as tumors or lesions, can considered relevant tokens. The patches showing enlarged lymph nodes or inflammatory signs, which often co-occur with cancer but are not definite, can be categorized as weakly relevant tokens. Meanwhile, the patches showing normal tissue or background anatomy, which are irrelevant the task, correspond irrelevant tokens.
Natural Language Dataset We conducted experiments sentence classification in natural language in addition to image data. The AG-news (Zhang et al. 2015) is the dataset for the topic classification of news articles. It has four largest classes: "world"", "sports", "business"" and "scienceetechhologyy The TREC (Li & Roth 2002) dataset is the dataset for question classification in 6 classes. Each question is labeled based on the content and the question type. The details of these datasets are summarized in Table[3] Note that the text data can also correspond to the data model in the paper based on the relevance of each word to the class, similar to the case of images.
F.3. Additional Real-World Experiments
In this section, we provide additional results under the real-world experiment setting described in Appendix F22
Table presents results corresponding to the experiment Table of the main text, where the label noise level v = varied to 77 0.0 and r 0.2. In the absence of label noise, the model fits the data across all settings except for BreastMNIST. Notably, in the case of AG-news, where increasing the sample size under 7 0.1 led to worse test accuracy, no such accuracy degradation is observed when 0.0. The case of v 0.2 shows similar trends to those observed with 7 0.1 in the main text.
Next, Figure presents the results as the label noise ratio is varied continuously from r = 0.0 to 0.2. In datasets such as CIFARIO and TREC, the training loss increases as the label noise increases, whereas in MNIST, PneumoniaMNIST. and AG-news, the model maintains relatively good fitting across all noise levels. Regarding test loss, we observe that language
83

---

Benign Overfitting in Token Sclection of Attention Mechanism
Train Loss
Test Loss
mnist
336
cifar10
pneumoniamnist
pneumoniamnist
breastmnist-
breastmnist
166
agnews
001
das
015
017
010
das
Noise Ratio
Noisc Ratio
Figure 8: meat-map of train loss and test loss when the label noise ratio v is varied from 0.0 to 0.2. The training size 200, and the training setup follows that Appendix
The yellow color indicates a higher loss value.
MNIST
CIFAR10
Max softma probabiitt MMa soltmax probab lkty Max softmox probabiltty
Max softmax probsiility Max seltmax probab lity Max softmax prebability
0.4
04
02
0.
0.0
PneumoniaMNIST
BreastMNIST
0.-
0.4
0.4
0.
0.
0.0
AG news
TREC
0.0
0.05
250
500
750
1600
1250
1500
1750
2000
250
300
750
1300
1250
1500
3750
2000
Training epochs
Training eaochs
Figure 9: Dynamics of the maximum softmax probability the position of the CLS token in the final attention layer. The experimental setup follows the real-world experiment described in Appendix F2 with label noise level v) 0.1.
datasets such AG-news and TREC are relatively more difficult to predict, and the increase in test loss is more severe as the noise ratio increases.
Finally, Figure 9 illustrates the evolution of the maximum softmax probability at the position of the CLS token in the real-world experimental setting. For all datasets, the softmax probability increases as training progresses. In MNIST and PneumoniaMNIIST, the model tends to focus on a single token, whereas in other datasets, attention does not necessarily concentrate on single token. This behavior may be attributed differences in the data distribution and architecture, such skip connections and non-linear feed-forward layers, which not captured by our analytical setting.
84


Reviews:
[{"review_id": "KTf1HDcp4D", "reviewer": "Reviewer_ETcq", "summary": "The paper develops a theoretical framework to analyze the dynamics and generalization properties of token selection in attention mechanisms under label noise, focusing on a one-layer attention network for binary classification. It demonstrates that with a high signal-to-noise ratio (SNR), the model selectively fits clean samples and generalizes well (not overfitting), while with a low SNR, the model overfits noisy training data yet still achieves low test error—a phenomenon known as benign overfitting. The analysis introduces key concepts such as the evolution of softmax probabilities and the \"attention gap\" metric, which quantify how the mechanism distinguishes between relevant and noisy tokens during training, and it highlights a delayed generalization process reminiscent of grokking. Extensive experiments on both synthetic and real-world datasets support the theoretical findings by showcasing transitions between harmful overfitting, benign overfitting, and non-overfitting regimes.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We hope that our answer has addressed all your concerns and pray this reply will change your decision.\n\nWe first answer your question because it is an important point relating to our contribution.\n\n---\n> 1: What is the specific difficulty of label noise setting?\n\nThe 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:\n1. Signal learning vs memorization in each token selection (Figure 2).\n2. Clean samples vs noisy samples in the learning direction of class signals​.\n\nThese 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]).\nFor 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.\nThus, 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.\nIn 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.\n\n---\n> 2: The novelty and impact are limited compared to [Jiang+, 2024].\n\nThe 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.\nAdditionally, we discuss finer differences starting from line 627 in the Appendix.\n\n---\n> 3: There is no sufficient comparison with [Magen+, 2024].\n\nWe discuss the differences from line 143 and a more detailed difference from line 650 in the Appendix, which clarifies the novelty of our work.\nIn response to the review, **we have added a row for [Magen+, 2024] in Table 1**.\nFurthermore, since the Appendix previously only provided a theoretical comparison, **we have also newly included a comparison of experimental validation**, as suggested.\n\n---\n> 4: There is no definition of “good run”.\n\nWe 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.\n\n---\n> 5: There are some concerns about the threshold between benign and harmful overfitting.\n\nAssumption (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.**\nAlthough 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].\nFurthermore, the boundary is demonstrated through synthetic experiments (Figure 4, right).\nIn response to your comment, **we have explicitly annotated Figure 1 with the specific values** instead of \"large\" and \"small\".\n\n---\n> 6: Fixing head $\\nu$ as in Assumption 3.3 limits the generality of the results.\n\nOur 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.\n\nAssumption 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**.\nIn 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.\n\nFinally, 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.", "claims_and_evidence": "The paper’s claims are largely supported by rigorous theoretical analysis and clear empirical evidence, with the real-world experiments being a particular highlight. However, when compared to Jiang et al. (2024), the paper appears relatively less impressive in terms of novelty and overall impact. Moreover, the comparison with Magen et al. (2024) is not sufficiently detailed; it would be beneficial to include a comprehensive comparison—preferably in Table 1—to clearly delineate the strengths and weaknesses of each approach in both the theoretical analysis and experimental validation.", "methods_and_evaluation_criteria": "The experiments are clear and can effectively support the theoretical results, with a primary focus on validating the theoretical claims.", "theoretical_claims": "I reviewed the theoretical proofs and did not find any major issues overall. However, I have two points for improvement. First, the appendices frequently refer to a “good run” without a precise definition; clarifying what qualifies as a good run would be beneficial. Second, while the paper distinguishes between benign and harmful overfitting, it would be clearer to provide an explicit formal threshold or result that delineates the boundary between these two regimes.\n\nAdditionally, I feel that directly assuming a fixed $\\nu$ (as in Assumption 3.3) is somewhat of a “shortcut”. This assumption sidesteps the challenge of analyzing the full dynamics when $\\nu$ is also learned, and may limit the generality of the results. Providing further justification or exploring a more general setting ould strengthen the paper’s theoretical contributions.", "experimental_designs_or_analyses": "I examined the experimental designs for both the synthetic and real-world experiments. The synthetic experiments are well-designed to validate the theoretical predictions.", "relation_to_broader_scientific_literature": "Overall, this is an interesting topic with some novel findings. The paper builds on and extends prior work on benign overfitting—studied in linear regression, two-layer neural networks, and kernel methods (e.g., Bartlett et al., 2020; Hastie et al., 2022; Liang & Rakhlin, 2020)—by applying these ideas to the attention mechanism in transformers. \n\nWhile the topic and findings are compelling, the contribution is somewhat diminished in comparison with recent works like Jiang et al. (2024) and Magen et al. (2024), which also explore similar phenomena in vision transformers.", "essential_references_not_discussed": "I agree that the paper seems to have provided all the essential references needed to understand its context and contributions. The citations cover the key areas of benign overfitting, attention mechanism dynamics, and related analyses in transformers. No critical works appear to be missing.", "other_strengths_and_weaknesses": "Other strengths of the paper include its clear and well-organized presentation, as well as its rigorous approach that blends theoretical analysis with empirical validation. The work extends existing benign overfitting analyses to the context of attention mechanisms in transformers, which is a creative and timely direction given the widespread use of these models.\n\nOn the other hand, some weaknesses are apparent. The novelty is somewhat limited when compared to concurrent works such as Jiang et al. (2024) and Magen et al. (2024), which explore similar phenomena in vision transformers. Additionally, certain assumptions—like directly fixing $\\nu$ (Assumption 3.3)—could be seen as a shortcut that may restrict the generality of the results. Clarifications regarding the definition of “good run” in the appendices and a more explicit delineation between benign and harmful overfitting regimes would also improve the paper's clarity and impact.", "other_comments_or_suggestions": "I suggest that in Figure 1 the paper quantitatively indicate how large or small the values are (for example, by annotating the scales or ranges on the plot). \n\nAdditionally, in the second part of Theorem 4.1, it would be beneficial to explicitly state the lower bound for SNR to clarify the precise conditions under which benign overfitting occurs.", "questions_for_authors": "What is the specific role of label noise in your training process, and how does its presence complicate the dynamics compared to a scenario without label noise?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3h9N0jbRrm", "reviewer": "Reviewer_ddf6", "summary": "The paper studies benign overfitting in token selection within the attention mechanism, using a data model consisting of signal and noise and a one-layer attention model. This work characterizes the conditions under which benign overfitting, harmful overfitting, or no overfitting occurs.\n\n## Update After Rebuttal\nMost of my concerns were addressed during the rebuttal, and I increased my score to 2 (weak reject). However, I believe that additional reviews are necessary to determine whether the revised manuscript fully addresses the concerns raised by all reviewers. For this reason, I still lean toward rejection.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We hope that our answer has addressed all your concerns and pray this reply will change your decision.\n\n---\n> 1: The proof techniques lack motivation. Section 4.1 lacks sufficient explanation or intuition in the current draft.\n\n**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.\n- The overview of Section 4.1 after left column, line 308\n- The reason for introducing Definition 4.4 after left column, line 310\n- The motivation of the function $g(x)$ at line 329 (see also Comment 9 below)\n- The brief explanation of how the dynamics of the attention gap are derived after Lemma 4.5\n- The new section titled \"Proof Sketch\" before Section C\n\n---\n> 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.\n\nIncreasing 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.\nTo 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$.\n\n---\n> 3: Yun et al. (2020) is an incorrect reference.\n\nYou are correct, and we have already fixed it. Additionally, we have reviewed the references again to ensure there are no errors.\n\n---\n> 4: Assumption 3.3, (A1)–(A8), and the first two lines in Theorem 4.1 seem overly restrictive, and their significance is unclear.\n\n**Assumption 3.3**  \nOur 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]).\nAssumption 3.3 is necessary to formulate the problem of token selection given the assigned token scores.\nTo clarify this point, **we have moved this part from Assumption to the Problem Setting**.\nFinally, **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.\n\n**(A1-A8)**  \nIn 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.\nSeveral papers express parameter assumptions using big-O notation ignoring logarithmic dependencies [Cao+, 2022; Jiang+, 2024], but we present them explicitly.  \nCould 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.\n\n**Assumption in Theorem 4.1**  \nThe 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.\n\n---\n> 5: Running head should be fixed.\n\nThank you for your comment. We've already fixed it.\n\n---\n> 6: The authors should provide the gradient update equations for $W$ and $p$.\n\nFor 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:**\n\"The weight updates with specifically calculated loss gradients are provided in Appendix B.\"\n\n---\n> 7: What happens in “Not Overfitting” case after the time step $\\tau$?\n\nIn the not-overfitting case, signal learning is dominant, and the model fits only the clean samples as stated in the theorem.\nThis is discussed in detail in Figure 2, the paragraph from line 291, and Appendix D.2.\n\n---\n> 8: Time step $\\tau$ should depend on uncertainty $\\delta$ but it seems there is no such dependency.\n\nUncertainty $\\delta$ indirectly affects the time step $\\tau$ through the parameter assumptions (A1)-(A8).\nA similar order notation for the time step is also found in existing benign overfitting studies [Cao+, 2022; Jiang+, 2024].\n\n---\n> 9 : Why does $g(x)=2x+2\\sinh(x− \\log⁡T)$ appear in the proof techniques?\n\nDue 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**.", "claims_and_evidence": "The main theorem statement is clearly presented. While I did not have time to read the full proof in the appendix, I find that the proof techniques described in the main text lack motivation. That is, it is difficult for the reader to understand the intuition behind the chosen proof techniques.", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "I did not check the full proof in the appendix and only read the proof techniques section in the main text. Due to its less-motivated presentation, I cannot confidently assess the correctness of the theorems.", "experimental_designs_or_analyses": "I find the experimental results, especially the real-setting experiments, somewhat trivial and expected—showing that more data leads to better loss. I highly encourage the authors to include additional experimental results that better highlight the intuition behind their findings.", "relation_to_broader_scientific_literature": "The paper studies benign overfitting in a more general data distribution and a more complex neural network architecture compared to previous works, such as Cao et al. (2022) and Oymak et al. (2023). This contribution broadens the understanding of benign overfitting in more general settings.", "essential_references_not_discussed": "I believe the paper includes most of the relevant works on benign overfitting and the theoretical aspects of attention and transformers. However, I found an incorrect reference. In line 107 (left), Yun et al. (2020) is cited as a related work on the training dynamics of transformers. However, this paper does not consider training dynamics; instead, it shows that transformers are universal approximators. I recommend the authors carefully review their citations, as there may be other misreferenced works that I did not notice.", "other_strengths_and_weaknesses": "The problem setting is novel and more general than in previous works. Additionally, the comparison to prior studies (Table 1) is helpful. However, there are several weaknesses that should be addressed:\n\n- Some assumptions are strong and not well-motivated. For example, Assumption 3.3, (A1)–(A8), and the first two lines in Theorem 4.1 seem overly restrictive, and their significance is unclear.\n- The proof techniques lack motivation. Section 4.1 is difficult to follow because it primarily consists of lemmas without sufficient explanation or intuition in the current draft.", "other_comments_or_suggestions": "* Running head should be fixed\n\n* I suggest that the authors provide the gradient update equations for $W$ and$p$ for the convenience of readers. Since these equations are not included in the current draft, it is difficult to follow the technical details. Also, I encourage the authors to clarify the reason why the technical assumptions are required in their analysis.", "questions_for_authors": "- In Theorem 4.1, what happens in “Not Overfitting” case after the time step $\\tau$?\n- In Theorem 4.1, I think time step $\\tau$ should depend on uncertainty $\\delta$ but it seems there is no such dependency.\n- Why does $g(x) = 2x+2\\mathrm{sing h}(x-\\log T)$ appeare in the proof techniques?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BaY0950Gq2", "reviewer": "Reviewer_vUF4", "summary": "This paper presents theoretical analysis of benign overfitting in token selection of the attention mechanism focusing on the training dynamics and the generalization performance. The author shows that under some conditions based on signal-to-noise ratio, \"benign overfitting\" phenomenon occurs in the attention mechanism.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We hope that our answer has addressed all your concerns and pray this reply will change your decision.\n\n---\n> 1: The comparison among not, benign, and harmful overfitting lacks in real-world experiments. It needs more explanation and experiments for harmful overfitting.\n\nThe 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.\nHowever, 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.\n\nIn 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.\n\n---\n> 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?\n\nIf 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**.\n\nSince 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]).  \nAssumption 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**.\nFurthermore, please note that this setting is practically plausible in the area of parameter-efficient fine-tuning, such as prompt tuning and LoRA.\n\nOur 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.\n\n---\n> 4: The assumptions (A1-A8) seem very strong.\n\nIn 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.\nWhile 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.\n\nCould 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.\n\n---\n> 5: Three distinct token groups and a consistent scale $\\rho$ are implausible.\n\nBuilding on previous research, we have carefully designed the analysis setting to better align with real-world scenarios.\n\nThe 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.\nOur study introduces several realistic elements to the analysis, including middle states termed \"weakly relevant tokens\" and non-orthogonality between signal and noise.\nTo 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.\n\nFinally, our analysis also holds if the scale $\\rho$ varies across examples. **We have added this clarification after Definition 3.1.**\n\n---\n> 6: What is the motivation for using $g(x)$?\n\nDue 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)$.**", "claims_and_evidence": "Both the theoretical analysis and experiments clearly support the claims.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "The proofs seem correct.", "experimental_designs_or_analyses": "In the real-world experiments, the comparison among \"not overfitting\", \"benign overfitting\", and \"harmful overfitting\" is not shown. Maybe it needs more explanation and experiments for harmful overfitting.", "relation_to_broader_scientific_literature": "This paper addresses the attention mechanism from the perspective of benign overfitting, which will guide the use of attention-based models.", "essential_references_not_discussed": "No more.", "other_strengths_and_weaknesses": "Strength:\n\n1. The paper understands the token selection ability of attention mechanism from the \"benign overfitting\" aspect.\n\n2. The mathematical analysis and proofs are solid and clear.\n\n3. The paper proposes a novel method of using an auxiliary function $g(x)$ to describe the training of attention gap. This method successfully deals with the difficulty caused by the softmax probabilities.\n\nWeakness:\n\n1. The model may be a bit simple, involving only one attention layer without any additional nonlinearity. What if there is FFN and nonlinearity on top of the attention layer?\n\n2. Fixing the value of $v$ and only training $W$ and $p$ during the training process is less challenging. Once $v$ is also included into the updating procedure, what will it influence the attention mechanism?\n\n3. The assumptions in the paper seem very strong (A1-A8). It is not clear why we need these assumptions. It is recommended to contain some explanations about these assumptions.\n\n4. The data distribution appears implausible due to exactly three distinct groups and a consistent scale $\\rho$ that separates the relevant token from the weakly relevant token.", "other_comments_or_suggestions": "No more.", "questions_for_authors": "Why does the author consider the function $g(x)$ (line 328) to analyze the attention gap? What's the motivation of utilizing this kind of function?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LqoL6CsB1D", "reviewer": "Reviewer_BNsU", "summary": "This paper explores benign overfitting in the token selection process of attention mechanisms, analyzing how transformers generalize despite fitting noisy training labels. It adopts feature learning framework to explain when models ignore noise (high SNR) or fit it while still generalizing well. Through theoretical analysis and experiments on both synthetic and real-world datasets, the paper demonstrates how training dynamics influence token selection.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review. We appreciate that you highly evaluate our work.\n\n---\n> 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.\n\nIn response to your comment, **we are making the following additions to Section G.2**:\n1. 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\n2. Heatmap illustrating how accuracies change as the label noise $\\eta$ varies continuously within the range $[0, 0.2]$ for each dataset.\nWhile 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.\n\n---\n> 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?\n\nTo address your concern and enhance the motivation, we have made the following improvements in Section 4.1 and Appendix:\n- To prevent the abrupt start of Section 4.1, **we have added the following overview after left column, line 308**:\n\t  \"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.\"\n- **We have clarified the reason for introducing Definition 4.4** by adding the following sentence after left column, line 310:\n\t  \"This quantity is useful for evaluating the softmax probability at each time step $\\tau$ for a given training example.\"\n- **We have added the motivation of the function $g(x)$ at line 329** as follows:\n\t  \"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.\"\n- **We have added a brief explanation of how the dynamics of the attention gap are derived after Lemma 4.5**:\n\t  \"This lemma is shown by tracking the gradient descent dynamics and conducting induction argument with several desirable properties.\"\n- **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.\n\nThe 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).\n\n---\n> 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?\n\nThank you for your question. We answer the questions as follows.\n1. 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.\n2. 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.\n3. 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.", "claims_and_evidence": "This paper is well written, with rigorous theoretical analysis and clear experimental evidence.", "methods_and_evaluation_criteria": "The methods and evaluation criteria makes sense for the problem.", "theoretical_claims": "The theoretical analysis is rigorous. By adopting feature learning framework, the authors analyzed the training dynamic systems and give the description of growth rate of the attention. I believe the rate is analyzed from some implicit equations but I did not check the details.", "experimental_designs_or_analyses": "The experiment results further supported the theory.", "relation_to_broader_scientific_literature": "I appreciate the authors efforts to compare their results with existing works and highlight their contributions.", "essential_references_not_discussed": "_", "other_strengths_and_weaknesses": "Several comments:\n\n[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.\n\n[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?", "other_comments_or_suggestions": "See above.", "questions_for_authors": "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?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The authors present an analysis of the optimization and generalization dynamics of a two layer transformer in the high dimensional regime.  They demonstrate how the signal-to-noise ratio in the token generation model can determine whether or not the transformer exhibits benign overfitting (low SNR setting) or classical generalization (high SNR setting).  

The reviews were mixed, with some reviewers having hesitation about the significance of the problem and with the improvement over Magen et al.  Magen et al's work was posted on arxiv on Oct 10, which was less than four months prior to the ICML submission deadline, and thus is a concurrent work and can't be used to rule against this.  On the one hand, I do think the results here could be improved by demonstrating whether or not *harmful* overfitting ever occurs in the high-SNR setting.  On the other, I think the current paper does advance our understanding of how attention mechanisms fit noisy examples, and how benign overfitting can occur with this architecture.  I think the positives of this paper clearly outweigh the negatives and I recommend acceptance.
Decision: Accept (poster)