Trained Transformer Classifiers Generalize and Exhibit Benign Overfitting In-Context

Thank you for the constructive feedback. We address specific points:

Task Complexity Compared to Wu et al.: This is a great question. Our setting is not directly comparable since in their setting the covariance matrix of the data ($E$ xx^T $$) does not depend on the underlying predictor, and so their bound in Theorem 4.1 shows a task complexity which depends explicitly on the trace of the data covariance matrix, and as this trace is larger the problem is harder (more variance, harder to learn). In the classification setting with class-conditional Gaussian data that we consider, the data covariance matrix takes the form $E$ xx^\top \= E (y\mu + z)(y\mu + z^\top) \= (R^2/d) I \+ \\Lambda, so its trace gets larger when either R gets larger (which makes learning easier) or when the cluster noise covariance $\\Lambda$ gets larger (which makes learning harder). It’s thus hard to directly compare these two settings.

Motivation of classification: Our goal is to eventually understand what’s happening in more real-world transformer models, which rely upon training on the cross-entropy loss and when the data comes from discrete tokens. We think a useful step towards this goal is to begin to understand the behavior of transformers which are trained on classification losses in settings which we understand well, like class-conditional Gaussians.

Signal-to-Noise Ratio: Different $\tilde R$ at test time is interesting because it shows transformers can generalize to more challenging settings than seen during training. This is particularly relevant for real-world applications where test conditions may be more difficult than training. We also think it’s notable that the transformer has very different behavior when $\\tilde R$ is small, medium, or large (see our new Figures 2 and 3).

We agree that understanding M vs. N is also an important distinction, and although this work provides a fairly complete picture of the dependence on M when N is large, we think it’s an interesting direction for future research to understand what happens when N is small.

Question re: easy/complex: Sorry for the lack of clarity here. By “easy” we were referring to both high SNR as well as lack of label-flipping noise, and by “complex” we meant low SNR and the presence of label-flipping noise.

Experiments: We have added new experiments to the paper, see Figures 2 and 3, which reveal the importance of $\\tilde R$ and high-dimensionality for benign overfitting, as well as the role of the number of pre-training tasks $B$. We are happy to provide our code or run further experiments if the reviewer is interested!

We appreciate the positive feedback on our work. Regarding the key concern on the data assumptions: we agree that the assumptions we make are relatively strong. However, we think these assumptions are reasonable for a theoretical work that gives a precise analysis of the effects of all of the core ingredients to the transformer training pipeline: the implicit regularization effect of gradient descent in pre-training; the number of pre-training tasks; the number of examples per pre-training task; the signal-to-noise ratio of the pre-training tasks; the number of in-context examples as test time; the signal-to-noise ratio of the test-time tasks; and the effect of label noise at test time. To our knowledge, no prior work has captured all of these aspects of pre-training/testing transformers simultaneously. Moreover, we discovered the possibility of benign overfitting in the forward pass of the transformer, which has not appeared in prior work.

We think this work paves the way for future works which we hope can develop as fine-grained of a characterization of the interplay between data, architecture, and optimizers as we did but for transformers with more complex architectures and realistic data distributions. As we mention in the discussion, we believe it should be possible to extend our analysis to non-convex linear transformer architectures which are homogeneous in the parameters, since the KKT condition analysis approach can work in that setting as well.

Regarding the question: we’ve revised our manuscript with experiments which demonstrate a setting which satisfies all of our assumptions and exhibits the type of behaviors that our theory suggests—see Figures 2 and 3 with details on the experiments in lines 377-400 and Appendix E. Namely, benign overfitting in-context occurs when the data is sufficiently high-dimensional and $\tilde R$ is in a sweet spot, and test error improves with greater pre-training tasks, especially when $\tilde R$ isn’t too large.

We thank the reviewer for their thoughtful comments. We address the main points below:

Simplified Architecture: While we analyze a 1-layer linear transformer, this simplification allows us to rigorously characterize fundamental properties of in-context learning that may extend to more complex architectures. Similar to how early theoretical works on neural networks focused on simplified architectures to establish foundational results (citation), our work provides a theoretical framework that can guide future analyses of more complex transformers. Moreover, prior works have utilized the same architecture to derive theoretical insights (Wu et al. 2024, Kim et al. 2024).

Regarding more complex architecture, as we mention in the conclusion we believe our techniques can be extended to more complex linear transformer architectures as the last-token prediction would still be homogeneous in the network parameters, enabling the usage of the implicit regularization of GD on the logistic loss for such architectures (Lyu & Li, 2020). We think an interesting but more challenging question is to extend our results to softmax attention architectures.

Numerical Validation: We agree that empirical validation with GD would be valuable, and we have added these in the revised PDF; see Figures 2 and 3 with details on the experiments in lines 377-400 and Appendix E. The experiments broadly support our theoretical findings, demonstrating that benign overfitting in-context occurs for particular settings of $\\tilde R$, $d$, $M$, and $B$.

However, like other foundational theoretical works in this area (Von Oswald et al 2022, Akyurek et al 2022, Zhang et al 2023), our focus is on establishing theoretical guarantees and understanding the principles of in-context learning. We believe our theoretical results are valuable on their own, as they:

1. Provide precise analysis of the test-time performance of a pre-trained model and how this depends on the number of pre-training tasks, the signal-to-noise ratio, the number of in-context examples, and the noise rate. To the best of our knowledge, prior works fail to capture all of these properties simultaneously, and none covered the linear classification setting.

2. Demonstrate the possibility of benign overfitting in the forward pass of transformers: we believe this is a novel setting and finding which is of interest to theorists working on understanding LLM generalization.

Pre-training with noisy labels: We think this is an interesting question. We think the theoretical analysis of pre-training with noisy labels would be significantly different. If we aimed to use the implicit regularization approach based on analyzing solutions to our equation (3), because in this setting there may not be a $U$ which satisfies the constraints of (3), and even if such a $U$ does exist then we would be investigating the behavior of a pre-trained transformer which has memorized noisy labels. In this setting it is not clear that generalization is possible, and we believe a resolution to this question would be very challenging. On the other hand, if we looked directly at the dynamics of GD, we conjecture that if GD is early-stopped before overfitting, we might expect similar behavior as to what we see for the max-margin in our paper. We have added comments on this to our discussion in lines 506-513.

Thank you for the detailed review. We address your concerns:

Readability: We appreciate the feedback about density and notation. We have added a comprehensive notation table in the appendix, see Table 1 and the added reference in line 138.

Experiments for early stopping: The reviewer raises an important point about KKT conditions not being satisfied under early stopping. Our analysis focuses on the theoretical limits of the limiting behavior of the transformer, similar to other works analyzing implicit regularization (citation). We agree that studying partial convergence would be valuable , and we have added experiments which investigate empirically what happens in early-stopped networks, which broadly shows that early-stopped networks have similar behavior to what our theory suggests (see Figures 2 and 3, lines 377-400 in the revised paper, and Appendix E for details on experiments).

Regarding stochasticity and our focus on KKT conditions: To clarify, our analysis holds for the limiting behavior of both GD and SGD with a constant step size, due to Nacson et al. 2019. In general we think there is a lot of value in investigating if generalization in neural networks can be explained solely via an analysis of the implicit regularization of the optimizer—if this is possible, then we have persuasive evidence for why this generalization behavior occurs. If it isn’t possible, then it points to other aspects of the neural net training process which we should focus on instead.