Feature Averaging: An Implicit Bias of Gradient Descent Leading to Non-Robustness in Neural Networks

## Abstract

Abstract In this work, we investigate a particular implicit bias in the gradient descent training process, which we term “Feature Averaging”, and argue that it is one of the principal factors contributing to non-robustness of deep neural networks. Despite the existence of multiple discriminative features capable of classifying data, neural networks trained by gradient descent exhibit a tendency to learn the average (or certain combination) of these features, rather than distinguishing and leveraging each feature individually. In particular, we provide a detailed theoretical analysis of the training dynamics of gradient descent in a two-layer ReLU network for a binary classification task, where the data distribution consists of multiple clusters with orthogonal cluster center vectors. We rigorously prove that gradient descent converges to the regime of feature averaging, wherein the weights associated with each hidden-layer neuron represent an average of the cluster centers (each center corresponding to a distinct feature). It leads the network classifier to be non-robust due to an attack that aligns with the negative direction of the averaged features. Furthermore, we prove that, with the provision of more granular supervised information, a two-layer multi-class neural network is capable of learning individual features, from which one can derive a binary classifier with the optimal robustness under our setting. Besides, we also conduct extensive experiments using synthetic datasets, MNIST and CIFAR-10 to substantiate the phenomenon of feature averaging and its role in adversarial robustness of neural networks. We hope the theoretical and empirical insights can provide a deeper understanding of the impact of the gradient descent training on feature learning process, which in turn influences the robustness of the network, and how more detailed supervision may enhance model robustness.

## 1 Introduction

Deep learning has achieved unprecedented success in a vast range of application domains. These models have been adopted in numerous applications including many safety-sensitive systems, such as autonomous driving and diagnostic assistance technologies. Despite the success of deep learning, a landmark study by Szegedy et al. () exposed that deep neural networks are extremely vulnerable to adversarial attacks. These attacks involve adding nearly imperceptible and carefully chosen perturbations to input data to confound deep learning models into making incorrect predictions. The perturbed inputs are termed adversarial examples and their existence has attracted significant attention from the research community.
Since then, various attacks (Biggio et al., ; Szegedy et al., ; Goodfellow et al., ; Madry et al., ) and defenses (Goodfellow et al., ; Madry et al., ; Shafahi et al., ; Pang et al., ) were developed, but the issue of adversarial robustness is still far from being resolved.

Gaining a deeper understanding of the adversarial robustness of neural networks is crucial not only for improving the reliability of deep learning systems in practice but also for illuminating the underlying theory of deep learning.
Daniely and Shacham (); Bubeck et al. (); Bartlett et al. (); Montanari and Wu () proved the existence of adversarial examples for neural networks with random weights across various architectures.
Tsipras et al. (); Zhang et al. ()
analyzed the fundamental trade-off between robustness and accuracy.
Bubeck et al. (); Bubeck and Sellke (); Li et al. (); Li and Li () proved that having a large model size is necessary for achieving robustness in many settings.
Ilyas et al. (); Tsilivis and Kempe (); Kumano et al. () studied the relationship between adversarial examples and the presence of non-robust but predictive features in the data distribution.

Another related line of works
studies the implicit bias of gradient descent,
showing that even without explicit regularization, gradient descent can implicitly bias the network towards solutions that enjoy favorable properties, in particular, those related to good generalization, such as norm minimization (Gunasekar et al., ), margin maximization (Soudry et al., ; Lyu and Li, ), equivalence to kernel methods (Arora et al., ), and having a flat local loss landscape (Blanc et al., ; Damian et al., ; Li et al., ).
However, a good generalization does not necessarily imply a good robustness with respect to inputs.
In fact, the seminal works by Vardi et al. () and Frei et al. () demonstrated that the implicit bias of gradient descent can be a “double-edged sword” in binary classification: for a multi-cluster data distribution where data is drawn from a fixed set of feature vectors and then perturbed by Gaussian noise,
gradient descent tends to learn a solution that generalizes well on clean data but is non-robust to small adversarial perturbations.
This interesting result is derived from the KKT-based characterization of the implicit bias proved in Lyu and Li (); Ji and Telgarsky (), but two limitations remain: (1) their characterization does not explicitly give the form of the learned non-robust solution, only an indirect characterization through the KKT conditions; (2) the KKT conditions hold exactly only after training the network for infinite many steps past the point where the training accuracy reaches $100\%$, but the network may be early-stopped in practice. See [Section 2](#S2) for more discussions.

In this paper, we adopt a similar multi-cluster data distribution and training setup as Vardi et al. () and Frei et al. (), but provide a detailed theoretical analysis that is able to explicitly characterize the entire training trajectory of gradient descent on two-layer ReLU networks and the form of the learned solution.
We rigorously prove that gradient descent exhibits a particular implicit bias, which we term feature averaging, that leads to generalizable but adversarially non-robust solutions soon after the training starts.
More specifically, we show that neurons in the hidden layer do not try to distinguish each individual cluster, but rather learn to having its weight vector as the average of some cluster centers (“features”) scaled by a positive factor. This feature-averaging bias can then be exploited by certain adversarial attacks, which provably leads to a wrong prediction even for small adversarial perturbations.

In our experiments, we observe similar phenomena empirically in several other settings. We argue that feature averaging it is one of the major factors contributing to non-robustness of deep neural networks, and show it is
closely related and collaborate several known phenomena and theoretical models in the study of adversarial robustness,
such as the robustness of the features (Tsipras et al., ; Ilyas et al., ; Allen-Zhu and Li, ; Tsilivis and Kempe, ), the Lipschitzness of the model (Bubeck et al., ; Bubeck and Sellke, ; Li et al., ; Li and Li, ), simplicity bias (Shah et al., ; Lyu et al., ),
and the dimpled manifold model (Shamir et al., ). A detailed discussion can be found in [Section 7](#S7).

In light of the feature-averaging bias, we then ask the following question: if we provide the model with more fine-grained supervision information on each individual feature, can gradient descent avoid the feature-averaging bias and learn a robust solution?
Theoretically, we affirmatively answer this question by showing that the following simple and natural method can enhance the robustness: the model is first trained as a multi-class classifier to predict the feature that each input belongs to, and then a binary classifier can be derived from the learned model. We show that this forces the model to learn each individual feature, leading to a robust binary classifier with optimal robust radius.
Empirically, we design several experiments, using synthetic and real datasets, and demonstrate that feature-level supervision information can be very effective in enhancing the robustness of the model (even with standard training).

Our technical contributions can be summarized as follows:

- 1.
(Section [4.1](#S4.SS1)) For certain multi-cluster data distributions (similar to Frei et al. ()), we prove that two-layer ReLU networks trained by gradient descent converge to feature-averaging solutions. In particular, we show that even though there exist solutions with optimal robust radius $O(\sqrt{d})$ ([Theorem D.3](#A4.Thmtheorem3)), gradient descent biases the weight vector of each neuron to the average of cluster-center features and leads to a solution that is non-robust for any radius $\Omega(\sqrt{d/k})$ ([Theorem 4.5](#S4.Thmtheorem5)), where $d$ is the input dimension and $k$ is the number of clusters. This result also resolves a version of the conjecture proposed by Min and Vidal () in our setting ([Theorem 4.6](#S4.Thmtheorem6)). A proof sketch is provided in [Section 5](#S5).
- 2.
(Section [4.2](#S4.SS2))
We show that if the model is trained with the feature-level labels,
a two-layer network can learn each individual feature, leading to a robust model with optimal robust radius $O(\sqrt{d})$ (Theorem [4.7](#S4.Thmtheorem7)).
- 3.
(Section [6](#S6))
We validate our theoretical results on synthetic data and real-world datasets such as MNIST and CIFAR-10. We empirically show that gradient descent learns the average feature. Our experiments also demonstrate enhanced robustness by incorpoarating finer-grained supervision information.

Finally, beyond our theoretical settings, we believe a similar feature-averaging bias may emerge in more complex scenarios. While there may not be a predefined set of features in these cases, gradient descent could still have a tendency to combine many localized, semantically meaningful (and thus, more robust (Ilyas et al., ; Tsilivis and Kempe, )) features into a single discriminative but non-robust feature. We hope our work sheds light on how this feature-averaging bias may arise and how it can be mitigated in broader contexts.

## 2 Related Work

Implicit Bias of Gradient Descent.
The implicit bias of gradient descent has been studied from various perspectives.
A prominent line of works establishes an equivalence between neural networks in certain training regimes to kernel regression with Neural Tangent Kernel (NTK) (Du et al., , ; Allen-Zhu et al., ; Zou et al., ; Chizat et al., ; Arora et al., ; Ji and Telgarsky, ; Cao and Gu, ), but the generalization of kernel regression is usually worse than that of real-world neural networks.
Other works prove other types of implicit biases beyond this NTK regime, including margin maximization (Soudry et al., ; Nacson et al., ; Lyu and Li, ; Ji and Telgarsky, ), parameter norm minimization (Gunasekar et al., , ; Arora et al., ) and sharpness reduction (Blanc et al., ; Damian et al., ; HaoChen et al., ; Li et al., ; Lyu et al., ; Gu et al., , ).
In particular, Lyu and Li (); Ji and Telgarsky (); Nacson et al. () proved that training homogeneous neural networks (including two-layer ReLU networks) with logistic loss and gradient flow leads to the convergence to a KKT solution of a margin maximization problem under mild assumptions. The aforementioned works by Vardi et al. () and Frei et al. () connected this line of works on margin to the non-robustness of neural networks, which we discuss shortly.

Feature Learning Theory for Two-Layer Networks. The feature learning theory of two-layer neural networks as proposed in various recent studies (Wen and Li, ; Allen-Zhu and Li, ; Chen et al., ; Cao et al., ; Zhou et al., ; Chidambaram et al., ; Allen-Zhu and Li, ; Kou et al., ; Simsek et al., ) aims to explore how features are learned in deep learning. This theory extends the theoretical optimization analysis beyond the scope of the neural tangent kernel (NTK) theory (Jacot et al., ; Du et al., , ; Allen-Zhu et al., ; Arora et al., ). Various data assumptions about feature-noise structure are made in these feature learning works. Assuming a sparse coding model, Wen and Li () studied the feature learning process of self-supervised contrastive learning, and Allen-Zhu and Li () proposed a principle called feature purification to explain the inner workings of
adversarial training.
Allen-Zhu and Li () assumed the so-called “multi-view” data with patch structure to understand the benefits of ensembles in deep learning. Following this multi-view data assumption, Chidambaram et al. () showed that data mix-up algorithm can provably learn diverse features to improve generalization. Cao et al. (); Kou et al. () explored the benign overfitting phenomenon of two-layer convolutional neural networks by leveraging a technique of signal-noise decomposition. Zhou et al. () studied feature condensation and prove that, for two-layer network with small initialization, input weights of hidden neurons condense onto isolated orientations at the initial training stage. Simsek et al. () focused on the regression setting and study the compression of the teacher network, and they found that weight vectors, whether copying an individual teacher vector or averaging a set of teacher vectors, are critical points of the loss function.

Comparisons to Vardi et al. (), and Frei et al. ().
Recently, Vardi et al. () and Frei et al. () demonstrated that for two-layer ReLU networks, any KKT solution to the maximum margin problem leads to non-robust solutions for binary classification on a multi-cluster distribution. Their finding highlights the significance of the optimization process in the (non-)robustness of neural networks. Our theoretical results are inspired by theirs, but
differ from theirs in the following important aspects:
(1) Conceptually, feature averaging is arguably more intuitive and concrete (in the feature level) than the set of KKT properties. Moreover, feature averaging (or its nonlinear extensions) may appear in more complex and general setting even when the solution is far from a KKT point. (2) Technically, we perform a detailed and finite-time analysis of the gradient descent dynamics, in contrast to their result about limiting behavior of gradient descent. In particular, our analysis of gradient descent dynamics reveals the feature learning process. Furthermore, we comment that the time complexity converging from an initialization point to a KKT solution can be slow, i.e., $\Omega(1/\operatorname{log}(t))$, as proven in Soudry et al. (); Lyu and Li (); Kou et al. (). (3) Our analysis of the GD dynamics requires small initialization, whereas their results depend on starting from a solution that already correctly classifies the training set, which is an assumption made in Lyu and Li () for achieving KKT points.

Comparisons to Min and Vidal ().
A very recent work by Min and Vidal () studied
a data model very similar to that in Vardi et al. () and Frei et al. (), and further conjectured that the weight vectors of a two-layer ReLU network trained by gradient flow converge to an average of the cluster centers. They empirically verified the conjecture through numerical simulations on a synthetic dataset. Our results resolve the conjecture under a moderately modified setup, where we fix readout weights of the second layer and only train the hidden weights of the first layer using the gradient descent method. Specifically, we show that the weight vector associated with each neuron aligns with a weighted average of cluster features, and that the ratio between the weights of distinct clusters is close to $1$ (Theorem [4.5](#S4.Thmtheorem5)), which furthermore implies their conjecture, as stated in Theorem [4.6](#S4.Thmtheorem6).

## 3 Problem Setup

In this section, we introduce some useful notations and concepts, including the multi-cluster data distribution, the two-layer neural network learner and the gradient descent algorithm.

Notations. We use bold-face letters to denote vectors, e.g., ${\bm{x}}=\left(x_{1},\ldots,x_{d}\right)$. For $\bm{x}\in\mathbb{R}^{d}$, we denote by $\|\bm{x}\|$ the Euclidean ($\ell_{2}$) norm. We denote by $\mathbbm{1}(\cdot)$ the standard indicator function.We denote $\operatorname{sgn}(z)=1$ if $z>0$ and -1 otherwise. For an integer $n\geq 1$, we denote $[n]=\{1,\ldots,n\}$. We denote by ${\mathcal{N}}\left(\mu,\sigma^{2}\right)$ the normal distribution with mean $\mu\in\mathbb{R}$ and variance $\sigma^{2}$, and by ${\mathcal{N}}(\bm{\mu},\bm{\Sigma})$ the multivariate normal distribution with mean vector $\bm{\mu}$ and covariance matrix $\bm{\Sigma}$. The identity matrix of size $d$ is denoted by $\bm{I}_{d}$. We use $\operatorname{Unif}(A)$ to denote the uniform distribution on the support set $A$. We use standard asymptotic notation ${\mathcal{O}}(\cdot)$ and $\Omega(\cdot)$ to hide constant factors, and $\tilde{{\mathcal{O}}}(\cdot),\tilde{\Omega}(\cdot)$ to hide logarithmic factors.

### 3.1 Data Distribution

Following Vardi et al. (); Frei et al. (), we consider binary classification on the following data distribution with multiple clusters.

###### Definition 3.1 (Multi-Cluster Data Distribution) .

Given $k$ vectors ${\bm{\mu}}_{1},\dots,{\bm{\mu}}_{k}\in\mathbb{R}^{d}$, called the cluster features, and a partition of $[k]$ into two disjoint sets $J_{\pm}=(J_{+},J_{-})$,
we define $\mathcal{D}(\{{\bm{\mu}}_{j}\}_{j=1}^{k},J_{\pm})$
as a data distribution on $\mathbb{R}^{d}\times\{-1,1\}$, where each data point $({\bm{x}},y)$ is generated as follows:

- 1.
Draw a cluster index as $j\sim\mathrm{Unif}([k])$;
- 2.
Set $y=+1$ if $j\in J_{+}$; otherwise $j\in J_{-}$ and set $y=-1$;
- 3.
Draw ${\bm{x}}:={\bm{\mu}}_{j}+{\bm{\xi}}$, where ${\bm{\xi}}\sim{\mathcal{N}}({\bm{0}},{\bm{I}}_{d})$.

For convenience, we write $\mathcal{D}$ instead of $\mathcal{D}(\{{\bm{\mu}}_{j}\}_{j=1}^{k},J_{\pm})$ if $\{{\bm{\mu}}_{j}\}_{j=1}^{k}$ and $J_{\pm}$ are clear from the context.
For $s\in\{\pm 1\}$, we write $J_{s}$ to denote $J_{+}$ if $s=+1$ and $J_{-}$ if $s=-1$.

To ease the analysis,
we make the following simplifying assumptions
on the distribution.

###### Assumption 3.2 (Orthogonal Equinorm Cluster Features) .

The cluster features $\{{\bm{\mu}}_{j}\}_{j=1}^{k}$ satisfy the properties that (1) $\|{\bm{\mu}}_{j}\|=\sqrt{d}$ for all $j\in[k]$; and (2) ${\bm{\mu}}_{i}\perp{\bm{\mu}}_{j}$ for all $1\leq i<j\leq k$.

###### Assumption 3.3 (Nearly Balanced Classification) .

The partition $J_{\pm}$ satisfies $c^{-1}\leq\frac{\lvert J_{+}\rvert}{\lvert J_{-}\rvert}\leq c$ for some absolute constant $c\geq 1$.

Our data distribution is similar to that in Vardi et al. () and Frei et al. (). In particular, Vardi et al. () consider a setting where data are comprised
of $k$ nearly orthogonal data points in $\mathbb{R}^{d}$. This assumption is further relaxed in Frei et al. (), where they assume $k$ clusters with nearly orthogonal cluster means $\{\bm{\mu}_{i}\}_{i=1}^{k}$ (i.e., they have that $\frac{|\langle\bm{\mu}_{i},\bm{\mu}_{j}\rangle|}{\|\bm{\mu}_{i}\|\|\bm{\mu}_{j}\|}=O\left(\frac{1}{k}\right)$ holds for all $i\neq j$). For simplicity, we focus on the exactly orthogonal cluster setting in our work.

###### Definition 3.1 (Multi-Cluster Data Distribution) .

###### Assumption 3.2 (Orthogonal Equinorm Cluster Features) .

###### Assumption 3.3 (Nearly Balanced Classification) .

### 3.2 Neural Network Learner

A training dataset ${\mathcal{S}}:=\left\{({\bm{x}}_{i},y_{i})\right\}_{i=1}^{n}\subseteq\mathbb{R}^{d}\times\{-1,1\}$ of size $n$ is randomly sampled from the data distribution $\mathcal{D}(\{{\bm{\mu}}_{j}\}_{j=1}^{k},J_{\pm})$ and is used to train a two-layer neural network.

#### Network Architecture.

We focus on learning two-layer ReLU networks. Such networks are usually defined as $f_{{\bm{\theta}}}({\bm{x}}):=\sum_{j=1}^{M}a_{j}\operatorname{ReLU}(\langle{{\bm{w}}_{j}},{{\bm{x}}}\rangle+b_{j})$,
where ${\bm{\theta}}:=\left(\{a_{j}\}_{j=1}^{M},\{{\bm{w}}_{j}\}_{j=1}^{M},\{b_{j}\}_{j=1}^{M}\right)$ are the parameters of the network, and $\operatorname{ReLU}(\,\cdot\,)$ is the ReLU activation function defined as $\operatorname{ReLU}(z)=\max(0,z)$.

For the sake of simplicity, we consider the case where $M=2m$ is even and fix the second layer as $a_{j}=\frac{1}{m}$ for $1\leq j\leq m$ and $a_{j}=-\frac{1}{m}$ for $m+1\leq j\leq 2m$, which is a widely adopted setting in the literature of feature learning theory (Allen-Zhu and Li, ; Cao et al., ; Kou et al., ). With this simplification, we focus on training only the first layer $(\{{\bm{w}}_{j}\}_{j=1}^{M},\{b_{j}\}_{j=1}^{M})$ and rewrite the network as

$$ $f_{{\bm{\theta}}}({\bm{x}}):=\frac{1}{m}\sum_{r\in[m]}\operatorname{ReLU}(\langle\bm{w}_{+1,r},\bm{x}\rangle+b_{+1,r})-\frac{1}{m}\sum_{r\in[m]}\operatorname{ReLU}(\langle\bm{w}_{-1,r},\bm{x}\rangle+b_{-1,r}),$ $$

where ${\bm{\theta}}=(\{{\bm{w}}_{+1,r}\}_{r=1}^{m},\{b_{+1,r}\}_{r=1}^{m},\{{\bm{w}}_{-1,r}\}_{r=1}^{m},\{b_{-1,r}\}_{r=1}^{m})$ are the trainable parameters, and ${\bm{w}}_{+1,r}$ and $b_{+1,r}$ correspond to the neurons with $a_{r}=\frac{1}{m}$, while ${\bm{w}}_{-1,r}$ and $b_{-1,r}$ correspond to the neurons with $a_{r}=-\frac{1}{m}$.

Training Objective and Gradient Descent. The neural network $f_{{\bm{\theta}}}(\cdot)$ is trained to minimize the following empirical loss on the training dataset ${\mathcal{S}}$: ${\mathcal{L}}({\bm{\theta}}):=\frac{1}{n}\sum_{i=1}^{n}\ell\left(y_{i}f_{\bm{\theta}}\left(\bm{x}_{i}\right)\right)$,
where $\ell(q):=\log(1+e^{-q})$ is the logistic loss.
We apply gradient descent to minimize this loss:

$$ $\bm{\theta}^{(t+1)}=\bm{\theta}^{(t)}-\eta\nabla{\mathcal{L}}(\bm{\theta}^{(t)}),$ (1) $$

where $\bm{\theta}^{(t)}$ denotes the parameters at $t$-th iteration for all $t\geq 0$, and $\eta>0$ is the learning rate. We specify the derivative of ReLU activation as $\operatorname{ReLU}^{\prime}(z)=\mathbbm{1}(z\geq 0)$ in backpropagation.
At initialization, we set $\bm{w}_{s,r}^{(0)}\sim\mathcal{N}(\bm{0},\sigma_{\mathrm{w}}^{2}\bm{I}_{d})$ and $b_{s,r}^{(0)}\sim\mathcal{N}(0,\sigma_{\mathrm{b}}^{2})$ for some $\sigma_{\mathrm{w}},\sigma_{\mathrm{b}}>0$.

#### Clean Accuracy and Robust Accuracy.

For a given data distribution ${\mathcal{D}}$ over $\mathbb{R}^{d}\times\{-1,1\}$, the clean accuracy of a neural network $f_{{\bm{\theta}}}:\mathbb{R}^{d}\rightarrow\mathbb{R}$ on ${\mathcal{D}}$
is defined as

$$ $\mathrm{Acc}_{\mathrm{clean}}^{{\mathcal{D}}}(f_{{\bm{\theta}}}):=\mathbb{P}_{(\bm{x},y)\sim{\mathcal{D}}}\left[\operatorname{sgn}(f_{\bm{\theta}}(\bm{x}))=y\right].$ $$

In this work, we focus on the $\ell_{2}$-robustness.
The $\ell_{2}$ $\delta$-robust accuracy of $f_{{\bm{\theta}}}$ on ${\mathcal{D}}$ is defined as

$$ $\mathrm{Acc}_{\mathrm{robust}}^{{\mathcal{D}}}(f_{{\bm{\theta}}};\delta):=\mathbb{P}_{(\bm{x},y)\sim{\mathcal{D}}}\left[\forall\bm{\rho}\in\mathbb{B}_{\delta}:\operatorname{sgn}(f_{\bm{\theta}}(\bm{x}+\bm{\rho}))=y\right],$ $$

where $\mathbb{B}_{\delta}:=\{\bm{\rho}\in\mathbb{R}^{d}:\|\bm{\rho}\|\leq\delta\}$ is the $\ell_{2}$-ball centered at the origin with radius $\delta$.
We say that a neural network $f_{{\bm{\theta}}}$ is $\delta$-robust if $\mathrm{Acc}_{\mathrm{robust}}^{{\mathcal{D}}}(f_{{\bm{\theta}}};\delta)\geq 1-\epsilon(d)$
for some function $\epsilon(d)$ that vanishes to zero, i.e., $\epsilon(d)\to 0$ as $d\to\infty$.

#### Robust Networks Exist.

In a very similar setting to ours, Frei et al. () show that there exists a two-layer ReLU network that can achieve nearly $100\%$ clean accuracy and $\Omega(\sqrt{d})$-robust accuracy on their data distribution.
In our setting, we can also construct a similar network that achieves nearly $100\%$ clean accuracy and $\Omega(\sqrt{d})$-robust accuracy.
In particular, such network utilizes one hidden neural to capture
one feature/cluster (i.e., the neural is activated only if the input
point is from the corresponding cluster).
See Theorem [D.3](#A4.Thmtheorem3) in Appendix [D.2](#A4.SS2) for the details and
Figure [1](#S4.F1) for an illustration.
However, we will soon show that, despite such $\Omega(\sqrt{d})$-robust network exists, gradient descent is incapable of learning such a robust network, but instead converges to a very different solution with a robust radius that is $\Theta(\sqrt{k})$ times smaller.

## 4 Main Results

Figure: Figure 1: Schematic illustration of feature-averaging and feature-decoupling: We consider a dataset with $5$ clusters. The first three clusters belong to $J_{+}$, and the other two to $J_{-}$. Denote $\bm{\mu}_{+}:={(\bm{\mu}_{1}+\bm{\mu}_{2}+\bm{\mu}_{3})}/{3},\bm{\mu}_{-}:={(\bm{\mu}_{4}+\bm{\mu}_{5})}/{2}$. For ease of illustration, we assume that $\sum_{j=1}^{5}\bm{\mu}_{j}=\bm{0}$. The feature-averaging classifier $f_{\mathrm{FA}}$ leverages two neurons with averaged features $\bm{\mu}_{+}$ and $\bm{\mu}_{-}$ to classify all data, which corresponds to a linear classifier (the gray line). The feature-decoupling classifier $f_{\mathrm{FD}}$ leverages individual features and has more complex polyhedral decision boundary (green lines). Note that the instance is high dimensional and this is only a schematic illustration. The distance between data points and the decision boundary of $f_{\mathrm{FD}}$ (green lines) is much larger than that of $f_{\mathrm{FA}}$ (gray line), which implies that the feature-decoupling classifier is more robust than the feature-averaging one.
Refer to caption: /html/2410.10322/assets/x1.png

In this section, we present our main technical results.
In Subsection [4.1](#S4.SS1), we first present the main result (Theorem [4.5](#S4.Thmtheorem5)) regarding feature averaging, that is standard gradient descent training finds feature averaging solutions
for the data distribution $\mathcal{D}$ and such feature averaging solution is non-robust.
In Subsection [4.2](#S4.SS2), we demonstrate that if more supervisory information can be obtained (specific cluster categories rather than just binary classification labels), we can achieve feature decoupling
via gradient descent on a similar two layer multi-class network. Consequently, we can obtain a binary classification network with optimal robust perturbation radius (Theorem [4.7](#S4.Thmtheorem7)).

### 4.1 Network Learner Provably Learns Feature-Averaging Solution

The prior work by Frei et al. () has showed that, under certain conditions, training a two-layer ReLU network for infinite time converges to a network that can achieve nearly $100\%$ clean accuracy on $\mathcal{D}$ but is only $o(\sqrt{d/k})$-robust. A subsequent work by Min and Vidal () conjectured that the network converges to a specific form of solution, which we refer to as the feature-averaging network.

###### Definition 4.1 (Feature-Averaging Network) .

We define $f_{\mathrm{FA}}({\bm{x}})$ as the following function:

$$ $f_{\mathrm{FA}}({\bm{x}}):=\lvert J_{+}\rvert\cdot\operatorname{ReLU}\left(\left\langle{{\bm{\mu}_{+}}},{{\bm{x}}}\right\rangle\right)-\lvert J_{-}\rvert\cdot\operatorname{ReLU}\left(\left\langle{{\bm{\mu}_{-}}},{{\bm{x}}}\right\rangle\right),$ $$

where ${\bm{\mu}_{+}}:=\frac{1}{|J_{+}|}\sum_{j\in J_{+}}{\bm{\mu}}_{j}$ is the average of cluster centers in the positive class, and similarly ${\bm{\mu}_{-}}:=\frac{1}{|J_{-}|}\sum_{j\in J_{-}}{\bm{\mu}}_{j}$ is that for the negative class.
We say that a two-layer ReLU network $f_{{\bm{\theta}}}({\bm{x}})$ is a feature-averaging network if $f_{{\bm{\theta}}}({\bm{x}})=C\cdot f_{\mathrm{FA}}({\bm{x}})$ for some $C>0$.

###### Remark 4.2 .

The feature-averaging network uses the first neuron to process all data within positive clusters, and the second neuron negative clusters. Thus, it can correctly classify clean data. However, it fails to robustly classify perturbed data for a radius larger than $\Omega(\sqrt{d/k})$: in particular, consider the attack vector $\bm{\rho}$ that aligns with the negative direction of the averaged features, i.e., $\bm{\rho}\parallel(-\sum_{j\in J_{+}}\bm{\mu}_{j}+\sum_{j\in J_{-}}\bm{\mu}_{j})$. One can easily check that with
$\|\bm{\rho}\|=\delta=\Omega(\sqrt{d/k})$,
the attack is successful, i.e.,
$\operatorname{sgn}(f_{\mathrm{FA}}({\bm{x}}+\bm{\rho}))\neq\operatorname{sgn}(f_{\mathrm{FA}}({\bm{x}}))$ due to the linearity of $f_{\mathrm{FA}}({\bm{x}}+\bm{\rho})$ over $\bm{\rho}$. See Appendix [D.1](#A4.SS1) for more details, and see Figure [1](#S4.F1) for an illustration.

Our first main result is a non-asymptotic analysis
of the training dynamics that explicitly characterizes the solution learned by gradient descent on distribution $\mathcal{D}$
after a finite number of iterations.
For theoretical analysis, we make the following assumptions about the hyper-parameters.

###### Assumption 4.3 (Choices of Hyper-Parameters) .

We assume that:

$$ $\displaystyle d$ $\displaystyle=\Omega(k^{10})$ $\displaystyle c$ $\displaystyle=\Theta(1)$ $\displaystyle n$ $\displaystyle\in[\Omega(k^{7}),\operatorname{exp}(O(\operatorname{log}^{2}(d)))]$ $\displaystyle m$ $\displaystyle=\Theta(k)$ $\displaystyle\eta$ $\displaystyle=O(d^{-2})$ $\displaystyle\sigma_{b}^{2}$ $\displaystyle=\sigma^{2}_{w}=O(\eta k^{-5}).$ $$

###### Remark 4.4 (Discussion of Hyper-Parameter Choices) .

We make specific choices of hyper-parameters for the sake of calculations, and we emphasize that these may not be the tightest possible choices. In particular, we need the data dimension $d$ to be significantly larger than the number of clusters $k$ to ensure all $k$ cluster features are orthogonal within $\mathbb{R}^{d}$. We further require that the number of samples $n$ is a large polynomial of $k$ to ensure that the network can learn all $k$ cluster features. We assume the learning rate $\eta$ and the initialization magnitude $\sigma_{w},\sigma_{b}$ are sufficiently small, which helps the network to be trained in the feature learning regime (Lyu et al., ; Cao et al., ; Allen-Zhu and Li, ; Kou et al., ).

Now, everything is ready to state the first main theorem of our paper, which characterizes the weights of the learned network and shows that after a certain number of iterations, the network can be closely approximated by the feature-averaging network (defined in Definition [4.1](#S4.Thmtheorem1)).

###### Theorem 4.5 .

In the setting of training a two-layer ReLU network on the binary classification problem $\mathcal{D}(\{{\bm{\mu}}_{j}\}_{j=1}^{k},J_{\pm})$ as described in [Section 3](#S3),
under [Assumptions 3.2](#S3.Thmtheorem2), [3.3](#S3.Thmtheorem3) and [4.3](#S4.Thmtheorem3),
for some $\gamma=o(1)$, after $\Omega(\eta^{-1})\leq T\leq\exp(\tilde{O}(k^{1/2}))$ iterations, with probability at least $1-\gamma$, the neural network satisfies the following properties:

- 1.
The clean accuracy is nearly perfect: $\mathrm{Acc}_{\mathrm{clean}}^{\mathcal{D}}(f_{{\bm{\theta}}^{(T)}})\geq 1-\exp(-\Omega(\log^{2}d))$.
- 2.
Gradient descent leads the network to the feature-averaging regime: there exists a time-variant coefficient $\lambda^{(T)}\in[\Omega(1),+\infty)$ such that for all $s\in\{\pm 1\}$, $r\in[m]$, the weight vector $\bm{w}_{s,r}^{(T)}$ can be approximated as
$\bigg{\|}{\bm{w}}_{s,r}^{(T)}-\lambda^{(T)}\sum_{j\in J_{s}}\|\bm{\mu}_{j}\|^{-2}\bm{\mu}_{j}\bigg{\|}\leq o(d^{-1/2})$
and the bias terms are sufficiently small, i.e., $\left\lvert b_{s,r}^{(T)}\right\rvert\leq o(1)$.
- 3.
Consequently, the network is non-robust: for perturbation radius $\delta=\Omega(\sqrt{d/k})$, the $\delta$-robust accuracy is nearly zero, i.e., $\mathrm{Acc}_{\mathrm{robust}}^{\mathcal{D}}(f_{{\bm{\theta}}^{(T)}};\delta)\leq\exp(-\Omega(\log^{2}d))$.

We provide a proof sketch for [Theorem 4.5](#S4.Thmtheorem5) in [Section 5](#S5) (see the full proof in [Section B.2](#A2.SS2)). [Theorem 4.5](#S4.Thmtheorem5) suggests that the weight vector aligns with the average of cluster features: the
direction of the weight vector associated with a positive neuron converges to the average of positive cluster features ${\bm{\mu}_{+}}$, and that associated with a negative neuron to the average of negative cluster features ${\bm{\mu}_{+}}$. Moreover, the above feature-averaging property of learned network implies non-robustness, i.e., the learned network is only $o(\sqrt{d/k})$-robust although an $\Omega(\sqrt{d})$-robust solution exists as we proved in Section [3](#S3).

As a corollary of [Theorem 4.5](#S4.Thmtheorem5), we resolve the conjecture proposed in Min and Vidal () in our setting.

###### Theorem 4.6 (Conjecture 1 from Min and Vidal ( 2024 ) ) .

In the setting of [Theorem 4.5](#S4.Thmtheorem5), we have that $\inf_{C>0}\sup_{{\bm{x}}\in\mathbb{R}^{d}:\|{\bm{x}}\|_{2}=\sqrt{d}}\left\lvert Cf_{\mathrm{FA}}({\bm{x}})-f_{{\bm{\theta}}^{(T)}}({\bm{x}})\right\rvert=o(1)$,
where $f_{\mathrm{FA}}({\bm{x}})$ is the feature-averaging network ([Definition 4.1](#S4.Thmtheorem1)).

Under a similar orthogonal cluster data assumption, Min and Vidal () conjecture that two-layer neural network converges to the feature-averaging solution via gradient flow training with small initialization.
They empirically validate the conjecture via experiments on
synthetic datasets. [Theorem 4.6](#S4.Thmtheorem6) provides a rigorous proof
for the conjecture, although the original conjecture is stated under a slightly different setting from ours. In their setting, the second layer of the network is also trainable, but we fix the second layer for simplicity. We also require certain assumptions on the hyperparameters, which has been discussed in details in [Assumption 4.3](#S4.Thmtheorem3) and [Remark 4.4](#S4.Thmtheorem4).

###### Definition 4.1 (Feature-Averaging Network) .

###### Remark 4.2 .

###### Assumption 4.3 (Choices of Hyper-Parameters) .

###### Remark 4.4 (Discussion of Hyper-Parameter Choices) .

###### Theorem 4.5 .

###### Theorem 4.6 (Conjecture 1 from Min and Vidal ( 2024 ) ) .

### 4.2 Fine-Grained Supervision Improves Robustness

We have shown that gradient descent is unable to differentiate
individual cluster features, which causes non-robustness.
Hence, a natural question is what if
we provide more fine-grained feature level supervision,
can gradient descent learn a robust solution? We show that this is indeed possible in the case where each data point is labeled with the cluster it belongs to, rather than just a binary label.

#### Fine-Grained Supervision.

Following the setting in [Section 3](#S3), we consider the binary classification task with data distribution $\mathcal{D}(\{{\bm{\mu}}_{j}\}_{j=1}^{k},J_{\pm})$. But instead of training the model directly to predict the binary labels, we assume that we are able to label each data point with the cluster $\hat{y}\in[k]$ it belongs to, and then we train a $k$-class classifier to predict the cluster labels. More specifically, we first sample a training set ${\mathcal{S}}:=\{({\bm{x}}_{i},y_{i})\}_{i=1}^{n}\subseteq\mathbb{R}^{d}\times\{\pm 1\}$ from $\mathcal{D}$, along with the cluster labels $\{\tilde{y}_{i}\}_{i=1}^{n}$ for all data points. Then a $k$-class neural network classifier is trained on $\tilde{{\mathcal{S}}}:=\{({\bm{x}}_{i},\tilde{y}_{i})\}_{i=1}^{n}\subseteq\mathbb{R}^{d}\times[k]$.

Multi-Class Network Classifier. We train the following two-layer neural network for the $k$-class classification mentioned above: $\bm{F}_{\bm{\theta}}({\bm{x}}):=(f_{1}({\bm{x}}),f_{2}({\bm{x}}),\dots,f_{k}({\bm{x}}))\in\mathbb{R}^{k}$, where $f_{j}({\bm{x}}):=\frac{1}{h}\sum_{r=1}^{h}\operatorname{ReLU}(\langle\bm{w}_{j,r},\bm{x}\rangle)$, ${\bm{\theta}}:=({\bm{w}}_{1,1},{\bm{w}}_{1,2},\dots,{\bm{w}}_{k,h})\in\mathbb{R}^{khd}$ are trainable weights, and $h=\Theta(1)$. One can think $\{\operatorname{ReLU}(\langle{{\bm{w}}_{j,r}},{{\bm{x}}}\rangle)\}_{j\in[k],r\in[h]}$ as $kh=\Theta(k)$ neurons partitioned into $k$ groups, where the corresponding second layer weights are set in a way that the $j$-th group only contributes to the $j$-th output $f_{j}({\bm{x}})$ of the network.
The output $\bm{F}_{\bm{\theta}}({\bm{x}})$ is converted to probabilities using the softmax function, namely $p_{j}({\bm{x}}):=\frac{\exp(f_{j}({\bm{x}}))}{\sum_{i=1}^{k}\exp(f_{i}({\bm{x}}))}$ for $j\in[k]$. For predicting the binary label for the original binary classification task on $\mathcal{D}$, we take the difference of the probabilities of the positive and negative classes, i.e., $F^{\mathrm{binary}}_{{\bm{\theta}}}({\bm{x}}):=\sum_{j\in J_{+}}p_{j}({\bm{x}})-\sum_{j\in J_{-}}p_{j}({\bm{x}})$. The clean accuracy $\mathrm{Acc}_{\mathrm{clean}}^{\mathcal{D}}(F^{\mathrm{binary}}_{{\bm{\theta}}})$ and $\delta$-robust accuracy $\mathrm{Acc}_{\mathrm{robust}}^{\mathcal{D}}(F^{\mathrm{binary}}_{{\bm{\theta}}};\delta)$ are then defined similarly as before.

Training Objective and Gradient Descent with Fine-Grained Supervision.
We train the multi-class network ${\bm{F}}_{\bm{\theta}}({\bm{x}})$ to minimize the cross-entropy loss ${\mathcal{L}}_{\mathrm{CE}}({\bm{\theta}}):=-\frac{1}{n}\sum_{i=1}^{n}\log p_{\tilde{y}_{i}}({\bm{x}}_{i})$. Similar to [Section 3](#S3), we use gradient descent to minimize the loss function ${\mathcal{L}}_{\mathrm{CE}}({\bm{\theta}})$ with learning rate $\eta$, i.e., $\bm{\theta}^{(t+1)}=\bm{\theta}^{(t)}-\eta\nabla_{\bm{\theta}}\mathcal{L}_{\textit{CE}}(\bm{F}_{\bm{\theta}^{(t)}})$. At initialization, we set ${\bm{w}}_{j,r}^{(0)}\sim{\mathcal{N}}(0,\sigma_{\mathrm{w}}^{2}{\bm{I}}_{d})$ for some $\sigma_{\mathrm{w}}>0$.

GD Finds Robust Networks. In contrast to the feature-averaging implicit bias in our previous setting (Theorem [4.5](#S4.Thmtheorem5)), the following theorem shows that with fine-grained supervision, gradient descent converges to a neural network that learns decoupled features, i.e., the weight of each neuron is aligned with one
cluster feature.

###### Theorem 4.7 .

In the setting of training a multi-class network on the multiple classification problem $\tilde{{\mathcal{S}}}:=\{({\bm{x}}_{i},\tilde{y}_{i})\}_{i=1}^{n}\subseteq\mathbb{R}^{d}\times[k]$ as described in the above,
under [Assumptions 3.2](#S3.Thmtheorem2), [3.3](#S3.Thmtheorem3) and [4.3](#S4.Thmtheorem3),
for some $\gamma=o(1)$, after $\Omega(\eta^{-1}k^{8})\leq T\leq\exp(\tilde{O}(k^{1/2}))$ iterations, with probability at least $1-\gamma$, the neural network satisfies the following properties:

- 1.
The clean accuracy is nearly perfect: $\mathrm{Acc}_{\mathrm{clean}}^{\mathcal{D}}(F^{\mathrm{binary}}_{{\bm{\theta}}^{(T)}})\geq 1-\exp(-\Omega(\log^{2}d))$.
- 2.
The network converges to the feature-decoupling regime: there exists a time-variant coefficient $\lambda^{(T)}\in[\Omega(\operatorname{log}k),+\infty)$ such that for all $j\in[k]$, $r\in[h]$, the weight vector $\bm{w}_{j,r}^{(T)}$ can be approximated as
$\bigg{\|}{\bm{w}}_{j,r}^{(T)}-\lambda^{(T)}\|\bm{\mu}_{j}\|^{-2}\bm{\mu}_{j}\bigg{\|}\leq o(d^{-1/2}).$
- 3.
Consequently, the corresponding binary classifier achieves optimal robustness: for perturbation radius $\delta=O(\sqrt{d})$, the $\delta$-robust accuracy is also nearly perfect, i.e., $\mathrm{Acc}_{\mathrm{robust}}^{\mathcal{D}}(F^{\mathrm{binary}}_{{\bm{\theta}}^{(T)}};\delta)\geq 1-\exp(-\Omega(\log^{2}d))$.

The detailed proof can be found in Appendix [C.3](#A3.SS3). Theorem [4.7](#S4.Thmtheorem7) manifests that the multi-class network learns the decoupled features, and the induced binary classifier achieves optimal robustness.
See Figure [1](#S4.F1) for an illustration.
Instead of leveraging the bias term to filter out cluster noise as the feature-decoupling classifier $f_{\mathrm{FD}}$ that we illustrated in Figure [1](#S4.F1) and Theorem [D.3](#A4.Thmtheorem3), the soft-max operator of $F^{\mathrm{binary}}_{{\bm{\theta}}}$ plays a similar role here.

It can be easily verified that $F^{\mathrm{binary}}_{{\bm{\theta}}^{(T)}}$ achieves optimal robustness radius (up to constant factor) since the distance between distinct cluster centers is at most $\Theta(\sqrt{d})$ (i.e. $\|\bm{\mu}_{i}-\bm{\mu}_{j}\|=\Theta(\sqrt{d}),\forall i\neq j$).

Convergence to Robust Networks Requires Implicit Bias. In fact, adding more fine-grained supervision signals does not trivially lead to decoupled features and robustness, since the above network found by gradient descent is not the only solution that can achieve $100\%$ clean accuracy. As a counterexample, we show that there exists a multi-class network that achieves perfect clean accuracy but is not $\Omega(\sqrt{d/k})$-robust, which is formally given in the following proposition.

###### Proposition 4.8 .

Consider the following multi-class network $\bm{F}_{\tilde{{\bm{\theta}}}}$: for all $j\in[k]$, the sub-network $f_{j}$ has only single neuron ($h=1$) and is defined as $f_{j}(\bm{x})=\operatorname{ReLU}\left(\left\langle\bm{\mu}_{j}+\sum_{l\in J_{s}}\bm{\mu}_{l},\bm{x}\right\rangle\right)$,
where cluster $j$ has binary label $s\in\{\pm 1\}$.
With probability at least $1-\operatorname{exp}(-\Omega(\operatorname{log}^{2}d))$ over $\tilde{S}$, we have that ${\mathcal{L}}_{\mathrm{CE}}(\tilde{{\bm{\theta}}})\leq\operatorname{exp}(-\Omega(d))=o(1)$, where $\tilde{{\bm{\theta}}}$ denotes the weights of $\bm{F}_{\tilde{{\bm{\theta}}}}$. Moreover, $\mathrm{Acc}_{\mathrm{clean}}^{\mathcal{D}}(F^{\mathrm{binary}}_{\tilde{{\bm{\theta}}}})\geq 1-\exp(-\Omega(\log^{2}d))$, $\mathrm{Acc}_{\mathrm{robust}}^{\mathcal{D}}(F^{\mathrm{binary}}_{\tilde{{\bm{\theta}}}};\Omega(\sqrt{d/k}))\leq\exp(-\Omega(\log^{2}d))$.

###### Theorem 4.7 .

###### Proposition 4.8 .

## 5 Analysis of Training Dynamics for Feature-Averaging Regime

In this section, we present a proof sketch of Theorem [4.5](#S4.Thmtheorem5), where we provide a detailed analysis of training dynamics in feature-averaging regime.

### 5.1 Deriving Dynamics of Coefficients From Gradient Descent

By rigorously analyzing the gradient descent iterations, we know that each neuron is situated within a span that encompasses the collective cluster features and the intrinsic noise of the training data points. This span is explicitly characterized by the weight-feature correlations, which is shown as:

###### Lemma 5.1 (Weight Decomposition) .

During the training dynamics, there exists the following normalized coefficient sequences $\lambda_{s,r,j}^{(t)}$ and $\sigma_{s,r,i}^{(t)}$ for each pair $s\in\{-1,+1\},r\in[m],j\in[k],i\in[n]$ such that

$$ $\bm{w}_{s,r}^{(t)}=\bm{w}_{s,r}^{(0)}+\sum_{j\in[k]}\lambda_{s,r,j}^{(t)}\|\bm{\mu}_{j}\|^{-2}\bm{\mu}_{j}+\sum_{i\in[n]}\sigma_{s,r,i}^{(t)}\|\bm{\xi}_{i}\|^{-2}\bm{\xi}_{i}.$ $$

In light of Lemma [5.1](#S5.Thmtheorem1) and
the second item of the above theorem indicates that
$\bm{w}_{s,r}^{(t)}$ is approximately proportional to
the average of features in $J_{s}$ (the coefficients from the same class are large and approximately the same, and those from the opposite
class are small).

In order to deal with the behavior of ReLU activation, we define $S_{s,i}^{(t)}:=\{j\in[m]:\langle\bm{w}_{s,j}^{(t)},\bm{x}_{i}\rangle+\bm{b}_{s,j}^{(t)}>0\}$, for $s\in\{-1,+1\}$ and $i\in[n]$, denoting the set of indices of neurons in positive or negative class (determined by $s$) which is activated by training data point $\bm{x}_{i}$ at time step $t$. Then, we apply Lemma [5.1](#S5.Thmtheorem1) to the gradient descent iteration ([1](#S3.E1)), deriving the following result.

###### Lemma 5.2 (Updates of Coefficients $\lambda_{s,r,j}^{(t)},\sigma_{s,r,i}^{(t)}$ ) .

For each pair $s\in\{-1,+1\},r\in[m],j\in[k],i\in[N]$ and time $t\geq 0$, we have the following update equations:

$$ $\displaystyle\lambda_{s,r,j}^{(t+1)}=\lambda_{s,r,j}^{(t)}-\dfrac{s\eta}{nm}\cdot\sum\limits_{i\in I_{j}}\ell_{i}^{\prime(t)}\|\bm{\mu_{j}}\|^{2}\mathbbm{1}\left(r\in S_{s,i}^{(t)}\right),$ (2) $\displaystyle\sigma_{s,r,i}^{(t+1)}=\sigma_{s,r,i}^{(t)}-\dfrac{s\eta}{nm}\cdot\ell_{i}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}\mathbbm{1}\left(r\in S_{s,i}^{(t)}\right),$ (3) $$

where $\ell_{i}^{\prime(t)}:=\ell^{\prime}(y_{i}f_{\bm{\theta}^{(t)}}(\bm{x}_{i}))$ denotes the point-wise loss derivative at point $\bm{x}_{i}$, and $I_{j}:=\{i\in[n]:\bm{x}_{i}\textit{ in cluster }j\}$ denotes the set of the training points in the $j$-th cluster.

According to equations ([2](#S5.E2)) and ([3](#S5.E3)) from Lemma [5.2](#S5.Thmtheorem2), we know

$$ $\lambda_{s,r,j}^{(t)}=\sum_{i\in I_{j}}\frac{\|\bm{\xi}_{i}\|^{2}}{\|\bm{\mu}_{j}\|^{2}}\sigma_{s,r,i}^{(t)}\approx\sum_{i\in I_{j}}\sigma_{s,r,i}^{(t)},$ (4) $$

where we also use $\lambda_{s,r,j}^{(0)}=\sigma_{s,r,i}^{(0)}=0$ and the fact that, w.h.p., we have $\|\bm{\xi}_{i}\|\approx\sqrt{d}=\|\bm{\mu}_{j}\|$. It suggests that we only need to focus on the dynamics of the noise coefficients $\sigma_{s,r,i}^{(t)}$ (i.e., equation ([3](#S5.E3))).

###### Lemma 5.1 (Weight Decomposition) .

###### Lemma 5.2 (Updates of Coefficients λ s , r , j ( t ) , σ s , r , i ( t ) superscript subscript 𝜆 𝑠 𝑟 𝑗 𝑡 superscript subscript 𝜎 𝑠 𝑟 𝑖 𝑡 \lambda_{s,r,j}^{(t)},\sigma_{s,r,i}^{(t)} ) .

### 5.2 Two Key Techniques about Loss Derivative and Activation Region

It seems that the main difficulty in analyzing the iteration ([3](#S5.E3)) is addressing the time-variant loss derivative $\ell_{i}^{\prime(t)}$ and $\operatorname{ReLU}$ activation region $S_{s,i}^{(t)}$. To overcome these two challenges, we provide two corresponding key techniques (Lemma [5.3](#S5.Thmtheorem3) and Lemma [5.5](#S5.Thmtheorem5)) as follows, which can usefully simplify the analysis of noise coefficients’ dynamics.

Key Technique 1: Bounding Loss Derivative Ratio.
We will establish the connection between loss derivative ratio $y_{i}\ell_{i}^{\prime(t)}/y_{j}\ell_{j}^{\prime(t)}$ and the training data margin gap $\Delta_{q}^{(t)}(i,j):=q_{i}^{(t)}-q_{j}^{(t)}$, where $q_{i}^{(t)}$ denotes the margin of the $i$-th training data at iteration $t$ defined as $q_{i}^{(t)}:=y_{i}f_{\theta^{(t)}}\left(\bm{x}_{i}\right)$. Then, We have:

###### Lemma 5.3 (Training data margin are balanced during training dynamics) .

There exists a time threshold $T_{0}$ such that, for any time $1\leq t\leq T_{0}$ and distinct data points $(\bm{x}_{i},\bm{x}_{j})$, it holds that

$$ $\Delta_{q}^{(t)}(i,j)\leq\epsilon(k),$ (5) $$

where we use $\epsilon(k)$ to denote a time-independent error term satisfying $\epsilon(k)\rightarrow 0$ as $k\rightarrow\infty$.

According to Lemma [5.3](#S5.Thmtheorem3), for any distinct training data points $\bm{x}_{i}$ and $\bm{x}_{j}$ with the same label, the loss derivative ratio can be bounded as:

$$ $y_{i}\ell_{i}^{\prime(t)}/y_{j}\ell_{j}^{\prime(t)}\approx\operatorname{exp}(\Delta_{q}^{(t)}(j,i))\approx 1+\Delta_{q}^{(t)}(j,i)\overset{(\ref{equ:delta})}{=}1\pm o(1),$ (6) $$

where the first approximation holds due to $\ell^{\prime}(z)=1/(1+\operatorname{exp}(z))$ and we use the fact $e^{z}\approx 1+z$ for small $z$ in the second approximation.

###### Remark 5.4 .

This method was initially proposed by Chatterji and Long () in the context of benign overfitting for linear classification and was subsequently extended to networks with non-linear activation (Frei et al., ; Kou et al., ). In this paper, we extend the auto-balance technique of Kou et al. () from the single-feature case to our multi-cluster scenario to prove Lemma [5.3](#S5.Thmtheorem3).

Key Technique 2: Analyzing ReLU Activation Regions. Then, we turn to the analysis of the activation regions $S_{s,i}^{(t)}$. In fact, after the first gradient descent update, the set of activated neurons can be described in the following lemma.

###### Lemma 5.5 (Each training data can activate all its corresponding neurons) .

For the same time threshold $T_{0}$ as that in Lemma [5.3](#S5.Thmtheorem3) and all time $1\leq t\leq T_{0}$, it holds that
$S_{1,i}^{(t)}=[m]$ for all $i\in I_{+}$ and $S_{-1,i}^{(t)}=[m]$ for all $i\in I_{-}$, where $I_{+}:=\{i:i\in I,y_{i}=1\}$ and $I_{-}:=\{i:i\in I,y_{i}=-1\}$.

We rewrite our model as $f_{\bm{\theta}^{(t)}}=f_{1}^{(t)}+f_{-1}^{(t)}$, where $f_{s}^{(t)}:=\frac{s}{m}\sum_{r\in[m]}\operatorname{ReLU}(\langle\bm{w}_{s,r}^{(t)},\bm{x}\rangle+b_{s,r}^{(t)}),s\in\{-1,1\}$. Then, Lemma [5.5](#S5.Thmtheorem5) manifests that $f_{s}^{(t)}$ is linear in the training data point $(\bm{x}_{i},y_{i})$ with label $y_{i}=s$. If we show $f_{-y_{i}}^{(t)}$ keeps small, we will have the linearization $f_{\bm{\theta}^{(t)}}\approx f_{y_{i}}^{(t)}=\frac{y_{i}}{m}\sum_{r\in[m]}(\langle\bm{w}_{y_{i},r}^{(t)},\bm{x}_{i}\rangle+b_{y_{i},r}^{(t)})$, which allows us to approximate the data margin by noise coefficients (applying Lemma [5.1](#S5.Thmtheorem1) and ([4](#S5.E4))).

###### Remark 5.6 .

Indeed, we use induction to prove Lemma [5.3](#S5.Thmtheorem3) and Lemma [5.5](#S5.Thmtheorem5) together (see Lemma [B.7](#A2.Thmtheorem7) in the appendix), where we show the case when $t=1$ by our small initialization assumption and use the auto-balance technique to complete the inductive step (see the full proof in Appendix [B.1](#A2.SS1)).

###### Lemma 5.3 (Training data margin are balanced during training dynamics) .

###### Remark 5.4 .

###### Lemma 5.5 (Each training data can activate all its corresponding neurons) .

###### Remark 5.6 .

### 5.3 Proof Sketch of Theorem 4.5

Now, based on the two key techniques above, we provide a proof
sketch of Theorem [4.5](#S4.Thmtheorem5), which consists of
five steps.

Step 1: Proving that feature coefficient ratio $\lambda_{s_{1},r_{1},j_{1}}^{(T)}/\lambda_{s_{2},r_{2},j_{2}}^{(T)}(j_{1}\in J_{s_{1}},j_{2}\in J_{s_{2}})$ is close to $1$.
By Lemma [5.5](#S5.Thmtheorem5), for all $s\in\{-1,1\},i\in I_{s}$, we know

$$ $\sigma_{s,r,i}^{(t+1)}=\sigma_{s,r,i}^{(t)}-\frac{s\eta}{nm}\ell_{i}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}.$ (7) $$

Combined with the loss derivative ratio bound ([6](#S5.E6)), it furthermore implies that the noise coefficient ratio is close to $1$, i.e., for any $r_{1},r_{2}\in[m],i_{1},i_{2}\in I$, we have

$$ ${\sigma_{y_{i_{1}},r_{1},i_{1}}^{(t)}}/{\sigma_{{y_{i_{2}}},r_{2},i_{2}}^{(t)}}\overset{(\ref{equ:noise_update})}{\approx}{\sum_{t^{\prime}=0}^{t}\ell_{i_{1}}^{\prime(t^{\prime})}}/{\sum_{t^{\prime}=0}^{t}\ell_{i_{2}}^{\prime(t^{\prime})}}\overset{(\ref{equ:loss_bound})}{\approx}1\pm o(1).$ (8) $$

Thus, for any $s_{1},s_{2}\in\{-1,1\},j_{1}\in J_{s_{1}},j_{2}\in J_{s_{2}}$ and time $t\leq T$, we can derive

$$ ${\lambda_{s_{1},r_{1},j_{1}}^{(t)}}/{\lambda_{s_{2},r_{2},j_{2}}^{(t)}}\overset{(\ref{equ:lam_sig})}{\approx}{\sum_{i_{1}\in I_{j_{1}}}\sigma_{s_{1},r_{1},i_{1}}^{(t)}}/{\sum_{i_{2}\in I_{j_{2}}}\sigma_{s_{2},r_{2},i_{2}}^{(t)}}\overset{(\ref{equ:noise_ratio})}{\approx}{|I_{j_{1}}|}/{|I_{j_{2}}|}=1\pm o(1).$ $$

Step 2: Proving that $\lambda_{s,r,j}^{(T)}$ attains $\Omega(1)$ for $j\in J_{s}$, and keeps $o(1)$ for $j\in J_{-s}$.
By induction, we can show that both bias terms $b_{s,r}^{(t)}$ and $\lambda_{s,r,j}^{(t)}(j\in J_{-s})$ keep $o(1)$-order during the learning process (Lemma [B.10](#A2.Thmtheorem10) and Corollary [B.16](#A2.Thmtheorem16)), which thereby implies the following approximation, i.e., for any $s\in\{-1,1\},r\in[m],i\in I_{j}$, we have

$$ $\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\approx\lambda_{s,r,j}^{(t)}+\sigma_{s,r,i}^{(t)},$ (9) $$

where we also need time $t$ satisfying $t=\operatorname{exp}(\tilde{O}(k^{0.5}))$ (see details in Lemma [B.13](#A2.Thmtheorem13)).

Then, for any $s\in\{-1,1\},i\in I_{j}$ and data point $(\bm{x}_{i},y_{i})$ satisfying $y_{i}=-s$, we know

$$ $\operatorname{ReLU}(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)})\overset{(\ref{equ:linear_appr})}{=}\operatorname{ReLU}(\lambda_{s,r,j}^{(t)}+\sigma_{s,r,i}^{(t)}+o(1))\leq o(1),$ (10) $$

where the last inequality holds due to $\lambda_{s,r,j}^{(t)},\sigma_{s,r,i}^{(t)}\leq 0$ (Lemma [B.5](#A2.Thmtheorem5)).

Next, we approximate the model output for training data point $(\bm{x}_{i},y_{i})$ belonging to the $j$-th cluster as

$$ $\displaystyle f_{\bm{\theta}^{(t)}}(\bm{x}_{i})$ $\displaystyle=\sum_{s\in\{-1,1\}}\sum_{r\in[m]}\frac{s}{m}\operatorname{ReLU}(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)})\overset{(\ref{equ:neg_small})}{\approx}\frac{y_{i}}{m}\sum_{r\in[m]}(\langle\bm{w}_{y_{i},r}^{(t)},\bm{x}_{i}\rangle+b_{y_{i},r}^{(t)})$ (11) $\displaystyle\overset{(\ref{equ:linear_appr})}{\approx}\frac{y_{i}}{m}\sum_{r\in[m]}(\lambda_{y_{i},r,j}^{(t)}+\sigma_{y_{i},r,i}^{(t)})\overset{(\ref{equ:lam_sig})}{\approx}\frac{y_{i}}{m}\sum_{r\in[m]}\bigg{(}\sum_{i^{\prime}\in I_{j}}\sigma_{y_{i},r,i^{\prime}}^{(t)}+\sigma_{y_{i},r,i}^{(t)}\bigg{)}\overset{(\ref{equ:noise_ratio})}{\approx}y_{i}(|I_{j}|+1)\sigma_{y_{i},1,i}^{(t)}.$ $$

Therefore, we derive the following approximate update w.r.t. $\sigma_{y_{i},1,i}^{(t)}$, i.e., for any iteration $t\in[0,\operatorname{exp}(\tilde{O}(k^{1/2}))]$, we have $\sigma_{y_{i},1,i}^{(t+1)}\overset{(\ref{equ:noise_update})(\ref{equ:func})}{\approx}\sigma_{y_{i},1,i}^{(t)}+\frac{\eta d}{nm}\operatorname{exp}\left(-\frac{n}{k}\sigma_{y_{i},1,i}^{(t)}\right)$ (Lemma [B.21](#A2.Thmtheorem21)). By leveraging $\operatorname{log}(z+1)-\operatorname{log}(z)\approx\frac{1}{z}$, we inductively prove $\sigma_{y_{i},1,i}^{(t)}\approx\frac{k}{n}\operatorname{log}(\eta t)$ (Lemma [B.7](#A2.Thmtheorem7)) and $\lambda_{s,r,j}^{(t)}\approx\operatorname{log}(\eta t),j\in J_{s}$ (Lemma [B.23](#A2.Thmtheorem23)). When the assumption $T=\Omega(\eta^{-1})$ holds, we have $\lambda_{s,r,j}^{(T)}=\Omega(1),j\in J_{s}$.

Step 3: Gradient descent leads the network to the feature-averaging regime.
We can choose $\lambda^{(T)}=\lambda_{1,1,j_{0}}^{(T)}$ for some $j_{0}\in J_{+}$ as the representative of $\Lambda^{(T)}=\{\lambda_{s,r,j}^{(T)}:s\in\{-1,+1\},r\in[m],j\in J_{s}\}$.

Combining the result in Step 1 and Step 2, we know that for any $\lambda_{s,r,j}^{(T)}\in\Lambda^{(T)}$,

$$ $\lambda^{(T)}\approx\lambda_{s,r,j}^{(T)}\approx\log(\eta T).$ $$

We can also prove that the $\bm{w}_{s,r}^{(T)}$ is minimally affected by the coefficient $\sigma_{s,r,i}^{(T)}$ in weight decomposition. Thus, we have ([Lemma B.25](#A2.Thmtheorem25))

$$ $\bigg{\|}{\bm{w}}_{s,r}^{(T)}-\lambda^{(T)}\sum_{j\in J_{s}}\|\bm{\mu}_{j}\|^{-2}\bm{\mu}_{j}\bigg{\|}\leq o(d^{-1/2}).$ $$

Step 4: Proving that the clean accuracy is perfect.
For a randomly-sampled test data point $(\bm{x}=\bm{\mu}_{j}+\bm{\xi},y)\sim\mathcal{D}$ within cluster $j\in J_{y}$, we can prove that, with probability at least $1-\operatorname{exp}(\Omega(\log^{2}d))$, it holds that $|\langle\bm{w}_{s,r}^{(T)},\bm{\xi}\rangle|=o(1)$ for all $s\in\{\pm 1\},r\in[m]$ (Lemma [B.26](#A2.Thmtheorem26)). Then, for data satisfying the above condition, we can calculate the data margin as $yf_{\bm{\theta}^{(T)}}(\bm{x})\approx\frac{1}{m}\sum_{r\in[m]}\lambda_{s,r,j}^{(T)}=\Omega(1)>0$, which implies that $\mathrm{Acc}_{\mathrm{clean}}^{\mathcal{D}}(f_{{\bm{\theta}}^{(T)}})\geq 1-\exp(-\Omega(\log^{2}d))$.

Step 5: Proving that the robust accuracy is poor.
We consider the perturbation $\bm{\rho}=-\frac{2(1+c)}{k}\left(\sum_{j\in J_{+}}\bm{\mu}_{j}-\sum_{j\in J_{-}}\bm{\mu}_{j}\right)$. By applying Lemma [B.26](#A2.Thmtheorem26) again, we can derive that $\operatorname{sgn}(f_{\bm{\theta}^{(T)}}(\bm{x}+\bm{\rho}))\neq\operatorname{sgn}(f_{\bm{\theta}^{(T)}}(\bm{x}))$, which means $\mathrm{Acc}_{\mathrm{robust}}^{\mathcal{D}}(f_{{\bm{\theta}}^{(T)}};2(1+c)\sqrt{d/k})\leq\exp(-\Omega(\log^{2}d))$ and finishes the proof of Theorem [4.5](#S4.Thmtheorem5).

## 6 Experiments

We validate our theoretical results through a series of experiments. In [Section 6.1](#S6.SS1), we empirically verify that a binary classifier, trained on a binary classification task within a multi-cluster data setting, converges to a feature averaging solution, while incorporating fine-grained supervision information enables the model to learn decoupled features. Furthermore, in [Section 6.2](#S6.SS2), we demonstrate that incorporating fine-grained supervision information can enhance model robustness through the learning of decoupled features.

### 6.1 Empirical Verification of Feature Learning Process

oretical results about feature learning process,
We conduct numerical simulations on the synthetic dataset as we mentioned in Section [3](#S3) and also run real-world experiments on CIFAR-10 dataset. For both setups, we observe that the model trained with 2-class labels learns the feature-averaging solution (Figure [2(a)](#S6.F2.sf1) and [2(c)](#S6.F2.sf3)), while the model trained with finer-grained class labels learns decoupled features (Figure [2(b)](#S6.F2.sf2) and [2(d)](#S6.F2.sf4)).

Experiments on Multi-Cluster Data.
We generate the synthetic data following the data distribution in Section [3](#S3). Specifically, we choose the hyper-parameters as $k=10,d=3072,m=5,n=1000,\alpha=\sigma=1,\eta=0.001,\sigma_{w}=\sigma_{b}=0.00001,T=100$. For simplicity, we denote the weights of the two-layer network as $\bm{w}_{1},\bm{w}_{2},\dots,\bm{w}_{10}$ (where the first five weights correspond positive neurons and the other five weights correspond negative neurons). We also set the first five clusters as positive and the others as negative. Additionally, we provide an ablation study for other choices of hyper-parameters (see the details in Appendix [E.1](#A5.SS1)).

Experiments on CIFAR.
Beyond verifying the alignment between our theoretical findings and the results of numerical simulation on the synthetic multi-cluster data setup, as described in [Section 3](#S3), we also consider a more realistic setting where the multi-cluster structure of data naturally occurs. Specifically, we focus on a transfer learning setting, under which we utilize a pre-trained CLIP ViT-B-32 model (Radford et al., ) to obtain the image embedding for the CIFAR-10 dataset. We found that the embeddings of CIFAR-10 images approximately satisfy the multi-cluster structure, where the correlation between embeddings of images from the same class is significantly higher than that between embeddings of images from different classes (see detailed experimental results in [Section E.2](#A5.SS2)). We create a $2$-classification task from the CIFAR-10 dataset by merging the first $5$ classes into one class and the other $5$ classes into the other class. We apply two training strategies for training a two-layer neural network on this $2$-classification task: one is to train directly on the image embedding labeled for $2$-classification, and the other is to first train on the image embedding labeled for $10$ classes and then convert it to $2$-classification, where the two-layer network is as described in our theory (we fixe second layer as diagonal form, i.e., $f_{j}(\bm{z}):=\frac{1}{h}\sum_{r=1}^{h}\operatorname{ReLU}(\langle\bm{w}_{j,r},\bm{z}\rangle),\forall j\in[2]$ or $[10]$, and $\bm{z}$ denotes the image embedding). We set the width of the first layer to be $30$ ($h=3$) to ensure that the accuracy of the pre-trained model was not compromised.

For $10$-classification, we use $\bm{w}_{j}:=\frac{1}{h}\sum_{r=1}^{h}\bm{w}_{j,r}$ as the equivalent weight of $f_{j}$. For the $15$ positive weights and $15$ negative weights in the binary classification network, we equally divide them into $5$ positive classes and $5$ negative classes to ensure a fair comparison, between the two figures which ensures that two models both have the same form $\bm{F}:=(f_{1},f_{2},\dots,f_{10})\in\mathbb{R}^{10}$ and each sub-network $f_{j}$ corresponds to a weight vectors $\bm{w}_{j}$.

Figure: (a) SynData: 2-class
Refer to caption: /html/2410.10322/assets/x2.png

Experiment Results.
See experiment results in Figure [2](#S6.F2). It demonstrates our theory: model trained by only 2-class labels learns averaged features while model with fine-grained 10-class labels learns decoupled features. Figure [2(a)](#S6.F2.sf1) and Figure [2(c)](#S6.F2.sf3) correspond to feature-averaging result in Theorem [4.5](#S4.Thmtheorem5), where the correlations between each weight vector of positive (negative) neuron and all positive (negative) cluster features are uniformly larger than those between each weight vector of positive (negative) neuron and all negative (positive) cluster features; Figure [2(b)](#S6.F2.sf2) and Figure [2(d)](#S6.F2.sf4) correspond to feature-decoupling result in Theorem [4.7](#S4.Thmtheorem7), where the correlation matrix is nearly diagonal.

### 6.2 Robustness Improvement From Fine-Grained Supervision Information

To study the effect of adding more fine-grained supervision information, we imagine a scenario where a human annotator is asked to annotate the binarized dataset with 10 classes, which happens to result in the original dataset. We then compare the performance of directly training a binary classifier on the binarized dataset with that of training a 10-class classifier on the original dataset and then using its output to do binary classification.

Experiment Settings.
Following [Section 4.2](#S4.SS2),
to convert a 10-class model to a binary classifier,
we sum the probabilities corresponding to the 5 positive classes and subtract it by the sum of the probabilities corresponding to the 5 negative classes from the 10-class model’s output.
The robust accuracy is measured by using the standard PGD attacks (Madry et al., ) with different $\ell_{2}$-pertubation radius. We run experiments in the following datasets:

- •
Synthetic Dataset. We generate synthetic data and train the networks as the same as that in [Section 6.1](#S6.SS1).
- •
Binary Classification on MNIST and CIFAR-10. To further verify our theory in deep neural networks, on both MNIST and CIFAR-10 datasets, we train ResNet18 models from scratch with normal 10-classification labels and 2-classification labels (the 2-classification labels are generated by merging the first five and the last five classes from the 10-classification labels). We use stochastic gradient descent (SGD) with a batch size of 128, a momentum of 0.9, and a weight decay of $5\times 10^{-4}$ for training on both CIFAR-10 and MNIST datasets. We apply cosine learning rate scheduling with an initial learning rate of 0.1 for 200 epochs for training on CIFAR-10 and a fixed learning rate of 0.001 for 100 epochs for training on MNIST.

Experiment Results. The results are presented in Figure [3](#S6.F3). With the perturbation radius increasing, we can see that the models trained with 10-class labels have higher robust test accuracy than those trained with 2-class labels in all datasets. This collaborates
with our theoretical results (Theorem [4.5](#S4.Thmtheorem5) and [4.7](#S4.Thmtheorem7)) that models achieve better robustness with more supervised information.

Figure: Figure 3: Verifying robustness improvement: We compare adversarial robustness between model trained by 2-class labels (red line) and model trained by 10-class labels (blue line) on synthetic data (the left), MNIST (the middle) and CIFAR-10 (the right).
Refer to caption: /html/2410.10322/assets/x6.png

## 7 Connections of Our Results with Other Explanations of Adversarial Examples

(1) Approximate Linearity of the Model: Earlier hypothesis about the origin of adversarial examples (e.g., Goodfellow et al. ()) had proposed the idea that
the existence of adversarial examples is related to the fact the model $f_{\bm{\theta}}(\bm{x})$
is approximately linear. Subsequently, there is a sequence of theoretical studies
showing that adversarial examples exists abundantly in the input space for neural networks with random weights (without training) and a main insight is that such random networks is approximately linear and with high probability an input point is close to the decision boundary (by isoperimetry argument)
(see e.g., Gilmer et al. (); Bubeck et al. (); Bartlett et al. (); Montanari and Wu ()). Our Theorem [4.5](#S4.Thmtheorem5) proves similar approximate linearity (see details in the proof intuition of Theorem [D.2](#A4.Thmtheorem2))
(^1^11 Think of the special case that the weight vector corresponding to each neuron is exactly the average of the cluster means (${\bm{\mu}_{+}}:=\frac{1}{|J_{+}|}\sum_{j\in J_{+}}\bm{\mu}_{j}$ or ${\bm{\mu}_{-}}:=\frac{1}{|J_{-}|}\sum_{j\in J_{-}}\bm{\mu}_{j}$) and $\sum_{j=1}^{k}\bm{\mu}_{j}=\bm{0}$. In this case, the two-layer network reduces to a simple linear model w.r.t. the perturbation.)
and show it leads to adversarial examples for trained neural network (albeit with different data distribution from the aforementioned work).
Our result is also related to the dimpled manifold hypothesis (Shamir et al., ), which proposed that during training a neural network first finds a simple decision boundary
that is close to most training points.

Figure: Figure 4: A schematic illustration of the construction in Li et al. (): The positive class consists of blue points and the negative class the red points. In their lower bound, there are in fact exponentially blue points slightly above the hyperplane and exponentially many red ones slightly below it. The hyperplane has perfect clean accuracy but is non-robust, while a more robust classifier exists (by classifying the blue balls from the red balls). One can observe the conceptual similarity with Figure [1](#S4.F1).
Refer to caption: /html/2410.10322/assets/linear_separable.png

(2) Non-robust Features: Another appealing point of view was developed in Ilyas et al. (), which proposed that adversarial examples are related to the presence of non-robust features.
They showed empirically that neural networks learn both robust and non-robust features that are useful to classify clean images.
In image classification tasks, Ilyas et al. () visualized both robust and nonrobust features. While robust features are more perceptually meaningful for human, nonrobust features resemble
noise and artifacts. Interestingly, they showed that nonrobust feature
can be leveraged to construct adverserial examples for DNN.
Our paper presents a theoretical setting in which neural networks provably learn non-robust features (due to feature averaging), despite the existence of more robust features. Moreover, we prove that
the learnt non-robust feature (${\bm{\mu}_{+}}$ or ${\bm{\mu}_{-}}$) can be utilized to attack the feature-averaging network.

(3) Relation to the Lower Bound Examples in Li et al. ():
From the perspective of expressivity, Li et al. () constructed a lower bound example (see an illustration in Figure [4](#S7.F4)), for which there is non-robust linear classifier, but the set of robust solutions requires a hypothesis class of a much larger (in fact exponentially large) VC-dimension. This partially explains why neural networks are non-robust (unless they are exponentially large). The construction of our data distribution (as well as that in Vardi et al. (); Frei et al. ()) echos the essence of this lower bound example in spirit, and our results can be seen as an explanation from the perspective of optimization.

## 8 Conclusion

This paper exposes “Feature Averaging” as an implicit bias in gradient descent that may compromise the robustness of deep neural networks. Theoretical insights from a two-layer ReLU network reveal a tendency for gradient descent to average/combine individually meaningful features, which can lead to a loss of distinct discriminative information. We demonstrate that with more detailed feature level supervision, the networks can learn to differentiate these features, enhancing model robustness. This is supported by empirical evidence from both synthetic and real-world data, including MNIST and CIFAR-10. Our findings not only deepen our understanding of adversarial examples in deep learning but also suggest that fine-grained supervision can enhance the robustness of deep neural networks against adversarial attacks.

## Appendix A Preliminary Properties

In this section, we provide some useful properties of our training dataset and neural network learner at the initialization. These properties hold with high probability under our assumptions. Our subsequent proofs will be based on the validity of these properties.
The proofs of the inlined claims are concluded with the $\blacksquare$ symbol, while the proofs of the overarching results are concluded with the $\square$ symbol.

### A.1 Detailed Data Model and Assumptions

First, we recall the definition of multi-class feature data distribution that we defined in Section [3](#S3).

###### Definition A.1 (Multi-Cluster Data Distribution) .

Given $k$ vectors ${\bm{\mu}}_{1},\dots,{\bm{\mu}}_{k}\in\mathbb{R}^{d}$, called the cluster features, and a partition of $[k]$ into two disjoint sets $J_{\pm}=(J_{+},J_{-})$,
we define $\mathcal{D}(\{{\bm{\mu}}_{j}\}_{j=1}^{k},J_{\pm})$
as a data distribution on $\mathbb{R}^{d}\times\{-1,1\}$, where each data point $({\bm{x}},y)$ is generated as follows:

- 1.
Draw a cluster index as $j\sim\mathrm{Unif}([k])$;
- 2.
Set $y=+1$ if $j\in J_{+}$; otherwise $j\in J_{-}$ and set $y=-1$;
- 3.
Draw ${\bm{x}}:={\bm{\mu}}_{j}+{\bm{\xi}}$, where ${\bm{\xi}}\sim{\mathcal{N}}({\bm{0}},{\bm{I}}_{d})$.

For convenience, we write $\mathcal{D}$ instead of $\mathcal{D}(\{{\bm{\mu}}_{j}\}_{j=1}^{k},J_{\pm})$ if $\{{\bm{\mu}}_{j}\}_{j=1}^{k}$ and $J_{\pm}$ are clear from the context.
For $s\in\{\pm 1\}$, we write $J_{s}$ to denote $J_{+}$ if $s=+1$ and $J_{-}$ if $s=-1$.

To ease the analysis,
we make the following simplifying assumptions
on the distribution.

###### Assumption A.2 (Orthogonal Equinorm Cluster Features) .

The cluster features $\{{\bm{\mu}}_{j}\}_{j=1}^{k}$ satisfy the properties that (1) $\|{\bm{\mu}}_{j}\|=\sqrt{d}$ for all $j\in[k]$; and (2) ${\bm{\mu}}_{i}\perp{\bm{\mu}}_{j}$ for all $1\leq i<j\leq k$.

###### Assumption A.3 (Nearly Balanced Classification) .

The partition $J_{\pm}$ satisfies $c^{-1}\leq\frac{\lvert J_{+}\rvert}{\lvert J_{-}\rvert}\leq c$ for some absolute constant $c\geq 1$.

Next, we summarize the assumptions of these hyper-parameters that we mentioned in the main text, as listed below.

###### Assumption A.4 (Choices of Hyper-Parameters) .

We state the range of parameters for our proofs in the appendix to hold.

- •
$d=\Omega(k^{10})$ (recall $d$ is the data dimension)
- •
$c=\Theta(1)$ (recall $c$ is the balance ratio)
- •
$n\in[\Omega(k^{7}),\operatorname{exp}(O(\operatorname{log}^{2}(d)))]$ (recall $n$ is the number of samples)
- •
$m=\Theta(k)$ (recall $2m$ is the width of network learner)
- •
$\eta\leq O(d^{-2})$ (recall $\eta$ is the learning rate)
- •
$\sigma_{b}^{2}=\sigma^{2}_{w}\leq\eta k^{-5},$ (recall $\bm{w}_{s,r}^{(0)}\sim\mathcal{N}(\bm{0},\sigma_{w}^{2}\bm{I}_{d}),b_{s,r}^{(0)}\sim\mathcal{N}(0,\sigma_{b}^{2})$ give the initialization)

###### Remark A.5 (Discussion of Hyper-Parameter Choices) .

In this paper, we make specific choices of hyper-parameters for the sake of calculations (and we emphasize that these may not be the tightest possible choices), which is a widely-applied simplicity in the literature of feature learning works (Wen and Li, ; Chen et al., ; Allen-Zhu and Li, , ; Chidambaram et al., ). Namely, we need the data dimension $d$ to be a significantly larger polynomial in the number of clusters $k$ to ensure all $k$ cluster features $\bm{\mu}_{1},\bm{\mu}_{2},\dots,\bm{\mu}_{k}$ can be orthogonal within the space $\mathbb{R}^{d}$. The balance ratio $c$ is an absolute constant that is independent with $d$ and $k$. Our results can be extended to ${\bm{x}}:=\alpha{\bm{\mu}}_{j}+\sigma{\bm{\xi}}$ for some parameters $\alpha=\Theta(1)$ and $\sigma=\Theta(1)$, but here we set $\alpha=\sigma=1$ for simplicity. And we further require that $n$ is a large polynomial in $k$ due to our choice of large signal-noise-ratio (recall that we have $\alpha/\sigma=\Theta(1)$, which implies that $\|\bm{\mu}_{j}\|=\sqrt{d}\approx\|\bm{\xi}\|$ with high probability). We need $m\in[\operatorname{max}\{|J_{+}|,|J_{-}|\},k]$ for the existence of robust solution (Theorem [D.3](#A4.Thmtheorem3)). We assume the learning rate $\eta$ and the initialization magnitude $\sigma_{w},\sigma_{b}$ are sufficiently small, which helps the network to be trained in the feature learning regime (Lyu et al., ; Cao et al., ; Allen-Zhu and Li, ; Kou et al., ).

###### Definition A.1 (Multi-Cluster Data Distribution) .

###### Assumption A.2 (Orthogonal Equinorm Cluster Features) .

###### Assumption A.3 (Nearly Balanced Classification) .

###### Assumption A.4 (Choices of Hyper-Parameters) .

###### Remark A.5 (Discussion of Hyper-Parameter Choices) .

### A.2 Useful Properties of the Training Dataset

Now, we introduce some useful notations, for simplifying our proof.

- •
Denote $I=\{i:i=1,2,\cdots,n\}$ as the set of indices
of all training data points.
- •
Define $c(\cdot)$ as the map $I\rightarrow J$ where $c(i)$ represents the index of the cluster to which point $\bm{x}_{i}$ belongs.
- •
$I_{j}=\{i:i\in I,c(i)=j\}$ denotes the set of the training points in the $j$-th cluster.
- •
$I_{+}=\{i:i\in I,c(i)\in J_{+}\}$ and $I_{-}=\{i:i\in I,c(i)\in J_{-}\}$ denote the index sets of all positive class data points and negative class data points respectively.

Then, under Assumption [A.4](#A1.Thmtheorem4), we show that the following properties of the training dataset hold with high probability. Since our data distribution is similar to Frei et al. (), and most of the proofs of these properties follow standard approaches, and are similar to the ones presented in Frei et al. ().

Recall the cumulative distribution function (CDF) of the standard normal distribution, usually denoted as $\Phi(x)$, which is defined as the integral

$$ $\Phi(x):=\frac{1}{\sqrt{2\pi}}\int_{-\infty}^{x}e^{-t^{2}/2}dt=\frac{1}{\sqrt{2\pi}}\int_{-x}^{\infty}e^{-t^{2}/2}dt.$ $$

Additionally, we have the following commonly used bounds on the tail probabilities of the standard normal distribution.

$$ $\displaystyle\left(\dfrac{x^{2}-1}{\sqrt{2\pi}x^{3}}\right)\exp(-x^{2}/2)\leq 1-\Phi(x)\leq\dfrac{1}{\sqrt{2\pi}x}\exp(-x^{2}/2)$ (12) $$

###### Proposition A.6 .

Let $\Delta=4\sqrt{d}\ln(d)$ and $\delta=8n^{2}d^{-\frac{\ln(d)}{2}}+2k\left(n/k\right)^{-\ln(n/k)}$.
With probability at least $1-\delta$ over sampled training dataset $\mathcal{S}\sim\mathcal{D}^{n}$, we have the following properties:

- 1.
For every $i\in I$ we have $\sqrt{d}-2\ln(d)\leq\left\|\bm{\xi}_{i}\right\|\leq\sqrt{d}+2\ln(d)$.
- 2.
For every $i\in I$ we have $\left\|\bm{x}_{i}\right\|\leq 3\sqrt{d}$.
- 3.
For every $i,j\in I,i\neq j$ we have $\left|\left\langle\bm{\xi}_{i},\bm{\xi}_{j}\right\rangle\right|\leq\Delta$.
- 4.
For every $i\in I$ and $j\in J$ we have $\left|\left\langle\bm{\mu}_{j},\bm{\xi}_{i}\right\rangle\right|\leq\Delta$.
- 5.
For every $i,j\in I$ with $c(i)\neq c(j)$ we have $\left|\left\langle\bm{x}_{i},\bm{x}_{j}\right\rangle\right|\leq\Delta$.
- 6.
For every $i,j\in I$ with $c(i)=c(j)$ we have $\dfrac{1}{2}d\leq\left\langle\bm{x}_{i},\bm{x}_{j}\right\rangle\leq 2d$.
- 7.
For every $j\in J$ we have $\dfrac{n}{k}-\sqrt{\dfrac{2n}{k}}\ln(\dfrac{n}{k})\leq|I_{j}|\leq\dfrac{n}{k}+\sqrt{\dfrac{2n}{k}}\ln(\dfrac{n}{k}).$

###### Remark A.7 .

Property 1 and Property 2 show that the data-wise noise and training data point are bounded, i.e., $\|\bm{\xi}_{i}\|\approx\sqrt{d}$ and $\|\bm{x}_{i}\|=O(\sqrt{d})$. Property 3 and Property 4 show that the correlation between different noises (or between cluster center feature and random noise) is very small. Property 5 and Property 6 suggest that the correlation between training data points of different clusters is very small, but the correlation between training data points within the same cluster is very large. Property 7 manifests that the training dataset $\mathcal{S}$ approximately includes $\frac{n}{k}(=\Omega(k^{6}))$ examples from each cluster.

###### Proof of Proposition A.6 .

Now, we prove Property 1-7 one by one.

Property 1: We notice that $\left\|\bm{\xi}\right\|^{2}$ follows the Chi-squared distribution.

The concentration bound in Lemma 1 by Laurent and Massart () implies that for all $t\geq 0$, we have

$$ $\displaystyle\operatorname{Pr}\left[\left\|\bm{\xi}\right\|^{2}-d\geq 2\sqrt{dt}+2t\right]$ $\displaystyle\leq e^{-t},$ $\displaystyle\operatorname{Pr}\left[\left\|\bm{\xi}\right\|^{2}-d\leq-2\sqrt{dt}\right]$ $\displaystyle\leq e^{-t}.$ $$

Plugging in $t=\ln^{2}(d)$, we can see that

$$ $\displaystyle\operatorname{Pr}\left[\|\bm{\xi}\|^{2}\geq(\sqrt{d}+2\ln(d))^{2}\right]\leq\operatorname{Pr}\left[\left\|\bm{\xi}\right\|^{2}-d\geq 2\sqrt{d}\ln(d)+2\ln^{2}(d)\right]\leq d^{-\ln(d)},$ $\displaystyle\operatorname{Pr}\left[\|\bm{\xi}\|^{2}\leq(\sqrt{d}-2\ln(d))^{2}\right]\leq\operatorname{Pr}\left[\left\|\bm{\xi}\right\|^{2}-d\leq-2\sqrt{d}\ln(d)\right]\leq d^{-\ln(d)}.$ $$

Thus, we have

$$ $\operatorname{Pr}\left[\sqrt{d}-2\ln(d)\leq\left\|\bm{\xi}_{i}\right\|\leq\sqrt{d}+2\ln(d)\right]\leq 2d^{-\ln(d)}.$ (13) $$

Then by union bound, we have Property 1 holds for every $i\in I$ with probability at least $1-2nd^{-\ln(d)}$.

Property 2: When Property 1 holds, by triangle inequality, we know Property 2 holds:

$$ $\|\bm{x}_{i}\|\leq\|\bm{\mu}_{c(i)}\|+\|\bm{\xi}_{i}\|\leq\sqrt{d}+\sqrt{d}+\ln(d)\leq 3\sqrt{d}.$ $$

The proofs for other properties require calculations pertaining to Gaussian distribution. We first introduce a useful lemma below.

###### Lemma A.8 .

Let $\bm{\xi}\sim\mathrm{N}\left(\mathbf{0},I_{d}\right)$. For any $\bm{x}\in\mathbb{R}^{d}$ we have

$$ $\operatorname{Pr}\left[\left|\left\langle\bm{x},\bm{\xi}\right\rangle\right|\geq\|\bm{x}\|\ln(d)\right]\leq 2d^{-\frac{\ln(d)}{2}}.$ $$

###### Proof of Lemma A.8 .

Note that $\left\langle\dfrac{\bm{x}}{\|\bm{x}\|},\bm{\xi}\right\rangle$ has the distribution $\mathcal{N}\left(0,1\right)$.

By standard Gaussian tail bound, we have for every $t\geq 0$ that $\operatorname{Pr}\left[\left|\left\langle\frac{\bm{x}}{\|\bm{x}\|},\bm{\xi}\right\rangle\right|\geq t\right]\leq 2\exp\left(-\dfrac{t^{2}}{2}\right)$.

Plugging in $t=\ln(d)$, we can see that

$$ $\operatorname{Pr}\left[\left|\left\langle\frac{\bm{x}}{\|\bm{x}\|},\bm{\xi}\right\rangle\right|\geq\ln(d)\right]\leq 2\exp\left(-\frac{\ln^{2}(d)}{2}\right)=2d^{-\frac{\ln(d)}{2}}.$ $$

∎

Property 3: Next, we prove Property 3 using the result in Lemma [A.8](#A1.Ex19).

Noting that if $\left|\left\langle\bm{\xi}_{i},\bm{\xi}_{j}\right\rangle\right|\geq\sqrt{2d}\ln(d)$, we have that at least one of the following holds:

- 1.
$\|\bm{\xi}_{i}\|\geq\sqrt{2d}$;
- 2.
$\left|\left\langle\frac{\bm{\xi}_{i}}{\|\bm{\xi}_{i}\|},\bm{\xi}_{j}\right\rangle\right|\geq\ln(d)$.

Now we bound the probabilities of these two events separately.
By Property 1, we have

$$ $\operatorname{Pr}[\|\bm{\xi_{i}}\|\geq\sqrt{2d}]\leq 2d^{-\ln(d)}.$ $$

Next, by Lemma [A.8](#A1.Ex19), we have

$$ $\operatorname{Pr}\left[\left|\left\langle\frac{\bm{\xi}_{i}}{\|\bm{\xi}_{i}\|},\bm{\xi}_{j}\right\rangle\right|\geq\ln(d)\right]\leq 2d^{-\frac{\ln(d)}{2}}.$ $$

Then, by union bound, we know that,

$$ $\operatorname{Pr}\left[\left|\left\langle\bm{\xi}_{i},\bm{\xi}_{j}\right\rangle\right|\geq\sqrt{2d}\ln(d)\right]\leq 2d^{-\frac{\ln(d)}{2}}+2d^{-\ln(d)}\leq 4d^{-\frac{\ln(d)}{2}}$ $$

Then, applying the union bound for all pairs $i,j\in I,i\neq j$,
we have $\left|\left\langle\bm{\xi}_{i},\bm{\xi}_{j}\right\rangle\right|\leq\sqrt{2d}\ln(d)^{2}$ holds with probability at least $1-4n^{2}d^{-\frac{\ln(d)}{2}}$.

Property 4: Applying Lemma [A.8](#A1.Ex19), we have

$$ $\operatorname{Pr}\left[\left|\left\langle\bm{\mu_{j}},\bm{\xi}\right\rangle\right|\geq\sqrt{d}\ln(d)\right]\leq 2d^{-\frac{\ln(d)}{2}}.$ $$

Then for all pairs $i\in I,j\in J$,
applying union bound we have that $\left|\left\langle\bm{\mu_{j}},\bm{\xi}_{i}\right\rangle\right|\leq\sqrt{d}\ln(d)$ for all $i\in I,j\in J$ holds with probability at least $1-2n^{2}d^{-\frac{\ln(d)}{2}}$.

Property 5: By using the results above, we have

$$ $\displaystyle|\langle\bm{x}_{i},\bm{x}_{j}\rangle|$ $\displaystyle\leq\left|\langle\bm{\mu}_{c(i)},\bm{\mu}_{c(j)}\rangle\right|+\left|\langle\bm{\mu}_{c(i)},\bm{\xi}_{j}\rangle\right|+\left|\langle\bm{\mu}_{c(j)},\bm{\xi}_{i}\rangle\right|+\left|\langle\bm{\xi}_{i},\bm{\xi}_{j}\rangle\right|$ $\displaystyle=\left|\langle\bm{\mu}_{c(i)},\bm{\xi}_{j}\rangle\right|+\left|\langle\bm{\mu}_{c(j)},\bm{\xi}_{i}\rangle\right|+\left|\langle\bm{\xi}_{i},\bm{\xi}_{j}\rangle\right|$ $\displaystyle\leq 4\sqrt{d}\ln(d)=\Delta.$ $$

Property 6: By using the results above and noting that $\langle\mu_{c(i)},\mu_{c(j)}\rangle=d$ for $i,j$ with $c(i)=c(j)$, we have

$$ $\displaystyle\left|\langle\bm{x}_{i},\bm{x}_{j}\rangle-d\right|$ $\displaystyle\leq\left|\langle\bm{\mu}_{c(i)},\bm{\xi}_{j}\rangle\right|+\left|\langle\bm{\mu}_{c(j)},\bm{\xi}_{i}\rangle\right|+\left|\langle\bm{\xi}_{i},\bm{\xi}_{j}\rangle\right|$ $\displaystyle\leq 4\sqrt{d}\ln(d)=\Delta.$ $$

Thus, we have

$$ $\dfrac{1}{2}d\leq\left\langle\mathbf{x}_{i},\mathbf{x}_{j}\right\rangle\leq 2d.$ $$

Property 7: We define $X_{i}$ for $i\in I$
as the indicator random variable that the $i$-th point is in the
$j$-th cluster.
It takes value $1$ with probability $\frac{1}{k}$, and 0 with
probability $1-\frac{1}{k}$.

Then we know that $|I_{j}|=\sum\limits_{i\in I}X_{i}$. Applying Chernoff bound, we have that

$$ $\displaystyle\operatorname{Pr}\left[\dfrac{n}{k}-\sqrt{\dfrac{2n}{k}}\ln(\dfrac{n}{k})\leq|I_{j}|\leq\dfrac{n}{k}+\sqrt{\dfrac{2n}{k}}\ln(\dfrac{n}{k})\right]\geq 1-2\exp(-\ln^{2}(n/k))=1-2(n/k)^{-\ln(n/k)}.$ $$

Then for all $j\in J$, applying union bound, we have Property 7 holds with probability at least $1-2k(n/k)^{-\ln(n/k)}$.

Combining all of the above together, we have Proposition [A.6](#A1.Thmtheorem6) holds with probability at least $1-8n^{2}d^{-\frac{\ln(d)}{2}}-2k(n/k)^{-\ln(n/k)}$.
∎

###### Proposition A.6 .

###### Remark A.7 .

###### Proof of Proposition A.6 .

###### Lemma A.8 .

###### Proof of Lemma A.8 .

### A.3 Useful Properties of the Network Initialization

The proofs for properties of the network initialization require the range of the maximum value obtained from multiple independent samples drawn from a Gaussian distribution.
We first present the following useful lemma.

###### Lemma A.9 (Concentration of Maximum of Gaussians) .

Let $X_{i}\sim\mathcal{N}(0,1),1\leq i\leq l$ be i.i.d. random variables.
Denote $X_{\text{max}}=\max_{1\leq i\leq l}\{X_{i}\},X_{\text{min}}=\min_{1\leq i\leq l}\{X_{i}\}$. For any $t\geq 0$, we have

- •
$\operatorname{Pr}\left[X_{\text{min}}\leq-\sqrt{2\log(l)}-t\right]\leq\frac{1}{2}\exp\left(-t^{2}/2\right)$,
- •
$\operatorname{Pr}\left[X_{\text{max}}\leq\sqrt{2\log(l)-t}\right]\leq\exp\left(-\frac{e^{t/2}}{\sqrt{2\pi}(\sqrt{2\log(l)}+1)}\right)$.

###### Proof of Lemma A.9 .

The proof is standard and similar to Proposition A.1 and A.2 from Chidambaram et al. ().
We include it for convenience of the readers.

For any $a\geq 0$, we have

$$ $\operatorname{Pr}\left[X_{\text{min}}\geq-a\right]=(\Phi(a))^{l}=(1-Q(a))^{l},$ $$

where $Q(x)=1-\Phi(x)$. By using $(1-x)^{n}\geq 1-nx$ for any $x\in[0,1]$ and $n\in\mathbb{N}$, we then get

$$ $\operatorname{Pr}\left[X_{\text{min}}\geq-a\right]\geq 1-lQ(a).$ $$

Now we use the elementary inequality for the tail of the normal distribution:

$$ $Q(a)\leq\frac{1}{2}e^{-\frac{a^{2}}{2}},$ $$

so that

$$ $\operatorname{Pr}\left[X_{\text{min}}\geq-a\right]\geq 1-\dfrac{1}{2}le^{-\frac{a^{2}}{2}}.$ $$

Plugging $a=\sqrt{2\log(l)}+t$, we get

$$ $\displaystyle\operatorname{Pr}\left[X_{\text{min}}\geq-\sqrt{2\log(l)}-t\right]$ $\displaystyle\geq 1-l\exp\left(-\frac{(\sqrt{2\log(l)}+t)^{2}}{2}\right)$ $\displaystyle=1-\dfrac{1}{2}l\exp\left(-\frac{2\log(2l)+t^{2}}{2}\right)$ $\displaystyle=1-\frac{1}{2}e^{-\frac{t^{2}}{2}}.$ $$

Similar to the previous proof, we know that

$$ $\operatorname{Pr}\left[X_{\text{max}}\leq a\right]=\left(\Phi\left(a\right)\right)^{l}\leq\left(1-\frac{a}{\sqrt{2\pi}\left(a^{2}+1\right)}\exp\left(-\dfrac{a^{2}}{2}\right)\right)^{l}.$ $$

Plugging $a=\sqrt{2\log(l)-t}$,

$$ $\displaystyle\operatorname{Pr}\left[X_{\text{max}}\leq\sqrt{2\log(l)-t}\right]$ $\displaystyle\leq\left(1-\frac{\sqrt{2\log(l)-t}}{\sqrt{2\pi}l(2\log(l)-t+1)}\exp\left(\dfrac{t}{2}\right)\right)^{l}$ $\displaystyle\leq\left(1-\frac{\exp\left(t/2\right)}{\sqrt{2\pi}l(\sqrt{2\log(2l)}+1)}\right)^{l}$ $\displaystyle\leq\exp\left(-\frac{\exp\left(t/2\right)}{\sqrt{2\pi}(\sqrt{2\log(2l)}+1)}\right)$ $$

∎

By applying Proposition [A.6](#A1.Thmtheorem6), we can derive the following result, which gives the range of network parameters at the initialization.

###### Proposition A.10 .

With probability at least $1-4md^{-\ln(d)}-2m^{-3}$ , we have the following properties for our network initialization:

- •
For any $s\in\{-1,+1\},r\in[m]$, we have $\sigma_{w}\left(\sqrt{d}-2\ln(d)\right)\leq\|\bm{w}_{s,r}^{(0)}\|\leq\sigma_{w}\left(\sqrt{d}+2\ln(d)\right).$
- •
For any $s\in\{-1,+1\},r\in[m]$, we have $|b_{s,r}^{(0)}|\leq 2\sigma_{b}\sqrt{2\ln(m)}.$

###### Proof of Proposition A.10 .

For $\bm{w}_{s,r}^{(0)}$, reusing the same argument as the proof of Property (1) in Proposition [A.6](#A1.Thmtheorem6), we know that

$$ $\sigma_{w}\left(\sqrt{d}-2\ln(d)\right)\leq\|\bm{w}_{s,r}^{(0)}\|\leq\sigma_{w}\left(\sqrt{d}+2\ln(d)\right)$ $$

holds for all $s\in\{-1,+1\},r\in[m]$ with probability at least $1-4md^{-\ln(d)}$.

For $b_{s,r}^{(0)}$, by standard Gaussian tail bound, we know that

$$ $\displaystyle\text{Pr}\left[|b_{s,r}^{(0)}|\geq 2\sigma_{b}\sqrt{2\ln(m)}\right]\leq\exp\left(-4\ln(m)\right)=m^{-4}.$ $$

Then using union bound, we know that $|b_{s,r}^{(0)}|\geq 2\sigma_{b}\sqrt{2\ln(m)}$ holds for all $s\in\{-1,+1\},r\in[m]$ with probability at least $1-2m^{-3}$.

In conclusion, we know that these properties hold with probability at least $1-2md^{-\ln(d)}-2m^{-3}$.
∎

We then show that each neuron is activated by at least one training data point in each cluster upon initialization with high probability. We first formally define the notion of activation region as
follows.

###### Definition A.11 (Activation Region over Data Input) .

Let $T_{s,r,j}:=\{i\in I_{j}:\langle\bm{w}_{s,r}^{(0)},\bm{x}_{i}\rangle+b_{s,r}^{(0)}\geq 0\}$ be the set of indices of training data points in the $j$-th cluster which can activate the $r$-th neuron with weight $\bm{w}_{s,r}$ at time step 0.

Then, we give the following result about the activation region $T_{s,r,j}$.

###### Proposition A.12 .

Assuming Proposition [A.6](#A1.Thmtheorem6) and Proposition [A.10](#A1.Thmtheorem10) holds. Then with probability at least $1-m^{-0.01}-2mk\exp\left(-\frac{n}{9km^{2}}\right)$, for all $s\in\{-1,+1\},r\in[m],j\in J$, we have

$$ $|T_{s,r,j}|\geq\dfrac{n}{3km^{2}}.$ $$

###### Proof of Proposition A.12 .

Given $\bm{\mu}_{j}$ for $j\in J$, we have $\langle\bm{w}_{s,r}^{(0)},\bm{\mu}_{j}\rangle\sim\mathcal{N}(0,\sigma_{w}\sqrt{d})$.

Using the conclusion in Lemma [A.9](#A1.Thmtheorem9), with probability at least $1-m^{0.01}$, we have

$$ $\displaystyle\operatorname{min}_{s\in\{-1,+1\},r\in J}\langle\bm{w}_{s,r}^{(0)},\bm{\mu}_{j}\rangle$ $\displaystyle\geq-1.1\sigma_{w}\sqrt{d}\sqrt{2\ln(2m)}.$ $$

In the following proof, we assume that the above conclusion holds.

Given $\bm{w}_{s,r}^{(0)}$, we have $\langle\bm{w}_{s,r}^{(0)},\bm{\xi}_{i}\rangle\sim\mathcal{N}(0,\|\bm{w}_{s,r}^{(0)}\|)$.

Then by Gaussian tail bound ([12](#A1.E12)), we have

$$ $\displaystyle\text{Pr}\left[\langle\bm{w}_{s,r}^{(0)},\bm{\xi}_{i}\rangle\geq 1.2\sigma_{w}\sqrt{d}\sqrt{2\ln(2m)}\right]$ $\displaystyle=1-\Phi\left(\frac{1.2\sigma_{w}\sqrt{d}\sqrt{2\ln(2m)}}{\|w_{s,r}^{(0)}\|}\right)$ $\displaystyle\geq 1-\Phi(\sqrt{3\ln(2m)})$ $\displaystyle\geq\dfrac{1}{2\sqrt{3\ln(2m)}}\exp\left(-\frac{3\ln(2m)}{2}\right)$ $\displaystyle\geq m^{-2}.$ $$

We denote $X_{s,r,i}=\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(0)},\bm{\xi}_{i}\rangle\geq 1.2\sigma_{w}\sqrt{d}\sqrt{2\ln(2m)}\right),T^{\prime}_{s,r,j}=\sum\limits_{i\in I_{j}}X_{s,r,i},$

$Y_{s,r,i}=\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(0)},\bm{x}_{i}\rangle+b_{s,r}^{(0)}\geq 0\right)$ and we know that $|T_{s,r,j}|=\sum\limits_{i\in I_{j}}Y_{s,r,i}.$

If $\langle\bm{w}_{s,r}^{(0)},\bm{\xi}_{i}\rangle\geq 1.2\sigma_{w}\sqrt{d}\sqrt{2\ln(2m)}$, then

$$ $\displaystyle\langle\bm{w}_{s,r}^{(0)},\bm{x}_{i}\rangle+b_{s,r}^{(0)}$ $\displaystyle\geq\operatorname{min}_{s\in\{-1,+1\},r\in J}\langle\bm{w}_{s,r}^{(0)},\bm{\mu}_{p}\rangle+\langle\bm{w}_{s,r}^{(0)},\bm{\xi}_{i}\rangle-2\sigma_{b}\sqrt{2\ln(m)}$ $\displaystyle\geq 0.1\sigma_{w}\sqrt{d}\sqrt{2\ln(2m)}-2\sigma_{b}\sqrt{2\ln(m)}\geq 0.$ $$

That is to say $Y_{s,r,i}=1$ if $X_{s,r,i}=1$. Thus

$$ $\displaystyle\text{Pr}\left[|T_{s,r,j}|\geq\dfrac{n}{2m^{2}k}\right]\geq\text{Pr}\left[T^{\prime}_{s,r,j}\geq\dfrac{n}{2m^{2}k}\right]$ $$

For any given $s\in\{-1,+1\},r\in[m]$ and $i\in I_{j}$, we know that $X_{s,r,i}$ are i.i.d. and $\mathbb{E}[X_{s,r,i}]\geq m^{-2}.$

Then by Chernoff bound, we have

$$ $\displaystyle\text{Pr}\left[T^{\prime}_{s,r,j}\geq\dfrac{|I_{j}|}{2m^{2}}\right]$ $\displaystyle\geq 1-\exp\left(\dfrac{|I_{j}|}{8m^{2}}\right)\geq 1-\exp\left(\dfrac{n}{9km^{2}}\right)$ $$

Then

$$ $\displaystyle\text{Pr}\left[|T_{s,r,j}|\geq\dfrac{n}{3km^{2}}\right]$ $\displaystyle\geq\text{Pr}\left[T^{\prime}_{s,r,j}\geq\dfrac{n}{3km^{2}}\right]$ $\displaystyle\geq\text{Pr}\left[T^{\prime}_{s,r,j}\geq\dfrac{|I_{j}|}{2m^{2}}\right]$ $\displaystyle\geq 1-\exp\left(\dfrac{n}{9km^{2}}\right).$ $$

Then by union bound, we know that $|T_{s,r,j}|\geq\dfrac{n}{3km^{2}}$ holds for all $s\in\{-1,+1\},r\in[m],p\in J$ with probability at least $1-2mk\exp\left(\dfrac{n}{9km^{2}}\right)$.

Combing the above together, with probability at least $1-m^{-0.01}-2mk\exp\left(\dfrac{n}{9km^{2}}\right)$, we have $|T_{s,r,j}|\geq\dfrac{n}{3km^{2}}$.
∎

Next, we show that the pre-activation output of the network is very small, at the initialization.

###### Lemma A.13 .

For any $i\in I,r\in[m],s\in\{-1,+1\}$, we have $|\langle\bm{w}_{s,r}^{(0)},\bm{x}_{i}\rangle+b_{s,r}^{(0)}|\leq\dfrac{\eta\sqrt{d}}{nm}.$

###### Proof of Lemma A.13 .

By Property (2) in Proposition [A.6](#A1.Thmtheorem6) and Lemma [A.8](#A1.Thmtheorem8), we have

$$ $\|\bm{x}_{i}\|\leq 3\sqrt{d},\|\bm{w}_{s,r}^{(0)}\|\leq\sigma_{w}\left(\sqrt{d}+2\ln(d)\right),\|b_{s,r}^{(0)}\|\leq 2\sigma_{b}\sqrt{2\ln(2m)}.$ $$

Therefore, by triangle inequality, we know that

$$ $\displaystyle|\langle\bm{w}_{s,r}^{(0)},\bm{x}_{i}\rangle+b_{s,r}^{(0)}|\leq\|\bm{x}_{i}\|\|\bm{w}_{s,r}^{(0)}\|+\|b_{s,r}^{(0)}\|\leq 3\sqrt{d}\cdot 2\sigma_{w}\sqrt{d}+2\sigma_{b}\sqrt{2\ln(2m)}\leq\dfrac{\eta\sqrt{d}}{nm}$ $$

∎

Finally, we present the following two lemmas about the range of
the loss derivative.

Denote $\ell_{i}^{\prime(t)}:=\nabla_{f_{\bm{\theta}^{(t)}}(x_{i})}\ell(y_{i}f_{\bm{\theta}^{(t)}}(x))=-\dfrac{y_{i}\exp\left(-y_{i}f_{\bm{\theta}^{(t)}}(\bm{x}_{i})\right)}{1+\exp\left(-y_{i}f_{\bm{\theta}^{(t)}}(\bm{x}_{i})\right)}=-\dfrac{y_{i}}{1+\exp\left(y_{i}f_{\bm{\theta}^{(t)}}(\bm{x}_{i})\right)}$.

###### Lemma A.14 .

For each $i\in I$, we have$-\dfrac{2}{3}\leq y_{i}\ell_{i}^{\prime(0)}\leq-\dfrac{1}{3}$.

###### Proof of Lemma A.14 .

By applying Lemma [A.13](#A1.Thmtheorem13), we know

$$ $\displaystyle|f_{\bm{\theta}^{(0)}}(\bm{x}_{i})|\leq\dfrac{1}{m}\left(\sum\limits_{s\in\{-1,+1\}}\sum\limits_{r\in[m]}|\langle\bm{w}_{s,r}^{(0)},\bm{x}_{i}\rangle+b_{s,r}^{(0)}|\right)\leq\dfrac{2\eta\sqrt{d}}{nm}\leq\ln 2$ $$

Then, we have $1/2\leq\exp\left(y_{i}f_{\bm{\theta}^{(0)}}(\bm{x}_{i})\right)\leq 2$, and we can derive that

$$ $\displaystyle-\dfrac{2}{3}\leq y_{i}\ell_{i}^{\prime(0)}=-\dfrac{1}{1+\exp\left(y_{i}f_{\bm{\theta}^{(0)}}(\bm{x}_{i})\right)}\leq-\dfrac{1}{3}.$ $$

∎

###### Lemma A.15 .

For each $i\in I$ and any time step $t$, we have $-1\leq y_{i}\ell_{i}^{\prime(t)}\leq 0$.

###### Proof of Lemma A.15 .

It can be easily checked as follows.

$$ $-1\leq y_{i}\ell_{i}^{\prime(t)}=-\dfrac{1}{1+\exp\left(y_{i}f_{\bm{\theta}^{(t)}}(\bm{x}_{i})\right)}\leq 0.$ $$

∎

###### Lemma A.9 (Concentration of Maximum of Gaussians) .

###### Proof of Lemma A.9 .

###### Proposition A.10 .

###### Proof of Proposition A.10 .

###### Definition A.11 (Activation Region over Data Input) .

###### Proposition A.12 .

###### Proof of Proposition A.12 .

###### Lemma A.13 .

###### Proof of Lemma A.13 .

###### Lemma A.14 .

###### Proof of Lemma A.14 .

###### Lemma A.15 .

###### Proof of Lemma A.15 .

## Appendix B Proof for Section 4 : Feature-Averaging Regime

In this section, we provide the proof of Theorem [4.5](#S4.Thmtheorem5).
We analyze the training dynamics of gradient descent, which constitutes the main part of our proof.

### B.1 Analysis of Training Dynamics

We first assume that all the properties and lemmas mentioned in Appendix [A](#A1) hold with high probability over the sampled training dataset and the network initialization.

Now, we introduce some useful notations.
Denote $S_{s,i}^{(t)}$ as the set of indices of neurons in positive or negative class (determined by $s$) which has been activated by training data point $\bm{x}_{i}$ at time step $t$. Formally, we define it as $S_{s,i}^{(t)}=\{r\in[m]:\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\}$ for $s\in\{-1,+1\}$ and $i\in I$. The following lemma describes the set of activated neurons after the first gradient descent update.

###### Lemma B.1 .

$S_{1,i}^{(1)}=[m]$ for all $i\in I_{+}$ and $S_{-1,i}^{(1)}=[m]$ for all $i\in I_{-}$.

###### Proof of Lemma B.1 .

Without loss of generality, we consider the case when $\bm{x}_{i}$ belongs to the positive class. We show that
$\langle\bm{w}_{1,r}^{(1)},\bm{x}_{i}\rangle+b_{1,r}^{(1)}\geq 0$ for all $r\in[m]$. By applying the gradient descent update and Lemma [A.13](#A1.Thmtheorem13),
we have

$$ $\displaystyle\langle\bm{w}_{1,r}^{(1)},\bm{x}_{i}\rangle+b_{1,r}^{(1)}=$ $\displaystyle\langle\bm{w}_{1,r}^{(0)},\bm{x}_{i}\rangle+b_{1,r}^{(0)}-\eta\left(\langle\nabla_{\bm{w}_{1,r}}\mathcal{L}(\bm{\theta}^{(0)}),\bm{x}_{i}\rangle+\nabla_{b_{1,r}}\mathcal{L}(\bm{\theta}^{(0)})\right)$ $\displaystyle\geq$ $\displaystyle\dfrac{\eta}{nm}\left(-d^{1/2}-mn\left(\left\langle\nabla_{\bm{w}_{1,r}}\mathcal{L}(\bm{\theta}^{(0)}),\bm{x}_{i}\right\rangle+\nabla_{b_{1,r}}\mathcal{L}(\bm{\theta}^{(0)})\right)\right)$ $$

First, we examine the update of linear terms as follows:

$$ $\displaystyle-mn\langle\nabla_{\bm{w}_{1,r}}\mathcal{L}(\bm{\theta}^{(0)}),\bm{x}_{i}\rangle$ $\displaystyle=$ $\displaystyle-\langle\sum\limits_{p\in I_{+}}\ell_{p}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(0)},\bm{x}_{p}\rangle+b_{1,r}^{(0)}\geq 0\rangle\right)\bm{x}_{p}-\sum\limits_{p\in I_{-}}\ell_{p}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(0)},\bm{x}_{i}\rangle+b_{1,r}^{(0)}\geq 0\rangle\right)\bm{x}_{p},\bm{x}_{i}\rangle$ $\displaystyle\geq$ $\displaystyle-\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(0)},\bm{x}_{p}\rangle+b_{1,r}^{(0)}\geq 0\rangle\right)\langle\bm{x}_{p},\bm{x_{i}}\rangle+\sum\limits_{p\notin I_{c(i)}}\ell_{p}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(0)},\bm{x}_{p}\rangle+b_{1,r}^{(0)}\geq 0\rangle\right)|\langle\bm{x}_{p},\bm{x_{i}}\rangle|$ $\displaystyle\geq$ $\displaystyle\ \ \dfrac{1}{3}\sum\limits_{p\in I_{c(i)}}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(0)},\bm{x}_{p}\rangle+b_{1,r}^{(0)}\geq 0\rangle\right)\langle\bm{x}_{p},\bm{x_{i}}\rangle-\dfrac{2}{3}\sum\limits_{p\notin I_{c(i)}}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(0)},\bm{x}_{i}\rangle+b_{1,r}^{(0)}\geq 0\rangle\right)|\langle\bm{x}_{p},\bm{x_{i}}\rangle|$ $\displaystyle\geq$ $\displaystyle\ \ \dfrac{d}{6}\sum\limits_{p\in I_{c(i)}}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(0)},\bm{x}_{p}\rangle+b_{1,r}^{(0)}\geq 0\rangle\right)-\dfrac{2\Delta}{3}\sum\limits_{p\notin I_{c(i)}}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(0)},\bm{x}_{p}\rangle+b_{1,r}^{(0)}\geq 0\rangle\right)\ \ \ \ \ \left(\textit{Recall}\ \ \Delta=4\sqrt{d}\ln(d)\right)$ $\displaystyle=$ $\displaystyle\ \ \dfrac{1}{6}d|T_{1,r,c(i)}|-\dfrac{2\Delta}{3}\sum\limits_{l\in r/\{c(i)\}}|T_{1,r,l}|$ $\displaystyle\geq$ $\displaystyle\ \ \dfrac{nd}{18km^{2}}-\dfrac{2n\Delta}{3}\ \ \ \ \ \ \ (\text{By Proposition \ref{prop:_T_s_r_p}})$ $\displaystyle\geq$ $\displaystyle\ \ n\Delta$ (14) $$

Next, we examine the update of the bias term as follows:

$$ $\displaystyle-mn\nabla_{b_{1,r}}\mathcal{L}(\bm{\theta}^{(0)})=$ $\displaystyle-\sum\limits_{p\in I_{+}}\ell_{p}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(0)},\bm{x}_{p}\rangle+b_{1,r}^{(0)}\geq 0\rangle\right)+\sum\limits_{p\in I_{-}}\ell_{p}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(0)},\bm{x}_{p}\rangle+b_{1,r}^{(0)}\geq 0\rangle\right)$ $\displaystyle\geq$ $\displaystyle-\dfrac{2}{3}\sum\limits_{p\in I}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(0)},\bm{x}_{p}\rangle+b_{1,r}^{(0)}\geq 0\rangle\right)$ $\displaystyle\geq$ $\displaystyle-\dfrac{2n}{3}$ (15) $$

Combining ([14](#A2.E14)) and ([15](#A2.E15)) together, we know

$$ $\displaystyle\langle\bm{w}_{1,r}^{(1)},\bm{x}_{i}\rangle+b_{1,r}^{(1)}\geq$ $\displaystyle\dfrac{\eta}{nm}\left(-d^{1/2}-mn\left(\left\langle\nabla_{\bm{w}_{1,r}}\mathcal{L}(\bm{\theta}^{(0)}),\bm{x}_{i}\right\rangle+\nabla_{b_{1,r}}\mathcal{L}(\bm{\theta}^{(0)})\right)\right)$ $\displaystyle\geq$ $\displaystyle\dfrac{\eta}{nm}(-d^{1/2}+n\Delta-\dfrac{2n}{3})$ $\displaystyle\geq$ $\displaystyle 0$ $$

Since this inequality holds for all $r\in[m]$, we have that $S_{1,i}^{(1)}=[m]$. For the case when $\bm{x}_{i}$ belongs to the negative class, we have $S_{-1,i}^{(1)}=[m]$ using the same argument.
∎

Now, we analyze the dynamics of the coefficients in the training process, where we first give the following weight-decomposition lemma.

###### Lemma B.2 (Weight Decomposition) .

During the training dynamics, there exists the following coefficient sequences $\lambda_{s,r,j}^{(t)}$ and $\sigma_{s,r,i}^{(t)}$ for each $s\in\{-1,+1\},r\in[m],j\in J,i\in I$ such that

$$ $\displaystyle\bm{w}_{s,r}^{(t)}=\bm{w}_{s,r}^{(0)}+\sum_{j\in J}\lambda_{s,r,j}^{(t)}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}+\sum_{i\in I}\sigma_{s,r,i}^{(t)}\bm{\xi}_{i}\|\bm{\xi}_{i}\|^{-2}$ $$

###### Proof of Lemma B.2 .

First, we construct a set of $\{\hat{\lambda}_{s,r,j}^{(t)}\}$ and $\{\hat{\sigma}_{s,r,i}^{(t)}\}$ according to the following recursive formulas:

$$ $\displaystyle\hat{\lambda}_{s,r,j}^{(t+1)}=\hat{\lambda}_{s,r,j}^{(t)}-\dfrac{s\eta}{nm}\cdot\sum\limits_{i\in I_{j}}\ell_{i}^{\prime(t)}\|\bm{\mu_{j}}\|^{2}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right).$ $\displaystyle\hat{\sigma}_{s,r,i}^{(t+1)}=\hat{\sigma}_{s,r,i}^{(t)}-\dfrac{s\eta}{nm}\cdot\ell_{i}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right).$ $\displaystyle\hat{\lambda}_{s,r,j}^{(0)}=0,\hat{\sigma}_{s,r,i}^{(0)}=0.$ $$

Now, we prove
by induction on $t$
that $\{\hat{\lambda}_{s,r,j}^{(t)}\}$ and $\{\hat{\sigma}_{s,r,i}^{(t)}\}$ constructed as above
satisfy that

$$ $\displaystyle\bm{w}_{s,r}^{(t)}=\bm{w}_{s,r}^{(0)}+\sum_{j\in J}\hat{\lambda}_{s,r,j}^{(t)}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}+\sum_{i\in I}\hat{\sigma}_{s,r,i}^{(t)}\bm{\xi}_{i}\|\bm{\xi}_{i}\|^{-2}$ $$

The base case when $t=0$ the conclusion holds trivially.
Assuming the inductive hypothesis holds at time step $t$, we now consider the case at time step $t+1$. By the update equation in gradient descent, we know that

$$ $\displaystyle\bm{w}_{s,r}^{(t+1)}$ $\displaystyle=\bm{w}_{s,r}^{(t)}-\dfrac{s\eta}{nm}\sum\limits_{i\in I}\ell_{i}^{\prime(t)}\bm{x}_{i}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right)$ $\displaystyle=\bm{w}_{s,r}^{(t)}-\dfrac{s\eta}{nm}\sum\limits_{i\in I}\ell_{i}^{\prime(t)}\left(\bm{\mu}_{c(i)}+\bm{\xi}_{i}\right)\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right)$ $\displaystyle=\bm{w}_{s,r}^{(t)}-\dfrac{s\eta}{nm}\left(\sum\limits_{j\in J}\bm{\mu}_{j}\sum\limits_{i\in I_{j}}\ell_{i}^{\prime(t)}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right)+\sum\limits_{i\in I}\bm{\xi}_{i}\ell_{i}^{\prime(t)}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right)\right)$ $\displaystyle=\bm{w}_{s,r}^{(0)}+\sum\limits_{j\in J}\bm{\mu_{j}}\|\bm{\mu_{j}}\|^{-2}\left(\hat{\lambda}_{s,r,j}^{(t)}-\dfrac{s\eta}{nm}\cdot\sum\limits_{i\in I_{j}}\ell_{i}^{\prime(t)}\|\bm{\mu_{j}}\|^{2}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right)\right)$ $\displaystyle+\sum\limits_{i\in I}\bm{\xi}_{i}\|\bm{\xi}_{i}\|^{-2}\left(\hat{\sigma}_{s,r,i}^{(t)}-\dfrac{s\eta}{nm}\cdot\ell_{i}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right)\right)$ $\displaystyle=\bm{w}_{s,r}^{(0)}+\sum_{j\in J}\hat{\lambda}_{s,r,j}^{(t+1)}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}+\sum_{i\in I}\hat{\sigma}_{s,r,i}^{(t+1)}\bm{\xi}_{i}\|\bm{\xi}_{i}\|^{-2}$ $$

This concludes the inductive step and the proof of the lemma.
∎

Naturally, we have the following corollaries.

###### Corollary B.3 .

The coefficients $\lambda_{s,r,j}^{(t)},\sigma_{s,r,i}^{(t)}$ for $s\in\{-1,+1\},r\in[m],j\in J,i\in I$ defined in Corollary ([B.2](#A2.Thmtheorem2)) satisfy the following update equations:

$$ $\displaystyle\lambda_{s,r,j}^{(t+1)}=\lambda_{s,r,j}^{(t)}-\dfrac{s\eta}{nm}\cdot\sum\limits_{i\in I_{j}}\ell_{i}^{\prime(t)}\|\bm{\mu_{j}}\|^{2}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right),$ $\displaystyle\sigma_{s,r,i}^{(t+1)}=\sigma_{s,r,i}^{(t)}-\dfrac{s\eta}{nm}\cdot\ell_{i}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right),$ $\displaystyle\lambda_{s,r,j}^{(0)}=0,\sigma_{s,r,i}^{(0)}=0.$ $$

Indeed, we can only focus on the dynamics of noise coefficients due to the following lemma.

###### Corollary B.4 .

The coefficient sequences $\lambda_{s,r,j}^{(t)}$ and $\sigma_{s,r,i}^{(t)}$ for each pair $s\in\{-1,+1\},r\in[m],j\in J,i\in I$ defined in Lemma [C.6](#A3.Thmtheorem6) satisfy:

$$ $\lambda_{s,r,j}^{(t)}\|\bm{\mu}_{j}\|^{-2}=\sum\limits_{i\in I_{j}}\sigma_{s,r,i}^{(t)}\|\bm{\xi}_{i}\|^{-2}.$ $$

###### Proof of Corollary B.4 .

Using the result in Corollary [B.3](#A5.EGx28), we know that

$$ $\displaystyle\left(\lambda_{s,r,j}^{(t^{\prime}+1)}-\lambda_{s,r,j}^{(t^{\prime})}\right)\|\bm{\mu}_{j}\|^{-2}=$ $\displaystyle-\dfrac{s\eta}{nm}\cdot\sum\limits_{i\in I_{j}}\ell_{i}^{\prime(t^{\prime})}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t^{\prime})},\bm{x}_{i}\rangle+b_{s,r}^{(t^{\prime})}\geq 0\right)$ $\displaystyle=$ $\displaystyle\sum\limits_{i\in I_{j}}\left(\sigma_{s,r,i}^{(t^{\prime}+1)}-\sigma_{s,r,i}^{(t^{\prime})}\right)\|\bm{\xi}_{i}\|^{-2}.$ $$

Then summing up the above equations from $t^{\prime}=0$ to $t^{\prime}=t-1$, we have

$$ $\lambda_{s,r,j}^{(t)}\|\bm{\mu}_{j}\|^{-2}=\sum\limits_{i\in I_{j}}\sigma_{s,r,i}^{(t)}\|\bm{\xi}_{i}\|^{-2}.$ $$

∎

We show that the sign of each feature/noise coefficient remains unchanged during the full training process as in the following lemma.

###### Corollary B.5 .

The coefficient sequences $\lambda_{s,r,j}^{(t)}$ and $\sigma_{s,r,i}^{(t)}$ for each pair $s\in\{-1,+1\},r\in[m],j\in J,i\in I,t\geq 0$ defined in Lemma [B.2](#A2.Thmtheorem2) satisfy:

$$ $\displaystyle\lambda_{s,r,j}^{(t)}\begin{cases}&\geq 0\text{ if }i\in J_{s},\\ &<0\text{ if }i\notin J_{s}.\end{cases}$ $\displaystyle\sigma_{s,r,i}^{(t)}\begin{cases}&\geq 0\text{ if }i\in I_{s},\\ &<0\text{ if }i\notin I_{s}.\end{cases}$ $$

###### Proof of Corollary B.5 .

Using the results in Corollary [B.3](#A5.EGx28) and Lemma [A.15](#A1.Thmtheorem15), we know that

$$ $\displaystyle\operatorname{sgn}\left(\sigma_{s,r,i}^{(t+1)}-\sigma_{s,r,i}^{(t)}\right)=$ $\displaystyle\operatorname{sgn}\left(-\dfrac{s\eta}{nm}\cdot\ell_{i}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right)\right)$ $\displaystyle=$ $\displaystyle\operatorname{sgn}\left(-s\ell_{i}^{\prime(t)}\right)$ $\displaystyle=$ $\displaystyle\operatorname{sgn}\left(sy_{i}\right)$ $$

Noting that $\sigma_{s,r,i}^{(0)}=0$, we know that

$$ $\displaystyle\operatorname{sgn}\left(\sigma_{s,r,i}^{(t)}\right)=\operatorname{sgn}\left(sy_{i}\right)$ $$

That is to say

$$ $\sigma_{s,r,i}^{(t)}\begin{cases}&\geq 0\text{ if }i\in I_{s},\\ &<0\text{ if }i\notin I_{s}.\end{cases}$ $$

Then using the result in Corollary [B.4](#A2.Thmtheorem4), we know

$$ $\lambda_{s,r,j}^{(t)}\begin{cases}&\geq 0\text{ if }j\in J_{s},\\ &<0\text{ if }j\notin J_{s}.\end{cases}$ $$

∎

In the following proof, we need the concept of margin.
We denote the margin of training data point $\bm{x}_{i}$ at time step $t$ as $q_{i}^{(t)}=y_{i}f_{\theta^{(t)}}\left(\bm{x}_{i}\right)$ and the margin gap between training data points $\bm{x}_{i}$ and $\bm{x}_{j}$ at time step $t$ as $\Delta_{q}^{(t)}(i,j)=y_{i}f_{\theta^{(t)}}\left(\bm{x}_{i}\right)-y_{j}f_{\theta^{(t)}}\left(\bm{x}_{j}\right)$.

First, we analyze the relationship between the margin gap and the loss derivatives’ ratio for the two training data points in the following lemma.

###### Lemma B.6 .

For any time step $t$ and two training data points $\bm{x}_{i},\bm{x}_{j}$, if $q_{i}\geq q_{j}$, we have

$$ $e^{\Delta_{q}^{(t)}(i,j)/2}\leq\dfrac{y_{j}\ell^{\prime(t)}_{j}}{y_{i}\ell^{\prime(t)}_{i}}\leq e^{\Delta_{q}^{(t)}(i,j)}.$ $$

###### Proof of Lemma B.6 .

Recall $\ell^{\prime}(x)=-\dfrac{e^{-x}}{1+e^{-x}}=\dfrac{1}{1+e^{x}}$. Then

$$ $\displaystyle\dfrac{\ell^{\prime(t)}_{j}}{\ell^{\prime(t)}_{i}}\cdot e^{-\Delta_{q}^{(t)}(i,j)}=$ $\displaystyle\dfrac{1+e^{q_{i}}}{e^{\Delta_{q}^{(t)}(i,j)}+e^{q_{i}}}\leq 1.$ $$

Since the exponential function is convex, we know that

$$ $\displaystyle\dfrac{\ell^{\prime(t)}_{j}}{\ell^{\prime(t)}_{i}}\cdot e^{-\Delta_{q}^{(t)/2}(i,j)}=$ $\displaystyle\dfrac{1+e^{q_{i}}}{e^{(q_{i}-q_{j})/2}+e^{(q_{i}+q_{j})/2}}\geq 1.$ $$

∎

Next, we establish the relationship between the coefficient $\lambda_{s,r,j}^{(t)}$’s and $\sigma_{s,r,i}^{(t)}$’s in Corollary [B.2](#A2.Thmtheorem2) and margin $q$ we defined before. We denote

$$ $\hat{q}_{i}^{(t)}=\dfrac{1}{m}\left(\sum\limits_{r\in[m]}\lambda_{y_{i},r,c(i)}^{(t)}+\sum\limits_{r\in[m]}\sigma_{y_{i},r,i}^{(t)}\right),$ $$

$$ $\hat{\Delta}^{(t)}_{q}(i,j)=\hat{q}_{i}-\hat{q_{j}}=\dfrac{1}{m}\sum\limits_{r\in[m]}\left(\lambda_{y_{i},r,c(i)}^{(t)}-\lambda_{y_{j},r,c(j)}^{(t)}\right)+\dfrac{1}{m}\sum\limits_{r\in[m]}\left(\sigma_{y_{i},r,i}^{(t)}-\sigma_{y_{j},r,j}^{(t)}\right).$ $$

Later, we will show that $\hat{q}_{i}^{(t)}$ is a good approximation for margin $q_{i}$ and $\hat{\Delta}^{(t)}_{q}(i,j)$ is a good approximation for margin gap $\Delta^{(t)}_{q}(i,j)$ in Lemma [B.14](#A5.EGx47) and Corollary [B.15](#A2.Thmtheorem15).

Next, we arrive at the main part of the proof. Inspired by Kou et al. (), we can prove that the training data’s margin tends to balance automatically.

Denote $\epsilon=\text{max}\left\{\dfrac{2\ln(\frac{n}{k})}{\sqrt{\frac{n}{k}}-\ln(\frac{n}{k})},\dfrac{k^{2}\Delta}{d},\dfrac{k^{2}}{n}\right\}$.
We know that $\epsilon=o(k^{-2.5})$ according to our hyper-parameter [Assumption A.4](#A1.Thmtheorem4).

###### Lemma B.7 .

For $t\leq T_{0}=\operatorname{exp}(\tilde{O}(k^{0.5})),i,j\in I,s\in\{-1,1\}$,
the following statements hold:

- 1.
$\dfrac{k}{2n}\ln(t\eta)\leq\sigma_{y_{i},r,i}^{(t)}\leq\dfrac{2k}{n}\ln(t+1)$,
- 2.
$\hat{\Delta}_{q}^{(t)}(i,j)\leq 5\epsilon$ when $c(i)=c(j),$
- 3.
$\hat{\Delta}_{q}^{(t)}(i,j)\leq 63\epsilon,$
- 4.
$\Delta_{q}^{(t)}(i,j)\leq 65\epsilon,$
- 5.
$y_{i}\ell_{i}^{\prime\left(t\right)}/y_{j}\ell_{j}^{\prime\left(t\right)}\leq 1+14\epsilon$ when $c(i)=c(j)$,
- 6.
$y_{i}\ell_{i}^{\prime\left(t\right)}/y_{j}\ell_{j}^{\prime\left(t\right)}\leq 1+130\epsilon$,
- 7.
$S_{s,i}^{(t)}=[m]$ for $i\in I_{s}$,
- 8.
$|\lambda^{(t)}_{-s,r,c(i)}|\leq\epsilon,|\sigma_{-s,r,i}^{(t)}|\leq 2\epsilon$ for $i\in I_{s}$.

###### Remark B.8 .

Property 1 of Lemma [B.7](#A2.Thmtheorem7) shows that the growth rate of noise coefficient is the logarithm of time, i.e., $\sigma_{s,r,i}^{(t)}\approx\operatorname{log}t$. Property 2 and Property 3 show that the approximate margin gap is small, and the gap between two training data points with the same cluster index is smaller. Property 4 suggests that the exact margin gap is also small. Property 5 and Property 6 provide upper bounds for the loss derivative ratio. Property 7 manifests that the positive training data points can activate all positive neurons and the negative training data points can activate all negative neurons, respectively.

###### Proof of Lemma B.7 .

Without loss of generality, we assume that $i\in I_{+}$.

We use induction to prove this lemma.

Step 1: We first consider the base case when $t=1$ for the induction.

Property 7: Property 7 is exactly the conclusion of Lemma [B.1](#A2.Thmtheorem1).

Indeed, other properties can be easily verified because we adopted a small initialization.

Property 1 and 8: By Lemma [B.3](#A5.EGx28), $\eta\leq d^{-2}$ and noting that $\sigma_{s,r,i}^{(0)}=0$, we have

$$ $\displaystyle|\sigma_{s,r,i}^{(1)}|=|\sigma_{s,r,i}^{(0)}-\dfrac{\eta}{nm}\cdot\ell_{i}^{\prime(0)}\|\bm{\xi_{i}}\|^{2}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(0)},\bm{x}_{i}\rangle+b_{s,r}^{(0)}\geq 0\right)|\leq\dfrac{2\eta d}{nm}\leq\dfrac{2k}{n}\ln 2\leq 2\epsilon.$ $$

and by Corollary [B.5](#A2.Thmtheorem5), for $i\in I_{s}$ we have

$$ $\sigma_{s,r,i}^{(1)}\geq 0\geq\dfrac{k}{2n}\ln(\eta).$ $$

Property 2, 3 and 8: By applying Lemma [B.3](#A5.EGx28), $\eta\leq d^{-2}$ and noting that $\lambda_{s,r,j}^{(0)}=0$, we have

$$ $\displaystyle|\lambda_{s,r,j}^{(1)}|$ $\displaystyle=|\lambda_{s,r,j}^{(0)}-\dfrac{s\eta}{nm}\cdot\sum\limits_{i\in I_{j}}\ell_{i}^{\prime(0)}\|\bm{\mu_{j}}\|^{2}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(0)},\bm{x}_{i}\rangle+b_{s,r}^{(0)}\geq 0\right)|\leq\dfrac{2\eta d}{m}\leq\dfrac{\epsilon}{2m}(\text{Property 8})$ $$

Then, by the definition of $\hat{\Delta}^{(t)}_{q}(i,j)$, we know that

$$ $\displaystyle\hat{\Delta}^{(1)}_{q}(i,j)$ $\displaystyle=\sum\limits_{r\in[m]}\left(\lambda_{1,r,c(i)}^{(1)}-\lambda_{1,r,c(j)}^{(1)}\right)+\sum\limits_{r\in[m]}\left(\sigma_{1,r,i}^{(1)}-\sigma_{1,r,j}^{(1)}\right)$ $\displaystyle\leq 2m\dfrac{\epsilon}{2m}+2m\dfrac{2\eta d}{nm}$ $\displaystyle\leq 5\epsilon\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ (\text{Property 2})$ $\displaystyle\leq 63\epsilon\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ (\text{Property 3})$ $$

Property 4: By applying Lemma [A.15](#A1.Thmtheorem15), we know that

$$ $\displaystyle\left\|\bm{w}_{s,j}^{(1)}\right\|$ $\displaystyle=\left\|\bm{w}_{s,j}^{(0)}-\dfrac{\eta}{nm}\sum\limits_{i\in I}\ell_{i}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right)\bm{x}_{i}\right\|$ $\displaystyle\leq\left\|\bm{w}_{s,j}^{(0)}\right\|+\dfrac{\eta}{nm}\sum\limits_{i\in I}\left\|\bm{x}_{i}\right\|$ $\displaystyle\leq\left\|\bm{w}_{s,j}^{(0)}\right\|+\dfrac{3\sqrt{d}\eta}{m}$ $\displaystyle\leq\dfrac{\epsilon}{4\sqrt{d}}$ $\displaystyle\left\|b_{s,j}^{(1)}\right\|$ $\displaystyle=\left\|b_{s,j}^{(0)}-\dfrac{\eta}{nm}\sum\limits_{i\in I}\ell_{i}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right)\right\|$ $\displaystyle\leq\left\|b_{s,j}^{(0)}\right\|+\dfrac{\eta}{m}$ $\displaystyle\leq\dfrac{\epsilon}{10}$ (16) $$

Then, we have

$$ $\displaystyle|q_{i}^{(1)}|=$ $\displaystyle|f_{\theta^{(1)}}(\bm{x}_{i})|$ $\displaystyle\leq$ $\displaystyle\dfrac{1}{m}\left(\sum\limits_{s\in\{-1,+1\}}\sum\limits_{r\in[m]}|\langle\bm{w}_{s,j}^{(0)},\bm{x}_{i}\rangle+b_{s,j}^{(0)}|\right)$ $\displaystyle\leq$ $\displaystyle\dfrac{1}{m}\left(2m\left(\dfrac{\epsilon}{4}+\dfrac{\epsilon}{10}\right)\right)$ $\displaystyle\leq$ $\displaystyle\epsilon$ $$

By triangle inequality, we have

$$ $\Delta_{q}^{(1)}(i,j)\leq|q_{i}^{(1)}|+|q_{j}^{(1)}|\leq 2\epsilon$ $$

Property 5 and 6: By using the inequality above, Lemma [B.6](#A2.Ex102) and noting that $e^{x}\leq 1+2x$ for small $x$, we have

$$ $\dfrac{\ell^{\prime(1)}_{j}}{\ell^{\prime(1)}_{i}}\leq e^{\Delta_{q}^{(1)}(i,j)}\leq 1+2\Delta_{q}^{(1)}(i,j)\leq\underbrace{1+14\epsilon}_{\textit{Property 5}}\leq\underbrace{1+130\epsilon}_{\textit{Property 6}}.$ $$

Now we complete the proof of the base case when $t=1$ for induction.

Step 2: Assuming that the inductive hypothesis at time step $t$ holds, we consider time step $t+1$. We first give some useful lemmas based on the inductive hypotheses, and then go on to inductive proofs based on these lemmas.

###### Lemma B.9 .

Assuming the inductive hypotheses hold before time step $t$ and $i\in I_{s}$, the update equations in Corollary [B.3](#A5.EGx28) can be simplified as follows:

$$ $\displaystyle\lambda_{s,r,c(i)}^{(t+1)}=\lambda_{s,r,c(i)}^{(t)}-\dfrac{s\eta}{nm}\cdot\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}.$ $\displaystyle\sigma_{s,r,i}^{(t+1)}=\sigma_{s,r,i}^{(t)}-\dfrac{s\eta}{nm}\cdot\ell_{i}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}.$ $$

###### Proof of Lemma B.9 .

By Property 7 in Lemma [B.7](#A2.Thmtheorem7) (inductive hypothesis) and Corollary [B.3](#A5.EGx28), the conclusion is straightforward.
∎

Then, we demonstrate that the bias term remains consistently small as the following lemma.

###### Lemma B.10 .

For every $s\in\{-1,+1\},r\in[m],i\in I$, we have

$$ $|b_{s,r}^{(t)}|\leq\dfrac{\epsilon}{6}.$ $$

###### Proof of Lemma B.10 .

Without loss of generality, we assume that $s=1$. Then, for any $t^{\prime}\leq t$, by using Lemma [B.9](#A5.EGx40), we know that

$$ $\displaystyle\sigma_{1,r,i}^{(t^{\prime}+1)}$ $\displaystyle=\sigma_{1,r,i}^{(t^{\prime})}-\dfrac{\eta}{nm}\cdot\ell_{i}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}.$ $\displaystyle b_{1,r}^{(t^{\prime}+1)}$ $\displaystyle=b_{1,r}^{(t^{\prime})}-\dfrac{\eta}{nm}\sum\limits_{p\in I}\ell_{p}^{\prime(t)}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(t)},\bm{x}_{p}\rangle+b_{1,r}^{(t)}\geq 0\right)$ $\displaystyle\leq b_{1,r}^{(t^{\prime})}-\dfrac{\eta}{nm}\sum\limits_{p\in I^{+}}\ell_{p}^{\prime(t)}$ $$

Thus, we derive

$$ $\sigma_{1,r,i}^{(t^{\prime}+1)}-\sigma_{1,r,i}^{(t^{\prime})}\geq\dfrac{\ell_{i}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}}{\sum\limits_{p\in I^{+}}\ell_{p}^{\prime(t)}}\left(b_{1,r}^{(t^{\prime}+1)}-b_{1,r}^{(t^{\prime})}\right)\geq\dfrac{d}{2n}\left(b_{1,r}^{(t^{\prime}+1)}-b_{1,r}^{(t^{\prime})}\right).$ $$

Summing up the above inequality from $t^{\prime}=1$ to $t^{\prime}=t-1$, we have

$$ $b_{1,r}^{(t)}-b_{1,r}^{(1)}\leq\dfrac{2n}{d}\left(\sigma_{1,r,i}^{(t)}-\sigma_{1,r,i}^{(1)}\right).$ $$

Then by inequality ([16](#A2.E16)) and Property (1) in the inductive hypotheses, we know that

$$ $b_{1,r}^{(t)}\leq\dfrac{2n}{d}\sigma_{1,r,i}^{(t)}-\dfrac{2n}{d}\sigma_{1,r,i}^{(1)}+b_{1,r}^{(1)}\leq\dfrac{4k\ln(t+1)}{d}+\dfrac{\epsilon}{10}\leq\dfrac{\epsilon}{6}.$ $$

Reusing the same argument as in the previous proof, we know that $b_{1,r}^{(t)}\geq-\dfrac{\epsilon}{6}.$
∎

Next, we prove that the true value of the margin $q_{i}^{(t)}$ is close to its estimated value $\hat{q}_{i}^{(t)}$. To estimate the margin, we first estimate the value of each neuron in the following Lemma [B.11](#A2.Thmtheorem11), Lemma [B.12](#A2.Thmtheorem12) and Lemma [B.13](#A2.Thmtheorem13).

###### Lemma B.11 .

Assuming the inductive hypotheses hold before time step $t$, for all $s\in\{-1,+1\},r\in[m],j\in J$, we have

$$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(t)},\bm{\mu}_{j}\rangle-\lambda_{s,r,j}^{(t)}\right|\leq\dfrac{\epsilon}{6}.$ $$

###### Proof of Lemma B.11 .

We bound the gap between the inner product $\langle\bm{w}_{s,r}^{(t)},\bm{\mu}_{j}\rangle$ and feature coefficient $\lambda_{s,r,j}^{(t)}$ as follows:

$$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(t)},\bm{\mu}_{j}\rangle-\lambda_{s,r,j}^{(t)}\right|$ $\displaystyle=$ $\displaystyle\left|\left\langle\bm{w}_{s,r}^{(0)}+\sum_{p\in J}\lambda_{s,r,p}^{(t)}\bm{\mu}_{p}\|\bm{\mu}_{p}\|^{-2}+\sum_{q\in I}\sigma_{s,r,q}^{(t)}\bm{\xi}_{q}\|\bm{\xi}_{q}\|^{-2},\bm{\mu}_{j}\right\rangle-\lambda_{s,r,j}^{(t)}\right|$ $\displaystyle\leq$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(0)},\bm{\mu}_{j}\rangle\right|+\left|\langle\lambda_{s,r,j}^{(t)}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2},\bm{\mu}_{j}\rangle-\lambda_{s,r,j}^{(t)}\right|+\sum\limits_{p\neq j}\lambda_{s,r,p}^{(t)}\|\bm{\mu}_{p}\|^{-2}\left|\langle\bm{\mu_{p}},\bm{\mu}_{j}\rangle\right|+\sum\limits_{q\in I}\sigma_{s,r,q}^{(t)}\|\bm{\xi}_{q}\|^{-2}\left|\langle\bm{\xi}_{q},\bm{\mu}_{j}\rangle\right|$ $\displaystyle\leq$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(0)},\bm{\mu}_{j}\rangle\right|+\sum\limits_{p\neq j}\lambda_{s,r,p}^{(t)}\dfrac{2\Delta}{d}+\sum\limits_{q\in I}\sigma_{s,r,q}^{(t)}\dfrac{2\Delta}{d}$ $\displaystyle=$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(0)},\bm{\mu}_{j}\rangle\right|+\dfrac{2\Delta}{d}\left(\sum\limits_{p\neq j}\sum\limits_{q\in I_{p}}\sigma_{s,r,q}^{(t)}\dfrac{\|\bm{\xi}_{q}\|^{2}}{\|\bm{\mu}_{p}\|^{2}}+\sum\limits_{q\in I}\sigma_{s,r,q}^{(t)}\right)$ $\displaystyle\leq$ $\displaystyle\sqrt{d}\|\bm{w}_{s,r}^{(0)}\|+\dfrac{12\Delta\ln(t+1)k}{d}\leq\dfrac{\epsilon}{6}.$ $$

The first equation employs the weight decomposition in Lemma [B.2](#A2.Thmtheorem2); the second inequality expands the inner product and applies the triangle inequality; the third inequality utilizes the properties from Proposition [A.6](#A1.Thmtheorem6); the fourth equation utilizes Corollary [B.4](#A2.Thmtheorem4); the fifth inequality utilizes Property 1 and Property 8 in the inductive hypotheses.
∎

###### Lemma B.12 .

Assuming the inductive hypotheses hold before time step $t$, for all $s\in\{-1,+1\},r\in[m],i\in I$, we have

$$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(t)},\bm{\xi}_{i}\rangle-\sigma_{s,r,i}^{(t)}\right|\leq\dfrac{\epsilon}{6}.$ $$

###### Proof of Lemma B.12 .

We bound the gap between the inner product $\langle\bm{w}_{s,r}^{(t)},\bm{\xi}_{i}\rangle$ and noise coefficient $\sigma_{s,r,i}^{(t)}$ as follows:

$$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(t)},\bm{\xi}_{i}\rangle-\sigma_{s,r,i}^{(t)}\right|$ $\displaystyle=$ $\displaystyle\left|\left\langle\bm{w}_{s,r}^{(0)}+\sum_{p\in J}\lambda_{s,r,p}^{(t)}\bm{\mu}_{p}\|\bm{\mu}_{p}\|^{-2}+\sum_{q\in I}\sigma_{s,r,q}^{(t)}\bm{\xi}_{q}\|\bm{\xi}_{q}\|^{-2},\bm{\xi}_{i}\right\rangle-\sigma_{s,r,i}^{(t)}\right|$ $\displaystyle\leq$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(0)},\bm{\xi}_{i}\rangle\right|+\left|\langle\sigma_{s,r,i}^{(t)}\bm{\xi}_{i}\|\bm{\xi}_{i}\|^{-2},\bm{\xi}_{i}\rangle-\sigma_{s,r,i}^{(t)}\right|+\sum\limits_{p\in J}\lambda_{s,r,p}^{(t)}\|\bm{\mu}_{p}\|^{-2}\left|\langle\bm{\mu_{p}},\bm{\xi}_{i}\rangle\right|+\sum\limits_{q\neq i}\sigma_{s,r,q}^{(t)}\|\bm{\xi}_{q}\|^{-2}\left|\langle\bm{\xi}_{q},\bm{\xi}_{i}\rangle\right|$ $\displaystyle\leq$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(0)},\bm{\xi}_{i}\rangle\right|+\sum\limits_{p\in J}\lambda_{s,r,p}^{(t)}\dfrac{2\Delta}{d}+\sum\limits_{q\neq i}\sigma_{s,r,q}^{(t)}\dfrac{2\Delta}{d}$ $\displaystyle=$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(0)},\bm{\mu}_{j}\rangle\right|+\dfrac{2\Delta}{d}\left(\sum\limits_{p\neq j}\sum\limits_{q\in I_{p}}\sigma_{s,r,q}^{(t)}\dfrac{\|\bm{\xi}_{q}\|^{2}}{\|\bm{\mu}_{p}\|^{2}}+\sum\limits_{q\in I}\sigma_{s,r,q}^{(t)}\right)$ $\displaystyle\leq$ $\displaystyle\sqrt{d}\|\bm{w}_{s,r}^{(0)}\|+\dfrac{12\Delta\ln(t+1)k}{d}\leq\dfrac{\epsilon}{6}.$ $$

The first equation employs the weight decomposition in Lemma [B.2](#A2.Thmtheorem2); the second inequality expands the inner product and applies the triangle inequality; the third inequality utilizes the properties from Proposition [A.6](#A1.Thmtheorem6); the fourth equation utilizes Corollary [B.4](#A2.Thmtheorem4); the fifth inequality utilizes Property 1 and Property 8 in the inductive hypotheses.
∎

###### Lemma B.13 .

Assuming the inductive hypotheses hold before time step $t$, for all $s\in\{-1,+1\},r\in[m],i\in I$, we have

$$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle-\lambda_{s,r,c(i)}^{(t)}-\sigma_{s,r,i}^{(t)}\right|\leq\dfrac{\epsilon}{3}.$ $$

###### Proof of Lemma B.13 .

Using the conclusion in Lemma [B.11](#A2.Thmtheorem11)
and Lemma [B.12](#A2.Thmtheorem12) and triangle inequality, we can directly obtain the conclusion in this lemma.

$$ $\left|\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle-\lambda_{s,r,c(i)}^{(t)}-\sigma_{s,r,i}^{(t)}\right|\leq\left|\langle\bm{w}_{s,r}^{(t)},\bm{\mu}_{c(i)}\rangle-\lambda_{s,r,c(i)}^{(t)}\right|+\left|\langle\bm{w}_{s,r}^{(t)},\bm{\xi}_{i}\rangle-\sigma_{s,r,i}^{(t)}\right|\leq\dfrac{\epsilon}{3}.$ $$

∎

###### Lemma B.14 .

Assuming the inductive hypotheses hold before time step $t$ , for all $i\in I$, we have

$$ $\displaystyle|q_{i}^{(t)}-\hat{q}_{i}^{(t)}|\leq\epsilon$ $$

###### Proof of Lemma B.14 .

Without loss of generality, we assume that $i\in I_{+}$.

Using Property 7 in the inductive hypotheses, we know that

$$ $\displaystyle\left|q_{i}-\hat{q_{i}}\right|$ $\displaystyle=\dfrac{1}{m}\left|y_{i}f_{\theta^{(t)}}\left(\bm{x}_{i}\right)-\left(\sum\limits_{r\in[m]}\lambda_{1,r,c(i)}^{(t)}+\sum\limits_{r\in[m]}\sigma_{1,r,i}^{(t)}\right)\right|$ $\displaystyle\leq\dfrac{1}{m}\sum\limits_{r\in[m]}\left|\langle\bm{w}_{1,r}^{(t)},\bm{x}_{i}\rangle+b_{1,r}^{(t)}-\lambda_{1,r,c(i)}^{(t)}-\sigma_{1,r,i}^{(t)}\right|+\dfrac{1}{m}\sum\limits_{r\in[m]}\text{ReLU}\left(\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\rangle+b_{-1,r}^{(t)}\right)$ $\displaystyle\leq\underbrace{\dfrac{1}{m}\sum\limits_{r\in[m]}\left|\langle\bm{w}_{1,r}^{(t)},\bm{x}_{i}\rangle-\lambda_{1,r,c(i)}^{(t)}-\sigma_{1,r,i}^{(t)}\right|}_{\mathcal{L}_{1}}+\underbrace{\dfrac{1}{m}\sum\limits_{r\in[m]}\text{ReLU}\left(\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\rangle\right)}_{\mathcal{L}_{2}}+\underbrace{\dfrac{1}{m}\sum\limits_{r\in[m]}\left(|b_{-1,r}^{(t)}|+|b_{1,r}^{(t)}|\right)}_{\mathcal{L}_{3}}.$ $$

For $\mathcal{L}_{1}$ term, using the conclusion in Lemma [B.13](#A2.Thmtheorem13), we know that
$\mathcal{L}_{1}\leq\dfrac{\epsilon}{3}$.

For $\mathcal{L}_{2}$ term, we consider each term in the summation by distinguishing between two scenarios..

Case(I): $\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\rangle\leq 0.$

Then we know that $\text{ReLU}\left(\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\rangle\right)=0<\dfrac{\epsilon}{3}.$

Case(II): $\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\rangle\geq 0.$

Using the conclusion in Lemma [B.13](#A2.Thmtheorem13) and Corollary [B.5](#A2.Thmtheorem5), we know that

$$ $\text{ReLU}\left(\left\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\right\rangle\right)=\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\rangle\leq\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\rangle-\lambda_{-1,r,c(i)}^{(t)}-\sigma_{-1,r,i}^{(t)}\leq\dfrac{\epsilon}{3}.$ $$

Combining Case (I) and (II) together, we know that $\mathcal{L}_{2}\leq\dfrac{\epsilon}{3}.$

For $\mathcal{L}_{3}$ term, using the conclusion in Lemma [B.10](#A2.Thmtheorem10), we know that $\mathcal{L}_{3}\leq\dfrac{\epsilon}{3}$.

Combining the above together, we know that

$$ $|q_{i}^{(t)}-\hat{q}_{i}^{(t)}|\leq\epsilon.$ $$

∎

Then, we can estimate the margin gap between two training data points using a simple triangle inequality.

###### Corollary B.15 .

Assuming the inductive hypotheses hold before time step $t$, for all $i,j\in I$, we have

$$ $\left|\Delta_{q}^{(t)}(i,j)-\hat{\Delta}_{q}^{(t)}(i,j)\right|\leq 2\epsilon.$ $$

###### Proof of Corollary B.15 .

By Lemma [B.14](#A5.EGx47) and triangle inequality, we have

$$ $\left|\Delta_{q}^{(t)}(i,j)-\hat{\Delta}_{q}^{(t)}(i,j)\right|\leq|q_{i}^{(t)}-\hat{q}_{i}^{(t)}|+|q_{j}^{(t)}-\hat{q}_{j}^{(t)}|\leq 2\epsilon.$ $$

∎

Then we will analyze update equations for $\hat{\Delta}_{q}^{(t)}(i,j)$.

By Lemma [B.9](#A5.EGx40), we know that

$$ $\displaystyle\sum\limits_{r\in[m]}\left(\lambda_{y_{i},r,c(i)}^{(t+1)}-\lambda_{y_{j},r,c(j)}^{(t+1)}\right)=\sum\limits_{r\in[m]}\left(\lambda_{y_{i},r,c(i)}^{(t)}-\lambda_{y_{j},r,c(j)}^{(t)}\right)-\dfrac{\eta}{2n}\left(\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}-\sum\limits_{p\in I_{c(j)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(j)}\|^{2}\right)$ $\displaystyle\sum\limits_{r\in[m]}\left(\sigma_{y_{i},r,i}^{(t+1)}-\sigma_{y_{j},r,j}^{(t+1)}\right)=\sum\limits_{r\in[m]}\left(\sigma_{y_{i},r,i}^{(t)}-\sigma_{y_{j},r,j}^{(t)}\right)-\dfrac{\eta}{2n}\left(\ell_{i}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}-\ell_{j}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right)$ $$

Combining the above two equations together, we get the update equation for $\hat{\Delta}_{q}^{(t)}(i,j)$.

$$ $\displaystyle\hat{\Delta}_{q}^{(t+1)}(i,j)=\hat{\Delta}_{q}^{(t)}(i,j)-\dfrac{\eta}{2n}\left(\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}-\sum\limits_{p\in I_{c(j)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(j)}\|^{2}+\ell_{i}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}-\ell_{j}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right)$ (17) $$

###### Lemma B.16 ( Property 8 ) .

For $s\in\{-1,+1\},i\in I_{s}$, we have $|\lambda^{(t)}_{-s,r,c(i)}|\leq\epsilon,|\sigma_{-s,r,i}|\leq 2\epsilon$.

###### Proof of Lemma B.16 .

We prove that for $j\in J^{+},|\lambda^{(t)}_{-1,r,j}|\leq\epsilon$ for all $r\in[m]$ and the proof of the other part is similar.

We distinguish between two scenarios.

Case(I):
For all $i\in I_{j}$, $\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\rangle+b_{-1,r}^{(t)}<0$.

Then by Corollary [B.3](#A5.EGx28) and inductive hypothesis, we know that

$$ $|\lambda_{-1,r,j}^{(t+1)}|=|\lambda_{-1,r,j}^{(t)}|\leq\epsilon.$ $$

Case(II):
There exists $i\in I_{j}$ such that $\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\rangle+b_{-1,r}^{(t)}\geq 0.$

By Lemma [B.13](#A2.Thmtheorem13), we know that

$$ $\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\rangle-\lambda_{-1,r,j}^{(t)}-\sigma_{-1,r,i}^{(t)}\leq\dfrac{\epsilon}{3}.$ $$

Then by Lemma [B.10](#A2.Thmtheorem10) and noting that $\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\rangle+b_{-1,r}^{(t)}\geq 0$ and $\sigma_{-1,r,i}^{(t)}\leq 0$, we have

$$ $\lambda_{-1,r,j}^{(t)}\geq\langle\bm{w}_{-1,r}^{(t)},\bm{x}_{i}\rangle-\sigma_{-1,r,i}^{(t)}-\dfrac{\epsilon}{3}\geq-b_{-1,r}^{(t)}-\dfrac{\epsilon}{3}\geq-\dfrac{2\epsilon}{3}.$ $$

Then using the conclusion in Lemma [B.3](#A5.EGx28), we know that

$$ $\displaystyle|\lambda_{-1,r,j}^{(t+1)}-\lambda_{-1,r,j}^{(t)}|\leq\dfrac{\eta d}{nm}\sum\limits_{i\in I_{j}}|\ell_{i}^{\prime(t)}|\leq\dfrac{\eta d}{m}\leq\dfrac{\epsilon}{3}.$ $$

Thus we have $\lambda_{-1,r,j}^{(t+1)}\geq-\epsilon$. Noting that $\lambda_{-1,r,j}^{(t+1)}\leq 0$, we have $|\lambda_{-1,r,j}^{(t+1)}|\leq\epsilon$.
Then by [Corollary B.4](#A2.Thmtheorem4), we know that $|\sigma_{s,r,i}^{(t)}|\leq 2|\lambda_{s,r,j}^{(t+1)}|\leq 2\epsilon.$
∎

The lemmas we used for the inductive proof have all been proved, and now we can begin the main part of our proof.

Property 2: We first prove that Property 2 as the following lemma.

###### Lemma B.17 (Property 2 of Lemma B.7 ) .

Assuming the inductive hypotheses hold before time step $t$ and $c(i)=c(j)$, we have

$$ $\left|\hat{\Delta}_{q}^{(t+1)}(i,j)\right|\leq 5\epsilon.$ $$

###### Proof of Lemma B.17 .

Without loss of generality, we assume that
$\hat{q}^{(t+1)}_{i}\geq\hat{q}^{(t+1)}_{j}$. We distinguish between two scenarios., one is when $\left|\hat{\Delta}_{q}^{(t)}(i,j)\right|$ is relatively small and the other is when $\left|\hat{\Delta}_{q}^{(t)}(i,j)\right|$ is relatively large.

Case(I): $\hat{\Delta}_{q}^{(t)}(i,j)\leq 4\epsilon.$

By equation ([17](#A2.E17)), Lemma [A.15](#A1.Thmtheorem15) and $\eta\leq d^{-2}$, we know that

$$ $\displaystyle\left|\hat{\Delta}_{q}^{(t+1)}(i,j)-\hat{\Delta}_{q}^{(t)}(i,j)\right|$ $\displaystyle=$ $\displaystyle\dfrac{\eta}{2n}\left|\ell_{i}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}-\ell_{j}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right|\leq\dfrac{\eta}{2n}\left(\left|\ell_{i}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}\right|+\left|\ell_{j}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right|\right)\leq\dfrac{2\eta d}{n}\leq\epsilon$ $$

So we have

$$ $\hat{\Delta}_{q}^{(t+1)}(i,j)\leq 5\epsilon.$ $$

Case(II) : $\hat{\Delta}_{q}^{(t)}(i,j)\geq 4\epsilon.$

By Corollary [B.15](#A2.Thmtheorem15), we know that

$$ $\Delta_{q}^{(t)}(i,j)\geq\hat{\Delta}_{q}^{(t)}(i,j)-2\epsilon\geq 2\epsilon.$ $$

By Lemma [B.6](#A2.Ex102), we know that

$$ $\displaystyle\dfrac{\ell^{\prime(t)}_{j}}{\ell^{\prime(t)}_{i}}\geq e^{\Delta_{q}^{(t)}(i,j)/2}\geq 1+\Delta_{q}^{(t)}(i,j)/2\geq 1+\epsilon.$ (18) $$

By Equation ([17](#A2.E17)) and $c(i)=c(j)$, we know that

$$ $\displaystyle\hat{\Delta}_{q}^{(t+1)}(i,j)-\hat{\Delta}_{q}^{(t)}(i,j)$ $\displaystyle=$ $\displaystyle-\dfrac{\eta}{2n}\left(\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}-\sum\limits_{p\in I_{c(j)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(j)}\|^{2}+\ell_{i}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}-\ell_{j}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right)$ $\displaystyle=$ $\displaystyle-\dfrac{\eta}{2n}\left(\ell_{i}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}-\ell_{j}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right)$ $\displaystyle=$ $\displaystyle-\dfrac{\eta}{2n}\|\bm{\xi}_{j}\|^{2}\ell_{i}^{\prime(t)}\left(\underbrace{(1+\epsilon)-\dfrac{\ell_{j}^{\prime(t)}}{\ell_{i}^{\prime(t)}}}_{<0\textit{, by inequality (\ref{ineq_ratio_lower_bound})}}\right)+\dfrac{\eta}{2n}\ell_{i}^{\prime(t)}\left((1+\epsilon)\|\bm{\xi}_{j}\|^{2}-\|\bm{\xi}_{i}\|^{2}\right)$ $\displaystyle\leq$ $\displaystyle\ \ \ 0$ (19) $$

Furthermore, due to the inductive hypothesis,

$$ $\hat{\Delta}_{q}^{(t+1)}(i,j)\leq\hat{\Delta}_{q}^{(t)}(i,j)\leq 5\epsilon.$ $$

∎

Property 5: Using the result in this lemma and Corollary [B.15](#A2.Thmtheorem15), we know that

$$ $\Delta_{q}^{(t+1)}(i,j)\leq 7\epsilon.$ $$

Using the above inequality and Lemma [B.6](#A2.Ex102) and noting that $e^{x}\leq 1+2x$ for small $x$ we know that

$$ $\dfrac{\ell^{\prime(t+1)}_{j}}{\ell^{\prime(t+1)}_{i}}\leq e^{\Delta_{q}^{(t+1)}(i,j)}\leq 1+2\Delta_{q}^{(t+1)}(i,j)\leq 1+14\epsilon.$ $$

At this point, we have completed the inductive proofs for Property 2 and 5 in Lemma [B.7](#A2.Thmtheorem7).

Property 3, 4 and 6: Next, we consider the general case where the two training data points $\bm{x}_{i},\bm{x}_{j}$ are not necessarily in the same cluster to prove Property 3 and 6 in Lemma [B.7](#A2.Thmtheorem7). This part of the proof overlaps significantly with the previous one, with the main difference being the addition of an extra term in the update equation of $\hat{\Delta}_{q}^{(t)}(i,j)$.

###### Lemma B.18 (Property 3) .

Assuming the inductive hypotheses hold before time step $t$, we have

$$ $\left|\hat{\Delta}_{q}^{(t+1)}(i,j)\right|\leq 63\epsilon.$ $$

###### Proof of Lemma B.18 .

We distinguish between two scenarios, one is when $\left|\hat{\Delta}_{q}^{(t)}(i,j)\right|$ is relative small and the other is when $\left|\hat{\Delta}_{q}^{(t)}(i,j)\right|$ is relative large.

Case(I): $\hat{\Delta}_{q}^{(t)}(i,j)\leq 62\epsilon.$

By equation ([17](#A2.E17)), Lemma [A.15](#A1.Thmtheorem15) and $\eta\leq d^{-2}$, we know that

$$ $\displaystyle\left|\hat{\Delta}_{q}^{(t+1)}(i,j)-\hat{\Delta}_{q}^{(t)}(i,j)\right|$ $\displaystyle=$ $\displaystyle\dfrac{\eta}{2n}\left|\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}-\sum\limits_{p\in I_{c(j)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(j)}\|^{2}+\ell_{i}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}-\ell_{j}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right|$ $\displaystyle\leq$ $\displaystyle\dfrac{\eta}{2n}\left(\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}+\sum\limits_{p\in I_{c(j)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(j)}\|^{2}+\left|\ell_{i}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}\right|+\left|\ell_{j}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right|\right)$ $\displaystyle\leq$ $\displaystyle\dfrac{\eta}{2n}\left(\dfrac{2nd}{k}+\dfrac{2nd}{k}+2d+2d\right)\leq\dfrac{5\eta d}{2k}\leq\epsilon$ $$

So we have

$$ $\hat{\Delta}_{q}^{(t+1)}(i,j)\leq 63\epsilon.$ $$

Case(II): $\hat{\Delta}_{q}^{(t)}(i,j)\geq 62\epsilon.$

By Lemma [B.14](#A5.EGx47), we know that

$$ $\Delta_{q}^{(t)}(i,j)\geq\hat{\Delta}_{q}^{(t)}(i,j)-2\epsilon\geq 60\epsilon.$ $$

By Lemma [B.6](#A2.Ex102), we know that

$$ $\displaystyle\dfrac{\ell^{\prime(t)}_{j}}{\ell^{\prime(t)}_{i}}\geq e^{\Delta_{q}^{(t)}(i,j)/2}\geq 1+\Delta_{q}^{(t)}(i,j)/2\geq 1+30\epsilon.$ (20) $$

Furthermore, due to the inductive hypothesis, for any $p\in I_{c(i)},q\in I_{c(j)}$, we know that

$$ $\displaystyle\dfrac{\ell_{p}^{\prime(t)}}{\ell_{q}^{\prime(t)}}\leq\dfrac{(1+14\epsilon)\ell_{i}^{\prime(t)}}{\ell_{j}^{\prime(t)}/(1+14\epsilon)}\leq\dfrac{(1+14\epsilon)^{2}}{1+30\epsilon}\leq\dfrac{1}{1+\epsilon}.$ (21) $$

By Equation [17](#A2.E17) and $c(i)=c(j)$, we know that

$$ $\displaystyle\hat{\Delta}_{q}^{(t+1)}(i,j)-\hat{\Delta}_{q}^{(t)}(i,j)$ $\displaystyle=$ $\displaystyle-\dfrac{\eta}{2n}\left(\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}-\sum\limits_{p\in I_{c(j)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(j)}\|^{2}+\ell_{i}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}-\ell_{j}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right)$ $\displaystyle=$ $\displaystyle-\dfrac{\eta}{2n}\left(\ell_{i}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}-\ell_{j}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right)-\dfrac{\eta}{2n}\left(\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}-\sum\limits_{p\in I_{c(j)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(j)}\|^{2}\right)$ $$

We analyze each of these two terms separately.

Similar to the proof of inequality ([19](#A2.E19)),

$$ $\displaystyle-\dfrac{\eta}{2n}\left(\ell_{i}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}-\ell_{j}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right)$ $\displaystyle=$ $\displaystyle-\dfrac{\eta}{2n}\|\bm{\xi}_{j}\|^{2}\ell_{i}^{\prime(t)}\left(\underbrace{(1+30\epsilon)-\dfrac{\ell_{j}^{\prime(t)}}{\ell_{i}^{\prime(t)}}}_{<0\textit{, by inequality (\ref{ineq: li/lj different cluster})}}\right)+\dfrac{\eta}{2n}\ell_{i}^{\prime(t)}\left((1+30\epsilon)\|\bm{\xi}_{j}\|^{2}-\|\bm{\xi}_{i}\|^{2}\right)$ $\displaystyle\leq$ $\displaystyle\ \ \ 0$ (22) $$

By inequality ([21](#A2.E21)) and Property 7 in Proposition [A.6](#A1.Thmtheorem6), we have

$$ $\displaystyle\dfrac{\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}}{\sum\limits_{q\in I_{c(j)}}\ell_{q}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}}=\dfrac{\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}}{\sum\limits_{q\in I_{c(j)}}\ell_{q}^{\prime(t)}}\leq\dfrac{|I_{i}|}{(1+\epsilon)|I_{j}|}\leq 1.$ $$

Thus we know that

$$ $\displaystyle-\dfrac{\eta}{2n}\left(\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}-\sum\limits_{p\in I_{c(j)}}\ell_{p}^{\prime(t)}\|\bm{\mu}_{c(j)}\|^{2}\right)\leq 0.$ (23) $$

By combining ([22](#A2.E22)) and ([23](#A2.E23)) with the inductive hypothesis, we know that

$$ $\hat{\Delta}_{q}^{(t+1)}(i,j)\leq\hat{\Delta}_{q}^{(t)}(i,j)\leq 63\epsilon.$ $$

∎

Using the result in this lemma and Corollary [B.15](#A2.Thmtheorem15), we know that

$$ $\Delta_{q}^{(t+1)}(i,j)\leq 65\epsilon.$ $$

Using the above inequality and Lemma [B.6](#A2.Ex102) and noting that $e^{x}\leq 1+2x$ for small $x$ we know that

$$ $\dfrac{y_{j}\ell^{\prime(t+1)}_{j}}{y_{i}\ell^{\prime(t+1)}_{i}}\leq e^{\Delta_{q}^{(t+1)}(i,j)}\leq 1+2\Delta_{q}^{(t+1)}(i,j)\leq 1+130\epsilon.$ $$

Now, we have completed the inductive proofs for Property 2, 3, 4, 5 and 6 in Lemma [B.7](#A2.Thmtheorem7).

Property 1: To prove Property 1, we need to analyze the update equation for $\sigma_{s,r,i}^{(t)}$. We first prove that $\sigma_{s,r,i}^{(t)}$’s are balanced as follows.

###### Lemma B.19 .

For all $i_{1},i_{2}\in I,r_{1},r_{2}\in[m]$, we have

$$ $\left(1-200\epsilon\right)\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t)}-\dfrac{1}{nd}\leq\sigma_{y_{i_{1}},r_{1},i_{1}}^{(t)}\leq\left(1+200\epsilon\right)\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t)}+\dfrac{1}{nd}.$ $$

###### Proof of Lemma B.19 .

We first prove the right-hand side of the inequality and the proof for the left-hand side is similar.
By Lemma [B.9](#A5.EGx40), for any $t^{\prime}<t$, we know that

$$ $\displaystyle\dfrac{\sigma_{y_{i_{1}},r_{1},i_{1}}^{(t^{\prime}+1)}-\sigma_{y_{i_{1}},r_{1},i_{1}}^{(t^{\prime})}}{\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t^{\prime}+1)}-\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t^{\prime})}}=$ $\displaystyle\dfrac{y_{i_{1}}\ell_{i_{1}}^{\prime(t^{\prime})}\|\bm{\xi_{i_{1}}}\|^{2}}{y_{i_{2}}\ell_{i_{2}}^{\prime(t^{\prime})}\|\bm{\xi_{i_{2}}}\|^{2}}$ $\displaystyle\leq$ $\displaystyle\left(1+130\epsilon\right)\left(\dfrac{\sqrt{d}+\ln(d)}{\sqrt{d}-\ln(d)}\right)^{2}(\text{Applying Property 1 in Proposition \ref{prop:A_1}})$ $\displaystyle\leq$ $\displaystyle\ 1+200\epsilon$ $$

That is to say

$$ $\sigma_{y_{i_{1}},r_{1},i_{1}}^{(t^{\prime}+1)}-\sigma_{y_{i_{1}},r_{1},i_{1}}^{(t^{\prime})}\leq\left(1+200\epsilon\right)\left(\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t^{\prime}+1)}-\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t^{\prime})}\right).$ $$

Summing the above inequality from $t^{\prime}=1$ to $t^{\prime}=t-1$, we have

$$ $\sigma_{y_{i_{1}},r_{1},i_{1}}^{(t)}-\sigma_{y_{i_{1}},r_{1},i_{1}}^{(1)}\leq\left(1+200\epsilon\right)\left(\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t)}-\sigma_{y_{i_{2}},r_{2},i_{2}}^{(1)}\right).$ $$

Then, we can derive that

$$ $\displaystyle\sigma_{y_{i_{1}},r_{1},i_{1}}^{(t)}$ $\displaystyle\leq\left(1+200\epsilon\right)\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t)}+\sigma_{y_{i_{1}},r_{1},i_{1}}^{(1)}$ $\displaystyle=\left(1+200\epsilon\right)\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t)}-\dfrac{\eta}{nm}\cdot\ell_{i}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle+b_{s,r}^{(t)}\geq 0\right)$ $\displaystyle\leq\left(1+200\epsilon\right)\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t)}+\dfrac{2d\eta}{nm}$ $\displaystyle\leq\left(1+200\epsilon\right)\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t)}+\dfrac{1}{nd}.$ $$

Reusing the logic of the above proof, we know that

$$ $\left(1-200\epsilon\right)\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t)}-\dfrac{1}{nd}\leq\sigma_{y_{i_{1}},r_{1},i_{1}}^{(t)}.$ $$

∎

Then, we estimate the margin $q_{i}^{(t)}$ only using $\sigma_{s,r,i}^{(t)}$, which is presented as the following lemma.

###### Lemma B.20 .

For every $i\in I,r_{0}\in[m]$, we have
$\dfrac{n}{2k}\sigma_{y_{i},r_{0},i}^{(t)}-2\epsilon\leq q_{i}^{(t)}\leq\dfrac{2n}{k}\sigma_{y_{i},r_{0},i}^{(t)}+2\epsilon.$

###### Proof of Lemma B.20 .

Without loss of generality, we assume that $i\in I_{+}$.

We first prove the right-hand side of the inequality.

$$ $\displaystyle m\hat{q}_{i}^{(t)}$ $\displaystyle=\sum\limits_{r\in[m]}\lambda_{1,r,c(i)}^{(t)}+\sum\limits_{r\in[m]}\sigma_{1,r,i}^{(t)}$ $\displaystyle=\sum\limits_{r\in[m]}\sum\limits_{p\in I_{c(i)}}\dfrac{\|\bm{\xi}_{i}\|^{-2}}{\|\bm{\mu}_{c(i)}\|^{-2}}\sigma_{1,r,p}^{(t)}+\sum\limits_{r\in[m]}\sigma_{1,r,i}^{(t)}$ $\displaystyle\leq\sum\limits_{r\in[m]}\sum\limits_{p\in I_{c(i)}}\dfrac{\|\bm{\xi}_{i}\|^{-2}}{\|\bm{\mu}_{c(i)}\|^{-2}}\left(\left(1+200\epsilon\right)\sigma_{1,r_{0},i}^{(t)}+\dfrac{1}{nd}\right)+\sum\limits_{r\in[m]}\left(\left(1+200\epsilon\right)\sigma_{1,r_{0},i}^{(t)}+\dfrac{1}{nd}\right)$ $\displaystyle\leq\dfrac{2mn}{k}\sigma_{1,r_{0},i}^{(t)}+\dfrac{2m}{kd}$ $\displaystyle\leq\dfrac{2mn}{k}\sigma_{1,r_{0},i}^{(t)}+m\epsilon.$ $$

Then, by Lemma [B.14](#A5.EGx47),
we know that

$$ $q_{i}^{(t)}\leq\hat{q}_{i}^{(t)}+\epsilon\leq\dfrac{2n}{k}\sigma_{1,r_{0},i}^{(t)}+2\epsilon.$ $$

Reusing the argument of the above proof, we prove the left-hand side of the inequality.

$$ $\displaystyle m\hat{q}_{i}^{(t)}$ $\displaystyle=\sum\limits_{r\in[m]}\lambda_{1,r,c(i)}^{(t)}+\sum\limits_{r\in[m]}\sigma_{1,r,i}^{(t)}$ $\displaystyle=\sum\limits_{r\in[m]}\sum\limits_{p\in I_{c(i)}}\dfrac{\|\bm{\xi}_{i}\|^{-2}}{\|\bm{\mu}_{c(i)}\|^{-2}}\sigma_{1,r,p}^{(t)}+\sum\limits_{r\in[m]}\sigma_{1,r,i}^{(t)}$ $\displaystyle\geq\sum\limits_{r\in[m]}\sum\limits_{p\in I_{c(i)}}\dfrac{\|\bm{\xi}_{i}\|^{-2}}{\|\bm{\mu}_{c(i)}\|^{-2}}\left(\left(1-200\epsilon\right)\sigma_{1,r_{0},i}^{(t)}-\dfrac{1}{nd}\right)+\sum\limits_{r\in[m]}\left(\left(1-200\epsilon\right)\sigma_{1,r_{0},i}^{(t)}-\dfrac{1}{nd}\right)$ $\displaystyle\geq\dfrac{mn}{2k}\sigma_{1,r_{0},i}^{(t)}-\dfrac{m}{2kd}$ $\displaystyle\geq\dfrac{mn}{2k}\sigma_{1,r_{0},i}^{(t)}-m\epsilon.$ $$

Then by Lemma [B.14](#A5.EGx47),
we know that

$$ $q_{i}^{(t)}\geq\hat{q}_{i}^{(t)}-\epsilon\geq\dfrac{n}{2k}\sigma_{1,r_{0},i}^{(t)}-2\epsilon.$ $$

∎

Furthermore, we also need to estimate $\ell_{i}^{\prime(t)}$ using $\sigma_{s,r,i}^{(t)}$ as the following lemma.

###### Lemma B.21 .

For every $i\in I,r\in[m]$, we have

$$ $\dfrac{1}{3}\exp\left(-\dfrac{2n}{k}\sigma_{y_{i},r,i}^{(t)}\right)\leq-y_{i}\ell_{i}^{\prime(t)}\leq 2\exp\left(-\dfrac{n}{2k}\sigma_{y_{i},r,i}^{(t)}\right).$ $$

###### Proof of Lemma B.21 .

Without loss of generality, we assume that $i\in I_{+}$.

By Lemma [B.20](#A2.Thmtheorem20), we know that

$$ $\displaystyle-y_{i}\ell_{i}^{\prime(t)}=\dfrac{1}{1+\exp\left(q_{i}^{(t)}\right)}\geq\dfrac{1}{2\exp\left(q_{i}^{(t)}\right)}\geq\dfrac{1}{2}\exp\left(-\dfrac{2n}{k}\sigma_{1,r,i}^{(t)}-2\epsilon\right)\geq\dfrac{1}{3}\exp\left(-\dfrac{2n}{k}\sigma_{1,r,i}^{(t)}\right).$ $$

$$ $\displaystyle-y_{i}\ell_{i}^{\prime(t)}=\dfrac{1}{1+\exp\left(q_{i}^{(t)}\right)}\leq\dfrac{1}{\exp\left(q_{i}^{(t)}\right)}\leq\exp\left(-\dfrac{n}{2k}\sigma_{1,r,i}^{(t)}+2\epsilon\right)\leq 2\exp\left(-\dfrac{2n}{k}\sigma_{1,r,i}^{(t)}\right).$ $$

∎

Then, we can prove Property 1 based on the inductive hypothesis.

Without loss of generality, we assume that $i\in I_{+}$.

We first prove the left-hand side of the inequality.

$$ $\displaystyle\sigma_{1,r,i}^{(t+1)}$ $\displaystyle=\sigma_{1,r,i}^{(t)}-\dfrac{\eta}{nm}\cdot\ell_{i}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}$ $\displaystyle\leq\sigma_{1,r,i}^{(t)}+\dfrac{2\eta d}{nm}\exp\left(-\dfrac{n}{2k}\sigma_{1,r,i}^{(t)}\right)\quad\quad\quad(\text{Applying Lemma \ref{lem: loss sigma representaion}})$ $\displaystyle\leq\dfrac{2k}{n}\ln(t+1)+\dfrac{2\eta d}{nm}\dfrac{1}{t+1}\quad\quad\quad(\text{Monotone with respect to $\sigma_{1,r,i}^{(t)}$})$ $\displaystyle\leq\dfrac{2k}{n}\ln(t+2),$ $\displaystyle\sigma_{-1,r,i}^{(t+1)}$ $\displaystyle\leq 0\leq\dfrac{2k}{n}\ln(t+2)(\text{Corollary \ref{cor: cof sign}}).$ $$

Then we prove the right-hand side of the inequality.

$$ $\displaystyle\sigma_{1,r,i}^{(t+1)}$ $\displaystyle=\sigma_{1,r,i}^{(t)}-\dfrac{\eta}{nm}\cdot\ell_{i}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}$ $\displaystyle\geq\sigma_{1,r,i}^{(t)}+\dfrac{\eta d}{3nm}\exp\left(-\dfrac{2n}{k}\sigma_{1,j,i}^{(t)}\right)\quad\quad\quad(\text{Applying Lemma \ref{lem: loss sigma representaion}})$ $\displaystyle\geq\dfrac{k}{2n}\left(\ln(t)+\ln(\eta)\right)+\dfrac{d}{3nm}\dfrac{1}{t}\quad\quad\quad(\text{Monotonic with respect to $\sigma_{1,r,i}^{(t)}$})$ $\displaystyle\geq\dfrac{k}{2n}\left(\ln\left((t+1)\right)+\ln(\eta)\right)$ $\displaystyle=\dfrac{k}{2n}\ln((t+1)\eta).$ $$

Finally, we prove Property 7. The proof is very similar to the proof of Lemma [B.1](#A2.Thmtheorem1). We show that $\langle\bm{w}_{1,r}^{(t+1)},\bm{x}_{i}\rangle+b_{1,r}^{(t+1)}\geq 0$ for all $i\in I_{+}$. By the inductive hypothesis, we know that $\langle\bm{w}_{1,r}^{(t)},\bm{x}_{i}\rangle+b_{1,r}^{(t)}\geq 0$.

$$ $\displaystyle\langle\bm{w}_{1,r}^{(t+1)},\bm{x}_{i}\rangle+b_{1,r}^{(t+1)}=$ $\displaystyle\langle\bm{w}_{1,r}^{(t)},\bm{x}_{i}\rangle+b_{1,r}^{(t)}-\eta\left(\langle\nabla_{\bm{w}_{1,r}}\mathcal{L}(\bm{\theta}^{(t)}),\bm{x}_{i}\rangle+\nabla_{b_{1,r}}\mathcal{L}(\bm{\theta}^{(t)})\right)$ $\displaystyle\geq$ $\displaystyle-\eta\left(\left\langle\nabla_{\bm{w}_{1,r}}\mathcal{L}(\bm{\theta}^{(t)}),\bm{x}_{i}\right\rangle+\nabla_{b_{1,r}}\mathcal{L}(\bm{\theta}^{(t)})\right).$ $$

Denote $\ell^{\prime(t)}:=\ell^{\prime(t)}_{1}$.
By Property (6) in inductive hypotheses, we know that for all $i\in I$

$$ $\displaystyle\dfrac{1}{1+130\epsilon}\ell^{\prime(t)}_{i}\leq\ell^{\prime(t)}\leq(1+130\epsilon)\ell^{\prime(t)}_{i}$ $$

We examine the update of linear term first.

$$ $\displaystyle\begin{aligned} -\langle\nabla_{\bm{w}_{1,r}}\mathcal{L}(\bm{\theta}^{(t)}),\bm{x}_{i}\rangle=&-\langle\sum\limits_{p\in\mathcal{I_{+}}}\ell_{p}^{\prime(t)}\bm{x}_{p},\bm{x}_{i}\rangle\\ \geq&-\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}\langle\bm{x}_{p},\bm{x_{i}}\rangle+\sum\limits_{p\notin I_{c(i)}}\ell_{p}^{\prime(t)}|\langle\bm{x}_{p},\bm{x_{i}}\rangle|\\ \geq&-\dfrac{d}{2}\sum\limits_{p\in I_{c(i)}}\ell_{p}^{\prime(t)}+\Delta\sum\limits_{p\notin I_{c(i)}}\ell_{p}^{\prime(t)}\\ \geq&-\dfrac{dn\ell^{\prime(t)}}{4k}+2n\Delta\ell^{\prime(t)}\\ \geq&-n\Delta\ell^{\prime(t)}.\end{aligned}$ (24) $$

Then we examine the update of bias term.

$$ $\displaystyle-\nabla_{b_{1,r}}\mathcal{L}(\bm{\theta}^{(t)})=-\sum\limits_{p\in\mathcal{I}}\ell_{p}^{\prime(t)}\mathbbm{1}\left(\langle\bm{w}_{1,r}^{(t)},\bm{x}_{i}\rangle+b_{1,r}^{(t)}\geq 0\right)\geq n\ell^{\prime(t)}.$ (25) $$

Combining ([24](#A2.E24)) and ([25](#A2.E25)) together, we know that

$$ $\displaystyle\langle\bm{w}_{1,r}^{(t+1)},\bm{x}_{i}\rangle+b_{1,r}^{(t+1)}\geq 0$ $$

Thus we know that $S_{1,i}^{(t+1)}=[m]$.

For the case when $\bm{x}_{i}$ belongs to the negative class, we can obtain $S_{-1,i}^{(t+1)}=[m]$ using the same argument.
Now, we have completed the proof of Lemma [B.7](#A2.Thmtheorem7).
∎

###### Lemma B.1 .

###### Proof of Lemma B.1 .

###### Lemma B.2 (Weight Decomposition) .

###### Proof of Lemma B.2 .

###### Corollary B.3 .

###### Corollary B.4 .

###### Proof of Corollary B.4 .

###### Corollary B.5 .

###### Proof of Corollary B.5 .

###### Lemma B.6 .

###### Proof of Lemma B.6 .

###### Lemma B.7 .

###### Remark B.8 .

###### Proof of Lemma B.7 .

###### Lemma B.9 .

###### Proof of Lemma B.9 .

###### Lemma B.10 .

###### Proof of Lemma B.10 .

###### Lemma B.11 .

###### Proof of Lemma B.11 .

###### Lemma B.12 .

###### Proof of Lemma B.12 .

###### Lemma B.13 .

###### Proof of Lemma B.13 .

###### Lemma B.14 .

###### Proof of Lemma B.14 .

###### Corollary B.15 .

###### Proof of Corollary B.15 .

###### Lemma B.16 ( Property 8 ) .

###### Proof of Lemma B.16 .

###### Lemma B.17 (Property 2 of Lemma B.7 ) .

###### Proof of Lemma B.17 .

###### Lemma B.18 (Property 3) .

###### Proof of Lemma B.18 .

###### Lemma B.19 .

###### Proof of Lemma B.19 .

###### Lemma B.20 .

###### Proof of Lemma B.20 .

###### Lemma B.21 .

###### Proof of Lemma B.21 .

### B.2 Proof of Theorem 4.5

Now, we start to prove the main result Theorem [4.5](#S4.Thmtheorem5).

###### Theorem B.22 (Restatement of Theorem 4.5 ) .

In the setting of training a two-layer ReLU network on the binary classification problem $\mathcal{D}(\{{\bm{\mu}}_{j}\}_{j=1}^{k},J_{\pm})$ as described in [Section 3](#S3),
under [Assumptions 3.2](#S3.Thmtheorem2), [3.3](#S3.Thmtheorem3) and [4.3](#S4.Thmtheorem3),
for some $\gamma=o(1)$, after $\Omega(\eta^{-1})\leq T\leq\exp(\tilde{O}(k^{1/2}))$ iterations, with probability at least $1-\gamma$, the neural network satisfies the following properties:

- 1.
The clean accuracy is nearly perfect: $\mathrm{Acc}_{\mathrm{clean}}^{\mathcal{D}}(f_{{\bm{\theta}}^{(T)}})\geq 1-\exp(-\Omega(\log^{2}d))$.
- 2.
Gradient descent leads the network to the feature-averaging regime: there exists a time-variant coefficient $\lambda^{(T)}\in[\Omega(1),+\infty)$ such that for all $s\in\{\pm 1\}$, $r\in[m]$, the weight vector $\bm{w}_{s,r}^{(T)}$ can be approximated as
$\displaystyle\bigg{\|}{\bm{w}}_{s,r}^{(T)}-\lambda^{(T)}\sum_{j\in J_{s}}\|\bm{\mu}_{j}\|^{-2}\bm{\mu}_{j}\bigg{\|}\leq o(d^{-1/2}),$
and the bias term keeps sufficiently small, i.e., $\left\lvert b_{s,r}^{(T)}\right\rvert\leq o(1)$.
- 3.
Consequently, the network is non-robust: for perturbation radius $\delta=\Omega(\sqrt{d/k})$, the $\delta$-robust accuracy is nearly zero, i.e., $\mathrm{Acc}_{\mathrm{robust}}^{\mathcal{D}}(f_{{\bm{\theta}}^{(T)}};\delta)\leq\exp(-\Omega(\log^{2}d))$.

###### Proof of Theorem B.22 .

We first prove that gradient descent leads the network to the feature-averaging regime (Property 2).

###### Lemma B.23 .

For all $s\in\{-1,+1\},j\in J_{s},r\in[m]$, we have $\dfrac{\ln(T\eta)}{4}\leq\lambda_{s,r,j}^{(T)}\leq 4\ln(T+1).$

###### Proof of Lemma B.23 .

Without loss of generality, we assume that $s=1$.

Using Property 1 in Lemma [B.7](#A2.Thmtheorem7) and Corollary [B.4](#A2.Thmtheorem4), we know that

$$ $\displaystyle\lambda_{1,r,j}^{(T)}$ $\displaystyle=\sum\limits_{p\in I_{j}}\dfrac{\|\bm{\xi}_{p}\|^{2}}{\|\bm{\mu}_{j}\|^{2}}\sigma_{1,r,p}^{(T)}\geq\dfrac{n}{2k}\sigma_{1,r,1}^{(T)}\geq\dfrac{\ln(T\eta)}{4}$ $\displaystyle\lambda_{1,r,j}^{(T)}$ $\displaystyle=\sum\limits_{p\in I_{j}}\dfrac{\|\bm{\xi}_{p}\|^{2}}{\|\bm{\mu}_{j}\|^{2}}\sigma_{1,r,p}^{(T)}\leq\dfrac{2n}{k}\sigma_{1,r,1}^{(T)}\leq 4\ln(T+1)$ $$

∎

###### Lemma B.24 .

For $r_{1},r_{2}\in[m],s_{1},s_{2}\in\{-1,+1\},j_{1}\in J_{s_{1}},j_{2}\in J_{s_{2}}$, we have

$$ $\displaystyle\dfrac{\lambda_{s_{1},r_{1},j_{1}}^{(T)}}{\lambda_{s_{2},r_{2},j_{2}}^{(T)}}\leq 1+204\epsilon.$ $$

###### Proof.

By Lemma [B.19](#A2.Thmtheorem19), we know that for any $i_{1},i_{2}\in I$

$$ $\sigma_{y_{i_{1}},r_{1},i_{1}}^{(t)}\leq\left(1+200\epsilon\right)\sigma_{y_{i_{2}},r_{2},i_{2}}^{(t)}+(nd)^{-1}\leq(1+201\epsilon)\sigma_{y_{i_{1}},r_{2},i_{2}}^{(t)}.$ $$

$$ $\displaystyle\dfrac{\lambda_{s_{1},r_{1},j_{1}}^{(T)}}{\lambda_{s_{2},r_{2},j_{2}}^{(T)}}$ $\displaystyle=\dfrac{\sum\limits_{p\in I_{j_{1}}}\|\bm{\xi}_{p}\|^{2}\sigma_{s_{1},r_{1},p}^{(t)}}{\sum\limits_{p\in I_{j_{2}}}\|\bm{\xi}_{p}\|^{2}\sigma_{s_{2},r_{2},p}^{(t)}}$ $\displaystyle\leq(1+201\epsilon)\dfrac{\|\sqrt{d}+\ln(d)\|^{2}}{\|\sqrt{d}-\ln(d)\|^{2}}\dfrac{|I_{j_{1}}|}{|I_{j_{2}}|}$ $\displaystyle\leq(1+201\epsilon)(1+\epsilon)(1+\epsilon)$ $\displaystyle\leq 1+204\epsilon.$ $$

∎

We denote $\lambda^{(T)}=\lambda_{1,1,j_{0}}^{(T)}$ for some $j_{0}\in J_{+}$ as the representative of $\{\lambda_{s,r,j}:s\in\{-1,+1\},r\in[m],j\in J_{s}\}$.

By Lemma [B.23](#A2.Thmtheorem23) and Lemma [B.24](#A2.Thmtheorem24) ,for all $s\in\{-1,+1\},r\in[m],j\in J_{s}$, we have

$$ $\displaystyle|\lambda^{(T)}-\lambda_{s,r,j}^{(T)}|$ $\displaystyle\leq 204\epsilon\lambda^{(T)},$ $\displaystyle\lambda^{(T)}$ $\displaystyle\leq 4\ln(T+1),$ $\displaystyle\lambda^{(T)}$ $\displaystyle\geq\dfrac{\ln(T\eta)}{4}=\Omega(1).$ $$

###### Lemma B.25 .

For all $s\in\{-1,+1\},r\in[m]$, We have

$$ $\displaystyle\sqrt{d}\left\|\bm{w}_{s,r}^{(T)}-\lambda^{(T)}\sum\limits_{j\in J_{s}}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}\right\|=o(1).$ $$

###### Proof of Lemma B.25 .

Recall the weight decomposition in Lemma [B.2](#A2.Thmtheorem2).

$$ $\displaystyle\sqrt{d}\left(\bm{w}^{(T)}_{s,r}-\lambda^{(T)}\sum\limits_{j\in J_{s}}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}\right)$ $\displaystyle=\underbrace{\sqrt{d}\bm{w}_{s,r}^{(0)}}_{\mathcal{L}_{1}}+\underbrace{\sqrt{d}\sum_{j\in J_{s}}\left(\lambda_{s,r,j}^{(T)}-\lambda^{(T)}\right)\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}}_{\mathcal{L}_{2}}$ $\displaystyle+\underbrace{\sqrt{d}\sum_{j\in J_{-s}}\lambda_{s,r,j}^{(T)}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}}_{\mathcal{L}_{3}}+\underbrace{\sqrt{d}\sum_{i\in I}\sigma_{s,r,i}^{(T)}\bm{\xi}_{i}\|\bm{\xi}_{i}\|^{-2}}_{\mathcal{L}_{4}}$ $$

For $\mathcal{L}_{1}$ term, using the conclusion in Lemma [A.10](#A1.Thmtheorem10), we know that

$$ $\displaystyle\|\sqrt{d}\bm{w}_{s,r}^{(0)}\|\leq 2d\sigma_{w}\leq\epsilon=o(1)$ $$

For $\mathcal{L}_{2}$ term, using the conclusion in Lemma [B.23](#A2.Thmtheorem23), Lemma [B.24](#A2.Thmtheorem24) and noting that $\mu_{j}$ are pairwise orthogonal, we know that

$$ $\displaystyle\left\|\sqrt{d}\sum_{j\in J}\left(\lambda_{s,r,j}^{(T)}-\lambda^{(T)}\right)\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}\right\|$ $\displaystyle=\sqrt{\sum\limits_{j\in J_{s}}\left(\lambda_{s,r,j}^{(T)}-\lambda^{(T)}\right)^{2}}$ $\displaystyle\leq\sqrt{\sum\limits_{j\in J_{s}}(204\epsilon)^{2}\left(\lambda^{(T)}\right)^{2}}$ $\displaystyle\leq 204\epsilon\sqrt{k}\lambda^{(T)}$ $\displaystyle\leq 816\epsilon\sqrt{k}\ln(T+1)$ $\displaystyle\leq 900k\epsilon=o(1).$ $$

For $\mathcal{L}_{3}$ term, by [Lemma B.16](#A2.Thmtheorem16) and triangle inequality, we know that

$$ $\displaystyle\left\|\sqrt{d}\sum_{j\in J_{-s}}\lambda_{s,r,j}^{(T)}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}\right\|\leq k\epsilon=o(1).$ $$

For $\mathcal{L}_{4}$ term, by Property (1) in Lemma [B.7](#A2.Thmtheorem7), we have

$$ $\displaystyle\left\|\sqrt{d}\sum_{i\in I}\sigma_{s,r,i}^{(T)}\bm{\xi}_{i}\|\bm{\xi}_{i}\|^{-2}\right\|^{2}$ $\displaystyle=d\sum\limits_{i\in I}\left(\sigma_{s,r,i}^{(T)}\right)^{2}\|\bm{\xi}_{i}\|^{-2}+d\sum\limits_{i_{1}\neq i_{2}}\sigma_{s,r,i_{1}}^{(T)}\sigma_{s,r,i_{2}}^{(T)}\langle\bm{\xi}_{i_{1}},\bm{\xi}_{i_{2}}\rangle\|\bm{\xi}_{i}\|^{-4}$ $\displaystyle\leq 2\sum\limits_{i\in I}\left(\sigma_{s,r,i_{1}}^{(T)}\right)^{2}+\dfrac{2\Delta}{d}\sum\limits_{i_{1}\neq i_{2}}\sigma_{s,r,i_{1}}^{(T)}\sigma_{s,r,i_{2}}^{(T)}$ $\displaystyle\leq\dfrac{8k^{2}\ln^{2}(T+1)}{n}+\dfrac{8k^{2}\ln^{2}(T+1)\Delta}{d}$ $\displaystyle\leq\dfrac{8k^{3}}{n}+\dfrac{8k^{3}\Delta}{d}\leq 16k\epsilon=o(1).$ $$

Combining the above together, we know that

$$ $\displaystyle\sqrt{d}\left\|\bm{w}_{s,r}^{(T)}-\lambda^{(T)}\sum\limits_{j\in J_{s}}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}\right\|=o(1).$ $$

∎

By [Lemma B.10](#A2.Thmtheorem10), we know that $|b_{s,r}^{(T)}|\leq\epsilon.$

Then, we prove that the clean accuracy is nearly perfect (Property 1). We first need to prove the following lemma, which shows that the correlation between network weight and random noise is small.

###### Lemma B.26 .

Let $\bm{\xi}\sim\mathcal{N}(0,I_{d})$. Then, with probability at least $1-2nd^{-\ln(d)/2}$, for all $s\in\{-1,+1\},r\in[m]$ we have

$$ $\displaystyle|\langle\bm{w}_{s,r}^{(t)},\bm{\xi}\rangle|\leq\dfrac{\epsilon}{6}.$ $$

###### Proof of Lemma B.26 .

Reusing the argument of proof of Property (3) and (4) in Proposition [A.6](#A1.Thmtheorem6). We know that with probability at least $1-2nd^{-\ln(d)/2}$, for all $i\in I,j\in J$,

$$ $|\langle\bm{\mu}_{j},\bm{\xi}\rangle|\leq\Delta,|\langle\bm{\xi}_{i},\bm{\xi}\rangle|\leq\Delta.$ $$

The remaining part of the proof is similar to the proof of Lemma [B.11](#A2.Thmtheorem11) and Lemma [B.12](#A2.Thmtheorem12).

$$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(t)},\bm{\xi}\rangle\right|=$ $\displaystyle\left|\left\langle\bm{w}_{s,r}^{(0)}+\sum_{p\in J}\lambda_{s,r,p}^{(t)}\bm{\mu}_{p}\|\bm{\mu}_{p}\|^{-2}+\sum_{q\in I}\sigma_{s,r,q}^{(t)}\bm{\xi}_{q}\|\bm{\xi}_{q}\|^{-2},\bm{\xi}\right\rangle\right|$ $\displaystyle\leq$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(0)},\bm{\mu}_{j}\rangle\right|+\sum\limits_{p\in J}\lambda_{s,r,p}^{(t)}\|\bm{\mu}_{p}\|^{-2}\left|\langle\bm{\mu_{p}},\bm{\mu}_{j}\rangle\right|+\sum\limits_{q\in I}\sigma_{s,r,q}^{(t)}\|\bm{\xi}_{q}\|^{-2}\left|\langle\bm{\xi}_{q},\bm{\mu}_{j}\rangle\right|$ $\displaystyle\leq$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(0)},\bm{\mu}_{j}\rangle\right|+\dfrac{2\Delta}{d}\left(\sum\limits_{p\in J}\lambda_{s,r,p}^{(t)}+\sum\limits_{q\in I}\sigma_{s,r,q}^{(t)}\right)\leq\dfrac{\epsilon}{6}.$ $$

∎

Assume $(\bm{x},y)$ is randomly sampled from the data distribution $\mathcal{D}$. Without loss of generality, we assume that $\bm{x}=\bm{\mu}_{j}+\bm{\xi},y=1$. Using the conclusion in Lemma [B.11](#A2.Thmtheorem11), Lemma [B.26](#A2.Thmtheorem26) and Lemma [B.10](#A2.Thmtheorem10), we know that

$$ $\displaystyle\langle\bm{w}_{1,r}^{(T)},\bm{x}\rangle+b_{1,r}^{(T)}$ $\displaystyle=\langle\bm{w}_{1,r}^{(T)},\bm{\mu}_{j}\rangle+\langle\bm{w}_{1,r}^{(T)},\bm{\xi}\rangle+b_{1,r}^{(T)}$ $\displaystyle\geq\lambda_{1,r,j}^{(T)}-\dfrac{\epsilon}{6}-\dfrac{\epsilon}{6}-\dfrac{\epsilon}{3}$ $\displaystyle\geq\lambda_{1,r,j}^{(T)}-\epsilon.$ $$

$$ $\displaystyle\langle\bm{w}_{-1,r}^{(T)},\bm{x}\rangle+b_{1,r}^{(T)}$ $\displaystyle=\langle\bm{w}_{-1,r}^{(T)},\bm{\mu}_{j}\rangle+\langle\bm{w}_{-1,r}^{(T)},\bm{\xi}\rangle+b_{-1,r}^{(T)}$ $\displaystyle\leq\lambda_{-1,r,j}^{(T)}+\dfrac{\epsilon}{6}+\dfrac{\epsilon}{6}+\dfrac{\epsilon}{3}$ $\displaystyle\leq\lambda_{-1,r,j}^{(T)}+\epsilon$ $\displaystyle\leq\epsilon.$ $$

Then, we have

$$ $\displaystyle f_{\bm{\theta}^{(T)}}(x)$ $\displaystyle=\dfrac{1}{m}\sum\limits_{r\in[m]}\text{ReLU}\left(\langle\bm{w}_{1,r}^{(T)},\bm{x}\rangle+b_{1,r}^{(T)}\right)-\dfrac{1}{m}\sum\limits_{r\in[m]}\text{ReLU}\left(\langle\bm{w}_{-1,r}^{(T)},\bm{x}\rangle+b_{-1,r}^{(T)}\right)$ $\displaystyle\geq\dfrac{1}{m}\sum\limits_{r\in[m]}\left(\lambda_{1,r,j}^{(T)}-\epsilon\right)-\dfrac{1}{m}\sum\limits_{r\in[m]}\epsilon$ $\displaystyle=\dfrac{1}{m}\sum\limits_{r\in[m]}\lambda_{1,r,j}^{(T)}-2\epsilon\geq 0.$ $$

Thus $f_{\bm{\theta}^{(T)}}$ has perfect standard accuracy.

Finally, we prove that the network is non-robust (Property 3).

We consider the following perturbation

$$ $\bm{\rho}=-\dfrac{2(1+c)}{k}\left(\sum\limits_{j\in J_{+}}\bm{\mu}_{j}-\sum\limits_{j\in J_{-}}\bm{\mu}_{j}\right),$ $$

where $c$ is a constant such that $c^{-1}\leq|J_{+}|/|J_{-}|\leq c$. This is to say $|J_{+}|,|J_{-}|\geq\dfrac{k}{1+c}$.

Then, we have

$$ $\displaystyle\langle\bm{w}_{1,r}^{(T)},\bm{x}+\bm{\rho}\rangle+b_{1,r}^{(T)}$ $\displaystyle\leq$ $\displaystyle\langle\bm{w}_{1,r}^{(T)},\bm{\mu}_{j_{0}}\rangle+\langle\bm{w}_{1,r}^{(T)},\bm{\xi}\rangle-\dfrac{2(1+c)}{k}\sum\limits_{j\in J_{+}}\langle\bm{w}_{1,r}^{(T)},\bm{\mu}_{j}\rangle+\dfrac{2(1+c)}{k}\sum\limits_{j\in J_{-}}\langle\bm{w}_{1,r}^{(T)},\bm{\mu}_{j}\rangle+\dfrac{\epsilon}{3}$ $\displaystyle\leq$ $\displaystyle\lambda_{1,r,j_{0}}^{(T)}+\dfrac{\epsilon}{6}+\dfrac{\epsilon}{6}-\dfrac{2(1+c)}{k}\sum\limits_{j\in J_{+}}\left(\lambda_{1,r,j}^{(T)}-\dfrac{\epsilon}{6}\right)+\dfrac{2(1+c)}{k}\sum\limits_{j\in J_{-}}\left(\lambda_{1,r,j}^{(T)}+\dfrac{\epsilon}{6}\right)+\dfrac{\epsilon}{3}$ $\displaystyle\leq$ $\displaystyle\lambda_{1,r,j_{0}}^{(T)}-\dfrac{2(1+c)}{k}\sum\limits_{j\in J_{+}}\lambda_{1,r,j}^{(T)}+\dfrac{(3+c)\epsilon}{3}$ $\displaystyle\leq$ $\displaystyle\lambda_{1,r,j_{0}}^{(T)}-\dfrac{2(1+c)}{k}\sum\limits_{j\in J_{+}}\dfrac{3}{4}\lambda_{1,r,j_{0}}^{(T)}+\dfrac{(3+c)\epsilon}{3}$ $\displaystyle\leq$ $\displaystyle\lambda_{1,r,j_{0}}^{(T)}\left(1-\dfrac{3(1+c)|J_{+}|}{2k}\right)+\dfrac{(3+c)\epsilon}{3}$ $\displaystyle\leq$ $\displaystyle-\dfrac{1}{2}\lambda_{1,r,j_{0}}^{(T)}+\dfrac{(3+c)\epsilon}{3}<0.$ $$

The first equation expands $\bm{x}$ and $\bm{\rho}$ and uses the conclusion in Lemma [B.10](#A2.Thmtheorem10); the second inequality uses the conclusion in Lemma [B.11](#A2.Thmtheorem11) and Lemma [B.26](#A2.Thmtheorem26); the third inequality rearranges the terms and uses the conclusion in Corollary [B.5](#A2.Thmtheorem5); the fourth inequality uses conclusion in Theorem [4.5](#S4.Thmtheorem5).

Reusing the logic of the above inequality, we have

$$ $\displaystyle\langle\bm{w}_{-1,r}^{(T)},\bm{x}+\bm{\rho}\rangle+b_{-1,r}^{(T)}$ $\displaystyle\geq$ $\displaystyle\langle\bm{w}_{-1,r}^{(T)},\bm{\mu}_{j_{0}}\rangle+\langle\bm{w}_{-1,r}^{(T)},\bm{\xi}\rangle-\dfrac{2(1+c)}{k}\sum\limits_{j\in J_{+}}\langle\bm{w}_{-1,r}^{(T)},\bm{\mu}_{j}\rangle+\dfrac{2(1+c)}{k}\sum\limits_{j\in J_{-}}\langle\bm{w}_{-1,r}^{(T)},\bm{\mu}_{j}\rangle-\dfrac{\epsilon}{3}$ $\displaystyle\geq$ $\displaystyle\lambda_{-1,r,j_{0}}^{(T)}-\dfrac{\epsilon}{6}-\dfrac{\epsilon}{6}-\dfrac{2(1+c)}{k}\sum\limits_{j\in J_{+}}\left(\lambda_{-1,r,j}^{(T)}+\dfrac{\epsilon}{6}\right)+\dfrac{2(1+c)}{k}\sum\limits_{j\in J_{-}}\left(\lambda_{-1,r,j}^{(T)}-\dfrac{\epsilon}{6}\right)-\dfrac{\epsilon}{3}$ $\displaystyle\geq$ $\displaystyle\lambda_{-1,r,j_{0}}^{(T)}+\dfrac{2(1+c)}{k}\sum\limits_{j\in J_{-}}\lambda_{-1,r,j}^{(T)}-\dfrac{(3+c)\epsilon}{3}$ $\displaystyle\geq$ $\displaystyle\dfrac{2(1+c)}{k}\sum\limits_{j\in J_{-}}\lambda_{-1,r,j}^{(T)}-\dfrac{(6+c)\epsilon}{3}$ $\displaystyle\geq$ $\displaystyle\dfrac{2(1+c)|J_{-}|}{k}-\dfrac{(6+c)\epsilon}{3}\geq 0.$ $$

By combining the two inequalities above, we can obtain that

$$ $f_{\bm{\theta}^{(T)}}(\bm{x}+\bm{\rho})=\dfrac{1}{m}\sum\limits_{r\in[m]}\operatorname{ReLU}\left(\langle\bm{w}_{1,r}^{(T)},\bm{x}+\bm{\rho}\rangle+b_{1,r}^{(T)}\right)-\dfrac{1}{m}\sum\limits_{r\in[m]}\operatorname{ReLU}\left(\langle\bm{w}_{-1,r}^{(T)},\bm{x}+\bm{\rho}\rangle+b_{-1,r}^{(T)}\right)<0.$ $$

This is to say $\operatorname{sgn}(f_{\bm{\theta}^{(T)}}(\bm{x}+\bm{\rho}))\neq\operatorname{sgn}(f_{\bm{\theta}^{(T)}}(\bm{x}))$, which means $\mathrm{Acc}_{\mathrm{robust}}^{\mathcal{D}}(f_{{\bm{\theta}}^{(T)}};2(1+c)\sqrt{d/k})=o(1)$.
∎

###### Theorem B.22 (Restatement of Theorem 4.5 ) .

###### Proof of Theorem B.22 .

###### Lemma B.23 .

###### Proof of Lemma B.23 .

###### Lemma B.24 .

###### Proof.

###### Lemma B.25 .

###### Proof of Lemma B.25 .

###### Lemma B.26 .

###### Proof of Lemma B.26 .

### B.3 Proof of Theorem 4.6

###### Theorem B.27 (Restatement of Theorem 4.6 ) .

In the setting of [Theorem 4.5](#S4.Thmtheorem5),

$$ $\inf_{C>0}\sup_{{\bm{x}}\in\mathbb{R}^{d}:\|{\bm{x}}\|_{2}=\sqrt{d}}\left\lvert Cf_{\mathrm{FA}}({\bm{x}})-f_{{\bm{\theta}}^{(T)}}({\bm{x}})\right\rvert=o(1),$ $$

where $f_{\mathrm{FA}}({\bm{x}})$ is the feature-averaging network ([Definition 4.1](#S4.Thmtheorem1)).

###### Proof of Theorem B.27 .

By [Lemma B.25](#A2.Thmtheorem25), we have

$$ $\displaystyle\left|\text{ReLU}\left(\langle\bm{w}_{s,r}^{(T)},\bm{x}\rangle\right)-\text{ReLU}\left(\left\langle\lambda^{(T)}\sum\limits_{j\in J_{s}}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2},\bm{x}\right\rangle\right)\right|$ $\displaystyle\leq$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(T)},\bm{x}\rangle-\langle\lambda^{(T)}\sum\limits_{j\in J_{s}}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2},\bm{x}\rangle\right|$ $\displaystyle\leq$ $\displaystyle\|\bm{x}\|\left\|\bm{w}_{s,r}^{(T)}-\lambda^{(T)}\sum\limits_{j\in J_{s}}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}\right\|=o(1).$ $$

Thus we have

$$ $\displaystyle\left|\dfrac{1}{m}\sum\limits_{r\in[m]}\text{ReLU}\left(\langle\bm{w}_{1,r}^{(T)},\bm{x}\rangle+b_{1,r}^{(T)}\right)-\text{ReLU}\left(\left\langle\lambda^{(T)}\sum\limits_{j\in J_{+}}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2},\bm{x}\right\rangle\right)\right|$ $\displaystyle=$ $\displaystyle\left|\dfrac{1}{m}\sum\limits_{r\in[m]}\text{ReLU}\left(\langle\bm{w}_{1,r}^{(T)},\bm{x}\rangle\right)-\dfrac{1}{m}\sum\limits_{r\in[m]}\text{ReLU}\left(\left\langle\lambda^{(T)}\sum\limits_{j\in J_{+}}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2},\bm{x}\right\rangle\right)\right|+\dfrac{1}{m}\sum\limits_{r\in[m]}|b_{1,r}^{(T)}|$ $\displaystyle\leq$ $\displaystyle\dfrac{1}{m}\sum\limits_{r\in[m]}\left|\langle\bm{w}_{1,r}^{(T)},\bm{x}\rangle-\langle\lambda^{(T)}\sum\limits_{j\in J_{+}}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2},\bm{x}\rangle\right|+\epsilon$ $\displaystyle\leq$ $\displaystyle\dfrac{1}{m}\sum\limits_{r\in[m]}\|\bm{x}\|\left\|\bm{w}_{1,r}^{(T)}-\lambda^{(T)}\sum\limits_{j\in J_{+}}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}\right\|+\epsilon=o(1).$ $$

Similarly, we have

$$ $\displaystyle\left|\dfrac{1}{m}\sum\limits_{r\in[m]}\text{ReLU}\left(\langle\bm{w}_{-1,r}^{(T)},\bm{x}\rangle+b_{-1,r}^{(T)}\right)-\text{ReLU}\left(\left\langle\lambda^{(T)}\sum\limits_{j\in J_{-}}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2},\bm{x}\right\rangle\right)\right|=o(1)$ $$

Combining these two inequalities together, we have

$$ $\displaystyle\sup_{{\bm{x}}\in\mathbb{R}^{d}:\|{\bm{x}}\|_{2}=\sqrt{d}}\left\lvert\dfrac{\lambda^{(T)}}{d}f_{\mathrm{FA}}({\bm{x}})-f_{{\bm{\theta}}^{(T)}}({\bm{x}})\right\rvert=o(1).$ $$

∎

###### Theorem B.27 (Restatement of Theorem 4.6 ) .

###### Proof of Theorem B.27 .

## Appendix C Proof for Section 4 : Feature-Decoupling Regime

First, we recall the fine-Grained supervision, multi-Class network classifier and training algorithm.

#### Fine-Grained Supervision.

Following the setting in [Section 3](#S3), we consider the binary classification task with data distribution $\mathcal{D}(\{{\bm{\mu}}_{j}\}_{j=1}^{k},J_{\pm})$. But instead of training the model directly to predict the binary labels, we assume that we are able to label each data point with the cluster $\hat{y}\in[k]$ it belongs to, and then we train a $k$-class classifier to predict the cluster labels. More specifically, we first sample a training set ${\mathcal{S}}:=\{({\bm{x}}_{i},y_{i})\}_{i=1}^{n}\subseteq\mathbb{R}^{d}\times\{\pm 1\}$ from $\mathcal{D}$, along with the cluster labels $\{\tilde{y}_{i}\}_{i=1}^{n}$ for all data points. Then a $k$-class neural network classifier is trained on $\tilde{{\mathcal{S}}}:=\{({\bm{x}}_{i},\tilde{y}_{i})\}_{i=1}^{n}\subseteq\mathbb{R}^{d}\times[k]$.

Multi-Class Network Classifier. We train the following two-layer neural network for the $k$-class classification mentioned above: $\bm{F}_{\bm{\theta}}({\bm{x}}):=(f_{1}({\bm{x}}),f_{2}({\bm{x}}),\dots,f_{k}({\bm{x}}))\in\mathbb{R}^{k}$, where $f_{j}({\bm{x}}):=\frac{1}{m}\sum_{r=1}^{h}\operatorname{ReLU}(\langle\bm{w}_{j,r},\bm{x}\rangle)$, and ${\bm{\theta}}:=({\bm{w}}_{1,1},{\bm{w}}_{1,2},\dots,{\bm{w}}_{k,h})\in\mathbb{R}^{khd}$ are trainable weights, and $h=\Theta(1)$ is the width of each sub-network. The outputs $\bm{F}_{\bm{\theta}}({\bm{x}})$ are then converted to probabilities using the softmax function, namely $p_{j}({\bm{x}}):=\frac{\exp(f_{j}({\bm{x}}))}{\sum_{i=1}^{k}\exp(f_{i}({\bm{x}}))}$ for $j\in[k]$. For predicting the binary label for the original binary classification task on $\mathcal{D}$, we take the difference of the probabilities of the positive and negative classes, i.e., $F^{\mathrm{binary}}_{{\bm{\theta}}}({\bm{x}}):=\sum_{j\in J_{+}}p_{j}({\bm{x}})-\sum_{j\in J_{-}}p_{j}({\bm{x}})$. The clean accuracy $\mathrm{Acc}_{\mathrm{clean}}^{\mathcal{D}}(F^{\mathrm{binary}}_{{\bm{\theta}}})$ and $\delta$-robust accuracy $\mathrm{Acc}_{\mathrm{robust}}^{\mathcal{D}}(F^{\mathrm{binary}}_{{\bm{\theta}}};\delta)$ are then defined similarly as before.

Training Objective and Gradient Descent. We train the multi-class network ${\bm{F}}_{\bm{\theta}}({\bm{x}})$ to minimize the cross-entropy loss ${\mathcal{L}}_{\mathrm{CE}}({\bm{\theta}}):=-\frac{1}{n}\sum_{i=1}^{n}\log p_{\tilde{y}_{i}}({\bm{x}}_{i})$. Similar to [Section 3](#S3), we use gradient descent to minimize the loss function ${\mathcal{L}}_{\mathrm{CE}}({\bm{\theta}})$ with learning rate $\eta$, i.e., $\bm{\theta}^{(t+1)}=\bm{\theta}^{(t)}-\eta\nabla_{\bm{\theta}}\mathcal{L}_{\textit{CE}}(\bm{F}_{\bm{\theta}^{(t)}})$. At initialization, we set ${\bm{w}}_{j,r}^{(0)}\sim{\mathcal{N}}(0,\sigma_{\mathrm{w}}^{2}{\bm{I}}_{d})$ for some $\sigma_{\mathrm{w}}>0$.

Denote $\ell_{i,j}^{\prime(t)}:=\nabla_{f_{j}(\bm{x}_{i})}\mathcal{L}_{\textit{CE}}(\bm{F}^{(t)})=-\mathbbm{1}(\bm{x}_{i}\in I_{j})+\dfrac{\exp(f_{j}^{(t)}(\bm{x}_{i}))}{\sum\limits_{p\in J}\exp(f_{p}^{(t)}(\bm{x}_{i}))}.$

Since many of the proofs in this section are very similar to those in Appendix [B](#A2), we reuse the logic of the proofs and present the key steps.

We also assume that the properties of the training dataset(Proposition [A.6](#A1.Thmtheorem6)) in Appendix [A](#A1) hold.

### C.1 Propositions of Network Initialization

###### Proposition C.1 .

With probability at least $1-4hmd^{-\ln(d)}-2h^{-3}m^{-3}$ , we have the following properties for our network initialization:

- •
For any $r\in[m]$, we have $\sigma_{w}\left(\sqrt{d}-2\ln(d)\right)\leq\|\bm{w}_{s,r}^{(0)}\|\leq\sigma_{w}\left(\sqrt{d}+2\ln(d)\right).$

The proof of Proposition [C.1](#A3.Thmtheorem1) is the same as the proof of [A.10](#A1.Thmtheorem10).

###### Definition C.2 (Activation Region over Data Input) .

Let $T_{s,r,j}:=\{i\in I_{j}:\langle\bm{w}_{s,r}^{(0)},\bm{x}_{i}\rangle+b_{s,r}^{(0)}\geq 0\}$ be the set of indices of training data points in the $j$-th cluster which can activate the neuron with weight $\bm{w}_{s,r}$ at time step 0.

Then, we give the following result about the activation region $T_{s,r,j}$.

###### Proposition C.3 .

Assuming Proposition [A.6](#A1.Thmtheorem6) and Proposition [A.10](#A1.Thmtheorem10) holds. Then with probability at least $1-(hk)^{-0.01}-2hk^{2}\exp\left(-\frac{n}{9k^{3}h^{2}}\right)$, for all $r\in[h],s,j\in J$, we have

$$ $|T_{s,r,j}|\geq\dfrac{n}{3k^{3}h^{2}}.$ $$

The proof of this lemma is the same as the proof of Proposition [A.12](#A1.Ex43).

###### Lemma C.4 .

Assuming Proposition [C.1](#A3.Thmtheorem1) holds, for all $i\in I,s\in J$, we have

$$ $-\mathbbm{1}(\bm{x}_{i}\in I_{s})+\dfrac{1}{2k}\leq\ell_{i,s}^{\prime(0)}\leq-\mathbbm{1}(\bm{x}_{i}\in I_{s})+\dfrac{2}{k}.$ $$

###### Proof of Lemma C.4 .

By Proposition [C.1](#A3.Thmtheorem1) and Property 2 in [A.6](#A1.Thmtheorem6), for every $s\in J$, we have

$$ $\left|f_{s}^{(0)}(\bm{x}_{i})\right|\leq\dfrac{1}{h}\sum\limits_{r\in[h]}\left|\langle\bm{w}_{s,r}^{(0)},\bm{x}_{i}\rangle\right|\leq\dfrac{1}{h}\sum\limits_{r\in[h]}|\bm{w}_{s,r}^{(0)}|\left|\bm{x}_{i}\right|\leq 4\sigma_{w}d\leq\ln(2).$ $$

Thus $1\leq\exp(f_{s}(\bm{x}_{i}))\leq 2$. Then we have

$$ $\dfrac{1}{2k}\leq\dfrac{\exp(f_{s}^{(0)}(\bm{x}_{i}))}{\sum\limits_{p\in J}\exp(f_{p}^{(0)}(\bm{x}_{i}))}\leq\dfrac{2}{k}.$ $$

$$ $-\mathbbm{1}(\bm{x}_{i}\in I_{s})+\dfrac{1}{2k}\leq\ell_{i,s}^{\prime(0)}=-\mathbbm{1}(\bm{x}_{i}\in I_{s})+\dfrac{\exp(f_{s}^{(0)}(\bm{x}_{i}))}{\sum\limits_{p\in J}\exp(f_{p}^{(0)}(\bm{x}_{i}))}\leq-\mathbbm{1}(\bm{x}_{i}\in I_{s})+\dfrac{2}{k}.$ $$

∎

###### Proposition C.1 .

###### Definition C.2 (Activation Region over Data Input) .

###### Proposition C.3 .

###### Lemma C.4 .

###### Proof of Lemma C.4 .

### C.2 Analysis of Training Dynamics

Denote $S_{i,s}^{(t)}:=\{r\in[h]:\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle\geq 0\}$ for $i\in I,s\in J$.

###### Lemma C.5 .

For every $i\in I$, we have $S_{i,c(i)}^{(1)}=[h]$.

###### Proof of Lemma C.5 .

This proof is similar to the proof of Lemma
[B.1](#A2.Thmtheorem1).

For every $i\in I,r\in[h]$, we have

$$ $\displaystyle\langle\bm{w}_{c(i),r}^{(1)},\bm{x}_{i}\rangle$ $\displaystyle=$ $\displaystyle\langle\bm{w}_{c(i),r}^{(0)},\bm{x}_{i}\rangle-\eta\langle\nabla_{\bm{w}_{c(i),r}}\mathcal{L}(\bm{\theta}^{(0)}),\bm{x}_{i}\rangle$ $$

We examine the update term $\langle\nabla_{\bm{w}_{c(i),r}}\mathcal{L}(\bm{\theta}^{(0)}),\bm{x}_{i}\rangle$.

$$ $\displaystyle\begin{aligned} &-hn\langle\nabla_{\bm{w}_{c(i),r}}\mathcal{L}(\bm{\theta}^{(0)}),\bm{x}_{i}\rangle\\ =&-hn\langle\sum\limits_{j\in I_{c(i)}}\ell_{j,c(i)}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{c(i),r}^{(0)},\bm{x}_{j}\rangle\geq 0\rangle\right)\bm{x}_{j}+\sum\limits_{j\notin I_{c(i)}}\ell_{j,c(i)}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{c(i),r}^{(0)},\bm{x}_{j}\rangle\geq 0\rangle\right)\langle\bm{x}_{j},\bm{x}_{i}\rangle\\ \geq&-hn\sum\limits_{j\in I_{c(i)}}\ell_{j,c(i)}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{c(i),r}^{(0)},\bm{x}_{j}\rangle\geq 0\rangle\right)\langle\bm{x}_{j},\bm{x_{i}}\rangle-\sum\limits_{j\notin I_{c(i)}}\ell_{j,c(i)}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{c(i),r}^{(0)},\bm{x}_{j}\rangle\geq 0\rangle\right)|\langle\bm{x}_{j},\bm{x_{i}}\rangle|\end{aligned}$ $$

By Lemma [C.4](#A3.Thmtheorem4) and Lemma [C.3](#A3.Thmtheorem3), we know that

$$ $\displaystyle-\sum\limits_{j\in I_{c(i)}}\ell_{j,c(i)}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{c(i),r}^{(0)},\bm{x}_{j}\rangle\geq 0\rangle\right)\langle\bm{x}_{j},\bm{x_{i}}\rangle\geq(1-\dfrac{2}{k})\dfrac{d}{2}|T_{c(i),r,c(i)}|\geq\dfrac{dn}{12k^{3}h^{2}}.$ $\displaystyle-\sum\limits_{j\notin I_{c(i)}}\ell_{j,c(i)}^{\prime(0)}\mathbbm{1}\left(\langle\bm{w}_{c(i),r}^{(0)},\bm{x}_{j}\rangle\geq 0\rangle\right)|\langle\bm{x}_{j},\bm{x_{i}}\rangle|\geq-\dfrac{2n\Delta}{k}.$ $$

Combining the two inequalities above, we have

$$ $\displaystyle\langle\bm{w}_{c(i),r}^{(1)},\bm{x}_{i}\rangle$ $\displaystyle\geq\langle\bm{w}_{c(i),r}^{(0)},\bm{x}_{i}\rangle+\dfrac{\eta}{h}\left(\dfrac{d}{12k^{3}h^{2}}-\dfrac{2\Delta}{k}\right)$ $\displaystyle\geq-2\sigma_{w}d+\dfrac{\eta\Delta}{h}\geq 0.$ $$

Therefore, we have $S_{i,c(i)}^{(1)}=[h]$ for every $i\in I$.

∎

Recall the definition of weight decomposition we will use in multi-classification tasks.

###### Lemma C.6 (Weight Decomposition) .

During the training dynamics, there exists the following coefficient sequences $\lambda_{s,r,j}^{(t)}$ and $\sigma_{s,r,i}^{(t)}$ for each neuron $s,j\in J,r\in[h]$ such that

$$ $\displaystyle\bm{w}_{s,r}^{(t)}=\bm{w}_{s,r}^{(0)}+\sum_{j\in J}\lambda_{s,r,j}^{(t)}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}+\sum_{i\in I}\sigma_{s,r,i}^{(t)}\bm{\xi}_{i}\|\bm{\xi}_{i}\|^{-2}.$ $$

###### Corollary C.7 .

The coefficient sequences$\lambda_{r,j}^{(t)}$ and $\sigma_{i,j}^{(t)}$ for each pair $i\in I,r,j\in J$ defined in Lemma [C.6](#A3.Thmtheorem6) satisfy:

$$ $\lambda_{s,r,j}^{(t)}\|\bm{\mu}_{j}\|^{-2}=\sum\limits_{i\in I_{j}}\sigma_{s.r,i}^{(t)}\|\bm{\xi}_{i}\|^{-2}.$ $$

###### Corollary C.8 .

For all $i\in I,r,j\in J$, we have the following update equation for $\lambda_{s,r,j}$ and $\sigma_{s,r,i}$.

$$ $\displaystyle\lambda_{s,r,j}^{(t+1)}$ $\displaystyle=\lambda_{s,r,j}^{(t)}-\dfrac{\eta}{nh}\sum\limits_{p\in I_{j}}\ell_{p,s}^{\prime(t)}\|\bm{\mu}_{j}\|^{2}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{p}\rangle\geq 0\right),$ $\displaystyle\sigma_{s,r,i}^{(t+1)}$ $\displaystyle=\sigma_{s,r,i}^{(t)}-\dfrac{\eta}{nh}\ell_{i,s}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}\mathbbm{1}\left(\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle\geq 0\right),$ $\displaystyle\lambda_{s,r,j}^{(0)}$ $\displaystyle=0,\sigma_{s,r,i}^{(0)}=0.$ $$

###### Corollary C.9 .

The coefficient sequences $\lambda_{s,r,j}^{(t)}$ and $\sigma_{s,r,i}^{(t)}$ for each pair $s,j\in J,i\in I,r\in[h]$ defined in Lemma [C.6](#A3.Thmtheorem6) satisfy:

$$ $\displaystyle\lambda_{s,r,j}^{(t)}$ $\displaystyle\geq 0\ \ \textit{iff $s=j$},$ $\displaystyle\sigma_{s,r,i}^{(t)}$ $\displaystyle\geq 0\ \ \textit{iff $s=c(i)$}.$ $$

Then we reuse the logic of the proof of Lemma [B.7](#A2.Thmtheorem7) to prove the main result in our multi-classification setting.

Denote $q_{i}^{(t)}=f_{c(i)}^{(t)}(\bm{x}_{i}),\hat{q}_{i}^{(t)}=\dfrac{1}{h}\sum\limits_{r\in[h]}\left(\lambda_{c(i),r,c(i)}^{(t)}+\sigma_{c(i),r,i}^{(t)}\right)$.

$\Delta_{q}^{(t)}(i,j)=q_{i}^{(t)}-q_{j}^{(t)},$

$\hat{\Delta}_{q}^{(t)}(i,j)=\hat{q}_{i}^{(t)}-\hat{q}_{j}^{(t)}=\dfrac{1}{h}\sum\limits_{r\in[h]}\left(\lambda_{c(i),r,c(i)}^{(t)}-\lambda_{c(j),r,c(j)}^{(t)}+\sigma_{c(i),r,i}^{(t)}-\sigma_{c(j),r,j}^{(t)}\right)$.

Denote $\epsilon=\text{max}\left\{\dfrac{2\ln(\frac{n}{k})}{\sqrt{\frac{n}{k}}-\ln(\frac{n}{k})},\dfrac{k^{2}\Delta}{d},\dfrac{k^{2}}{n}\right\}$.
We know that $\epsilon=o(k^{-2.5})$ according to our hyper-parameter [Assumption A.4](#A1.Thmtheorem4).

###### Lemma C.10 .

For $t\leq T_{0}=\operatorname{exp}(\tilde{O}(k^{0.5}),i,j\in I,r\in[h]$, we have

- 1.
$\dfrac{k}{2n}\ln(t\eta)\leq\sigma_{c(i),r,i}\leq\dfrac{2k}{n}\ln(t+1)$
- 2.
$\hat{\Delta}_{q}^{(t)}(i,j)\leq 5k\epsilon$ when $c(i)=c(j),$
- 3.
$\hat{\Delta}_{q}^{(t)}(i,j)\leq 32k^{2}\epsilon,$
- 4.
$\Delta_{q}^{(t)}(i,j)\leq 33k^{2}\epsilon,$
- 5.
$\ell_{i,c(i)}^{\prime\left(t\right)}/\ell_{j,c(j)}^{\prime\left(t\right)}\leq 1+14k\epsilon$ when $c(i)=c(j)$,
- 6.
$\ell_{i,c(i)}^{\prime\left(t\right)}/\ell_{j,c(j)}^{\prime\left(t\right)}\leq 1+67k^{2}\epsilon$,
- 7.
$S_{i,c(i)}^{(1)}=[h]$,
- 8.
$\left|\lambda_{s,r,c(i)}^{(t)}\right|\leq\epsilon,\left|\sigma_{s,r,c(i)}^{(t)}\right|\leq 2\epsilon$ for $s\in J,s\neq c(i)$.

###### Proof of Lemma C.10 .

Since the proof of this lemma follows exactly the same logic as Lemma [B.7](#A2.Thmtheorem7), we omit some details and only outlined the necessary lemmas and the key steps of the proof.

First, the base case of the induction is simple, so we only consider the inductive step.

###### Lemma C.11 .

Assuming the inductive hypotheses hold before time step $t$, for all $r\in[h],s,j\in J$, we have

$$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(t)},\bm{\mu}_{j}\rangle-\lambda_{s,r,j}^{(t)}\right|\leq\dfrac{\epsilon}{6}.$ $$

###### Lemma C.12 .

Assuming the inductive hypotheses hold before time step $t$, for all $r\in[h],s,j\in J$, we have

$$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(t)},\bm{\xi}_{i}\rangle-\sigma_{s,r,i}^{(t)}\right|\leq\dfrac{\epsilon}{6}.$ $$

###### Lemma C.13 .

Assuming the inductive hypotheses hold before time step $t$, for all $r\in[h],s,j\in J$, we have

$$ $\displaystyle\left|\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle-\lambda_{s,r,c(i)}^{(t)}-\sigma_{s,r,i}^{(t)}\right|\leq\dfrac{\epsilon}{3}.$ $$

###### Lemma C.14 .

Assuming the inductive hypotheses hold before time step $t$, for all $i\in I$, we have

$$ $\displaystyle|q_{i}^{(t)}-\hat{q}_{i}^{(t)}|\leq\epsilon$ $$

The proofs of these three lemmas are identical to the proofs of [Lemma B.11](#A2.Thmtheorem11), [Lemma B.12](#A2.Thmtheorem12), [Lemma B.13](#A2.Thmtheorem13) and [Lemma B.14](#A5.EGx47) in Appendix [Appendix B](#A2), except that in the previous proof, there was an additional subscript $s$ used to indicate 2-classification label, whereas here it is used to represent a fine-grained $k$-classification label.

###### Corollary C.15 .

Assuming the inductive hypotheses hold before time step $t$, for all $i,j\in I$, we have

$$ $\left|\Delta_{q}^{(t)}(i,j)-\hat{\Delta}_{q}^{(t)}(i,j)\right|\leq 2\epsilon.$ $$

Next, we present the key steps of the auto-balance process for $\hat{\Delta}_{q}(i,j)$.

###### Lemma C.16 (Property 8) .

Assuming the inductive hypotheses hold before time step $t$, for $i\in I,s\in J,s\neq c(i),r\in[m]$, we have
$|\lambda_{s,r,c(i)}^{(t)}|\leq\epsilon,|\sigma_{s,r,i}^{(t)}|\leq 2\epsilon$.

###### Proof of Lemma C.16 .

We distinguish between two scenarios.

Case(I):
For all $i\in I_{j}$, $\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle<0$.

Then by Corollary [C.8](#A5.EGx102) and inductive hypothesis, we know that

$$ $|\lambda_{s,r,j}^{(t+1)}|=|\lambda_{s,r,j}^{(t)}|\leq\epsilon.$ $$

Case(II):
There exists $i\in I_{j}$ such that $\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle\geq 0.$

By Lemma [C.13](#A3.Thmtheorem13), we know that

$$ $\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle-\lambda_{s,r,j}^{(t)}-\sigma_{s,r,i}^{(t)}\leq\dfrac{\epsilon}{3}.$ $$

Then noting that $\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle\geq 0$ and $\sigma_{s,r,i}^{(t)}\leq 0$, we have

$$ $\lambda_{s,r,j}^{(t)}\geq\langle\bm{w}_{s,r}^{(t)},\bm{x}_{i}\rangle-\sigma_{s,r,i}^{(t)}-\dfrac{\epsilon}{3}\geq-\dfrac{\epsilon}{3}.$ $$

Then using the conclusion in Corollary [C.8](#A5.EGx102), we know that

$$ $\displaystyle|\lambda_{s,r,j}^{(t+1)}-\lambda_{s,r,j}^{(t)}|\leq\dfrac{\eta d}{nm}\sum\limits_{i\in I_{j}}-\ell_{i,s}^{\prime(t)}\leq\dfrac{\eta d}{m}\leq\dfrac{\epsilon}{3}.$ $$

Thus we have $\lambda_{r,j}^{(t+1)}\geq-\epsilon$. Noting that $\lambda_{s,r,j}^{(t+1)}\leq 0$, we have $|\lambda_{s,r,j}^{(t+1)}|\leq\epsilon$.
Then by [Corollary C.7](#A3.Thmtheorem7), we know that $|\sigma_{s,r,i}^{(t)}|\leq 2|\lambda_{s,r,j}^{(t+1)}|\leq 2\epsilon.$
∎

###### Lemma C.17 .

Assuming the inductive hypotheses hold before time step $t$, for all $s\in J,i\in I,s\neq c(i)$, we have

$$ $\displaystyle f_{s}(\bm{x_{i}})\leq\epsilon.$ $$

###### Proof.

By [Lemma C.13](#A3.Thmtheorem13) and [Corollary C.9](#A3.Thmtheorem9), we know that

$$ $\displaystyle f_{s}(\bm{x_{i}})$ $\displaystyle=\dfrac{1}{h}\sum\limits_{r\in[h]}\text{ReLU}(\langle\bm{w}_{s,r}^{(t)},\bm{x_{i}}\rangle)$ $\displaystyle\leq\dfrac{1}{h}\sum\limits_{r\in[h]}|\langle\bm{w}_{s,r}^{(t)},\bm{x_{i}}\rangle|$ $\displaystyle\leq\dfrac{1}{h}\sum\limits_{r\in[h]}(\dfrac{\epsilon}{3}+\lambda_{s,r,c(i)}^{(t)}+\sigma_{s,r,i}^{(t)})$ $\displaystyle\leq\dfrac{1}{h}\sum\limits_{r\in[h]}\dfrac{\epsilon}{3}$ $\displaystyle\leq\epsilon.$ $$

∎

###### Lemma C.18 .

Assuming the inductive hypotheses hold before time step $t$, for any two training data points $\bm{x}_{i},\bm{x}_{j}$, if $q_{i}\geq q_{j}$, we have

$$ $\dfrac{1}{1+2\epsilon}\dfrac{e^{\Delta_{q}^{(t)}(i,j)}+k-1}{k}\leq\dfrac{\ell^{\prime(t)}_{j,c(j)}}{\ell^{\prime(t)}_{i,c(i)}}\leq e^{\Delta_{q}^{(t)}(i,j)}(1+2\epsilon).$ $$

###### Proof of Lemma C.18 .

By [Lemma C.17](#A3.Thmtheorem17) and noting that $\exp(\epsilon)\leq 1+2\epsilon$, we know that

$$ $\displaystyle\dfrac{k-1}{k-1+\exp(q_{i}^{(t)})}\leq-\ell^{\prime(t)}_{i,c(i)}=\dfrac{\sum\limits_{p\neq c(i)}\exp(f_{p}^{(t)}(\bm{x}_{i}))}{\sum\limits_{p\in J}\exp(f_{p}^{(t)}(\bm{x}_{i}))}\leq\dfrac{(k-1)\exp(\epsilon)}{(k-1)\exp(\epsilon)+\exp(q_{i}^{(t)})}\leq\dfrac{(k-1)(1+2\epsilon)}{k-1+\exp(q_{i}^{(t)})}$ $$

Thus we know that

$$ $\displaystyle\dfrac{1}{(1+2\epsilon)}\dfrac{k-1+\exp(q_{i}^{(t)})}{k-1+\exp(q_{j}^{(t)})}\leq\dfrac{\ell^{\prime(t)}_{j,c(j)}}{\ell^{\prime(t)}_{i,c(i)}}\leq(1+2\epsilon)\dfrac{k-1+\exp(q_{i}^{(t)})}{k-1+\exp(q_{j}^{(t)})}$ $$

$$ $\displaystyle\dfrac{k-1+\exp(q_{i}^{(t)})}{k-1+\exp(q_{j}^{(t)})}$ $\displaystyle=1+\dfrac{\exp(q_{i}^{(t)})-\exp(q_{j}^{(t)})}{k-1+\exp(q_{j}^{(t)})}$ $\displaystyle=1+\dfrac{\exp(q_{j}^{(t)})}{k-1+\exp(q_{j}^{(t)})}\left(\exp\left(\Delta_{q}^{(t)}(i,j)\right)-1\right)$ $$

For the second term,

$$ $\dfrac{\exp\left(\Delta_{q}^{(t)}(i,j)\right)-1}{k}\leq\dfrac{\exp(q_{j}^{(t)})}{k-1+\exp(q_{j}^{(t)})}\left(\exp\left(\Delta_{q}^{(t)}(i,j)\right)-1\right)\leq\exp\left(\Delta_{q}^{(t)}(i,j)\right)-1.$ $$

Thus we know

$$ $\dfrac{1}{1+2\epsilon}\dfrac{e^{\Delta_{q}^{(t)}(i,j)}+k-1}{k}\leq\dfrac{\ell^{\prime(t)}_{j,c(j)}}{\ell^{\prime(t)}_{i,c(i)}}\leq e^{\Delta_{q}^{(t)}(i,j)}(1+2\epsilon).$ $$

∎

We first consider the case when $c(i)=c(j)$. We distinguish between two scenarios, one is when $\left|\hat{\Delta}_{q}^{(t)}(i,j)\right|$ is relatively small and the other is when $\left|\hat{\Delta}_{q}^{(t)}(i,j)\right|$ is relatively large.

Case(I): $\hat{\Delta}_{q}^{(t)}(i,j)\leq 4k\epsilon.$

In this case, we have
$\hat{\Delta}_{q}^{(t+1)}(i,j)\leq 5k\epsilon$ due to small learning rate $\eta$.

Case(II): $\hat{\Delta}_{q}^{(t)}(i,j)\geq 4k\epsilon.$

By Lemma [C.14](#A5.EGx107), we know that

$$ $\Delta_{q}^{(t)}(i,j)\geq\hat{\Delta}_{q}^{(t)}(i,j)-2\epsilon\geq 3k\epsilon.$ $$

By Lemma [C.18](#A3.Ex330), we know that

$$ $\displaystyle\dfrac{\ell^{\prime(t)}_{j,c(j)}}{\ell^{\prime(t)}_{i,c(i)}}\geq\dfrac{1}{1+2\epsilon}\dfrac{e^{\Delta_{q}^{(t)}(i,j)}+k-1}{k}\geq\dfrac{1}{1+2\epsilon}(1+\Delta_{q}^{(t)}(i,j)/k)\geq 1+\epsilon/2.$ $$

Noting that $c(i)=c(j)$, we know that

$$ $\displaystyle\hat{\Delta}_{q}^{(t+1)}(i,j)-\hat{\Delta}_{q}^{(t)}(i,j)$ $\displaystyle=$ $\displaystyle-\dfrac{\eta}{nh}\left(\ell_{i,c(i)}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}-\ell_{j,c(j)}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right)$ $\displaystyle\leq$ $\displaystyle\ \ \ 0$ $$

Then due to the inductive hypothesis,

$$ $\hat{\Delta}_{q}^{(t+1)}(i,j)\leq\hat{\Delta}_{q}^{(t)}(i,j)\leq 5k\epsilon.$ $$

By Corollary [C.15](#A3.Thmtheorem15), we can get Property 4

$$ $\Delta_{q}^{(t)}(i,j)\leq\Delta_{q}^{(t)}(i,j)+2\epsilon\leq 6k\epsilon.$ $$

By Lemma [C.18](#A3.Ex330) and noting that $e^{x}\leq 1+2x$ for small $x$ we know that

$$ $\dfrac{\ell^{\prime(t+1)}_{j,c(j)}}{\ell^{\prime(t+1)}_{i,c(i)}}\leq(1+2\epsilon)e^{\Delta_{q}^{(t+1)}(i,j)}\leq(1+2\epsilon)(1+2\Delta_{q}^{(t+1)}(i,j))\leq 1+14k\epsilon.$ $$

Next, we consider the case when $c(i)\neq c(j)$. We also distinguish between the two scenarios.

Case(I): $\hat{\Delta}_{q}^{(t)}(i,j)\leq 31k^{2}\epsilon$.
In this case, we have
$\hat{\Delta}_{q}^{(t+1)}(i,j)\leq 32k^{2}\epsilon$ due to small learning rate $\eta$.

Case(II): $\hat{\Delta}_{q}^{(t)}(i,j)\geq 31k^{2}\epsilon.$

By Lemma [C.14](#A5.EGx107), we know that

$$ $\Delta_{q}^{(t)}(i,j)\geq\hat{\Delta}_{q}^{(t)}(i,j)-2\epsilon\geq 30k^{2}\epsilon.$ $$

By Lemma [C.18](#A3.Ex330), we know that

$$ $\displaystyle\dfrac{\ell^{\prime(t)}_{j,c(j)}}{\ell^{\prime(t)}_{i,c(i)}}\geq\dfrac{1}{1+2\epsilon}\dfrac{e^{\Delta_{q}^{(t)}(i,j)}+k-1}{k}\geq\dfrac{1}{1+2\epsilon}(1+\Delta_{q}^{(t)}(i,j)/k)\geq 1+29k\epsilon.$ $$

Furthermore, due to the inductive hypothesis, for any $p\in I_{c(i)},q\in I_{c(j)}$, we know that

$$ $\displaystyle\dfrac{\ell_{p,c(p)}^{\prime(t)}}{\ell_{q,c(q)}^{\prime(t)}}\leq\dfrac{(1+14\epsilon)\ell_{i,c(i)}^{\prime(t)}}{\ell_{j,c(j)}^{\prime(t)}/(1+14k\epsilon)}\leq\dfrac{(1+14k\epsilon)^{2}}{1+29k\epsilon}\leq\dfrac{1}{1+\epsilon}.$ $$

We know that

$$ $\displaystyle\hat{\Delta}_{q}^{(t+1)}(i,j)-\hat{\Delta}_{q}^{(t)}(i,j)$ $\displaystyle=$ $\displaystyle-\dfrac{\eta}{nh}\left(\sum\limits_{p\in I_{c(i)}}\ell_{p,c(i)}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}-\sum\limits_{p\in I_{c(j)}}\ell_{p,c(j)}^{\prime(t)}\|\bm{\mu}_{c(j)}\|^{2}+\ell_{i,c(i)}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}-\ell_{j,c(j)}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right)$ $\displaystyle=$ $\displaystyle-\dfrac{\eta}{nh}\left(\ell_{i,c(i)}^{\prime(t)}\|\bm{\xi}_{i}\|^{2}-\ell_{j,c(j)}^{\prime(t)}\|\bm{\xi}_{j}\|^{2}\right)-\dfrac{\eta}{nh}\left(\sum\limits_{p\in I_{c(i)}}\ell_{p,c(i)}^{\prime(t)}\|\bm{\mu}_{c(i)}\|^{2}-\sum\limits_{p\in I_{c(j)}}\ell_{p,c(j)}^{\prime(t)}\|\bm{\mu}_{c(j)}\|^{2}\right)$ $\displaystyle\leq$ $\displaystyle 0.$ $$

By inductive hypothesis, we know that

$$ $\hat{\Delta}_{q}^{(t+1)}(i,j)\leq\hat{\Delta}_{q}^{(t)}(i,j)\leq 32k^{2}\epsilon.$ $$

Now, we have completed the main part of the proof, the inductive proofs of Properties 2 and 3. Subsequently, Properties 4, 5, and 6 can be directly derived from [Lemma C.18](#A3.Ex330).

By Corollary [C.15](#A3.Thmtheorem15), we can get Property 4

$$ $\Delta_{q}^{(t)}(i,j)\leq\Delta_{q}^{(t)}(i,j)+2\epsilon\leq 33k^{2}\epsilon$ $$

By Lemma [C.18](#A3.Ex330) and noting that $e^{x}\leq 1+2x$ for small $x$ we know that

$$ $\dfrac{\ell^{\prime(t+1)}_{j,c(j)}}{\ell^{\prime(t+1)}_{i,c(i)}}\leq(1+2\epsilon)e^{\Delta_{q}^{(t+1)}(i,j)}\leq(1+2\epsilon)(1+2\Delta_{q}^{(t+1)}(i,j))\leq 1+67k^{2}\epsilon.$ $$

###### Lemma C.19 .

For all $i_{1},i_{2}\in I,r_{1},r_{2}\in[h]$, we have

$$ $\left(1-200k^{2}\epsilon\right)\sigma_{c(i_{2}),r_{2},i_{2}}^{(t)}-\dfrac{1}{nd}\leq\sigma_{c(i_{1}),r_{1},i_{1}}^{(t)}\leq\left(1+200k^{2}\epsilon\right)\sigma_{c(i_{2}),r_{2},i_{2}}^{(t)}+\dfrac{1}{nd}.$ $$

###### Lemma C.20 .

For every $i\in I,r_{0}\in[h]$, we have
$\dfrac{n}{2k}\sigma_{c(i),r_{0},i}^{(t)}-2\epsilon\leq q_{i}^{(t)}\leq\dfrac{2n}{k}\sigma_{c(i),r_{0},i}^{(t)}+2\epsilon.$

The proof of these lemmas are the same as the proof of [Lemma B.19](#A2.Thmtheorem19) and [Lemma B.20](#A2.Thmtheorem20).

###### Lemma C.21 .

For every $i\in I,r\in[h]$, we have

$$ $\dfrac{1}{2}\exp\left(-\dfrac{2n}{k}\sigma_{1,r,i}^{(t)}\right)\leq-\ell_{i,c(i)}^{\prime(t)}\leq 2k\exp\left(-\dfrac{n}{2k}\sigma_{1,r,i}^{(t)}\right).$ $$

###### Proof of Lemma C.21 .

By Lemma [C.20](#A3.Thmtheorem20), we know that

$$ $\displaystyle-\ell_{i,c(i)}^{\prime(t)}$ $\displaystyle\geq\dfrac{k-1}{k-1+\exp\left(q_{i}^{(t)}\right)}$ $\displaystyle\geq\dfrac{k-1}{k-1+\exp\left(\dfrac{2n}{k}\sigma_{1,r,i}^{(t)}+2\epsilon\right)}$ $\displaystyle\geq\dfrac{k-1}{k-1+2\exp\left(\dfrac{2n}{k}\sigma_{1,r,i}^{(t)}\right)}$ $\displaystyle\geq\dfrac{1}{2}\exp\left(-\dfrac{2n}{k}\sigma_{1,r,i}^{(t)}\right).$ $$

By [Lemma C.17](#A3.Thmtheorem17) and [Lemma C.20](#A3.Thmtheorem20), we have

$$ $\displaystyle-\ell_{i,c(i)}^{\prime(t)}$ $\displaystyle\leq\dfrac{(k-1)\exp(\epsilon)}{(k-1)\exp(\epsilon)+\exp\left(q_{i}^{(t)}\right)}$ $\displaystyle\leq(1+2\epsilon)\dfrac{k-1}{k-1+\exp\left(\dfrac{n}{2k}\sigma_{1,r,i}^{(t)}-2\epsilon\right)}$ $\displaystyle\leq\dfrac{2(k-1)}{k-1+\exp\left(\dfrac{n}{2k}\sigma_{1,r,i}^{(t)}\right)}$ $\displaystyle\leq 2k\exp\left(-\dfrac{n}{2k}\sigma_{1,r,i}^{(t)}\right)$ $$

∎

Next, we prove Property 1.

$$ $\displaystyle\sigma_{c(i),r,i}^{(t+1)}$ $\displaystyle=\sigma_{c(i),r,i}^{(t)}-\dfrac{\eta}{nh}\cdot\ell_{i,c(i)}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}$ $\displaystyle\leq\sigma_{c(i),r,i}^{(t)}+\dfrac{2k\eta d}{nh}\exp(-\dfrac{n}{2k}\sigma_{c(i),r,i}^{(t)})$ $\displaystyle\leq\dfrac{2k}{n}\ln(t+1)+\dfrac{2k}{n}\dfrac{1}{2t}$ $\displaystyle\leq\dfrac{2k}{n}\ln(t+2)$ $$

$$ $\displaystyle\sigma_{c(i),r,i}^{(t+1)}$ $\displaystyle=\sigma_{c(i),r,i}^{(t)}-\dfrac{\eta}{nh}\cdot\ell_{i,c(i)}^{\prime(t)}\|\bm{\xi_{i}}\|^{2}$ $\displaystyle\geq\sigma_{c(i),r,i}^{(t)}+\dfrac{\eta d}{2nh}\exp(-\dfrac{k}{2n}\sigma_{c(i),r,i}^{(t)})$ $\displaystyle\geq\dfrac{k}{2n}\ln(t\eta)+\dfrac{k}{2n}\dfrac{2}{t}$ $\displaystyle\geq\dfrac{k}{2n}\ln((t+1)\eta)$ $$

Finally, we prove Property 7.

By Property 6 in the inductive hypotheses, we know that for all $i\in I$,

$$ $-\dfrac{1}{1+67k^{2}\epsilon}\ell_{1,c(1)}^{\prime(t)}\leq|\ell_{i,c(i)}^{\prime(t)}|\leq-(1+67k^{2}\epsilon)\ell_{1,c(1)}^{\prime(t)}.$ $$

Then we know that

$$ $\displaystyle-\langle\nabla_{\bm{w}_{p}}\mathcal{L}(\bm{\theta}^{(t)}),\bm{x}_{i}\rangle\geq$ $\displaystyle-\sum\limits_{p\in I_{c(i)}}\ell_{p,c(i)}^{\prime(t)}\langle\bm{x}_{p},\bm{x}_{i}\rangle-\sum\limits_{p\notin I_{c(i)}}|\ell_{p,c(i)}^{\prime(t)}|\langle\bm{x}_{p},\bm{x}_{i}\rangle|$ $\displaystyle\geq$ $\displaystyle-\sum\limits_{p\in I_{c(i)}}\ell_{p,c(p)}^{\prime(t)}\langle\bm{x}_{p},\bm{x}_{i}\rangle-\sum\limits_{p\notin I_{c(i)}}|\ell_{p,c(p)}^{\prime(t)}|\langle\bm{x}_{p},\bm{x}_{i}\rangle|$ $\displaystyle\geq$ $\displaystyle-\ell_{1,c(1)}^{(t)}\left(\sum\limits_{p\in I_{c(i)}}\dfrac{d}{2(1+67k^{2}\epsilon)}-\sum\limits_{p\notin I_{c(i)}}(1+67k^{2}\epsilon)\Delta\right)$ $\displaystyle\geq$ $\displaystyle-\ell_{1,c(1)}^{(t)}\left(\dfrac{dn}{4k(1+67k^{2}\epsilon)}-(1+67k^{2}\epsilon)n\Delta\right)\geq 0.$ $$

By Property 7 in the inductive hypotheses, we know that $\langle\bm{w}_{c(i)}^{(t+1)},\bm{x}_{i}\rangle\geq\langle\bm{w}_{c(i)}^{(t)},\bm{x}_{i}\rangle\geq 0$.

We complete the proof of Lemma [C.10](#A3.Thmtheorem10).

∎

###### Lemma C.5 .

###### Proof of Lemma C.5 .

###### Lemma C.6 (Weight Decomposition) .

###### Corollary C.7 .

###### Corollary C.8 .

###### Corollary C.9 .

###### Lemma C.10 .

###### Proof of Lemma C.10 .

###### Lemma C.11 .

###### Lemma C.12 .

###### Lemma C.13 .

###### Lemma C.14 .

###### Corollary C.15 .

###### Lemma C.16 (Property 8) .

###### Proof of Lemma C.16 .

###### Lemma C.17 .

###### Proof.

###### Lemma C.18 .

###### Proof of Lemma C.18 .

###### Lemma C.19 .

###### Lemma C.20 .

###### Lemma C.21 .

###### Proof of Lemma C.21 .

### C.3 Proof of Theorem 4.7

###### Theorem C.22 (Restatement of Theorem 4.7 ) .

In the setting of training a multi-class network on the multiple classification problem $\tilde{{\mathcal{S}}}:=\{({\bm{x}}_{i},\tilde{y}_{i})\}_{i=1}^{n}\subseteq\mathbb{R}^{d}\times[k]$ as described in the above,
under [Assumptions 3.2](#S3.Thmtheorem2), [3.3](#S3.Thmtheorem3) and [4.3](#S4.Thmtheorem3),
for some $\gamma=o(1)$, after $\Omega(\eta^{-1}k^{8})\leq T\leq\exp(\tilde{O}(k^{1/2}))$ iterations, with probability at least $1-\gamma$, the neural network satisfies the following properties:

- 1.
The clean accuracy is nearly perfect: $\mathrm{Acc}_{\mathrm{clean}}^{\mathcal{D}}(F^{\mathrm{binary}}_{{\bm{\theta}}^{(T)}})\geq 1-\exp(-\Omega(\log^{2}d))$.
- 2.
The network converges to the feature-decoupling regime: there exists a time-variant coefficient $\lambda^{(T)}\in[\Omega(\operatorname{log}k),+\infty)$ such that for all $j\in[k]$, $r\in[h]$, the weight vector $\bm{w}_{j,r}^{(T)}$ can be approximated as
$\displaystyle\bigg{\|}{\bm{w}}_{j,r}^{(T)}-\lambda^{(T)}\|\bm{\mu}_{j}\|^{-2}\bm{\mu}_{j}\bigg{\|}$
$\displaystyle\leq o(d^{-1/2}).$
- 3.
Consequently, the corresponding binary classifier achieves optimal robustness: for perturbation radius $\delta=O(\sqrt{d})$, the $\delta$-robust accuracy is also nearly perfect, i.e., $\mathrm{Acc}_{\mathrm{robust}}^{\mathcal{D}}(F^{\mathrm{binary}}_{{\bm{\theta}}^{(T)}};\delta)\geq 1-\exp(-\Omega(\log^{2}d))$.

###### Proof of Theorem Theorem C.22 .

We first prove that the network converges to the feature-decoupling regime(Property 2).

###### Lemma C.23 .

For all $j\in J,r\in[h]$, we have $\dfrac{\ln(T\eta)}{4}\leq\lambda_{j,r,j}^{(T)}\leq 4\ln(T+1).$

###### Proof of Lemma C.23 .

Using Property 1 in [Lemma C.10](#A3.Thmtheorem10) and Corollary [C.7](#A3.Thmtheorem7), we know that

$$ $\displaystyle\lambda_{j,r,j}^{(T)}$ $\displaystyle=\sum\limits_{p\in I_{j}}\dfrac{\|\bm{\xi}_{p}\|^{2}}{\|\bm{\mu}_{j}\|^{2}}\sigma_{j,r,p}^{(T)}\geq\dfrac{\ln(T\eta)}{4}$ $\displaystyle\lambda_{j,r,j}^{(T)}$ $\displaystyle=\sum\limits_{p\in I_{j}}\dfrac{\|\bm{\xi}_{p}\|^{2}}{\|\bm{\mu}_{j}\|^{2}}\sigma_{j,r,p}^{(T)}\leq 4\ln(T+1)$ $$

∎

###### Lemma C.24 .

For $r_{1},r_{2}\in[h],j_{1},j_{2}\in J$, we have

$$ $\displaystyle\dfrac{\lambda_{j_{1},r_{1},j_{1}}^{(T)}}{\lambda_{j_{2},r_{2},j_{2}}^{(T)}}\leq 1+204k^{2}\epsilon.$ $$

###### Proof.

By Lemma [C.19](#A3.Thmtheorem19), we know that for any $i_{1}\in I_{j_{1}},i_{2}\in I_{j_{2}}$

$$ $\sigma_{j_{1},r_{1},i_{1}}^{(t)}\leq\left(1+200k^{2}\epsilon\right)\sigma_{j_{2},r_{2},i_{2}}^{(t)}+(nd)^{-1}\leq(1+201k^{2}\epsilon)\sigma_{j_{1},r_{2},i_{2}}^{(t)}.$ $$

$$ $\displaystyle\dfrac{\lambda_{j_{1},r_{1},j_{1}}^{(T)}}{\lambda_{j_{2},r_{2},j_{2}}^{(T)}}$ $\displaystyle=\dfrac{\sum\limits_{p\in I_{j_{1}}}\|\bm{\xi}_{p}\|^{2}\sigma_{j_{1},r_{1},p}^{(t)}}{\sum\limits_{p\in I_{j_{2}}}\|\bm{\xi}_{p}\|^{2}\sigma_{j_{2},r_{2},p}^{(t)}}$ $\displaystyle\leq(1+201k^{2}\epsilon)\dfrac{\|\sqrt{d}+\ln(d)\|^{2}}{\|\sqrt{d}-\ln(d)\|^{2}}\dfrac{|I_{j_{1}}|}{|I_{j_{2}}|}$ $\displaystyle\leq(1+201k^{2}\epsilon)(1+\epsilon)(1+\epsilon)$ $\displaystyle\leq 1+204k^{2}\epsilon.$ $$

∎

We denote $\lambda^{(T)}=\lambda_{1,1,1}^{(T)}$ as the representative of $\{\lambda_{j,r,j}:r\in[m],j\in J\}$.

By Lemma [C.23](#A3.Thmtheorem23) and Lemma [C.24](#A3.Thmtheorem24) ,for all $r\in[h],j\in J$, we have

$$ $\displaystyle|\lambda^{(T)}-\lambda_{j,r,j}^{(T)}|$ $\displaystyle\leq 204k^{2}\epsilon\lambda^{(T)},$ $\displaystyle\lambda^{(T)}$ $\displaystyle\leq 4\ln(T+1),$ $\displaystyle\lambda^{(T)}$ $\displaystyle\geq\dfrac{\ln(T\eta)}{4}\geq 2\ln(k)=\Omega(\log(k)).$ $$

###### Lemma C.25 .

For all $s\in J,r\in[h]$, We have

$$ $\displaystyle\sqrt{d}\left\|\bm{w}_{s,r}^{(T)}-\lambda^{(T)}\bm{\mu}_{s}\|\bm{\mu}_{s}\|^{-2}\right\|=o(1).$ $$

###### Proof of Lemma B.25 .

Recall weight decomposition in Lemma [B.2](#A2.Thmtheorem2).

$$ $\displaystyle\sqrt{d}\left(\bm{w}^{(T)}_{s,r}-\lambda^{(T)}\bm{\mu}_{s}\|\bm{\mu}_{s}\|^{-2}\right)$ $\displaystyle=\underbrace{\sqrt{d}\bm{w}_{s,r}^{(0)}}_{\mathcal{L}_{1}}+\underbrace{\sqrt{d}\left(\lambda_{s,r,s}^{(T)}-\lambda^{(T)}\right)\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}}_{\mathcal{L}_{2}}$ $\displaystyle+\underbrace{\sqrt{d}\sum_{j\neq s}\lambda_{s,r,j}^{(T)}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}}_{\mathcal{L}_{3}}+\underbrace{\sqrt{d}\sum_{i\in I}\sigma_{s,r,i}^{(T)}\bm{\xi}_{i}\|\bm{\xi}_{i}\|^{-2}}_{\mathcal{L}_{4}}$ $$

For $\mathcal{L}_{1}$ term, using the conclusion in Lemma [A.10](#A1.Thmtheorem10), we know that

$$ $\displaystyle\|\sqrt{d}\bm{w}_{s,r}^{(0)}\|\leq 2d\sigma_{w}\leq\epsilon=o(1).$ $$

For $\mathcal{L}_{2}$ term, using the conclusion in Lemma [C.23](#A3.Thmtheorem23), Lemma [C.24](#A3.Thmtheorem24), we know that

$$ $\displaystyle\left\|\sqrt{d}\left(\lambda_{s,r,s}^{(T)}-\lambda^{(T)}\right)\bm{\mu}_{s}\|\bm{\mu}_{s}\|^{-2}\right\|=|\lambda_{s,r,s}^{(T)}-\lambda^{(T)}|\leq 204k^{2}\epsilon\lambda^{(T)}\leq 816k^{2}\epsilon\ln(T+1)\leq 900k^{2.5}\epsilon=o(1).$ $$

For $\mathcal{L}_{3}$ term, by [Lemma C.16](#A3.Thmtheorem16) and triangle inequality, we know that

$$ $\displaystyle\left\|\sqrt{d}\sum_{j\neq s}\lambda_{s,r,j}^{(T)}\bm{\mu}_{j}\|\bm{\mu}_{j}\|^{-2}\right\|\leq k\epsilon=o(1).$ $$

For $\mathcal{L}_{4}$ term, by Property (1) in Lemma [B.7](#A2.Thmtheorem7), we have

$$ $\displaystyle\left\|\sqrt{d}\sum_{i\in I}\sigma_{s,r,i}^{(T)}\bm{\xi}_{i}\|\bm{\xi}_{i}\|^{-2}\right\|^{2}$ $\displaystyle=d\sum\limits_{i\in I}\left(\sigma_{s,r,i}^{(T)}\right)^{2}\|\bm{\xi}_{i}\|^{-2}+d\sum\limits_{i_{1}\neq i_{2}}\sigma_{s,r,i_{1}}^{(T)}\sigma_{s,r,i_{2}}^{(T)}\langle\bm{\xi}_{i_{1}},\bm{\xi}_{i_{2}}\rangle\|\bm{\xi}_{i}\|^{-4}$ $\displaystyle\leq 2\sum\limits_{i\in I}\left(\sigma_{s,r,i_{1}}^{(T)}\right)^{2}+\dfrac{2\Delta}{d}\sum\limits_{i_{1}\neq i_{2}}\sigma_{s,r,i_{1}}^{(T)}\sigma_{s,r,i_{2}}^{(T)}$ $\displaystyle\leq\dfrac{8k^{2}\ln^{2}(T+1)}{n}+\dfrac{8k^{2}\ln^{2}(T+1)\Delta}{d}$ $\displaystyle\leq\dfrac{8k^{3}}{n}+\dfrac{8k^{3}\Delta}{d}\leq 16k\epsilon=o(1).$ $$

Combining the above together, we know that

$$ $\displaystyle\sqrt{d}\left\|\bm{w}_{s,r}^{(T)}-\lambda^{(T)}\bm{\mu}_{s}\|\bm{\mu}_{s}\|^{-2}\right\|=o(1).$ $$

∎

Then we prove that the the clean accuracy is nearly perfect(Property 1).

Assume $(\bm{x},y)$ is randomly sampled from the data distribution $\mathcal{D}$. Without loss of generality, we assume that $\bm{x}=\bm{\mu}_{j}+\bm{\xi},y=j$.

###### Lemma C.26 .

Let $\bm{\xi}\sim\mathcal{N}(0,I_{d})$. Then, with probability at least $1-2nd^{-\ln(d)/2}$, for all $s\in J,r\in[h]$ we have

$$ $\displaystyle|\langle\bm{w}_{s,r}^{(T)},\bm{\xi}\rangle|\leq\dfrac{\epsilon}{6}.$ $$

The proof of this lemma is the same as the proof of [Lemma B.26](#A2.Thmtheorem26).

Using the conclusion in Lemma [C.13](#A3.Thmtheorem13) and Lemma [C.26](#A3.Thmtheorem26), we know that for $s\in J,r\in[h],s\neq j$

$$ $\displaystyle\langle\bm{w}_{s,r}^{(T)},\bm{x}\rangle=\langle\bm{w}_{s,r}^{(T)},\bm{\mu}_{j}\rangle+\langle\bm{w}_{s,r}^{(T)},\bm{\xi}\rangle$ $\displaystyle\leq\lambda_{s,r,j}^{(T)}+\dfrac{\epsilon}{6}+\dfrac{\epsilon}{6}\leq\dfrac{\epsilon}{3},$ $\displaystyle\langle\bm{w}_{j,r}^{(T)},\bm{x}\rangle=\langle\bm{w}_{j,r}^{(T)},\bm{\mu}_{j}\rangle+\langle\bm{w}_{j,r}^{(T)},\bm{\xi}\rangle$ $\displaystyle\geq\lambda_{j,r,j}^{(T)}-\dfrac{\epsilon}{6}-\dfrac{\epsilon}{6}\geq\lambda_{j,r,j}^{(T)}-\dfrac{\epsilon}{3}\geq\dfrac{\epsilon}{3}.$ $$

Thus we know that $f_{j}(\bm{x})=\textbf{max}_{s\in J}\{f_{s}(\bm{x})\}.$ So $\bm{F}_{\bm{\theta}^{(T)}}$ has standard perfect accuracy.

Finally, we prove that the corresponding binary classifier achieves optimal robustness(Property 3).

By [Lemma C.25](#A3.Thmtheorem25), we know that $\sqrt{d}\|\bm{w}_{s,r}^{(T)}\|\leq\lambda^{(T)}+o(1)\leq 2\lambda^{(T)}.$

Then for any perturbation $\bm{\rho}$ with $\bm{\rho}\leq\dfrac{\sqrt{d}}{10}$.

We know that

$$ $\displaystyle\langle\bm{w}_{j,r}^{(T)},\bm{x}+\bm{\rho}\rangle$ $\displaystyle=\langle\bm{w}_{j,r}^{(T)},\bm{\mu_{i}}\rangle+\langle\bm{w}_{j,r}^{(T)},\bm{\xi}\rangle+\langle\bm{w}_{j,r}^{(T)},\bm{\rho}\rangle$ $\displaystyle\geq\lambda_{j,r,j}^{(T)}-\dfrac{\epsilon}{6}-\dfrac{\epsilon}{6}-\|\bm{w}_{j,r}^{(T)}\|\|\bm{\rho}\|$ $\displaystyle\geq\dfrac{3\lambda^{(T)}}{4}.$ $$

For $s\in J,s\neq j$, we know that

$$ $\displaystyle\langle\bm{w}_{s,r}^{(T)},\bm{x}+\bm{\rho}\rangle$ $\displaystyle=\langle\bm{w}_{s,r}^{(T)},\bm{\mu_{i}}\rangle+\langle\bm{w}_{s,r}^{(T)},\bm{\xi}\rangle+\langle\bm{w}_{s,r}^{(T)},\bm{\rho}\rangle$ $\displaystyle\leq\lambda_{s,r,j}^{(T)}+\dfrac{\epsilon}{6}+\dfrac{\epsilon}{6}+\|\bm{w}_{s,r}^{(T)}\|\|\bm{\rho}\|$ $\displaystyle\leq\epsilon+\dfrac{\epsilon}{3}+\dfrac{\lambda^{(T)}}{5}$ $\displaystyle\leq\dfrac{3\lambda^{(T)}}{4}-\ln(k).$ $$

Thus we know that $f_{j}(\bm{x}+\bm{\rho)}\geq\dfrac{3\lambda^{(T)}}{4}$ and $f_{s}(\bm{x}+\bm{\rho})\leq\dfrac{3\lambda^{(T)}}{4}-\ln(k).$

let $G(\bm{x})$ denote the numerator of $F_{\bm{\theta}^{(T)}}^{\textit{binary}}(\bm{x})$, where denominator is $\sum_{s\in J}e^{f_{s}(\bm{x})}$. We know

$$ $\operatorname{sgn}(F_{\bm{\theta}^{(T)}}^{\textit{binary}})=\operatorname{sgn}(G).$ $$

Thus we have

$$ $\displaystyle G(\bm{x}+\bm{\rho})$ $\displaystyle=\sum\limits_{j\in J_{+}}\exp\left(f_{j}(\bm{x}+\bm{\rho})\right)-\sum\limits_{j\in J^{-}}\exp\left(f_{j}(\bm{x}+\bm{\rho})\right)$ $\displaystyle\geq\exp\left(3\lambda^{(T)}/4\right)-\sum\limits_{j\in J_{-}}\exp\left(3\lambda^{(T)}/4-\ln(k)\right)$ $\displaystyle\geq 0.$ $$

That is to say $\operatorname{sgn}(G(\bm{x}+\bm{\rho}))=\operatorname{sgn}(G(\bm{x}))$, which means $F_{\bm{\theta}^{(T){}}}^{\textit{binary}}$ is robust under any perturbation with radius smaller than $\dfrac{\sqrt{d}}{10}$.
∎

###### Theorem C.22 (Restatement of Theorem 4.7 ) .

###### Proof of Theorem Theorem C.22 .

###### Lemma C.23 .

###### Proof of Lemma C.23 .

###### Lemma C.24 .

###### Proof.

###### Lemma C.25 .

###### Proof of Lemma B.25 .

###### Lemma C.26 .

## Appendix D Two Feature Learning Regimes: Feature Averaging and Feature Decoupling

In this section, we present two distinct parameter regimes for our two-layer network learner: feature averaging and feature decoupling.
The former means the weights associated with each neuron is a linear average of features, while the latter indicates that distinct features will be learned by separate neurons.
Our construction is similar to that in Frei et al. () and Min and Vidal (). We illustrate how a feature averaging solution leads to non-robustness,
while a feature decoupling solution exists and is more robust (w.r.t. to a much larger robust radius).

### D.1 Feature-Averaging Two-Layer Neural Network

Now, we begin by presenting the following example of a feature-averaging two-layer neural network, which is a more general version (including a bias term) than the one we mentioned in Definition [4.1](#S4.Thmtheorem1).

Feature-Averaging Two-Layer Neural Network. Consider the following two-layer neural network with identical positive neurons and identical negative neurons, which can be simplified as (i.e., we merge identical neurons as one neuron):

$$ $f_{\bm{\theta}_{\textit{avg}}}(\bm{x}):=\underbrace{\operatorname{ReLU}\bigg{(}\bigg{\langle}\sum_{j\in J_{+}}\bm{\mu}_{j},\bm{x}\bigg{\rangle}+b_{+}\bigg{)}}_{\textit{deals with all positive clusters}}-\underbrace{\operatorname{ReLU}\bigg{(}\bigg{\langle}\sum_{j\in J_{-}}\bm{\mu}_{j},\bm{x}\bigg{\rangle}+b_{-}\bigg{)}}_{\textit{deals with all negative clusters}},$ $$

where we choose weight $\bm{w}_{s,r}=\sum_{j\in J_{s}}\bm{\mu}_{j}$ for $s\in\{-1,+1\},r\in[m]$ and bias $b_{s,r}=b_{s}$ for $s\in\{-1,+1\},r\in[m]$.

Indeed, the feature-averaging network uses the first neuron to process all data within positive clusters, and it uses the second neuron to process all data within negative clusters. Thus, it can correctly classify clean data, which is shown as the following proposition.

###### Theorem D.1 .

There exist values of $b_{+}$ and $b_{-}$ such that the feature-averaging network $f_{\bm{\theta}_{\textit{avg}}}$ achieves $1-o(1)$ standard accuracy over $\mathcal{D}$.

###### Proof of Theorem D.1 .

Let $b_{+}=b_{-}=0$, and then we know, for data point $(\bm{x}=\alpha\bm{\mu}_{i}+\bm{\xi},y)\sim\mathcal{D}$ within cluster $i$ (w.l.o.g. we assume cluster $i$ is a positive cluster), with high probability, it holds that

$$ $\displaystyle f_{\bm{\theta}_{\textit{avg}}}(\bm{x})$ $\displaystyle\geq\langle\bm{\mu}_{i},\alpha\bm{\mu}_{i}\rangle+\sum_{j\in J_{+}\setminus\{i\}}\langle\bm{\mu}_{j},\bm{\xi}\rangle-\sum_{j\in J_{-}}\langle\bm{\mu}_{j},\bm{\xi}\rangle$ $\displaystyle\geq\Theta(d)-O(k\Delta)=\Theta(d)-O(k\sigma\sqrt{d}\operatorname{ln}(d))\geq 0,$ $$

which implies that $f_{\bm{\theta}_{\textit{avg}}}$ correctly classifies data $(\bm{x},y)$ with high probability.
∎

However, it fails to robustly classify perturbed data no matter what the bias term is, shown in
the following theorem.

###### Theorem D.2 .

For any values of $b_{+}$ and $b_{-}$ such that $f_{\bm{\theta}_{\textit{avg}}}$ has $1-o(1)$ standard accuracy, it holds that the feature-averaging network $f_{\bm{\theta}_{\textit{avg}}}$ has zero $\delta-$robust accuracy for perturbation radius $\delta=\Omega(\sqrt{d/k})$.

###### Proof of Theorem D.2 .

Indeed, we can choose the adversarial attack as $\bm{\rho}\propto-\sum_{j\in J_{+}}\bm{\mu}_{j}+\sum_{l\in J_{-}}\bm{\mu}_{l}$ and $\|\epsilon\|=\delta$. Then, for averaged features $\bm{w}_{s,r}=\sum_{j\in J_{s}}\bm{\mu}_{j}$, this perturbation can activate almost all of ReLU neurons, which w.h.p. leads a linearization over the perturbation $\bm{\rho}$

$$ $f_{\bm{\theta}_{\textit{avg}}}(\bm{x}+\bm{\rho})=f_{\bm{\theta}_{\textit{avg}}}(\bm{x})+\langle\nabla_{\bm{x}}f_{\bm{\theta}_{\textit{avg}}}(\bm{x}),\bm{\rho}\rangle.$ $$

Since $f_{\bm{\theta}_{\textit{avg}}}$ has $1-o(1)$ standard accuracy, we know that the bias term satisfy that $b_{+},b_{-}=O(d)$, which manifests that the classifier achieves a positive margin, i.e.

$$ $0<yf_{\bm{\theta}_{\textit{avg}}}(\bm{x})\leq O(d),$ $$

w.h.p. over $(\bm{x},y)$ sampled from $\mathcal{D}$.

Then, due to a large gradient norm over data input, i.e.

$$ $\|\nabla_{\bm{x}}f_{\bm{\theta}_{\textit{avg}}}(\bm{x})\|=\|\sum_{j\in J_{+}}\bm{\mu}_{j}-\sum_{l\in J_{-}}\bm{\mu}_{l}\|=\Omega(\sqrt{kd}),$ $$

we derive that the feature-averaging network $f_{\bm{\theta}_{\textit{avg}}}$ has zero $\delta-$robust accuracy for perturbation radius $\delta=\Omega(\sqrt{d/k})$.
∎

###### Theorem D.1 .

###### Proof of Theorem D.1 .

###### Theorem D.2 .

###### Proof of Theorem D.2 .

### D.2 Robust Two-Layer Neural Network Exists

In this section, we show a robust two-layer network exists for
$\mathcal{D}$, using a similar construction in Frei et al. ().

###### Theorem D.3 .

There exists a two-layer network $f_{{\bm{\theta}}_{\textit{dec}}}$ that is $\frac{\sqrt{d}}{3}$-robust for $\mathcal{D}$.

###### Proof of Theorem D.3 .

The construction is similar to that in Frei et al. ().
We define $f_{{\bm{\theta}}_{\textit{dec}}}:\mathbb{R}^{d}\to\mathbb{R}$ is a network that represents a positive constant times the following function:

$$ $f_{{\bm{\theta}}_{\textit{dec}}}({\bm{x}})\propto\sum_{j\in J_{+}}\operatorname{ReLU}\left(\langle{{\bm{\mu}}_{j}},{{\bm{x}}}\rangle-\frac{d}{2}\right)-\sum_{l\in J_{-}}\operatorname{ReLU}\left(\langle{{\bm{\mu}}_{l}},{{\bm{x}}}\rangle-\frac{d}{2}\right),$ $$

In particular, we set a two-layer width-$k$ ReLU network with ${\bm{w}}_{s,j}=\mathbbm{1}\left(j\in J_{s}\right){\bm{\mu}}_{j}$, $b_{s,j}=-\mathbbm{1}\left(j\in J_{s}\right)\frac{d}{2}$ for $s\in\{\pm 1\}$ and $j\in[k]$.

In this network, each neuron $\operatorname{ReLU}(\langle\bm{\mu}_{j},\bm{x}\rangle-\frac{d}{2})$ (or $\operatorname{ReLU}(\langle\bm{\mu}_{l},\bm{x}\rangle-\frac{d}{2})$) deals with one certain positive cluster $j$ (or negative cluster $l$), and we also apply the bias term to filter out intra/inter cluster noise. In this regime, for each data point $(\bm{x},y)$ belonging to cluster $i$ (we assume cluster $i$ is a positive cluster and $y=1$) and any perturbation $\bm{\rho}$ ($\|\bm{\rho}\|\leq\frac{\sqrt{d}}{3}$), we have the following linearization, w.h.p.

$$ $f_{\bm{\theta}_{\textit{dec}}}(\bm{x}+\bm{\rho})=\frac{1}{m}\langle\bm{\mu}_{i},\bm{x}+\bm{\rho}\rangle.$ $$

Then, we know the network $f_{\bm{\theta}_{\textit{dec}}}$ has $1-o(1)$ $\delta-$robust accuracy for $\delta\leq\frac{\sqrt{d}}{3}$.
∎

Note that $f_{{\bm{\theta}}_{\textit{dec}}}$ leverages
individual decoupled features, which
is a natural and robust solution to the binary classification on $\mathcal{D}$.
In fact, one can easily verify that the robustness of
$f_{{\bm{\theta}}_{\textit{dec}}}$
is optimal up to a constant factor,
as the distance between distinct cluster centers is $\Theta(\sqrt{d})$, i.e., $\|\bm{\mu}_{i}-\bm{\mu}_{j}\|=\Theta(\sqrt{d})$, for all $i\neq j$.
However, as we show in our main result that gradient descent does not learn this feature-decoupled network directly from $\mathcal{D}$, and instead converges to a different solution that is $\Theta(\sqrt{k})$ times less robust.

###### Theorem D.3 .

###### Proof of Theorem D.3 .

### D.3 Non-Robust Multi-Class Network Exists

Similar to the feature-averaging binary-class network as that we mentioned in Definition [4.1](#S4.Thmtheorem1), the non-robust multi-class network also exists, which is shown as the following proposition.

###### Theorem D.4 (Restatement of Proposition 4.8 ) .

Consider the following multi-class network $\bm{F}_{\tilde{{\bm{\theta}}}}$: for all $j\in[k]$, the sub-network $f_{j}$ has only single neuron ($h=1$) and is defined as $f_{j}(\bm{x})=\operatorname{ReLU}\left(\left\langle\bm{\mu}_{j}+\sum_{l\in J_{s}}\bm{\mu}_{l},\bm{x}\right\rangle\right)$,
where cluster $j$ has binary label $s\in\{\pm 1\}$.
With probability at least $1-\operatorname{exp}(-\Omega(\operatorname{log}^{2}d))$ over $\tilde{S}$, we have that ${\mathcal{L}}_{\mathrm{CE}}(\tilde{{\bm{\theta}}})\leq\operatorname{exp}(-\Omega(d))=o(1)$, where $\tilde{{\bm{\theta}}}$ denotes the weights of $\bm{F}_{\tilde{{\bm{\theta}}}}$. Moreover, $\mathrm{Acc}_{\mathrm{clean}}^{\mathcal{D}}(F^{\mathrm{binary}}_{\tilde{{\bm{\theta}}}})\geq 1-\exp(-\Omega(\log^{2}d))$, $\mathrm{Acc}_{\mathrm{robust}}^{\mathcal{D}}(F^{\mathrm{binary}}_{\tilde{{\bm{\theta}}}};\Omega(\sqrt{d/k}))\leq\exp(-\Omega(\log^{2}d))$.

###### Proof of Theorem D.4 .

Consider data point $(\bm{x},y)$ that is randomly sampled from the data distribution $\mathcal{D}$. Without loss of generality, we assume that $\bm{x}=\bm{\mu}_{j_{0}}+\bm{\xi},j_{0}\in J_{+}$. Reusing the argument of proof of Property (3) and (4) in Proposition [A.6](#A1.Thmtheorem6). We know that, with probability at least $1-2kd^{-\ln(d)/2}$, for all $j\in J$, we have

$$ $|\langle\bm{\mu}_{j},\bm{\xi}\rangle|\leq\Delta.$ $$

First, we prove the network has perfect clean accuracy when the above properties hold. Indeed, we calculate the output value of each sub-network as follows.

$$ $\displaystyle f_{j_{0}}(\bm{x})=\operatorname{ReLU}\left(\left\langle\bm{\mu}_{j_{0}}+\sum_{l\in J_{+}}\bm{\mu}_{l},\bm{x}\right\rangle\right)\geq 2d-k\Delta.$ $$

For $j\in J_{s},j\neq j_{0}$,

$$ $\displaystyle f_{j}(\bm{x})=\operatorname{ReLU}\left(\left\langle\bm{\mu}_{j}+\sum_{l\in J_{s}}\bm{\mu}_{l},\bm{x}\right\rangle\right)\leq d+k\Delta<2d-k\Delta=f_{j_{0}}(\bm{x}).$ $$

Thus, we know that $\bm{F}_{\tilde{{\bm{\theta}}}}(\bm{x})=j_{0}$ with probability at least $1-2kd^{-\ln(d)/2}$.

Then, with probability at least $1-\operatorname{exp}(-\Omega(\operatorname{log}^{2}d))$ over $\tilde{S}$ sampled from $\mathcal{D}$, for all $i\in I$, we have

$$ $\displaystyle-\operatorname{log}p_{\tilde{y}_{i}}(\bm{x}_{i})$ $\displaystyle=-\operatorname{log}\frac{\exp(f_{\tilde{y}_{i}}(\bm{x}_{i}))}{\sum_{j\in[k]}\exp(f_{j}(\bm{x}_{i}))}$ $\displaystyle\leq-\operatorname{log}(1-\operatorname{exp}(-\Omega(d)))$ $\displaystyle\leq\operatorname{exp}(-\Omega(d)),$ $$

where the last inequality holds due to $\operatorname{log}(1-z)\geq\Omega(z)$ for sufficiently small $z$. Therefore, we derive that

$$ ${\mathcal{L}}_{\mathrm{CE}}({\bm{\theta}}_{\tilde{{\bm{\theta}}}})=\frac{1}{n}\sum_{i=1}^{n}\operatorname{log}p_{\tilde{y}_{i}}(\bm{x}_{i})\leq\operatorname{exp}(-\Omega(d))=o(1).$ $$

Finally, we prove that the network has at most $2kd^{-\ln(d)/2}$ robust test accuracy against perturbation radius $\delta=\Omega(\sqrt{d/k})$.

Consider perturbation $\bm{\rho}=\dfrac{3(1+c)}{k}\left(\sum\limits_{l\in J_{+}}\bm{\mu}_{l}-\sum\limits_{l\in J-}\bm{\mu}_{l}\right)$.

For any $j\in J_{+}$, we know that

$$ $\displaystyle\left\langle\bm{\mu}_{j}+\sum_{l\in J_{+}}\bm{\mu}_{l},\bm{x}-\bm{\rho}\right\rangle$ $\displaystyle=\left\langle\bm{\mu}_{j}+\sum_{l\in J_{+}}\bm{\mu}_{l},\bm{\mu}_{j_{0}}+\bm{\xi}-\bm{\rho}\right\rangle$ $\displaystyle\leq 2d+(k+1)\Delta-\dfrac{k}{1+c}\dfrac{3(1+c)d}{k}$ $\displaystyle<0.$ $$

For any $j\in J_{-}$, we know that

$$ $\displaystyle\left\langle\bm{\mu}_{j}+\sum_{l\in J_{-}}\bm{\mu}_{l},\bm{x}-\bm{\rho}\right\rangle$ $\displaystyle=\left\langle\bm{\mu}_{j}+\sum_{l\in J_{+}}\bm{\mu}_{l},\bm{\mu}_{j_{0}}+\bm{\xi}-\bm{\rho}\right\rangle$ $\displaystyle\geq d+(k+1)\Delta-\dfrac{k}{1+c}+\dfrac{3(1+c)d}{k}$ $\displaystyle>0.$ $$

This is to say for any $j\in J_{+}$,

$$ $\displaystyle f_{j}(\bm{x}-\bm{\rho})$ $\displaystyle=\operatorname{ReLU}\left(\left\langle\bm{\mu}_{j}+\sum_{l\in J_{+}}\bm{\mu}_{l},\bm{x}-\bm{\rho}\right\rangle\right)=0.$ $$

For any $j\in J_{-}$,

$$ $\displaystyle f_{j}(\bm{x}-\bm{\rho})$ $\displaystyle=\operatorname{ReLU}\left(\left\langle\bm{\mu}_{j}+\sum_{l\in J_{-}}\bm{\mu}_{l},\bm{x}-\bm{\rho}\right\rangle\right)>0.$ $$

Thus, we obtain that
$\mathrm{Acc}_{\mathrm{robust}}^{\mathcal{D}}(F^{\mathrm{binary}}_{\mathrm{FA}};\delta)\leq\exp(-\Omega(\log^{2}d))$

∎

###### Theorem D.4 (Restatement of Proposition 4.8 ) .

###### Proof of Theorem D.4 .

## Appendix E Additional Experiments

### E.1 Ablation Study for Synthetic Data Setup

We conducted several additional experiments on synthetic datasets, as an ablation study for choices of hyper-parameters. The goal is to show that feature averaging happens in different settings.

Baseline Setting. We choose the hyper-parameters as $k=10,d=3072,m=5,n=1000,\alpha=\sigma=1,\eta=0.001,\sigma_{w}=\sigma_{b}=0.00001,T=100$. We denote the weights of the two-layer network as $\bm{w}_{1},\bm{w}_{2},\dots,\bm{w}_{10}$ (where the first five weights correspond positive neurons and the other five weights correspond negative neurons). We also set that the first five clusters are positive and the others are negative. Each element in the matrix, located at position $(i,j)$, represents the average cosine value of the angle between the feature vector $\bm{\mu_{i}}$ and the weight vector $\bm{w}_{j}$. The experiment result under baseline setting is presented as Figure [5](#A5.F5) (a), Figure [6](#A5.F6) (b), Figure [7](#A5.F7) (c) and Figure [8](#A5.F8) (b).

Effect of the number of samples. We vary the number of samples as $n=1000,10000,50000$. See results in Figure [5](#A5.F5). It shows that feature-averaging can not be mitigated via more training data.

Effect of the learning rate. We vary the learning rate as $\eta=0.01,0.001,0.0001$. See results in Figure [6](#A5.F6). It shows that the assumption about small learning rate is necessary for feature averaging.

Effect of the initialization. We vary the initialization magnitude as $\sigma_{w}=\sigma_{b}=0.001,0.0001,0.00001$. See results in Figure [7](#A5.F7). It shows that small initialization is necessary for feature averaging.

Effect of the signal-to-noise ratio. We vary the signal-to-noise ratio as $\operatorname{SNR}:=\alpha/\sigma=0.5,1,2$. See results in Figure [8](#A5.F8). It shows that our results can also apply to $\operatorname{SNR}=\Theta(1)$ case.

Effect of the orthogonal condition. We vary the cosine value of the angle between different cluster center features as $\operatorname{cos}(\bm{\mu}_{i},\bm{\mu}_{j})=0.00001,0.001,0.09,\forall i\neq j$. See results in Figure [9](#A5.F9). It shows that the exact orthogonal condition can be relaxed to a nearly orthogonal setting, under which feature averaging still happens.

### E.2 Additional Pre-Trained Model Experiments on CIFAR-10

Here, we verify the orthogonality of the extracted features (i.e., image embeddings of CLIP model) by calculating the correlation between them, as shown in Figure [10](#A5.F10).

Figure: (a) $n=1000$
Refer to caption: /html/2410.10322/assets/x9.png

Figure: (a) $\eta=0.01$
Refer to caption: /html/2410.10322/assets/x12.png

Figure: (a) $\sigma_{w}=\sigma_{b}=0.001$
Refer to caption: /html/2410.10322/assets/x15.png

Figure: (a) $\operatorname{SNR}=0.5$
Refer to caption: /html/2410.10322/assets/x18.png

Figure: (a) $\operatorname{cos}(\bm{\mu}_{i},\bm{\mu}_{j})=0.0001$
Refer to caption: /html/2410.10322/assets/x21.png

Figure: Figure 10: Verifying Orthogonal Condition: We plot the extracted feature correlation as a colormap, where each pixel represents some $\operatorname{cos}(\bm{z}_{i},\bm{z}_{j})$ between two extracted features $\bm{z}_{i},\bm{z}_{j}$ of two data $\bm{x}_{i},\bm{x}_{j}$ from CIFAR-10 training dataset (here, we sample $100$ instances for each class).
Refer to caption: /html/2410.10322/assets/x24.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. The feature average is clearly defined in two-layer networks - the linear input combinations. However, in the deeper layer, such a definition does not seem reasonable. How do we analyse features in this case? Can the authors provide one or more examples, conceptually or mathematically, about how features are averaged in deeper neural networks?

2. In image tasks, reasonable and robust features are often local and isolated. However, there are other scenarios where the features are required to have global or long-term dependency. Such as time series, and natural language. In those cases, it seems like the real interpretable and robust features are some "mixing" of the original data. Therefore:
 a) How should we define the feature mixing in those cases?
 b) Would the theory developed in the paper be effective in those circumstances? If so, how to measure or validate? If not, should there be possible modification can improve the adaptability of the theory?

3. There are researches that have out that, the neural network tends to learn low-frequency features (which are nonlocal and varied smoothly) more effectively than high-frequency features. Some connect such behaviour to the generalization capacity of over-parameterized neural networks. Do you think there is any relation between this theory and the ones developed in this paper? I want to hear about the comments of the authors. 
I list a few papers for reference:
Xu, Zhi-Qin John, et al. "Frequency principle: Fourier analysis sheds light on deep neural networks." arXiv preprint arXiv:1901.06523 (2019).
Xu, Zhi-Qin John, Yaoyu Zhang, and Tao Luo. "Overview frequency principle/spectral bias in deep learning." Communications on Applied Mathematics and Computation (2024): 1-38.
soundness: 3
strengths: The paper gives very thorough theoretical study of the problem. The hypothesis, assumptions and derivation are reasonable and sound. The results matched with the intuition well.
summary: This paper provides a theoretical analysis of the relation between feature dependency and adversarial robustness. The authors studied a two-layer ReLU network on the simulated data and derived that the gradient descent tends to learn averaged features instead of isolated/robust ones. The results were proved thoroughly. Then, the paper provides a test on the derived theory on real-world datasets, showing that by incorporating more precise training labels, the model can learn isolated features rather than mixed ones, therefore achieving better adversarial robustness.
weaknesses: The theory developed in this paper is phenomenological. It can be further improved if systematic and quantifying criteria can be built to design an optimization method that can help improve the robustness.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 4
questions: -Minor issues in notations: Big $O$ notation on line 185 is different from the others; $Theta$ notation is not defined.

-Why does strong correlations indicate the existence of feature averaging? Is there a rigorous proof? How do you define the threshold of having a strong correlation?

-Proposition 4.8 is used as an example of “adding more fine-grained supervision signals does not trivially lead to decoupled features and robustness”, but it does not follow the setup of your analysis, since $h=1$. How is this multi-class network obtained? Why is its parameter assumed to be the combination of $mu_s$?
soundness: 3
strengths: -This paper focuses on explaining the adversarial robustness of neural networks in classification problem, which of course is an important topic in ML/DL community.

-This paper is generally well-written and logically organized with clear notations, assumptions, and detailed references to prior works with comparisons

-This paper’s idea is intuitive: this paper characterizes the “approximate linearity” in networks under NTK assumptions as adversarial non-robustness in terms of Feature averaging.
summary: This paper analyzed the training dynamics of gradient descent on a simplified two-layer ReLU networks in a binary supervised classification task, where the feature vector is sample from a gaussian mixture model with equal probability for K>2 components and the label are masked as 2 categories. The main result in section 4 showed that, under some regulatory assumptions and choices of hyper-parameters, the learned network turns to a feature averaging network, which has low generalization error but is not robust to perturbation of dataset.
weaknesses: - Too many strong and unrealistic assumptions are used. First, the number of neurons $M=2m$ assumed to be finite as the number of clusters of features in assumption 4.3. It is against the purpose of using the neural tangent kernel (NTK) theory to approximate infinite wide neural networks. Thus, the result in this paper is limited since “training deep neural network is a highly non-convex and over-parametrized optimization problem”. Second, as the solution to adversarial robustness, “fine-grained supervision” suggests including the original cluster labels for all data points, which is extremely unrealistic, because the response is simply assigned from the cluster labels. It is not surprising that there is an improvement in adverbial robustness as the exact latent information is provided. Feature decoupling is a simple result as the network can learn each cluster’s mean given the cluster labels. 

- The evidence of feature averaging is shown Figure 2 as the strong correlations/cosine value between cluster’s mean and network weights, but the scale of correlations is not consistent. In particular, the results from CIFAR-10 are not convincing, where the average cosine value can be lower than 0.1.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. In the fine-grained supervision case of the multi-cluster data, I understand that the binary classification using the multi-class network results in a robust solution. But, I wonder if the learned network would still be non-robust for the multi-class classification. 
2. It is not an ask to prove the following. I am only curious to understand and know the authors thought. The fine-grained supervision setting can also be seen as learning a more expressive/complex classifier (learn k classes to do binary classification). Then, can the authors comment on the adversarial min-max training? Would similar feature decoupling, ie the weight of each neuron is aligned with one cluster feature, happen naturally if one does adversarial training of the binary class data?
soundness: 3
strengths: * The theoretical analysis of the finite time gradient descent dynamics is thorough in the feature learning regime. 
* This paper is an interesting contribution in understanding the robustness of neural networks from optimization perspective. The analysis also aided in proving a conjecture. 
* The paper is well written and the presentation is clear.
summary: This paper studies the implicit bias of gradient descent, which the authors call 'Feature Averaging' by theoretically analyzing two layer ReLU network with only a learnable first layer, under a specific data distribution - binary data having multiple clusters with orthogonal cluster centers. 
Under gradient descent, weights of neurons in the hidden layer converge to an average of the cluster centers, resulting in a non-robust network. Additionally, they also prove that with more fine-grained cluster information, the network can learn a robust solution for the binary classification problem.
weaknesses: * The assumptions on the data (orthonormal equinorm and multi-cluster) is restrictive.
* The proof is detailed and thorough for two layer ReLU with only one learnable hidden layer. Extending to deeper networks seems challenging.

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 3
questions: Q1: Can the fact that the network reaches an "averaging" solution be deduced from the KKT equations in Frei et al. (2022)?

Q2: The implicit bias of GD for this problem in 2-layer neural networks is undesirable for $\ell_2$ robustness. What about $\ell_\infty$ perturbations? It seems that this bias would help for $\ell_\infty$ perturbations, in contrast to the feature decoupling solution $NN_{FD}$. I would appreciate your thoughts on this.

Q3: Relatedly, would it be possible to measure robustness in the experiments on CIFAR10/MNIST against $\ell_\infty$ perturbations? If there is evidence for larger $\ell_\infty$ robustness when training with additional supervision, then this would make the authors' claims much stronger.

Q4: Figure 3: How do you measure perturbation radius? If, for example, you normalise MNIST images to be between 0 and 1, then a perturbation radius of 2.0 does not make any sense. Can you mention the ranges of the input for MNIST and CIFAR10?

Q5: In line 882, you mention that "In image classification tasks, Ilyas et al. (2019) visualized both robust and nonrobust features.". However, this is not true to the best of my knowledge. I believe there are 2 works that have provided such visualizations [1, 2].

[1]. G. Goh. A discussion of ’adversarial examples are not bugs, they are features’: Two examples of useful, non-robust features. Distill, 2019.
[2] Tsilivis, N. and Kempe, J. (2022). What can the neural tangent kernel tell us about adversarial robustness? Advances in Neural Information Processing Systems 35.
soundness: 4
strengths: This is a good paper, and it was a pleasure to read and review. I would like to highlight the following contributions:

- Important topic: The interplay between the implicit bias of optimization and robustness is, in my opinion, an important area of study.
- Finite time theoretical analysis of neural networks trained in the various settings: The paper rigorously analyzes the training trajectory of the networks. I enjoyed reading and learning about interesting techniques, such as the ones outlined in Section C.2. In general, the Appendix of the paper is exceptionally polished and the results are easy to follow and verify (to the extent that this was possible during the review process - the paper is 67pages long).
- Interesting suggestion on providing additional supervision during training for improved robustness: I found this observation interesting and the experimental results seem to suggest that these ideas may generalize beyond the toy distribution considered in the paper's analysis.
summary: This paper theoretically studies the robustness of two layer neural networks under a specific data distribution that consists of $k$ clusters in $\mathbb{R}^d$ and the target consists of deciding on which group of clusters the input belongs to (binary classification). The authors show that these networks tend to learn an "averaging" solution (i.e., the first-layer weights encode the averages of the clusters), making the networks susceptible to $\Omega(\sqrt{d / k})$ perturbations of the input (despite the existence of networks capable of tolerating $O(\sqrt{d})$ perturbations). Motivated by this observation, they propose modifying the training process by using each example's cluster vector as the label, and they prove that this approach indeed yields optimally robust networks. Finally, the paper presents experiments with both synthetic data and image classification benchmarks that contextualize the theoretical findings.
weaknesses: 1) The main conceptual observation does not appear to be new. To my understanding, Frei et al. (2022) showed that neural networks trained with gradient descent on this distribution are not robust to $\ell_2$ perturbations of the input due to the implicit bias of optimization. In lines 153-171, the authors argue that their work differs from prior work because they do not explicitly analyze the properties of the KKT point of the max-margin problem and instead perform a finite-time analysis. Furthermore, they argue that convergence to the KKT points might be slow. However, we know from Liu & Li (2020) that these conditions are approximately satisfied as long as there is a positive margin. Thus, approximate results are conceptually consistent with the KKT analysis. Additionally, I believe that the fact that the network reaches an "averaging" solution can be deduced from the KKT equations in Frei et al. (2022)—see also (Q1).

2) The paper argues that "feature averaging is a principal factor contributing to non-robustness of neural networks" (line 014). However, it focuses on one type of perturbations. While the authors provide evidence in this direction for $\ell_2$ perturbations in a simple distribution (& experimental results in more complicated settings),  it is not clear that this holds more broadly. Note that the model’s solution is exceptionally contrived, as all neurons (within each group) converge to the same solution. See also (Q2, Q3).

2) While the paper is well-written, I found the discussion on prior work somewhat insufficient. I would appreciate a bit more clarity on what is novel about the first part of the contributions (i.e., the proof that networks reach feature-averaging solutions). Furthermore, in line 110, the authors mention the work of Min & Vidal (2024) and their conjecture for the first time. It would be more appropriate to introduce this work and its conjecture earlier in the introduction, making it easier for readers to appreciate the contribution.

4) A list of non-significant typos/grammatical errors I spotted:
- line 091: features-> feature
- line 107: exists-> exist
- line 167: analyze -> analysis
- line 261: neural -> neuron
- line 893: echos -> echoes
- line 1098: We -> we
- line 1100: margin -> margins
- lines 1307-1310 require rephrasing (they do not make sense),

Meta-review: This paper studies the convergence of gradient descent for a two-layer ReLU network on a data distribution characterized by localized and separated clusters, with binary labels. In particular, the authors show that GD provides a solution, for appropriate assumptions, where the network weights converge to the mean of the label-wise clusters, in a phenomenon they term "feature averaging". While this observation (even in this synthetic setting) is not new (including being presented as a conjecture in a recent work) this work constitutes the first proof of this behavior. Consequently, as pointed out by other works, this solution is susceptible to adversarial perturbations, and thus non-robust. In turn, the paper shows that extra granular supervision alleviates this issue, converging to a solution with improved robustness (and no feature-averaging). The theory is supported by experimental results. 

The contribution of this paper is important, formally showing that GD can have an implicit bias that conveys a lack of adversarial robustness. Some weaknesses were noted by the reviewers, concerning mostly clarification on the connection to other closely related work, and on the technicalities of their setup.

In all, the reviewers have a consensus that this is a worthy publication, and I concur.

Decision: Accept (Poster)