ICML 2025

Benign Overfitting in Token Selection of Attention Mechanism
Keitaro Sakamoto 1 Issei Sato 1
Abstract
Attention mechanism is a fundamental component
of the transformer model and plays a 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 dynamics 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
an initial phase of overfitting. Finally, we provide
experiments to support our theoretical analysis
using both synthetic and real-world datasets.
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 variable-length input sequences
and flexibly select important tokens based on the input. The
seminal work on the training dynamics of this token selec-
tion mechanism (Tarzanagh et al., 2023a;b; Vasudeva et al.,
2024; Sheen et al., 2024) studied the implicit bias of gradi-
ent descent to the max-margin token separator. However,
1Department of Computer Science, The University of
Tokyo, Tokyo, Japan. Correspondence to: Keitaro...