Research map / Architectures and efficiency / Attention and transformers

Benign Overfitting in Token Selection of Attention Mechanism

ICML 2025 · Architectures and efficiency › Attention and transformers · approach: Sparse and structured attention · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Theoretical analysis of one-layer softmax token selection trained by gradient descent under label noise, showing benign overfitting depends on SNR and that generalization can be delayed exponentially after fitting training data.

Problem

Understanding how attention learns to select tokens, and whether the resulting solutions generalize, when training data contain label noise (benign vs harmful overfitting).

Approach

Analyzes gradient descent on the key-query matrix and tunable token of a one-layer attention model with a fixed linear head, under a data model with relevant, weakly relevant (confusing) and irrelevant tokens plus label noise; tracks attention gaps via a function g(x)=2x+2sinh(x-log T) to show SNR-dependent regimes (not overfitting, benign overfitting) and exponentially delayed generalization.

Key results

• Theorem: with SNR^2 = omega(1/n) the model fits only clean data and generalizes (not overfitting); with SNR^2 = o(1/n) it fits all training data including noisy labels while test error is low (benign overfitting), with exponentially large training time needed for generalization.
• On real data with final attention layer of pretrained ViT/BERT trained on label-noisy data (eta=0.1, 2000 epochs), MNIST test accuracy 88.5 to 91.4% across n=20 to 1000 and CIFAR-10 about 95%, with near-zero training loss.
• AG-news shows harmful overfitting: training loss near 0 while test accuracy drops from 83.6% to 77.6% as n grows to 1000.

Measured effects reported in the paper

Surprising findings

• Generalization is delayed exponentially in SNR after training data are already fit, a grokking-like effect.
• AG-news exhibits harmful overfitting in the real-world experiments, with test accuracy decreasing as training size grows.

Limitations

• Analysis restricted to a one-layer attention model with fixed linear head, binary classification, and T=Theta(1) tokens.
• Real-world experiments train only the final attention layer on frozen features; SNR cannot be controlled directly, so the link to theory is indirect.
• Assumptions require over-parameterization and small learning rate and specific initialization scales.

Experimental setup

Datasets: MNIST CIFAR-10 PneumoniaMNIST BreastMNIST AG-news TREC

Benchmarks: MNIST CIFAR-10 PneumoniaMNIST BreastMNIST AG-news TREC

Metrics: training loss test accuracy softmax probability dynamics loss heat-map over d and signal norm

Baselines: Tarzanagh et al. max-margin attention Jiang et al. ViT benign overfitting Kim et al. attention generalization

Builds on: implicit bias of attention (max-margin token separator) benign overfitting in linear and two-layer models Li et al. ViT benign overfitting analysis

Closest papers by meaning

• Benign Overfitting in Single-Head Attention (NeurIPS 2025 · similarity 0.88)
• Hierarchical Self-Attention: Generalizing Neural Attention Mechanics to Multi-Scale Problems (NeurIPS 2025 · similarity 0.86)
• From Attention to Activation: Unraveling the Enigmas of Large Language Models (ICLR 2025 · similarity 0.86)
• Unveiling Memorization–Generalization Coexistence: A Case Study on Arithmetic Tasks with Label Noise (NeurIPS 2026 · similarity 0.85)
• Selective Attention Improves Transformer (ICLR 2025 · similarity 0.85)
• Token Statistics Transformer: Linear-Time Attention via Variational Rate Reduction (ICLR 2025 · similarity 0.85)