Title: Benign Overfitting in Single-Head Attention
Venue: NeurIPS 2025
Accepted: yes
Keywords: benign overfitting

single-head attention

transformers
Abstract: The phenomenon of benign overfitting, where a trained neural network perfectly fits noisy training data but still achieves near-optimal test performance, has been extensively studied in recent years for linear models and fully-connected/convolutional networks. In this work, we study benign overfitting in a single-head softmax attention model, which is the fundamental building block of Transformers. We prove that under appropriate conditions, the model exhibits benign overfitting in a classification setting already after two steps of gradient descent. Moreover, we show conditions where a minimum-norm/maximum-margin interpolator exhibits benign overfitting. We study how the overfitting behavior depends on the signal-to-noise ratio (SNR) of the data distribution, namely, the ratio between norms of signal and noise tokens, and prove that a sufficiently large SNR is both necessary and sufficient for benign overfitting.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the benign overfitting phenomenon in single head attention. Specifically, when signal-to-noise ratio (SNR) is larger than $\\Omega(1/\\sqrt{n})$, the test loss will be small after training.\nTo show that the requirement on SNR is tight, this paper provides a harmful overfitting result when SNR is smaller than $O(1/\\sqrt{n})$ in the max-margin scenario.\nThe experimental results verifies the theoretical results.\n\nSTRENGTHS AND WEAKNESSES:\nstrengths:\n\n1. This paper demonstrates the importance of SNR for generalisation through rigorous theoretical analysis, and characterizes the boundary between benign overfitting and harmful overfitting.\n\n2. The theoretical result are verified by the experiments (especially the heatmap).\n\n3. The writing is clear and easy to follow, all the assumptions and theorem are well structured.\n\nweaknesses:\n\n1. The assumptions are oversimplified: The original attention mechanism consists of the QKV matrices, but this paper reduces them to two vectors, making it misaligned with reality.\n\n2. Lack of training dynamics: This paper only consider two step of GD in Section 3. It is better to consider a smaller step size to characterize the dynamics of $v$ and $p$.\n\n3. The benign overfitting phenomenon has been studied in [1] and [2], thus this paper seems not novel.\n\n4. In terms of gradient descent (Section 3), no harmful overfitting results are provided in this paper.\n\nreferences\n\n[1] Jiang, J., Huang, W., Zhang, M., Suzuki, T., & Nie, L. Unveil benign overfitting for transformer in vision: Training dynamics, convergence, and generalization. In NeurIPS 2024.\n\n[2] Sakamoto, K., & Sato, I. Benign or not-benign overfitting in token selection of attention mechanism. In ICML 2025.\n\nQUESTIONS:\nQ1: Line 172, this paper requires $d \\ge C \\cdot n^2 \\log (n / \\delta) / \\eta^2$, then as $\\eta \\rightarrow 0$, we have $d \\rightarrow \\infty$. Is the method in this paper able to analyze the situation when $\\eta = 0$ (no lable flipping) ?\n\n\nQ2: The condition of benign overfitting is determined by SNR and n. Would it be related to label flipping rate $\\eta$ ? \n\nIn my intuition, the $\\Omega(1/\\sqrt{n})$ requirement may become more specific if $\\eta$ is considered.\n\n\nQ3: This paper only validate the theoretical results on sythetic datasets.It is recommended that the authors consider more experiments on real datasets, such as MNIST and CIFAR10.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies benign overfitting in a model for single head softmax attention. The attention layer is adapted to perform binary classification by adopting a vector valued value matrix and considering its sign. The data consists of $n$ points, or inputs, where input consists of $T$ vectors or tokens: one of these tokens correlates with an output label paired with the input while the rest are gaussian noise. The attention layer is trained by tuning the output value vector along with a prompt vector inside the softmax function and benign overfitting is established in a certain regime and under a variety of training processes.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper is well written and of high quality. The assumptions made are clearly stated and remarks are provided throughout to add interpretation and highlight key aspects of the work, as such I rate the clarity also highly. The technical contributions are the study of benign overfitting in the context of an attention layer, which to the best of my knowledge is new, but under pretty standard data assumptions. Therefore, and although not exceptional, I think the originality of the paper is good. Overall I think this is a good technically solid paper. If I had to identify the weakest dimension of the paper I might say significance: in particular, the paper seems to adapt single head attention to binary classification in a manner which feels a little contrived. The wider implications of the results also do not seem particularly well discussed.\n\nQUESTIONS:\n- Is there a typo where $W \\in R^{T \\times d}$ should be  $W \\in R^{d \\times T}$\n- On the model:\n\n     i) In equation 2 why do we suddenly ignore / remove the $S(XWX^T)$ term? \n\n     ii) It seems a little odd to me that on the one hand you argue that $W$ and $q$ are fixed relative to one another and therefore you can simplify $Wq$ to $q$, but on the other you train $v$ and $q$ together. If $q$ is a prompt while $W$ and $v$ are parameters then does it not seem more relevant perhaps to consider them trained via separate processes?\n\n- Assumptions\n\n   i) The required dimension of the tokens seems very large, i.e., $d$ needs to scale like $n^2$. What is the reason for this in this setting and could it be relaxed? For example, in some of the works you cite for linear classifiers and ReLU networks then analogous results can be proved for $d$ scaling like $n$.\n\n  ii) What breaks down if the number of tokens is not bounded by a constant? Intuitively a larger $T$ feels like it would help one achieve benign overfitting?\n\n- General comments: for any data set you might sample is it always true that a solution pair $(v,q)$ exist that solve the constraints of your optimization problem?\n\n- Theorem 3.3\n\n   i) Item 3 confuses me a bit: shouldn't for clean data the result imply that the attention mass is focused on the first token while for the noisy data points the mass on the first token should be very small and the mass distributed amongst the other near orthogonal tokens? For large $T$ I don't see how the result states this. Maybe this is why you only allow for a few tokens?\n\n  ii) What happens if you keep training for more steps?\n\n\n- Can you comment on any differences between benign overfitting across the different models people have studied, i.e., linear, shallow feedforward, CNN and single head attention? In particular, are some models in some sense more or less biased towards the benign overfitting in some or any sense you can identify?\n\nLIMITATIONS:\nNot applicable.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates benign overfitting in a single-head softmax attention model trained with noisy labels. The authors prove that when the signal-to-noise ratio is $\\Theta(1/\\sqrt{n})$, the model can interpolate noisy training data while generalizing well to test data. This is shown both for early-stage gradient descent after two iterations and for max-margin solutions. They provide tight conditions on the SNR for benign versus harmful overfitting and support the theory with synthetic experiments.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The paper addresses benign overfitting in attention-based models, a relatively underexplored setting compared to linear regression, kernel models, and CNNs. This aligns with growing interest in understanding generalization in Transformers.\n\nWeaknesses:\n1. The theoretical results for gradient descent stop after just two iterations, and do not analyze convergence behavior or connection to the max-margin solution. As a result, the relationship between early dynamics and implicit bias remains unclear.\n2. Although each input has $n$ tokens, only one is a signal token and the rest are orthogonal noise. In clean samples, attention rapidly concentrates on the signal; in noisy samples, it shifts to a single dominant noise token. This dynamic resembles a two-mode selection process, so the attention mechanism effectively reduces to selecting between signal and noise, limiting the richness of attention interactions typically seen in real-world applications.\n3. Signal and noise vectors are assumed orthogonal, noise tokens are Gaussian, and all vectors have fixed norms. These stylized assumptions are typical in theoretical work on overparameterized models but limit practical relevance.\n4. While the authors emphasize that softmax attention leads to token selection, it remains somewhat ambiguous how this differs fundamentally from a model that linearly interpolates signal and noise under logistic loss. In the current setup, the attention head largely acts as a soft selector, and it’s unclear what inductive bias is unique to attention beyond this selection.\n\nQUESTIONS:\n1. Can you extend the gradient descent analysis beyond two iterations? Does it eventually converge to the max-margin solution, or does it behave differently?\n2. In the multi-token setting (e.g., multiple noise tokens), how robust is the token selection pattern? Does the model consistently attend to one dominant noise token for noisy samples, or does attention fragment across them?\n3. What fundamentally distinguishes the attention model in your setup from a standard linear classifier trained on convex combinations of signal and noise tokens?\n\nLIMITATIONS:\nThis paper offers a rigorous and technically interesting analysis of benign overfitting in single-head softmax attention models under noise. However, the theoretical results are limited to early-stage training and rely on strong simplifying assumptions about both data and architecture. While the paper improves our theoretical understanding, its insights into realistic attention models remain preliminary. Stronger justification of the attention mechanism's unique role, extended training dynamics, or theoretical treatment of more general token structures would significantly enhance the impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors study benign overfitting for a single softmax attention model. Using a specific data distribution (Section 2.1), a simplified attention model (Section 2.2)  and specific assumptions (Assumptions 3.1) which include a given signal-to-noise ratio, they show that benign overfitting occurs at the 2nd iteration of gradient descent. This is the main result and verified empirically for synthetic data from the specified data distribution. Complementary results for a max-margin learning process are developed and proven.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- This seems to be a solid theoretical analysis of benign overfitting in the studied context.\n- New perspectives on the conditions required for benign overfitting to occur for a single attention head.\n- Empirical results confirming theoretical predictions for the specific setting: both in the main paper, and the additional results in the appendix, Section A.4.\n\nWeaknesses:\n- In addition to the assumptions listed in 3.1. which can be considered ‘appropriate conditions’, the specific data distribution studied is very specific: two signals orthogonal, and noise also orthogonal with the signal. While this is a valid setting to investigate, it is not clear from the paper how this restricts the applicability of the analysis, or how significant the findings are. The study context itself can be better motivated.\n- 74 pages of supplementary material is excessive – very little of this was reviewed.\n- I struggled to follow the proofs, and was unable to do so in the time available.\n\nMinor comments:\n- Fig 2: check top/bottom vs left/right panels; ‘exsists’ line 487\n- Even if obvious, please consider adding the longer form of abbreviations on first use (GD, SVM etc).\n\nQUESTIONS:\n- To what extent does the studied data distribution restrict the value and significance of the findings?\n- The simplified attention model includes a combined key-query. Given that having these separate or combined displayed different optimization characteristics in [5], is this choice worth commenting on in the context of the current paper?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}