Benign Overfitting in Single-Head Attention

We thank the reviewer for their thoughtful feedback. Below, we respond to each concern and question in order.

W1: First, we chose to combine the key and query matrices to simplify the analysis, which is already technically involved. We note that this simplification is very common in prior works studying attention mechanisms. For example, [1,2,3,4,5] (as well as many other papers) also used a combined key-query matrix. Then, following many prior works on classification (or regression) tasks in transformer, we introduce a classification token (CLS token or prompt), whose corresponding output is used for prediction, which leads us to Eq. 2, that is, $f(X; W, v) = v^\top X^\top \mathbb{S}(X W^\top q)$, that also has been studied in prior works (e.g., [1,2,6] and more). We then follow the approach of [5], which showed that optimizing W with a fixed q is equivalent to optimizing q and removing W, as done in our paper. We note that in Section 3 of that paper, they are optimizing both v and q, exactly as we do. We also note that our model in Eq. 3 can be obtained by fixing $q=(1,0,\ldots,0)^\top$ in Eq. 2, and then the vector $p$ from Eq. 3 is simply the first column of $W$ in Eq. 2.

W2: We agree that considering a smaller step size and analyzing the trajectory for more than two steps is interesting. However, analyzing even two steps of GD in our setting is challenging. This analysis gives us many insights into how attention handles noisy label data, how the attention focuses on the signal token for clean samples and on the noisy tokens for noisy samples, sufficient and necessary conditions to allow benign overfitting, and the relation between different parameters. It also provides insights into how the trajectory looks after more steps of GD (see Remark 5.1, line 371), which we validate in our experiments (Figure 3 in the appendix). Our experimental results suggest that our theoretical findings extend beyond the two-step setting, demonstrating consistency with longer training and a broader range of learning rates. While a full trajectory analysis is an interesting direction, our results already highlight fundamental phenomena.

Moreover, we prove that min-norm/max-margin solutions also exhibit benign overfitting, where these solutions roughly correspond to the convergence of GD with (a small) weight decay. Indeed, GD with weight decay tends (after enough iterations) to solutions that minimize the loss while preserving a small weight norm. We also refer the reviewer to lines 232-238, which provide additional motivation for analyzing min-norm/max-margin solutions.

W3: Benign overfitting has been studied across a variety of architectures, including CNNs, MLPs, and linear models. Given the widespread interest in Transformers, we believe it is natural and valuable for multiple studies, including ours, to explore benign overfitting in this context. We view this as a sign of growing interest in the topic, and believe our contributions offer a significant and distinctive perspective.

Although [7], [6], and our paper all study the benign overfitting phenomenon in attention mechanisms, there are several fundamental differences in setting, assumptions, and theoretical scope. Below we provide a detailed comparison to clarify the novelty of our contribution.

• Jiang et al. [7] do not incorporate label noise in their analysis. The presence of label noise is not a minor extension but a crucial element in understanding whether interpolation is compatible with generalization. In fact, label noise is central to the definition of benign overfitting in the literature and appears in almost all relevant works, with only a few exceptions such as [8], and [9]. Without label noise, many existing results on benign overfitting can be trivially explained through standard uniform convergence arguments. For example, consider the classical result of Bartlett et al. [10] on benign overfitting in linear regression, where they proved under what conditions on the data covariance matrix minimum-norm interpolators exhibit benign overfitting. In linear regression, generalization in a realizable setting with a sufficiently large dataset follows from standard norm-based uniform convergence bounds. Hence, [10]’s result is challenging and interesting mainly because it considers label noise.

• The work of Sakamoto et al. [6] (which was developed in parallel to our work) fixes the attention head $v$ and performs gradient descent only on the attention weights. They also make assumptions on the correlation between the fixed attention head $v$ and the signal token to ensure benign overfitting happens. Their assumption on $v$ holds with exponentially small probability (in the dimension $d$) if $v$ is drawn from a spherically symmetric distribution. In contrast, our setting allows both the attention vector $p$ and the prediction head $v$ to be fully trainable, and we analyze their joint optimization dynamics.

Finally, both works show that there exists a training step where the model achieves a low test loss and zero training error, but they do not characterize behavior beyond that point. In contrast, we go further by analyzing the behavior of the min-norm/max-margin solution, which can be viewed as a fully converged solution, since it corresponds to the convergence of GD with (a small) weight decay. We prove that the min-norm or max-margin interpolator also exhibits benign overfitting, and thus our results describe not just a transient stage but also the asymptotic behavior of a trained model.

W4: We can show harmful overfitting results in the first two steps (when SNR = $o(1/\sqrt{n})$). Roughly speaking, the attention will be focused on the noisy tokens, both for clean and noisy examples, which leads to a small training error but poor generalization. We will add a remark about that in the revised version. Moreover, our experiments (specifically Figure 1) indicate that for SNR = $o(1/\sqrt{n})$, harmful overfitting persists even after many training steps.

Q1: Our GD analysis focuses on the case where there is at least one noisy sample, which occurs with high probability as long as $\eta \gg 1/n$, where $n$ is the number of training samples. When there are no noisy examples, the problem becomes significantly simpler and has already been addressed in prior work, such as the GD result of Jiang et al.. More broadly, the typical regime studied in the benign overfitting literature is when the label flipping rate $\eta$ is large enough (typically, a constant), so that noisy samples exist in the training set (with high probability).

Q2: In our work, we assume the label flipping rate is bounded by a certain constant $1/C$ (Assumptions 3.1 and 4.1, Item 4) to ensure that the label flipping rate is relatively low so the noisy samples will not outweigh clean samples in deciding the convergence direction of model parameters. This assumption is typical in the study of benign overfitting with label flipping noise (e.g. [6, 12, 13]) and we follow this line of work to study it in attention mechanisms. We agree that increasing $\eta$ will add difficulty to achieve benign overfitting, and the SNR threshold might also grow. We leave it as an intriguing topic to study in the future.

Q3: To further validate our theoretical findings, we conducted additional experiments on real-world datasets, including MNIST and CIFAR-10. In both cases, we trained a one-layer Transformer model ($d=1024$) to perform binary classification. Since the signal-to-noise ratio (SNR) is fixed for each dataset, we varied the training sample size $n$ to examine how train and test accuracy evolve with $n$. In both cases, the results indicate that while training accuracy remains near 100%, the test accuracy improves as $n$ increases. This corresponds to our SNR threshold $\Theta(1/\sqrt{n})$ that determines the transition between benign and harmful overfitting.

MNIST, $\eta=0.1$, 500 itr, 2 heads:

CIFAR-10, $\eta=0.05$, 500 itr, 4 heads:

[1] Tarzanagh et al. Max-margin token selection in the attention mechanism. NeurIPS 23.

[2] Vasudeva et al. "Implicit bias and fast convergence rates for self-attention."

[3] Johannes Von Oswald et al. Transformers learn in-context by gradient descent.

[4] Ahn et al. Transformers learn to implement preconditioned gradient descent for in-context learning.

[5] Zhang et al. Trained transformers learn linear models in-context.

[6] Sakamoto et al, I. Benign or not-benign overfitting in token selection of attention mechanism. In ICML 2025.

[7] Jiang et al. Unveil benign overfitting for transformer in vision.

[8] Cao et al. Benign overfitting in two-layer convolutional neural networks.

[9] Wang et al. Benign overfitting in binary classification of gaussian mixtures.

[10] Bartlett P.L. et al. Benign overfitting in linear regression.

[11] Shang et al. Initialization Matters: On the Benign Overfitting of Two-Layer ReLU CNN with Fully Trainable Layers

[12] Kou et al. Benign overfitting in two-layer relu convolutional neural networks

[13] Meng et al. Benign overfitting in two-layer relu convolutional neural networks for xor data

We thank the reviewer for their thoughtful feedback. Below, we respond to each concern and question in order.

Q1: Thank you for pointing this out. There is indeed a typo in the expression, although it should be $W \in \mathbb{R}^{d \times d}$. We will fix it in our revised version.

Q2:

2.1: The output of the softmax function in Eq. 1 is a matrix in $\mathbb{R}^{(T+1)\times(T+1)}$, representing interactions between all pairs of tokens. However, for classification (or regression) purposes, the model must produce a scalar as an output. To achieve this, both in theory and practice (e.g. [1], see Figure 1), it is common to introduce a classification token (CLS token or prompt), whose corresponding output is used for prediction. Thus, in Eq. 2, we focus on the row corresponding to the "classification" token, which leads directly to Equation 3.

[1] Dosovitskiy et al. "An image is worth 16x16 words: Transformers for image recognition at scale."

2.2: If q is a fixed vector (i.e. corresponds to a fixed prompt), then only $W$ and $v$ need to be trained. In this case, we follow the approach of [2], which showed that optimizing W with a fixed q is equivalent to optimizing q and removing W, as done in our paper. We note that in Section 3 of that paper, they are optimizing both v and q, exactly as we do. We agree that this approach can be viewed as a kind of simplification. We also note that our model in Eq. 3 can be obtained by fixing $\mathbf{q}=(1,0,\ldots,0)^\top$ in Eq. 2, and then the vector $\mathbf{p}$ from Eq. 3 is simply the first column of $W$ in Eq. 2.

[2] Tarzanagh et al. Max-margin token selection in attention mechanism.

Q3:

Q3.1: The $\Omega(n^2)$ dimension requirement arises from our need to ensure sufficient separation between token inner products and large dimension will help the model to overfit all training data more easily. This assumption is crucial for our theoretical analysis of token selection and margin gaps, and is also common in prior works on benign overfitting (e.g., [3, 4, 5]).

Nevertheless, we acknowledge that refining the lower bound on $d$ with respect to $n$ is an interesting direction, as explored in works like [6]. However, such relaxations are typically more tractable in linear models or ReLU networks, where feature mappings are fixed and the inductive biases are simpler. In contrast, the softmax attention mechanism involves data-dependent and dynamically evolving token weights, making the analysis more delicate. We have also done relevant experiments (see Figure 4 in the Appendix) and note that benign overfitting occurs for $d = 2n \ll n^2$, which suggests that the assumptions on $d$ in our theorems are loose. We view our current dimensionality assumption as a step toward a better understanding of benign overfitting in attention models. Investigating whether this condition can be relaxed while preserving the phenomenon would be a valuable direction for future work.

[3] Frei S., Chatterji N., Bartlett P. Benign Overfitting without Linearity: Neural Network Classifiers Trained by Gradient Descent for Noisy Linear Data.

[4] Kou Y., Chen Z., Chen Y., Gu Q. Benign overfitting in two-layer relu convolutional neural networks.

[5] Jiang J., Huang W., Zhang M., Suzuki T., Nie L. Unveil benign overfitting for transformer in vision.

[6] Karhadkar K., George E., Murray M., Montúfar G., Needell D. Benign overfitting in leaky ReLU networks with moderate input dimension.

Q3.2: While increasing $T$ introduces more tokens to interpolate, it also raises the risk of the attention mechanism focusing on high-scoring noise tokens due to softmax sensitivity. Besides, the definition of signal-to-noise ratio also needs to be redefined when the token number increases dramatically. In our current analysis, we fix $T=\Theta(1)$ to ensure analytical tractability to study the token-selection mechanism in the attention model. Extending our results to the large-$T$ regime, where a phase transition might occur, is an interesting direction for future work.

Q4: In our data model and architecture, with high probability over the training data, we can construct a pair $(v, p)$ that interpolates the training data and satisfies our constraints, and we even showed that GD reaches such a solution (Thm. 3.3).

Q5:

5.1: Item 3 in Thm. 3.3 shows that the attention mass for noisy examples is focused on the noisy tokens as expected. We therefore believe that the reviewer's concern is more about the clean examples, which indeed yield a relatively small lower bound on the softmax probability of the first token, especially when the number of tokens $T$ is large. This lower bound is actually more closely related to the SNR threshold. Our proof suggests that for clean examples, when the SNR is large (i.e. $c_\rho$ is large), the attention head $v$ has a significant component in the direction of the first (signal) token, compared to the direction of the noisy tokens. In such cases, even if the softmax probability of the signal token is not very large, the model still prioritizes the signal token. Conversely, when the SNR is smaller, we can obtain a stronger lower bound on the softmax probability of the first token for clean examples, as stated in Remark 3.4 and formally proved in the appendix.

5.2: While our current analysis of GD is limited to two steps, we believe that benign overfitting persists well beyond this point: First, our experimental results suggest that the theoretical findings extend to longer training, showing consistent behavior with our analysis even after many GD steps.

Second, we proved that min-norm/max-margin solutions also exhibit benign overfitting, when these solutions roughly correspond to the convergence of GD with (a small) weight decay. Indeed, it is natural to expect that GD with weight decay tends (after enough iterations) to a solution that minimizes the loss while preserving a small weight norm. We also refer the reviewer to lines 232-238, which provide additional motivation for analyzing min-norm/max-margin solutions.

Taken together, these points reinforce our belief that benign overfitting persists beyond the two-step regime.

Q6:

• Linear vs single-head attention: Benign overfitting was first studied in the context of linear models (Bartlett et al.[7]). However, for linear models and MLPs the data has a different structure (without tokens) and the implicit bias is better understood in these cases.

• CNN vs single-head attention: Compared with two-layer CNN with ReLU$^q$ activation (Cao et al.[8]), when the signal-to-noise ratio is small (SNR$\le 1$), our results show that the single-head attention model requires a smaller number of samples to achieve benign overfitting, which reflects the advantage of the attention mechanism. Compared with two-layer CNN with ReLU activation (Kou et al.[9]), they only analyze the case where SNR is small ($\|\mu\|/\sqrt{d} = O(1/n^{1/2})$) as they impose the lower bound of $d$ is related with signal strength $\|\mu\|$. So their results cannot apply to the case when SNR increases while our results show that attention model can generalize well without an upper bound for SNR due to its token-selection mechanism.

[7] Bartlett P.L., Long P.M., Lugosi G., Tsigler A. Benign overfitting in linear regression.

[8] Cao Y., Chen Z., Belkin M., Gu Q. Benign overfitting in two-layer convolutional neural networks

[9] Kou Y., Chen Z., Chen Y., Gu Q. Benign overfitting in two-layer relu convolutional neural networks

We thank the reviewer for their thoughtful feedback. Below we respond to each concern and question in order.

W1: See the answer for Q1.

W2: Indeed, under our setup, the attention mechanism exhibits a two-mode behavior, selecting between signal and noise tokens in clean and noisy samples respectively. This token-selection behavior is a reflection of an important inductive bias in softmax attention. It aligns with recent empirical findings in real-world models, where attention often becomes highly sparse or concentrates on a few dominant tokens([1, 2, 3]).

Besides, the real-world attention settings are highly complex and currently out of reach for comprehensive theoretical analysis. We were inspired by image data to design this simplified but still representative data model to capture essential aspects of token selection while remaining analytically tractable. Similar distributions are common in theoretical studies of the benign overfitting phenomenon (see [4, 5, 6, 7])

Finally, even in this simplified setting, the two-mode behavior poses significant technical challenges. Our analysis requires careful handling of randomness across noise tokens and a detailed understanding of how parameter norms relate to generalization in the overparameterized regime. While we agree that richer token interactions are important in practical applications, our focus here is to isolate and understand this core mechanism. We believe that this is a meaningful step toward more general theoretical frameworks, and we view extending the analysis to broader data distributions as an important direction for future work.

[1] Ataee Tarzanagh D., Li Y., Zhang X., Oymak S. Max-margin token selection in attention mechanism

[2] Jelassi S., Sander M., Li Y. Vision transformers provably learn spatial structure.

[3] Wang Z., Wei S., Hsu D., Lee J. Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannot

[4] Cao Y., Chen Z., Belkin M., Gu Q. Benign overfitting in two-layer convolutional neural networks

[5] Kou Y., Chen Z., Chen Y., Gu Q. Benign overfitting in two-layer relu convolutional neural networks

[6] Meng X., Zou D., Cao Y. Benign overfitting in two-layer relu convolutional neural networks for xor data

[7] Jiang J., Huang W., Zhang M., Suzuki T., and Nie L. Unveil benign overfitting for transformer in vision: Training dynamics, convergence, and generalization.

W3: We adopt this specific signal–noise data distribution to enable a clear and tractable theoretical analysis. This setup is motivated by image data, where inputs are composed of distinct tokens and only some of these tokens are related to the image's class label.

As the reviewer mentioned, such a distribution is typical in theoretical studies on benign overfitting in convolutional networks (where the input has several patches) and transformers (where the input has several tokens). For convolutional neural networks (CNN), we are aware of three papers that studied benign overfitting [8,9,10], and all of them consider a data distribution where there is a single signal patch and a single noise patch, and the signal is either $\pm \mu$ for some vector $\mu$ [8,9], or has an XOR pattern [10]. Moreover, in [8,10] the noise patch is exactly orthogonal to the signal patch. The prior work [11], which studied benign overfitting in transformers (albeit without label noise), considered a distribution similar to ours: one signal token and several noise tokens, where the signal is either $\mu_+$ or $\mu_{-}$ for orthogonal $\mu_+,\mu_-$, and the noise tokens are orthogonal to the signal. Hence, in our paper we follow a common setting, which is also the focus of previous works on related topics.

[8] Cao Y., Chen Z., Belkin M., Gu Q. Benign overfitting in two-layer convolutional neural networks

[9] Kou Y., Chen Z., Chen Y., Gu Q. Benign overfitting in two-layer relu convolutional neural networks

[10] Meng X., Zou D., Cao Y. Benign overfitting in two-layer relu convolutional neural networks for xor data

[11] Jiang J., Huang W., Zhang M., Suzuki T., and Nie L. Unveil benign overfitting for transformer in vision: Training dynamics, convergence, and generalization.

W4: See the Answer for Q3.

Q1: In general, there is no known characterization of the implicit bias in softmax transformers that applies to our setting, so it remains unclear whether GD converges to the max-margin solution. This is indeed a very interesting open direction. We note that min-norm/max-margin solutions roughly correspond to the convergence of GD with (a small) weight decay. Indeed, GD with weight decay tends to prefer solutions that minimize the loss while preserving a small weight norm. We also refer the reviewer to lines 232-238, which provide additional motivation for analyzing min-norm/max-margin solutions, even in settings where the implicit bias of GD does not necessarily lead to a min-norm solution (e.g., Savarese et al. [28], Ongie et al. [29], Ergen and Pilanci [30], Hanin [31], Debarre et al. [ 32], Boursier and Flammarion [33]).

While we don't know how to extend the analysis beyond a constant number of steps. We note that analyzing even two steps of GD in our setting is challenging. This analysis gives us many insights into how attention handles noisy label data, how the attention focuses on the signal token for clean samples and on the noisy tokens for noisy samples, sufficient and necessary conditions to allow benign overfitting, and the relation between different parameters. It also provides insights into how the trajectory looks after more steps of GD (see Remark 5.1, line 371), which we validate in our experiments (Figure 3 in the appendix). Our experimental results suggest that our theoretical findings extend beyond the two-step setting, demonstrating consistency with longer training and a broader range of learning rates. While a full trajectory analysis is an interesting direction, our results already highlight fundamental phenomena.

Q2: In our current setup, even with multiple noise tokens, attention mechanism consistently concentrates on a single dominant noise token in the noisy samples. This is due to the softmax amplification and randomness in token scores. Since all noise tokens are sampled independently from the same Gaussian distribution, their dot products with the learned query vector exhibit random fluctuations. As a result, one noise token typically achieves a noticeably higher score than the others, and the softmax function amplifies this gap exponentially, which leads to sharply concentrated attention on that token.

However, which specific noise token is selected does not significantly affect generalization, as long as the model avoids assigning high attention weight to the signal token in noisy samples. Therefore, in our theoretical analysis, we focus on the aggregate attention weight assigned to all noise tokens, rather than tracking the exact distribution across them. This abstraction allows for a cleaner and more interpretable characterization of benign overfitting behavior in the attention model.

Q3: First, we note that our motivation in this work is to understand benign overfitting in attention mechanisms. Thus, our focus is on analyzing the behavior of our model, rather than understanding whether the same data distribution is also learnable with other models.

Regarding the question, while the attention output is indeed a convex combination of token embeddings, our model differs fundamentally from a standard linear classifier trained on convex combinations of signal and noise tokens in several key ways. First, in our attention model, the combination weights (attention scores) are learned via a softmax over dot products, and are therefore data-dependent and nonlinear in the output. This differs from a linear classifier, where the convex combination is fixed and independent of the data. Second, the attention model does not assume that the position of the signal token is fixed. Hence, it is more flexible compared to non-attention models like linear models and MLPs.

We thank the reviewer for their thoughtful feedback. Below we respond to each concern and question in order.

W1: We adopt this specific signal–noise data distribution to enable a clear and tractable theoretical analysis. This setup is motivated by image data, where inputs are composed of distinct tokens and only some of these tokens are related to the image's class label.

Such a distribution is typical in theoretical studies on benign overfitting in convolutional networks (where the input has several patches) and transformers (where the input has several tokens). For convolutional neural networks (CNN), we are aware of three papers that studied benign overfitting [1,2,3], and all of them consider a data distribution where there is a single signal patch and a single noise patch, and the signal is either $\pm \mu$ for some vector $\mu$ [1,2], or has an XOR pattern [3]. Moreover, in [1,3] the noise patch is exactly orthogonal to the signal patch. The prior work [4], which studied benign overfitting in transformers (albeit without label noise), considered a distribution similar to ours: one signal token and several noise tokens, where the signal is either $\mu_+$ or $\mu_{-}$ for orthogonal $\mu_+,\mu_-$, and the noise tokens are orthogonal to the signal. Hence, in our paper we follow a common setting, which is also the focus of previous works on related topics.

[1] Cao Y., Chen Z., Belkin M., Gu Q. Benign overfitting in two-layer convolutional neural networks

[2] Kou Y., Chen Z., Chen Y., Gu Q. Benign overfitting in two-layer relu convolutional neural networks

[3] Meng X., Zou D., Cao Y. Benign overfitting in two-layer relu convolutional neural networks for xor data

[4] Jiang J., Huang W., Zhang M., Suzuki T., and Nie L. Unveil benign overfitting for transformer in vision: Training dynamics, convergence, and generalization.

Q1: see the Answer for W1

Q2: We chose to combine the key and query representations to simplify the analysis, which is already technically involved. We note that this simplification is common in prior works studying attention mechanisms. For example, [6,7,8,9,10] (as well as many other papers) also used a combined key-query matrix.

As the reviewer noted, Tarzanagh et al. [5] explore optimization behavior under separate key and query matrices, but their analysis focuses only on the optimization of the softmax weights, with the attention head (denoted by $v$ in both papers) held fixed. In contrast, our setup involves optimizing both the softmax weights and the attention heads, so their technique does not directly apply to our setting. We will add a comment on this issue in the final version.

[5] Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines.

[6] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in the attention mechanism.

[7] Vasudeva, Bhavya, Puneesh Deora, and Christos Thrampoulidis. "Implicit bias and fast convergence rates for self-attention."

[8] Johannes Von Oswald et al. Transformers learn in-context by gradient descent.

[9] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning.

[10] Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context.