Understanding Generalization in Transformers: Error Bounds and Training Dynamics Under Benign and Harmful Overfitting

TL;DR

This work tackles how transformers generalize under benign and harmful overfitting when trained with labeled flip noise. It develops a generalization theory for a two-layer transformer, separating training dynamics into three stages for each regime and providing stage-specific error bounds under varying SNR. The authors also present a comprehensive experimental study that confirms the theoretical predictions, reveals a phase transition boundary governed by data size and signal-to-noise ratio, and analyzes factors such as learning rate and V-matrix initialization. The results extend benign overfitting theory from linear and CNN settings to transformers, relax certain assumptions, and offer practical insights into training dynamics and generalization in over-parameterized transformer models.

Abstract

Transformers serve as the foundational architecture for many successful large-scale models, demonstrating the ability to overfit the training data while maintaining strong generalization on unseen data, a phenomenon known as benign overfitting. However, research on how the training dynamics influence error bounds within the context of benign overfitting has been limited. This paper addresses this gap by developing a generalization theory for a two-layer transformer with labeled flip noise. Specifically, we present generalization error bounds for both benign and harmful overfitting under varying signal-to-noise ratios (SNR), where the training dynamics are categorized into three distinct stages, each with its corresponding error bounds. Additionally, we conduct extensive experiments to identify key factors that influence test errors in transformers. Our experimental results align closely with the theoretical predictions, validating our findings.

Figures (7)

• Figure 1: Training stage analysis of benign overfitting and harmful overfitting under label flipping conditions (Experimental Design Reference jiang2024unveilbenignoverfittingtransformer), (a)(b): Test loss and training loss vary over time; (c)(d): Signal atten and noise atten vary over time; (e)(f): Signal V and noise V vary over time.
• Figure 2: The test loss w.r.t. different label-flipping probability $\alpha$. Figure (a) represents the heat map drawn at a label flipping probability of $\alpha = 0.001$, Figure (b) represents the heat map drawn at a label flipping probability of $\alpha = 0.01$, Figure (c) represents the heat map drawn at a label flipping probability of $\alpha = 0.1$, and Figure (d) represents the heat map drawn at a label flipping probability of $\alpha = 0.2$(Experimental design originated from jiang2024unveilbenignoverfittingtransformer).
• Figure 3: Phase transition between benign and harmful overfitting based on Figure \ref{['heatmap']}, Map the parts with benign overfitting and the parts without benign overfitting to opposite colors.
• Figure 4: (a) shows the variation of $C$ value with $\alpha$, while (b) calculates the similarity between the two images $\alpha=0.1$ and $\alpha=0.001$, with higher scores indicating higher similarity.
• Figure 5: The test loss w.r.t. different leaning rates.

Theorems & Definitions (15)

• Definition 3.1: Data Generation Model
• Theorem 4.2: Benign overfitting in transformers
• Theorem 4.3: Harmful overfitting in transformers
• Definition 5.1: Splitting of V vector
• Lemma 5.2: The test loss of benign overfitting
• Lemma 5.3: Update rules of $\mathbf{W}_V$
• Theorem 2.1
• Lemma 2.2: The test loss of benign overfitting in Lemma \ref{['lem_the_test_loss_benign_overfitting']}
• Lemma 2.3: Relationship of constants
• Theorem 4.1: (First part of Theorem \ref{['thm:4.1']})