Feature Averaging: An Implicit Bias of Gradient Descent Leading to Non-Robustness in Neural Networks
Authors: Binghui Li
Zhixuan Pan
Kaifeng Lyu
Jian Li
Venue: ICLR 2025 Poster (2025)
Keywords: deep learning theory
feature learning
adversarial robustness
implicit bias
Primary area: learning theory
TL;DR: In this paper, we provide a theoretical explanation for non-robustness in trained neural networks.

Abstract: In this work, we investigate a particular implicit bias in gradient descent training, which we term “Feature Averaging,” and argue that it is one of the principal factors contributing to the non-robustness of deep neural networks. We show that, even when multiple discriminative features are present in the input data, neural networks trained by gradient descent tend to rely on an average (or a certain combination) of these features for classification, rather than distinguishing and leveraging each feature individually. Specifically, we provide a detailed theoretical analysis of the training dynamics of two-layer ReLU networks on a binary classification task, where the data distribution consists of multiple clusters with mutually orthogonal centers. We rigorously prove that gradient descent biases the network towards feature averaging, where the weights of each hidden neuron represent an average of the cluster centers (each corresponding to a distinct feature), thereby making the network vulnerable to input perturbations aligned with the negative direction of the averaged features. On the positive side, we demonstrate that this vulnerability can be mitigated through more granular supervision. In particular, we prove that a two-layer ReLU network can achieve optimal robustness when trained to classify individual features rather than merely the original binary classes. Finally, we validate our theoretical findings with experiments on synthetic datasets, MNIST, and CIFAR-10, and confirm the prevalence of feature averaging and its impact on adversarial robustness. We hope these theoretical and empirical insights deepen the understanding of how gradient descent shapes feature learning and adversarial robustness, and how more detailed supervision can enhance robustness.

Published as a conference paper at ICLR 2025
FEATURE AVERAGING:
AN IMPLICIT BIAS OF GRADIENT DESCENT LEADING TO NON-ROBUSTNESS IN NEURAL NETWORKS
Binghui Li Zhixuan Pan Kaifeng Lyu² Jian Li
CCenter for Machine Learning Research, Peking University
Institute for Interdisciplinary Information Sciences, Tsinghua University "Simons Institute, UC Berkeley
i 1s. Libinghui@pkir edu..d... en, pax2Demails.tssnhhhaa on kai wwweeeelleeereeeeeeee.....ee edu, lapordge@gmail.com
ABSTRACT
In this work, we investigate a particular implicit bias in gradient descent training, which we term "Feature Averaging," and argue that is one of the principal factors contributing to the non-robustness of deep neural networks. We show that, even when multiple discriminative features are present in the input data, neural networks trained by gradient descent tend to rely of these features for classification, rather than distinguishing and leveraging each feature individually. Specifically, the training dynamics of two-layer ReLU networks on a binary classification task, where the data distribution consists of multiple clusters with mutually orthogonal centers. We rigorously prove that gradient descent biases the network towards feature averaging, where the weights of each hidden neuron represent average of the cluster centers (each corresponding to distinct feature), thereby making the network vulnerable input perturbations aligned with the negative direction of the averaged features. On the positive side, we demonstrate that this vulnerability can be mitigated through more granular supervision. two-layer ReI.U network can achieve optimal robustness when trained to classify individual features rather than merely the original binary classes. Finally, we validate our theoretical findings with experiments synthetic datascts, MNIST, and CIFAR-10, and confirm the prevalence feature averaging and its impact on adversarial robustness. We hope these theoretical and empirical insights deepen the understanding of how gradient descent shapes feature learning and adversarial robustness, and how more detailed supervision can enhance robustness.
average (or certain combination)
provide detailed theoretical analysis of
particular, we prove that a
INTRODUCTION
Deep learning has achieved unprecedented success across a wide range of application domains, includ- ing many safety-critics systems such as autonomous driving and diagnostic assistance technologies. Despite these successes, landmark study by Szegedy et al. (2013) 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., 2013; Szegedy et al., 2013; Goodfellow et al., 2014; Madry et al., 2018) and defenses (Goodfellow al., 2014; Madry et al., 2018; Shafahi et al., 2019; Pang et al., 2022) were developed, but the issue of adversarial robustness is still far from being resolved.
Previous attempts to explain the adversarial robustness of neural networks have been made from various thcoretical perspectives. Daniely & Shacham (2020); Bubeck et al. (2021a); Bartlett et al.
"Equal contribution. alphabet ordering. Corresponding author.

---

Published as a conference paper at ICLR 2025
(2021); Montanari & Wu (2023) proved the existence of adversarial examples for neural networks with random weights across various architectures. Tsipras et al. (2019); Zhang et al. (2019) analyzed the fundamental trade-off between robustness and accuracy. Bubeck et al. (2021b); Bubeck & Sellke (2021); Li et al. (2022a); Li & Li (2023) proved that having a large model size is necessary for achieving robustness in many settings. Ilyas et al. (2019); Tsilivis & Kempe (2022); Kumano et al. (2024); Li & Li (2024) studied the relationship between adversarial examples and the presence of non-robust but predictive features in the data distribution.
A related line of work in deep learning theory studies the implicit bias of gradient descent to explain why neural networks generalize so well. Training deep neural networks highly non-convex and over-param peramarrmetrize optimization problem, in which there are many solutions that fit the training data correctly. Recent studies suggest that, without explicit regularization, gradient descent seems to implicitly bias towards solutions that enjoy favorable properties, particularly good generalization. Hence, characterizing various implicit biases in favor of better generalization has been extensively studied in recent years (Gunasckar et al., 2017: Soudry et al., 2018; Arora et al., 2019b; Lyu & Li, 2020; Blanc et al., 2020). However, good generalization properties do not necessarily imply good robustness with respect to inputs. Indeed, even to adversarial examples. In fact, recent studies by Vardi et al. (2022) and Frei et al. (2024) proved that the implicit bias of gradient descent be "double-edged sword", in the sense that it leads to generalizable solutions with perfect clean accuracy, but being non-robust (susceptible to small adversarial l2 2--peruubbttions), even though there exist robust networks with perfect robust accuracy. Under a similar data setup, Min & Vidal (2024) further conjectured that the weight vectors of two-layer ReLU network trained by gradient flow converge to an average of the cluster centers.
neural networks are vulnerable
In this paper, we perform detailed analysis of the training dynamics of gradient descent on two-layer ReLU networks (under data distributions similar to Vardi et al. (2022), Frei et al. (2024) and Min & Vidal (2024), and the detailed discussion about the comnection between their works and our paper is deferred to Section 2), and rigorously prove that the learned weights exhibit a particular implicit bias, which we term feature averaging. In our setting, feature averaging refers to a particularly simple form: the network trained by gradient descent tends to learn the average of useful features, in the sense that the weight vector associated with each hidden-layer neuron is a weighted average of feature vectors. This also resolves the conjecture made by Min & Vidal (2024). Further, one can easily show that such average is more susceptible small adversarial perturbations than individual features, rendering the learned solution non-robust.
In our experiments, we empirically observe similar phenomena in several other settings. We argue that feature averaging is a key factor contributing to the non-robustness of deep neural networks, and demonstrate its close relationship with several known phenomena and theoretical models in adversarial robustness research. These include the observation that neural networks tend to leverage both robust and non-robust features for classification (Tsipras et al., 2019; Ilyas et al., 2019; Allen-Zhu & Li, 2022; Tsilivis & Kempe, 2022; Li & Li, 2024), the connection between model Lipschitzness (smoothness) and over pver-parameterization (Bubeck et al., 2021b; Bubeck & Sellke, 2021; Li al., 2022a; Li & Li, 2023), the simplicity bias of gradient descent that leads to non-robustness (Shah et al., 2020; Lyu et al., 2021), and the dimpled manifold hypothesis (Shamir et al., 2021). Beyond the linear average behavior studied in this work, we conjecture that feature averaging may appear in more complex forms in real-world settings. For example, the neural network may tend to combine many localized, semantically meaningful (hence more robust (Ilyas et al., 2019; Tsilivis & Kempe, 2022)) features into one discriminative but non-robust feature.
In light of the feature averaging phenomena, propose to enhance the robustness by learning individual features. In particular, we explore a natural and simple, yet less explored method in the study of adversarial robustness, which is provide more granular supervised information related to individual features and force the model to learn the individual features. Theoretically, we prove that training a two-layer ReLU network with feature-level labels leads binary classifier with optimal robustness. Empirically, we design several experiments, using synthetic and real datasets, and the experimental results demonstrate that feature-level supervised information can be very effective in enhancing the robustness of the model (even with standard training). These results are consistent with the empirical findings by Sitawarin et al. (2022): Li et al. (2024), which showed that incorporating fine-grained annotations, such as part-level segmentation, can substantially enhance the adversarial robustness of object recognition systems. See Appendix A for a more detailed discussion of these
2

---

Published as a conference paper at ICLR 2025
connections, including the relationship of feature averaging existing robustness phenomena, and how more granular supervision may improve adversarial robustness.
Our technical contributions can be summarized follows:
1. (Section 4.1) Under certain multi-cluste data distributions (similar to that in Frei et al. (2024)), we prove that two-layer ReLU networks trained by gradient descent converge to feature-averaging solutions (Theorem 4.5). In particular, we show that the weight vector associated with each hidden-layer neuron converges to the average of cluster-center features and the crature-averagiin solution is non-robust w.n.. the radius (V/k), while there exist solutions with optimal robust radius O(vd) (where d is the data dimension, k is the number of clusters, and the existence of such optimal robust solutions is shown in Theorem G.3). This result also solves the conjecture of Min & Vidal (2024) under our settings (Theorem 4.6).
2. (Section 4.2) We show that if the model is provided with the feature level labels (in fact a multi-class classification problem in our multi-cluster data distribution setting), two-layer network can learn the individual features, which furthermore can induce a robust model with optimal robust radius O(vd) (Theorem 4.7).
3. (Section 5) 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 averaged features. Our experiments also demonstrate enhanced robustness through the incorporation of fine-grained supervisory information.
RELATED WORK
Implicit Bias of Gradient Descent. The implicit bias of gradient descent has been studied from various perspectives. The most prominent line of works establishes an equivalence between neural networks in certain training regimes to kernel regression with Neural Tangent Kernel (NTK) (Du al., 20;9b;a; Allen-Zhu et al., 2019a; Zou et al., 2020; Chizat et al., 2019; Arora et al., 2019b; Ji & Telgarsky, 2020b; Cao & Gu, 2019), 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., 2018; 8; Nacson et al., 2019; Lyu & Li, 2020; & Telgarsky, 2020a). parameter norm minimization (Gunasekar al., 2017: 2018: Arora et al.. 2019a) and sharpness reduction (Blanc et al., 2020; Damian et al., 2021; HaoChen et al., 2021; Li al., 2022b; Lyu et al., 2022: Gu et al., 2023). All these works focus on implicit biases that may lead to good generalization except that Vardi et al. (2022) and Frei et al. (2024) connected the 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 & Li, 2021; Allen-Zhu & Li, 2022: Chen et al., 2022; Cao et al., 2022; Zhou et al., 2022; Chidambaram et al., 2023; Allen-Zhu & Li, 2023; Kou et al., 2023a; Simsek et al., 2023) 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., 2018; Du et al., 2019b;a; Allen-Zhu et al., 2019 9019b; Arora et al., 2019b). Among these feature learning works, there exist various data assumptions about feature-noise structure. Based on the data assumption of sparse coding model, Wen & Li (2021) study feature learning process of self-supervised contrastive learning, and Allen-Zhu & Li (2022) propose a principle called feature purification to explain the workings of adversarial training. Allen-Zhu & Li (2023) utilize multi-view-buat patch-structured data assumption to understand the benefits of ensembles in deep lcarning. Following the multi-view data proposed in Allen-Zhu & Li (2023), Chidambaram et al. (2023) show that data mix-up algorithm can provably learn diverse features to improve generalization. Cao et al. (2022); Kou et al. (2023a) explore the benign overfitting phenomenon two-layer convolutional neural networks by leveraging a technique of signal- -oise decomposition. Zhou et al. (2022) study 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 al. (2023) focus on the regression setting and study the compression of the teacher network, and they find that weight vectors, whether copying an individual teacher vector averaging set of teacher vectors, are critical points of the loss function.
3

---

Published as a conference paper at ICLR 2025
Comparisons with Vardi et al. (2022), Frei et al. (2024) and Min & Vidal (2024). Recently, Vardi et al. (2022) and Frei et al. (2024) demonstrated that for two-layer ReLU networks, any KKT solution to the maximum margin program (it is known that gradient flow converges to such KKT solution (Lyu & Li, 2020; Ji & Telgarsky, 2020a)) leads to non-robust solutions under the assumption of synthetic cluster data, and Min & Vidal (2024) further conjectured that the weight vectors of two-layer ReLU network converge to an average of cluster-center vectors. Their finding highlights the significance of the optimization process in the Jobuutobus of neural networks. Our theoretical results are inspired by theirs, but differ from theirs in the following important aspects: (1) Conceptually, feature averaging 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 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 mil 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., (1(1/ log(t)) proven in Soudry et al. (2018); Lyu & Li (2020); Kou et al. (2023b))- (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 (an assumption made in (Lyu & Li, 2020) for achieving KKT points). (4) Our result (Theorem 4.5) solves the conjecture proposed by Min & Vidal (2024), where we show that the weight vector associated with each neuron aligns with weighted average of cluster features, and the ratio between weights of distinct clusters is close to 1.
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., x (Zi...., T3). For x Rd we denote by ||>||| the Euclidean (&2) norm. We denote by 1()) the standard indicator function. We denote sgn (2) = 1f 0 and -1 otherwise. For integer n 2 1, we denote [n] = f1,....,7) and variance a², and by Ni, Σ) the
We
denote by N (u,0²) the normal distribution with mean μ multivariate normal distribution with mean vector μ and covariance matrix E. The identity matrix of size is denoted by Ia. We use (n) to denote the uniform distribution on the support set A. We use standard asymptotic notation O(-) and 22-) to hide constant factors, and Ó(-), ②(()) to hide logarithmic factors.
3.1 DATA DISTRIBUTION
Following Vardi et al. (2022); Frei et al. (2024), we consider binary classification on the following data distribution with multiple clusters.
Definition 3.1 (Multi-Cluster Data Distribution). Given k vectors ....... μa E Rd, called the cluster features, and a partition of [k] into two disjoint sets J+ (J+,J_), we define data distribution on Rd x {1,,1}, where each data point (xx y) is generated as follows:
as
1. Draw cluster index as Unif(
2. Set y if E J4: otherwise j J_ and set y -1;
3. Draw x μ, + E, where & Ia).
For convenience, we write DD instead of (((11)))1: if (1g (1 and Jr are clear from the context. For [+1}, we write J, denote J if = +1 and J_ if 1.
To ease the analysis, we make the following simplifying assumptions on the distribution. Assumption 3.2 (Orthogonal Equinorm Cluster Features). The cluster features (M11) satisfy the properties that (1) || = vd for all j E [k]; and (2) μ4 1 μ, for all <i< < k.
Assumption 3.3 (Nearly Balanced Classification). The partition J= satisfies c < some absolute constant c 1.
[4
< c for
[_
4

---

Published as a conference paper at ICLR 2025
Our data distribution is similar to that in Vardi et al. (2022) and Frei al. (2024). In particular, Vardi et al. (2022) consider a setting where data are comprised of k nearly orthogonal data points in IR This assumption is further relaxed in Frei et al. (2024), where they assume & clusters with nearly ((.))))) orthogonal cluster means { (i.e., they have that || - (+) holds for + j). For simplicity, our work focuses on the setting with clusters exactly orthogonal to each other.
3.2 NEURAL NETWORK LEARNER
training dataset S = (10 y1) F- C C (-1,1) of size n is randomly sampled from the data distribution D([ J+) and used to train two-layer neural network.
Network Architecture. We focus on learning two-layer ReLU networks. Such networks are usually ((w), defined as fo Σド= و;=1 aj ReLU(((,, 2.) + bj), where 0 == {يك}حم {uny } {hy }.a are
the parameters of the network, and ReLU(-)) is the ReLU activation function defined as ReLU(=) max(0, z).
For the sake of simplicity, we consider the case where M 2m is even and fix the second layer as aj = 1 for S S m and a, = for m - 1 1 j S 2m, which is a widely adopted setting in the literature of feature learning theory (Allen-Zhu & Li, 2022; Cao et al., 2022; Kou al., 2023a). With this simplification, we focus on training only the first layer (1w,) 1==1: (b, } and rewrite the network as
)
fa(x) =
Σ
ReLU((W+1.r) b41,r) 1
Σ
ReLU( (w-1,,,)) + b-
(E(m)
rE(m]
where 0 = ((W+1, (b+1-) -- -.--1-=-=1 (b_1.- ) are the trainable parameters, and 1+1,r and 1, correspond to the neurons with a, = m while w 1.r and b_1,- correspond to the neurons with
Training Objective and Gradient Descent. The neural network fe(-) is trained to minimize the following empirical loss on the training dataset S: L(0) == Σi= E (v: (æ:)), where E(g) == log(1 + e=)) is the logistic loss. We apply gradient descent to minimize this loss:
((+1)) ( LO
where (()) denotes the parameters t-th iteration for all 0, and n is the learning rate. We specify the derivative of ReLU activation as ReLU'(z) = [(z 2 0) in backpropagation. At
initialization, we set w.! (O, of Ia) and 600 NO0 of. for some dw, σb > 0.
Clean Accuracy and Robust Accuracy. For given data distribution D over Rd (-1, 1), the clean accuracy of neural network fo IRd + R on D is defined as
Acc (fe) = ((rwy)> (e (x)) = vl
In this work, we focus on the ezrrouussssss. The E2 d-robust accuracy of fo on is defined as Acc (Se; 8) = IP (( (V' BB sgn (So (a: + p)) = y), where Bş := (p Rd || < 5} is the ball centered at the origin with radius 8. We say that neural network fe is 8-robust if Pcc (fo; 8) 2 1 c(d) for some function c(d) that vanishes to zero, i.e., e(d) + as oc.
In very similar setting to ours, Frei et al. (2024) show that there exists two-layer ReLU network that can achieve nearly 100% clean accuracy and Qi -rouustuust accuracy on their data distribution. In our setting, we can also construct a similar network that achieves nearly 100% clean accuracy and R ((vd))robus accuracy. neuron to capture feature/cluster (i.e.. the neural is activated only if the input point is from the corresponding cluster). See Theorem G.3 in Appendix G.2 for the details and Figure for an illustration. However, we will soon show that, despite such l( -(ouusro network exists, gradient descent is incapable of learning such robust network, but instead converges very different solution with robust radius that ((Vk) times smaller.
Robust Networks Exist.
particular, such network utilizes one hidden
5

---

Published as a conference paper at ICLR 2025
14 = (41 + H2 + H3)/3
J (1,2,3,/= (4,5), 5,
= (1+ 142 + H33/3,
+
H. = = (14
125)
0
FFA 3ReLU ((,,,))) 2 ((4_
Pra
fro: f1 f2 f3- - f/ fs. fj(x)= RetU (1x + bj).i E [5]
=- ( ps)/2
and feature-decoupling: We consider dataset with clusters. The first three clusters belong to J4, and the other two to J_ Denote μt == (M) + p22 + pa)/3, pa- == (1u + / For ease of illustration, we assume that Σ 14 = 0. The feature-averaging classifier fFA leverages two neurons with averaged features to classify all data, which corresponds to a linear classifier (the gray line). The feature-decoupling classifier fFD leverages individual features and has more complex polyhedral decision boundary (green lines). Note that the instance is high dimensional and this is only schematic illustration. The distance between data points and the decision boundary of fFD (green lines) is much larger than that of fEA (gray line), which implies that the feature- -ectuul-e classifier is more robust than the reature-averagin one.
Figure 1: Schematic illustration of feature-average
MAIN RESULTS
In this section, present our main technical results. In Section 4.1, we first present the main result (Theorem 4.5) regarding feature averaging, that is standard gradient descent training finds feature averaging solutions for the data distribution D and such feature averaging solution is non- robust. In Section 4.2, 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 similar two layer multi-class network. Consequently, we can obtain binary classification network with optimal robust perturbation radius (Theorem 4.7).
4.1 NETWORK LEARNER PROVABLY LEARNS PEATURE-AVERAGIIN SOLUTION
The prior work by Frei et al. (2024) has showed that, under certain conditions, training two-layer ReI.U network for infinite time converges to a network that can achieve nearly 100% clean accuracy on 'D but is only o -robuust. A subsequent work by Min & Vidal (2024) conjectured that the network converges to specific form of solution, which we refer to the feature-averaging network. Definition 4.1 (Feature Averaiig Network). We define (ax) as the following function:
fFA(x) |J+] ReLU ((44, x)) -J-1 ReLU ((,,z)),
where p J.. Ejer, HÚ the average of cluster centers the positive class, and similarly ___ [5_1 Ejer_ Hj is that for the negative class. We say that two-layer ReLU network fo(x) is a feature-averaginn network if fo(æ) = C fFX(#) for some Remark 4.2. The cevtare-averaging network uses the first neuron 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 S2 /d/k); in particular, consider the attack vector P that aligns with the negative direction of the averaged features, i.e., P & -EEJ, H2 + ESEJ_ Hÿ. One can easily check that with || = 5 22ldd/k, the attack is successful, i.e., sgn(fra( + pp) + (on(fxx(x)) due to the linearity of fvx(x + p) over p. See Appendix G.I for the details, and see Figure 1 for illustration.
0.
Our first main result a asympooti analysis of the training dynamics that explicitly characterizes the solution learned by gradient descent on distribution D after finite number of iterations. For theoretical analysis, we make the following assumptions about the hyper-parameters.
Assumption 4.3 (Choices of
We assume that:
d = S2k" 19)
(1)
[Ω((k)), exp(OOloo
(Ik)
1 = (
of = =(nkk-))

---

Published as a conference paper at ICLR 2025
Remark 4.4 (Discussion of Hyper-Parameee Choices). We make specific choices of hyper- papper-arrnneerss for the sake of calculations, and we emphasize that these may not be the tightest possible choices. In need the data dimension d to be significantly larger than the number of clusters k to ensure all ls cluster features are orthogonal within TR. We further require that the number of samples is a large polynomial of k to ensure that the network can learn all k. cluster features. We assume the learning rate and the initialization magnitude ow, Tb are sufficiently small, which helps the network to be trained in the feature learning regime (Lyu et al., 2021; Cao al., 2022; Allen-Zhu & Li, 2023; Kou et al., 2023a).
particular,
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-vveragin network (defined in Definition 4.1).
Theorem 4.5, In the setting of training a two-layer ReLU network on the binary classification problem P(1u;)==11J1 as described in Section 3, under Assumptions 3.2, 3.3 and for some =(1), after 1 <T S exp(((k)/2)) iterations, with probability at least 1 - "y the neural network satisfies the following properties:
1. The clean accuracy is nearly perfect: Acc (forr) 2 exx(eeexp(-(og d)).
2. Gradient descent leads the network to the -everagnav regime: there exists time- variant coefficient X(T) E [9((1,, +00) such that for all E {+1}, r E [m], the weight
vector why (T) can be approximated as
w (T) (22) | 313
S 0(d-2/22
JEJ.
and the bias terms sufficiently small, i.e.,
b5.7 S o(1.
3. Consequently, the network is non-robust: for perturbation radius 5 = 52( 22(/d/k). the 8-robust accuracy is nearly zero, i.e., ( S exp(-=Q(log* d)).
We provide a proof sketch for Theorem 4.5 in Appendix C (see the full proof in Appendix E.2). Theorem 4.5 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 [4+, and that associated with negative neuron the average of negative cluster features μ4. Moreover, the above feature-averaging property of learned network implies non-robustness, i.e.,
the learned network is only of / d/k) -roustobuust although an ((vd) dd-rouusttt solution exists as we proved in Section
As corollary of Theorem 4.5, we resolve the conjecture proposed in Min & Vidal (2024) in our setting.
Theorem 4.6 (Conjecture from Min & Vidal (2024)). In the setting Theorem 4.5, we have that fermy (z) (C)
infoso sup sup (CfFa(x)
o(1), where frA(#) the feature-averaging
nerwork (Defimition 4.1).
Under a similar orthogonal cluster data assumption, Min & Vidal (2024) 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. The- orem 4.6 provides rigorous proof for the conjecture, although the original conjecture stated under slightly different setting from ours. In their setting, the second layer of the network is also trainable, but we the second layer for simplicity. We also require certain assumptions on the hyperparameters, which has been discussed in details in Assumption 4.3 and Remark 4.4.
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, natural question is: if we provided more fine- grained feature-level supervision, can gradient descent learn 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 binary label.
7

---

Published as a conference paper at ICLR 2025
Fine-Grained Supervision. Following the setting in Section 3, consider the binary classification task with data distribution D( (( ( 1)-1, 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 ý E [R] it belongs to, and then train k-class classifier to predict the cluster labels. More specifically, we first sample a training set S := {( y:) S IR x {+1} from D, along with the cluster labels {ūr }==1 for all data points. Then k-class neural network classifier is trained on S == {(z,, Rt) S R x [k].
Multi-Class Network Classifier. We train the following two-layer neural network for the k-class classification mentioned above: Fo(x) = (fi(z), f2
E R*, where fa()) == $ Σー ReLU((w,,.. 2)), 0 == (w 1, w1,2..... WK.A) E Rkhdd are trainable weights, and h = @(1) One can think (ReI.U((un,a)) x))
SE(M) as kh (k)) neurons partitioned into
k groups, where the corresponding second layer weights are in way that the j-th group only contributes to the j-th output (ą) of the network. The output Fe (a) is converted to probabilities exp(fx(m))
using the softmax function, namely Py(x) Σiz (x)) for j E [k]. For predicting the binary label for the original binary classification task on D, we take the difference of the probabilities of the positive and negative classes, i.c., Jrbiuary (x) EjEJ, P2(x) -E _)t)_)(<)))
(z)) The clean accuracy Aecde (=binary) and 8-robust accuracy Accrobust (Fobinary, ;6) are then defined similarly before.
Training Objective and Gradient Descent with Fine-Grained Supervision. We train the multi- class network Fe (z) minimize the cross-entropy loss LCe
>E= log Pq. (z;). Similar to Section 3, we use gradient descent to minimize the loss function LCE z5 ()) with learning rate f). i.e.,
((++1) (()) VoLce Foco). At initialization, we set w ~N,0, a Ia) for some Ow O.
GD Finds Robust Networks. In contrast to the feature-averaging implicit bias in our previous setting (Theorem 4.5), the following theorem shows that with fine-grained supervision, gradient descent converges to 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 multi-class network on the multiple classification problem S : (x) T.) C _2 x [k] as described in the above, under Assumptions 3.2, 3.3 and 4.3. some 7 o(1), after 2 ST S exp(O(k) iterations, with probability at least 7, the neural network satisfies the following properties:
1. The clean accuracy is nearly perfect: Acc (Fains 21-exp(-221onn d))-
2. The network converges to the feature-decoupling regime: there exists time-variant coeffi- cient X(T E [S2(log k), +00) such that for all E [R4), r E [h], the weight vector w can be approximated as
(T)
(T)) w - |1 31, S old
3. Consequently, the corresponding binary classifier achieves optimal robustness: for perturbation radius ō = O(v)), the 6-robust accuracy is also nearly perfect, i.e. AccPommst (F Anbinary, ;8) > 1 - exp( 2Qlgg d)).
The detailed proof can be found in Appendix F.3. Theorem 4.7 manifests that the multi-class network learns the decoupled features, and the induced binary classifier achieves optimal robustness. See Figure 1 for an illustration. Instead of leveraging the bias term to filter out cluster noise as the feature-decoupling classifier fFD that we illustrated in Figure and Theorem G.3, the soft-max operator of F Fhinary plays similar role here.
can be easily verified that hinary achieves optimal robustness radius (up to constant factor) since the distance between distinct cluster centers is at most @ Vd) (i.e. || |, = @(vd), Vi # 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 counterexample, we show that there exists a multi-class network that achieves perfect clean accuracy but is not Q( -d/b)- which is formally given in the following proposition.
8

---

Published as a conference paper at ICLR 2025
3.6
08
98
0.4
02
0.4
01
40-
80
P1
M1
Rb
Mb-
P4
Pi:
Mr
Mr
My W.W.W MyMalo
WWW. PPL
WMW WWY Welh
SynData: 2-class
(b) SynData: 10-class
(c) CIFAR- 10: 2-class
(d) CIFAR-I0- 10: 10-class
Figure 2: Illustration of feature averaging and feature decoupling on synthetic dataset (a.b) and CIFAR-10 dataset (c.d). Figure (a) and Figure (c) correspond to models trained using 2-class labels, and Figure (b) and Figure (d) correspond to models trained using 10-class labels, respectively, Each element in the matrix, located position (i. represents the average cosine value of the angle between the feature vector μe of the i-th feature and the equivalent weight vector wy of the fz(-).
Proposition 4.8. Consider the following multi-class network Fg: for all j E [k], the sub-network fj has only single neuron (h = 1) and defined as fz(2) = ReLU ((1, + EieJ, Hi, x)), where cluster has binary label
{+1}. With probability least - cxp( exp(-9((lg' d)) over S, we have that CCE(O) (0) S ((d)) o(1), where 0 denotes the weights of Fø. Moreover Acc (1 (hbimary) 2 1 exll-expp(-ppp exp(-S((lg" d)), Accrobust (1' ((Vprrk)((K)))) S exp(-S(log' d)).
5 EXPERIMENTS
5.1 FEATURE AVERAGING AND FEATURE DECOUPLING
To validate our theoretical results, we conduct experiments on the synthetic dataset as we mentioned in Section 3 and the CIFAR-10 -00 dataset, described as follows:
Synthetic Dataset. We generate the synthetic data following the data distribution in Section 3. Specifically, we choose the hyper-parameter:
& 10, d 3072, m 5, n 1000, = = 1, r 0.001, an = d = 0.00001, T = 100. For simplicity, we denote the weights of the two-layer network W1, W22,.... ... w1o (where the first five weights correspond positive neurons and the other five weights correspond negative neurons). We also the first five clusters as positive and the others negative. Additionally, we provide an ablation study for other choices of hyper-ypaaeee parameters (see the details in Appendix B).
Binary Classification on CIFAR-10, We create a binary classification task from the CIFAR- 10 dataset by merging the first five classes into class and the other five classes into the other class. We use the normal 10-classification on CIFAR-10 as the 10-class task.
Pre-trained Feature Extractor. We utilize a ResNet18 model pre-trained on CIFAR- 10 and replaced the model's final layer with two-layer ReLU neural network described in our theory (fixed second layer as diagonal form, i.c., (z( := 1 Σ ReLU((W,, 2)), Vj [10]. where z denotes the hidden representation of the penultimate layer). We only train the last two layers. We choose the width of the first layer to be 30 (h 3) ensure that the accuracy of the pre-trained model was not compromised. Then, inspired by the theoretical study about neuron collapse (Papyan al., 2020), we calculate the average value of the penultimate layer output of the neural network for each class the corresponding wy
Σ/_ Wj.r as the equivalent weight of fz. For the 15 positive weights and 15 negative weights in the binary classification network, equally divide them into 5 positive classes and 5 negative classes to ensure fair comparison, which ensures that two models both have the same form f2. ... f10) IR and each sub-network f, corresponds
feature μu; of that class [10). For 10-classification,
weight vectors w,,
Experiment Results. The results shown in Figure 2, which demonstrates theoretical findings: 2-classification model learns ceature-averaging solution while 10-classification model learns the feature-decoupling solution. Concretely, Figure 2a and Figure 2c correspond to feature-averaging

---

Published as a conference paper at ICLR 2025
Synthetic
Robust Test Accurccy
Robust Rost Accursyy
10
Perturbttion
Porturbation
Perturbation Radius
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-I0 10 (the right).
regime in Theorem 4.5, where the correlations between each weight vector of positive (negative) neuron and all positive (negative) cluster features are nearly-euually larger than those between each weight vector of positive (negative) neuron and all negative (positive) cluster features; Figure 2b and Figure 2d correspond to our feature-decoupling regime in Theorem 4.7, where the correlation matrix nearly diagonal. Additionally, we also provide a transfer learning experiment based on the CLIP model (Radford et al., 2021) to further verify our theory (see the details in the Appendix H).
5.2 ROBUSTNESS IMPROVEMENT FROM GRAINER SUPERVISION INFORMATION
Moreover, we aim verify whether the model trained with fine-grained supervision information (i.e., 10-class labels) is more robust compared to the model trained with only binary (2-class) labels.
Experiment Settings. To ensure fairness in the comparison, we sum the logits corresponding the 5 positive classes and subtracted the sum of the logits corresponding to the 5 negative classes from the 10-class model's output. This result used as the binary classification output for the 10-class model. The robust accuracy is measured by using the standard PGD attacks (Madry et al., 201 8) with different L (y-pertubation radius. We run experiments in the following datasets:
Synthetic Dataset. We generate synthetic data as the same as that in Section 5.1.
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 (MNIST: parity- classification; CIFAR-10: 10: clasy-classification that we mentioned in Section 5.1).
Experiment Results. The results are presented in Figure 3. 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 all datasets. This collaborates with our theoretical results (Theorem 4.5 and Theorem 4.7) that the models can achieve better robustness with more supervised information.
6 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 two-layer ReLU network reveal tendency for gradient descent to average/combine individually meaningful features, which can lead 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 supported by empirical evidence from both synthetic and real-world data, including MNIST and CIFAR-10. 10. Our findings not only deepen our understanding of adversarial examples in deep learning also suggest that fine-grained supervision can enhance the robustness of deep neural networks against adversarial attacks.
ACKNOWLEDGMENTS
Binghui Li is partially supported by the Elite Ph.D. Program in Applied Mathematics for PhD Candidates in Peking University. We thank anonymous reviewers for their valuable suggestions.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Zeyuan Allen-Zhu and Yuanzhi Li. Feature purification: How adversarial training performs robust deep learning. In 2021 IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS), pp. 977-988. IEEE, 2022.
Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self- distillation in deep learning. In The Eleventh International Conference on Learning Representations, 2023.
Zeyuan Alu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 242-252. PMLR, 09-15 Jun 2019a.
Zeyuan Allen- Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pp. 242-252. PMLR, 2019b.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bout Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019a.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. International Conference on Machine Learning, pp. 322-332. PMLR, 2019b.
Peter Bartlett, Sébastien Bubeck, and Yeshwanth Cherapanamjeri. Adversarial examples in multi- layer random relu networks. Advances in Neural Information Processing Systems, 34:9241-9252. 2021.
Battista Biggio, Igino Corona, Davide Maiorca, Blaine Nelson, Nedim Šindié, Pavel Laskov, Giorgio Giacinto, and Fabio Roli. Evasion attacks against machine learning at test time. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2013, Prague, Czech Republic, September 23-27, 2013, Proceedings, Part III 13, pp. 387-402. Springer, 2013.
Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an Constin-Uhllenbee like process. Jacob Abernethy and Shivani Agarwal (eds.), Proceedings of Thirty Third Conference on Learning Theory, volume 125 of Proceedings of Machine Learning Research, pp. 483-513, PMLR, 09-12 Jul 2020.
Sébastien Bubeck and Mark Sellke. A universal law of robustness via isoperimetry. Advances in Neural Information Processing Systems, 34:2881 11122811-200 1-28822, 2021.
Sébastien Bubeck, Yeshwanth Cherapanamjeri, Gauthier Gidel, and Remi Tachet des Combes. single gradient step finds adversarial examples on random two-layers neural networks. Advances in Neural Information Processing Systems, 140109, 2021a.
Sébastien Bubeck, Yuanzhi Li, and Dheeraj M Nagaraj. A law of robustness for two-layers neural networks. In Conference on Learning Theory, pp. 804-820. PMLR, 2021b.
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent wide and deep neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bun E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Yuan Cao, Zixiang Chen, Misha Belkin, and Quanquan Gu. Benign overfitting in two-layer convo- lutional neural networks. Advances in neural information processing systems, 35:25237-25250 2022.
Niladri S Chatterji and Philip M Long. Finite sample analysis of interpolating linear classifiers in the overparameterized regime. Journal of Machine Learning Research, 22(129))1-30, 2021.
11

---

Published as a conference paper at ICLR 2025
Zixiang Chen, Yihe Deng, Yue Wu, Quanquan Gu, and Yuanzhi Li. Towards understanding the mixture-of-experts layer in deep learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Muthu Chidambaram, Xiang Wang, Chenwei Wu, and Rong Ge. Provably learning diverse features in -iew data with midpoint mixup. In International Conference on Machine Learning, Pp.
5563-5599. PMLR, 2023.
Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable programming. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-But Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Alex Damian, Tengyu Ma, and Jason D Lee. Label noise SGD provably prefers flat global minimizers. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 27449-27461. Curran Associates, Inc., 2021.
Amit Daniely and Hadas Shacham. Most relu networks suffer from \ell adversarial perturbations. Advances in Neural Information Processing Systems, 33:6629-6636, 2020.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pp. 1675-1685 PMLR, 2019a.
Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-paramettrrized neural networks. In International Conference on Learning Representations, 2019b.
Logan Engstrom, Justin Gilmer, Gabriel Goh, Dan Hendrycks, Andrew Ilyas, Aleksander Madry, Reiichiro Nakano, Preetum Nakkiran, Shibani Santurkar, Brandon Tran, Dimitris Tsipras, and Eric Wallace. A discussion of addveraaiil examples are not bugs, they are features'. Distill, 2019. doi: 10.2391 5/distill.00011 https://issilll.ub//011//aaaaa wwwppp//////////////pppaaaaeaa wuw.eee--buu-dddussioooo
Spencer Frei, Niladri S Chatterji, and Peter Bartlett. Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. In Conference on Learning Theory, pp.
PMLR, 2022.
Spencer Frei, Gal Vardi, Peter Bartlett, and Nati Srebro. The double-edged sword of implicit bias: Generalization vs. robustness in relu networks. Advances in Neural Information Processing Systems, 36, 2024.
Yossi Gandelsman, Alexei A Efros, and Jacob Steinhardt. Interpreting the second-order effects of neurons in clip. arXiv preprint arXiv:2406.0433 04341, 2024.
Justin Gilmer, Luke Metz, Fartash Faghri. Samuel S Schoenholz, Maithra Raghu, Martin Wattenberg, and Ian Goodfellow. Adversarial spheres. arXiv preprint arXiv:1801.02771 2018
Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv:1412.6571 2014.
Xinran Gu, Kaifeng Lyu, Longbo Huang, and Sanjeev Arora. Why (and when) does local SGD gen- eralize better than SGD? In The Eleventh International Conference Learning Representations, 2023.
Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro. Implicit regularization in matrix factorization. In Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1832-1841. PMLR, 10-15 Jul 2018.
12

---

Published as a conference paper at ICLR 2025
Jeff Z HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. In Conference on Learning Theory. pp. 2315-2357. PMLR, 2021.
Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Logan Engstrom, Brandon Tran, and Aleksander Madry, Adversarial examples are not bugs, they are features. Advances in neural information processing systems, 32, 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
Ziwei and Matus Telgarsky. Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems, 33:17176-17186, 2020a.
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. International Conference on Learning Representations, 2020b.
Yiwen Kou, Zixiang Chen, Yuanzhou Chen, and Quanquan Gu. Benign overfitting in two-layer relu convolutional neural networks. In International Conference on Machine Learning, pp. 17615- 17659. PMLR, 2023a.
Yiwen Kou, Zixiang Chen, and Quanquan Gu. Implicit bias of gradient descent for two-layer reLU and leaky reLU networks on oothhoootthognal data. In Thirly-seventh Conference on Neural Information Processing Systems, 2023b.
Soichiro Kumano, Hiroshi Kera, and Toshihiko Yamasaki. Theoretical understanding of learning from adversarial perturbations. In The Twelfth International Conference on Learning Representations, 2024.
Beatrice Laurent and Pascal Massart. Adaptive estimation of a quadratic functional by model selection. Annals of statistics, pp. 1302-1338, 2000.
Binghui Li and Yuanzhi Li. Why clean generalization and robust overfitting both happen in adversarial training, arXiv preprint arXiv:2306.0127// 2023.
Binghui Li and Yuanzhi Li. Adversarial training can provably improve robustness: Theoretical analysis of feature learning process under structured data. arXiv preprint arXiv:2410.08503, 2024.
Binghui Li, Jikai Jin, Han Zhong, John Hopcroft, and Liwei Wang. Why robust generalization in deep learning is difficult: Perspective of expressive power. Advances in Neural Information Processing Systems, 35:4370-4384, 2022a.
Xiao Li, Yining Liu, Na Dong, Sitian Qin, and Xiaolin Hu. Partimagenet++ dataset: Scaling up part- based models for robust recognition. European Conference on Computer Vision, pp. 396-46 066-414. Springer, 2024.
Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? -a mathematical framework. In International Conference on Learning Representations, 2022b.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogencous neural networks. In International Conference on Learning Representations, 2020.
Kaifeng Lyu, Zhiyuan Li, Runzhe Wang, and Sanjeev Arora. Gradient descent two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 34:12978-12991. 2021.
Kaifeng Lyu, Zhiyuan Li. and Sanjeev Arora. Understanding the generalization benefit of normaliza- tion layers: Sharpness reduction. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances Neural Information Processing Systems, volume 35, pp. 146899347088 Curran Associates, Inc., 2022.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. In International Conference on Learning Representations, 2018
13

---

Published as a conference paper at ICLR 2025
Hancheng Min and Rene Vidal. Can implicit bias imply adversarial robustness? In Forty-first International Conference on Machine Learning, 2024.
Andrea Montanari and Yuchen Wu. Adversarial examples in random neural networks with general activations. Mathematical Statistics and Learning. 6(1):143-200 2023.
Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro, and Daniel Soudry. Lexicographic and depth-sensitive margins in homogeneous and non-homogeneour deep models. Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 4683-4692. PMLR, 09-15 Jun 2019.
Tianyu Pang, Min Lin, Xiao Yang. Jun Zhu, and Shuicheng Yan. Robustness and accuracy could be reconcilable by (proper) definition. In International Conference on Machine Learning, pp. 17258-172777 PMLR, 2022.
Vardan Papyan, XY Han, and David I. Donoho. Prevalence of neural collapse during the terminal phase of deep learning training. Proceedings of the National Academy of Sciences, 117(40):
24652-24663, 2020.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark. Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8748-8763. PMLR, 18-24 Jul 2021.
Ali Shafahi, Mahyar Najibi, Mohammad Amin Ghiasi, Zheng Xu, John Dickerson, Christoph Studer, Larry S Davis, Gavin Taylor, and Tom Goldstein. Adversarial training for free! Advances in Neural Information Processing Systems, 32, 2019.
Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateck Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. Advances Neural Information Processing Systems, 33:9573-9558 2020.
Adi Shamir, Odelia Melamed, and Oriel BenShmuel. The dimpled manifold model of adversarial examples in machine learning. arXiv preprint urXiv:2106. 10151, 2021.
Berfin Simsek, Amire Bendjeddou, Wulfram Gerstner, and Johanni Brea. Should ander-parameterizn student networks copy average teacher weights? In Thirty-seentth Conference Neural Information Processing Systems, 2023.
Chawin Sitawarin, Kornrapat Pongmala, Yizheng Chen, Nicholas Carlini, and David Wagner. Part- based models improve adversarial robustness. arXiv preprint 2=XX::2209009111 2022.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19(1): 2822-2878, 2018
Christian Szegedy, Wojeiech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfellow, and Rob Fergus. Intriguing properties of neural networks. arXiv preprint arXiv:1312.6199 2013.
Nikolaos Tsilivis and Julia Kempe. What can the neural tangent kernel tell us about adversarial robustness? Advances Neural Information Processing Systems, 15:18116-1818 2022.
Dimitris Tsipras, Shibani Santurkar, Logan Engstrom, Alexander Turner, and Aleksander Madry. Ro- bustness may be at odds with accuracy. In International Conference on Learning Representations, 2019.
Gal Vardi, Gilad Yehudai, and Ohad Shamir. Gradient methods provably converge to non-robust networks. Advances in Neural Information Processing Systems, 35:20921-20932, 2022.
Zixin Wen and Yuanzhi Li. Toward understanding the feature learning process of self- contrastive learning. In International Conference on Machine Learning, pp. 11112-11122, PMLR, 2021.
14

---

Published as a conference paper at ICLR 2025
Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma. Frequency principle: Fourier analysis sheds light on deep neural networks. arXiv preprint arXiv: wwXiv:1901.06523 2019.
Zhi-Qin John Xu, Yaoyu Zhang, and Tao Luo. Overview frequency principle/spectral bias in deep learning. Communications on Applied Mathematics and Compulation, pp. 1-38, 2024.
Hongyang Zhang. Yaodong Yu, Jiantao Jiao, Eric Xing, Laurent El Ghaoui, and Michael Jordan. Theoretically principled trade-off between robustness and accuracy. In International conference on machine learning. pp. 1472-7482 PMLR, 2019.
Hanxu Zhou, Qixuan Zhou, Tao Luo, Yaoyu Zhang, and Zhi-Qin John Xu. Towards understanding the condensation of neural networks initial training. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep ReLU networks. Machine learning, 109:467-492, 2020.
15

---

Published as a conference paper at ICLR 2025
CONTENTS
Introduction
2 Related Work
3 Problem Setup
3.1 Data Distribution
3.2
Neural Network Learner
4
Main Results
4.1
Network Learner Provably Learns FeatrreAAveragng Solution
4.2
Fine-Grained Supervision Improves Robustness
5
Experiments
5.1
Feature Averaging and Feature Decoupling
5.2
Robustness Improvement From Fine-Grained Supervision Information
10
Conclusion
10
Connections Our Results with Other Explanations of Adversarial Examples
18
Additional Synthetie Experiments
20
B.1 Ablation Study on Synthetic Datasets
20
B.2 Adversarial Training on Synthetic Datasets
20
Analysis of Training Dynamies for Feature-Averaging Regime
23
C.I Deriving Dynamics of Coefficients From Gradient Descent
23
C.2 Two Key Techniques about Loss Derivative and Activation Region
24
C.3 Proof Sketch of Theorem C.2
25
D Preliminary Properties
27
D.I. Detailed Data Model and Assumptions
27
D.2 Useful Properties of the Training Dataset..
28
D.3 Useful Properties of the Network Initialization
30
Proof for Section 4: Feature-Averaginn Regime
35
E. Analysis of Training Dynamics
35
E.2 Proof of Theorem 4.5
50
E.3 Proof of Theorem 4.6
54
Proof for Section 4: Feature-Decoupling Regime
56
F1
Propositions of Network Initialization
56
16

---

Published as a conference paper at ICLR 2025
F.2 Analysis of Training Dynamies
57
F3 Proof of Theorem 4.7
64
Two Feature Learning Regimes: Feature Averaging and Feature Decoupling
68
G.1
Feature-Averaging Two-Layer Neural Network
68
G.2
Robust Two-Layer Neural Network Exists
69
G.3 Non-Robust Multi-Class Network Exists
69
H Additional Real-World Experiments
71
H.1 Empirical Verification of Feature Learning Process
71
H.2 Infinity-Norm Case
17

---

Published as a conference paper at ICLR 2025
A CONNECTIONS OF OUR RESULTS WITH OTHER EXPLANATIONS OF ADVERSARIAL EXAMPLES
(I) Approximate Linearity of the Model: Earlier hypothesis about the origin of adversarial examples (e.g., Goodfellow et al. (2014)) had proposed the idea that the existence of adversarial examples is related to the fact the model fo(a>) is approximately linear. Subsequently, there is a sequence of theoretical studies showing that adversarial examples exist abundantly in the input space for neural networks with random weights (without training), and a main insight is that such random networks are approximately linear and with high probability (by isoperimetry argument) (see e.g.. (Gilmer et al., 2018; Bubeck et al., 2021a; Bartlett et al., 2021; Montanari & Wu, 2023)). Our Theorem 4.5 proves similar approximate linearity. See details in the proof intuition of Theorem G.2. Think of the special case that the weight vector corresponding to each neuron is exactly the average of the cluster means (14+ == μ. Ejer, My or p._ == 1o_ Σjer_ 14)
input point is close to the decision boundary
and Σj= 14, = O. In this case, the two-layer network reduces to a simple linear model w.r.t. the perturbation. Theorem 4.5 also shows 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., 2021), which proposed that during training neural network first finds simple decision boundary that is close to most training points.
Another appealing point of view was developed in Ilyas et al. (2019), which proposed that adversarial examples are related to the presence of non-robust features. Here, feature refers to an individual measurable property or characteristic of the data used by the model to make predictions. Robust features refer to those that are not easily disturbed or affected by variations or noise in the data, allowing the model to maintain stable performance. Non-robust features, on the other hand, are features derived from patterns in the data distribution that are highly predictive but fragile. making them often incomprehensible to humans. 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. (2019), Engstrom et al. (2019), Tsilivis & Kempe (2022) and Li & Li (2024) visualized both robust and non-robust features. While robust features are more perceptually meaningful for human, non-robust features resemble noise and artifacts. Interestingly. they showed that non-robust 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 learned non-robust feature (1++ or )4_ can be utilized to attack the feature-averaging network.
(2) Non-Robust Features:
(3) Relation to the Lower Bound Examples in Li et (2022a): From the perspective of expressivity, Li et al. (2022a) constructed a lower bound example (see an illustration in Figure 4), 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., 2022; Frei et al., 2024)) echoes the essence of this lower bound example in spirit, and our results can be seen as an explanation from the perspective of optimization.
(4) Relation to Frequency Bias in Xu et al. (2019) and Xu et al. (2024): These works refer the phenomena that deep neural networks generally learn lower-freqeency features first, and then the higher-frequency ones. This can be seen as a particular form of simplicity bias. The feature averaging bias studied is also a form of simplicity bias: under our theoretical setup, the simplicity refers to the linear combination of cluster features, and it is closely related the approximate linearity of the decision boundary, as discussed in (1). Hence, both studies assert that neural networks tend to favor simplicity during the initial stages of training, sharing a similar underlying spirit.
(5) Relation to Superposition in Gandelsman et al. (2024): Similar "averaging" or "superposition" effects are also observed in the work of Gandelsman et al. (2024). In particular, they proposed an automated interpretation of individual neurons in CLIP models (Radford et al., 2021) by generating textual descriptions of their functions. They observed that in CLIP models, each neuron may encode several distinct and unrelated concepts, and they leveraged this effect to construct "semantic" adversarial examples. Specifically, they decomposed a recond-order effect of each neuron into a sparse set of word directions in the joint text-image space, and co-oceurring words in these sets can be used for the mass generation of semantic adversarial images. For example, the text prompt "A cat lounging in the sun, with a group of elephants in the background and a value sign in the foreground"
18

---

Published as a conference paper at ICLR 2025
Pg -
Figure 4: A schematic illustration of the construction Li al. (2022a): 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 more robust classifier exists (by classifying the blue balls from the red balls). One can observe the conceptual similarity with Figure 1.
generates an image where CLIP's prediction assigns a 65% probability to the label "dog" and 35% probability to the label "cat." Similarly, in our setting, neurons also encode multiple features, and our adversarial example is constructed by combining multiple features together.
(6) Relation to Fine-Grained Annotation in Sitawarin al. (2022) and Li al. (2024): These works have demonstrated that combining human prior knowledge with end-to-end learning can enhance the robustness of deep neural networks for object classification through part-based model that makes predictions by recognizing the parts of the object in bottom-up manner. By utilizing richer annotations, this approach helps networks learn more robust features without requiring additional samples larger models. Specifically, part- based models, which provide detailed segmentations of objects, allow the network to focus on discriminative object parts, thereby improving its ability resist adversarial attacks. Consistent with these findings, our theory and experiments show that fine-grained supervisory information can similarly improve the performance robust classification. We believe that incorporating additional supervision signals to further enhance robustness could be an important direction for future research.
19

---

Published as a conference paper at ICLR 2025
B ADDITIONAL SYNTHETIC EXPERIMENTS
B.1 ABLATION STUDY ON SYNTHETIC DATASETS
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-paramenn o = 1, = 0.001, Tw 0b = 0.00001, T = 100. We denote the weights of the two-layer network as w1, 272, who (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 μe; and the weight vector wj. The experiment result under baseline setting is presented as Figure 5 (a), Figure (b), Figure (c) and Figure 8 (b).
as 10, 3072, m = 5, n 1000, a
Effect of the number of samples. We vary the number of samples as 1000, 10000, 50000. See results in Figure 5. It shows that feature- averragng can not be mitigated via more training data.
Effect of the learning rate. We vary the learning rate as 7 0.01, 0.001, 0.0001. See results in Figure 6. It shows that the assumption about small learning rate is necessary for feature averaging.
Effect of the initialization magnitude. We vary the initialization magnitude ou ob 0.001, 0.0001, 0.00001. See results in Figure 7. shows that the assumption about small initialization is necessary for feature averaging.
Effect of the signal-to-noise ratio. We vary the signal-to to-is-st ratio as SNR := a/o = 0.5, 1,2. See results in Figure 8. It shows that our results can also apply to SNR (11) case.
Effect of the orthogonal condition. We vary the cosine value of the angle between different cluster center features cos ( 14j) = 0.00001, 0.001, 0.09, Vi & j. See results in Figure 9. It shows that the exact orthogonal condition can be relaxed to nearly orthogonal setting, under which feature averaging still happens.
Effect of the equinorm condition. We vary the minimal and maximal norms of cluster centers as min, || max; |||||| L.0/L. 0.8/1.2, 0.6/ 1.4 (from the smallest norm the largest norm, an arithmetic progression is formed). See results in Figure 10. It shows that the exact equinorm condition can be relaxed to general non-equinorm setting, where feature averaging still occurs.
B.2 ADVERSARIAL. TRAINING ON SYNTHETIC DATASETS
To investigate the relationship between our theoretical results and adversarial training, conducted the following experiments on synthetic data.
Experiment Settings. We employ hyper- pppee-prraeetess that are largely consistent with the baseline settings: k- = 10, = 3072, m = 50, n = 1000, a = o = 1, 0.001, ow = a = 0.00001, and T = 100. The only modification is increasing the network width m, adversarial training requires a wider network for optimal performance. We utilize PGD-based adversarial training. During adversarial training, we need select the E2 radius for adversarial attacks. Through experiment, find that training with an attack radius of 40 achieved the best robustness. We visualize the weights and feature correlations of the network trained with this radius using the same methods.
Additionally, we compare the robustness of the network obtained via adversarial training with that of networks trained using binary and 10-class labels. To illustrate the effect of different attack radii on the robustness of networks trained with adversarial Training, we include results for a network trained with an attack radìus of 20 the robustness experiments.
See experiment results in Figure 11. The results indicate that networks trained with adversarial training do exhibit tendency to learn feature decoupling solutions. However, the degree of decoupling is less pronounced compared to networks trained with fine- fine-grained supervision. In terms of robustness, adversarial training does provide significant improvements, but it remains slightly inferior to the robustness achieved through fine-grained supervision.
Experiment Results.
20

---

Published as a conference paper at ICLR 2025
Tm00200
n=50200
(a) 1000
(b)
10000
(c)
50000
Figure 5: Illustration of feature averaging on synthetic dataset, when varying the number of samples
24
12
11
(a) 17
0.01
(b) 17 0.001
(c) 17 0.0001
Figure 6: Illustration of feature averaging synthetic dataset, when varying learning rate
21

---

Published as a conference paper at ICLR 2025
Inbe0.001
hnb00.00022
06 0.001 (b) 0b 0.0001
(c) σ 0.00001
Figure 7: Illustration of feature averaging on synthetic dataset, when varying the initialization magnitude (of a
12
SNR 0.5
(h) SNR
(c) SNR 2
Figure 8: Illustration of feature averaging synthetic dataset, when varying the signal-to-noise ratio (SNR).
Inaully orhsppod
Tharry atcegeeu
meary athhogau
32
11
(a) cos M4)) 0.0001
(b) cos(p;; (4)) = 0.01
(c) cos (1i; (4)) 0.09
Figure Illustration of feature averaging synthetic dataset, when varying the orthogonal condition.
22

---

Published as a conference paper at ICLR 2025
irin mww.e//mn nomm-111
min nommmax nomm=38/12
nommimax nont -
[4
M4
0.5
44
0.5
M2
0.4
M2
M5
0.4
A4:
Mc
Mc
14:
Mi
0.2
Mc
0.2
p4
0.1
146
0.0
14
0.0
M10
P10
MLo
WW
Mig Mww.WW.
(a) min/max norm 1.0/1.0
(b) min/max norm 00.8/12
(c) min/max norm 0.6/1.4
Figure 10: Illustration of feature averaging on synthetic dataset, when varying the equinorm condition.
Synthetic Data
10
Robust Test Accuracy
62
adv40
56
42
1.
-04
-06
Perturbation Radius (Ly)
Figure 11: Adversarial Training Synthetic Datasets. The left: compare adversarial robust- ness among model trained by 2-class labels (red line), models trained by adversarial training with perturbation radius = 20 (purple) and 40 (green) and model trained by 10-class labels (blue line) on synthetic data. The right: each element in the matrix, located at position (i,jj,, represents the average cosine value of the angle between the feature vector by of the i-th feature and the equivalent weight vector w, of the (-).)
C ANALYSIS OF TRAINING DYNAMICS FOR FEATURE-AVERAGING REGIME
In this section, present proof sketch of Theorem 4.5, where provide detailed analysis of training dynamics in feature-as caattrr-vveraging regime.
C.I DERIVING DYNAMICS OF COEFFICIENTS FROM GRADIENT DESCENT
By rigorously analyzing the gradient descent iterations, we know that each neuron is situated within span that encompasses the collective cluster features and the intrinsic noise of the training data points. This span explicitly characterized by the weight-feature correlations. which is shown as:
Lemma C.1 (Weight Decomposition). During the training dynamics, there exists the following (t) normalized coefficient sequences s.nj and ; each pair E {-1,+13, r E [m].]] E [b), [m] such that
w (t)
(())
Σ 10 |1, M4
Σ sEim]
(0) en. 1 -
JELA]
Then, we give the restatement of Theorem 4.5 as follows.
Theorem C.2 (Restatement of Theorem 4.5), the setting of training two-layer ReLU network on the binary classification problem D((u;) J=) as described in Section 3, under Assumptions 3.2, 3.3 and 4.3, for some 7 = o(1), after S2nn
-1fn-1 11 < T S exp(Ō(22/21) iterations, with probability at least 1 7i the neural network satisfies the following properties:
1. The clean accuracy is nearly perfect: Acc, (Jor;) 2
d)).
23

---

Published as a conference paper at ICLR 2025
2. Gradient descent leads the network to the feature-averaginn regime: there exists time- variant coefficient X E [9((1), -00) such that for all {+1}, r E [m], the weight
(T' vector Warr can be approximated
!T
A())
Σ JEJ.
|
S sodd1/3
and the bias terms are sufficiently small, i.e.,
ba. S 0(1).
3. Consequently, the network is non-robust: for perturbation radius = 92( ((V///k) the 5-robust accuracy is nearly zero, i.e., Accrobust (forr); I exp(-9((lo) d)).
In light of Lemma C.I and the second item of the above theorem indicates that w is approximately proportional to the average of features in Js (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 == (i E [m] (w σ.))
b > 0}. for {-1,+1} and i E [n)] denoting the set of indices of neurons in positive or negative class (determined by s) which is activated by training data point I; time step t. Then, we apply Lemma C.1 to the gradient descent iteration (1), deriving the following result.
Lemma C.3 (Updates of Coefficients x...j. [k), é [N] and time 2 we have the following update equations:
(1) e.r)). (.
For each pair s E {-1,+ e [m].j
A (f+1) s,r.
(1)
8n.
Σ
|1g1121 (s))
(2)
iET
0 a,r. = G
(t+1)
(1)
8T)
||
(reS!!)
(3)
nm
where ( : (yifor (x,)) denotes the point-wise loss derivative at point w. and 1, = {i [n]
x; in cluster j} denotes the set of the training points in the j-th cluster.
According to equations (2) and (3) from Lemma C.3, know
E
s,..
Σ EZ
|1 a.
(1)
R
Σ
S.r.J1
(t)
(4)
iEL,
where we also use s.rrj 0 e.r, = and the fact that, w.h.p., we have ||E/|] = Vd = || ||-
(0)
(t) s,r,s (i.e., equation
suggests that we only need to focus on the dynamics of the noise coefficients (3)).
C.2 TWO KEY TECHNIQUES ABOUT LOss DERIVATIVE AND ACTIVATION REGION
It seems that the main difficulty in analyzing the iteration (3) is addressing the time-variant loss derivative ((()) and ReLU activation region S: To overcome these challenges, provide two corresponding key techniques (Lemma C.4 and Lemma C.6) 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 (,e,(t) ( /y, and the training data margin gap A (i,j) 9; 9 where q;
denotes the margin of the i-th training data at iteration defined as q.¹) yifo (æx). Then, have:
Lemma C.4 (Training data margins are balanced during training dynamics). There exists time threshold To such that, for any time & To and distinct data points (æ:; D))), it holds that
(1. $ e(k),
(5)
where we e(k) to denote time-independent term satisfying e(k) 0 as h: Co.
24

---

Published as a conference paper at ICLR 2025
According to Lemma C.4, for any distinct training data points æ; and 2 with the same label, the loss derivative ratio can be bounded as:
yr /2, /use (10) so exp( ~ 1+ A (i, o1).
(6)
where the first approximation holds due to (2) = 1/ + exp(z)) and the fact e ss 1 + for small 2 in the second approximation.
Remark C.5. This method initially proposed by Chatterji & Long (2021) in the context of benign overfitting for linear classification and was subsequently extended to networks with non- non-linear activation (Frei al., 2022; Kou al., 2023a). In this paper, we extend the aulo-balance technique of Kou et al. (2023a) from the single-feature case to our multi-cluster scenario to prove Lemma C.4.
Key Technique 2: Analyzing ReLU Activation Regions. Then, we turn to the analysis of the 2.2- In fact, after the first gradient descent update, the set of activated neurons can be described in the following lemma.
activation regions s(O.
Lemma C.6 (Each training data can activate all its corresponding neurons). For the same time threshold To as that in Lemma C.4 and all time 1 < 1 S To. it holds that S [m] all i E 1+ -1,3 for all i E I_,, where I i E I, yi = =] and I_ {i 1:iEl E.E, y/ -1..
and S ((t)
We rewrite our model as fau) -10 + p(() where f!² = nn Er((m) ReLU((W.. 12 x)) + b. E
(-1,1): 1). Then, Lemma C.6 manifests that f!²) is linear in the training data point (æ), V/ with label yr = 8. If we show {"g keeps small, we will have the linearization foco ~ f Mn
m Erg(m) ((w)) x.) + b(f) which allows us to approximate the data margin by noise coef- ficients (applying Lemma C.1 and (4)).
Remark C.7. Indeed, we use induction to prove Lemma C.4 and Lemma C.6 together (see Lemma E.7 in the appendix), where we show the case when t = 1 by our small initialization assumption and use the auto balance technique complete the inductive step (see the full proof in Appendix E.
C.3 PROOF SKETCH OF THEOREM C.2
Now, based on the two key techniques above, we provide proof sketch of Theorem C.2, which consists of five steps.
Step 1: Proving that feature coefficient ratio X (T)
(T) s1 22 (ir E Joii E Jo.) is close to 1.
By Lemma C.6, for all s E (-1, 1}, E Ig. we know
0 a.r,
(1+1)
0 s.r.a
(t)
s7)
1E.11..
ver
Combined with the loss derivative ratio bound (6), it furthermore implies that the noise coefficient
ratio is close to 1, i.e., for any T1, r2 E [m], $1, l2 E we have
10 (1)
o1).
(8)
P=0
Thus, for any 81, 82 E (-1,1],5 E $2 E Jog and time t S.T. T. we can derive
(
Σ
I /11//,,/===000
Step 2: Proving that
(T) A.n attains ?(1) for j Jo. and keeps o(1) for j J_..
By induction, we can show that both bias terms ((0) and 1! A (j E J_.) keep or1)-rrder dur- ing the learning process (Lemma 10 and Corollary E.16), which thereby implies the following approximation, i.e., for any s E (1,, E [m]. E Ij. we have
(1) (9 mn) + 1!0) R X s.n. + 8,
25

---

Published as a conference paper at ICLR 2025
where we also need time satisfying t exp(O(x9.5)) (see details in Lemma E.13).
Then, for any & E {-1,1},ře E 1 and data point (2), y4) satisfying yi = -s, we know
ReLU((W m.) B ReLU
ann and o0 + - o(1),
(10)
(1) where the last inequality holds due to A e,ri S (Lemma E.5).
Next, we approximate the model output for training data point (æ;, Ms) belonging to the j-th cluster as
(x;)
Σ
ReLU((W
x.) + b(t))
(10)
y1
((w())
xxi) + b(t)
Weirs
m
m
.E(-
rE[m]
rE(m]
(9) Vi m rE[m]
Σ (1
Ding
(0) w..
Σ
Σ del
Darr +
(1))
Bennd
(8) ys(1 + 1)0 y.1.
2
rE[m]
(11)
Therefore, we derive the following approximate update w.r.t.
Villis i.e., for any iteration
(t)
t E 10, expx(((())) we have (e+1) (7)(1)
y, ~ 1. + nd exp (-x
(Lemma E.21). By leveraging log(2 + 1) log(z) = 1, 2' we inductively prove o 9i.1, s & log(n)) (Lemma E.7) and
= (
(t)
aar = log(nt), E J, (Lemma E.23). When the assumption T = (()=1) holds, we have (7)
ann = 2(1,, E J..
Step Gradient descent leads the network to the feature-averaging regime.
We can choose (() = (2T
1,1,jo for some jo E J+ as the representative of A(T) = (1.77) E
{-1,+1),r E jm],j E J.).
Combining the result in Step and Step 2, we know that for any ( A(2)
((T)
(T) snn log(1TT).
We can also prove that the (T) is minimally affected by the coefficient 0 s,r,3 in weight decomposition. Thus, we have (Lemma E.25)
w (T)
(T)
||m 1uu;l121p M)
S o(d
jEJ,
Step 4: Proving that the clean accuracy is perfect.
For a randomly-sampler test data point (x = μ + {. D within cluster j E Jy. we can prove that, with probability least exp((2 d)), it holds that kw!? E o(1) for all 8 E {+1}, E [m]] (Lemma E.26). Then, for data satisfying the above condition, we can calculate the data margin yforr (x) = m Eremm] Aears 2(1) O, which implies that Acc (forn) 2
1 1 D-cxp(-S2(ogg d)).
Step Proving that the robust accuracy is poor.
M3 11)). By applying Lemma E.26 again, we can derive that sgn(feerr (xc + p)) sgu(S (x)), which means Acc (Sorr); 2(1 c) c) d/k) S exp( sppo=22 d)) and finishes the proof of Theorem C.2.
2(1 c)
We consider the perturbation
h
26

---

Published as a conference paper at ICLR 2025
D 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 assumptions. Our subsequent proofs will be based on the validity of these properties. The proofs of the inlined claims are concluded with the symbol, while the proofs of the overarching results are concluded with the O symbol.
D.1 DETAILED DATA MODEL AND ASSUMPTIONS
First, we recall the definition of multi-class feature data distribution that we defined in Section 3. Definition D.1 (Multi- Cluster Data Distribution). Given & vectors Pi..... M& IR called the clus- ter features, and a partition of [k] into two disjoint sets J+ = (J+, J-), we define D (( 1/:
data distribution Rd {-1, 1). where each data point (ar, y) is generated follows:
1. Draw a cluster index as j Unif([k));
2. Set = +1 if E J+; otherwise j E J_ and set y = 11;
3. Draw x p4j + &, where & ~ NN Ia).
For convenience, we write D instead (( (/11)1111:01 (1) if {rL [1,})=1 and Jr are clear from the context. For s E [+1], we write J, to denote J+ if = +1 and J_ if == 1.
To ease the analysis, make the following simplifying assumptions the distribution.
Assumption D.2 (Orthogonal Equinorm Cluster Features). The cluster features {1385=1 satisfy the properties that (1) ||wy||| vd for all j [k]; and (2) 14; My for all S ン k.
Assumption D.3 (Nearly Balanced Classification). The partition J7 satisfies c < M some absolute constant 2 1.
S c for
Next, we summarize the assumptions of these hyper-parameters that we mentioned in the main text, as listed below.
Assumption D.4 (Choices of Hyper-Parametersy We state the range parameters for our proofs in the appendix to hold.
d = ((Kk2))
(recall di the data dimension) (recall c is the balance ratio) (recall is the number of samples) (recall 2m is the width of network learner)
c c= ((1)
e [S(())), exp(O(log"(())
= ((k)
" s (
(recall r) is the learning rate)
6.0 (N of) give the initialization) Remark D.5 (Discussion of Hyper-Paramee 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 widely-applied simplicity in the literature of feature learning works (Wen Li. 2021; Chen et al., 2022; Allen-Zhu Li, 2022; 2023; Chidambaran et al., 2023). Namely, we need the data dimension d to be significantly larger polynomial in the number of clusters k to ensure all cluster features 14, 22, ... 14s. can orthogonal within the space TRd.. The balance ratio e is an absolute constant that is independent with d and k. Our results can be extended to == oxjy + of for some parameters a ①11) and =(1), but here we set = a 1 for simplicity. And further require that is large polynomial in k due our choice of large signal-noise-rati (recall that we have a/o (1), which implies that || I| = vd x ||E|| with high probability). We need m E [max[[]+]] [ k:) for the existence of robust solution (Theorem G.3). We assume the learning rate and the initialization magnitude ow, Gg sufficiently small, which helps the network to be trained in the feature learning regime (Lyu et al., 2021; Cao et al, 2022; Allen Zhu Li, 2023;
. of= I 1 (recall !.!
Kou al., 2023a).
27

---

Published as a conference paper at ICLR 2025
D.2 USEFUL PROPERTIES OF THE TRAINING DATASET
Now, we introduce some useful notations, for simplifying our proof.
Denote / = {i i = 1, -. , as the set of indices of all training data points.
Define c(-) as the map
where c(i) represents the index of the cluster to which point
belongs.
1, = (i:ie c(i) = j} denotes the set of the training points the j-th cluster.
. 4 = i E J+1 } and I- = {i i E I,c(i) E J_} denote the index sets of all positive class data points and negative class data points respectively.
Then, under Assumption D.2,D.3.D.4, we show that Proposition D.6 of the training dataset hold with high probability. It is worth noting that most of the proofs of Proposition D.6 follow standard methodologies and closely resemble those presented in Frei et al. (2024), as our data distribution is similar theirs.
Recall the cumulative distribution function (CDF) of the standard normal distribution, usually denoted as ((r), which is defined the integral
(xx)
1/
P /2
2 -00
V2T -x
Additionally, we have the following commonly used bounds on the tail probabilities of the standard normal distribution.
exp(- -1 - S V2TT
exp(-x²
(12)
2
Proposition D.6. Let A 4yd and = 8n? d m + 2k (n/k) In(n:/),, With probability at least 1 over sampled training dataset Dr,, we have the following properties:
1. For every E I we have vd - 2n((d) S ||E/l| S vd+ 2ln(dd.
2. For every E I we have |ao.|| S 3vd.
3. For every i, E 1, + j we have (E:s Ei>l S A
4. For every E I and j E we have Kp,; Ei) S
5. For every i,j E with c(i) + c(f) we have (az , S A.
6. For every i,j E with c(i) = c(j) we have 1 S (x,, 2d.
2n
2n In(
7. For every we have
In R S | S
V
k
Remark D.7. Property I and Property 2 show that the data-wise noise and training data point are bounded, i.e. ||&,]] RJ Vd and ||>;||| = O(v)). 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 S approximately includes Y (= ((kk)) examples from each cluster.
Proof of Proposition D.6. Now, we prove Property 1-7 one by one.
Property 1: We notice that | follows the Chi-squared distribution.
The concentration bound in Lemma by Laurent & Massart (2000) implies that for all / 2 0, we have
Pr ||E| d2 2v + 2/
<e_
Pr [|E|| - d < -2vdt <e
28

---

Published as a conference paper at ICLR 2025
Plugging in t (d), that
Pr |El| 2 2 = Pr |612- - 2vdlm(d) +2)n
- d in(d)
Pr || (vd 2ln(d)) s Pr ||E|| - d -2vdin(d))
S d In(d)
Thus, we have
Pr
Vd - - ||& S vd + 2
s 2d In(d)
(13)
Then by union bound, we have Property holds for every with probability at least 1 2nd In(d)
Property 2: When Property holds, by triangle inequality, we know Property 2 holds:
Mo(0 |I + ||&/|| S vd + vd In(d) S 3vd.
The proofs for other properties require calculations pertaining to Gaussian distribution. We first introduce useful lemma below.
Lemma D.8. Let & ~ N (0, Ia). For any x E Rd we have
Pr (r,, 4) 2 || In(d)] S In(c)
Proof of Lemma D.8. Note that
||||||
<)
has the distribution A (0, 1).
By standard Gaussian tail bound, have for every 2 0 that Pr
KK ||||
<)
2
2 exp
(-2)
(-2))
12
Plugging in f = (nd,, we can see that
In²(d)
Pr
2 In(d)
< 22exp
2d In(e))
x
2
Property 3: Next, we prove Property using the result in Lemma D.8.
Noting that if KE,, E,)l 2 IndIn(d), we have that least one of the following holds:
| 2 v2d;
2.
É |TE
In(d).
Now we bound the probabilities of these two events separately, By Property 1, we have
Pr[Ilaa 2 V2d S 2d In(d)
Next, by Lemma D.8, we have
Pr KK
&: ||E
Ei
2 In(d)
< 2d in(())
Then, by union bound, know that,
Pr [K: 6,)l 2 V2dln(d) < 2d
Im(d)
+ 2d In(d) < 4d In(d)
Then, applying the union bound for all pairs i E * j. we have (E: (, 22d holds with probability at least 4n²d
29

---

Published as a conference paper at ICLR 2025
Property 4: Applying Lemma D.8, we have
Pr [[,s vdin(d)) 2d indd)
Then for all pairs i & I, E J. applying union bound we have that ((m),:))l S vdln(d) for all i I,j E J holds with probability at least 2nnd Infa)
Property 5: By using the results above, we have
( 22 1 - (MeG)) McGi)) + (Maco) E5) + KMaco siil + +((E,
KMeet) En> (MEG) E:)l + K(E,
S <4VdIn(d)= =A
Property 6: By using the results above and noting that (Me()) Me(s) for i, with c(i) c(j). we have
(0,, .z)) dl S (Ma(S) Ex) + (MOG) Ei) Ke)
S <dV = A.
Thus, have
1 2d dd (x,, s 2d.
Property 7: We define X, for E the indicator random variable that the i-th point the j-th cluster. It takes value with probability T. and 0 with probability - F.
Then we know that [
Σ X,. Applying Chernoff bound, we have that
1El
Pr
2n
In
1 17
2n k
In
(exp((In'(n//a)) = 2(nn/k) ln(n/k)
k
Then for all E J. applying union bound, we have Property 7 holds with probability at least
1 1 (k(n/) In(nk/))
Combining all of the above together, we have Proposition D.6 holds with probability least 8n2d loca) 2k(nnkk))nn ln(n/k)
D.3
USEFUL. PROPERTIES OF THE NETWORK INITIALIZATION
The proofs for propertics of the network initialization require the range of the maximum value obtained from multiple independent samples drawn from Gaussian distribution. We first present the following useful lemma.
Lemma D.9 (Concentration of Maximum of Gaussians). Let X, NNO, < sisl be i.i.d. random variables. Denote Xmax (ix Xmin min For any t 2 0. we have
Pr Xmin S -V22lo(( t } exp (-t?/2).
Pr Xamax S (2log(() -7] S Eexp
( 22
2/2
Proof of Lemma D.9. The proof is standard and similar to Proposition A.I and A.2 from Chi- dambaram et (2023). We include for convenience of the readers.
For any a 2 0, we have
Pr Pr(Xmin min 2 -a] = (o(a)) (1 (1-a())))
30

---

Published as a conference paper at ICLR 2025
where Q(x) - (x)).. By using (1 z)" 2 næ for any [0, 1] and n E N, then get
Pr min 2 -a] 2 - (Q(a).
Now we use the elementary inequality for the tail of the normal distribution:
1
Q(a) 2
2
so that
1
Pr mina 2 aa] 2
le 2
2
Plugging (2log()) + L, get
(V/2lg(()) +)²
Pr
Xmin
2lolg(( - 2 lexp
2
1
2 log((1)) +
1
Texp
2
1
1
e
Similar to the previous proof, we know that
Pr [mxxx S =a (a))" $ (1-
a
(- a²
2 ))
exp
(-2
V2T (a2 +1)
Plugging / [2lo((1)- 11
v 2 2log(()) - 22 (og(()) 1)
(1))
(1)
Pr Xmna S 22log(() -t <
exp
2
exp (1/2)
V2 /2 loog(21)
exp (t/2)
I exp
V / (2ggg(22)))) + +1)
By applying Proposition D.6, we can derive the following result, which gives the range of network parameters the initialization.
Proposition D.10. With probability at least - 4mmd- In(d) 2mm- have the following properties for our network initialization:
For any s +11,+1),r [m], we have ow
5
S
ow (Vd+22n(()))
For any E {-1, (-1,+1},r E [m], we have b00 & 20, /0s/2In(i In()).
Proof of Proposition D.10. For w reusing the same argument as the proof of Property (1) in
Proposition D.6, we know that
Tv (Va- lnn(d) || < am (vā +
(Vā+2ln(d)) 2ln(d))
holds for all (-1,+111 E [m] with probability at least 4mn--l(dd)
31

---

Published as a conference paper at ICLR 2025
For barrs by standard Gaussian tail bound, we know that
Pr [600 2 20% (2l((m)) S exp((4ln(m) = m
Then using union bound, we know that [b__]] 2 20% 2m)(m) holds for all E {-1,+1), e [m] with probability at least
2m -3
In conclusion, we know that these properties hold with probability at least 1 1-112-000 2mdd In(d) 22m
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 D.11 (Activation Region over Data Input). Let Tenj == {i L (wo.., x.) + (0) 2 0} be the set of indices of training data points in the j-th cluster which can activate the r-th neuron with weight war at time step
Then, we give the following result about the activation region Tarrá
Proposition D.12. Assuming Proposition D.6 and Proposition D. 10 holds. Then with probability at least 1 0..11 22mk exp
(- 9km
R for all E +-1,+ [m]] we have
[aari 2 3km2
Proof of Proposition D.12. Given H4 for J, we have (w m4) ow ow
Using the conclusion in Lemma D.9, with probability least 0.01, have
min, (w: (0) m4) 2 1110 vd / 2 (2m).
In the following proof, we assume that the above conclusion holds.
Given w we have (wi?) Ex) ~ (O, ( (0) I)) Then by Gaussian tail bound (12), we have
Pr ww (0) s.rt Ei) 2 1.20 Vdv22nd2 (nmmm =1-0 - + 1.20 Vd / 1m(mm))
| ||
((V3ln (nmm)
3 3mmm
exp
2/33n 31m(2m) m ?
2
We denote Xani = (ww....)) E:) 2 1.20 ln(mm))
TTrrr Σ X
IET
Yend ((w!))
1006.rx0;) xxi) + b.o0t 2 0) and we know that [Taral
E Yanni.
ieT,
(w {.) 1.20 vd 2 In (2m), then
(w0 xzi) + ( > min mingg(-1
(w w 14)) + (no!)) Ei) - 20 / 211(m)
(0)
2 0.10 v d 2lm(2m) 20h v 21m(m) 2 0.
That is to say Y.rri
1ff X
a, 1. Thus
Pr
To Tarai
>
n
> Pr Ty
n
2m
2
32

---

Published as a conference paper at ICLR 2025
For any given {-1, +1), [m] and E Iy. we know that X x,,,, are ii.d. and B[X 2 m=
Then by Chernoff bound, we have
[r 2m?
Pr T, 2
2 cxp
exp
8m3
9km?
Then
Pr
T (arrul
2
3km
2 Pr
T
air.
3km 2
Pr
T" T
2rm²
exp
9km
Then by union bound, we know that To.r.l 2
holds for all f-1,++1], E [m], E J
3km²
with probability at least - 2mk exp 9km²
Combing the above together, with probability at least m -0.01 2mnk exp
have
9km
[Tar. 3km²
Next, we show that the pre-activation output of the network is very small, at the initialization.
nvd
Lemma D.13. For any I, E [m], {-1,+1}, we have (w., (0)
x;) + (0)
nm
Proof of Lemma D.13. By Property (2) in Proposition D.6 and Lemma D.8, we have
| s 3vd, (01l < aw (Vd+211g) + 21n(())))) 1162 ll & 20b /2 l2(m).
Therefore, by triangle inequality, we know that
nVd
(w 20 xx) + 61% < ||x w 10ll |666 <3vd 20 Vd + 20, / 1n(2m) <
nm
Finally, we present the following two lemmas about the range of the loss derivative.
y/ exp (-yifors (a::)
Denote e((+) Sigoo (n,) (1n) (x))
Di
1 + exp (opp--Mee -- (xx))
+ exp (VrSorn
2 vie, 3
Lemma D.14. For each I, we have-
3
Proof of Lemma D.14. By applying Lemma D.13, we know
fo (x:) <
1
Σ
Σ
(w x.) + b9o)
2n nm
In 2
+E|-1+11 rE]m]
Then, we have 1/2 < exp (yafor (xx)) < 2, and we can derive that
< y 11)
<
1 exp (vi (x.))
3
Lemma D.15. For each and any time step we have -1 E y:ee!!) < 0.
33

---

Published as a conference paper at ICLR 2025
Proof of Lemma D.15. It can be easily checked as follows.
-15 S y,e, + exp (Vifow (2xi) S
34

---

Published as a conference paper at ICLR 2025
E PROOF FOR SECTION 4: FEVTURE-AVERAGING REGIME
In this section, we provide the proof of Theorem 4.5. We analyze the training dynamics of gradient descent, which constitutes the main part of our proof.
E.1 ANALYSIS OF TRAINING DYNAMICS
We first assume that all the properties and lemmas mentioned in Appendix D hold with high probability over the sampled training dataset and the network initialization.
Now, we introduce some useful notations. Denote 5 as the set of indices of neurons in positive or negative class (determined by s) which has been activated by training data point æ; at time step t. Formally, we define it as ((t)) = fr E [mm] (w a.) + b.9 2 0} for E [-1,+1} and E 7. The following lemma describes the set of activated neurons after the first gradient descent update.
Lemma E.1. s for all E 1+ and so) 1, = for all i I_.
Proof of Lemma E.I. Without loss of generality, we consider the case when æ; belongs to the positive (1)
class. We show that (wf vi) + fir 2 0 for all E [m]. By applying the gradient descent update and Lemma D.13, we have
(1)
(1 (1) x.) -b 6/
(w
(w 1. x.i) + 610 17 ((00)) x.) + V bi. c()()))
(0)
1 (-dl/2- mn ((V,,,, .2.) + Ve, C (()
nm
First, we examine the update of linear terms as follows:
mn ( wi., C(((()), x.))
= Σ (1))
(w (0)
+ b1.rr 2 0)) xp-
Σ
((O)
(wo x.) + -1 2 01) xp, xi)
p
PEL+
PEI-
Σ
((0) 1
(w
f.r xp) + (0) 2 01)
(0)
(Ip: xt4)
Σ ((00 (ww mp) 61 2 20)) (x)p
DEJo(
pGlo())
Σ
1
mm
[
1 Fp) + +b110 (0) 2 01)
(Tp: mz)
2
Σ 1 (1w xz) + 11 2 20) (sr,,
3
3
pEla()
pfTai
d
Σ
1
(w)r
Lp) + 1,0 2 0))
24
Σ
1
(1w
xp)) +1100 2 01)
(Recall A= AVdln(d))
6
pETa()
PFI0(
2A
d(Ti.rote)
d
TT
6"
3
lEr/fe(1))
nd
2nA
(By Proposition D.12)
18km2
3
nA
(14)
Next, we examine the update of the bias term as follows:
-mn Vo,, V.. ((0 =
Σ
((O)
(w 1,r kp) + 610 2 01)
+ Σ
(())
(w0 xp) + 61000 2 20)
#
PEL
pel
Σ1 1
(w 1,r Fpp) 110 2 0)
2
3
PEl
2n
(15)
3
35

---

Published as a conference paper at ICLR 2025
Combining (14) and (15) together, know
(1) xz) + 1rr [wf
(1)
-dd/22 mm
(V 2.) +
+.y+VVon Vb, ((((0)))))))
nm
2m,
+ +nd 2 (-d
3
20
Since this inequality holds for all r E [m²] we have that s [m]] For the case when æ; belongs to the negative class, we have s s(1) 1,, = [m] using the same argument.
Now, we analyze the dynamics of the coefficients in the training process, where we first give the following reight-decomposition lemma.
Lemma E.2 (Weight Decomposition). During the training dynamics, there exists the following a,dd and o(t) d,,,,, each E {-1,+1],7 E [m].] E J,i E such that
coefficient sequences (t)
w!! ! (0)
jEJ
iES
Proof of Lemma E.2. First, we construct a set of ( and (o (0 according to the following recursive formulas:
(1+1) (())
ST7
E49 | ((w)") xx) b 20)
and =,rd
iEI
0 (t+1) srd
8 (1)
ST]
4000 ||: (w) x.) + b00 2 0)
a,r.
(0)
0, s.r,i
0.
Now, we prove by induction on & that (1 and (o! constructed above satisfy that
((( 3,r
(0) (()
|
E:1.111111
2,""
JEJ
iET
The base case when f = 0 the conclusion holds trivially. Assuming the inductive hypothesis holds at time step 1, we now consider the case time step + 1. By the update equation in gradient descent, we know that
(t+1) 6,
(t)
ST]
20001111 ((w))) + 610 20
TUT
(1)
8T]
(Mo()) + Ei) 1 (no!))) (10000000 x.)) + (1() 2 0)
u
5.7
mm
SEI
(1)
81]
Σ pi Σ {()
(w
(t) s.r æz) + 1(() 2 0) + Σ ( (fw! x.) 650 2 20)
((w!!).
w
JEJ
KE,
iET
w
+ Σ
M |111ا1||1
(1)) A.T.
s7
((()) | ||ا 1 (no!! x2,) + 150 2 20)
nm
JEJ
NEZ
Σ E1115:11 (6(1)
S7.
(1o0 xx)) + 1000 2 0)) (10 ||
nm
w + Σ
(1+1)
M4y
Σ
((+1)
E.
JEJ
iEl
This concludes the inductive step and the proof of the lemma.
36

---

Published as a conference paper at ICLR 2025
Naturally, we have the following corollaries.
Corollary E.3. The coefficients (1)
s.r for E (-1,+1)rr E [m], E J, i E defined in
Corollary (E.2) satisfy the following update equations:
(++1)
(1)
817
Σ 400
| (ww! xei) + 2 20).
s,r.j
s.r.
7171
(++1)
ST]
(w
( xz,) + 60 2 0)
0 8,
(0)
0,
0 8,
0.
Indeed, we only focus the dynamics of noise coefficients due the following lemma.
Corollary E.4. The coefficient sequences X X.) and o (t) J defined Lemma F66 satisfy:
s,r for each pair E {-1,+1), E [m]. E
A
11l -2 Σ
1--21
EEI
Proof of Corollary E.4. Using the result in Corollary E.3, we know that
and - >l?2)) ||,|||
s77 nm
Σ
p(())
((wo(t) mx) + b 2 0)
iEI,
Σ
0 s,T.
(t'+1)
(t')
||5: -2
iel,
Then summing up the above equations from t' 0
have
(1) s.n
|1g
-2= Σ
(9
||: -2
NE
We show that the sign of each feature/noise coefficient remains unchanged during the full training process as in the following lemma.
Corollary E.5. The coefficient sequences (t) and (1) for each pair E {11, r e [m].j e I,t 2 0 defined in Lemma E.2 satisfy:
2 0ifi Js 0ffi $ J..
aind
(t) Sand
> 0ift Is,
< 0 0 L.-
Proof of Corollary E.5. Using the results Corollary E.3 and Lemma D.15, we know that
(t+1)
(t)
STj
(t)
| 2 (wo()) 2.)) + b 2 0))
sgn (g 0
sgn
nm
(-s(.t)
sgn
sgn
(8y:)
Noting that a
s.r,á 0, we know that
(0)
sgn
(t)
sgn (sys)
37

---

Published as a conference paper at ICLR 2025
That is to say
(0)
0ifi I, 0fff $ I..
S.r.
Then using the result in Corollary E.4, we know
Of E Ja, Of j $ J.-
(0)
In the following proof, need the concept of margin. We denote the margin of training data point 2; at time step as q. yiJoo (ær)) and the margin gap between training data points 2; and æy at time step A (i, ye (an:) yj gets (ır))-
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 E.6. For any time step and two training data points 2., B. if qi 2 9j. have
y
(65/12
S (1.)
yill())
Proof of Lemma E.6. Recall (x)
e
Then
1
lte
(t)
(is)
1 + el
S 1.
(()
(1,) + et:
Since the exponential function convex, we know that
(1) e (15)
1+ c%
clu 2,1/2 1 elactay1/2 2 1.
(())
Next, we establish the relationship between the coefficient x()) 's and (4) margin g defined before. We denote
o.r. 's in Corollary E.2 and
10 Veirsefi)
+
Σ
(4) Beirii
rE[m]
re(m)
(i.)) d. 4
(1 (ur.o())
(1)
-
((())
m
m
rE(m]
rE(m)
Later, will show that q is good approximation for margin qi and ( (i,j) is good approxi-
(1))
mation for margin gap A (i, in Lemma 1. and Corollary E.
Next, arrive the main part of the proof. Inspired by Kou et al. (2023a), can prove that the training data's s margin tends to balance automatically.
2 2 (
APA
22
Denote max
We know that o(k -2.5) according to our hyper-
VT
(n(
parameter Assumption D.4.
Lemma E.7. For Tn (x..0(()0.5))) i,j E, (-1,1) the following statements hold:
38

---

Published as a conference paper at ICLR 2025
1. In(tn) S yun.i S n
k
(1)
2k In(t + 1),
2n
2. ( (1,5) - 5c when c(i) = c(j),
3. A (1.5) 63e,
4. 4 10 (i.5) & 65c,
5. v /use/ (()) < 1 + 14c when c(i) - c(i).
6. y.C." 1/,,"/15( / (10 2 1 130e,
7. (00) for Is.
8. 1 (0)
-a) S c, - 2c for E Is.
Remark E.8. Property of Lemma E.7 shows that the growth rate of noise coefficient is the logarithm s.r.à = logt. t. Property and Property show that the approximate margin gap small,
of time, i.e, (t) and the gap between 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.
training data points with the same cluster index is smaller. Property
Proof of Lemma E.7. Without loss of generality, we assume that i E I+.
We use induction prove this lemma.
Step 1: We first consider the base case when = 1 for the induction.
Property 7: Property 7 exactly the conclusion of Lemma E.I.
Indeed, other properties can be easily verified because we adopted a small initialization.
Property and 8: By Lemma E.3, 17 d=2 and noting that o s.r.. O, we have
(0)
"
2 (0) 11:11111111 (ww!)) xx) + + ((0) 2 0) 1
2nd
2k
<
In 2 - 2c.
nm
nm
and by Corollary E.5, for T, we have
(1) s.r.
0 2
k
In(
2n
Property 2, 3 and 8: By applying Lemma E.3, 17 S d and noting that x 0, we have
[X (1)
[A
(0)
ST)
||
(w0 x.)) + (1) 2 0)
15
2nd
m < (Property 8)
s.n.j
2m
Then, by the definition of ( (i.j), we know that
A ((,)))
Σ (20) (nnd
+ Σ
(o
rE]m]
re(m]
2nd
2m 2m
+ 2m
nm
< 5c
(Property 2)
63c
(Property 3)
39

---

Published as a conference paper at ICLR 2025
Property 4: By applying Lemma D.15, we know that
(1)
(0)
(:00)
(10
22)) +b2 2 20)=
s.j
1
TuT
SET
W s.j
(0)
[
TUTI
KET
3vd
(0)
3.
m
4yd
(0)
Σ
10)
+ +00 2 0)
nm
EE!
(0)
71
m
(16)
10
Then, we have
10:
= foci (x:))
Σ
Σ
/ x.) + 620
[E(-1,+1] rE[m)
1
2m ( )
m
<e
By triangle incquality, we have
(1,5) S q. 01 11,++1,010 10, S 2c
Property 5 and 6: By using the inequality above, Lemma E.6 and noting that e² < 1 + 2x for small J, we have
Gas S SPA
< 24+42 < 1+ < 1 130e.
Property
Property
Now we complete the proof of the base case when 1 for induction.
Step 2:
Assuming that the inductive hypothesis at time step holds, we consider time step 1 1. We first give some useful lemmas based on the inductive hypotheses, and then go on inductive proofs based on these lemmas.
Lemma E.9. Assuming the inductive hypotheses hold before time step and E Ig. the update equations in Corollary E.3 can be simplified as follows:
(1+))
(1)
Σ g
PET.())
(1+1)
||E
3.r..
Proof of Lemma E.9. By Property 7 in Lemma E.7 (inductive hypothesis) and Corollary E.3, the conclusion is straightforward.
40

---

Published as a conference paper at ICLR 2025
Then, demonstrate that the bias term remains consistently small as the following lemma.
Lemma E.10. For every s E (-1,++ +1).re E [m],
i E I, we have
16 1
Proof of Lemma 10. Without loss of generality, we assume that s Then, for any t' S t. by using Lemma E.9, we know that
(1++1) Lrs = 1.r
(t')
17
16:112.
11) = 1,7
7
g((t)
1
(10
mp) + + 2 20)
p
nm
PEI
(t)
((t)
nm
PEII
Thus, we derive
11++1)
(t)
(1)
12
(b((+)1)
(20)
2 d (b((+11) - b(:2)
1rs
1,r,8
Σ
(()
1,r
1,
2n
1,
PET+
Summing up the above incquality from = 1 to = =t 1, we have
b 61,0 -
2n
1,r - f.f..)
d
Then by inequality (16) and Property (1) in the inductive hypotheses, we know that
(
2n (t)
2n (1) d T + +60
4k 4k + 1)
d
10
Reusing the same argument as the previous proof, we know that b,.
Next, we prove that the true value of the margin q; the margin, we first estimate the value of each neuron in the following Lemma E.11, Lemma E. and Lemma 13
(t) close to its estimated value q( To estimate
Lemma E.11. Assuming the inductive hypotheses hold before time step I, for all 8 [m],j E J, we have
(w) s.r: (2))
Proof of Lemma E. We bound the gap between the inner product (wif)) (4)) and feature coefficient (t)
s.r.j follows:
(t)
s,r: (4))
x s.rvj
(1) A Mp |Ppl 2 +
Σ
(1)
Eg |5g11= -2 14: 1aal1--111)
100
un s,r
+
ann
PEJ
9E8
(w) s. mj)
(t)
(1g)
+ Σ
(1)
|1p -1 11 M331 + Σ
(1)
|,1
KKgg HY)
ann
Surpp
airg
GET
2A
o (t) 2A
s mو3)
(0)
(t)
+
Σ
d
s,rv
d
P
GET
2A
||E, 1
(w 143))
ΣΣ
(0)
(1) s.rg
d
|1pp|
PPI
9E1
124 In(t + 1)k
svd|ww!!! +
d
6
41

---

Published as a conference paper at ICLR 2025
The first equation employs the weight decomposition in Lemma E.2; the second inequality expands the inner product and applies the triangle inequality; the third inequality utilizes the properties from Proposition D.6; the fourth equation utilizes Corollary E.4; the fifth inequality utilizes Property and Property 8 in the inductive hypotheses.
Lemma E.12. Assuming the inductive hypotheses hold before time step for all (11,+1),re [m], i E I, we have
(unl.) E:)
((() a..r. S 6"
Proof of Lemma E. We bound the gap between the inner product (wo?7 E.) and noise coefficient (b)
a s,r as follows:
(w (1)
sirt Ei)
(0)
+
s.pppp 2 Σ
(1)
(e)
PEJ
9-6
(w (0)
{.)
+ 10 1000
Sills -1 - σ a,
+ Σ
(1)
8 s. 11 1 : +
Σ
(1) s.rg
2 E.)
1
PEJ
934
(w s,ro 6:)
(()
2A
Σ
(t)
2A
2,14 d
PEJ
24
I
= (w m4j)
Σ
(t)
Σ
(1) sirg
d
8,7.0
|
pFj
gEtp
gEl
12A
In(t 1)k
<vd|w 2011
+
d
The first equation employs the weight decomposition in Lemma E.2; the second inequality expands the inner product and applies the triangle inequality; the third inequality utilizes the properties from Proposition D.6; the fourth equation utilizes Corollary E.4; the fifth inequality utilizes Property 1 and Property 8 in the inductive hypotheses.
Lemma E.13. Assuming the inductive hypotheses hold before time step for all +1r, [m], i E L, we have
(wo ( 2.)) - -110
(0) 0
s.r.(fi)
3
Proof of Lemma 13 Using the conclusion in Lemma E.11 and Lemma E.12 and triangle inequality, we can directly obtain the conclusion in this lemma.
(() (w!()), v.)) s.mo(() -...
-a!0)
(2)
(ww (c(0)
sonoco) + +1u0 Ex) (!) S
Lemma E.14. Assuming the inductive hypotheses hold before time step t, for all E I, we have
46 (e)
95 1
Proof of Lemma 1. Without loss of generality, we assume that 4+-
42

---

Published as a conference paper at ICLR 2025
Using Property 7 in the inductive hypotheses, we know that
19:
Grl
1 m
yefoco (x:)
Σ
(
Lnoe))
Lorg
(8)
rE(m)
re(m]
(w 1,r: Dii + B
1 (,,,((()
0 1,rs
+
Σ
ReLU (w xzi + +b
m
re [m
rEmm]
w
1,r x.)) - (4)
(1)
1,r
(t)
ReLU
(8)
1
(6 -,,, + 10
m
m re[m]
rE[m]
re/m]
La
Ls
For L1 term, using the conclusion in Lemma E.13, we know that L S
For L2 term, we consider each term in the summation by distinguishing between two scenarios..
Case(I):
(t) (w 1,r xi)) S O.
Then we know that ReLU ((w0 22))
3
Case(II): (w (w-1,r 1,r. x.) 2
Using the conclusion in Lemma E.13 3 and Corollary E.5, we know that
ReLU ((w="),,:
"i,,, xx,) S (w (t)
1, xei) (()
(1)) (8)
-o S
3
Combining Case (I) and (II) together, know that La S 3
For La term, using the conclusion in Lemma E0 know that L3 S Combining the above together, we know that
lg (1) S.
Then, we can estimate the margin gap between two training data points using simple triangle inequality. Corollary E.15. Assuming the inductive hypotheses hold before time step t, for all i.j E I, we have (4.1)) (1.5) 2c.
Proof of Corollary E.15. By Lemma E.14 and triangle incquality, we have (8)
(1,)) A 100 (1,
la (t)
0 + -1g5"
1/01 < 2e.
Then we will analyze update equations for A (1) (i,j). By Lemma E.9, know that
Σ
(((++1)
((++1) ((tt1oos))=
Σ
(1(0) (()) -200mo6))
17
Σ
pl
Σ
|
vi.rs((()
2n
p
rE(m)
rE(m)
PETe(1)
DET.G)
Σ
(o(tt))
Σ
(o(1)
Br.re 0 (()
(5(0)) ||E:
E/ (,1 (1E,112)
2n
rE[m]
rE(m]
43

---

Published as a conference paper at ICLR 2025
Combining the above two equations together, we get the update equation for (.)(i.)).
(i,)) A (1,3)
7
Σ
Σ
goo 11 + (, 11 - (5 115
2n
p
pETc
pelcon
(17)
Lemma E.16 (Property 8). For 8 E {-1,+1},ie Is we have (1 Memolll S s,,lo-andd S 2e.
Proof of Lemma 1. We prove that for j J- -- 1nng for all E [m]] and the proof of the other part is similar.
We distinguish between two scenarios.
Case(D): For all E I, (w 2i) 604 1-1. < O.
Then by Corollary E.3 and inductive hypothesis, we know that
| = 1 11n, S.
Case(II); There exists E I such that (w (0)
-1.: Di) + ( 2
By Lemma E.13, we know that
( (w -1--1 xi) - (1) -1rr -1,1 3
- (8)
Then by Lemma E.10 and noting that (w 1.r, xzi + (() 2 and (1)
011,,, S 0, we have
(1) 1,ro 2 (w -w 1,r xzi (1)
(t)
1r 2 -b
(1)
2c
1,
Then using the conclusion in Lemma E.3, we know that
(*++1) (00 >
nd nm
nd
E
m
3
Thus we have (1+1) 2 -6. Noting that (1+1) we know that |o (1)
A 1,,..
S O. we have [ S e. Then by Corollary E.4,
s.r. S 2
(++1)
S 2e.
s..
The lemmas 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 the following lemma.
Lemma E.17 (Property 2 of Lemma E.7). Assuming the inductive hypotheses hold before time step t. and o(i) c(j), we have
A (i.5) < 5c.
(++1) 2 d;
(t+1) We distinguish
Proof of Lemma E. Without loss of generality, we assume that q between two scenarios., one is when | (i.5) is relatively small and the other is when | (i, is relatively large.
Case(I): A (i,j) 4e.
By equation (17), Lemma D.15 and 17 S 2 we know that
(1+1) (i,j) (i.j)
||:11 1500 |1, |
77
(1
||
- 11))11))
2nd
2n
2n
44

---

Published as a conference paper at ICLR 2025
So have
(1,1) S Se.
q
Case(II): (i.)) 2 4c.
By Corollary E.15, we know that
(1,)) (1,)) -2E 2e.
By Lemma E.6, we know that
2 e
A (1J)/2 2
1
(18)
((()
By Equation (17) and c(i) = c(j), we know that
(++1) A (i,j) - (1)
Σ gn(t) ||00)1
Σ
12 11
2n
PETe(0
PEToC)
(41) 11:/1 - (1)))))))
2n
g(o)
(2(0)(()
(1+
+
17 ((9 ((1 1) 2
2m
(()))
2n
<0, inequality (18))
S0 0
(19)
Furthermore, due the inductive hypothesis.
À (++1) (i.)) S (1.s) 5c.
Property 5: Using the result in this lemma and Corollary E.15, we know that
(i S 7e.
Using the above inequality and Lemma E.6 and noting that et < + 2x for small we know that
(++1) ( S (A $1+ 21 (1+1) (1,5) $ 14c.
At this point, we have completed the inductive proofs for Property 2 and 5 in Lemma E.7.
Property 3, 4 and 6: Next, we consider the general case where the two training data points æ, 2; are not necessarily in the same cluster to prove Property and 6 in Lemma E.7. 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 Á (i.j).
Lemma E.18 (Property Assuming the inductive hypotheses hold before time step i, we have
A (1.5) 63c.
Proof of Lemma E.18. We distinguish between two scenarios, one is when A (110(1,)) (1. small and the other is when ( (i, is relative large.
relative
Case(I): (i,)) < 62c.
(i,
45

---

Published as a conference paper at ICLR 2025
By equation (17), Lemma D.15 and 1, S d d=?
we know that
11+1) (i,5) - (1.5))
Σ
E
ge |0
11 11,0
2n
p
pele())
(0 11
((())
Σ
prt 1|0c)) 110001111/11 (3/01151111
2n
pEfo(s))
pefct
2nd
2nd
5nd
-
+ 2d + 2d
c
22
k
2k
So we have
A (1++)) (i.j) 63c.
Case(II): (i,j) 2 62c.
By Lemma E. 14, we know that
(1,)) 2 (1,5) 2x 60c.
By Lemma E.6, we know that
(())
2
(1.112 21 + A (1) (2,1//2 1 30c.
(20)
((B)
Furthermore, due to the inductive hypothesis, for any E Le(0) 8 E La(j), we know that
(
(1
146
(1 + 14c)
1
$
(21)
40
(1+
1 + 30c
1
By Equation 17 and c(1) = c(j), we know that
Å (1+1) (i. - (() (i,j)
Σ
(x(() (1 13-- Σ
p (
1 E|
2n
pE
pEla())
(10) 11
-2
-220011121 115011))))))
Σ (r(e)
Σ
2n
2n
pETe()
We analyze each of these two terms separately.
Similar to the proof of inequality (19),
||&
(000 11))
2n
" |k, (2((
(1 + (1+30c) 30c)
11 (()) ((1 + 30c) 11E)
2n
((0
2n
<0, by Inequulity (20)
< 0
(22)
By inequality (21) and Property 7 in Proposition D.6, we have
Σ
10)
|1
Σ pETc()
p
pElo())
44
$ 1.
Σ gETo6)
e
Σ
10)
(1
e)|4g
gET,C)
46

---

Published as a conference paper at ICLR 2025
Thus know that
7
Σ
20 p
E
10.
(23)
2n
pela())
PEtes))
By combining (22) and (23) with the inductive hypothesis, we know that A ( (i, - ( 63e.
Using the result in this lemma and Corollary E.15, we know that
A (++ (i, & 65c.
Using the above inequality and Lemma E.6 and noting that e" 1 11+ + 2x for small a we know that
yj 2 ((+)))
A (c.) <1+ (4(11) (i.j) S 130e.
y. vee
Now, we have completed the inductive proofs for Property 2, 3, 4, 5 and 6 in Lemma E.7. Property 1: To prove Property 1, we need to analyze the update equation for ( ",y We first prove (t) that 0 s.r, 's are balanced as follows.
Lemma E.19. For all i,,12 I,ri, 12 [m], we have
(t) (1 - -02000c 0 N2da nd S (e) Migiii S (1 + 200c)
(t) Mig M2dda
nd
Proof of Lemma 1. We first prove the right-hand side of the inequality and the proof for the left-hand side is similar. By Lemma E.9, for any < t, we know that
11)+1)
(t')
Vi, p((()) |1 ||
miids
(1++1) (t')
(())) ||E:2 1|
Viz
Vd + In(d) )
2 (Applying Property in Proposition D.6)
S (1 130c)
vd - In(d)
< 1 + 200
That is to
(f'+1)
( (f'+1)
(t)) Mig (2ikk
Miii ti S (1 + 200c)
-
Summing the above inequality from
= we have
(a
Piri Mayh < (1 + 200c)
Then, we can derive that
a (t)
Ui Mi S (1 + 200c)
(t)
+ a (1)
Big Pavia
Big 1i.h
(1 200c)
(0)
1
16 (no!)) (0o()).2)) x2)) b60 2 0)
22..2 Vig
nm
2dg
< (1 200c) (t)
Vig ra. 71m
ndd Reusing the logic of the above proof, we know that
< (1 + 200c)
Vig 12.i2
(t)
+
(t) (1 200c) 12yii
1
S (1)
nd
47

---

Published as a conference paper at ICLR 2025
Then, estimate the margin qi (6) only using s.r. which is presented the following lemma.
(6)
Lemma E.20. For every i 1, [m]) we have
n (0)
22 S Sa S
2n (t)
Weroois + 2e.
2k Bor
Proof of Lemma E.20. Without loss of generality, we assume that E 1+
We first prove the right-hand side of the inequality.
mij
Σ
Inne(
+
Lins
rE(m]
rE(m]
||&:|1 -2
(t)
(t)
Ma(())
irr
rE[m] PE Te())
rE [m
Σ
Σ
Me(all 1
|E
(1 + 200c) a
1.ro.t
+
nd
1
+ Σ
(11+2000)
(1++000)
+
1 nd
[m[m] pEfc())
re(m]
2mn (1)
2m
0 1.ro. +
k
kd
2mn
(1)
1.ro,i + me.
Then, by Lemma E.14, we know that
q 0 s 4;
2n (1) +Es
k 1,7o.. + 2e.
Reusing the argument of the above proof, we prove the left-hand side of the inequality.
mg! (1)
Σ
Σ
+
rE[m]
rE[m]
|&&l| 2
Σ
(6)
Mc(
Lrpp
+
re(m] PE PETe())
rE [m)
Σ
Σ
||&.|1
(1 - 200c)
0 1.ro.t
1
ofro nd + Σ (11-200e)
200c)
Mc(s)||
nd
rE[m] PE Te))
rE[m]
mn
m
2k 1.ro 2kd
mn
2k 1.ro.t me.
Then by Lemma E.14, we know that
(1) (1)
(t)
q q
2k 1,ro. 2c.
Furthermore, we also need to estimate ((() using o s.n as the following lemma.
Lemma E.21. For every i E 1, E [m]] we have
2n
< - yiel!) < 2eepp
n
3
exp
--
ho Virr,
2h yr.r.
Proof of Lemma E.21. Without loss of generality, we assume that E 4..
By Lemma E.20, we know that
((1)
1
2n
2n
(2)
(2)
2
2
exp
h
1.r,i -
2e
>
3
exp
k
0
1,r.
+ exp
9
exp
9
48

---

Published as a conference paper at ICLR 2025
-Ve(
S exp
(- 2k 1,n + 22
(- n2 (t)
S eexp
(-7
(22 (t)
T (
+ exp
qi
exp 9i
Then, can prove Property 1 based the inductive hypothesis.
Without loss of generality, we assume that i E L.
We first prove the left-hand side of the inequality.
(++1)
(t)
7
11:
1,,,
(1)
2nd
n (c) exp (-
1,r,1
2k 1.r,a)
(Applying Lemma E.21)
2k
2nd
In(1 + 1) +
(Monotone with respect to df (.r,))
n
7mm t
2k
In(t + 2),
(++1)
1,
0
<
2k
In(t + 2) (Corollary B.5).
Then we prove the right-hand side of the inequality.
(1+1)
(t)
1
nm
(1)
nd
exp (-7
2n (1)
1.r
+ 3rm
k 01..1
(Applying Lemma E.21)
h
(In(t) In((n))
d
1
(Monotonic with respect ( (().))
2m
3nm
h
(In ((t + 1)) + In(n))
2n
k In((t + 1)m).
2n
Finally, we prove Property 7. The proof is very similar to the proof of Lemma E.1. We show that (w ((+1)
v.) + (1+1) 2 0 for all i I.. By the inductive hypothesis, we know that (w (1)
xx,)
((0) 20 O.
(t+1) (w)r xx)) + (1+1) ( 1, xx)) + b/00 - 17 ((V ((0))). xz) + Vot,.-(011
2 11 ( (V Vio xi + Vo., ((O(t)))
Denote g(()) = (()), By Property (6) in inductive hypotheses, we know that for all i E
((((0)
1 (0 S e(0) < (1 + 130
1 130c
We examine the update of linear term first
- -V wn,, ((0)))
xz) = Σ
(())
Ep. Ri)
DEI,
Σ
p((t) p
(Ip, xi) + Σ
(d(t) (2p. xt)
pFIc
d
Σ ((())
+ A
(())
(24)
d1
2
+ (m)
4k
> nl)())
49

---

Published as a conference paper at ICLR 2025
Then examine the update of bias term.
Vbb bb. (((
Σ
efol (2o xii + b1.tt 20) 2 nL(
(25)
PEI
Combining (24) and (25) together, know that
((+1) (wi, x.i) + oj, 2
((+1))
Thus we know that s((++1) 1,1 [m]]
For the case when æ; belongs to the negative class, we can obtain 5 [m]] using the same argument. Now, we have completed the proof of Lemma F.7.
E.2 PROOF OF THEOREM 4.5
Now, we start to prove the main result Theorem 4.5.
Theorem E.22 (Restatement Theorem 4.5). In the setting of training a two-layer ReLU network on the binary classification problem P(uu))5=1> J=) as described in Section 3, under Assumptions 3.2, 3.3 and for some 7 = o(1), after Sin I T S exp(((x2/21) iterations, with probability at least 1 " the neural network satisfies the following properties:
1. The clean accuracy is nearly perfect: Acc (Sorr) ) exp(-S((og' d)).
2. Gradient descent leads the network to the feature-averaging regime: there exists time- variant coefficient X(T) E [I(1), +20) such that for all E {+1}, E [m], the weight
vector way can be approximated as
((T)
o(d-1/1))
JEJ.
and the bias term keeps sufficiently small, i.e.,
|,077 S 0((.).
3. Consequently, the network is non-robust: for perturbation radius = 52 22(vd/k) d/k), the 8-robust accuracy is nearly zero, i.e., Acc (foer) exp(-S((lg" d)).
Proof of Theorem E.22. We first prove that gradient descent leads the network to the feature- averaging regime (Property 2).
(T) Lemma F.23. For all {-1,+1}, Ja,r E [m], we have In(Tr) S s.r. S 54In(T
Proof of Lemma E.23. Without loss of generality, we assume that = 1.
Using Property in Lemma E.7 and Corollary E.4, we know that
X
1
(T)
Σ
|
2
1rr.p
(T)
21
1.r.1
(T)
In(Tr)
pEL,
A
(T)
Σ
||
(T)
2n (T)
<4nn(T +1)
Lnn
|e,
11r.p
k
Lemma E.24. For T1,72 E [m], $1, 82 E +-1,+1),jr E Jon. E Jozi have
(7)
< 204c.
(2)
Sg.ro.ja
50

---

Published as a conference paper at ICLR 2025
Proof. By Lemma E.19, we know that for any i1,
Mig 111d S (1 + 200c) of (4)
12212 + (nd) + 201c)) (4)
Mi 12.12""
Σ
11pp11008111 061.r1p
(2)
$1.11111
Ae2.2adaa
(T)
Σ
||Ep11
(e)
|V + In(d) 11 1
$ (1 + 201c)
|| - [2 1/
(1 + 201c) (1)
+ 204c.
We denote ((2)_ (2)
1,1.0 for some jo E J the representative of (darrs : 8 E (-1,+1]
[m]],j E J,}. By Lemma E.23 and Lemma E.24 for all E {-1,+1],r E [m],j E J., we have
((T)
s.yJ - 204c
(T)
1(7) < 41n(T + +1),
A(T) > In(Tr) R(1).
Lemma E.25. For all s E {-1, +1,, E [m]. We have
Vd w
A
Σ
= o(1)
jeJ.
Proof of Lemma E.25.
Recall the weight decomposition in Lemma E.2.
Vd
(T)
(27)
Σ
My 1 11-
Vdw
+
Vd.
(T)
x(T))
3,
s.n
3EJ.
jEJ.
La
+ Vd Σ
X (T)
+ Vd Σ
&ilel -2
JEJ_.
Vel
For C1 term, using the conclusion in Lemma D.10, we know that
VV < 2do <c= o(1)
For L2 term, using the conclusion in Lemma E.23, Lemma E.24 and noting that My are pairwise orthogonal, we know that
Vd jeJ
aind 2(T) 1ullm,11
(1T))
- 2
=
jEJ.
S Σ (204))2 (AC3))2 JEJ.
(TT)
I ll +1)
900kc o(1)
51

---

Published as a conference paper at ICLR 2025
For Ly term, by Lemma E.16 and triangle inequality, we know that
vd Σ -
ke o(1)
JEJ_,
For L4 term, by Property (1) in Lemma E.7, we have
Vd
(T) s,
||Ei||
f + Σ
s.n s.ndda (St, (.)llell
(T) (T)
iE!
iET
117/2
(T)
2A
(T)
0 8,rib1
+
2, 2,ri
VEI
8k2 In (T +
8k In (T 1)
n
d
8
84 A
S 16ke = o(1)
Combining the above together, we know that
d
w
|
=
jEJ.
By Lemma E.10, we know that | I
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 E.26. Let & ~ (O, Ta). Then, with probability at least 2nd In(d)/2 for all s E [-1,+13, [m] have
Kwly EN S 6'
Proof of Lemma E.26. Reusing the argument of proof of Property (3) and (4) in Proposition D.6. We know that with probability least 1 2nd In(d)/2 for all E T,j J,
Ke 4) s A, KE:, 4 A.
The remaining part of the proof is similar to the proof of Lemma E. and Lemma E.
(ww E)
(0)
(0)
2
Σ
( sin
E, 11 11 5,6)
=
PEJ
GEI
(1) s p2y)
100) 11ppl KHpp M.ح1 + Σ
(()
oohhallsall 10salll-(((6) -1( my)
PEJ
9E1
2A
< (w
(0)
+
(4)
(t)
d
s.r.p
s.rg
6
PEJ
1E1
Assume (x, y) randomly sampled from the data distribution D. Without loss of generality, we assume that μ4 + E.y 1. Using the conclusion in Lemma E. 11, Lemma E.26 and Lemma
52

---

Published as a conference paper at ICLR 2025
E.10, we know that
(w) (T)
1,r x)) + 1,7
(w0 (T)
1b) + (m (2)
5 + b/7
(4) c
11.r
6
6
2 1rr
(w 1, x) + +b TT
(T)
(w (T)
1, M4) + (w -
(w (T)
{) +b
S x 1,n +
(T)
+ c
+
6 3
$ (T)
Lnd te
6.
Then, we have
focr (x))
Σ
ReLUU (10 xx) b1.T)
Σ
ReLU (ur ((
-1,,, x)) + b..r))
m
m
rE(m [m]
rE[m)
Anng -c)
(T)
Σ
rE[m]
rE[m]
(2) 1,5 26 2 O.
m
rE[m]
Thus focz) has perfect standard accuracy. Finally, we prove that the network is non-robust (Property 3). We consider the following perturbation
p = c(1+c)
13- 11 Σ 11: (1)
Σ JEJ+
k
JEJ-
k
where e is constant such that c- < |J41/J-1
This is say |J+||_-12
1+c
Then, we have
ww i,r x + p) + U1.
(T)
(T)
<(w(() 1 13,) + (w (T)
(T)
E)
2(1 + c)
Σ
(w (T)
T :4y)
2(1 + c)
(w (T)
101,uu) +
k.
JEJ+
JEJ
(1(
1,rdo
+
6
E
6
2(1 + c)
k
1
(T)
6
E
+
2(1 + c)
k
Σ
(2 + - + 3
JE
jEJ_
(T)
2(1 + c)
(3 c)c
X +
L,7do
k
3
JEJI
(T)
2(1 + c)
Σ
3 (2)
(3 + c)e
b
A.
3
JEJ
3(1 + e)|J+1
(3 + c)e
<X
2k
3
(3 + c)E
<
< 0.
3
The first equation expands and p and uses the conclusion in Lemma E010; the second inequality uses the conclusion in Lemma E. 11 and Lemma E.26; the third inequality rearranges the terms and uses the conclusion in Corollary F.5; the fourth inequality uses conclusion in Theorem 4.5.
53

---

Published as a conference paper at ICLR 2025
Reusing the logic of the above inequality, have
(w0 1, + p) + 0-1.r
(T)
(7)
(w 1, Mo0 + (w -1r: 4)
(T)
2(1 +
Σ
(w (wo (T)
1,m
-1:t Me} + 2(1 +c)
Σ
(w (T)
1,
3
JEJ+
JEJ-
>1 (T)
2(1 + c)
(X(T) (T)
) + ((1
Σ
(2)
6
1,ddo
k
k
6
3
JEJ+
jEJ
(1(7)
2(1 + c)
Σ
(3+c)e
-1,rd0
k
3
JEJ
2(1+ c)
(6+
Σ
(T)
-1,n
1,.j
3
jEJ-
2(1 + c)|J-1
(6+ c)6
2 0.
k
3
By combining the two inequalities above, we can obtain that
fo
Σ
RcLU (1w177 +p: + b1
1
Σ
ReLU
(w
(T) -1,; p) bl.,))
m rE/m]
m rElm
This is to say (f (x+ p)) # (forr) (z)), which means Acc ( o(1).
E.3 PROOF OF THEOREM 4.6
Theorem E.27 (Restatement of Theorem 4.6). In the setting of Theorem 4.5,
inf C50
sup -
CCA - forrs (x)) o(1),
where fFA(x) is the feature-averaging network (Definition 4.1).
Proof of Theorem E.27. By Lemma E.25, we have
((100)
ReLU ((w) .x))
ReLU
Σ
gEJ.
<
(w! (T) x)) (T)
Σ
M1 ا -2, x0)
JEJ.
<|||||
(T
((T)
1 |1 -
= =
jeJ.
54

---

Published as a conference paper at ICLR 2025
Thus have
ReLU ((1T7)
ReLU (wi x)) + b1.))
(T)
Σ
H ,
My
M
rE|m]
JEJ+
ReLU (x(r
ReLU (10 m))
(T)
Σ
Σ
|ay
||-
Σ
m rE]m]
m
rE m
SEJI
EE M
ww (2)
x) - (2(T)
Σ
Mi
||1 [,,,))
t
rm(m]
SEJ+
Σ
wir
(T)
(7)
Σ
Mi || 1-2
+c= o1).
rE[m]
JEJ+
Similarly, we have
ReLU (w0T),00) x) + +bG..)) - ReLU ((x07)
My Mylkee ,2) ))
Σ
= =(())
m rE(m]
JEJ-
Combining these inequalities together, we have
((T)
sup
fEA(2) - ferm = o(1).
d
=
55

---

Published as a conference paper at ICLR 2025
F PROOF FOR SECTION 4: FEATURE-DECOULLING REGIME
First, we recall the fine-Grained supervision, multi-Clast network classifier and training algorithm.
Fine-Grained Supervision. Following the setting in Section 3, we consider the binary classification task with data distribution D (((pe)) } 1; Jz). 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 1 E [k] it belongs to, and then train k-class classifier to predict the cluster labels. More specifically,
first sample
a training set S == (( y:) S IR x {+1} from D, along with the cluster labels {ūr } for all
data points. Then k-class neural network classifier is trained on S == {(2,, Bc));= S R x [k].
Multi-Class Network Classifier. We train the following two-layer neural network for the k-class classification mentioned above: (x = (2())) (x(()))).......)))))))))).. fx())
E Rt, where f)( == E h (ur3,, and 0 == (W),1; w1,2..... wk.e) Ih are trainable weights, and == is the width of each sub-network. The outputs (x( are then converted to probabilities (x)( for j E [k]. For predicting the binary label for the original binary classification task on D, we take the difference of the probabilities of the positive and negative classes, i.e., phinary (a=) EjeJ, Px( __ee_
using the softmax function, namely ()(2)) Σi (x))
(z)).. The clean accuracy Accelean (Fhimary and 8-robust accuracy Accrobust [Pbinary, are then defined similarly before.
Training Objective and Gradient Descent. We train the multi-class network Fa to minimize the cross-entropy loss (0) ;= - - log Pú, (z.). Similar to Section 3, gradient descent minimize the loss function Lce (0) (0) with learning rate 17, i... ((++1) ((t) Ve (VeCee(Fgu)) (Fa.)). At
initialization, we set w ~N(0, o Ia) for some ow > O.
(2)
Denote ((t)) Ve 1ax E 1,)++
Σ exp(F (x,))
PEJ
Since many of the proofs in this section are very similar to those Appendix E, we reuse the logic of the proofs and present the key steps.
We also assume that the properties of the training dataset(Propositioon D.6) in Appendix D hold.
F.1 PROPOSITIONS OF NETWORK INITIALIZATION
Proposition F.1. With probability at least 1 - 4hmd- In(d) 22 22-33- we have the following
properties for our network initialization:
For any [m], we have ow o/d-21m(s < ||w.ffl S ow (Vd+2lm(dd)
The proof of Proposition F. is the same as the proof of D.10.
Definition F.2 (Activation Region over Data Input). Let Tard == (i [j: (w.., m;))
x;) + 6 2 0}
be the set of indices of training data points in the j-th cluster which can activate the neuron with weight Sar r at time step O.
Then, we give the following result about the activation region Tay Proposition F.3. Assuming Proposition D.6 and Proposition D. 10 holds. Then with probability at least 1 (hk) 0.00 2hik² exp 9k3ne for all r E [h], s, E J, we have
( n
T 2
31,322"
The proof of this lemma is the same as the proof of Proposition D.12. Lemma E4. Assuming Proposition F.I holds, for all i E I, E J, we have
-1(2; E E.)++
2k < ((0) < 1 E EL) +
2 k
56

---

Published as a conference paper at ICLR 2025
Proof of Lemma F.4. By Proposition F.1 and Property in D.6, for every J, we have
5.00 (1x;)
-
(w
si))
1
Σ
S 40 S In(2).
h
rE[h]
rE[A]
Thus exp(fa(x:)) < Then have
1
S exp(S: (ax))
- k
2k
Σ exp(Fr (0) (2:))
PEJ
= 1(x, E 1,) + exp ff= (0) (x:))
11xx E 1,) +
(() 2k 6,
- E 1,) + k
Σ
exp(Fp (0)
(x,))
PEJ
F.2 ANALYSIS OF TRAINING DYNAMICS
Denote 5/02 = {r E [b] (wer 2.) 2 03 for
Lemma F.5. For every? I, we have (.o(s) = [h].
Proof of Lemma E.5. This proof is similar to the proof of Lemma E.1.
For every i I,r E [h], we have
(w c(s).r* Fi.)
= (w x.) - 17 V
) x.,)
Wen. ((0(0)),
We examine the update term (V Voneca ((0(0),,,))
hn wotsy
L( L(B
hm
Σ
(0)
ie(1)
(un (0)
of Bz)) 20) n.
Σ
Eno00 (ww 1z5) 2 2)) (org, mi)
(0)
jela(s
jeJo(
hm Σ
(1)) j.c(s)
1
(w)
c()).r
(0)
x5) 2 0) (x))
E
(0)) 22000 (1w (() 2z,) 2 2)) (00, x4)1
jETa()
By Lemma E.4 and Lemma F.3, we know that
Σ jEl((a)
(o00) 1 (w) xz) 20) (oz, x2) 2 (1
(0)
}) Tai)) 2 12
2 d
dn
Σ
((()) 1 3.0(1)
(w (0)
((9).) 2z,) 2 0)) (x)1
2nd
k
Combining the two inequalities above, we have
{w (1)
of)).r xz) 2 (w ((),.)
xz) + 17
2A
h 12k-322
7
2 -02 d+ +
2 0.
h
Therefore, have (1)
[h] for every E T.
57

---

Published as a conference paper at ICLR 2025
Recall the definition of weight decomposition we will use in multi-classfficttinn tasks. Lemma F66 (Weight Decomposition). During the training dynamics, there exists the following o for each neuron 8, E J,r [h) such that
coefficient sequences s. and
( D.r
Σ Cow E.
jEJ
Corollary E.7. The coefficient sequences . and (t) F.6 satisfy:
oig for each pair i E I,r,j E J defined in Lemma
ugl 2 == 0 1
0000
s.ny
VEl,
Corollary F88 For all i E 1,r,j E J, we have the following update equation for Nanrá and aani
(t+1)
((w! kpp) 2 0).
Arr Asrr. nh
PEI,
(++1) (t)
(()) 16:1121 (w (c)
( xc. 2 0).
0 5,y srr, nh
(0)
= 0, o (0)
s,rsi =
Corollary F.9. The coefficient sequences defined in Lemma F.6 satisfy:
X and o for each pair 8,j E J, E I,r [M]
(1)
0 iff
0 iff c(i).
Then we reuse the logic the proof of Lemma E.7 to prove main result in our multi-classifictiion setting.
Denote di (0) - [ (x:)). (0)
1 Σ
(1
((l)ino()) + toll d()))a))
h rE(N]
( (i.j) = di 4j
(1)
(1.j) = d. 4j
(t)
1 Σ
(1(0)
d()).ma()) - (10
(e))mol)) + (t)
c(),,, - (()
-
h rE(h]
21n (
1:2 122
Denote max
We know that e okk 2.5) according to our hyper-
VT In
parameter Assumption D.4.
Lemma F.10. Fort To exp( ) i.j Irr [h], we have
k In((n)) S oc(0).n I In(1 + 1)
2k
1.
2n
2. A (i,j) 5ke when c(i) c(j),
3. A (i.j) 3222e,
4. ((9(1)) (i.j) < <3312e,
5. G.00 /100)
(00
jc(s) < 14ke when c(1) c(j),
6. (1) ((0)
/ j.cl)) < 67k2e.
7. (1))
[h].
(.0(1)
58

---

Published as a conference paper at ICLR 2025
10 (.rss S e, sme(()) $ for J, x
8.
Proof of Lemma F.10. Since the proof of this lemma follows exactly the same logic as Lemma E.7, 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 F.11. Assuming the inductive hypotheses hold before time step t, for all s,] e. we have
(1o mb) 5.r. S
6*
Lemma F.12. Assuming the inductive hypotheses hold before time step t. for all r [h] E J, have
(anl) Ei)
0 o S
(1) e
Lemma F.13. Assuming the inductive hypotheses hold before time step t, for all r E [M), E J. have
(w!)) - (10
e
s.rso(f) scr -
Lemma F.14. Assuming the inductive hypotheses hold before time step t, for all E I, we have
(c) 15
9
The proofs of these three lemmas are identical to the proofs of Lemma E.11, Lemma E.12, Lemma E.13 and Lemma E.14 in Appendix Appendix 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 fine-grained k-classification label.
Corollary F15. Assuming the inductive hypotheses hold before time step for all i,j E 1, we have
A - (1.5) 2E.
(
Next, we present the key steps of the auto-balance process for (i.j)
Lemma F.16 (Property 8), Assuming the inductive hypotheses hold before time step t, for E 1,
(1) J,8 # c(i)r r [m], we have (1) ( S e, aaaa I 2e.
Proof of Lemma 1.16. We distinguish between two scenarios.
Case(I): For all E I,. (wo.)) m.) 0.
(1)
Then by Corollary F.8 and inductive hypothesis, we know that
[ (1 S.c.
Case(II): There exists E I, such that (wh)). x.) 2 0. By Lemma F.13, we know that
(un mx) A
(()
S
3
Then noting that (w!), x.) 2 0 and o (t)
aari - 0. we have
s.r.j
(t)
> (w) s,yr
(1 x,) - σ s.r.i
(t)
3
6
2
3
59

---

Published as a conference paper at ICLR 2025
Then using the conclusion in Corollary F.8, we know that
1 (c+1) (1)
ad
(())
nd
6
s,njj A.
mum
m
3
iE
Thus we have (1+1)) 2 -6. Noting that s...j Ar.
I O, we have
1A
ann
S
Then by Corollary F.7,
we know that l0 s.r,i S 2 I 2c.
(6)
Lemma E.17. Assuming the inductive hypotheses hold before time step for all E J,i I, s # cii). we have
fx(x)) S.
Proof. By Lemma F.13 and Corollary F.9, know that
fa(w:)
1
Σ
ReLU((W
e az))
h
rE[M]
Kwo
at))
h
rE
+ (e)
(..ro(()) + +o.2,))
h
rE[M]
rE(M]
Lemma F.18. Assuming the inductive hypotheses hold before time step t. for any two training data points Hi Fj. if 2 qj. have
(1.) +
((0)
1.0())
(1d)
1+ 2e
k
(())
(,((s)
Proof of Lemma F.18. By Lemma F.17 and noting that exp(e) $ $1+ 2e, we know that
Σ exp(Ft (x:))
k
-(10) (.(())
pfc(1)
(k 1) exp(c)
(K - 1)(1 22)
<
exp(() exp(g (1)) 5
k + exp(1:")
(11)
Σ exp
1! (xx;)
(k. (k-1) 1)
1+ -
PEJ
Thus we know that
1 k - 1 + exp(x
(1)
(t)
h 1 + exp(g:")
< (1 + 2c)
(1 + 2c) k -1+exp(1)""
(0)
(())
カ- 1 + exp(g))
(1(1))
k - + exp
explg. (1) - exp(1)") (1)
+ exp (1, ())
h - 1+ --1+exp(1""" (1)
exp(1,)) (0)
(oxp (t.)) (xx.)))()))))))) -1)
=1 -- 1 + expla)")
60

---

Published as a conference paper at ICLR 2025
For the second term,
exp ( (1) (6.3)
-1
exp(d,") (0)
exp (4. (1.3) -1) s exp (4. 1.
k
k + exp(o!")
Thus we know
(1,)
(i,) + k
(()
(,(())
Se 0 ( (1+
1+ 2e
k
(()
1,0(s)
We first consider the case when c(i) c(j). We distinguish between two scenarios, one is when
A (1) (i.5) is relatively small and the other is when A (i.5) is relatively large.
Case(I): A (i, S 4kie.
In this case, have A (i,)) S 5kt due to small learning rate 7).
Case A (1, 2 4kc.
By Lemma F. F.14, know that
(1,5)) (1,5) 2c 2 3kc.
By Lemma F.18, we know that
((t)
e (6.5) + k
Je((i)
(1+ (1)
1 + 2e
k
1 + 2e
1,0(1)
Noting that e(i) = c(j), we know that
A (1,)) (1,5))
g(t) nh (e - ([,,1
Then due the inductive hypothesis,
(++ 1) (1,5) < (i.)) S 5ke.
By Corollary F.15, we can get Property 4
A (i.j) < ( (i.j) 2e < 6ke.
By Lemma F.18 and noting that e" <1+ + 2x for small x we know that
(1+1))
j.((
20)e ( < (1 21) + 22 2A(++) < 11 14ke. < <11 21+22e)
((++1
(,(())
Next, we consider the case when c(i) clj). We also distinguish between the two scenarios.
Case(T): A (i.j) < 3.. In this case, we have Á ((+1) (i, & 321.2 due to small learning rate
Case(I)): A (i.j) 2 31A2e.
By Lemma F.14, we know that
((.)) ((( 2c 2 30.
61

---

Published as a conference paper at ICLR 2025
By Lemma F.18, know that
p(te)
(b.)) + k -1
(1 + 4 (i,5)/k 2 1+ 29ke.
((t)
+ 2E
k
1 + 2E
Le(())
Furthermore, due to the inductive hypothesis, for any p E Ic(s): 9 E Io(y), we know that
((())
p.c(p)
(1+ 14c)
2,((()
(1 + 14kc)²
(())
((8)
/1 )
1 + 29ke
1 +
9.(())
We know that
A A(41)
(i,j)
(((t)
Me(s)) 3
Σ
plets) (1)))
(((t) (oct))lE:1 1,0 112
-
rih
p.clf)
pete(s)
(c10 (1ot))lls:11
((e)) (1ot)llss11
1
Σ
(())
p.ote) p.cce) Mc(c))
E
nh
nh
pelois)
pElcoi
<0.
By inductive hypothesis, we know that
A (1,)) (.) S 32
Now, 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 F.18.
By Corollary F.15, we can get Property
A (i.j) - (1 + 2e < 336
By Lemma E.18 and noting that e" 11 2x for small x we know that
(((+1)
S (1+
(0t (t,)
(1 + 21) + 22 24 (i.)) S 1
[[+1)
(,(())
Lemma F.19, For all i E 1,r1,72 E [h], we have
- (1-200k2c claalirasa nd S (1)
clh)m $ (1 + 200k2e) (8)
1
c(12).r2.z nd.
Lemma F220. For every e I,ro E [h], we have
_ (0)
2E < qi
(1)
2n (t)
+ De.
c(0).ooa
The proof of these lemmas are the same as the proof of Lemma E.19 and Lemma E.20.
Lemma F21. For every E I,r [h], we have
2
1
exp
2n (t)
k
1,r5
-2(()
p(1)
(,((1)
S 2k exp
(- n (t)
2h 1,r:a
62

---

Published as a conference paper at ICLR 2025
Proof of Lemma F.21. By Lemma F.20, we know that
k 1 - exp
qi
k-1
2
2n (2)
4: 1 + exp
k
1,r,
+ 2c
k: 1
2n (0)
h 1 - 2 eexp
k Iinni
1 exp
(-x
2n (1)
k l.nes
By Lemma E.17 and Lemma F.20, we have
(A: - 1) exp(()
I
L.(()
(k - 1)exp(e) + exp
(t)
k-1
S (1+ + 2E)
Y1
(t)
k - 1 + exp
2k
1,r,
2e
2(k: -1)
<
k + exp
(t)
2k
n2 (t)
S 2k exp
2h
1.r,a
Next, we prove Property 1.
(t+1)
- (1)
c(f).n nh (d()) EEll
(()
(1)
2knd
n (t)
c(0).n
nh exp(
2k (())..)))
2k
2k 1
In(t + 1)
n 26
2k
< In(1 + 2)
(++)
(1)
11 (t)
12
di),
d(),
nh 1,((3)
(0)
+ nd
k (2)
(()).r
2nh exp(
2n ..)
k: 2
In(tn) +
2n
2n
In((t + 1)m)
2n
Finally, we prove Property 7.
By Property 6 in the inductive hypotheses, we know that for all i E I,
1
(0)
1,0(1) S ((()))
Go0g S -11 67k e(1) 1,0(1))
+ 67122
63

---

Published as a conference paper at ICLR 2025
Then know that
-V ((0(0),
2- Σ
((t)
Σ
(e(())
p.ote) p.ott) (2
Pio(s)
pElets)
péta(a
Σ
(())
(pp 2.))
Σ
plo(e) Kapp 2i:1
Pio(p)
pElo()
n fate
(1)
d
11+
1.c(1)
2(1 + 67h2c)
pEteto
(t) (,((I)
dn
(1+ 67k
20.
4k:(1 + 67h2e)
By Property 7 in the inductive hypotheses, we know that (w (t+1) We complete the proof of Lemma F.10.
o(i) (xi) 2 (w (2)
c(s)s
xii) 20.
F.3 PROOF OF THEOREM 4.7
Theorem F.22 (Restatement of Theorem 4.7). In the setting of training multi-class network
the multiple classification problem $ == {(2,, b:)1 S Rd x [k] as described in the above, under Assumptions 3.2, 3.3 and 4.3, for some iterations, with probability at least 1 1- "/ the neural network satisfies the following properties:
o1), after 1i(n-1 < T S exp(O(x)/2))
1. The clean accuracy is nearly perfect:
Acc (1 Login
2 1 111-exp(-((eee) d)).
2. The network converges to the feature- dectuul-d regime: there exists time-variant coeffi- cient ((T) E [22(lgg k), +oo) such that for all [k], r E [h], the weight vector w (T) be approximated as
can
(T) W ( 112 S 1/2-
3. Consequently the corresponding binary classifier achieves optimal robustness: for perturbation radius 5 O(vd), the 8-robust accuracy is also nearly perfect, iee Acc Acc (F EEhinary, ;8) 2 1 - c1p1-ppp----- (exp(-Q((lo)) ((glog?d)) d)).
g(T)
Proof of Theorem Theorem F222. We first prove that the network converges to the feature-decoupling regime(Property 2).
Lemma E223. For all j E J,r (M), we have In(TT)
(T) S 41n(T + 1).
4
Proof of Lemma F223, Using Property in Lemma F.10 and Corollary F.7, we know that
|
2
(n(T))
X 3,r
= Σ
(T)
[4 2
j.m.
4
PET,
|
1 1 = Σ
(T)
< 4n + +1)
| 2
j.rpp
pEl,
Lemma F24. For T1, E [1]. E J, we have
(T)
< + 204
(2)
jairada
64

---

Published as a conference paper at ICLR 2025
Proof. By Lemma F.19, know that for any E 11 12 L/2
Jurii I (1 + 20002e)
har2id + (nd) $ + 2012? () 2..22.2"
Σ
||Ep
(t)
(7)
PEln1
(2)
Σ
115p
J2.T2.J
Ja.12.P
pEl
|| + Ind] 1
$ (1 201k è)
|| - Ind]]l 57
< (1+ 2014 e1 + +e((te))
< + 2041
We denote X(T) -1,,11 as the representative of (Arn : [m].j E J..
By Lemma F.23 and Lemma F.24 ffor all [h].j E J, have
(XTT) - / - 204k2 EA EA
x v 4 4nn(T + 1),
\ In(TT))
2 2 ln()) (((og(())).
4
Lemma F.25, For all 8 E J,r [h]. We have
va w A(T)
pe, 10 o(1)
Proof of Lemma E.25. Recall weight decomposition in Lemma E.2.
Vd ww!?
X Malllaal11
Vdw (0) 8.
+
vd
(T) s
13l1pas11
C1
C2
s.n
(T)
+
(T)
Ei
jtt
C4
For L1 term, using the conclusion in Lemma D.10, we know that
Vdw ( 2do S ===
For L2 term, using the conclusion in Lemma F:23, Lemma F.24, we know that
Va (T) A.r.s x(T) u=ll
(T s.r.s X 20442 z 8 In(T + 1) s c=
For Ly term, by Lemma F.16 and triangle inequality, we know that
Va va>x
ske= o(1)
65

---

Published as a conference paper at ICLR 2025
For L1 term, by Property (1) in Lemma E.7, we have
Vi S:
0
(2))
2
||E:1|
2
+ d
Σ
0
s,rh s.rdi (Si, 1
(T) (T)
iEl
VEV
(T)
24
(2) (2)
+
d
iE!
$1712
8h In (T +1) 1)
8k
In (T + 1)A
d
8k
8K A
S 16kc d(1)
n
Combining the above together, we know that
vd
w ( - 1
p. well = o1)).
Then we prove that the the clean accuracy is nearly perfect(Property 1).
Assume (x, y) is randomly sampled from the data distribution 'D. Without loss of generality, we assume that pe; + E. = j.
Lemma F22. Let g N0, Ta)- Then, with probability least 1 1-2nd 1n(d/2 for foralls E J,r E [M] we have
(w(T)
(T)
The proof of this lemma is the same as the proof of Lemma E.26.
Using the conclusion in Lemma F.13 and Lemma F.26, we know that for s E J,r E [M]. Fj
(T) (w x) = (w srr 14)) 8, 5) s sn.
(T)
6
3
(T) (w) J.r x) = (w) (T)
14) (w) j,r E) jin
3
3
Thus know that (x) max.ex mf (z)) So Forry has standard perfect accuracy. Finally, we prove that the corresponding binary classifier achieves optimal rorostys 3).
By Lemma F.25, we know that vdlww!!
w. (T) S (()) + o(1) - (T)
Vd
Then for any perturbation P with p 5
10"
We know that
(w (T) j.r x + p) = (w (T)
14)) + (w
(T) j.r
-) + (10) p) 100 (T)
2 find
(T)
c
||||||
6 6
(2
For E J, 8 + j. we know that
(w() (T)
x + p) = (w (T)
m4) + (w..
.) + (wi7.. (T) p)
< (T) + 6 +
<
+ -| (
|||p||
s,r.j
6
6
(T)
Set
+ 5
3
(1')
In(k).
66

---

Published as a conference paper at ICLR 2025
3A(T)
32(T)
Thus we know that fz + p) 2
and fz + p) =
4 -l((k).
let (ax) denote the numerator of Frinan
(z), where denominator is Ever efulz)) We know
sgn( (g sgn(G).
Thus have
G(x + p) = Σ cxp (fz(2 + p)) Σ
exp (fj(a + p))
JEJ
JEJ
2 exp (31 /4)
Σ
(21 /4 - In(())
exp
JEJ
0.
That is to say sgn(Gg (G(x p)) = sgn(G(z)), which means Feirs is robust under any perturbation
Vd
with radius smaller than
10
67

---

Published as a conference paper at ICLR 2025
G TWO FEATURE LEARNING REGIMES: FEATURE AVERAGING AND FEATURE DECOUPLING
In this section, we present two distinct parameter regimes for two-layer network learner: feature averaging and and feature decoupling. The former means the weights associated with each neuron is 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. (2024) and Min & Vidal (2024), We illustrate how feature averaging solution leads to non-robustness, while feature decoupling solution exists and is more robust (w.r.t to much larger robust radius).
G.1 FEATURE-AVERAGING TWO-LAYER NEURAL NETWORK
Now, we begin by presenting the following example of feature-averaging two-layer neural network, which is more general version (including bias term) than the one we mentioned in Definition 4.1.
Feature-A Aveerriig Two-Layer Neural Network. Consider 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):
fo (x) ReLU
(
Σ Mj + +b+) ) jEJ+
ReLU (< Σ
M.. 13 (En, iEJ_
+ + )
deals with all positive clusters
deals with all negutive clusters
where we choose weight Ws,r = Eje,, My for 8 E (-1,+13,r E [m] and bias b = by for E (-1,++1,r E [m].
Indeed, the feature-averaging network uses the first neuron 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 G.1. There exist values of b+ and b_ such that the feature-averaging network fe... achieves 1 - 11)( standard accuracy over DD
Proof of Theorem G.I. Let b4 = b_ = 0, and then we know, for data point (æ = aj,, +. D within cluster i (w.Lo.g. assume cluster is positive cluster), with high probability, it holds that
fo (x) 2 (p: ap,, +
Σ (pz, E)
Σ (u
jEJ (0)
jEJ
(()) - 0 ( (koVVll(d)) 20, 0,
which implies that fe.., correctly classifies data (z, y) with high probability.
However, it fails to robustly classify perturbed data matter what the bias term is, shown in the following theorem.
Theorem G.2. For any values of b4 and b_ such that fe.., has 1- o(1) standard accuracy, holds that the feature- network fo... has zero r-robsst accuracy for perturbation radius 8 (V/k)).
Proof of Theorem G.2. Indeed, can choose the adversarial attack
- -EjE, μ
Eier_ fi and ||e]| 8. Then, for averaged features ws,r ΣjeJ, My, this perturbation can activate almost all of ReLU neurons, which w.h.p. leads a linearization over the perturbation p
Se_,, (ar p) So.., (x)) + <V (x), p)-
Since fe.., has 1 - o()) standard accuracy, we know that the bias term satisfy that 64, b_ = O(d). which manifests that the classifier achieves positive margin, i.c.
< y (x) < Od),,
w.h.p. over (æ, y) sampled from D.
68

---

va Published as a conference paper at ICLR 2025
Then, due to a large gradient norm data imput, i.e.
V V (x))
Σ
p4 E
=ulkdd((Vkd),
JEJ+
1EJ_
we derive that the feature-averaging network fe.,, has zero 8 robust accuracy for perturbation radius
5 = l(Vdk)).
G.2 ROBUST TWO-LAYER NEURAL NETWORK EXISTS
In this section, we show a robust two-layer network exists for D, using a similar construction in Frei et al. (2024).
Theorem G.3. There exists a two-layer network fo.. that is 3
-robustt for D.
Proof Theorem G.3. The construction similar to that in Frei et (2024). We define So.. _ RR a network that represents a positive constant times the following function:
d
Σ ReLU (IuI 2))
J0 (m)
Σ 1EJ+
RReLU
ReLU
(11),
2
LEJ_
In particular, we set two-layer width-k: ReLU network with Waj = 1 1 -1 (j E J,) & for E {+1} and E [A:]
Jo) Jo) pg. bed
In this network, each neuron ReLU((()) x+) 9) (or ReLU((M,, 4) deals with one certain positive cluster j (or negative cluster 1), and we also apply the bias term to filter out intra/inter cluster noise. In this regime, for each data point (æ. y) belonging to cluster (we assume cluster is
positive cluster and y = 1) and any perturbation Clell < 3 w.h.p.
we have the following linearization,
fem (xc + (p) - -11,, x + p)-
Then, we know the network Sane has - o(1 8-robust accuracy for 8 <
Vd 3
Note that fe., leverages individual decoupled features, which is natural and robust solution to the binary classification on D. In fact, one can easily verify that the robustness of fea. is optimal up to constant factor, as the distance between distinct cluster centers is O ((vd), i.e. ||p, - 1,l ((dd). for all i * j. However, as we show in our main result that gradient descent does not learn this feature-decoupled network directly from D, and instead converges different solution that is (vK)) times less robust.
G.3 NOB-ROBUST MULTI-CLASS NETWORK EXISTS
Similar the feature-averaging binary-class network as that mentioned in Definition 4.1, the non-robust multi-class network also exists, which is shown as the following proposition.
Theorem G.4 (Restatement of Proposition 4.8), Consider the following multi-class network Fg: for all E [k), the sub-network fq has only single neuron (h 1) and is defined as fz(=)) ReLU ((r EieJ, m1,z)),, where cluster has binary label & {+1). With prob-
ability at least - exp(-S((lo' d)) over S, we have that Lce(() < exp(-9(d)) = o(1), where θ denotes the weights of Fø- Moreover, Acc (Fblnary Acc (phinary,s( (S( < exp(=S((log' d)).
1 exp( exp((S((log' d)),
Proof of Theorem G.4. Consider data point (x. that is randomly sampled from the data distribution D. Without loss of generality, assume that æ Mjo + E.jo E J4. Reusing the argument of proof of Property (3) and (4) in Proposition D.6. We know that, with probability at least 2kd In(d//2 for all j E J, we have
Kpg t < A.
69

---

Published as a conference paper at ICLR 2025
First, prove the network has perfect clean accuracy when the above properties hold. Indeed, calculate the output value each sub- sub-network as follows.
(x) ReLU
E M1, LEJ1
/
2 2d kA.
For j E J # jo.
(11++ (1 M.
f5(z) ReLU
S <d+ka KA
2d - A (z))
LEJ.
Thus, we know that Fa(Ø) = jo with probability at least 2kd 1n(d)/2
Then, with probability at least sep(--2(log' d)) over sampled from D, for all i E 7, we have exp(ff, (zx))
log Ps. (xx) log Σj{(M) exp(f
5 log - exp(-((d))) S exp())((d)),
2) 2 02(2) for sufficiently small z. Therefore, we where the last inequality holds due to log(1 derive that
(C Σ log Pa. (z,) exp(- o(1).
Finally, we prove that the network has at most 2kd Ind)/2 robust test accuracy against perturbation radius 8 = Q2 /d/k).
3(1 c)
Consider perturbation For any j J4, we know that
Σ
11 Σ )
k
IEJ+
LEJ
μ4 +
Σ
11,00 00 =
|4j + Σ
Mis Mijo + -
LEJ+
LEJ+
3(1 + c)d
$ 2d (k + 1)
+ k
< 0.
For any j J_, we know that
Mi + Σ
M1, x - P))
=
μij + Σ
ul: Hs + -
IEJ
(EJ4
k
3(1 c)d
d+ (k 1)A
I+c
>0
This is to say for any j E J4,
(11= +
M1, -0)
15( - pp) PRRU
Σ
VEJ+
For any j J_.
fz p) ReLU ((m (n1+ + Σ
p4l x -0000--))>
JEJ
F.innary, ;8) E exp( (olpp=2(log2 d)) Thus, we obtain that Ace (F
70

---

Published as a conference paper at ICLR 2025
H ADDITIONAL REAL- WEALDWORLD EXPERIMENTS
H.1 EMPIRICAL VERIFICATION OF FEATURE LEARNING PROCESS
Beyond verifying the alignment between our theoretical findings and the results of numerical simu- lation on the synthetic multi-cluster data setup, as described in Section 3, we also consider a more realistic setting where the multi-cluster structure of data naturally occurs.
Transfer Learning Based on CLIP Model. Here, we focus on transfer learning setting, under which we utilize pre-trained CLIP ViT-B-32 model (Radford et al., 2021) 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 experimental results in Figure 12, where we verify the orthogonality of the extracted features (i.c.. image embeddings of CLIP model) by calculating the correlation between them.
Binary Classification on C.FAR. We create a 2-classification task from the CIFAR-10 10 dataset by merging the first 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 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., [+)): := Σ (e 2)). Vj E [2] or [10], and z denotes the image embedding). We set the width of the first layer to be 1000 (h- 10) to ensure that the accuracy of the pre-trained model was not compromised. For 10-classification, we use w := b Wi,r as the equivalent weight of fz. For the 500 positive weights and 500 negative weights in the binary classification network, we equally divide them into 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 F := f2......... f10) E R 10 and each sub-network ƒ/ corresponds to weight vectors w0..
See experiment results in Figure 13. Indeed, deep neural networks can be viewed as consisting of two parts: feature extractor that is a mapping from the input space to the latent space, and a shallow classifier that predicts the classification results based on the extracted features (in the latent space). Feature averaging means that the shallow classifier mixes extracted features corresponding to different classes in the latent space. And our empirical results verify this.
Experiment Results.
H.2 INFINITY-NORM CASE
Here, we aim to verify whether the model trained with fine-grained supervision information (i.e., 10-class labels) is more -robustt compared to the model trained with only binary (2-class) labels.
Experiment Settings. To ensure fairness in the comparison, we sum the logits corresponding to the 5 positive classes and subtracted the sum of the logits corresponding to the 5 negative classes from the 10-class model's output. This result is used as the binary classification output for the 10-class model. The robust accuracy measured by using the standard PGD attacks (Madry et al., 2018) with different lo pertubatioon radius. We run experiments in the following datasets:
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 (MNIST: parity-classfficatin; CIFAR-10: binary-llassffiattion as that we mentioned in Section 5.1).
Experiment Results. The results are presented in Figure 14. 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, which empirically shows that fine-grained supervision also improves foo robustness.
71

---

Published as a conference paper at ICLR 2025
100
200
0.00
600
0.25
0.50
800
200
400
500
800
Figure 12: Verifying Orthogonal Condition: We plot the extracted feature correlation as a colormap, where each pixel represents some cos(z;, Z,, 2j) between extracted features 2;, Zj of two data æ., z; from CIFAR-10 training dataset (here, we sample 100 instances for each class).
4
0.6
0.1
Us
M3
0.4
44
4s
M5
0.0
44
46
0.2
447
μ7
14
-0.1
0.0
H4
μg
M10
P10
0.2
w: W, W. W., W. We W,, We WgW10
W; W; W3 W. W. W. W. W. WgW10
(a) CIFAR-10: 2-class
(b) CIFAR-10: 10-class
Figure 13: Illustration of feature averaging and feature decoupling on CIFAR-10 dataset. Figure (a) corresponds models trained using 2-class labels, and Figure (b) corresponds models trained using 10-class labels, respectively. Each element in the matrix, located at position (i.j), represents the average cosine value of the angle between the feature vector μe of the i-th feature and the equivalent weight vector w, of the (-).
MNIST
CIFAR-10
100
Robust Test Accuracy
Robust Test Accuracy
60
60
>0
10
10
70-
2
33
51055
Perturbation Radius (L.))
Perturbation Radius(L=)
Figure 14: Verifying robustness improvement in lac case: We compare Eo adversarial robustness
between model trained by 2-class labels (red line) and model trained by 10-class labels (blue line) on MNIST (the left) and CIFAR-10 (the right).
72


Reviews:
[{"review_id": "OJ9VP4DBuk", "reviewer": "Reviewer_DJdJ", "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.", "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?\n\n2. 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:\n a) How should we define the feature mixing in those cases?\n 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?\n\n3. 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. \nI list a few papers for reference:\nXu, Zhi-Qin John, et al. \"Frequency principle: Fourier analysis sheds light on deep neural networks.\" arXiv preprint arXiv:1901.06523 (2019).\nXu, 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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AHSATGrsDp", "reviewer": "Reviewer_h25y", "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.", "questions": "-Minor issues in notations: Big $O$ notation on line 185 is different from the others; $Theta$ notation is not defined.\n\n-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?\n\n-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$?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 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.\n\n-This paper is generally well-written and logically organized with clear notations, assumptions, and detailed references to prior works with comparisons\n\n-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.", "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. \n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ek4Bq9VBmG", "reviewer": "Reviewer_dDY4", "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. \nUnder 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.", "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. \n2. 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The theoretical analysis of the finite time gradient descent dynamics is thorough in the feature learning regime. \n* This paper is an interesting contribution in understanding the robustness of neural networks from optimization perspective. The analysis also aided in proving a conjecture. \n* The paper is well written and the presentation is clear.", "weaknesses": "* The assumptions on the data (orthonormal equinorm and multi-cluster) is restrictive.\n* The proof is detailed and thorough for two layer ReLU with only one learnable hidden layer. Extending to deeper networks seems challenging.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KudbXtcy6c", "reviewer": "Reviewer_MXLz", "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.", "questions": "Q1: Can the fact that the network reaches an \"averaging\" solution be deduced from the KKT equations in Frei et al. (2022)?\n\nQ2: 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.\n\nQ3: 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.\n\nQ4: 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?\n\nQ5: 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].\n\n[1]. G. Goh. A discussion of ’adversarial examples are not bugs, they are features’: Two examples of useful, non-robust features. Distill, 2019.\n[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.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This is a good paper, and it was a pleasure to read and review. I would like to highlight the following contributions:\n\n- Important topic: The interplay between the implicit bias of optimization and robustness is, in my opinion, an important area of study.\n- 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).\n- 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.", "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).\n\n2) 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).\n\n2) 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.\n\n4) A list of non-significant typos/grammatical errors I spotted:\n- line 091: features-> feature\n- line 107: exists-> exist\n- line 167: analyze -> analysis\n- line 261: neural -> neuron\n- line 893: echos -> echoes\n- line 1098: We -> we\n- line 1100: margin -> margins\n- lines 1307-1310 require rephrasing (they do not make sense),", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)