Trained Transformer Classifiers Generalize and Exhibit Benign Overfitting In-Context

## Abstract

Abstract Transformers have the capacity to act as supervised learning algorithms: by properly encoding a set of labeled training (“in-context”) examples and an unlabeled test example into an input sequence of vectors of the same dimension, the forward pass of the transformer can produce predictions for that unlabeled test example. A line of recent work has shown that when linear transformers are pre-trained on random instances for linear regression tasks, these trained transformers make predictions using an algorithm similar to that of ordinary least squares. In this work, we investigate the behavior of linear transformers trained on random linear classification tasks. Via an analysis of the implicit regularization of gradient descent, we characterize how many pre-training tasks and in-context examples are needed for the trained transformer to generalize well at test-time. We further show that in some settings, these trained transformers can exhibit “benign overfitting in-context”: when in-context examples are corrupted by label flipping noise, the transformer memorizes all of its in-context examples (including those with noisy labels) yet still generalizes near-optimally for clean test examples.

## 1 Introduction

A key feature of transformer-based large language models (LLMs) is their ability to perform in-context learning: by providing a few labeled examples to a trained transformer with an unlabeled example for which the user wants a prediction, LLMs can formulate accurate predictions without any updates to their parameters \citepbrown2020language.
This ability to perform in-context learning is influenced by the interplay between tokenization of inputs (i.e. how to feed data into the transformer), pretraining datasets, optimization algorithms used for pre-training, the particular architecture of the transformer, as well as the properties of the in-context examples that are provided at test time.

A number of recent works have sought to develop a deeper theoretical understanding of in-context learning by investigating transformers trained on classical supervised learning tasks. This line of work was initiated by the experiments of \citetgarg2022can, who showed that when transformers are trained on random instances of supervised learning problems, they learn to implement supervised learning algorithms at test-time: e.g., when training GPT2 architectures \citepradford2019language on random linear regression tasks with weights sampled from a Gaussian prior, the transformer learns to implement an algorithm similar to ordinary least squares. This led to a series of follow-up works which analyzed the dynamics of gradient descent/flow over simplified linear transformer architectures when trained on linear regression tasks \citepvon2022transformers,akyurek2022learning, some providing guarantees for how many tasks or samples were needed to generalize well \citepzhang2024icl,ahn2023transformers,wu2024pretraining.

Figure: Figure 1: Benign overfitting in-context: after pre-training (not shown), when given a sequence of in-context examples, the transformer memorizes noisy labels yet still generalizes well to unseen test examples.
Refer to caption: x1.png

In this work, we analyze the behavior of linear transformer architectures which are trained by gradient descent on the logistic or exponential loss over random linear classification tasks. We consider a restricted linear attention model, a setting considered in prior works \citepwu2024pretraining,kim2024transformersminimaxoptimalnonparametric.
We assume that pre-training tasks are sampled from random instances of class-conditional Gaussian mixture model data, i.e. for some $R>0$, covariance matrix $\Lambda$ and for each task $\tau=1,\dots,B$ we have

$$ $\mu_{\tau}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathsf{Unif}(R\cdot \mathbb{S}^{d-1}),\quad y_{\tau,i}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{ \sim}}\mathsf{Unif}(\{\pm 1\}),\quad z_{\tau,i}\stackrel{{\scriptstyle\mathrm{ i.i.d.}}}{{\sim}}\mathsf{N}(0,\Lambda),\qquad x_{\tau,i}=y_{\tau,i}\mu_{\tau}+ z_{\tau,i}.$ $$

We assume that test-time in-context examples also come from class-conditional Gaussian mixture models as above but with two important differences. First, the signal-to-noise (determined by $R$) at test time can differ from those seen during pre-training. Second, we allow for label-flipping noise to be present at test time, i.e. we flip labels $y_{i}\mapsto-y_{i}$ with probability $p$.

Our first contribution is an analysis of how many pre-training tasks are needed in order to generalize well at test-time. We find that after pre-training the transformer can tolerate smaller SNR than those tasks which it was trained on. Moreover, we find that at test time the transformer can tolerate label-flipping noise, even though the pre-training data does not have label-flipping noise. Thus, even when pre-training on simple and easy-to-learn datasets, the transformer can generalize on more complex tasks. Our proof follows by an explicit analysis of the implicit regularization due to gradient descent during pre-training. To our knowledge, this is the first theoretical result on in-context learning for linear classification.

Our second contribution is the finding that the trained transformer can exhibit “benign overfitting in-context”: namely, if the test-time sequence is subject to label-flipping noise, then in some settings the transformer memorizes the in-context training data yet still generalizes near-optimally. The precise setting where this occurs requires the features to lie in a high-dimensional space and for the SNR to be relatively small, a set of conditions observed in prior works on benign overfitting in classification tasks \citepchatterji2020linearnoise,frei2022benign,frei2023benignkkt. Figure 1 illustrates what this phenomenon would look like in the language setting. To the best of our knowledge, no prior work demonstrated that transformers could exhibit benign overfitting.

## 2 Related Work

#### In-context learning for supervised learning tasks.

Following the initial experiments of \citetgarg2022can, a number of works sought to understand what types of algorithms are implemented by transformers when trained on supervised learning tasks. \citetakyurek2022learning and \citetvon2022transformers used approximation-theoretic and mechanistic approaches to understand how a variety of transformer architectures could implement linear regression algorithms. \citetbai2024transformers showed that transformers could implement a variety of classical algorithms used for different supervised learning tasks, e.g. ridge regression, Lasso, and gradient descent on two-layer networks. \citetzhang2024icl, \citetahn2023transformers and \citetmahankali2023stepgradientdescentprovably examined the landscape of single-layer linear transformers trained on linear regression tasks, with \citetzhang2024icl additionally developing guarantees for convergence of (non-convex) gradient flow dynamics. There are other works which examine training of transformers on data coming from hidden Markov chains \citepedelman2024evolutionstatisticalinductionheads or nonparametric function classes \citepkim2024transformersminimaxoptimalnonparametric. The most closely related work to this one is \citetwu2024pretraining, who focused on a convex linear transformer architecture, as we do in this work, trained on linear regression tasks with SGD. They provided task and (in-context) sample complexity guarantees when training by SGD. In contrast, we develop guarantees for the classification setting, and we develop guarantees via an analysis of the implicit regularization of GD.

#### Implicit regularization in transformers.

The starting point for our analysis of the task complexity of pre-training and the sample complexity of in-context learning is an analysis of the implicit regularization of the optimization algorithm. We outline the most related works here and refer the reader to the survey by \citetvardi2022implicit. The convex linear transformer architecture we study is linear in (the vectorization of) its parameters, which by \citetsoudry2018implicit implies that gradient descent converges in direction to the max-margin solution. The more general class of linear transformer architectures are non-convex in general, but certain subclasses of them are homogeneous in their parameters and hence converge (in direction) to KKT points of max-margin solutions \citeplyu2019gradient,ji2020directional. Another line of work seeks to understand the implicit bias of GD for softmax-based transformers \citeptarzanagh2023max,tarzanagh2023transformers,thrampoulidis2024implicitbiasnexttokenprediction,vasudeva2024implicit, typically with more stringent assumptions on the structure of the training data.

#### Benign overfitting.

The ability for neural networks to overfit to noise yet still generalize well \citepzhang2016understanding led to a significant line of work on reconciling this phenomenon with classical learning theory. We focus on the most directly relevant works here and point the reader to the surveys \citetbartlett2021statisticalviewpoint,belkin2021survey for a more extensive review. Prior works have shown that kernel-based methods \citepbelkin2018doesdatainterpolationcontradict and the ordinary least squares solution \citepbartlett2020benignpnas can exhibit benign overfitting in regression tasks, and that the max-margin solution can in classification tasks \citepchatterji2020linearnoise,frei2023benignkkt. Our proof technique comes from the works \citetfrei2023implicit,frei2023benignkkt, where it was shown that the max-margin solution over 2-layer neural nets behaves similarly to a nearly uniform average of the training data and that this average can exhibit benign overfitting in class-conditional Gaussian mixture models. We similarly show that the forward pass of a pre-trained transformer behaves similarly to an average over the in-context training examples, and use this to show benign overfitting.

## 3 Preliminaries

In this section we introduce the pre-training distribution, the particular transformer architecture we consider, and the assumptions on the pre-training data.

### 3.1 Notation

We denote matrices with capital letters $W$ and vectors and scalars in lowercase. The Frobenius norm of a matrix is denoted $\|W\|_{F}$, and its spectral norm as $\|W\|_{2}$. We use $a\wedge b:=\min(a,b)$ and $a\vee b=\max(a,b)$. We use the standard $O(\cdot),\Omega(\cdot)$, and $\Theta(\cdot)$ notations, where $\tilde{O},\tilde{\Omega},\tilde{\Theta}$ ignore logarithmic factors. We use $\mathbbm{1}(x)$ as the indicator function, which is 1 if $x$ is true and 0 otherwise. We denote $[M]=\{1,2,\dots,M\}$ for a positive integer $M$.

### 3.2 Setting

We consider classification tasks where the labels $y_{i}\in\{\pm 1\}$ and features $x\in\mathbb{R}^{d}$. We assume the pretraining data satisfies the following throughout the paper.

###### Assumption 3.1 .

Let $\Lambda\succ 0$ be a symmetric positive-definite $d\times d$ matrix and let $B,N\geq 1$ and $R>0$. Let $\tau=1,\dots B$ and $i=1,\dots,{N+1}$.

- 1.
Let $\mu_{\tau}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathsf{Unif}(R\cdot
\mathbb{S}^{d-1})$ be uniform on the sphere of radius $R$ in $d$ dimensions.
- 2.
Let $z_{\tau,i}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathsf{N}(0,\Lambda)$, and $y_{\tau,i}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathsf{Unif}(\{\pm 1\})$, where $\mu_{\tau}$, $z_{\tau,i}$ and $y_{\tau,i}$ are mutually independent.
- 3.
Set $x_{\tau,i}:=y_{\tau,i}\mu_{\tau}+z_{\tau,i}$.

We therefore assume each task is a random instance of a class-conditional Gaussian mixture model, where the signal-to-noise ratio is governed by how large $R$ is relative to $d$: when $R$ is larger, learning is easier.

We tokenize samples to be fed into the transformer in the following way: for a sequence of $N$ labeled examples $(x_{i},y_{i})_{i=1}^{N}$ which we use to formulate a prediction for an unlabeled test example $x_{N+1}$, we concatenate the $(x_{i},y_{i})$ pairs into a $d+1$ dimensional vector and tokenize the test example as $(x_{N+1},0)$,

$$ $E=\begin{pmatrix}x_{1}&x_{2}&\cdots&x_{N}&x_{N+1}\\ y_{1}&y_{2}&\cdots&y_{N}&0\end{pmatrix}\in\mathbb{R}^{(d+1)\times(N+1)}.$ (3.1) $$

In the standard softmax-based attention with a single head \citepvaswani2017attention, the transformer is defined in terms of parameters $W^{V}\in\mathbb{R}^{d_{e}\times d_{e}}$, $W^{K},W^{Q}\in\mathbb{R}^{d_{k}\times d_{e}}$. If $\theta=(W^{K},W^{Q},W^{V})$ then softmax attention computes

$$ $f(E;\theta)=E+W^{V}E\cdot\mathrm{softmax}\left(\frac{(W^{K}E)^{\top}W^{Q}E}{ \rho_{N,d_{e}}}\right),$ $$

where $\rho_{N,d_{e}}$ is a normalization factor which is not optimized (but which may depend on $N$ and $d_{e}$). In linear transformers, the softmax is replaced with the identity function. Following prior work \citepvon2022transformers,zhang2024icl,ahn2023transformers, we merge the $K$ and $Q$ matrices into $W^{KQ}:=(W^{K})^{\top}W^{Q}$ and we consider an objective function where the aim is to use the first $N$ columns of (3.1) to formulate predictions for $x_{N+1}$, whereby the bottom-right corner of the output matrix of $f(E;\theta)$ serves as this prediction. Writing $W^{\Delta}=\begin{pmatrix}W_{11}^{\Delta}&w_{12}^{\Delta}\\
(w_{21}^{\Delta})^{\top}&w_{22}^{\Delta},\end{pmatrix}$ for $\Delta\in\{V,KQ\}$, for the linear transformer architecture, this results in the prediction

$$ $\widehat{y}(E;\theta)=\begin{pmatrix}(w_{21}^{V})^{\top}&w_{22}^{V}\end{ pmatrix}\cdot\frac{1}{N}EE^{\top}\cdot\begin{pmatrix}W_{11}^{KQ}\\ (w_{21}^{KQ})^{\top}\end{pmatrix}x_{N+1}.$ $$

Due to the product of matrices appearing above, the resulting objective function is non-convex, which makes the analysis of its training dynamics complex. In this work, we follow \citetwu2024pretraining,kim2024transformersminimaxoptimalnonparametric and instead consider a convex parameterization of the linear transformer which results from taking $w_{21}^{KQ}=w_{21}^{V}=0$ and setting $w_{22}^{V}=1$. This results in the following prediction for the label of $x_{N+1}$,

$$ $\widehat{y}(E;W)=\left(\frac{1}{N}\sum_{i=1}^{N}y_{i}x_{i}\right)^{\top}Wx_{N+ 1}.$ (3.2) $$

That is, $\widehat{y}(E;W)$ corresponds to the predictions coming from a linear predictor trained by one step of gradient descent on the logistic or squared loss (initialized at 0) but where the parameter $W$ is a (learned) matrix preconditioner. Prior work by \citetzhang2024icl showed that when a single-layer linear transformer architecture is trained by gradient flow on linear regression tasks, it learns a function of the same type, where the preconditioner
was found to correspond to the inverse covariance matrix of the pretraining data.

For each task $\tau$, denote the embedding matrix $E_{\tau}$ as the one formed using labeled examples $(x_{\tau,i},y_{\tau_{i}})_{i=1}^{N}$ from Assumption 3.1 as per (3.1). We consider linear transformers which are trained on the prediction of the last token using the logistic or exponential loss, namely

$$ $\widehat{L}(W):=\frac{1}{B}\sum_{\tau=1}^{B}\ell\big{(}y_{\tau,N+1}\cdot \widehat{y}(E_{\tau};W)\big{)},\qquad\ell\in\{q\mapsto\log(1+\exp(-q)),\,q \mapsto\exp(-q)\}.$ $$

We are interested in the behavior of gradient descent on this objective,

$$ $W_{t+1}=W_{t}-\alpha\nabla\widehat{L}(W_{t}).$ $$

Note that $\widehat{y}(E;W)$ is homogeneous in $W$ (linear in $\mathrm{vec}(W)$) and so by \citetsoudry2018implicit we know that gradient descent has an implicit bias towards maximum-margin solutions, as we recall in the following theorem:

###### Theorem 3.2 ( \citet soudry2018implicit) .

Define the $\ell_{2}$-max margin solution as

$$ $W_{\mathsf{MM}}:=\mathop{\mathrm{argmin}}\left\{\|U\|_{F}^{2}:\left(\nicefrac{ {1}}{{N}}\textstyle\sum_{i=1}^{N}y_{\tau,i}x_{\tau,i}\right)^{\top}Uy_{\tau,N+ 1}x_{\tau,N+1}\geq 1,\,\forall\tau=1,\dots,B\right\}.$ (3.3) $$

If there exists $U$ which satisfies the constraints in (3.3), then provided $\alpha$ is sufficiently small, then $W_{t}$ converges in direction to $W_{\mathsf{MM}}$, i.e. for some constant $c>0$ we have $\nicefrac{{W_{t}}}{{\|W_{t}\|}}\to cW_{\mathsf{MM}}$.

This theorem shows that the max-margin solution (3.3) characterizes the limiting behavior of gradient descent. In the remainder, we will focus on the max-margin solution and derive all of the generalization guarantees as a consequence of the properties of the max-margin solution.

We next introduce the assumptions we need to ensure that the max-margin solution is well-behaved enough that we can later derive in-context learning guarantees. We fix a probability threshold $\delta\in(0,1)$, and we shall show that all of our results hold with probability at least $1-\delta$ provided the absolute constant $C>1$ appearing below is a large enough universal constant (independent of $d$, $B$, $\Lambda$, and $\delta)$.

- (A1)
$R$ is large enough so that
$\displaystyle R^{2}\geq C^{2}\sqrt{d\operatorname*{tr}(\Lambda^{2})}\vee C^{2}
\left(\frac{\operatorname*{tr}(\Lambda)}{d}\vee\sqrt{\frac{\operatorname*{tr}(
\Lambda^{2})}{N}}\vee\|\Lambda\|_{2}\right)\log(2B/\delta)$
- (A2)
For some $c_{B}>0$, we have $B\geq c_{B}d$, and
- (A3)
$d\geq C\log^{4}(2B^{2}/\delta)$.

The first assumption (A1) guarantees that the signal-to-noise ratio is sufficiently large (recall that $\Lambda$ is the covariance of $z_{\tau,i}$ in the identity $x_{\tau,i}=y_{\tau,i}\mu_{\tau}+z_{\tau,i}$); provided $B$ is polynomial in $d$, this assumptions is satisfied for $\Lambda=I$ when $R\gg\sqrt{d}$. The second assumption (A2) specifies how many pre-training sequences we (pre-)train on; we shall see later on that the generalization error achieved for in-context examples is determined in part by how small $c_{B}\wedge 1$ is ($c_{B}\wedge 1$ larger ensures better generalization). Note that we allow $c_{B}$ to be non-constant (so that $o_{d}(d)$ tasks are permitted). The assumption (A3) is needed for a technical reason, namely to ensure certain concentration bounds regarding sub-exponential random variables. Finally, note the near-independence of these assumptions on the number of examples $N$ per task during training: in many situations (e.g., $\Lambda=I$), we only require a single demonstration (i.e., $N=1)$ of the form $(x_{\tau,1},y_{\tau,1})$.

Finally, we introduce the distribution on the test-time in-context examples. We allow for a more general distribution at test-time than the one used during pre-training, where the in-context examples can have label-flipping noise and a potentially smaller cluster mean, corresponding to a potentially smaller signal-to-noise ratio. We assume, for $i=1,\dots,M$,

- 1.
$\tilde{y}_{i}\sim\mathsf{Unif}(\{\pm 1\})$.
- 2.
$y_{i}=\tilde{y}_{i}$ w.p. $1-p$ and $y_{i}=-\tilde{y}_{i}$ w.p. $p$ (label flipping noise).
- 3.
$\mu\sim\tilde{R}\cdot\mathsf{Unif}(\mathbb{S}^{d-1})$.
- 4.
$z_{i}\sim\mathsf{N}(0,\Lambda)$ and $x_{i}=\tilde{y}_{i}\mu+z_{i}$.

We assume the random variables $\tilde{y}_{i},\mu,z_{i}$ are all mutually independent. We allow for the size of the cluster mean $\tilde{R}$ at test-time to potentially be different than the size of the cluster mean $R$ from the pretraining data, and indeed we shall see in Theorem 4.1 below that we can permit a smaller $\tilde{R}$ than what we require for $R$. Also note that we allow for a potentially different sequence length at test-time ($M$) than was observed during training ($N$); prior work in the linear regression setting showed that the transformer’s behavior can depend quite differently on $M$ vs. $N$ \citepzhang2024icl.

###### Assumption 3.1 .

###### Theorem 3.2 ( \citet soudry2018implicit) .

## 4 Main Results

### 4.1 Generalization

Recall that the transformer with parameters $W$ makes predictions by embedding the set $\{(x_{i},y_{i})\}_{i=1}^{M}\cup\{(x_{M+1},0)\}$ into a matrix $E$ and then using $\mathrm{sign}(\widehat{y}(E;W))$ as our prediction for $y_{M+1}$.
Our goal is to understand the expected risk of the max-margin solution, i.e. the probability of misclassification of the test example $(x_{M+1},y_{M+1})$: using (3.2),

$$ $\displaystyle R(W_{\mathsf{MM}})$ $\displaystyle=\mathbb{P}_{(x_{i},y_{i})_{1}^{M+1},\ \mu}(\mathrm{sign}( \widehat{y}(E;W_{\mathsf{MM}}))\neq y_{M+1})$ $\displaystyle=\mathbb{P}\left(\left[\frac{1}{M}\sum_{i=1}^{M}y_{i}x_{i}\right] ^{\top}W_{\mathsf{MM}}y_{M+1}x_{M+1}<0\right).$ (4.1) $$

Our main result regarding generalization is the following theorem.

###### Theorem 4.1 .

Let $\delta\in(0,1)$ be arbitrary. Suppose that $p=1/2-c_{p}$ where $c_{p}\in(0,1/2]$ is an absolute constant. There are absolute constants $C>1,c>0$ depending only on $c_{p}$ such that if assumptions (A1), (A2), and (A3) hold, then with probability at least $1-10\delta$ over the draws of $\{\mu_{\tau},(x_{\tau,i},y_{\tau,i})_{i=1}^{N+1}\}_{\tau=1}^{B}$, when sampling a new task $\{\mu,(x_{i},y_{i})_{i=1}^{M+1}\}$, the max-margin solution (3.3) satisfies

$$ $\displaystyle\qquad\mathbb{P}_{(x_{i},y_{i})_{i=1}^{M+1},\ \mu}\big{(}\mathrm{ sign}(\widehat{y}(E;W_{\mathsf{MM}}))\neq y_{M+1}\big{)}$ $\displaystyle\leq\left(p+2\exp(-cM)\right)\mathbbm{1}(p>0)+2\exp\left(-c\rho \sqrt{d}\right)+4\exp\left(-\frac{c\rho\tilde{R}}{\|\Lambda^{1/2}\|_{2}}\right )+2\exp\left(-\frac{c\rho M^{1/2}\tilde{R}^{2}}{\|\Lambda\|_{2}\sqrt{d}}\right),$ $$

where $\rho:=\frac{c_{B}\wedge 1}{\log^{2}(2B^{2}/\delta)}$.

Let us make a few observations on the above theorem. For simplicity let us assume that the covariance matrix $\Lambda=I_{d}$, so that $\|\Lambda^{1/2}\|=\|\Lambda\|=1$, and let us assume we want the results to hold with a fixed probability threshold of $\delta=0.001$. Then assumptions (A1) through (A3) are satisfied for any $N\geq 1$ examples per task so long as $R=\|\mu_{\tau}\|\gg C\sqrt{d}$. For the test error, there are a number of different regimes which require different considerations.

- •
Due to the dependence on $\rho$, the test error degrades when $c_{B}$ gets smaller, and we require $c_{B}=\tilde{\Omega}(d^{-1/2})$ to achieve a non-vacuous generalization bound, i.e. $B=\tilde{\Omega}(d^{1/2})$ pre-training tasks suffices if $\tilde{R}$ and $M$ are large enough.(^1^11We also require $B=O(\mathrm{poly}(d))$, this comes from using a union bound over the pre-training tasks to ensure the max-margin solution is well-behaved.)
- •
While we require pre-training signal $R=\Omega(\sqrt{d})$, to achieve test error near the noise rate at test-time the signal $\tilde{R}$ can be as small as $\tilde{R}=\tilde{\Omega}(1)$ provided $M$ is large enough and $\rho=\tilde{\Omega}(1)$. This means the signal-to-noise ratio at test-time can be significantly smaller than what was observed during pre-training.
- •
When $p=0$, there is no label-flipping noise. Here it is possible to have as few as one example per task $(M=1)$ provided $\tilde{R}=\tilde{\Omega}(d^{1/4})$, and $\rho=\tilde{\Omega}(1)$.
- •
When $p>0$, there is label-flipping noise. The same analysis holds as in the $p=0$ setting except that we always require $M=\tilde{\Omega}(1)$, which makes intuitive sense (with label noise, one must see more than one example to guarantee generalization w.h.p.). Notably, there was no label noise during pre-training, yet the transformer can generalize under label noise at test time.

###### Theorem 4.1 .

### 4.2 Benign Overfitting

In this subsection we shall assume that the covariance matrix $\Lambda=I_{d}$. This is needed for technical reasons of the proof. We could more accommodate a covariance matrix of the form $\Lambda$ where $c_{1}I\preceq\Lambda\preceq c_{2}I$ but we just present the case $\Lambda=I$ for simplicity.

Our goal will be to show that the max-margin transformer can exhibit benign overfitting in-context. By this we mean that in its forward pass, the transformer memorizes all of the in-context examples, yet still generalizes well for new test examples. To make this more precise, for a sequence of training examples $\{(x_{i},y_{i})\}_{i=1}^{M}$, denote $E(x)$ as the embedding matrix corresponding to the sequence $(x_{i},y_{i})_{i=1}^{M},(x,0)$; recall that the examples $(x_{i},y_{i})_{i=1}^{M}$ are used to formulate predictions for $x$. We say that the transformer exhibits benign overfitting in-context if,

- •
All training examples are memorized: with high probability (over $\mu$ and $(x_{i},y_{i})_{i=1}^{M}$), for each example $(x_{k},y_{k})$, $k\leq M$, we have $\mathrm{sign}(\widehat{y}(E(x_{k});W))=y_{k}$.
- •
Test examples are classified near-optimally: $\mathbb{P}_{\mu,\,(x_{i},y_{i})_{i=1}^{M+1}}(\mathrm{sign}(\widehat{y}(E(x_{M+
1});W)=y_{M+1})\leq p+\varepsilon$ for some vanishing $\varepsilon$.

In Figure 1 we provide a schematic of what this phenomenon would look like in language tasks, where a user provides a sequence of country-capital pairs, some of which have errors, and the transformer repeats these errors at test time for those countries which it has seen with noisy labels, yet still generalizes well for never-before-seen countries.

To prove this phenomenon holds, we need to show both memorization and generalization. Theorem 4.1 in the previous subsection proved the generalization part, thus all that remains is to show that memorization can co-occur with generalization. The next theorem demonstrates that this is possible if the dimension $d$ is suficciently large with respect to the number of in-context examples $M$.

###### Theorem 4.2 .

Assume $\Lambda=I$ and let $\delta\in(0,1)$ be arbitrary. Suppose $p=1/2-c_{p}$ where $c_{p}\in(0,1/2)$ is an absolute constant. There are absolute constants $C>1$, $c>0$ depending only on $c_{p}$ such that if assumptions (A1) through (A3) hold,
then the following holds. With probability at least $1-10\delta$ over the draws of $\{\mu_{\tau},(x_{\tau,i},y_{\tau,i})_{i=1}^{N+1}\}_{\tau=1}^{B}$, for all $\tau\in[B]$, when sampling a new task $\{\mu,(x_{i},y_{i})_{i=1}^{M+1}\}$, the max-margin solution (3.3) satisfies,

- •
In-context training examples are memorized: for $\rho:=\frac{c_{B}\wedge 1}{\log^{2}(2B^{2}/\delta)}$,
$\displaystyle\qquad\mathbb{P}_{(x_{i},y_{i})_{i=1}^{M},\ \mu}\left(\exists k
\in[M]\ \text{ s.t.}\ \mathrm{sign}(\widehat{y}(E(x_{k});W_{\mathsf{MM}}))\neq
y
_{k}\right)$
$\displaystyle\leq 4M\exp\left(-\frac{c\rho\sqrt{d}}{\sqrt{M}}\right)+8M\exp
\left(-\frac{c\rho d}{M(\tilde{R}^{2}\vee\tilde{R})}\right).$
- •
In-context test example achieves test error close to the noise rate:
$\displaystyle\qquad\mathbb{P}_{(x_{i},y_{i})_{i=1}^{M+1},\ \mu}\big{(}\mathrm{
sign}(\widehat{y}(E(x_{M+1});W_{\mathsf{MM}}))\neq y_{M+1}\big{)}$
$\displaystyle\leq p+2\exp(-cM)+2\exp\left(-c\rho\sqrt{d}\right)+4\exp\left(-c
\rho\tilde{R}\right)+2\exp\left(-\frac{c\rho M^{1/2}\tilde{R}^{2}}{\sqrt{d}}
\right),$

The claim regarding generalization in the above theorem is the same as in Theorem 4.1 (we focus here on the case $p>0$ since label noise is essential for overfitting), and the same comments following that theorem apply here as well. But a natural question is whether one can satisfy both memorization and near-noise-rate test error, i.e. benign overfitting. This is indeed possible, and can be seen most easily in the following setting: let $\delta=0.001$, let $M$ be a large constant satisfying $p+2\exp(-cM)\leq p+0.0001$, and assume $B=d$ so $\rho=1/\log^{2}(2d^{2}/\delta)$. If $\tilde{R}=d^{\beta}$ for $\beta\in(1/4,1/2)$, then w.p. at least 99.9% over the pre-training data, since $\beta<1/2$ memorization occurs with probability $1-o_{d}(1)$, and since $\beta>1/4$ the test error for the fresh test example $x_{M+1}$ is at most $p+0.0001+o_{d}(1)$. More generally, one can see that memorization occurs when $d\gg M(\tilde{R}^{2}\vee\tilde{R})$ and $d\gg M$, i.e. when the dimension is large relative to the number of samples and when the signal of the test-time sequence is not too large relative to the dimension divided by the number of in-context samples.

###### Theorem 4.2 .

## 5 Proof Sketch

We will first discuss a sketch of the proof of generalization given in Theorem 4.1 when $\Lambda=I$. For notational simplicity let us denote $W=W_{\mathsf{MM}}$, $\widehat{\mu}:=\frac{1}{M}\sum_{i=1}^{M}y_{i}x_{i}$, and let us drop the $M+1$ subscript so that we denote $(x_{M+1},y_{M+1},\tilde{y}_{M+1})=(x,y,\tilde{y})$. Then the test error is given by the probability of the event,

$$ $\{\mathrm{sign}(\widehat{y}(E;W))\neq y_{M+1}\}=\{\widehat{\mu}^{\top}Wyx\leq 0 ,\,y=\tilde{y}\}\cup\{\widehat{\mu}^{\top}Wyx\leq 0,\,y=-\tilde{y}\}.$ $$

The test error can thus be bounded as

$$ $\displaystyle\mathbb{P}(\widehat{y}(E;W)\neq y_{M+1})\leq\mathbb{P}(\widehat{ \mu}^{\top}W\tilde{y}x\leq 0)+\mathbb{P}(y=-\tilde{y})=\mathbb{P}(\widehat{\mu }^{\top}W\tilde{y}x\leq 0)+p.$ $$

Let us denote the examples $(x_{i},y_{i})$ in the test-time sequence which have clean labels $y_{i}=\tilde{y}_{i}$ as $i\in\mathcal{C}$, while those with noisy labels $y_{i}=-\tilde{y}_{i}$ as $i\in\mathcal{N}$. Then $\sum_{i=1}^{M}y_{i}x_{i}=(|\mathcal{C}|-|\mathcal{N}|)\mu+\sum_{i=1}^{M}y_{i}z
_{i}=(M-2|\mathcal{N}|)\mu+\sum_{i=1}^{M}y_{i}z_{i}$. If we denote $\widehat{p}:=|\mathcal{N}|/M$ then by standard properties of the Gaussian, we have for $z,z^{\prime}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathsf{N}(0,I)$,

$$ $\displaystyle\widehat{\mu}$ $\displaystyle\stackrel{{\scriptstyle\mathrm{d}}}{{=}}\left(1-2\widehat{p} \right)\mu+M^{-1/2}z,$ $\displaystyle\tilde{y}x$ $\displaystyle\stackrel{{\scriptstyle\mathrm{d}}}{{=}}\mu+z^{\prime}.$ $$

We thus have

$$ $\displaystyle\qquad\mathbb{P}(\widehat{y}(E;W)\neq y_{M+1})$ $\displaystyle\leq p+\mathbb{P}\Big{(}\big{(}(1-2\widehat{p})\mu+M^{-1/2}z\big{ )}^{\top}W(\mu+z^{\prime})<0\Big{)}$ $\displaystyle=p+\mathbb{P}\Big{(}(1-2\widehat{p})\mu^{\top}W\mu<-(1-2\widehat{ p})\mu^{\top}Wz^{\prime}-M^{-1/2}z^{\top}W\mu-M^{-1/2}z^{\top}Wz^{\prime}\Big{ )}.$ $$

This decomposition allows for us to show the test error is near the noise rate $p$ if we can show that $(1-2\widehat{p})\mu^{\top}W\mu$ is large and positive while all of the other terms involving $\mu^{\top}Wz^{\prime}$, $z^{\top}W\mu$, and $z^{\top}Wz^{\prime}$ are small in comparison. Now if $\mu$ were a standard Gaussian, or if it were sub-Gaussian with independent components, then $\mu^{\top}W\mu$ would be close to its mean $\operatorname*{tr}(W)$, with fluctuations determined from its mean determined by $\|W\|_{F}$, via a standard Hanson-Wright inequality. However, since $\mu\sim\mathsf{Unif}(\tilde{R}\cdot\mathbb{S}^{d-1})$, the components of $\mu$ are not independent, and so we derive a modified version of Hanson-Wright (Lemma C.2) which applies in this setting where we show $\mu^{\top}W\mu$ is close to its mean $\frac{\tilde{R}^{2}}{d}\operatorname*{tr}(W)$ with fluctuations from its mean determined by $\|W\|_{F}$. Thus, in order to show that the max-margin solution achieves (in-context) test error close to the noise rate $p$, we must show the following:

- •
A lower bound for $\operatorname*{tr}(W)$, so that the mean of $\mu^{\top}W\mu$ is large and positive.
- •
An upper bound on $\|W\|_{F}$, so that we can control the fluctuations of $\mu^{\top}W\mu$ from its mean.
- •
An upper bound on $\widehat{p}$ so that $(1-2\widehat{p})$ is positive.
- •
Upper bounds on the quantities $|(1-2\widehat{p})\mu^{\top}Wz^{\prime}|$, $|M^{-1/2}z^{\top}W\mu|$ and $|M^{-1/2}z^{\top}Wz^{\prime}|$ so that these quantities are smaller than the positive term coming from the mean of $\mu^{\top}W\mu$. These, in turn, will require upper bounds on $\|W\|_{F}$.

Thus provided we have control over $\operatorname*{tr}(W)$ and $\|W\|_{F}$, we can apply straightforward concentration inequalities to show that the (in-context) test error is near the noise rate. So let us now describe how to derive such bounds.

The starting place for our analysis is the definition of the max-margin solution (3.3). For notational simplicity let us define $\widehat{\mu}_{\tau}:=\frac{1}{N}\sum_{i=1}^{n}y_{\tau,i}x_{\tau,i}$ and let us denote $x_{\tau,N+1},y_{\tau,N+1}$ as $x_{\tau},y_{\tau}$ so that the max-margin solution can be written

$$ $W:=\mathop{\mathrm{argmin}}\{\|U\|_{F}^{2}:\widehat{\mu}_{\tau}^{\top}Uy_{\tau }x_{\tau}\geq 1,\,\forall\tau=1,\dots,B\}.$ (5.1) $$

Since $\nabla\widehat{y}(E_{\tau};W)=\widehat{\mu}_{\tau}x_{\tau}^{\top}$, the KKT conditions imply that there exist $\lambda_{1},\dots,\lambda_{B}\geq 0$ such that

$$ $W=\sum_{\tau=1}^{B}\lambda_{\tau}y_{\tau}\widehat{\mu}_{\tau}x_{\tau}^{\top},$ (5.2) $$

and moreover we have $\lambda_{\tau}=0$ whenever $y_{\tau}\widehat{\mu}_{\tau}^{\top}Wx_{\tau}\neq 1$. By a careful analysis of the KKT conditions (see Lemma B.1), we can derive nearly-matching upper and lower bounds for $\sum_{\tau}\lambda_{\tau}$, $\operatorname*{tr}(W)$, and $\|W\|_{F}$: if we assume that $c_{B}>0$ is an absolute constant (i.e. the number of tasks $B$ is greater than a constant multiple of the input dimension $d$) then these bounds take the form

$$ $\sum_{\tau=1}^{B}\lambda_{\tau}=\tilde{\Theta}(d/R^{4}),\qquad\operatorname*{ tr}(W)=\tilde{\Theta}(d/R^{2}),\qquad\|W\|_{F}=\tilde{\Theta}(\sqrt{d}/R^{2}).$ (5.3) $$

These bounds, together with the concentration inequalities outlined above, suffice for proving Theorem 4.2.

As for the possibility of benign overfitting in-context shown in Theorem 4.2, since Theorem 4.1 guarantees generalization, the only task that remains to be shown is that overfitting occurs: namely, we want to ensure that for every example $(x_{k},y_{k})$, if we denote $E(x_{k})$ as the embedding matrix with input sequence $(x_{i},y_{i})_{i=1}^{M}$ but with test example $x_{k}$, then we want to show that

$$ $\mathrm{sign}\left(\widehat{y}(E(x_{k});W)\right)=y_{k}\Longleftrightarrow \widehat{\mu}^{\top}Wy_{k}x_{k}>0.$ $$

Now, since $\operatorname*{tr}(cI)=cd$ and $\|cI\|_{F}=c\sqrt{d}$, the identities in (5.3) suggest that $W$ has properties similar to that of a scaled identity matrix. If $W$ were indeed a scaled identity matrix, then the goal would be to show that

$$ $M\widehat{\mu}^{\top}y_{k}x_{k}=\left(\sum_{i=1}^{M}y_{i}x_{i}\right)^{\top}y_ {k}x_{k}=\|x_{k}\|^{2}+\sum_{i\neq k}\langle y_{i}x_{i},y_{k}x_{k}\rangle\geq \|x_{k}\|^{2}-\sum_{i\neq k}|\langle x_{i},x_{k}\rangle|>0.$ $$

That is, we would need to ensure that the examples $\{x_{i}\}_{i=1}^{M}$ are nearly-orthogonal in a particular sense \citepfrei2023implicit. This occurs provided the ambient dimension $d$ is much larger than the number of examples $M$ and the signal strength $\tilde{R}^{2}=\|\mu\|^{2}$, and has been shown in prior work on benign overfitting in neural networks \citepfrei2022benign,frei2023benignkkt. In our setting, $W$ is not a multiple of the identity matrix so this proof technique does not directly apply, but at a high level the proof ideas are similar, and appear in Section D.

## 6 Conclusion

In this work we developed task complexity and sample complexity guarantees for in-context learning class-conditional Gaussian mixture models with a single-layer linear transformer architecture. We analyzed the implicit regularization of gradient descent to characterize the algorithm implemented by the transformer after pre-training. This allowed for us to quantify how many in-context samples are needed in order to achieve small test error, which to the best of our knowledge has not been explored in the classification setting prior to this work. We also showed how the trained transformer can exhibit benign overfitting in-context, i.e. in its forward pass the transformer can memorize noisy examples yet still achieve near-optimal test error.

There are a number of natural directions for future research. We relied upon a convex linear transformer architecture which allows for us to identify the pre-trained transformer as the global max-margin solution in parameter space. More general linear transformer architectures are not convex but are often homogeneous in their parameters. In this setting we thus know \citeplyu2019gradient,ji2020directional that GD has an implicit bias towards first-order stationary (KKT) points of the max-margin problem in parameter space. It may be possible to analzye the consequences of the KKT conditions to develop generalization guarantees \citepsafran2022effective,frei2023doubleedgedswordimplicitbias,frei2023benignkkt. For softmax-based transformer architectures, it would be interesting to see if prior works on implicit regularization of GD over these architectures \citeptarzanagh2023max,tarzanagh2023transformers,thrampoulidis2024implicitbiasnexttokenprediction,vasudeva2024implicit can be used to understand in-context learning. Finally, we assumed that the signal-to-noise ratio in the pre-training data was quite large (see (A1)), and that the pre-training data did not have noisy labels. It would be interesting to understand if pre-training on more difficult or noisy data would result in a qualitatively different algorithm implemented by the pre-trained transformer.

## Appendix A Properties of the pre-training datasets

We begin by developing guarantees for various properties of the pre-training datasets, which will form the basis for understanding properties of the max-margin solution (3.3).

###### Lemma A.1 .

There is an absolute constant $c_{0}>1$ such that with probability at least $1-10\delta$ over the draws of $\{\mu_{\tau},(x_{\tau},y_{\tau}),(x_{\tau,i},y_{\tau,i})_{i=1}^{N}\}_{\tau=1}^
{B}$, for all $\tau\in[B]$ and $q\neq\tau$,

$$ $\displaystyle\left|\|\widehat{\mu}_{\tau}\|^{2}-R^{2}\right|$ $\displaystyle\leq\frac{c_{0}R\sqrt{\operatorname*{tr}(\Lambda)}\log(2B/\delta) }{\sqrt{Nd}}+4\frac{\operatorname*{tr}(\Lambda)\vee c_{0}\|\Lambda\|_{2}\log(2 B/\delta)}{N},$ $\displaystyle|\|x_{\tau}\|^{2}-R^{2}|$ $\displaystyle\leq\frac{2c_{0}R\sqrt{\operatorname*{tr}(\Lambda)}\log(2B/\delta )}{\sqrt{d}}+4\left(\operatorname*{tr}(\Lambda)\vee c_{0}\|\Lambda\|_{2}\log(2 B/\delta)\right),$ $\displaystyle|\langle\widehat{\mu}_{q},\widehat{\mu}_{\tau}\rangle|$ $\displaystyle\leq c_{0}\left(\frac{R^{2}}{\sqrt{d}}+\frac{R\sqrt{\operatorname *{tr}(\Lambda)}}{\sqrt{Nd}}+\frac{\sqrt{\operatorname*{tr}(\Lambda^{2})}}{N} \right)\log(2B^{2}/\delta),$ $\displaystyle|\langle x_{\tau},x_{q}\rangle|$ $\displaystyle\leq c_{0}\left(\frac{R^{2}}{\sqrt{d}}+\frac{R\sqrt{\operatorname *{tr}(\Lambda)}}{\sqrt{d}}+\sqrt{\operatorname*{tr}(\Lambda^{2})}\right)\log(2 B^{2}/\delta),$ $\displaystyle\left|\langle\widehat{\mu}_{\tau},y_{\tau}x_{\tau}\rangle-R^{2}\right|$ $\displaystyle\leq c_{0}\left(\left[1+\frac{1}{\sqrt{N}}\right]\frac{R\sqrt{ \operatorname*{tr}(\Lambda)}}{\sqrt{d}}+\frac{\sqrt{\operatorname*{tr}(\Lambda ^{2})}}{\sqrt{N}}\right)\log(2B/\delta)$ $$

###### Proof.

By definition of $\widehat{\mu}_{\tau}$ and properties of the Gaussian distribution, there is $z_{\tau}^{\prime}\sim\mathsf{N}(0,\Lambda)$ such that

$$ $\widehat{\mu}_{\tau}=\frac{1}{N}\sum_{i=1}^{N}y_{\tau,i}(y_{\tau,i}\mu_{\tau}+ z_{\tau,i})=\mu_{\tau}+\frac{1}{N}\sum_{i=1}^{N}y_{\tau,i}z_{\tau,i}=\mu_{\tau }+\frac{1}{\sqrt{N}}z_{\tau}^{\prime}.$ $$

Thus for $\tau\neq q$ there are $z_{\tau}^{\prime},z_{q}^{\prime}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim
}}\mathsf{N}(0,\Lambda)$ such that

$$ $\displaystyle\langle\widehat{\mu}_{q},\widehat{\mu}_{\tau}\rangle$ $\displaystyle\stackrel{{\scriptstyle\text{d}}}{{=}}\langle\mu_{\tau}+N^{-1/2}z _{\tau}^{\prime},\mu_{q}+N^{-1/2}z_{q}^{\prime}\rangle$ $\displaystyle=\langle\mu_{\tau},\mu_{q}\rangle+N^{-1/2}\langle z_{\tau}^{ \prime},\mu_{q}\rangle+N^{-1/2}\langle z_{q}^{\prime},\mu_{\tau}\rangle+N^{-1} \langle z_{\tau}^{\prime},z_{q}^{\prime}\rangle.$ (A.1) $$

We first derive an upper bound for this quantity when $q\neq\tau$.
Now, since $\mu_{q},\mu_{\tau}$ are independent and sub-Gaussian random vectors with sub-Gaussian norm at most $cR/\sqrt{d}$ \citep[Theorem 3.4.6]vershynincup for some absolute constant $c>0$, by \citet[Lemma 6.2.3]vershynincup, we have for some $c^{\prime}>0$ it holds that for any $\beta\in\mathbb{R}$, if $g,g^{\prime}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathsf{N}(0,I_{d})$,

$$ $\mathbb{E}[\exp(\beta\mu_{q}^{\top}\mu_{\tau})]\leq\mathbb{E}[\exp(c^{\prime}R ^{2}d^{-1}\beta g^{\top}g^{\prime})].$ $$

By \citet[Lemma 6.2.2]vershynincup, for some $c_{1}>0$, provided $c|\beta|R^{2}/d\leq c_{1}$, it holds that

$$ $\mathbb{E}[\exp(cR^{2}d^{-1}\beta g^{\top}g^{\prime})]\leq\exp(c_{1}\beta^{2}R ^{4}d^{-2}\|I_{d}\|_{F}^{2})=\exp(c_{1}\beta^{2}R^{4}d^{-1}).$ $$

Since $\mu_{q},\mu_{\tau}$ are mean-zero, by \citet[Proposition 2.7.1]vershynincup this implies the quantity $\mu_{q}^{\top}\mu_{\tau}$ is sub-exponential with $\|\mu_{q}^{\top}\mu_{\tau}\|_{\psi_{1}}\leq c_{2}R^{2}/\sqrt{d}$
for some absolute constant $c_{2}>0$. We therefore have, for some absolute $c_{3}>0$, w.p. at least $1-\delta$, for all $\tau\in[B]$,

$$ $\displaystyle|\langle\mu_{\tau},\mu_{q}\rangle|$ $\displaystyle\leq c_{3}R^{2}d^{-1/2}\log(2B/\delta).$ (A.2) $$

Again using Lemmas 6.2.2 and 6.2.3 from \citepvershynincup, since $\mu_{q}$ has sub-Gaussian norm at most $cR/\sqrt{d}$ and $z^{\prime}_{\tau}=\Lambda^{1/2}g^{\prime\prime}$ where $g^{\prime\prime}$ has sub-Gaussian norm at most $c$, we have

$$ $\displaystyle\mathbb{E}[\exp(\beta\mu_{q}^{\top}z_{\tau}^{\prime})]$ $\displaystyle\leq\mathbb{E}[\exp(cRd^{-1/2}\beta g^{\top}\Lambda^{1/2}g^{ \prime})],$ $$

and thus provided $cRd^{-1/2}|\beta|\leq c_{1}/\|\Lambda^{1/2}\|_{2}$ we have

$$ $\mathbb{E}[\exp(cRd^{-1/2}\beta g^{\top}\Lambda^{1/2}g^{\prime})]\leq\exp(c^{ \prime}R^{2}d^{-1}\beta^{2}\|\Lambda^{1/2}\|_{F}^{2})=\exp(c^{\prime}R^{2}d^{- 1}\beta^{2}\operatorname*{tr}(\Lambda)).$ $$

In particular, the quantity $\mu_{q}^{\top}z_{\tau}^{\prime}$ is sub-exponential with sub-exponential norm $\|\mu_{q}^{\top}z_{\tau}^{\prime}\|_{\psi_{1}}\leq c_{4}Rd^{-1/2}\sqrt{
\operatorname*{tr}(\Lambda)}$, and so for some absolute $c_{5}>0$ we have with probability at least $1-\delta$, for all $q,\tau\in[B]$ with $q\neq\tau$,

$$ $\displaystyle|\langle\mu_{q},z_{\tau}^{\prime}\rangle|\leq c_{5}Rd^{-1/2}\sqrt {\operatorname*{tr}(\Lambda)}\log(2B^{2}/\delta).$ (A.3) $$

For $\langle z_{q}^{\prime},z_{\tau}^{\prime}\rangle$ with $\tau\neq q$ we can directly use the MGF of Gaussian chaos \citep[Lemma 6.2.2]vershynincup: $\langle z_{q}^{\prime},z_{\tau}^{\prime}\rangle=g^{\top}\Lambda g^{\prime}$ for i.i.d. $g,g^{\prime}\sim\mathsf{N}(0,I_{d})$ so that for $\beta\leq c/\|\Lambda\|_{2}$,

$$ $\displaystyle\mathbb{E}[\exp(\beta\langle z_{q}^{\prime},z_{\tau}^{\prime} \rangle)]\leq\exp(c_{6}\beta^{2}\|\Lambda\|_{F}^{2})=\exp(c_{6}\beta^{2} \operatorname*{tr}(\Lambda^{2})).$ $$

In particular, $\left\|\langle z_{q}^{\prime},z_{\tau}^{\prime}\rangle\right\|_{\psi_{1}}\leq c
_{7}\sqrt{\operatorname*{tr}(\Lambda^{2})}$ so that sub-exponential concentration implies that with probability at least $1-\delta$, for any $q,\tau\in[B]$ with $q\neq\tau$,

$$ $|\langle z_{\tau}^{\prime},z_{q}^{\prime}\rangle|\leq c_{7}\sqrt{\operatorname *{tr}(\Lambda^{2})}\log(2B^{2}/\delta).$ (A.4) $$

Putting (A.2), (A.3), and (A.4) into (A.1) we get for $q\neq\tau$,

$$ $\displaystyle|\langle\widehat{\mu}_{q},\widehat{\mu}_{\tau}\rangle|$ $\displaystyle=c_{8}\left(\frac{R^{2}}{\sqrt{d}}+\frac{R\sqrt{\operatorname*{tr }(\Lambda)}}{\sqrt{Nd}}+\frac{\sqrt{\operatorname*{tr}(\Lambda^{2})}}{N}\right )\log(2B^{2}/\delta).$ (A.5) $$

As for $\|\widehat{\mu}_{\tau}\|^{2}$, from (A.1) we have

$$ $\displaystyle\|\widehat{\mu}_{\tau}\|^{2}$ $\displaystyle=\|\mu_{\tau}\|^{2}+2N^{-1/2}\langle z_{\tau}^{\prime},\mu_{\tau} \rangle+N^{-1}\|z_{\tau}^{\prime}\|^{2}$ (A.6) $$

From here, the same argument used to bound (A.3) holds since that bound only relied upon the fact that $\mu_{q}$ and $z_{\tau}^{\prime}$ are independent, while $\mu_{\tau}$ and $z_{\tau}^{\prime}$ are independent as well. In particular, with probability at least $1-\delta$, for all $\tau\in[B]$,

$$ $\displaystyle|\langle\mu_{\tau},z_{\tau}^{\prime}\rangle|\leq c_{5}Rd^{-1/2} \sqrt{\operatorname*{tr}(\Lambda)}\log(2B/\delta).$ (A.7) $$

Next, we have that $\|z_{\tau}^{\prime}\|^{2}\stackrel{{\scriptstyle\text{d}}}{{=}}g^{\top}\Lambda
g
=\|\Lambda^{1/2}g\|^{2}$ for $g\sim\mathsf{N}(0,I_{d})$. Therefore \citet[Theorem 6.3.2]vershynincup implies $\|\|z_{\tau}^{\prime}\|_{2}-\|\Lambda^{1/2}\|_{F}\|_{\psi_{2}}\leq c\|\Lambda^
{1/2}\|_{2}$ for some absolute constant $c>0$. Note that $\|\Lambda^{1/2}\|_{F}=\sqrt{\operatorname*{tr}(\Lambda)}$. And thus by sub-Gaussian concentration, we have for some constant $c_{9}>0$, with probability at least $1-\delta$, for all $\tau\in[B]$,

$$ $\displaystyle|\|z_{\tau}^{\prime}\|_{2}-\sqrt{\operatorname*{tr}(\Lambda)}| \leq c_{9}^{1/2}\|\Lambda^{1/2}\|_{2}\sqrt{\log(2B/\delta)}\implies\|z_{\tau}^ {\prime}\|\leq 2(\sqrt{\operatorname*{tr}(\Lambda)}\vee c_{9}^{1/2}\|\Lambda^{ 1/2}\|_{2}\sqrt{\log(2B/\delta)}).$ $$

In particular,

$$ $\displaystyle\|z_{\tau}^{\prime}\|^{2}\leq 4\left(\operatorname*{tr}(\Lambda) \vee c_{9}\|\Lambda\|_{2}\log(2B/\delta)\right).$ (A.8) $$

Putting (A.8) and (A.7) into (A.6) and using that $\|\mu_{\tau}\|^{2}=R^{2}$, we get with probability at least $1-2\delta$,

$$ $\displaystyle\left|\|\widehat{\mu}_{\tau}\|^{2}-R^{2}\right|\leq\frac{c_{5}R \sqrt{\operatorname*{tr}(\Lambda)}\log(2B/\delta)}{\sqrt{Nd}}+4\frac{ \operatorname*{tr}(\Lambda)\vee c_{9}\|\Lambda\|_{2}\log(2B/\delta)}{N}.$ $$

As for $\|x_{\tau}\|^{2}$, by definition,

$$ $\displaystyle\|x_{\tau}\|^{2}$ $\displaystyle=\|\mu_{\tau}\|^{2}+2\langle\mu_{\tau},z_{\tau}\rangle+\|z_{\tau} \|^{2}=R^{2}+2\langle\mu_{\tau},z_{\tau}\rangle+\|z_{\tau}\|^{2}.$ $$

Since $z_{\tau}\sim\mathsf{N}(0,\Lambda)$ has the same distribution as $z_{\tau}^{\prime}$, the same analysis used to prove (A.7) and (A.8) yields that with probability at least $1-2\delta$, for all $\tau\in[B]$,

$$ $\displaystyle|\langle\mu_{\tau},z_{\tau}\rangle|$ $\displaystyle\leq c_{5}Rd^{-1/2}\sqrt{\operatorname*{tr}(\Lambda)}\log(2B/ \delta),$ $\displaystyle\|z_{\tau}\|^{2}$ $\displaystyle\leq 4(\operatorname*{tr}(\Lambda)\vee c_{9}\|\Lambda\|_{2}\log(2 B/\delta).$ $$

Substituting these into the preceding display we get

$$ $\displaystyle|\|x_{\tau}\|^{2}-R^{2}|\leq\frac{2c_{5}R\sqrt{\operatorname*{tr} (\Lambda)}\log(2B/\delta)}{\sqrt{d}}+4\left(\operatorname*{tr}(\Lambda)\vee c_ {9}\|\Lambda\|_{2}\log(2B/\delta)\right)$ $$

Thus provided $R$ is sufficiently large, then we also have $\|x_{\tau}\|^{2}=\Theta(R^{2})$.

Next we bound $|\langle x_{\tau},x_{q}\rangle|$: There are $z_{\tau}^{\prime},z_{q}^{\prime}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim
}}\mathsf{N}(0,\Lambda)$ such that

$$ $\displaystyle\langle y_{\tau}x_{\tau},y_{q}x_{q}\rangle$ $\displaystyle\stackrel{{\scriptstyle\text{d}}}{{=}}\langle\mu_{\tau}+z_{\tau}^ {\prime},\mu_{q}+z_{q}^{\prime}\rangle.$ $$

It is clear that the same exact analysis we used to analyze (A.1) leads to the claim that with probability at least $1-\delta$, for all $q\neq\tau$:

$$ $\displaystyle|\langle x_{q},x_{\tau}\rangle|$ $\displaystyle\leq c_{9}\left(\frac{R^{2}}{\sqrt{d}}+\frac{R\sqrt{\operatorname *{tr}(\Lambda)}}{\sqrt{d}}+\sqrt{\operatorname*{tr}(\Lambda^{2})}\right)\log(2 B^{2}/\delta).$ (A.9) $$

Finally, we consider $y_{\tau}\widehat{\mu}_{\tau}^{\top}x_{\tau}$. Just as in the previous analyses, there are $z_{\tau},z_{\tau}^{\prime}\sim\mathsf{N}(0,\Lambda)$ such that

$$ $\displaystyle\langle\widehat{\mu}_{\tau},y_{\tau}x_{\tau}\rangle$ $\displaystyle\stackrel{{\scriptstyle\text{d}}}{{=}}\langle\mu_{\tau}+N^{-1/2}z _{\tau},\mu_{\tau}+z_{\tau}^{\prime}\rangle$ $\displaystyle=\|\mu_{\tau}\|^{2}+N^{-1/2}\langle z_{\tau},\mu_{\tau}\rangle+ \langle z_{\tau}^{\prime},\mu_{\tau}\rangle+N^{-1/2}\langle z_{\tau},z_{\tau}^ {\prime}\rangle.$ $$

Again using an analysis similar to that used for (A.1) yields that with probability at least $1-\delta$, for all $\tau\in[B]$,

$$ $\displaystyle\left|\langle\widehat{\mu}_{\tau},y_{\tau}x_{\tau}\rangle-R^{2} \right|\leq c_{10}\left(\left[1+\frac{1}{\sqrt{N}}\right]\frac{R\sqrt{ \operatorname*{tr}(\Lambda)}}{\sqrt{d}}+\frac{\sqrt{\operatorname*{tr}(\Lambda ^{2})}}{\sqrt{N}}\right)\log(2B/\delta).$ (A.10) $$

Taking a union bound over each of the events shows that all of the desired claims of Lemma A.1 hold with probability at least $1-10\delta$.

∎

The events of Lemma A.1 hold with probability at least $1-10\delta$, independently of the assumptions (A1), (A2), and (A3).(^2^22Note, however, that the quantities appearing on the right-hand sides of each inequality in the lemma are only small when these assumptions hold; this is the reason for these assumptions.) Our results will require this event to hold, and so we introduce the following to allow for us to refer to this in later lemmas.

###### Definition A.2 .

Let us say that a good run holds if the events of Lemma A.1 hold.

###### Lemma A.1 .

###### Proof.

###### Definition A.2 .

## Appendix B Analysis of max-margin solution

In this section we derive a number of properties of the max-margin solution (3.3).

###### Lemma B.1 .

On a good run, for any $c_{B}>0$ and for $C>1$ sufficiently large under Assumptions (A1) and (A2), the max-margin solution $W$ of Problem 3.3,

$$ $W=\sum_{\tau=1}^{B}\lambda_{\tau}y_{\tau}\widehat{\mu}_{\tau}x_{\tau}^{\top},$ $$

is such that the $\lambda_{\tau}\geq 0$ satisfy the following:

$$ $\frac{(c_{B}\wedge 1)d}{8c_{0}^{2}R^{4}\log^{2}(2B^{2}/\delta)}\leq\sum_{\tau= 1}^{B}\lambda_{\tau}\leq\frac{4d}{R^{4}},$ $$

where $c_{0}>1$ is the constant from Lemma A.1.
Further, we have the inequalities

$$ $\displaystyle\frac{(c_{B}\wedge 1)\sqrt{d}}{16c_{0}^{2}R^{2}\log^{2}(2B^{2}/ \delta)}\leq\|W\|_{F}\leq\frac{2\sqrt{d}}{R^{2}},$ $$

and

$$ $\frac{(c_{B}\wedge 1)d}{16c_{0}^{2}R^{2}\log^{2}(2B^{2}/\delta)}\leq \operatorname*{tr}(W)\leq\frac{6d}{R^{2}}.$ $$

###### Proof.

We first derive an upper bound on $\|W\|_{F}$ by showing that the matrix $U=I_{d}$ satisfies the constraints of the max-margin problem (3.3).

$$ $\displaystyle y_{\tau}\widehat{\mu}_{\tau}^{\top}Ix_{\tau}$ $\displaystyle=\langle\widehat{\mu}_{\tau},y_{\tau}x_{\tau}\rangle$ $\displaystyle\overset{(i)}{\geq}R^{2}\left(1-2c_{0}\left(\frac{\sqrt{ \operatorname*{tr}(\Lambda)}}{R\sqrt{d}}+\sqrt{\frac{\operatorname*{tr}( \Lambda^{2})}{NR^{4}}}\right)\log(2B/\delta)\right)$ $\displaystyle\overset{(ii)}{\geq}R^{2}\left(1-\frac{4c_{0}}{C}\right)$ $\displaystyle\geq\frac{1}{2}R^{2}.$ (B.1) $$

Inequality $(i)$ uses Lemma A.1, while inequality $(ii)$ holds for $C>8c_{0}$ sufficiently large via Assumption (A1). Thus the matrix $2I_{d}/R^{2}$ separates the training data with margin at least 1 for every sample. Since $W$ is the minimum Frobenius norm matrix which separates all of the training data with margin 1, this implies

$$ $\displaystyle\|W\|_{F}$ $\displaystyle\leq\|2I_{d}/R^{2}\|_{F}=\frac{2\sqrt{d}}{R^{2}}.$ (B.2) $$

On the other hand, by the variational definition of the norm, since $\|I_{d}/\sqrt{d}\|_{F}=1$ we know $\|W\|_{F}\geq\langle W,I_{d}/\sqrt{d}\rangle$, and hence

$$ $\displaystyle\|W\|_{F}$ $\displaystyle\geq\frac{1}{\sqrt{d}}\langle W,I_{d}\rangle$ $\displaystyle=\frac{1}{\sqrt{d}}\left\langle\sum_{\tau=1}^{B}\lambda_{\tau} \widehat{\mu}_{\tau}y_{\tau}x_{\tau}^{\top},I_{d}\right\rangle$ $\displaystyle=\frac{1}{\sqrt{d}}\sum_{\tau=1}^{B}\lambda_{\tau}\langle\widehat {\mu}_{\tau},y_{\tau}x_{\tau}\rangle$ $\displaystyle\geq\frac{1}{\sqrt{d}}\cdot\frac{1}{2}R^{2}\sum_{\tau=1}^{B} \lambda_{\tau},$ (B.3) $$

where the last line uses (B.1). Putting this and the preceding display together, we get

$$ $\displaystyle\sum_{\tau=1}^{B}\lambda_{\tau}\leq\frac{4d}{R^{4}}.$ (B.4) $$

Next, we note that by the feasibility conditions of the max-margin problem, we have for any $\tau$,

$$ $\displaystyle 1$ $\displaystyle\leq y_{\tau}\widehat{\mu}_{\tau}^{\top}Wx_{\tau}$ $\displaystyle=\widehat{\mu}_{\tau}^{\top}\left(\sum_{q=1}^{B}\lambda_{q}y_{q} \widehat{\mu}_{q}x_{q}^{\top}\right)y_{\tau}x_{\tau}$ $\displaystyle=\sum_{q=1}^{B}\lambda_{q}\langle\widehat{\mu}_{\tau},\widehat{ \mu}_{q}\rangle\langle y_{\tau}x_{\tau},y_{q}x_{q}\rangle$ $\displaystyle=\lambda_{\tau}\|\widehat{\mu}_{\tau}\|^{2}\|x_{\tau}\|^{2}+\sum_ {q:\ q\neq\tau}\lambda_{q}\langle\widehat{\mu}_{\tau},\widehat{\mu}_{q}\rangle \langle y_{q}x_{q},x_{\tau}y_{\tau}\rangle.$ (B.5) $$

Now by Lemma A.1, we have

$$ $\displaystyle\left|\frac{\|\widehat{\mu}_{\tau}\|^{2}}{R^{2}}-1\right|$ $\displaystyle\leq\frac{c_{0}\sqrt{\operatorname*{tr}(\Lambda)\log(2B/\delta)}} {R\sqrt{Nd}}+4\frac{\operatorname*{tr}(\Lambda)\vee c_{0}\|\Lambda\|_{2}\log(2 B/\delta)}{NR^{2}},$ $$

and

$$ $\displaystyle\left|\frac{\|x_{\tau}\|^{2}}{R^{2}}-1\right|$ $\displaystyle\leq\frac{c_{0}\sqrt{\operatorname*{tr}(\Lambda)}\log(2B/\delta)} {R\sqrt{d}}+4\frac{\operatorname*{tr}(\Lambda)\vee c_{0}\|\Lambda\|_{2}\log(2B /\delta)}{R^{2}}.$ $$

Using Assumption (A1) we see that for $C$ sufficiently large we have

$$ $\displaystyle\|\widehat{\mu}_{\tau}\|^{2}\|x_{\tau}\|^{2}\leq R^{4}\left(1+ \frac{c_{0}}{NC}+\frac{4c_{0}}{NC}\right)\cdot\left(1+\frac{c_{0}}{C}+\frac{4c _{0}}{C}\right)\leq\frac{3}{2}R^{4},$ (B.6) $$

where the inequality uses Assumption (A1). Likewise, we have

$$ $\displaystyle\|\widehat{\mu}_{\tau}\|^{2}\|x_{\tau}\|^{2}\geq\frac{1}{2}R^{4}.$ (B.7) $$

As for the cross terms,
again using Lemma A.1, for all $\tau\neq q$,

$$ $\displaystyle\quad|\langle\widehat{\mu}_{\tau},\widehat{\mu}_{q}\rangle|\cdot| \langle x_{\tau},x_{q}\rangle|$ $\displaystyle\leq c_{0}^{2}\left(\frac{R^{2}}{\sqrt{d}}+\frac{R\sqrt{ \operatorname*{tr}(\Lambda)}}{\sqrt{Nd}}+\frac{\sqrt{\operatorname*{tr}( \Lambda^{2})}}{N}\right)\cdot\left(\frac{R^{2}}{\sqrt{d}}+\frac{R\sqrt{ \operatorname*{tr}(\Lambda)}}{\sqrt{d}}+\sqrt{\operatorname*{tr}(\Lambda^{2})} \right)\log^{2}(2B^{2}/\delta)$ $\displaystyle=\frac{c_{0}^{2}R^{4}}{d}\left(1+\frac{\sqrt{\operatorname*{tr}( \Lambda)}}{R\sqrt{N}}+\frac{\sqrt{d\operatorname*{tr}(\Lambda^{2})}}{R^{2}N} \right)\cdot\left(1+\frac{\sqrt{\operatorname*{tr}(\Lambda)}}{R}+\frac{\sqrt{d \operatorname*{tr}(\Lambda^{2})}}{R^{2}}\right)\log^{2}(2B^{2}/\delta)$ $\displaystyle\overset{(i)}{\leq}\frac{c_{0}^{2}R^{4}}{d}\left(1+\frac{1}{C \sqrt{N}}+\frac{1}{4N}\right)\cdot\left(1+\frac{1}{C}+\frac{1}{4}\right)\log^{ 2}(2B^{2}/\delta)$ $\displaystyle\overset{(ii)}{\leq}\frac{2c_{0}^{2}R^{4}\log^{2}(2B^{2}/\delta)} {d}.$ (B.8) $$

Inequality $(i)$ uses Lemma A.1 and $(ii)$ uses that $C$ is large enough.
Putting this together with (B.6) and substituting these into the consequences of the feasibility condition (B.5) we get for any $\tau$,

$$ $\displaystyle 1$ $\displaystyle\leq\frac{3}{2}R^{4}\lambda_{\tau}+\frac{2c_{0}^{2}R^{4}\log^{2}( 2B^{2}/\delta)}{d}\sum_{q:\ q\neq\tau}\lambda_{q}.$ (B.9) $$

We now show that this implies $\sum_{q=1}^{B}\lambda_{q}\geq\frac{(c_{B}\wedge 1)d}{8c_{0}^{2}R^{4}\log^{2}(2
B^{2}/\delta)}$. Towards this end, let us first consider the case that there exists some $\tau$ with $\lambda_{\tau}\leq\frac{c_{B}d}{4R^{4}B}$. Then the preceding display implies

$$ $\displaystyle 1$ $\displaystyle\leq\frac{3}{2}R^{4}\cdot\frac{c_{B}d}{4R^{4}B}+\frac{2c_{0}^{2}R ^{4}\log^{2}(2B^{2}/\delta)}{d}\sum_{q:\ q\neq\tau}\lambda_{q}$ $\displaystyle\leq\frac{3}{4}+\frac{2c_{0}^{2}R^{4}\log^{2}(2B^{2}/\delta)}{d} \sum_{q:\ q\neq\tau}\lambda_{q}.$ $$

The final inequality uses Assumption (A2), i.e. $B\geq c_{B}d$. Rearranging we see that

$$ $\sum_{q=1}^{B}\lambda_{q}\geq\sum_{q:\ q\neq\tau}\lambda_{q}\geq\frac{d}{8c_{0 }^{2}R^{4}\log^{2}(2B^{2}/\delta)}\geq\frac{(c_{B}\wedge 1)d}{8c_{0}^{2}R^{4} \log^{2}(2B^{2}/\delta)}.$ $$

Let us now consider the only remaining case, whereby for each $\tau$ we have $\lambda_{\tau}>\frac{c_{B}d}{4R^{4}B}$. Then,

$$ $\displaystyle\sum_{q=1}^{B}\lambda_{q}>\sum_{q=1}^{B}\frac{c_{B}d}{4R^{4}B}= \frac{c_{B}d}{4R^{4}}\geq\frac{(c_{B}\wedge 1)d}{8c_{0}^{2}R^{4}\log^{2}(2B^{2 }/\delta)},$ $$

where the final inequality uses that $c_{0}\geq 1$. We therefore have $\sum_{q=1}^{B}\lambda_{q}\geq\frac{(c_{B}\wedge 1)d}{8c_{0}^{2}R^{4}\log^{2}(2
B^{2}/\delta)}$, which together with (B.4) completes the proof for the upper and lower bounds of $\sum_{q=1}^{B}\lambda_{q}$. The upper bound for the Frobenius norm of $W$ follow by (B.2), while the lower bound follows by $\sum_{q=1}^{B}\lambda_{q}\geq\frac{(c_{B}\wedge 1)d}{8c_{0}^{2}R^{4}\log^{2}(2
B^{2}/\delta)}$ and (B.3).

For the trace term, we have,

$$ $\displaystyle\operatorname*{tr}(W_{\mathsf{MM}})$ $\displaystyle=\operatorname*{tr}\left(\sum_{\tau=1}^{B}\lambda_{\tau}y_{\tau} \widehat{\mu}_{\tau}x_{\tau}^{\top}\right)$ $\displaystyle=\sum_{\tau=1}^{B}\lambda_{\tau}\operatorname*{tr}(y_{\tau} \widehat{\mu}_{\tau}x_{\tau}^{\top})$ $\displaystyle=\sum_{\tau=1}^{B}\lambda_{\tau}y_{\tau}x_{\tau}^{\top}\widehat{ \mu}_{\tau}$ $\displaystyle\overset{(i)}{\geq}\sum_{\tau=1}^{B}\lambda_{\tau}R^{2}/2$ $\displaystyle\overset{(ii)}{\geq}\frac{(c_{B}\wedge 1)d}{16c_{0}^{2}R^{2}\log^ {2}(2B^{2}/\delta)}$ (B.10) $$

Inequality $(i)$ uses (B.1) while inequality $(ii)$ uses the lower bound for $\sum_{\tau}\lambda_{\tau}$. Similarly, by Lemma A.1 we have

$$ $\displaystyle|\langle\widehat{\mu}_{\tau},y_{\tau}x_{\tau}\rangle|\leq R^{2} \left(1+c_{0}\left(\left[1+\frac{1}{\sqrt{N}}\right]\frac{\sqrt{\operatorname* {tr}(\Lambda)}}{R\sqrt{d}}+\frac{\sqrt{\operatorname*{tr}(\Lambda^{2})}}{\sqrt {R^{4}N}}\right)\log(2B/\delta)\right)\leq\frac{3}{2}R^{2},$ $$

where the last inequality uses Assumption (A1). Therefore,

$$ $\displaystyle\operatorname*{tr}(W_{\mathsf{MM}})$ $\displaystyle=\sum_{\tau=1}^{B}\lambda_{\tau}y_{\tau}x_{\tau}^{\top}\widehat{ \mu}_{\tau}$ $\displaystyle\leq\sum_{\tau=1}^{B}\lambda_{\tau}\cdot\frac{3}{2}R^{2}$ $\displaystyle\leq\frac{6d}{R^{2}},$ $$

where the last inequality uses the upper bound for $\sum_{\tau=1}^{B}\lambda_{\tau}$ proved earlier.
∎

###### Lemma B.1 .

###### Proof.

## Appendix C Proof of Theorem 4.1

Our goal is to bound

$$ $\displaystyle R(W_{\mathsf{MM}})$ $\displaystyle=\mathbb{P}_{(x_{i},y_{i})_{1}^{M+1},\ \mu}(\mathrm{sign}( \widehat{y}(E;W_{\mathsf{MM}}))\neq y_{M+1})$ $\displaystyle=\mathbb{P}\left(\left[\frac{1}{M}\sum_{i=1}^{M}y_{i}x_{i}\right] ^{\top}W_{\mathsf{MM}}y_{M+1}x_{M+1}<0\right).$ (C.1) $$

The probability above is over the draws of $\mu$ and of $(x_{i},y_{i})_{i=1}^{M+1}$ and $\mu$. The matrix $W_{\mathsf{MM}}$ is the max-margin solution when training on the in-context tasks (3.3). To prove the risk is close to the label-flipping noise rate $p$, it suffices to show that the error on clean examples is small:

$$ $\displaystyle R(W_{\mathsf{MM}})$ $\displaystyle=\mathbb{P}\left(y_{M+1}\cdot\mathrm{sign}(\widehat{y}(E;W_{ \mathsf{MM}})<0,\,y_{M+1}\neq\tilde{y}_{M+1}\right)$ $\displaystyle\qquad+\mathbb{P}\left(y_{M+1}\cdot\mathrm{sign}(\widehat{y}(E;W_ {\mathsf{MM}})<0,\,y_{M+1}=\tilde{y}_{M+1}\right)$ $\displaystyle\leq p+\mathbb{P}\left(y_{M+1}\cdot\mathrm{sign}(\widehat{y}(E;W_ {\mathsf{MM}})<0,\,y_{M+1}=\tilde{y}_{M+1}\right).$ (C.2) $$

For notational simplicity, let us denote $\widehat{y}:=\widehat{y}(E;W_{\mathsf{MM}})$. Let us also denote the event $\tilde{A}$ as

$$ $\tilde{A}:=\{y_{M+1}\widehat{y}<0\}=\{y_{M+1}\widehat{y}-\mathbb{E}[y_{M+1} \widehat{y}]<-\mathbb{E}[y_{M+1}\widehat{y}]\},$ $$

so that $R(W_{\mathsf{MM}})=\mathbb{P}(\tilde{A})$. Further, let us introduce the sets $\mathcal{C},\mathcal{N}\subset[M]$ as the clean and noisy test examples respectively ($|\mathcal{C}|+|\mathcal{N}|=M$) so that

$$ $y_{i}=\tilde{y}_{i}\,\forall i\in\mathcal{C},\quad y_{i}=-\tilde{y}_{i}\, \forall i\in\mathcal{N}.$ $$

Then we have the identities

$$ $\displaystyle y_{i}x_{i}=\mu+y_{i}z_{i},$ $\displaystyle\quad i\in\mathcal{C},$ $\displaystyle y_{i}x_{i}=-\mu+y_{i}z_{i},$ $\displaystyle\quad i\in\mathcal{N}.$ $$

Therefore there exist independent $\zeta,\zeta_{M+1}\sim\mathsf{N}(0,\Lambda)$ such that

$$ $\frac{1}{M}\sum_{i=1}^{M}y_{i}x_{i}\stackrel{{\scriptstyle\text{d}}}{{=}}\left (1-\frac{2|\mathcal{N}|}{M}\right)\mu+\frac{1}{\sqrt{M}}\zeta,\quad\tilde{y}_{ M+1}x_{M+1}\stackrel{{\scriptstyle\text{d}}}{{=}}\mu+\zeta_{M+1},$ $$

where we have used that $|\mathcal{C}|-|\mathcal{N}|=M-2|\mathcal{N}|$.

In particular, if we define

$$ $A:=\Big{\{}\left((1-\nicefrac{{2|\mathcal{N}|}}{{M}})\mu+M^{-1/2}\zeta\right)^ {\top}W(\mu+\zeta_{M+1})-\mathbb{E}[y_{M+1}\widehat{y}]<-\mathbb{E}[y_{M+1} \widehat{y}]\Big{\}}$ $$

then we have

$$ $\displaystyle\mathbb{P}\left(y_{M+1}\widehat{y}-\mathbb{E}[y_{M+1}\widehat{y}] <-\mathbb{E}[y_{M+1}\widehat{y}],\,y_{M+1}=\tilde{y}_{M+1}\right)=\mathbb{P}(A).$ $$

Continuing from (C.2) this means for any $\alpha_{0},\alpha_{1},\alpha_{2},\alpha_{3}>0$,

$$ $\displaystyle\quad R(W_{\mathsf{MM}})$ $\displaystyle\leq p+\mathbb{P}\left(\left((1-\nicefrac{{2|\mathcal{N}|}}{{M}}) \mu+M^{-1/2}\zeta\right)^{\top}W_{\mathsf{MM}}(\mu+\zeta_{M+1})-\mathbb{E}[y_{ M+1}\widehat{y}]<-\mathbb{E}[y_{M+1}\widehat{y}]\right)$ $\displaystyle\leq p+\mathbb{P}(|\mathcal{N}|/M\leq\alpha_{0})+\mathbb{P}(|M^{- 1/2}\zeta^{\top}W_{\mathsf{MM}}\mu|>\alpha_{1})$ $\displaystyle\quad+\mathbb{P}(|M^{-1/2}\zeta^{\top}W_{\mathsf{MM}}\zeta_{M+1}| >\alpha_{2})+\mathbb{P}((1-\nicefrac{{2|\mathcal{N}|}}{{M}})|\mu^{\top}W_{ \mathsf{MM}}\zeta_{M+1}|>\alpha_{3})$ $\displaystyle\quad+\mathbb{P}\Bigg{(}A\cap\{|\mathcal{N}|/M\leq\alpha_{0}\} \cap\{|M^{-1/2}\zeta^{\top}W_{\mathsf{MM}}\mu|\leq\alpha_{1}\}$ $\displaystyle\qquad\quad\cap\{|M^{-1/2}\zeta^{\top}W_{\mathsf{MM}}\zeta_{M+1}| \leq\alpha_{2}\}\cap\{|\mu^{\top}W_{\mathsf{MM}}\zeta_{M+1}|\leq\alpha_{3}\} \Bigg{)}.$ (C.3) $$

In particular, to derive an upper bound on the risk it suffices to derive a lower bound on $\mu^{\top}W_{\mathsf{MM}}\mu$ and upper bounds on each of the absolute values of the four quantities $|\mathcal{N}|/M$, $\zeta^{\top}W_{\mathsf{MM}}\mu,\zeta^{\top}W_{\mathsf{MM}}\zeta_{M+1}$ and $\mu^{\top}W\zeta_{N+1}$. We shall do so in what follows.

### C.1 | 𝒩 | / M 𝒩 𝑀 |\mathcal{N}|/M | caligraphic_N | / italic_M

###### Lemma C.1 .

There is some constant $c>0$ such that for any $t\geq 0$, we have

$$ $\displaystyle\begin{cases}|\mathcal{N}|=0,&p=0,\\ \mathbb{P}\left(\left|\frac{|\mathcal{N}|}{M}-p\right|\geq t\right)\leq 2\exp( -2t^{2}M),&p>0.\end{cases}$ (C.4) $$

###### Proof.

If $p=0$, there is no label flipping noise and so $|\mathcal{N}|=0$ deterministically. For the $p>0$ case, by definition, $|\mathcal{N}|=\sum_{i=1}^{M}\mathbbm{1}(y_{i}\neq\tilde{y}_{i})$ is a sum of $M$ independent random variables bounded between 0 and 1 with mean $p$. By Hoeffding’s inequality, for any $u\geq 0$ we have

$$ $\mathbb{P}(|\ |\mathcal{N}|-Mp\ |\geq u)=\mathbb{P}(|\ |\mathcal{N}|/M-p\ | \geq u/M)\leq 2\exp\left(-\frac{2u^{2}}{M}\right).$ $$

Setting $u=Mt$ completes the proof.
∎

###### Lemma C.1 .

###### Proof.

### C.2 μ ⊤ ⁢ W ⁢ μ superscript 𝜇 top 𝑊 𝜇 \mu^{\top}W\mu italic_μ start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_μ

The quantity $\mu^{\top}W\mu$ is a quadratic form, and we can provide a concentration inequality for this in the following lemma.

###### Lemma C.2 (Hanson-Wright for uniform on the sphere) .

Let $\tilde{R}>0$ and $Q\in\tilde{\mathbb{R}}^{d\times d}$ be a matrix. If $\mu\sim\mathsf{Unif}(\tilde{R}\cdot\mathbb{S}^{d-1})$, then for any $t\geq 0$,

$$ $\mathbb{P}\left(\left|\mu^{\top}Q\mu-\frac{\tilde{R}^{2}}{d}\operatorname*{tr} (Q)\right|\geq t\right)=\mathbb{P}\left(\left|\mu^{\top}Q\mu-\mathbb{E}[\mu^{ \top}Q\mu]\right|\geq t\right)\leq 2\exp\left(-c\min\left(\frac{t^{2}d^{2}}{ \tilde{R}^{4}\|Q\|_{F}^{2}},\frac{td}{\tilde{R}^{2}\|Q\|_{2}}\right)\right).$ $$

Since $\mu^{\top}Q\mu$ is a quadratic form, we would like to use the Hanson-Wright inequality here. But the trouble is that for $\mu\sim\tilde{R}\mathsf{Unif}(\mathbb{S}^{d-1})$, the components of $\mu$ are not independent, so the standard Hanson-Wright inequality is not directly applicable, but requires some additional work. Our proof instead leverages the following:

###### Theorem C.3 ( \citet adamczak2015hansonwright, Theorem 2.5) .

Let $X$ be a mean-zero random vector in $\mathbb{R}^{d}$. If $X$ satisfies the $K$-convex concentration property, i.e. there exists $K$ such that for every 1-Lipschitz convex function $\phi:\mathbb{R}^{d}\to\mathbb{R}$ we have $\mathbb{E}|\phi(X)|<\infty$ and for every $t>0$,

$$ $\mathbb{P}(|\phi(X)-\mathbb{E}\phi(X)|\geq t)\leq 2\exp(-t^{2}/K^{2}),$ $$

then there exists $c>0$ such that for any matrix $Q\in\mathbb{R}^{d\times d}$ and every $t>0$,

$$ $\mathbb{P}(|X^{\top}QX-\mathbb{E}[X^{\top}QX]|\geq t)\leq 2\exp\left(-c\min \left(\frac{t^{2}}{2K^{4}\|Q\|_{F}^{2}},\frac{t}{K^{2}\|Q\|_{2}}\right)\right).$ $$

###### Proof of Lemma C.2 .

By \citetadamczak2015hansonwright, a sufficient condition for a Hanson-Wright-type inequality to hold is that the random vector $\mu$ satisfies Lipschitz concentration.
By \citet[Theorem 5.1.4]vershynincup we have that for some constant $c>0$, for any $1$-Lipschitz function $f$ on the sphere $\sqrt{d}\cdot\mathbb{S}^{d-1}$ and for $X\sim\mathsf{Unif}(\sqrt{d}\mathbb{S}^{d-1})$,

$$ $\mathbb{P}(|f(X)-\mathbb{E}[f(X)]|\geq t)\leq 2\exp(-ct^{2}).$ $$

Thus if we consider $g:\tilde{R}\cdot\mathbb{S}^{d-1}\to\mathbb{R}$ which is 1-Lipschitz and is defined on the sphere of radius $\tilde{R}$, then $Y=\frac{\tilde{R}}{\sqrt{d}}X$ where $X$ is uniform on the sphere of radius $\tilde{R}$. In particular, the function $g(Y):=g(\frac{\tilde{R}}{\sqrt{d}}X)$ is the composition of a 1-Lipschitz function and a $\tilde{R}/\sqrt{d}$-Lipschitz function defined on the sphere of radius $\sqrt{d}$, and we therefore see that

$$ $\mathbb{P}(|g(Y)-\mathbb{E}g(Y)|\geq t)\leq 2\exp(-ct^{2}d/\tilde{R}^{2}).$ $$

That is, any 1-Lipschitz function on the sphere of radius $\tilde{R}$ satisfies the $K$-convex concentration property with $K=\tilde{R}/\sqrt{d}$. Thus Theorem C.3 implies for any matrix $Q$,

$$ $\displaystyle\mathbb{P}(|\mu^{\top}Q\mu-\mathbb{E}[\mu^{\top}Q\mu]|\geq t)\leq 2 \exp\left(-c\min\left(\frac{t^{2}d^{2}}{\tilde{R}^{4}\|Q\|_{F}^{2}},\frac{td}{ \tilde{R}^{2}\|Q\|_{2}}\right)\right).$ $$

Finally, note that by rotational symmetry we have $\mathbb{E}[\mu]=0$ and $\mathbb{E}[\mu_{i}^{2}]=\tilde{R}^{2}/d$ for each $i\in[d]$. Additionally, for distinct components $i\neq j$, $\mathbb{E}[\mu_{i}\mu_{j}]=\mathbb{E}[\mu_{i}(-\mu_{j})]=0$ and hence $\mathbb{E}[\mu^{\top}Q\mu]=\sum_{i=1}^{d}\mathbb{E}[\mu_{i}^{2}Q_{ii}]=\frac{
\tilde{R}^{2}}{d}\operatorname*{tr}(Q)$.
∎

###### Lemma C.2 (Hanson-Wright for uniform on the sphere) .

###### Theorem C.3 ( \citet adamczak2015hansonwright, Theorem 2.5) .

###### Proof of Lemma C.2 .

### C.3 μ ⊤ ⁢ W ⁢ ζ superscript 𝜇 top 𝑊 𝜁 \mu^{\top}W\zeta italic_μ start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_ζ and μ ⊤ ⁢ W ⁢ ζ M + 1 superscript 𝜇 top 𝑊 subscript 𝜁 𝑀 1 \mu^{\top}W\zeta_{M+1} italic_μ start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_ζ start_POSTSUBSCRIPT italic_M + 1 end_POSTSUBSCRIPT

###### Lemma C.4 .

Let $\mu\sim\tilde{R}\cdot\mathsf{Unif}(\mathbb{S}^{d-1})$, $g\sim\mathsf{N}(0,I_{d})$ (independent of $\mu$) and let $Q\in\mathbb{R}^{d\times d}$ be a matrix. There is an absolute constant $c>0$ such that for any $t\geq 0$,

$$ $\mathbb{P}\left(|\mu^{\top}Qg|\geq t\right)\leq 2\exp\left(-\frac{ct\sqrt{d}}{ \tilde{R}\|Q\|_{F}}\right).$ $$

###### Proof.

The uniform distribution on the sphere of radius $\tilde{R}$ is sub-Gaussian with sub-Gaussian norm satisfying $\|\mu\|_{\psi_{2}}\leq c\tilde{R}/\sqrt{d}$ by \citet[Theorem 3.4.6]vershynincup, while $\|g\|_{\psi_{2}}\leq c$, for some absolute constant $c>0$. By \citet[Lemma 6.2.3]vershynincup, this implies that for independent $g_{1},g_{2}\sim\mathsf{N}(0,I_{d})$ and any $\beta\in\mathbb{R}$,

$$ $\displaystyle\mathbb{E}\exp\left(\beta\frac{\sqrt{d}}{\tilde{R}}\mu^{\top}Qg \right)\leq\mathbb{E}\exp\left(c_{1}\beta g_{1}^{\top}Qg_{2}\right).$ $$

Then using the moment-generating function of Gaussian chaos \citep[Lemma 6.2.2]vershynincup, for $\beta$ satisfying $|\beta|\leq c_{2}/\|Q\|_{2}$ we have

$$ $\displaystyle\mathbb{E}\exp\left(\beta\frac{\sqrt{d}}{\tilde{R}}\mu^{\top}Qg\right)$ $\displaystyle\leq\mathbb{E}\exp\left(c_{1}\beta g_{1}^{\top}Qg_{2}\right)$ $\displaystyle\leq\exp(c_{3}\beta^{2}\|Q\|_{F}^{2}).$ $$

That is, the random variable $\tilde{R}^{-1}\sqrt{d}\mu^{\top}Qg$ is mean-zero and has sub-exponential norm at most $\max(c_{2}^{-1},c_{3})\|Q\|_{F}$. There is therefore a constant $c_{4}>0$ such that for any $u\geq 0$,

$$ $\mathbb{P}(|\tilde{R}^{-1}\sqrt{d}\mu^{\top}Qg|\geq u)=\mathbb{P}\left(|\mu^{ \top}Qg|\geq\frac{\tilde{R}u}{\sqrt{d}}\right)\leq 2\exp(-cu/\|Q\|_{F}).$ $$

Setting $u=t\sqrt{d}/\tilde{R}$ we get

$$ $\displaystyle\mathbb{P}\left(|\mu^{\top}Qg|\geq t\right)\leq 2\exp\left(-\frac {ct\sqrt{d}}{\tilde{R}\|Q\|_{F}}\right)$ $$

∎

###### Lemma C.4 .

###### Proof.

### C.4 ζ ⊤ ⁢ W ⁢ ζ M + 1 superscript 𝜁 top 𝑊 subscript 𝜁 𝑀 1 \zeta^{\top}W\zeta_{M+1} italic_ζ start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_ζ start_POSTSUBSCRIPT italic_M + 1 end_POSTSUBSCRIPT

###### Lemma C.5 .

Let $\zeta,\zeta^{\prime}\stackrel{{\scriptstyle\mathrm{i.i.d.}}}{{\sim}}\mathsf{N}
(0,\Lambda)$ and let $Q\in\mathbb{R}^{d\times d}$ be a matrix. There is a constant $c>0$ such that for all $t\geq 0$,

$$ $\displaystyle\mathbb{P}(|\zeta^{\top}Q\zeta^{\prime}|\geq t)\leq 2\exp\left(- \frac{ct}{\|\Lambda^{1/2}Q\Lambda^{1/2}Q\|_{F}}\right)\leq 2\exp\left(-\frac{ ct}{\|\Lambda\|_{2}\|Q\|_{F}}\right).$ $$

###### Proof.

We can write $\zeta^{\top}Q\zeta_{M+1}$ as $g_{1}^{\top}\Lambda^{1/2}Q\Lambda^{1/2}g_{2}$ where $g_{1},g_{2}\sim\mathsf{N}(0,I_{d})$ are independent, i.e. it is a Gaussian chaos random variable. By \citet[Lemma 6.2.2]vershynincup this random variable is sub-exponential with sub-exponential norm at most $c\|\Lambda^{1/2}Q\Lambda^{1/2}\|_{F}$. Then standard sub-exponential concentration (e.g. \citet[Proposition 2.7.1]vershynincup) completes the proof.

∎

###### Lemma C.5 .

###### Proof.

### C.5 𝔼 ⁢ [ μ ⊤ ⁢ W ⁢ μ ] 𝔼 delimited-[] superscript 𝜇 top 𝑊 𝜇 \mathbb{E}[\mu^{\top}W\mu] blackboard_E [ italic_μ start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_μ ]

###### Lemma C.6 .

On a good run, for any $c_{B}>0$ and for $C>1$ sufficiently large under assumptions (A1) through (A3), the max-margin solution $W_{\mathsf{MM}}$ satisfies

$$ $\frac{(c_{B}\wedge 1)\tilde{R}^{2}}{16c_{0}^{2}R^{2}\log^{2}(2B^{2}/\delta)} \leq\mathbb{E}[\mu^{\top}W_{\mathsf{MM}}\mu]\leq\frac{6\tilde{R}^{2}}{R^{2}}.$ $$

###### Proof.

By definition,

$$ $\displaystyle\mathbb{E}\left[\mu^{\top}W_{\mathsf{MM}}\mu\right]$ $\displaystyle=\frac{\tilde{R}^{2}}{d}\operatorname*{tr}(W_{\mathsf{MM}}).$ (C.5) $$

We therefore need only upper and lower bounds on $\operatorname*{tr}(W_{\mathsf{MM}})$. Using the bounds on $\operatorname*{tr}(W_{\mathsf{MM}})$ from Lemma B.1, we get the desired inequalities.
∎

###### Lemma C.6 .

###### Proof.

### C.6 Putting it all together

For notational simplicity let us write $W=W_{\mathsf{MM}}$. Returning to the decomposition (C.3) and using (C.2), we have

$$ $\displaystyle R(W)$ $\displaystyle\leq p+\mathbb{P}\Bigg{(}(1-\nicefrac{{2|\mathcal{N}|}}{{M}})\mu^ {\top}W\mu-\mathbb{E}[y\widehat{y}]$ $\displaystyle\qquad\qquad<-\mathbb{E}[y\widehat{y}]-M^{-1/2}\zeta^{\top}W\mu-( 1-\nicefrac{{2|\mathcal{N}|}}{{M}})\mu^{\top}W\zeta_{M+1}-M^{-1/2}\zeta^{\top} W\zeta_{M+1}\Bigg{)}$ (C.6) $$

We’ll consider two cases, depending upon whether the label-flipping noise rate $p=0$ or $p>0$. For notational simplicity let’s denote $\rho$ as the quantity

$$ $\rho:=\frac{c_{B}\wedge 1}{16c_{0}^{2}\log^{2}(2B^{2}/\delta)}\in(0,1).$ $$

Then Lemma C.6 states

$$ $\rho\cdot\frac{\tilde{R}^{2}}{R^{2}}\leq\mathbb{E}[\mu^{\top}W\mu]\leq\frac{6 \tilde{R}^{2}}{R^{2}}.$ (C.7) $$

#### Noiseless case ( p = 0 ) 𝑝 0 (p=0) ( italic_p = 0 ) .

Since $p=0$ we know $|\mathcal{N}|=0$. From (C.6) we have

$$ $\displaystyle\qquad R(W)$ $\displaystyle=\mathbb{P}\Bigg{(}\mu^{\top}W\mu-\mathbb{E}[\mu^{\top}W\mu]<- \mathbb{E}[\mu^{\top}W\mu]-M^{-1/2}\zeta^{\top}W\mu-\mu^{\top}W\zeta_{M+1}-M^{ -1/2}\zeta^{\top}W\zeta_{M+1}\Bigg{)}$ $\displaystyle\overset{(i)}{\leq}\mathbb{P}\Bigg{(}\mu^{\top}W\mu-\mathbb{E}[ \mu^{\top}W\mu]<-\frac{\rho\tilde{R}^{2}}{R^{2}}-M^{-1/2}\zeta^{\top}W\mu-\mu^ {\top}W\zeta_{M+1}-M^{-1/2}\zeta^{\top}W\zeta_{M+1}\Bigg{)}$ $\displaystyle\leq\mathbb{P}\Bigg{(}\mu^{\top}W\mu-\mathbb{E}[\mu^{\top}W\mu]<- \frac{\rho\tilde{R}^{2}}{2R^{2}}\Bigg{)}+\mathbb{P}\left(|M^{-1/2}\zeta^{\top} W\mu|\geq\frac{\rho\tilde{R}^{2}}{8R^{2}}\right)+\mathbb{P}\left(\left|\mu^{ \top}W\zeta_{M+1}\right|\geq\frac{\rho\tilde{R}^{2}}{8R^{2}}\right)$ $\displaystyle\qquad+\mathbb{P}\left(\left|M^{-1/2}\zeta^{\top}W\zeta_{M+1} \right|\geq\frac{\rho\tilde{R}^{2}}{8R^{2}}\right).$ (C.8) $$

Inequality $(i)$ uses Lemma C.6. We now proceed by bounding each of the remaining terms in the inequality above.

For the first term we can use Lemma C.2 with $t=\rho\tilde{R}^{2}/2R^{2}$. To do so we need to examine the quantity $\frac{td}{\tilde{R}^{2}\|W\|_{F}}$, which we have

$$ $\displaystyle\frac{td}{\tilde{R}^{2}\|W\|_{F}}$ $\displaystyle=\frac{\rho d}{2R^{2}\|W\|_{F}}\geq\frac{\rho d}{2R^{2}\cdot 2 \sqrt{d}R^{-2}}=\frac{\rho\sqrt{d}}{4}.$ $$

Let’s assume for the moment that $\rho\sqrt{d}/4>1$. We can then apply Lemma C.2 (noting that this lemma holds if we replace the $\|Q\|_{2}$ with $\|Q\|_{F}$),

$$ $\displaystyle\mathbb{P}\Bigg{(}\mu^{\top}W\mu-\mathbb{E}[\mu^{\top}W\mu]<- \frac{\rho\tilde{R}^{2}}{2R^{2}}\Bigg{)}$ $\displaystyle\leq 2\exp\left(-c\min\left(\frac{d^{2}}{\tilde{R}^{4}\|W\|_{F}^{ 2}}\cdot\frac{\rho^{2}\tilde{R}^{4}}{4R^{4}},\frac{d}{\tilde{R}^{2}\|W\|_{F}} \cdot\frac{\rho\tilde{R}^{2}}{2R^{2}}\right)\right)$ $\displaystyle\leq 2\exp\left(-\frac{cd\rho}{2\|W\|_{F}R^{2}}\right)$ $\displaystyle\overset{(i)}{\leq}2\exp\left(-\frac{c\rho\sqrt{d}}{4}\right).$ (C.9) $$

Note that if $\rho\sqrt{d}/4\leq 1$ then the above bound still holds since $\mathbb{P}(\cdot)\leq 1$ is a trivial inequality. This completes the bound for the first term in (C.8).

The second and third terms in (C.8) can be bounded using Lemma C.4: since $\zeta_{M+1}=\Lambda^{1/2}g$ for $g\sim\mathsf{N}(0,I_{d})$,

$$ $\displaystyle\mathbb{P}\left(|\mu^{\top}W\zeta_{M+1}|\geq\frac{\rho\tilde{R}^{ 2}}{8R^{2}}\right)$ $\displaystyle=\mathbb{P}\left(|\mu^{\top}W\Lambda^{1/2}g|\geq\frac{\rho\tilde{ R}^{2}}{8R^{2}}\right)$ $\displaystyle\leq 2\exp\left(-\frac{c\sqrt{d}}{\tilde{R}\|W\Lambda^{1/2}\|_{F} }\cdot\frac{\rho\tilde{R}^{2}}{8R^{2}}\right)$ $\displaystyle\leq 2\exp\left(-\frac{c\sqrt{d}}{\tilde{R}\|\Lambda^{1/2}\|_{2} \|W\|_{F}}\cdot\frac{\rho\tilde{R}^{2}}{8R^{2}}\right)$ $\displaystyle\overset{(i)}{\leq}2\exp\left(-\frac{c\rho\tilde{R}\sqrt{d}}{8R^{ 2}\|\Lambda^{1/2}\|_{2}\cdot 2\sqrt{d}R^{-2}}\right)$ $\displaystyle=2\exp\left(-\frac{c\rho\tilde{R}}{16\|\Lambda^{1/2}\|_{2}}\right).$ (C.10) $$

The inequality $(i)$ uses the upper bound $\|W\|\leq 2\sqrt{d}/R^{2}$ from Lemma B.1.

For the final term in (C.8) we can use Lemma C.5,

$$ $\displaystyle\mathbb{P}\left(\left|M^{-1/2}\zeta^{\top}W\zeta_{M+1}\right|\geq \frac{\rho\tilde{R}^{2}}{8R^{2}}\right)$ $\displaystyle=\mathbb{P}\left(\left|\zeta^{\top}W\zeta_{M+1}\right|\geq\frac{ \rho M^{1/2}\tilde{R}^{2}}{8R^{2}}\right)$ $\displaystyle\leq 2\exp\left(-\frac{c}{\|\Lambda\|_{2}\|W\|_{F}}\cdot\frac{ \rho M^{1/2}\tilde{R}^{2}}{8R^{2}}\right)$ $\displaystyle\overset{(i)}{\leq}2\exp\left(-\frac{c}{\|\Lambda\|_{2}\cdot 2 \sqrt{d}R^{-2}}\cdot\frac{\rho M^{1/2}\tilde{R}^{2}}{8R^{2}}\right)$ $\displaystyle=2\exp\left(-\frac{c\rho M^{1/2}\tilde{R}^{2}}{16\|\Lambda\|_{2} \sqrt{d}}\right).$ (C.11) $$

Again $(i)$ uses Lemma C.5.

Putting together (C.9), (C.10), (C.11) we get

$$ $\displaystyle R(W)$ $\displaystyle\leq 2\exp\left(-\frac{c\rho\sqrt{d}}{4}\right)+4\exp\left(-\frac {c\rho\tilde{R}}{16\|\Lambda^{1/2}\|_{2}}\right)+2\exp\left(-\frac{c\rho M^{1/ 2}\tilde{R}^{2}}{16\|\Lambda\|_{2}\sqrt{d}}\right).$ $$

#### Noisy case.

Returning to (C.6): as before, the ‘signal’ in the problem comes from the term $(1-2|\mathcal{N}|/M)\mu^{\top}W\mu$, which ideally is large and positive. It is natural that we should require more samples $M$ the closer the noise rate $p$ gets to $1/2$ (namely, the smaller $c_{p}$ is), since otherwise with nontrivial probability we will see more noisy examples than clean ones and learning should be impossible. To this end, let’s assume that $p\leq\frac{1}{2}-c_{p}$ for some absolute constant $c_{p}\in(0,1/2)$. Continuing from (C.4) which shows that

$$ $\mathbb{P}\left(1-2|\mathcal{N}|/M>\frac{1}{2}c_{p}\right)=\mathbb{P}\left( \frac{2|\mathcal{N}|}{M}-1+2c_{p}\leq\frac{6}{4}c_{p}\right)\leq 2\exp\left(- \frac{18c_{p}^{2}M}{16}\right)\leq 2\exp(-c_{p}^{2}M).$ (C.12) $$

Let’s call the event

$$ $\mathcal{E}:=\left\{1-\frac{2|\mathcal{N}|}{M}>\frac{1}{2}c_{p}\right\}.$ $$

Continuing from (C.6), since $1-2|\mathcal{N}|/M>0$ on $\mathcal{E}$ we have

$$ $\displaystyle\qquad R(W)\leq p$ $\displaystyle\quad+\mathbb{P}\left(\mu^{\top}W\mu-\frac{\mathbb{E}[y\widehat{y }]}{(1-\nicefrac{{2|\mathcal{N}|}}{{M}})}<-\frac{\mathbb{E}[y\widehat{y}]-M^{- 1/2}\zeta^{\top}W\mu-(1-\nicefrac{{2|\mathcal{N}|}}{{M}})\mu^{\top}W\zeta_{M+1 }-M^{-1/2}\zeta^{\top}W\zeta_{M+1}}{(1-\nicefrac{{2|\mathcal{N}|}}{{M}})}\right)$ $\displaystyle\leq p+\mathbb{P}(\mathcal{E}^{c})$ $\displaystyle\quad+\mathbb{P}\left(\mathcal{E},\ \mu^{\top}W\mu-\frac{\mathbb{ E}[y\widehat{y}]}{(1-\nicefrac{{2|\mathcal{N}|}}{{M}})}<-\frac{\mathbb{E}[y \widehat{y}]-M^{-1/2}\zeta^{\top}W\mu-(1-\nicefrac{{2|\mathcal{N}|}}{{M}})\mu^ {\top}W\zeta_{M+1}-M^{-1/2}\zeta^{\top}W\zeta_{M+1}}{(1-\nicefrac{{2|\mathcal{ N}|}}{{M}})}\right)$ $\displaystyle\overset{(i)}{\leq}p+2\exp(-c_{p}^{2}M)$ $\displaystyle\quad+\mathbb{P}\left(\mathcal{E},\ \mu^{\top}W\mu-\frac{\mathbb{ E}[y\widehat{y}]}{(1-\nicefrac{{2|\mathcal{N}|}}{{M}})}<-\frac{\mathbb{E}[y \widehat{y}]-M^{-1/2}\zeta^{\top}W\mu-(1-\nicefrac{{2|\mathcal{N}|}}{{M}})\mu^ {\top}W\zeta_{M+1}-M^{-1/2}\zeta^{\top}W\zeta_{M+1}}{(1-\nicefrac{{2|\mathcal{ N}|}}{{M}})}\right)$ $\displaystyle\leq p+2\exp(-c_{p}^{2}M)$ $\displaystyle\quad+\mathbb{P}\Bigg{(}\mathcal{E},\ \mu^{\top}W\mu-\frac{ \mathbb{E}[y\widehat{y}]}{(1-\nicefrac{{2|\mathcal{N}|}}{{M}})}$ $\displaystyle\qquad\qquad\qquad<\frac{-\mathbb{E}[y\widehat{y}]}{1-2|\mathcal{ N}|/M}+2c_{p}^{-1}\left[M^{-1/2}|\zeta^{\top}W\mu|+|1-\nicefrac{{2|\mathcal{N} |}}{{M}}|\cdot|\mu^{\top}W\zeta_{M+1}|+|M^{-1/2}\zeta^{\top}W\zeta_{M+1}| \right]\Bigg{)}.$ (C.13) $$

The inequality $(i)$ uses (C.12).
We want to deal with the quantity $\mu^{\top}W\mu-\frac{\mathbb{E}[\mu^{\top}W\mu]}{1-2|\mathcal{N}|/M}=\frac{
\mathbb{E}[y\widehat{y}]}{(1-2|\mathcal{N}|/M)}$ on the event $\mathcal{E}$. We have,

$$ $\displaystyle\mu^{\top}W\mu-\frac{\tilde{R}^{2}(1-2p)\operatorname*{tr}(W)}{d( 1-2|\mathcal{N}|/M)}$ $\displaystyle=\mu^{\top}W\mu-\frac{\tilde{R}^{2}\operatorname*{tr}(W)}{d}\cdot \frac{1-2p}{1-\nicefrac{{2|\mathcal{N}|}}{{m}}}$ $\displaystyle=\mu^{\top}W\mu-\frac{\tilde{R}^{2}\operatorname*{tr}(W)}{d}+ \frac{\tilde{R}^{2}\operatorname*{tr}(W)}{d}\cdot\left(1-\frac{1-2p}{1- \nicefrac{{2|\mathcal{N}|}}{{m}}}\right)$ $\displaystyle=\mu^{\top}W\mu-\frac{\tilde{R}^{2}\operatorname*{tr}(W)}{d}- \frac{\tilde{R}^{2}\operatorname*{tr}(W)}{d}\cdot\frac{-2p+\nicefrac{{2| \mathcal{N}|}}{{M}}}{1-\nicefrac{{2|\mathcal{N}|}}{{M}}}.$ $$

Continuing from (C.13) this means

$$ $\displaystyle\qquad R(W)\leq p+2\exp(-c_{p}^{2}M)$ $\displaystyle+\mathbb{P}\Bigg{(}\mathcal{E},\ \mu^{\top}W\mu-\frac{\tilde{R}^{ 2}\operatorname*{tr}(W)}{d}-\frac{\tilde{R}^{2}\operatorname*{tr}(W)}{d}\cdot \frac{-2p+\nicefrac{{2|\mathcal{N}|}}{{M}}}{1-\nicefrac{{2|\mathcal{N}|}}{{M}}}$ $\displaystyle\qquad\qquad<\frac{-\tilde{R}^{2}(1-2p)\operatorname*{tr}(W)}{d(1 -2|\mathcal{N}|/M)}+2c_{p}^{-1}\left[M^{-1/2}|\zeta^{\top}W\mu|+|1-\nicefrac{{ 2|\mathcal{N}|}}{{M}}|\cdot|\mu^{\top}W\zeta_{M+1}|+|M^{-1/2}\zeta^{\top}W \zeta_{M+1}|\right]\Bigg{)}$ $\displaystyle\qquad=p+2\exp(-c_{p}^{2}M)$ $\displaystyle+\mathbb{P}\Bigg{(}\mathcal{E},\ \mu^{\top}W\mu-\mathbb{E}[\mu^{ \top}W\mu]$ $\displaystyle\qquad\qquad<\frac{-\tilde{R}^{2}\operatorname*{tr}(W)}{d}+2c_{p} ^{-1}\left[M^{-1/2}|\zeta^{\top}W\mu|+|1-\nicefrac{{2|\mathcal{N}|}}{{M}}| \cdot|\mu^{\top}W\zeta_{M+1}|+|M^{-1/2}\zeta^{\top}W\zeta_{M+1}|\right]\Bigg{)}$ $\displaystyle\qquad\overset{(i)}{\leq}p+2\exp(-c_{p}^{2}M)$ $\displaystyle+\mathbb{P}\Bigg{(}\mathcal{E},\ \mu^{\top}W\mu-\mathbb{E}[\mu^{ \top}W\mu]$ $\displaystyle\qquad\qquad<-\rho\cdot\frac{\tilde{R}^{2}}{R^{2}}+2c_{p}^{-1} \left[M^{-1/2}|\zeta^{\top}W\mu|+|1-\nicefrac{{2|\mathcal{N}|}}{{M}}|\cdot|\mu ^{\top}W\zeta_{M+1}|+|M^{-1/2}\zeta^{\top}W\zeta_{M+1}|\right]\Bigg{)}$ $\displaystyle\leq p+2\exp(-c_{p}^{2}M)+\mathbb{P}(\mu^{\top}W\mu-\mathbb{E}[ \mu^{\top}W\mu]<-\frac{1}{2}\rho\cdot\frac{\tilde{R}^{2}}{R^{2}})+\mathbb{P} \left(2c_{p}^{-1}M^{-1/2}|\zeta^{\top}W\mu|>\frac{\rho\tilde{R}^{2}}{8R^{2}}\right)$ $\displaystyle\qquad+\mathbb{P}\left(|1-2|\mathcal{N}|/M|\cdot|\mu^{\top}W\zeta _{M+1}|\geq\frac{\rho\tilde{R}^{2}}{8R^{2}}\right)+\mathbb{P}\left(|M^{-1/2} \zeta^{\top}W\zeta_{M+1}|\geq\frac{\rho\tilde{R}^{2}}{8R^{2}}\right).$ (C.14) $$

Inequality $(i)$ uses (C.7). From here the proof is exactly the same as in the clean case. For the first term, we use similar arguments used derive (C.9). To apply Lemma C.2, we need to examine the quantity $\frac{td}{\tilde{R}^{2}\|W\|_{F}}$ when $t=\rho\tilde{R}^{2}/2R^{2}$: we have,

$$ $\displaystyle\frac{td}{\tilde{R}^{2}\|W\|_{F}}=\frac{\rho d}{2R^{2}\|W\|_{F}} \geq\frac{\rho d}{2R^{2}\cdot 2\sqrt{d}R^{-2}}=\frac{\rho\sqrt{d}}{4}.$ (C.15) $$

Again if $\rho\sqrt{d}/4>1$ then we have by Lemma C.2,

$$ $\displaystyle\mathbb{P}\Bigg{(}\mu^{\top}W\mu-\mathbb{E}[\mu^{\top}W\mu]<- \frac{\rho\tilde{R}^{2}}{2R^{2}}\Bigg{)}$ $\displaystyle\leq 2\exp\left(-c\min\left(\frac{d^{2}}{\tilde{R}^{4}\|W\|_{F}^{ 2}}\cdot\frac{\rho^{2}\tilde{R}^{4}}{4R^{4}},\frac{d}{\tilde{R}^{2}\|W\|_{F}} \cdot\frac{\rho\tilde{R}^{2}}{2R^{2}}\right)\right)$ $\displaystyle\leq 2\exp\left(-\frac{cd\rho}{2\|W\|_{F}R^{2}}\right)$ $\displaystyle\overset{(i)}{\leq}2\exp\left(-\frac{c\rho\sqrt{d}}{4}\right).$ (C.16) $$

Inequality $(i)$ uses (C.15). Note that if $\rho\sqrt{d}/4\leq 1$ then the above bound still holds since $\mathbb{P}(\cdot)\leq 1$ is a trivial inequality.

As for the $\mu^{\top}W\zeta_{M+1}$ and $\zeta^{\top}W\mu$ terms, similarly as to the analysis used to derive (C.10) via Lemma C.4, we have,

$$ $\displaystyle\mathbb{P}\left(2c_{p}^{-1}|\mu^{\top}W\zeta_{M+1}|\geq\frac{\rho \tilde{R}^{2}}{8R^{2}}\right)$ $\displaystyle=\mathbb{P}\left(|\mu^{\top}W\Lambda^{1/2}g|>\frac{c_{p}\rho \tilde{R}^{2}}{16R^{2}}\right)$ $\displaystyle(g\sim\mathsf{N}(0,I))$ $\displaystyle\leq 2\exp\left(-\frac{cc_{p}\rho\tilde{R}}{32\|\Lambda^{1/2}\|}\right)$ (C.17) $$

Note that $|1-2|\mathcal{N}|/M|\leq 1$ always so this takes care of the term involving $\mu^{\top}W\zeta_{M+1}$. Since the analysis used to derive (C.10) only relies upon upper bounds for $\|W\|_{F}=\|W^{\top}\|_{F}$, and since $M\geq 1$ and $c_{p}\in(0,1/2]$, the same analysis holds for the term $|\zeta^{\top}W\mu|$.

For the $\zeta^{\top}W\zeta_{M+1}$ term, the same analysis used for (C.11) applies here as well, since

$$ $\displaystyle\mathbb{P}\left(|M^{-1/2}\zeta^{\top}W\zeta_{M+1}|\geq\frac{\rho \tilde{R}^{2}}{8R^{2}}\right)$ $\displaystyle=\mathbb{P}\left(|\zeta^{\top}W\zeta_{M+1}|\geq\frac{\rho M^{1/2} \tilde{R}^{2}}{8R^{2}}\right)$ $\displaystyle\leq 2\exp\left(-\frac{c\rho M^{1/2}\tilde{R}^{2}}{16\|\Lambda\|_ {2}\sqrt{d}}\right)$ (C.18) $$

Plugging (C.16), (C.17), and (C.18) into (C.14) we get

$$ $\displaystyle R(W)$ $\displaystyle\leq p+2\exp(-c_{p}^{2}M)+2\exp\left(-\frac{c\rho\sqrt{d}}{4} \right)+4\exp\left(-\frac{cc_{p}\rho\tilde{R}}{32\|\Lambda^{1/2}\|}\right)+2 \exp\left(-\frac{c\rho M^{1/2}\tilde{R}^{2}}{16\|\Lambda\|_{2}\sqrt{d}}\right).$ $$

By adjusting the constant $c$ to depend on $c_{p}$, this completes the proof of Theorem 4.1.

## Appendix D Proof of Theorem 4.2

Recall that we assume that $\Lambda=I$ in this section, and for simplicity denote $W=W_{\mathsf{MM}}$. Our goal is to show that for each $k$,

$$ $\widehat{\mu}^{\top}Wy_{k}x_{k}>0.$ $$

We have:

$$ $\displaystyle\widehat{\mu}^{\top}Wy_{k}x_{k}$ $\displaystyle=\left(\frac{1}{M}\sum_{i=1}^{M}y_{i}x_{i}\right)^{\top}Wy_{k}x_{k}$ $\displaystyle=\frac{1}{M}\left[x_{k}^{\top}Wx_{k}+\sum_{i\neq k}y_{i}y_{k}x_{i }^{\top}Wx_{k}\right].$ $$

Our goal here will be to show that the first quantity $x_{k}^{\top}Wx_{k}$ is large and positive and dominates the second term involving $x_{i}^{\top}Wx_{k}$. By definition we have

$$ $\displaystyle x_{k}^{\top}Wx_{k}$ $\displaystyle=(\mu+y_{k}z_{k})^{\top}W(\mu+y_{k}z_{k})$ $\displaystyle=\mu^{\top}W\mu+y_{k}z_{k}^{\top}W\mu+\mu^{\top}Wy_{k}z_{k}+z_{k} ^{\top}Wz_{k}.$ (D.1) $$

In particular we have the identities

$$ $\displaystyle\quad\{\widehat{\mu}^{\top}Wy_{k}x_{k}<0\}$ $\displaystyle=\{x_{k}^{\top}Wx_{k}<-\textstyle\sum_{i\neq k}y_{i}y_{k}x_{i}^{ \top}Wx_{k}\}$ $\displaystyle=\{z_{k}^{\top}Wz_{k}-\operatorname*{tr}(W)<-\operatorname*{tr}(W )-y_{k}z_{k}^{\top}W\mu-\mu^{\top}Wy_{k}z_{k}-\mu^{\top}W\mu-\textstyle\sum_{i \neq k}y_{i}y_{k}x_{i}^{\top}Wx_{k}\}$ $\displaystyle\subset\{z_{k}^{\top}Wz_{k}-\operatorname*{tr}(W)<-\frac{1}{2} \operatorname*{tr}(W)\}\cup\{|z_{k}^{\top}W\mu|>\frac{1}{8}\operatorname*{tr}( W)\}$ $\displaystyle\qquad\cup\{|\mu^{\top}Wz_{k}|>\frac{1}{8}\operatorname*{tr}(W)\} \cup\{\mu^{\top}W\mu<-\frac{1}{8}\operatorname*{tr}(W)\}\cup\left\{\left| \textstyle\sum_{i\neq k}y_{i}y_{k}x_{i}^{\top}Wx_{k}\right|>\frac{1}{8} \operatorname*{tr}(W)\right\}.$ (D.2) $$

We will therefore proceed by bounding the probability each of these events. We will use the same notation from the proof of Theorem 4.1 presented in Appendix C, whereby $i\in\mathcal{C}$ refers to clean examples $(x_{i},y_{i})$ with $\tilde{y}_{i}=y_{i}$, while $i\in\mathcal{N}$ means $y_{i}=-\tilde{y}_{i}$.

### D.1 ∑ i ≠ k x i ⊤ ⁢ W ⁢ x k subscript 𝑖 𝑘 superscript subscript 𝑥 𝑖 top 𝑊 subscript 𝑥 𝑘 \sum_{i\neq k}x_{i}^{\top}Wx_{k} ∑ start_POSTSUBSCRIPT italic_i ≠ italic_k end_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_x start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT term.

We have,

$$ $\displaystyle\sum_{i\neq k}y_{i}x_{i}$ $\displaystyle=\sum_{i\in\mathcal{C},\ i\neq k}y_{i}x_{i}+\sum_{i\in\mathcal{N} ,\ i\neq k}y_{i}x_{i}$ $\displaystyle=\sum_{i\in\mathcal{C},\ i\neq k}(\mu+\tilde{y}_{k}z_{k})+\sum_{i \in\mathcal{N},\ i\neq k}(-\mu-\tilde{y}_{k}z_{k})$ $\displaystyle=(|\mathcal{C}\setminus\{k\}|-|\mathcal{N}\setminus\{k\}|)\mu+ \sum_{i\in\mathcal{C},\ i\neq k}\tilde{y}_{k}z_{k}-\sum_{i\in\mathcal{N},\ i \neq k}\tilde{y}_{k}z_{k}.$ $$

Now since the label-flipping noise is independent of $z_{k}$, the random variables $\tilde{y}_{k}z_{k}$ are i.i.d. standard normals. Therefore there is $g_{k}\sim\mathsf{N}(0,I_{d})$ such that

$$ $\sum_{i\in\mathcal{C},\ i\neq k}\tilde{y}_{k}z_{k}-\sum_{i\in\mathcal{N},\ i \neq k}\tilde{y}_{k}z_{k}=\sqrt{M-1}g_{k}.$ $$

If we denote $N_{k}:=|\mathcal{C}\setminus\{k\}|-|\mathcal{N}\setminus\{k\}|$ then clearly $|N_{k}|\leq M$ and thus we have

$$ $\sum_{i\neq k}y_{i}x_{i}=N_{k}\mu+\sqrt{M-1}g_{k},$ $$

where $|N_{k}|\leq M$ and $g_{k}$ is a standard normal and $g_{k}$ is independent of $z_{k}$ and $y_{k}$. Therefore

$$ $\displaystyle\left|\sum_{i\neq k}y_{i}x_{i}^{\top}Wy_{k}x_{k}\right|$ $\displaystyle=\left|(N_{k}\mu+\sqrt{M-1}g_{k})^{\top}W(\tilde{y}_{k}y_{k}\mu+y _{k}z_{k})\right|$ $\displaystyle\leq M|\mu^{\top}W\mu|+\sqrt{M}|g_{k}^{\top}W\mu|+M|\mu^{\top}Wz_ {k}|+\sqrt{M}|g_{k}^{\top}Wz_{k}|.$ (D.3) $$

In particular we have

$$ $\displaystyle\quad\mathbb{P}\left(\left|\sum_{i\neq k}y_{i}x_{i}^{\top}Wy_{k}x _{k}\right|>\frac{\operatorname*{tr}(W)}{8}\right)\leq\mathbb{P}(M|\mu^{\top}W \mu|>\frac{\operatorname*{tr}(W)}{32})$ $\displaystyle+\mathbb{P}(\sqrt{M}|g_{k}^{\top}W\mu|>\frac{\operatorname*{tr}(W )}{32})+\mathbb{P}(M|\mu^{\top}Wz_{k}|>\frac{\operatorname*{tr}(W)}{32})+ \mathbb{P}(\sqrt{M}|g_{k}^{\top}Wz_{k}|>\frac{\operatorname*{tr}(W)}{32}).$ (D.4) $$

We’ll proceed by bounding each of these terms in sequence.

#### μ ⊤ ⁢ W ⁢ μ superscript 𝜇 top 𝑊 𝜇 \mu^{\top}W\mu italic_μ start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_μ term.

Let $L$ be an integer (we will take $L=M$ in this proof but in a later proof we will take $L=1$). Since $\operatorname*{tr}(W)>0$ by Lemma B.1, we have

$$ $\displaystyle\{|L\mu^{\top}W\mu|>\frac{\operatorname*{tr}(W)}{32}\}$ $\displaystyle=\{|\mu^{\top}W\mu-\frac{\tilde{R}^{2}\operatorname*{tr}(W)}{d}|+ \frac{\tilde{R}^{2}\operatorname*{tr}(W)}{d}>\frac{\operatorname*{tr}(W)}{32L}\}$ $$

This means

$$ $\displaystyle\mathbb{P}\left(|L\mu^{\top}W\mu|>\frac{\operatorname*{tr}(W)}{32 }\right)$ $\displaystyle\leq\mathbb{P}\left(\left|\mu^{\top}W\mu-\frac{\tilde{R}^{2} \operatorname*{tr}(W)}{d}\right|>\frac{\operatorname*{tr}(W)}{32L}\left(1- \frac{32\tilde{R}^{2}L}{d}\right)\right).$ $$

Assume for now $d>128\rho^{-1}\tilde{R}^{2}L$ so that $32\tilde{R}^{2}L/d<1/2$. We then have by Lemma C.2,

$$ $\displaystyle\mathbb{P}\left(|L\mu^{\top}W\mu|>\frac{\operatorname*{tr}(W)}{32 }\right)$ $\displaystyle\leq\mathbb{P}\left(\left|\mu^{\top}W\mu-\frac{\tilde{R}^{2} \operatorname*{tr}(W)}{d}\right|>\frac{\operatorname*{tr}(W)}{64L}\right)$ $\displaystyle\leq 2\exp\left(-c\min\left(\frac{d^{2}}{\tilde{R}^{4}\|W\|_{F}^{ 2}}\left(\frac{\operatorname*{tr}(W)}{64L}\right)^{2},\frac{d}{\tilde{R}^{2}\| W\|_{F}}\left(\frac{\operatorname*{tr}(W)}{64L}\right)\right)\right).$ $$

As for the quantity appearing in the minimum we have by Lemma B.1

$$ $\displaystyle\frac{d\operatorname*{tr}(W)}{64L}\cdot\frac{1}{\tilde{R}^{2}} \geq\frac{d}{64L}\cdot\frac{\rho d}{R^{2}}\cdot\frac{R^{2}}{2\tilde{R}^{2} \sqrt{d}}=\frac{\rho d^{3/2}}{128L\tilde{R}^{2}}\geq\frac{\rho d}{128L\tilde{R }^{2}}.$ $$

We therefore have

$$ $\displaystyle\mathbb{P}\left(|L\mu^{\top}W\mu|>\frac{\operatorname*{tr}(W)}{32 }\right)$ $\displaystyle\leq 2\exp\left(-\frac{c\rho d}{128L\tilde{R}^{2}}\right).$ (D.5) $$

We see that the case $d\leq 128\rho^{-1}\tilde{R}^{2}L$ results in the same inequality to hold (albeit vacuously). This completes the proof for the $\mu^{\top}W\mu$ term in (D.4) by taking $L=M$.

#### g k ⊤ ⁢ W ⁢ μ superscript subscript 𝑔 𝑘 top 𝑊 𝜇 g_{k}^{\top}W\mu italic_g start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_μ and μ ⊤ ⁢ W ⁢ z k superscript 𝜇 top 𝑊 subscript 𝑧 𝑘 \mu^{\top}Wz_{k} italic_μ start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_z start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT terms.

The terms $\mu^{\top}Wg_{k}$ and $z_{k}^{\top}W\mu$ can be controlled using Lemma C.4. For an integer $L\in\mathbb{N}$ we have

$$ $\displaystyle\mathbb{P}\left(L|\mu^{\top}Wg_{k}|\geq\frac{\operatorname*{tr}(W )}{32}\right)$ $\displaystyle\leq 2\exp\left(-\frac{c\sqrt{d}}{\tilde{R}\|W\|_{F}}\cdot\frac{ \operatorname*{tr}(W)}{8L}\right).$ $$

We have

$$ $\displaystyle\frac{\sqrt{d}}{\tilde{R}\|W\|_{F}}\cdot\frac{\operatorname*{tr}( W)}{32L}\geq\frac{R^{2}\sqrt{d}}{\tilde{R}\cdot 2\sqrt{d}}\cdot\frac{\rho d}{3 2R^{2}L}=\frac{\rho d}{64\tilde{R}L}.$ $$

Note that the above argument applies to $W^{\top}$ as well as $W$, so that we have

$$ $\displaystyle\mathbb{P}\left(L|g_{k}^{\top}W\mu|\geq\frac{\operatorname*{tr}(W )}{32}\right)\vee\mathbb{P}\left(L|\mu^{\top}Wg_{k}|\geq\frac{\operatorname*{ tr}(W)}{32}\right)$ $\displaystyle\leq 2\exp\left(-\frac{c\rho d}{64\tilde{R}L}\right).$ (D.6) $$

In particular, we have

$$ $\displaystyle\mathbb{P}\left(\sqrt{M}|g_{k}^{\top}W\mu|\geq\frac{\operatorname *{tr}(W)}{32}\right)$ $\displaystyle\leq 2\exp\left(-\frac{c\rho d}{64\tilde{R}\sqrt{M}}\right),$ $\displaystyle\mathbb{P}\left(M|\mu^{\top}Wz_{k}|\geq\frac{\operatorname*{tr}(W )}{32}\right)$ $\displaystyle\leq 2\exp\left(-\frac{c\rho d}{64\tilde{R}M}\right).$ (D.7) $$

#### g k ⊤ ⁢ W ⁢ z k superscript subscript 𝑔 𝑘 top 𝑊 subscript 𝑧 𝑘 g_{k}^{\top}Wz_{k} italic_g start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_z start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT term.

The last term $g_{k}^{\top}Wz_{k}$ is a Gaussian chaos random variable. By \citet[Lemma 6.2.2]vershynincup this random variable is sub-exponential with sub-exponential norm at most $c\|W\|_{F}$. Therefore we have

$$ $\displaystyle\mathbb{P}(\sqrt{M}|g_{k}^{\top}Wz_{k}|>\frac{\operatorname*{tr}( W)}{32})$ $\displaystyle=\mathbb{P}(|g_{k}^{\top}Wz_{k}|>\frac{\operatorname*{tr}(W)}{32 \sqrt{M}})\leq 2\exp\left(-c\frac{\operatorname*{tr}(W)}{32\sqrt{M}\|W\|_{F}} \right).$ $$

Since $\operatorname*{tr}(W)/\|W\|_{F}\geq\rho\sqrt{d}/2$ by Lemma B.1 this implies

$$ $\displaystyle\mathbb{P}(\sqrt{M}|g_{k}^{\top}Wz_{k}|>\frac{\operatorname*{tr}( W)}{32})\leq 2\exp\left(-\frac{c\rho\sqrt{d}}{64\sqrt{M}}\right).$ (D.8) $$

#### Putting it all together.

Continuing from (D.4) we can use (D.5), (D.7) and (D.8) to get

$$ $\displaystyle\quad\mathbb{P}\left(\left|\sum_{i\neq k}y_{i}x_{i}^{\top}Wy_{k}x _{k}\right|>\frac{\operatorname*{tr}(W)}{8}\right)\leq 2\exp\left(-\frac{c\rho d }{128M\tilde{R}^{2}}\right)$ $\displaystyle+2\exp\left(-\frac{c\rho d}{64\tilde{R}\sqrt{M}}\right)+2\exp \left(-\frac{c\rho d}{64M\tilde{R}}\right)+2\exp\left(-\frac{c\rho\sqrt{d}}{64 \sqrt{M}}\right).$ (D.9) $$

### D.2 x k ⊤ ⁢ W ⁢ x k superscript subscript 𝑥 𝑘 top 𝑊 subscript 𝑥 𝑘 x_{k}^{\top}Wx_{k} italic_x start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_x start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT term.

It remains to bound the probability of the first three events in (D.2).

#### z k ⊤ ⁢ W ⁢ z k superscript subscript 𝑧 𝑘 top 𝑊 subscript 𝑧 𝑘 z_{k}^{\top}Wz_{k} italic_z start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_z start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT term.

First note that since $z_{k}$ is a standard Gaussian, by a standard Hanson-Wright inequality \citep[Lemma 6.2.1]vershynincup we have for some constant $0<c<1$, for any $t\geq 0$,

$$ $\displaystyle\mathbb{P}(|z_{k}^{\top}Wz_{k}-\operatorname*{tr}(W)|\geq t)\leq 2 \exp\left(-c\min\left(\frac{t^{2}}{\|W\|_{F}^{2}},\frac{t}{\|W\|_{F}}\right) \right).$ $$

So for $t=\frac{1}{2}\operatorname*{tr}(W)$ we have

$$ $\displaystyle\frac{t}{\|W\|_{F}}$ $\displaystyle=\frac{\operatorname*{tr}(W)}{2\|W\|_{F}}\geq\frac{\rho d}{R^{2}} \cdot\frac{R^{2}}{4\sqrt{d}}=\frac{\rho\sqrt{d}}{4}.$ $$

This means that if $\rho\sqrt{d}/4\geq 1$ we have by Lemma B.1

$$ $\displaystyle\mathbb{P}\left(|z_{k}^{\top}Wz_{k}-\operatorname*{tr}(W)|\geq \frac{1}{2}\operatorname*{tr}(W)\right)\leq 2\exp\left(-c\rho\sqrt{d}/4\right).$ (D.10) $$

Since the same inequality trivially holds if $\rho\sqrt{d}/4<1$ (by potentially taking $c$ to be a smaller constant), this completes this term.

#### μ ⊤ ⁢ W ⁢ μ superscript 𝜇 top 𝑊 𝜇 \mu^{\top}W\mu italic_μ start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_μ term.

This is covered by (D.5) with $L=1$:

$$ $\displaystyle\mathbb{P}\left(|\mu^{\top}W\mu|>\frac{\operatorname*{tr}(W)}{32}\right)$ $\displaystyle\leq 2\exp\left(-\frac{c\rho d}{128\tilde{R}^{2}}\right).$ (D.11) $$

#### z k ⊤ ⁢ W ⁢ μ superscript subscript 𝑧 𝑘 top 𝑊 𝜇 z_{k}^{\top}W\mu italic_z start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT italic_W italic_μ terms.

Since $g_{k}$ and $z_{k}$ are i.i.d., this is covered by (D.6) with $L=1$:

$$ $\displaystyle\mathbb{P}\left(|z_{k}^{\top}W\mu|\geq\frac{\operatorname*{tr}(W) }{32}\right)\vee\mathbb{P}\left(|\mu^{\top}Wz_{k}|\geq\frac{\operatorname*{tr} (W)}{32}\right)$ $\displaystyle\leq 2\exp\left(-\frac{c\rho d}{64\tilde{R}}\right).$ (D.12) $$

#### Putting it all together.

Continuing from (D.2) we have

$$ $\displaystyle\quad\mathbb{P}\left(\widehat{\mu}^{\top}Wy_{k}x_{k}<0\right)$ $\displaystyle\leq\mathbb{P}\left(z_{k}^{\top}Wz_{k}-\operatorname*{tr}(W)<- \frac{1}{2}\operatorname*{tr}(W)\right)+\mathbb{P}\left(|z_{k}^{\top}W\mu|> \frac{1}{8}\operatorname*{tr}(W)\right)$ $\displaystyle\qquad+\mathbb{P}\left(|\mu^{\top}Wz_{k}|>\frac{1}{8} \operatorname*{tr}(W)\right)+\mathbb{P}\left(\mu^{\top}W\mu<-\frac{1}{8} \operatorname*{tr}(W)\right)+\mathbb{P}\left(\left|\textstyle\sum_{i\neq k}y_{ i}y_{k}x_{i}^{\top}Wx_{k}\right|>\frac{1}{8}\operatorname*{tr}(W)\right)$ $\displaystyle\leq 2\exp\left(-\frac{c\rho\sqrt{d}}{4}\right)+4\exp\left(-\frac {c\rho d}{64\tilde{R}}\right)+2\exp\left(-\frac{c\rho d}{128M\tilde{R}^{2}}\right)$ $\displaystyle\qquad+2\exp\left(-\frac{c\rho d}{64\tilde{R}\sqrt{M}}\right)+2 \exp\left(-\frac{c\rho d}{64M\tilde{R}}\right)+2\exp\left(-\frac{c\rho\sqrt{d} }{64\sqrt{M}}\right)$ $\displaystyle\leq 4\exp\left(-\frac{c\rho\sqrt{d}}{64\sqrt{M}}\right)+8\exp \left(-\frac{c\rho d}{128M(\tilde{R}^{2}\vee\tilde{R})}\right).$ $$

Using a union bound this allows for us to conclude

$$ $\displaystyle\mathbb{P}\left(\exists k\in[M]\ \text{ s.t.}\ \widehat{\mu}^{ \top}Wy_{k}x_{k}<0\right)\leq 4M\exp\left(-\frac{c\rho\sqrt{d}}{64\sqrt{M}} \right)+8M\exp\left(-\frac{c\rho d}{128M(\tilde{R}^{2}\vee\tilde{R})}\right),$ $$

as claimed in the theorem.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Practical Implications: Recognizing the limitations of the simplified setting, can the authors discuss any potential practical insights for realistic problems?

Exploration of Alternative Architectures: How would the results generalize to more complex transformer models? 

Pre-Training with Noisy Labels: The study assumes clean input data and labels during pre-training. How would the result change once noise is introduced in the pre-training data? 

Numerical Validation of Theoretical Results: Have the authors tried conducting numerical experiments to validate your theoretical findings? It would be highly insightful to see how well the theoretical predictions align with empirical observations in both the simplified setting used for analysis and in more realistic settings involving complex transformer architectures and real-world datasets. For instance, could the authors present experimental results demonstrating the relationship between the number of pre-training tasks, the signal-to-noise ratio, and the generalization performance, as suggested by your theoretical bounds? Furthermore, showcasing empirical examples of benign overfitting in-context would be compelling evidence to support this key finding.
soundness: 3
strengths: Rigorous theoretical analysis: The paper offers a detailed analysis of the implicit regularization of gradient descent during pre-training and leverages the max-margin framework to derive generalization guarantees.

Novel insights: The paper reveals the intriguing ability of the trained transformer to generalize under lower SNR and label-flipping noise at test time, even when such conditions are absent during pre-training.

Demonstration of benign overfitting: The paper provides theoretical and experimental evidence for the previously unobserved phenomenon of benign overfitting in transformers.

This paper presents novel theoretical insights into the in-context learning capabilities of transformers for linear classification tasks. The demonstration of benign overfitting is particularly noteworthy. While the simplified assumptions and focus on a convex architecture limit the direct applicability of the findings, the paper offers a valuable foundation for future research in understanding in-context learning in more realistic transformer settings.
summary: This paper studies the ability of linear transformers to perform in-context learning for linear classification tasks. The authors consider a simplified, convex linear transformer architecture trained on a series of random, class-conditional Gaussian mixture models.

The paper investigates two key aspects:

1. the number of pre-training tasks required for the transformer to generalize well at test time, even when the test data has lower signal-to-noise ratios (SNR) and label-flipping noise not present in the pre-training data.

2. the phenomenon of “benign overfitting in-context,” where the transformer memorizes noisy in-context training examples while still achieving near-optimal generalization for clean test examples.

The authors claim this work to be the first theoretical analysis of in-context learning for linear classification using linear transformers, and also the first demonstration of benign overfitting capabilities.

This paper makes valuable theoretical contributions to the understanding of in-context learning in transformers, highlighting their generalization capabilities and the intriguing phenomenon of benign overfitting. While the simplified model and assumptions might limit direct practical implications, the paper opens up interesting avenues for future research on the theoretical foundations of in-context learning.
weaknesses: Simplified architecture: The analysis relies on a 1-layer linear transformer model, which is a simplification of standard multi-layer transformer models with softmax-based attention. It is unclear whether about the applicability of the results to more complex transformer architectures.

Noise-free pre-training: The requirement of clean data during pre-training restricts the practicality of the approach in noisy environments.

Strong assumptions on data distribution: The paper focuses on class-conditional Gaussian mixture models, which might not be representative of the complex data distributions encountered in natural language processing tasks.

Absence of Numerical Validation: While the theoretical analysis is valuable, the paper lacks numerical experiments to validate the theoretical findings. Providing empirical evidence to support the theoretical claims, and demonstrating how the results generalize to more realistic settings, would enhance the paper's impact and persuasiveness.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: - Could tighter bounds be established for stochastic (Gibbs) classifiers? This would be stochasticity in the loss that is not a result of label-flips on the test set, but is a result of stochastic classifier decision.


- Line 377 breaks a formula.
soundness: 3
strengths: - The work references and uses prior work to great effect. It builds up on ideas established around the implicit regularization effect of the optimization algorithms used to train transformer models. SGD-based algorithms applied to convex-linear transformers have a directional bias towards maximum-margin solutions. Hence once can apply techniques used to study such solutions to derive generalization bounds such as the one presented in 4.1 and 4.2.
summary: This work presents a sample complexity analysis for in-context learning of class-conditional Gaussian mixtures with a restricted linear attention transformer model (single layer). It uses the results established on how SGD on convex-linear transformer models has a bias toward maximum-margin solution (in direction). The work thus assumes the KKT conditions are satisfied at the SGD solution and uses them to quantify the sample requirement to achieve a small test error. It shows how --by choice of specific model and data parameters-- one can achieve benign overfitting for in-context training, a phenomenon in which noisy training examples are exactly memorized, yet the test accuracy remains near-optimal.
weaknesses: - Readability: The paper is very dense and in parts hard to follow, requiring multiple revisits to earlier sections for (non-standard) notation and technical definitions. This, unfortunately, severely cuts into the readability of the paper. Adding a table of notation either in the paper or the appendix can be helpful.

- Analysis: There is no empirical evaluation of the tightness of derived bounds. IMO the derived bounds are to be understood in the limit and to prove the theoretical possibility of benign over-fitting. They might be too lose in practice. Adding an study in which $M$ or $d$ can be varied can help establish the empirical tightness of the bound.  

- Scope: The work only focuses on a specific restricted linear model. It defers analysis on non-convex linear or softmax-based attention models to future work.

- Assumptions: Even with the SGD-based solutions being directionally biased towards max-margin solutions, it doesn’t necessarily follow that the KKT conditions are satisfied in real implementations with early stopping or partial convergence. It’s also a leap to assume this (line of) theoretical analysis can be extended to other settings beyond the restricted or convex linear case addressed in the paper. IMO these assumption should be studied/validated under realistic training settings. Adding such analysis to the paper can clear some of these points.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: In Figure 1, the author used a natural language example to show the phenomenon of benign overfitting. However, I don't think it satisfies the data assumption. Is there an example/illustration that precisely satisfies the data/model assumptions?
soundness: 3
strengths: - This work theoretically analyzed how many pre-training tasks and in-context examples are needed for linear transformers to effectively generalize in linear classification tasks.
- This work studied the phenomenon of benign overfitting in this context and bounded the generalization error.
- The paper is solid and well presented. All the symbols are properly defined. The assumptions and theoretical results are clearly stated. (I did not carefully check the proofs.)
- The paper is well contextualized. The author mentioned related work (that I'm not familiar with) and discussed future work directions, which helped me position this work.
summary: The paper investigated the behavior of linear transformer models trained on random classification tasks, extending the existing work on linear regression tasks. It studied how linear transformers can generalize and exhibit _benign overfitting_, meaning that even if a model memorizes noisy training examples, it still performs well on test examples. This paper is completely theoretical, and no empirical evidence was given.
weaknesses: My biggest concern is the same for all theoretical work on simple models (e.g., single-layer linear transformer) with strong data assumptions (e.g., class-conditional Gaussian mixture). It is perfectly fine to start with a simple model with strong assumptions, as long as it is a good proxy for real-world problems or a stepping stone to understanding more complex problems. However, we need either theoretical or empirical evidence showing it is the case. In the current paper, there is a lack of discussion (or references) on whether or to what extent real-world data satisfies the data assumptions. It is unclear whether the theoretical results can be generalized to more complex models. Overall, I'd like to know how the theoretical results can guide our practice, which is not completely clear in the current version.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. At row 72, 73 → ‘even when pre-training on simple and easy-to-learn datasets, the transformer can generalize on more complex tasks.’ What does easy and complex task mean in this context? Does it mean the signal-to-noise ratio?

Minor:
1. At line 331, a typo of ‘sufficiently’.
soundness: 3
strengths: 1. The study of in-context learning is important, and the investigation on linear classification setting is novel.
2. The results of benign overfitting of transformer, to the reviewer’s knowledge, is new.
3. The paper is well-written and easy to follow.
summary: This paper studies in-context learning of linear classification tasks on linear transformers. The authors characterize how many pre-training tasks and in-context samples are needed for generalization, with an analysis of the implicit regularization of gradient descent. They show that after pre-training the transformer can tolerate smaller SNR than those tasks on which it was trained. Additionally, they show these trained transformers can memorize all in-context samples yet still generalize near-optimally when the in-context samples’ labels are flipped.
weaknesses: 1. It would have been interesting to discuss the connection of task complexity results under linear classification to Wu et al. [1]. In addition, briefly mentioning the importance of investigating linear classification instead of regression in the paper would be beneficial.
2. Can you explain why it would be interesting to consider a different signal-to-noise ratio R\tilde at test time? Analyzing the difference between M and N seems to be more interesting. 
3. The reviewer personally would like to see consistent empirical observation with simple experiments (such as a simplified version of Raventós et al [2]) but understand this might be out of the scope of this work. Maybe the author can consider mentioning it as a limitation.

[1] How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression? Wu et al. ICLR 2024   
[2] Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression. Raventós et al. NeurIPS 2023

Meta-review: The paper provides a theoretical study of in-context learning in transformers. Specifically, the authors consider a single-layer linear attention transformer model, and assume that the pretraining data consists of linear classification problems sampled from a fixed task distribution. They then provide generalization bounds for in-context learning of the linear regression algorithm. The proof in particular relies on the implicit bias of gradient descent, as well as assumptions about the data. The authors derive generalization bounds in terms of the number of examples per task, number of pre-training tasks, signal-to-noise ratio. The authors also show that the models can generalize at test time to tasks with lower signal-to-noise ratio and also demonstrate benign overfitting, where the models memorize mislabeled examples but still generalize well. The authors also provide empirical support for their predictions.

Strengths:
- The paper derives the first generalization guarantees for in-context learning in transformers.
  + The bound depends on the number of pre-training tasks and in-context examples, which is qualitatively interesting.
- The authors show interesting generalization beyond pretraining task distribution: to higher SNR and benign overfitting.
+ The authors provide some empirical support for their derivations.

Weaknesses:
- The paper makes a number of strong assumptions
  + Very simple architecture: one linear attention layer. This architecture is very different from the transformer models used in practice.
  + Strong assumptions on the task / data distribution

Decision recommendation: The paper provides an interesting new analysis of in-context learning in a simplified theoretical setting. While the paper makes strong simplifying assumptions, I think it is still a valuable theoretical contribution and a starting point for understanding more practical models and data distributions. The reviewers are unanimous in accepting the paper. I recommend accepting the paper.

Decision: Accept (Poster)