Simplicity Bias in 1-Hidden Layer Neural Networks
Authors: Depen Morwani
jatin batra
Prateek Jain
Praneeth Netrapalli
Venue: NeurIPS 2023 poster (2023)
Keywords: Simplicity Bias
Gradient Descent
Implicit Bias
Neural Networks
TL;DR: Gradient Descent on 1-hidden-layer neural network learns a function of essentially a lower dimensional projection of the input.

Abstract: Recent works have demonstrated that neural networks exhibit extreme *simplicity bias* (SB). That is,  they learn *only the simplest* features  to solve a task at hand, even in the presence of other, more robust but more complex features. Due to the lack of a general and rigorous definition of *features*, these works showcase SB on *semi-synthetic* datasets such as Color-MNIST , MNIST-CIFAR where
 defining features is relatively easier. 

In this work, we rigorously define as well as thoroughly establish SB for *one hidden layer* neural networks in the infinite width regime. More concretely, (i) we define SB as the network essentially being a function of a low dimensional projection of the inputs 
(ii) theoretically, we show that when the data is linearly separable, the network primarily depends on only the linearly separable ($1$-dimensional) subspace even in the presence of an arbitrarily large number of other, more complex features which could have led to a significantly more robust classifier,  (iii) empirically, we show that models trained on *real* datasets such as Imagenet and Waterbirds-Landbirds indeed depend on a low dimensional projection of the inputs, thereby demonstrating SB on these datasets, iv) finally, we present a natural ensemble approach that encourages diversity in  models by training successive models on features not used by earlier models, and demonstrate that it yields models that are significantly more robust to Gaussian noise.

Simplicity Bias in 1-Hidden Layer Neural Networks
Depen Morwani Department of Computer Science Harvard University dmorwani@g wwh.aarrdddd.. edu
Jatin Batra
School of Technology and Computer Science
Tata Institute of Fundamental Research (TIFR)
jatin. batrattifr..ra.......e in
Prateek Jain' Google Research prajainggoogle.com
Praneeth Netrapalli'
Google Research
pnetrapalli@gode com
Abstract
Recent works (Shah et al., 2020 Chen 2021 have demonstrated that neural networks exhibit extreme simplicity bias (SB) That is. they learn only the simplest features to solve a task hand, even in the presence of other, more robust but more complex features. Due to the lack of general and rigorous definition of features, these works showcase SB on semi-syathetic datasets such as CIFAR where defining features is relatively easier.
MNIST-
In this work, rigorously define well as thoroughly establish SB for one hidden layer neural networks. More concretely, (i) we define SB as the network essentially being function of a low dimensional projection of the inputs (ii) theoretically, in the infinite width regime, we show that when the data linearly separable, the network primarily depends even in the presence of an arbitrarily large number of other, more complex features which could have led to a significantly more robust classifier, (iii) empirically, we show that models trained on real datasets such as Imagenet and Waterbirds- Landbirds indeed depend demonstrating SB on these datasets, iv) finally, we present natural ensemble approach that encourages diversity in models by training successive models on features not used by carlier models, and demonstrate that it yields models that are significantly more robust to Gaussian noise.
only the linearly separable (1-dimensional) subspace
low dimensional projection of the inputs, thereby
Introduction
It is well known that neural networks (NNs) vulnerable to distribution shifts well to adversarial examples (Szegedy et al.] 2014, Hendrycks et al. 2021). A recent line of work (Geirhos) al. 2018 (Shah al. [2020] Geirhos al. 2020) proposes that Simplicity Bias (SB) (or shortcut Tearning) i.e., the tendency of neural networks (NNs) to learn only the simplest features over other useful but more complex features, is key reason behind non-robustness of the trained networks. The argument is roughly as follows: for example, in the classification of swans vs bears, as illustrated in Figure there are many features such background, color of the animal, shape of the animal etc. that can be used for classification. However using only one or few of them can lead to models that are robust specific distribution shifts, while using all the features can lead more robust models.
Several recent works have demonstrated SB on variety of semi-real constructed datasets (Geirhos) ct al. 2018 Shah et al. 2020 Chen et al. 2021 and have hypothesized SB to be the key reason
part of the work done while Google Research, India Alphabetical Ordering Image source: Wikipedia swa bea
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

for NN's brittleness to distribution shifts (Shah et al. 2020). However, such observations are still only for specific semi-real datasets, and a general method that can identify SB on given dataset and given model is still missing in literature. Such a method would be useful not only to estimate the robustness of a model but could also help in designing more robust models.
A key challenge in designing such general method identify (and potentially fix) SB is that the notion feature itself is vague and lacks rigorous definition. Existing works Geirhos et al. (2018): Shah et al. (2020); Chen et al. (2021) avoid this challenge of vague feature defi- nition by using carefully designed datasets (c.g., concatenation of MNIST images and CIFAR images), where certain high level features (e.g.. MNIST features and CIFAR features, shape and texture features) are already baked in the dataset definition, and arguing about their simplicity is intuitively easy.
Classification of swans vs bears. There are several features such as background, color of the animal, shape of the animal etc., each of which is sufficient for classification but using all of them will lead more robust model.
Figure 1:
Contributions:
Our first contribution is to pro- vide a precise defimition of particular simplicity
bias- LD-SB- referring to low dimensional input dependence of the model.
Definition 1.1 (LD-SB). A model f R.d - Re with inputs x E Id and outputs f(x) E R (e.g..
logits for с classes), trained on distribution (x, v) D satisfies LD-SB if there exists projection matrix P E pddd satisfying:
rank (P) < d.
IPprred[ + P1 22)) = (xx)((( 1 - for (x(1), (0)), (z((2, 1(2)) ~ D, where pred(f(z)) represents the predicted label for
* An independent model g trained on (Piæ, y) where (I, ~ D satisfies |Acc(g) - Acc(f) S ez.
for some small e and 62. Here P1 is the projection matrix onto the subspace orthogonal to P, and A/e( represents the accuracy of f.
In words, LD-SB says that there exists a small k-dimensional subspace (given by the projection matrix P) in the input space Rd, which is the only thing that the model f considers in labeling any input point æ. In particular, if we mix two data points 2 and x2 by using the projection of x onto P and the projection of 2/2 onto the orthogonal subspace Pi, the output of f on this mixed point Pxi + Pi P2 is the same as that on I1. This would have been fine if the subspace Pt does not contain any feature useful for classification. However, the third bullet point says that P1 indeed contains features that are useful for classification since an independent model g trained on (Pix, v) achieves high accuracy.
Theoretically. we demonstrate LD-SB of 1-hidden layer NNs in the infinite width limit for a fairly general class of distributions called independent features model (IFM), where the features (i.e., coordinates) are distributed independently conditioned on the label. IFM has long history and is widely studied, especially in the context of naive- -Baee classifiers Lewis (1998). For IFM. we show that as long as there is even a single feature in which the data is linearly separable, NNs trained using SGD will learn models that rely almost exclusively on this linearly separable feature, even when there are an arbitrarily large number of features in which the data is separable but with a non-linear boundary. Empirically, we demonstrate LD-SB on three real world datasets: binary and multiclass version of Imagenette (FastAl 2021), vaterbirds-landdbin
(Sagawa et al., (020a)) as well as the ImageNet (Deng al. 2009) dataset. Compared the results in Shah al. (222)) our results (i)
theoretically show LD-SB SB in a fairly general setting and (ii) empirically show LD-SB on real datasets.
Finally, building upon these insights, we propose simple ensemble method OrthoP that sequentially constructs NNs by projecting out the input data directions that are used by previous NNs. We demonstrate that this method can lead significantly more robust ensembles for real-world datasets in presence of simple distribution shifts like Gaussian noise.
Why study 1-hidden layer networks in the infinite width regime?

---

1. From a practical standpoint, the dominant paradigm in machine learning right now is to pretrain large models on large amounts of data and then finetune on small target datasets. Given the large and diverse pretraining data seen by these models, it has been observed that they do learn rich features (Rosenfeld et al. 2022 Nasery et al. 2022). However, finetuning on target datasets might not utilize all the features in the pretrained model. Consequently, approaches that can train robust finetuning heads (such as a 1-hidden layer network on top) can be quite effective.
2. From a theoretical standpoint, there have been several works that analyze training dynamics of finite width networks (Ding et al., 2022), and show convergence to global minima on the training data. However, these results do not identify which among the many global minima, the training dynamics converge to, which is crucial in determining the nature of SB of the converged model. Such precise characterization of the final convergence point is known only for infinite width 1 1-hidden layer networks (Chizat et al. 2019 Chizat & Bach 2020).
3. While our theoretical analysis works in the setting of infinite width networks, our extensive experiments on several large scale datasets suggest that the results continue to hold even for finite width networks. Furthermore, Vyas al. (2023) show that the behavior of neural networks remains consistent with width in the feature learning regime.
To summarize, this paper characterizes the nature of SB in 1-hidden layer networks, and also proposes a novel ensemble training approach, called OrthoP, which leads to more robust ensembles. While the theoretical results in the infinite width regime, empirical results on several real world datasets show that these results continue to hold even for finite width networks.
Paper organization: This paper is organized as follows. Section 2 presents related work. Section presents preliminaries. Our main results on LD-SB are presented in Section 4 Section presents results on training diverse classifiers. We conclude in Section
Related Work
Simplicity Bias: Subsequent to Shah et al. (2020), there have been several papers investigating the presence/absence of SB in various networks as well as reasons behind SB (Scimeca et al. (021). Of these, Huh et al. (202)) is the most closely related work to ours. Huh et al. (2021) empirically observe that on certain synthetic datasets, the embeddings of NNs both at initialization as well as after training have a low rank structure. In contrast, we prove LD-SB theoretically on the IFM model as well as empirically validate this real datasets. Furthermore, our results show that while the network weights exhibit low rank structure in the rich regime (see Section 3.2 for definition). the manifestation of LD-SB is far more subtle in lazy regime. Moreover, also show how to train a second diverse model and combine it to obtain a robust ensemble. Gaalanti & Poggio (2022) provide theoretical intuition behind the relation between various hyperparameters (such as learning rate, batch size etc.) and rank of learnt weight matrices, and demonstrate it empirically, Pezeshki ee al. (202T) propose that gradient starvation the beginning of training is potential reason for SB in the lazy/NTK regime but the conditions are hard to interpret. In contrast, our results are shown for any dataset in the IFM model in the rich regime of training. Finally, Lyu et al. (2021) consider anti-symmetric datasets and show that single hidden layer input homogeneous networks (1.c., without bias parameters) converge to linear classifiers. However, our results hold for general datasets and do not require input homogeneity.
LDSBB
Learning diverse classifiers: There have been several works that attempt learn diverse classifiers. Most works try to learn such models by ensuring that the input gradients of these models do not align (Ross& & Doshi- Velez 2018 Teney al. 2022). Xu al. (2022) propose a way learn diverse/orthogonal classifiers under the assumption that a complete classifier, that uses all features is available, and demonstrates its utility for various downstream tasks such as style transfer. Lee et al. (2022) learn diverse classifiers by enforcing diversity on unlabeled target data.
Spurious correlations:
There has been large body of work which identifies reasons for spurious correlations in NNs (Sagawa et al. (020b) as well as proposing algorithmic fixes in different settings (Liu et al. 1 Chen et al. 2020b). Simplicity bias seems to be one of the primary reasons behind learning spurious correlations within NNs (Shah et al. 2020).
Implicit bias of gradient deseent: There is also large body of work understanding the implicit bias of gradient descent dynamics. Most of these works are for standard linear (Ji & Telgarsky 2019) or deep linear networks (Soudry et al. 2018: Gunasekar et al. 2018). For nonlinear neural networks,
3

---

one of the well- results is for the case of 1 1-hidden layer neural networks with homogeneous activation functions (Chizat & Bach 2020), which we crucially use in our proofs. More related works are provided in Appendix.
3 Preliminaries
In this section, we provide the notation and background on infinite width max-margin classifiers that required to interpret the results of this paper.
3.1 Basic notions
1-hidden layer neural networks and loss function. Consider instances I E R2 and labels y E {+1) jointly distributed as D. A 1-hidden layer neural network model (or fully connected network (FCN)) for predicting the label for a given instance x, is defined by parameters (w E Rmmxd b e Rm E Rm). For fixed activation function 6, given input instance v, the model is given as f((16, b, a), x) == (ā, Φ(ux + b)). where (6-) is applied elementwise. The cross entropy loss L. for a given model f. input x and label y is given as L (f(x), v) def log(1 + (app-y/((w;))))))) b, x))).
Margin. For data distribution D, the margin of model f(z) is given as min(rry) (z)).. (z).
Notation. Here some useful notation that we will use repeatedly. For matrix A, A(i,.) denotes the ith row of A. For any N. Sk-1 denotes the surface of the unit norn Euclidean sphere in dimension k.
3.2 Initializations
The gradient descent dynamics of the network depends strongly the scale of initialization. In this work, we primarily consider rich regime initialization.
Rich regime. In rich regime initialization, for any i e [mm] the parameters (i, of the first layer are sampled from uniform distribution on Sd. Each a(i) is sampled from Unifl -, }. and the output of the network scaled down by 1 (Chizat Bach 2020). This is roughly equivalent in scale to Xavier initialization Glorot & Bengio (2010), where the weight parameters in both the layers are initialized approximately as N (0, 2) when m d.
In addition, we also present some results for the lazy regime initialization described below.
Lazy regime. In the lazy regime, the weight parameters in the first layer initialized with A/O0, 1)) those of second layer are initialized with (O, ) and the biases are initialized to 0 (Bietti Mairal
2019 Lee et al. 2019). This is approximately equivalent in scale to Kaiming initialization (He et al. 2015
3.3
Infinite Width Case
For 1-hidden layer neural networks with ReLU activation in the infinite width limit i.e., as m -> oc, Jacot et al. (2018); Chizat et al. (2019); Chizat Bach (2020) gave interesting characterizations of the trained model. As mentioned above, the training process of these models falls into one of two regimes depending on the scale of initialization (Chizat al. 2019):
Rich regime. In the infinite width limit, the neural network parameters can thought of as a distribution over triples (w, b, a) E gd+1 where w Rd b, e R. Under the rich regime initialization, the function f computed by the model can be expressed as
f(v, x) = (ww. (a( x) + 6)]
Chizat & Bach (2020) showed that the training process with rich initialization be thought of as gradient flow on the Wasserstein-2 space and gave the following characterization of the trained model under the cross entropy loss E (x [C v, (x, y))]
Theorem 3.1 is an informal version of Chizat & Bach 2020 Theorem 5. For exact result, refer Theorem E.I
in Appendix

---

Theorem 3.1. (Informal) Chizat & Bach 2020 Under rich initialization in the infinite width limit with cross entropy loss. if gradient flow on 1-hidden layer NN with ReLU activation converges, converges to a maximum margin classifter v² given as
v
arg max min (u, vEP(S(+1) (a.
(2)
where P(Sd+1) denotes the space of distributions over gd+1
This training regime is known as the 'rich' regime since it learns data dependent features (wx,
Lazy regime. Jacot et al. (2018) showed that in the infinite width limit, the neural network behaves like a kernel machine. This kernel is popularly known as the Neural Tangent Kernel(NKK), and given by (xx, 2') aw 8W where W denotes the set of all trainable weight parameters. This initialization regime is called 'lazy' regime since the weights do not change much from initialization, and the NTK remains almost constant, i.c, the network does not learn data dependent features. We will use the following characterization of the NTK for 1-hidden layer neural networks.
(f()) 0f(²')
Theorem 3.2. Bietti & Mairal (2019) Under lazy regime initialization in the infinite width limit, the NTK for I-hidden layer neural networks with ReLU activation i.e., Φ(u) max(u, 0), is given
(x,²²) |<<||<<||
1 -2u((T cos (w)) + VI u²).
K ( x') = =|>0|||22||
where n(u) =
Lazy regime for binary classification. (Soudry al. (2018) showed that for linearly separable datasets, gradient descent for linear predictors logistic loss converges to the max-margin support vector machine (SVM) classifier. This implies that, any sufficiently wide neural network, when trained for finite time in the lazy regime on dataset that is separable by the finite-width induced NTK, will tend towards the L2 max-margin--lassinne given by
arg min Iifll s.t. uf(x) 2 14 (2.9) D. JEH
where represents the Reproducing Kernel Hilbert Space (RKHS) associated with the finite width kernel [Chizat] 2020). With increasing width, this kernel tends towards the infinite-width NTK (which universal (Ji et al. 2020)). Therefore, in lazy regime, we will focus on the L2 max- max-marrin--casssfie induced by the infinite-widt NTK.
4 Characterization of SB in 1-hidden layer neural networks
In this section, we first theoretically characterize the SB exhibited by gradient descent on linearly separable datasets in the independent features model (IFM). The main result, stated in Theorem 4.1 is that for binary classification of inputs in Rd, even there a single coordinate in which the data is linearly separable, gradient descent dynamics will learn a model that relies solely this coordinate, even when there are an arbitrarily large number but by non-linear classifier. In other words, the simplicity bias of these networks is characterized by low dimensional input dependence, which we denote by LD-SB, We then experimentally verify that NNs trained some real datasets do indeed satisfy LD-SB.
1 of coordinates in which the data is separable,
4.1 Dataset
We consider datasets in the independent features model (IFM), where the joint distribution over (xr, y) satisfies p(r, y) r(y) [ 9 ly). i.e the features are distributed independently conditioned on the label Here r(v) is a distribution over (-1,+1} and (z,ll)) denotes the conditional distribution of 1th cooooorrdiaaee x; given y. IFM is widely studied in literature, particularly the context of naive- Bayes classifiers Lewis (1998). We make the following assumptions which posit that there are at least two features of differing complexity for classification: one with a linear boundary and least one other with a non-linear boundary. See Figure for illustrative example.
5

---

Teskum
Figure 2: Illustration of an IFM dataset. Given class #1 represented by blue and red respectively, each coordinate value is drawn independently from the corresponding distribution. Shown above the supports of distributions on three different coordinates for an illustrative IFM dataset, for positive and negative labels.
One of the coordinates (say, the 1st coordinate WLOG) is separable by a linear decision boundary
with margin (see Figure 2 i.. By >> 0, such that Supp(4 (z) ly +1)) S bys oc) and
E Supp(gn (rily = 1)) C (-x0,- -7]. where Supp(-) denotes the support of a distribution.
None of the other coordinates linearly separable. More precisely, for all the other coordinates
i E [d] \ [1), E Supp(g: = -1) and {-1,+1} C Supp(y (2; +1)).
The dataset can be perfectly classified even without using the linear coordinate. This means, 3i * such that (xrlly has disjoint support for y +1 and y = -1
Though we assume axis aligned features, our results also hold for any rotation of the dataset. While results hold in the general IFM setting, in comparison, current results for SB e.g., Shah al.) (2020), are obtained for very specialized datasets within IFM and do not apply to IFM in general.
4.2
Main result
Our main result states that, for rich initialization (Section 3.2, NNs demonstrate LD-SB for any IFM dataset satisfying the above conditions. proof appears in Appendix A.I
Theorem 4.1. For any dataset in the IFM model with bounded density and bounded support, satisfy- ing the above conditions and x 2 1, and for 1-hidden layer networks with ReLU activation in the infinite width limit (i.e., Eqn. (7), there is unique max margin classifier v² (i.e., satisfying Eqn. (2)).
This v is given by: V = 0.5dg, + 0.50g, on Sd+1, where 8 represents the dirac-delta distribu-
V2(1+135* -1//2) and def [L. L. ... O denotes first standard basis vector. This implies f(v", (1)) + P1x = f(v" (1) (x(1), (()),, (x(2), (2) ~ D, where P represents the (rank-1) projection matrix on first coordinate.
tion, 01 = - (2((4 C1. /2(14155 1/V2), = (- (2(1477) CL.
Together with Theorem 3.1 this implies that if gradient flow converges, it converges to v² given above. Since P is rank-1 matrix and fiv", Pz( P x2) f(v" (x)),, 1/ satisfies the first two conditions of LD-SB (Definition [.I] with coordinates (2.... -.. d) has disjoint support for (<ally +1) and (x,ly - 1), P1 (x) can still perfectly classify the given dataset, thereby implying the third condition of LD-SB with 62 = 0.
and 6 = Moreover, since at least one of the
It is well known that the rich regime is more relevant for the practical performance of NNs since allows for feature learning, while lazy regime does not (Chizat al. 2019). Nevertheless, in the next section, we present theoretical evidence that LD-SB holds even in the lazy regime, by considering much more specialized dataset within IFM.
4.3
Lazy regime
In this regime, we will work with the following dataset within the IFM family:
For {+1} generate (Ir, y) E D as
+1 for y 1 y = -1
TV, Vi 2, d,
for
*Using linear probe for classifying pretrained representations is standard practice self-supervised learning
(Chen et al. 2020a Grill et al. 2020).
In Shah al. 2020 the theoretical results were obtained mainly for linear and non-linear coordinate

---

Table 1: Demonstration of LD-SB in the rich regime: This table presents P1 and P logit well as prediction changes on the five datasets. These results confirm that projection of input x onto the subspace spanned by P essentially determines the model's prediction I. f (resp. 4) indicates that LD-SB implies large (resp. small) value.
Dataset
rank (P) PI-IC (4)
(f)
PI-pC (4)
(t)
b- b-Imagenette
28.57 = 0.26
92.13
0.24
6.35 0.06
47.02++0.24
Imagenette
10
33.64 1.21
106.29 + 0.53
89.88 12.04+0.29
Waterbirds
25.24 1.03
102.35
6.78 0.15
35.96+0.02
MNIST-CIFAR
1
38.97
101.98 - 0.31
5.41 +.55
45.15+0044
Imagenet 150
15.78 + 0.05
132.05 0.06
13.05:10.03
99.76-10.01
Although the dataset above is point mass dataset, it still exhibits an important characteristic in common with the rich regime dataset only one of the coordinates is linearly separable while others not. For this dataset, we provide the characterization of max-margin NTK (as in Eqn. 3
Theorem 4.2. There exists do > 0 such that for every do, there exists an absolute constant N such that for all d > N and [7, 8vd), the L2 max-margin classifier for joint training of both the layers of I-hidden layer FCN with ReI.U activation in the NTK regime on the dataset D, ie., any f
satisfying Eqn. (3) satisfies:
(Fed(f(Px)) + PLx(())) pred(f(z())
A (x ( (x(2),y2)) E D
where P represents the projection matrix on the first coordinate and pred(f(z)) represents the predicted label by the model f on x.
The proof of this theorem is presented in Appendix A.2 The above theorem shows that the prediction on mixed example (z(1) + P1 22) is the same that (1)) (i.e., e1 0 in Definition 1.1 Furthermore, since there exists at least one coordinate & 1 which can be used to perfectly classify
the dataset, we have that Definition is satisfied with e2 0, thus establishing LD-SB.
4.4 Empirical verilication
In this section, we will present empirical results demonstrating LD-SB on 4 real datasets: Imagenette (FastAl 202T), binary version of Imagenette (b-Imagenette), wateebbrdss-lnddbird (Sagawa et 2020a) and Imagenet (Deng et al. (2009) as well as one designed dataset MNIST-CIFAR (Shah et al. 2020 More details about the datasets can be found in Appendix B.I
4.4.1
Experimental setup
We take Imagenet pretrained Resner-50 models, with 2048 features, for feature extraction and train a 1-hidden layer fully connected network, with ReL.U nonlinearity. During finetuning, we freeze the backbone Resnet-50 model and train only the 1-hidden layer head (details Appendix B.T))
Given model f(.), we establish its low dimensional SB by identifying small dimensional subspace, identified by projection matrix P, such that if we mix inputs æ and
Demonstrating LD-SB:
T2 as Px + Pizz, the model's output on the mixed input 5: def Pxr, + P1 22. f(ī) is always close
to the model's output on #1 i.e., fx1). We measure closeness in four metrics: (1) P1 logit change
(P PLC): relative change of logits wrt 2 i.e., [z f(x)) || [f (2)P logit change (P-LC):
relative change wrt logits of 222 i.c., [[ f(x2) (x>)) (3) P1 --reedicitionn change (P (PI:PC):
Pfpred(f(3)) * pred ( and (4) P-prediction change (P-pC) (P-pC):Pppred( Pfpred (z) * pred(f(xx)))
The quantities rank (P) and PI-PC correspond to and 6 in Definition I.1 respectively. To demonstrate that the subspace P1 has features that are useful for prediction,
also train model Spojj as follows. Given the initial model f and the corresponding projection matrix P, we then train another model fpoo by projecting the imput through P1 i.e., instead of using dataset (x(1) ()) for training, we use (P) ()) for training the second model (denoted by fpmj). We refer to this training procedure as OrthoP for orthogonal projection. The quantity |Ace(f) Ace(Smoj)
corresponds to Eg in Definition [.] We now describe how identify P in rich and lazy regimes.
7

---

Table 2: Demonstration of LD-SB in lazy regime: This table presents P1 and P logit well as prediction changes on the five datasets. These results confirm that the projection of input & onto the subspace spanned by P essentially determines the model's prediction on
Dataset
rank (P) (4-
P-LCC (f)
Pi-pC (4)
P-pC (t)
b Imagenette
36.9411.011
138.41#1.62
5.5 = 1.13
17.711555
Imagenette
15
55.9943.88
133.86+5.42
11.25+0036
89.75=0.15
Waterbirds
6
36.8915.18
20.7440.64 105.4117.06
45.96:0069
MNIST-CIFAR
2
24.9 0.61
141.12+1.86
0.53 0.24
49.83+0.78
Imagenet
200
32.7440.02
132.47=0.04
18.2
99.7440.01
100
100
MESTOCR
Iragretee E-PMeesdi
80
80
Elfective rank
Effective rank
60
60
4.0
40
MA65-09 Weartires
strageett
0
1.00
200
300
400
500
100
200
300
400
500
Epochs
Epochs
(a) Rich regime
(b) Lazy regime
Figure 3: Evolution of effective rank of first layer weight matrices in rich and lazy regimes.
4.4.2 Rich regime
Theorem 4.1 suggests that asymptotically, the first layer weight matrix will be low rank. However, since we train only for a finite amount of time, the weight matrix will only be approximately low rank. To quantify this, we use the notion of effective rank Roy Vetterli (207))
Definition 4.3. Given matrix M, its effective rank is defined e Σ, (M (M) where (M) denotes the th singular value of M and or (M)2 del a.(M)²
E.
One way to interpret the effective rank is that is the exponential of von-Neumann entropy Petz (001 of the matrix .MM" TTMMTTy. where Tr (-) denotes the trace of matrix. For illustration, the effective rank of projection matrix onto i dimensions equals k.
Figure 3a shows the evolution of the effective rank through training on the four datasets. We observe that the effective rank of the weight matrix decreases drastically towards the end of training. In this case, we set P to be the subspace spanned by the top singular directions of the first layer weight matrix. Table presents the results for P1 and P-LC as well pC, while Table 3 presents Acc(f) and Acc These results establish LD-SB in the rich regime.
Table 3: Accuracy of fpoj in rich regime
Table Accuracy of Spoj in lazy regime
Dataset
Acc(f)
Acc( ppojj
Dataset
Acc(f) Fpo))))
b b-mangenette
93.35
91.35 0.32
I-mageenetee
93.09
91.77 + 0.34
Imagenette
79.67
71.93 0.12
Imagenette
80.31
77.34 0.21
Waterbirds
90.29
89.92
Waterbirds
90.4
0.15
MNIST-CIFAR
99.69
98.95 0.02
MNIST-CIFAN
99.74
98.54 - 0.00
Imagenet
72.02
69.63 0.08
Imagenet
72.6
72.07 = 0.08

---

Table 5: Mistake diversity and class conditioned logit correlation of models trained independently (Mist-Div (f. fma) and CC-LogitCorr (f. fima) resp.) vs trained sequentially after projecting out important features of the first model (Mist-Div (f, fpwg) and CC-LogitCori (f, fpoj) resp.). The results demonstrate that f and Spmj are more diverse compared to ʃ and find-
Mist-Div (f. fina) (f)
Mist-Div
CC-LogitCor
CC-LogitCorr
Dataset
fprj) (1)
fin) (4)
(pojj (L)
B-Imagenette
3.87 1.54
21.15 + 1.57
99.88 + 0.01
90.86 + 1.08
Imagenette
6.6 + 0.46
11.44 + 0.65
99.31 + 0.12
91 + 0.59
Waterbirds
2.9 + 0.52
14.53 + 0.48
99.66 + 0.04
93.81 + 0.48
MNIST-CIFAR
0.0
5.56 + 7.89
99.76 0.17
78.74 = 2.28
Imagenet
6.97 #.06
12.31 + 0.16
99.5 = 0.0
92.52 = 0.01
4.4.3 Lazy regime
For the lazy regime, it turns out that the rank of first layer weight matrix remains high throughout training, as shown in Figure 3b However, we are able to find a low dimensional projection matrix P satisfying the conditions of LD-SB (as stated in 1.I as the solution to an optimization problem. More concretely, given pretrained model f and a rank r, we obtain a projection matrix P solving:
min- Σ (I(PP)) + dC (S(P) a (.)))) E=1
1
where [[[]] represents a uniform distribution over all the labels, (x(1),, y ..... (xx()), (()) are training examples and L (:,-) is the cross entropy loss. We reiterate that the optimization is only over P, while the model parameters f are unchanged.
words, the above function ensures that the neural network produces correct predictions along P and uninformative predictions along PI. Table presents the results for P1 and P-LC as well as pC, while Table presents Acc(f) and Acc(Sp These results again establish LD-SB in the lazy regime.
Training diverse classifiers using OrthoP
Our results above motivate a natural strategy to construct diverse ensembles i.e., use f and fyoj instead of two independently trained models. this section, we provide two natural diversity metrics and empirically demonstrate that OrthoP leads to diverse models in practice. We also demonstrate that an ensemble of f and fproj has higher robustness to Gaussian noise compared to an ensemble of independently trained models.
Diversity Metries: Given any two models f and f. we empirically evaluate their diversity using two He ()) (z
metrics. The first mistake diversity: Mist-Div (f. def (f (z
where abuse notation by using f(z;) (resp. f(x:)) to denote the class predicted by f (resp f) on Higher Mist-Div (f. means that there is very little overlap in the mistakes of f and f. The second is class conditioned logit correlation i.e., correlation between outputs of and f. wv)
conditioned on the class. More concretely, CC-LogitCorr (f.j) Evey Com (I((()))))
[3]
where corr([((x)) (f y, y) represents the empirical correlation between the logits of f and on the data points where the true label is y. Table compares the diversity of two independently trained models (f and fina) with that of two sequentially trained models (f and fpro)). The results demonstrate that f and froj more diverse compared to f and find- We have also compared OrthoP to another diverse training method Evading-St SB (Teney et al.) (022). The results are provided in Table 7 in the Appendix. As can be seen, results either better or comparable to Evading-SB.
Figure 4 shows the variation of test accuracy with the strength of gaussian noise added to the pretrained representations of the dataset. Here, ensemble is obtained by weighted averaging of the logits of multiple models, trained either independently (Sina or using OrthoP (Frro)). ). Morcover, also compare our method to another diversity training method (foa) termed as
Ensembling:

---

0.60
0.0
0.12
der
sid
(a) Imagenette
(b) Waterbirds-Landbiice
(c) Imagenet
Figure 4: Variation of accuracy vs standard deviation of Gaussian noise added the pretrained representations of the dataset. Here modell, ensemble-ind, mnsemble-proj and ensemble-esh refer to the original model f. ensemble of and independently trained model find. ensemble of f and fpoj trained using OrthoP, and ensemble of f and fesb obtained using Evading-SB method respectively.
Evading-Si SB Teney et al. 2022). ). We can see that, an ensemble of f and fpen is much more robust as compared to ensemble of and find, and generally comparable to and ensemble of f and fesb-
6 Conclusion: Summary, Limitations and Future Directions
In this work, we propose a rigorous definition of simplicity bias, which is believed to be a key reason for their brittleness (Shah et al. 2020). In particular, we prove that 1-hidden layer networks suffer from low dimensional input dependence (LD-SB), and empirically verify this phenomenon on several real world datasets. We also propose a novel approach OrthoP- train diverse models, and demonstrate that an ensemble consisting of such diverse models is more robust to Gaussian noise. Extending these inisights deeper models in the finite width setting are interesting directions for future work.
7 Acknowledgements
acknowledge support from Simons Investigator Fellowship, NSF grant DARPA grant INF2010022,, and DOE grant DE-SC0022199.
10

---

References
Bear image from http:///wwrr.....aaaaa///////aaa Grizzl/Bearreembaffffttt.... Acceeee
Accessed: 2022-09-26.
Swan image from wikipedia. httpp:///a...iiiippaaaaaoooo
(Crp (Cygnus.
olor) _nd_add_d__eeee._______________ jpg Accessed: 2022-09-26.
Arora, S., Cohen, N., Hu, W., and Luo, Y. Implicit regularization in deep matrix factor- ization. In Wallach, H., Larochelle. H., Beygelzimer. A., d'Alché-Buc; F., Fox, E., and Garnett, (eds.), Advances in Neural Information Processing Systems, volume 32. Cur- ran Associates, Inc., 2019. URL https: : pttt:::://poceeeingss. nwwsisseerip... ctlccc//ppppp cttpppaprre//19//teeeee. c00783365f hat.:///8000111114466e682222228 Paper pdf
Bietti, A. and Mairal, J. On the inductive bias neural tangent kernels. In Advances in Neural Information Processing Systems, volume 32, 2019. URL http:: :/ttps::/proceeeiinng neurinn pdf
Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of
visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020a.
Chen, T., Luo, C., and Li, L. Intriguing properties of contrastive losses. Advances in Neural Information Processing Systems, 34:1 11111133418444 2021.
Chen, Y.. Wei, C., Kumar, A., and Ma, I. Self-training avoids using spurious features under domain shift. Advances in Neural Information Processing Systems, 33:21061-21071, 2020b.
Chizat, L. ii:
Gradient descent for wide Generalization and implicit bias. htt.::/pee...oaeaaaaaotaotaaaaad
two-layer neural networks httpe::/ffraccisaacaa com/
2020.
Chizat, I.. and Bach, F. On the global convergence of gradient descent for over-parameterizee models using optimal transport. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa- Advanees Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 31.
Curran Associates, Inc., 2018. URL https: // http///proeeeee...e//e//// neurips cc/paperr/0118//1llll a1af 08866688440000079900000000
pdf
Chizat, L. and Bach, F. R. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on Learning Theory, 2020, volume 125, pp. 1305-1338. 2020.
Chizat, L.,, Oyallon, E.,, and Bach, F. On lazy training in differentiable programming. In Advances in Neural Information Processing Systems, volume 32, 2019. URI httttt:://rpceeeddnnni pdf
neurips
Cook. M., Zare, A., and Gader, P. Outlier detection through null space analysis of neural networks, 2020.
01263
Deng, J., Dong, W., Socher, R., Li, L.-J., Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. leee, 2009.
Ding, Z., Chen, S., Li, Q., and Wright, S. Overparamettrization of deep resnet: Zero loss and mean-field analysis. Journal of Machine Learning Research, 23(48):1 2022. URL attp:///mil org/ httg//apers/v23///2
0669 html
Fang. C., Lee, J., Yang, P., and Zhang, T. Modeling from features: mieldffield framework for over-parameterized deep neural networks. In Belkin, M. and Kpotufe, S. (eds.), Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pp. 1887-1936. PMLR, 15-19 Aug 2021. URL v134/fang211 hhmll html
www.iinnnmmmm.....mmmmmmm mlr mlr.press/
FastAI. Imagenette dataset. httts:://githni http://ttw...mtttttaaaaaaannnnee
2021. 2021.

---

Frei, S., Vardi, G., Bartlett, P.. Srebro. N.. and Hu, W. Implicit bias in leaky reLU networks trained on high-dimensional data. In The Eleventh International Conference on Learning Representations. 2023. WRR.https://// /totpp:///penrevieew nww/forumm htt.:ooommmmdddpppppy......
Galanti, and Poggio, T. Sgd noise and implicit low-rank bias in deep neural networks, 2022. URL org/ abs/ 0tlbb//226.00999
Geirhos, R., Rubisch, P., Michaelis, C., Bethge, M., Wichmann, F. A., and Brendel, W. Imagenet- trained cnns are biased towards texture; increasing shape bias improves accuracy and robustness. In International Conference on Learning Representations, 2018.
Geirhos, R., Jacobsen, J.-H., Michaelis, C., Zemel, R., Brendel, W., Bethge, M., and Wichmann, F. A. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(1 226665-673 2020.
Glorot, X. and Bengio. Y. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statisties, pp. 249-2566 2010.
Grill, J.-B.. Strub, F., Altché, F., Tallec, C., Richemond, P., Buchatskaya, E., Doersch, C., Avila Pires, B., Guo, 7.,, Gheshlaghi Azar, M., et al. Bootstrap your own latent-a new approach to self- supervised learning. Advances in neural information processing systems, 33:21271-21284, 2020.
Gunasekar, S., Woodworth, B. E.. Bhojanapalli, S.. Neyshabur, B., and Srebro, N. Implicit regular- ization matrix factorization. Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems. volume 30. Curran Associates, Inc., 2017. URL http:tttp//ppoeeediiii 2017/file/581912220 446dae66376666666434 hat.ae/e66..111cddddd1d11101-11. Paper. .pff
wwwdddiinnnnuuiipiioe neurips.cc/paper.
Gunasekar, S., Lee, J. D., Soudry, D., and Srebro, N. Implicit bias of gradient descent on linear convolutional networks. In Advances in Neural Information Processing Sys- tems, volume 31, 2018. URL https: //pp://///oocce//p/ neurips httpe/ccccpppe////////ee (osSeebbfaacc5122eeeeeeaa222
pdf
Hacohen, G. and Weinshall, D. Principal components bias in over-parameterrzed linear models, and its manifestation in deep neural networks. Journal of Machine Learning Research, 23(155):1-46, 2022. URL http:://:/
html
Hacohen, G., Choshen, L., and Weinshall, D. Let's agree to agree: Neural networks share clas- sification order on real datasets. In III, H. D. and Singh, A. (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 3950-3960. PMLR, 13-18 Jul 2020. URL /ttppptts////roteeeee v119 2119/hacohen20a.nnto html
mlr, .uress/
He, K., Zhang, X., Ren, S., and Sun, Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE international conference computer vision, pp. 1026-10344 2015.
Hendrycks, D., Basart, S., Mu, N., Kadavath, S., Wang, F., Dorundo, E., Desai, R., Zhu, T., Parajuli, S., Guo, M., et al. The many faces robustness: A critical analysis of ofutt--diitributtin generalization. In International Conference Computer Vision, pp. 8340-8349, 2021.
Huh, M., Mobahi, H., Zhang, R., Cheung, B., Agrawal, P., and Isola, The low-rank simplicity bias in deep networks. arXiv preprint arXiv:2103. 110.22003.. 2021.
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems, pp. 8580-8589, 2018. 2018
Ji, and Telgarsky, M. The implicit bias of gradient descent on nonseparable data. In Conference on Learning Theory, pp. 1772-1798, 2019.
Ji, 7.,, Telgarsky, M., and Xian, Neural tangent kernels, transportation mappings, and universal approximation. In International Conference on Learning Representations, 2020.
12

---

Kalimeris, D., Kaplun, G., Nakkiran, P., Edelman, B., Yang, T., Barak, B., and Zhang, H. Sgd on neural networks learns functions of increasing complexity. In Wallach, H.. Larochelle, H.. Beygelzimer, A., d' d'Alché-Bue, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL, Utt.htttss//proceddiie Paper. pdff
neurips. http:ps.cccppapp///2311 htt.://ee...2224.545577766066000
Lee, J., Xiao, I.,, Schoenholz, S., Bahri, Y., Novak, R., Sohl-Diccsteenn, J., and Pennington, J. Wide neural networks of any depth evolve as linear models under gradient descent. In Advances in Neural Information Processing Systems, volume 32, 2019. URL htthsttps:/ /tppp/ee/// neuriiinsnnn.uuss........ cc/ Paper .
Lee, Y., Yao, H., and Finn, C. Diversify and disambiguate: Learning from underspecified data. arXiv preprint arXiv:2202.03418, 2022.
Lewis, D. D. Naive (bayes) at forty: The independence assumption in information retrieval. In European conference on machine learning, pp. 4-15, 1998.
Li, Z., Luo, Y., and Lyu, K. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. In International Conference on Learning Representations, 2021. URL
ntt/forumm id= AHOs>Smmmi
Liu, E. Z., Haghgoo, B., Chen, A. S., Raghunathan, A., Koh, P. W., Sagawa, S., Liang, P., and Finn, C. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, pp. 6781-6792 2021.
Lyu, K., Li, Z., Wang, R., and Arora, S. Gradient descent two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 34: 14:18-98-1299 2021.
Melamed, O., Yehudai, G., and Vardi, G. Adversarial examples exist in two-layer relu networks for low dimensional data manifolds, 2023.
Nasery, A., Addepalli, S., Netrapalli, P., and Jain, Daft: Distilling adversarially fine-tuned models for better ood generalization. arXiv preprint 1xXx:22088.09133 2022.
Ndiour, I., Ahuja, N., and Tickoo, O. Out-of-distribution detection with subspace techniques and probabilistic modeling of features, 2020. htttttttt///xiiii org/abs abaaaabs//122 04250
Petz, D. Entropy, von neumann and the von neumann entropy. In John von Neumunn and the foundations of quantum physics, pp. 83-96. 2001.
Pezeshki, M., Kaba, O., Bengio, Y., Courville, A. C., Precup, D., and Lajoie, G. Gradient starvation: A learning proclivity in neural networks. Advances in Neural Information Processing Systems, 34: 1256-1272, 2021.
Phuong, M. and Lampert, C. H. The inductive bias of re(lu) networks on orthogonally separable data. In International Conference on Learning Representations, 2021. URL https://openrrrrirr
Rahaman, N., Baratin, A., Arpit, D., Draxler, F., Lin, M., Hamprecht, F. A., Bengio. Y., and Courville, A. On the spectral bias of neural networks. 2018. doi: 10.48500AARXIV.1866.07777 URL /ttts://arxiv.cor org/ abs/ 1816.00080088 08734
Razin, N. and Cohen, N. Implicit regularization in deep learning may not be explainable by norms. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33. pp. 21174-21187, Cur- Associates, Inc., 2020. URL
nessroeeeddnnggsnnrriips neurips cc/ http//pppr/220000 paapr//000//ille/ pdf
F21 222255599ad222accceeeeeeee6
Ronen, B., Jacobs, D., Kasten, Y., and Kritchman, S. The convergence rate of neural networks for learned functions of different frequencies. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc. F. Fox, E.,, and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. JRL htttlpp:///prcceeeee http///ooeeeenn..eeee// neur caaurip..oo ec/
pdf
13

---

Rosenfeld, E., Ravikumar, P., and Risteski, A. Nomain-adjussen regression or. Emm may already learn features sufficient for outtof-ddstributionn generalization. arXiv preprint arXiv:2202.0683 066566 2022.
Ross, A. and Doshi Velez, F. Improving the adversarial robustness and interpretability of deep neural networks by regularizing their input gradients. In AAAI Conference on Artificial Intelligence, volume 32, 2018.
Roy, O. and Vetterli, M. The effective rank: A measure of effective dimensionality. In European signal processing conference, pp. 606-610, 2007.
Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P. Distributionally robust neural networks. In International Conference on Learning Representations, 2020a.
Sagawa, S., Raghunathan, A., Koh, P. W. and Liang, P. An investigation of why overparameterization exacerbates spurious correlations. In International Conference on Machine Learning, pp. 8346- 8356, 2020b.
Santambrogio, F. Euclidean, Metric, and Wasserstein gradient flows: an overview, 2016. URL 03890
https://arxiv.orgg
Scimeca, L., Oh, S. J., Chun, S., Poli, M., and Yun, S. Which shorteut cues will dnns choose? a study from the parameter-space perspective. In International Conference on Learning Representations, 2021.
Shah, H., Tamuly, K., Raghunathan, A., Jain, P., and Netrapalli, P. The pitfalls of simplicity bias in neural networks. Advances in Neural Information Processing Systems, 33:9573-9585 2020.
Soudry, D., Hoffer, E., Nacson, M. S.,, Gunasekar, S., and Srebro, N. The implicit bias of gradient descent on separable data. J. Mach. Learn. Res., 19(1):2822-28)) 2878, jan 2018.
Szegedy, C., Zaremba, W., Sutskever, I., Bruna, J., Erhan, D., Goodfellow, I., and Fergus, R. Intriguing properties of neural networks. In International Conference on Learning Representations, 2014.
Teney, D., Abbasnejad, E., Lucey, S., and van den Hengel, A. Evading the simplicity bias: Training a diverse set of models discovers solutions with superior ood generalization. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16761-16772, 2022.
Timor, N., Vardi, G., and Shamir, O. Implicit regularization towards rank minimization in relu networks. In Agrawal, S. and Orabona, F. (eds.), Proceedings of The 34th International Conference on Algorithmic Learning Theory, volume 201 of Proceedings of Machine Learning Research, pp. 1429-1459. PMLR, 20 Feb-23 Feb 2023. URL httt:::t//ooteeeeeeeeeeeee .ml...s......... mlr molr.press/v201)
timor23a. html
Vardi, G., Yehudai, G., and Shamir, O. Gradient methods provably converge to non-robust networks. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL /tRL:tttts//opeereviies wtt::/forommmmddddxXiiagggg
Vyas, N., Atamasov, A., Bordelon, B., Morwani, D., Sainathan, S., and Pehlevan, C. Feature-leaming networks are consistent across widths at realistic scales, 2023.
Wang, H., Li, Z., Feng, L., and Zhang, W. Vim: -utttf-disttibuttin with virtual-logit matching. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18 24. 2022, pp. 4911-4920. IEEE, 2022. doi: 10. 0IWCVVPR52688.2022000422 URL https: /doi. org/ 10. 1111 CVwO9//VP22288.... 2022. 000877
Xu, Y.. He, H., Shen, T., and Jaakkola, T. Controlling directions orthogonal to a classifier. arXiv preprint arXiv:2201.1 /259, 2022.
Yang, G. and Hu, E. J. Tensor programs iv: Feature learning in infinitewwidth neural networks. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 11727-11737. PMLR, 18-24 Jul 2021.
mlr. wtt.ss//113/yaagg2..... htm]
Zaeemzadeh, A., Bisagno. N., Sambugaro, Z., Conci, N., Rahnavard, N., and Shah, M. Out-of- distribution detection using union of I-dimensional subspaces. In Proceedings of the IEFF/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 9452-9461, June 2021.
14

---

A Proofs for rich and lazy regime
A.1 Rich regime
We restate Theorem 4.1 below and prove it.
Theorem A.I. For any dataset in IFM model satisfying the conditions in Section 4.1 2 1 and 2V) as in Eqn. (1). the distribution v 0.50g, + -.500z on Su+1 is the unique max-margin classifier satisfying Eqn.
2 where 01 (1=1211170 22(1+T²) C1 (2(++²)) 1/v2).0; 02
(2(11+-)) 01 /211+17) -1/v22) and e1 def [1,0......0) denotes first standard basis vector. In
particular, this implies that if gradient flow for 1-hidden layer FCN with ReLU activation under rich initialization in the infinite width limit with cross entropy loss converges, and satisfies the technical conditions in Theorem E.1 then it converges to v² satisfying f(v", + P1r2)) = fv" 21)V(E1, y1), (22,, V2) E D, where P represents the (rank-1) projection matrix the first coordinate.
Proof of Theorem A.7 The proof relies showing that '" is max-margin classifier in Theorem 3.1 To this end, we employ a primal-dual characterization max- margin classifiers and construct dual certificate that proves the optimality of margin of v², Chizat & Bach (2020) showed the following puimal-dual characterization of maximum margin classifiers in eqn. ②)
Lemma A.2. (Chizat & Bach 2020 v² satisfies eqn. 2 if there exists data distribution p such that the following two complementary slackness conditions hold:
Supp(v") S
arg max vlal
b))] and
(4)
(w (0))44))1
Supp(p") C arg min (B [a(o( + b))]
(5)
The plan is to construct distribution p" that satisfies the conditions of the above Lemma. Uniqueness. Note further that for a fixed p", R(z, yf x) is an upper bound for the margin min( ff z) of any classifier v. Hence, for uniqueness, it suffices show that d8, de, are the unique maximizers of the objective on the RHS of eqn. 4 and that the unique maximum margin convex combination of da, 50, over D is v*.
We first describe the support D of p". For y E {+1} we generate (x. y) E D as
X1 yy
ViE .., d, .,,,,xx,x, {
=1 for =
0 for ! = -1
Now for (I,z)= E D, define
{
0.5 for y= 0.5 for -1
p (z. 3)
(6)
Note that p" is supported on 2d-1 positive instances and one negative instance. We begin by showing eqn. ( Claim A.3. p" as in eqn. satisfies eqn. combination of dø, 6g, is v
Further, the unique maximum margin convex
Proof. Let find the minimizers (x,y) D of yf(u,z) " (a(o((w, + +])] for any
Aoo, (1 Ajooz. 0 < X 1. ff x) for (x. v) with y -1 (denoting 2 by -a, where a 27)
2f(,)))= = -1 -1[A* (
1
eT eT(aaje1) + /2(1 + /2(1+
/2 +
V2
1
+ (1 2) Φ
2211 ej (-aiee))
22
15

---

and for (x, v) with y = (denoting 21 by a2, where a2 2 7) is
vf(u,x)= = 1 =1x0
e7 (a(01) +
1
v2((1
v2(1+>2)
V/2
+ (1 *
7 ej (a2e1) 2( + 72)
+ 72(
As 7 2 the expressions above equal +1 and (1 A) vo
respectively, and hence are
minimized at @1 02 7.- Hence, the margin of is min(), 1 A) VI which is uniquely maximized at X = 1/2. Further for À = 1/2, all points in have the same value vf
2
In the rest of the proof we show eqn. (4) Let denote by g(w, b, == Be sla(a((w, +b)) We show that 50, 50, are the only maximizers of g(w, b, a) over gd+1,
We first find g(01),9(02).
((01) = Pr(y = 11
1
7
ef (7()))) +
1
V2
122(1 +2)
12(11 72))
( (V2(1+ (-((-
1
+1+
+ =
V2 o
22(1+
4
where the first term is because W2, W3..--:: Wi zero for lj. Similarly, g(02)
We now
show that g(w, b) < v/2+1 for b) & {01, 02).
We begin by showing the following simple but useful claim.
Claim A.4. All maximizers of g(w, b, a) over gd+1 satisfy |a| 1/v2.
Proof. The proof essentially follows from the 1- homogennttyy of the ReLU function 6 and sepa-
rability of g(w, b, a). Note that g(w, b, a) V| 2 + ba g(w', b,,, 1) where ||w'| + b2 1. Maximizing g(w, a) is equivalent to maximizing g(w', 1) over Sd and / | + 1 over gd+1 respectively. The second of these has its unique maximum at |a| = 1/v2, completing the proof. D
Now express g(w, b, a)
(w, b, a) = a (Pr(y- = 1)E((((" (i0 x b) = 1] -Pr(y- -1)E(((1 +bll =-1]) -1])
Eg [d(rw) + +b+ Σ o,w;)] o(b -var)
2
i=2
where a; independent Rademacher random variables. We have two cases
Case 1: 1/v2. By eqn.
7 we have
g(w,b.1/V2) S E. [o(Tw] + + Σ oww))]
2/2
1=2
To simplify the above, define the random variable X 0,W;; and denote ywu + b by a. Note
that |a|| hyw: + bl S /
211 which follows from ||m|| + b² = 1/2. The expectation in the last
expression above becomes
E(p(X + a)] = E((X + aX - a
-a)
= =E 2 a]] aal - Pr(X 2 a)) x E 2 atl ta,
16

---

where the last equality follows from symmetry of X. Note that Var X) = Ei wi
which is at most
(using T + a and ||w|? b² 1/2). Using A.5 to upper bound [|X1 a))
147² we have
Blo(X + +a]] a+
1
min
2
147
(1
a²
2
2'
20
1 1 7
We check that the RHS of the above has its unique maximizer at V - for |a| S
1t,²
Hence g(w, b, a) S /+7² in this case. We are now done since any (wi, b) satisfying ywi
1+>2 and w 6² 1/2 has b
v2(1+>")
Case 2: = -1/22. Using eqn. we have g(w, b, -1/V2) S o(b - 2w1)/2v2 which for 6² + wi < 1/2 attains its unique maximum
b (2(1+)))
Finally, note that the weights of the trained network (w, b, a) are sampled from v², Hence, the final claim in the theorem about f(v", Px1 + P1zz) follows since the distribution of only has support on C₁ and Cc.
Note: The proof above also holds for any rotation of the dataset, with the weights corresponding to the first standard basis vector, i.e, the direction linearly separable coordinate in that basis.
A.1.1 Auxiliary lemmas for rich regime
Lemma A.5. For any symmetric discrete random variable X with bounded variance, for 0.
1
1
Var(X)
E[XI(X 2 a)] S V
min
Var(X).
2
2a²
Proof.
E[ 2 a)]
Σ rp(z) Σ zza
Voll))Vploia p(x 2 a) Σ x²p(x),
(8)
z2a
x2a
where the last inequality by Schwartz. Also by Chebyshev's inequality, x( a)
<
V (X)X/222 Combining this with eqn. (8) and using symmetry of X and non-negativity of a gives the required lemma.
A.1.2 OrthoP method on IFM
Here, we theoretically establish that ff and fproj obtained via OrthoP rely on different features for any dataset within IFM. Consequently, by the definition of IFM, and fpoj have independent logits conditioned on the class.
Proposition A.6. Consider any IFM dataset as described in Section 4.1 Let f be the model described Theorem 3.1 and Fpmoj be the second model obtained by OrthoP. Then, the outputs and fpmoj on x i.e., ) and fpmm (r) depend only on 2 and {z2,-.. ... .ra) respectively. Let the model obtained in Theorem 3.1 be denoted by f. Consider the projection matrix P along the top singular vector of the first layer weight matrix of f. Then, the dataset obtained by projecting the input through P. is not separable along the linear coordinate.
Proof. As shown in Theorem
3.1 the final distribution of the weights is given by D² = 0.560, + +0.50
where 0g = (2(1+)²) G1 /2 1/V2),0,
( C1, /2(14555' -1/V2) and e def
[1, 0, ... 0j denotes first standard basis vector.
17

---

As the first layer weight matrix only has support along the e₁ direction, therefore its top singular vector also points along the e direction. Hence, = ej and Pi the identity matrix. Thus, the dataset obtained by projecting the input through Pa has value 0 for the linear coordinate, for both y = +1 and y = -1. Hence, not separable along the linear coordinate. Thus, the second model fproj relies on other coordinates for classification.
e where I denotes
L
A.2
Lazy regime
Theorem
4.2 is corollary of the following more general theorem.
Theorem A.7. Consider point E D. For sufficiently small there exist absolute constant N such that for all d > N.y < cvd and "y 2 7, for the joint training of both the layers of 1-hidden layer FCN with ReLU activation in the NTK regime, the prediction of any point of the form (C, (5,22::)
satisfies the following:
1. For Ç 2 0.73, the prediction is positive.
2. For Ç < 00.95% the prediction negative.
The above theorem establishes that perturbing x1 by O(7) changes pred(f(z) for D (whereas classifier exists that achieves margin of (2(v)) D, as D has margin 1 for coordinates {2....l d)). As = od), this shows that the learned model is adversarially vulnerable.
Proof of Theorem A.7 The idea of the proof is to obtain an analytical expression for f(z) using KKT conditions for the max-margin SVM for the NTK kernel (as Theorem 3.2).
We begin with some preliminaries. We will refer to the first coordinate of the instance as the 'linear' coordinate, and the rest as 'non-linear' coordinates. Also, henceforth append extra coordinate with value to all our instances (corresponding to bias term) as is standard for working with unbiased SVM without loss of generality.
Explicit expression for f. Using representer theorem for max margin kernel SVM, we know that f can be expressed
f(x)
Σ
Xay K (x)
(2(9)
for some A; 2 0 (that are known as Lagrange multipliers). Further by KKT conditions, a function possessing such a representation (that correctly classifies D) has maximum margin if fxx(0) = whenever A (training points satisfying 0 are called support vectors).
We begin with useful claim.
Claim A.8. The max margin kernel SVM for D with the NTK kernel has all points in D as support vectors.
Proof. By the above discussion, it suffices to show that the (unique) solution a E IR PI Ka satisfies sagn(a:) y(9) for all i, where K is the [D] |D] Gram matrix with (i.j)th entry (f) ,00i and y: (0) (the Lagrange multipliers X, are then given by y;a,).
Structure of Gram matrix. Order D so that the positive instances appear first. Then the Gram
B C) where B 122 x2 and R the CT
matrix K has block structure of the form
R
Gram matrices for the positive and negative instances respectively, and C R 2d x1 represents the K(x(0,, x D)) values for i < (D].
Recall that for the NTK kernel, K (x) x.)) has the form ||x (|<(((()))))))) (())(()))) Note all the positive instances have the same norm (denoted by A1 V /d+ and the inner product between two positive instances depends only on the number of non- mot-hin non-linear coordinates (denoted by 6 for 0 < i < d - 1). Hence, the rows of B are permutations of each other, with the entry p appearing ( (d-1) times. Similarly, the entries in C' are all equal and denoted by P1p2 where Ba
18

---

denotes k(x(1), xx (I) for any t < [D] and = ||.| - 1
The only entry in R is p2r(1).
In particular,
( -
d
1 2
8: = R
for i [D] 1].
and
= = R
d 72
Vd+
Now are ready to solve Ka y. By symmetry in the structure of K, a looks like (a, a ............. b). where the first D - entries are the same.
Expanding Ka y, we get two equations given by
d-1 1
api
Σ
1
Bi )²
+ bpip2Ba=1 =1 and 2d ap1p28a + p2x(1)b = -1.
i
0=
Solving, we get
(2t((1) + PLi3
2d aod-tapip2y
app 111223d
a
0ip2 Σ ==
[x(L)B) [3]
and b=
(2)((1)
We now show that > and < O. Note that for sufficiently large 3d can be made arbitrarily close to x()) = 1/4 (since K is smooth around 0). Hence, a >> 0 implies b < O. We in fact give the following estimate for a:
a = 21-d (2k(1) + P1Ba
2
+ +0 ( $
where
(9)
EpiP2
For the lower bound on g. write
[d/2] Σ
d 1)
dd
(B, + Ba-1)) -2t
i
==0
[d/2]
d - 1) i (7 2Ba/2 - 22-1 2/
2d-1
i=0
where for the first inequality we used convexity and for the second inequality we used Ba/2 <(0) + (1//),(1)))))))))))))))))))) 3a = x(0) + O(1/Vd). For the upper bound on &, write
2i
|kt (x(1);, - 331 S <(1)
i
i
k
d + 72
1=0
2i
x(1)(d 1)2d-1
2 d + 72
x(1)2
d 722
where for the second inequality we used k(u) < 1 + u (which holds by convexity and x(-1) x1] = 2).
Now we analyze predicted labels for points of the form (G, (2d41)) where o E D. We make two cases depending on the label of x.
Predicted label for point (0. where 2(()) D has positive label
(t)
Our point (denoted by æ) has the form (C. C2,.., Ca, 1) where & +1. The idea of the proof is to write f explicitly as a function of ( and work with its first order Taylor expansion around Ç with some additional work to take care of non-smoothness of f.
Explicit form for f. Let Ti def (x,x") a²) /[|]=//]]]aa/a/ [|=||] for a positive instance E D, where æ and a' have exactly i non-matching non-lineat coordinates (for 0 < i < - 1). Similarly denote by za the quantity (x, Pl) [|]=]]| l). In particular,
d 2i + 75
75
Ti =
PI ||0|||
and Ta =
p2 |||
19

---

By the above discussion, we have
DI-1
f(z) = Σ ((x, ) + bk (x, = ap1 121|
Σ
k((T:)
+ bp2 11 (1((Ta))
Substituting b and denoting
f<))/|/=]| by 9(5) we get
s(5) api Σ d - ( M(TT(((())
2d x(1)
k(Ta(()) p2x(1)
x(Ta(())
(10)
Now try expand s({) using the Taylor series around Ç = (note that g(7) 1/p1). Note that k' can however be unbounded around -1 and 1. To get around this, write = + where h has bounded first and second derivative, and q has lower order than h for of interest. In particular,
3d/4
d d-1 M((r(())) )
2a
m(Ta(()) P2n(1I)
h(() ap1
Σ dd//
k(Ta(()))
and
x(1)
(T))m(r((s))
((5) = ap1
Σ
b////-<//dd/1
( for and concentration for sums of Observe that g(() = o(ct) for using the estimate eqn. independent Bernoullis. By Taylor's theorem,
s(C) hh())(()++))) #1 -7) h²" ((00)((-)))) - + +c(),,
(11)
for some 0 [), where h(T) 1/vd. will turn out that [1((T) = @(IVVa). [4)(s) =
ol1/v)). This will allow us to complete the proof using the linear approximation of s({) by neglecting the second order term and g(C). We now compute h², h²," treating ||>||| d + Ç2 as constant for exposition (the proof works without this approximation or the reader may think of 7 as o(vā)). Using TC) R p
((())
(((() R apı
A6 (T((C)
2d-13a
B3 K
(Ta(S))
7
+
x'(Ta(())
i
-1||0:|
k(1)
P2 1|al
p2n(1) p2/|=r|
d
72
2d-13a
3a xs1)
72
(T((((() 72
h((() ~ apl
k"
(T((S)
(Ta(C))
o
0 10
(2x(1) p |
1i=0
Plugging |||||||
R3 2 v and substituting a from eqn.
9
((() (1 +
h²(4) 0 ( and
(3x(1)vd
which substituted in eqn.
11 with Ta(C) 20 0, k(0), K' (Ta(C) xs k0()) gives
s(C) 1
1 (1 + (1))(()) (0) (S (1)
Vd
x(1)p²
Hence, g(C) > whenever the coefficient of 1/Vd above is bounded above zero, and similar
condition holds for g({) < O. Using the estimates of from eqn. (9) and x'(0) 1, x(0)
1/ x() = 2, / + in the above gives that g({) > for Ç > -0.68- 1.68/7 and g()) <
for Ç -0.905- 1.905/7-
Predicted label for point (2, where (1)) D has negative label
(t)
Following the same plan, write our point (denoted by x) as (ç. .................... 0, 1).
20

---

Explicit form for f. Begin by finding
T: = + 75
1 1-79 75
and
P1 ||:|||
P2 |ar|
eqn. 10
now gives
g())= 2 |n(To())) Ban(Ta(()) x(1)
k(Ta(S))
p2x(1)
Expanding k(To(()) using Taylor series around ( = -1177
<(To(()) x(0) + (TO0))== + +
for some 0 E [-1, For large d, (0) and T/((0) 0(1/Vo). Hence have
g(C) p2x(1) + B1
Ben(Ta((2)
k(Fa(S)) p2x(1)
x(0) O
EpIP2
Vd x(1)
1 (0)
(²)(0)
1
M(Ta(()) + o(1).
P2 5
Ex(1)
x(1)
As before g(() > whenever the coefficient of 1/p2 above is bounded above zero which happens for < 2 0.73 (for 2 3). Similarly, s(s) < 0 for ( 0.
21

---

B Experiments
In this section, provide experimental details, including hyperparameter tuning setup and some additional experiments.
B.1 Details on the experimental setting
We will first describe the four datasets that have been used in this work.
Imagenette [FastAl 2021): This is a subset of 10 classes Imagenet, that are comparatively easier to classity.
2. b- D-Imagenette: This is binarized version of Imagenette, where only subset of two classes (tench and English springer) is used.
3. Waterbirds-Landdirds (Sagawa al. 2020a This is minioitty group dataset, consisting of waterbirds on water and land background, as well as landbirds on land and water background. This dataset serves on the spurious background feature when predicting the bird class, as most of the training examples have waterbirds water and landbirds land background.
baseline for checking the dependence of model
4. Imagenet (Deng et al. 2009): This is the standard benchmark for large scale image classification.
5. MNIST-CIFAR (Saah al. 2020) This is collage dataset, created by concatenating MNIST and CIFAR images along an axis. This synthetic dataset for evaluating the simplicity bias of trained model.
Setup Throughout the paper, we work with the pretrained representations of the above datasets, obtained by using an Imagenet pretrained Resnet 50. We finetune a 1-hidden layer FCN with hidden dimension of 100 (8000 for imagenet) on top of these representations (keeping the backbone fixed) using SGD with momentum of 0.9. Every model is trained for 20000 (100000 for Imagenet) steps with a warmup and cosine decay learning rate scheduler. For each the runs, we tune the batch size, learning rate and weight decay using validation accuracy. Below are the hyperparameter tuning details:
Batch size E {128, 256}
Learning rate:
Rich regime: E {0.5, 1.0] (for imagenet, E {5.00, 10.0) as learning rate rich regime needs to scale up with the hidden dimension)
Lazy regime: E 0.00),
Weight decay: E {0, 1]
The final numbers reported are averaged across 3 independent runs with the selected hyperparameters.
Evaluation For Imagenette, b-Imagenette, Imagenet and MNIST-CIFAR we report the standard test accuracy in all the experiments. For waterbirds, we report train-adjusted test accuracy, as reported in Sagawa et al. (2020a). Precisely, accuracy for each group present in the test data is individually calculated and then weighed by the proportion of the corresponding group in the train dataset.
B.2 Additional experimental results
In this section, we present few additional experimental results.
Results on Imagenet The evolution of effective rank of the first layer weight matrix shown in Figure As can be seen, the weight matrix becomes sufficiently low rank in thre rich regime the training progresses.
Singular value decay In Figure 6 we provide the singular value decay of the weight matrix for the first model trained in rich regime. can be seen, the top few singular values capture most the Frobenius norm of the matrix.
22

---

Imaganet
maegene
1000
1000
Effective rank
Effective rank
750
750
500
500
0
10
20
30
10
20
30
40
Epochs
Epochs
(a) Rich regime
(b) Lazy regime
Figure 5: Evolution of effective rank of first layer weight matrix (dimension 2048 x 2000) for Imagenet dataset in rich and lazy regime.
Fraction of frobenius norm
Imagenette
0.6
Worcrblinds
Imagenette
10 15
Singular value index
Figure Fraction Frobenius norm captured by the top 7th singular value i.e., of /// of vsiviio if the first layer weight matrix trained in rich regime various datasets.
MNIST-CIFAR In Figure 7 show that ensemble of and Spruj has better gaussian robustness than an ensemble of f and Jind on MNIST-CIFAR dataset.
Non-linearity of decision boundary Figure g shows the decision boundary of f and Spro 2-dimensional subspace spanned by top two singular vectors of the weight matrix. We observe that the decision boundary of the second model is more non-linear compared to that of the first model. We also report a quantitative measure of non-linearity of the decision boundary along the top two singular vectors for and fproj. Basically, we linear classifier to the decision boundary and report its accuracy. As shown in Table the test accuracy obtained by the linear classifier for fproj is less than f.
Variation of LD-SB with depth In Figure 9 and 10 we show the evolution of effective rank of weight matrices for depth-2 and ReLU networks. can be seen, the rank still decreases with training, however the effect is less pronounced for the initial layers. Note that the initialization used these runs was the feature learning initialization proposed Yang Hu (021).
Comparison of Mistake diversity with Evading SB Table shows the results of comparing OrthoP with Evading-SB (Teney et al. 2022). with regards to diversity metrics. As can be seen, OrthoP is either better or comparable to Evading-SB.
23

---

1.00
model essemble-pry
onsonel: Ind
0.95
accuracy
0.90
0.85
0.80
0.75
2
std dev
Figure 7: Variation of test accuracy with the standard deviation of Gaussian noise added to the pretrained representations of MNIST-CIFAR dataset. Model is kept fixed, and values for both the ensembles averaged across runs.
(a) b-Imagenette (f)
(b) b-Imagenette (fpoj)
(c) Waterbirds (f)
(d) Waterbirds (Sros)
Figure 8: Decision boundaries for f and fpnj for B-magnenette and Waterbirds datasets, visualized in the top 2 singular directions of the first layer weight matrix. The decision boundary of fproj is more non-linear compared that of f.
Table 6: Quantitative measurement of non-linearity of decision boundary accuracy of fitted linear classifier to the decision boundary
Dataset Linear- Claaaa--aaasfff-AAcc/]
Sinear-Classifie--ACee
b- D-Imagenette Waterbirds
96.12
95.28 + 0.2
97.28
93.24 0.24
MISTCCA 2ick
MISSTOTAA
Eftective rank
Effectivve rank
40
40
20
20
Epochs
Epochs
(a) Layer 1
(b) Layer 2
Figure 9: Evolution of effective rank of the weight matrices for depth-2 ReLU network on Resnet-50 pretrained representations of the dataset
24

---

1.0
Encotive rank
Effective
62
60
200
500
200
Layer
Layer
(c) layer
Figure 10: Evolution of effective rank of the weight matrices for depth-3 ReLU network on Resnet- 50 pretrained representations of the dataset
Table 7: Mistake diversity and class conditioned logit correlation of models trained sequentially after projecting out important features of the first model (Mist-Div (f, fpej) and CC-LogitCor (f. Fmoj) resp.) and using the method mentioned in Evading-SB ((Mist-Div (J, Jesh) and CC-LogitCory (J, Jesh) resp.). Numbers have been bolded if they are better than the other even after taking one standard deviation into account.
Mist-Div
Mist-Do
CC-LogitCorr
CC-LogitCor (f. (1)
Dataset
(f. (t)
fpri)) (t)
(f. (1)
I-mmanenette
17.86 + 0.0
21.15
1.57
90.61 0.48
90.86 1.08
Imagenette
11.23 + 0.24
11.44 + 0.65
92.48 + 0.68
91 + 0.59
Waterbirds
13.6 0.36
14.53 0.48
95.41
0.41
93.81 0.48
MNIST-CIFAR
0.0
5.56 7.89
76.19 2.38
78.74 :: 2.28
Imagenet
4.95
12.31 0.16
98.83 0.0
92.52 0.01
25

---

C Extended Related Works
In this section, we provide extensive literature survey of various topics that the paper is based on.
Low rank Simplicity Bias in Linear Networks Multiple works have established low rank sim- plicity bias for gradient descent linear networks, both for squared loss well cross-entropy loss. For squared loss, Gunasekar et al. (017)) conjectured that the network is biased towards finding minimum nuclear norm solutions for two-layer linear networks. Arora et al. (2019) refuted the conjecture and instead argued that the network is biased towards finding low rank solutions. Razin & Cohen 2020) provided empirical support the low rank conjecture, by providing synthetic examples where the network drives nuclear norm to infinity, but minimizes the rank of the effective linear mapping. Li et al. 2021) established that for small enough initialization, gradient flow linear networks follows greedy low-rank learning trajectory. For binary classification on linearly separable data, Ji & Telgarsky (2019) showed that the weight matrices of a linear network eventually become rank- I as training progresses.
Low rank Simplicity Bias in Non-Linear Networks There have been few works establishing low- rank simplicity bias for non-linear networks. Phuong & Lampert (202T) showed that, for orthogonally separable datasets, the rank of first layer weights for two-layer ReLU networks eventually tends to two under gradient flow on logistic loss. Freii et al. (2023) extend this result to nearly-orthogon
datasets and Leaky-Reel ReUU networks. Timor al. (2023) showed that, maximum margin ReLU networks of sufficiently large depth implicitly also have low rank. Huh et al. (021) empirically established that the rank of the embeddings learnt by a neural network with ReLU activations goes down as training progresses. Galanti & Poggio (2022) provided an intuition behind the relation between the rank of the weight matrices and various hyperparameter such as batch size, weight decay etc. In contrast to these works, for layer nets, we theoretically and empirically establish that the network depends on an extremely low dimensional projection of the input, and this bias can be utilized to develop a robust classifier.
A Adversarial Robustness Some of the recent works have also shown that gradient flow on shallow networks generally leads to non-robust classifiers. Vardi et al. (2022) showed that for nearly orthogo- nal datasets, gradient flow on 2-layer ReLU networks converges non-robust networks, even when robust classifiers are possible for the dataset. Melamed et al. (2023) show similar results for datasets which lie on a cow-dimesiional linear subspace.
Relation to OOD Many recent works in OOD detection (Cook et al. 2020 Zaeemzadeh et al.
2021) explicitly create low-rank embeddings so that it is easier to discriminate them for OOD point. Other works also implicitly rely on the low-rank nature of the embeddings, [Ndiour et al. (2020) use PCA on the learnt features, and only model the likelihood along the small subspace spanned by the top few directions. Wang et al. (202)) utilise the low rank nature of the embeddings to estimate the perpendicular projection of a given data point to this low rank subspace and combine with logit information to detect OOD datapoints. While there have been works implicitly utilizing the low rank property of embeddings, we note that our paper (i) demonstrates low rank property of the weights, rather than that of embeddings, and (ii) shows that it is a consequence of SB.
Other Simplicity Bias There have been many works exploring the nature of simplicity bias in neural networks, both empirically and theoretically. Kalimeris et al. (2019) empirically demonstrated that SGD on neural networks gradually learns functions of increasing complexity. [Rahaman al. (2018) empirically demonstrated that neural networks tend to learn lower frequency functions first. Ronen et al. (2019) theoretically established that in NTK regime, the convergence rate depends on the eigenvalues of the kernel spectrum. Hacohen et al. (2020) showed that neural networks always learn train and test examples almost in the same order, irrespective of the architecture. Pezeshki et al.) (2021) proposes that gradient starvation at the beginning of training is a potential reason for SB in the lazy/NTK regime but the conditions are hard to interpret. In contrast, our results are shown for any dataset in the IFM model in the rich regime of training. [yu et al. (202T) consider anti- symmetric datasets and show that single hidden layer input homogeneous networks (i.e., without bias parameters) converge to linear classifiers. However, such networks have strictly weaker expressive power compared to those with bias parameters. Hacchen & Weinshall (2022) showed that for deep linear networks, in NTK regime, they learn the higher principal components of the input data first.
26

---

Most of the previous works used simplicity bias as reason behind better generalization of neural nets. However, (Shah al. (2020) showed that extreme simplicity bias could also lead to worse OOD performance.
Learning diverse classifiers: There have been several works that attempt learn diverse classifiers. Most works try to learn such models by ensuring that the input gradients of these models do not align (Ross& & Doshi- Velez 2018 Teney et 2022). Xu al. (2022) proposes a way learn diverse/orthogonal classifiers under the assumption that a complete classifier, that uses all features is available, and demonstrates its utility for various downstream tasks such as style transfer, Lee et al. (2022) learns diverse classifiers by enforcing diversity on unlabeled target data.
Spurious correlations: There has been large body of work which identifies the reasons for spurious correlations in NNs (Sagawa et al. (020b) as well as proposing algorithmic fixes in different settings (Liu et al. 202T Chen et al. 2020b).
There also a large body of work understanding the implicit bias of gradient descent dynamics. Most of these works are for standard linear (Ji & Telgarsky) 2019) or deep linear networks (Soudry et al. 2018: Gunasekar et al. 2018). For nonlinear neural networks, one of the well-known results is for the case of 1-hidden layer neural networks with homogeneous activation functions (Chizat & Bach 2020), which we crucially use in our proofs.
Implicit bias of gradient descent:
D More discussion on the extension of results to deep nets
Extending our theoretical results to deep nets a very exciting and challenging research direction. For shallow as well as deep nets, even in the mean field regime of training, results regarding convergence to global minima have been established (Chizat Bach 2018 8 Fang al. 2021). However, the best of our knowledge, only for I-hidden layer FCN (Chizat & Bach, 2020), a precise characterization of the global minima which gradient flow converges has been established. Understanding this implicit bias of gradient flow is still an open problem for deep nets, which we think is essential for extension of our results deep nets.
E
Convergence to F1 mi--mxx-margnn classifier for ReLU networks
In this section, will provide brief background on Wasserstein gradient flow and state the precise result Chizat & Bach (2020) regarding the asymptotic convergence point of gradient flow on ReLU networks. We will follow the notation Chizat & Bach (2020) for ease of the reader. In this entire section, we will consider that a neural network is parameterized by a probability measure on the neurons and is given by
Mp, - f
((w, x d))(w))
where x)) = b(a (a, 1))+ (+ denotes the positive component, i.e the ReLU activation) with w = (a, E Rd-2
E.1
Wasserstein gradient flow
Gradient flow be defined for many functions f over general metric space X. For given step size define
1
Th+1 E arg min (xx + d(z xx2
21
where d is the metric associated with X. With appropriate interpolation schemes (Santambrogio 2016), this curve converges to the gradient flow curve step size tends O.
Wasserstein metric on the space probability measures is defined as
[/ 11lx-v1'ddy v11 1/P
Wplu, 22) inf
where V1, U2 are two probability measures and 7 coupling between them (i.e marginals o are #1 and 12). Here, will be particularly concerned with the case p = 2. For two discrete measures
27

---

Σ 8z, and 1/2 Σ وyy their Wasserstein distance is defined as
1
(21, 12) min [|:; Vo(()) 1|
over all permutations {1, {1...m) m) neighborhood of 1/1, then the mapping would remain the same within that small neighborhood. Thus, Wasserstein gradient flow on the discrete measure would be the same gradient flow the particles as the notion of distance the same locally. This intuition leads to the proof that gradient flow on 2-layer neural net converges to Wasserstein gradient flow the probability measure js width tends to infinity. This is made formal Chizat & Bach (2018).
m.. {1,...,m. Notice that if we are considering small
E.2 Asymptotic convergence point of gradient flow on ReLU networks
A neural network is parameterized by probability measure on the neurons and is given by
h(,, = 6(w, x ddd((w)
baaaaaaaa where o(w,r) b(a" (xr, 1))+ (+ denotes the positive component, i.e the ReLU activation) with w = (a, b) Rd+2 As the network is 2-homogeneous, a projection the measure on the unit sphere can be defined. The projection operator (IIy) on the sphere for measure μ is defined such that for any continuous function on the sphere,
J Jat (()d[[a(()))))
-/ (tl////o(u//l/lllddddo
J242
Now, let denote the input distribution on the input space N' and let the labeling function ; x - y be deterministic. Then, consider the population objective given by
-log log [ (()((,
F(u) -
Note that log doesn't affect the direction of the gradients, thus, the trajectory of gradient flow this loss is the same as on exponential loss. Also, let the population smooth margin be given by
S(N) = -log (Lr (xpp(-((a)dpp())
For this particular case, [(x) v()) Denote y(x) h(,,x)) by h(n).
Theorem E.1. Suppose that has bounded density and bounded support, and labeling function is continuous, then there exists a Wasserstein gradient flow (He) on F with μi = U($ & 1), i.e, input (resp. output) weights uniformly distributed on the sphere (resp. on {}) If (h((m) converges weakly in ((X). if D. = I1((,,) ([I)))
((+)))))))))) (8)+1)) converges weakly in P( and
F' ('(p)) converges in Clo to I' that satisfies the Morse-Sard property, then h(Poo, is maximizer for maxij 1r,51 $1 mingex (x))
where P(X) denotes the space of probability distributions on X and (($+))) denotes the total mass of the measure l22 gd+1.
To parse the theorem, note that
exp(-f(z))do(z) VS() Jx exp(=](x"))dp(r))
Thus, Vf()) convergence means that the exponentiated normalized margins converge. Also, 8 is similar to the directional convergence of weights. however, in this case, weights are replaced by directions in gd+ For explanation of the Morse-Sard property and the metric Cl please refer to Appendix H Chizat & Bach (2020).
28


Reviews:
[{"review_id": "FwunRQq8xj", "reviewer": "Reviewer_PpnG", "summary": "The paper studies the simplicity bias (SB) for one hidden layer neural networks. Specifically, it I) provides a formal definition of SB for 1-hidden layer NN, II) proves that simplicity bias takes place when a infinite-width 1-hidden layers NNs are trained on linearly separable datasets with a gradient descent algorithm, and III) demonstrates SB on real datasets, motivating a new ensemble approach to encourages diversity in models. The paper contributes on both the theory behind the why simplicity bias happens, and the practices of how to avoid SB in real datasets.", "questions": "1. Could you please describe how the theorems will fail if the 1-hidden layer NN is not of infinite width?\n\n2. Could you please further elaborate how Shah et al.'s work fits in your findings as a special case?\n\n3. Why not empirically validate the LD-SB in other artificially created data in the IFM than the ones used in Shah et al. (2020)'s work?", "limitations": "The authors have adequately addressed the limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "1. The research targets at simplicity bias, a crucial problem facing neural networks.\n\n2. The paper defines and proves existence of a more general SB on a more general type of datasets than the existing work, significantly advancing the theory of SB.\n\n3. The paper makes a connection between theory and practice by empirically studying the Low-dimensional SB on five real datasets.", "weaknesses": "1. The contributions of the paper are poorly presented in the introduction. The definition 1.1 is hard to understand due to the missing notations (e.g. x(1), y(1), P()).  It should be moved to the later chapters. Also, the clarification on why the focus is set on infinite-width regime is not convincing, especially the practical standpoint. \n\n2. The preliminaries are a bit too technical for the general audience. For instance, equations 1), 2) and 3) could be explained with plain and easy-to-understand language.\n\n3. It is claimed that the achieved result is a generalization of Shah et al. (2020)'s work. However, it is not discussed which special case of this work does Shah's work fall into.\n\n4.  The empirical validation part is not as promising as the theoretical part. For instance, the laze regime seems to be a general regime for LD-SB instead of a specific strategy for the lazy initialization of weight; the Rich regime is not comprehensively evaluated by compared with the same two-step approach but other ways of constructing P (and P_ortho).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback. The specific concerns are addressed below:\n\n1. **Presentation of Introduction** - Thanks for the suggestion about presentation in Definition 1.1. We will add exposition and unpack it for readers’ convenience. We would like to clarify that the focus on infinite width setting is primarily because of theoretical tractability. As we mention in lines 97-100, in the feature learning regime, recent works show that finite width neural networks are well approximated by their infinite width counterparts and our empirical results confirm that our theoretical results are valid for finite width networks used in practice.\n\n2. **Preliminaries** - We introduced the technical preliminaries to present our results rigorously. We are happy to add more exposition, but are constrained by the page limits.\n\n3. **Relation to Shah et.al** - Thanks for the question. The dataset used in Shah et al, is a particular instance of the IFM dataset, where one coordinate is linearly separable and other coordinates are separable using 2 linear pieces. Instead, our proof works for any linearly separable dataset within the IFM family. We will include this in the updated draft.\n\n4. **Lazy regime and Rich regime** - Sorry, we couldn’t completely understand the comment of the reviewer. Lazy regime training is indeed dependent on the initialization, and we showcase that it still exhibits LD-SB even when the weights themselves are not low rank. For rich regime, since the weights themselves are low rank, they clearly specify the low dimensional subspace used by the model. Thus, we consider it to be an appropriate choice for P.\n\n5. **Infinite width** - Our theoretical results rely on the characterization provided in Theorem 3.1 by Chizat and Bach (2020). Their proof works only for the infinite width, and hence our theoretical results also hold only for the infinite width. However, our empirical results clearly show that our theory is also valid for the finite width models used in practice.\n\n6. **Other synthetic datasets** - Given that earlier works such as Shah et al. already demonstrated in specialized, synthetic datasets within IFM, the focus of our work is to extend the generality of these results both theoretically and empirically. So, empirically, we focused on real world datasets. We can include results on other synthetic datasets in the appendix, but are not sure if it will teach us anything new. \n\nWe hope that we addressed the reviewers questions and concerns and would appreciate if the reviewer would update their score accordingly.\n"}, {"review_id": "VSGXytFPTD", "reviewer": "Reviewer_aqCb", "summary": "This work characterizes of a type of simplicity bias (SB) in 1 hidden layer neural networks (NNs), termed as LD-SB. A model exhibits LD-SB when its predictions almost entirely depend on a low-dimensional projection of the input.\n\nFor the theoretical analyses, the authors consider NNs with infinite width, in both rich and lazy initialization regimes, and the independent features model. They show that when the data is linearly separable in one co-ordinate, the model almost exclusively relies on a single feature, even in the presence of several other complex features that are non-linearly separable.\n\nThe authors conduct experiments to verify that 1 hidden layer NNs exhibit LD-SB. They consider four real datasets, including Imagenette, Imagenet, Waterbird and employ various quantifiers to demonstrate LD-SB.\n\nFinally, they propose a novel ensemble method called OrthoP. This method encourages feature learning by sequentially training NNs on projections of input data, by separating out directions already learned by previous NNs. The authors demonstrate that this ensemble method learns models that leads to improved robustness to Gaussian noise.", "questions": "$1.$ Contribution in relation to prior works on simplicity bias (SB):\n\nPrevious works on SB, such as [1], focus on a data setting where one feature is linearly separable with a smaller margin, alongside other features that are non-linearly separable. In that setting, using all the features leads to a predictor with a better margin.\n\nHowever, this work states that there exists a unique max-margin predictor that solely relies on the linearly separable feature (theorem 4.1). Based on my understanding, the predictor that uses the non-linearly separable feature in this setting would have a smaller margin. This is also consistent with the experimental results in Fig. 4, which illustrate that an ensemble of $f$ and $f_{\\text{proj}}$ has comparable or worse performance compared to using only $f$, when there is no Gaussian noise. Consequently, it appears reasonable to expect the neural network to learn the linear predictor.\n\nIn contrast, [1] and [2] demonstrate (empirically and theoretically, respectively) that 1 hidden layer NNs learn the max-margin linear predictor even when the global max-margin predictor is non-linear.\n\nI suggest including a discussion to clarify this distinction in comparison to other works. This would help elucidate that the SB characterization examined in this work is slightly different, and this form of SB is not necessarily surprising or detrimental.\n\n$2.$ Regarding Definition 1.1:\n\nThe definition of LD-SB does not explicitly mention that model $f$ obtains high accuracy on samples from distribution $D$. Lines 66-68 state that the third condition implies model $g$ obtains high accuracy, which may not always be true (e.g., if accuracy of $f$ is low). The definition implicitly assumes that $f$ has high accuracy, but this should be mentioned clearly.\n\n$3.$ Section 5: Validation on additional datasets and comparison of OrthoP with other methods:\n\n- The empirical validation of LD-SB in 1 hidden layer NNs on real datasets can be made more extensive, by including more datasets, such as the CelebA dataset or variants of the Imagenet dataset considered in [3].\n\n- Consider comparing OrthoP with other rich feature learning methods, such as gradient starvation [4], etc.\n\n- Fig. 3 should include results for Imagenet dataset, and the MNIST-CIFAR results from the Appendix (Fig. 7) can be incorporated into Fig. 4.\n\n$4.$ Presentation:\n\n- Regarding notation:\n\n    - The notation for indexing two datapoints seems inconsistent (compare line 56 and line 63).\n\n    - Consider using standard symbols for the set of real numbers (line 145), matrices (line 147: use $\\mathbf{W}$ instead of $w$), vectors, etc. Also, moving lines 152-154 to the beginning of Section 3.1 could enhance clarity.\n\n    - The notation $\\bar{(\\cdot)}$ is not defined. $f(x)$ and $f((\\bar{w},\\bar{b},\\bar{a}),x)$ are used interchangeably, but without stating that these are normalized parameters, due to which the definition of margin is not clear.\n\n- Figures:\n    \n    The quality of the figures should be improved. In most cases, the axis labels and legends are too small and difficult to read. In Fig. 4, the blue and green colors are hard to distinguish, and thicker lines in Figs. 3 and 4 could enhance comprehension. \n    \n- Other points:\n\n    - Consider relocating Table 1 to page 8 so that those values can be compared with the values in Table 2.\n\n    - The Waterbirds dataset is referred to as waterbirds-landbirds dataset in some instances.\n\n    - Consider adding related work [5], which empirically demonstrates that using SGD with a large step size encourages neural networks to learn sparse solutions.\n\nReferences:\n\n[1] Shah, H., Tamuly, K., Raghunathan, A., Jain, P., and Netrapalli, P. The pitfalls of simplicity bias in neural networks. Advances in Neural Information Processing Systems, 33:9573–9585, 2020.\n\n[2] Lyu, K., Li, Z., Wang, R., and Arora, S. Gradient descent on two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 34:12978–12991, 2021.\n\n[3] P. Kirichenko, P. Izmailov, and A. G. Wilson. Last layer re-training is sufficient for robustness to spurious correlations. arXiv preprint arXiv:2204.02937, 2022.\n\n[4] Pezeshki, M., Kaba, O., Bengio, Y., Courville, A. C., Precup, D., and Lajoie, G. Gradient starvation: A learning proclivity in neural networks. Advances in Neural Information Processing Systems, 34:1256–1272, 2021.\n\n[5] Andriushchenko, M., Varre, A., Pillaud-Vivien, L., and Flammarion, N. SGD with large step sizes learns sparse features. In ICML, 2023.", "limitations": "[+] The authors have provided the code and experimental details to support reproducibility. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. This work gives a precise characterization of LD-SB, and the quantifiers utilized in Section 4.4.1 to demonstrate LD-SB are well-motivated based on the definition of LD-SB (Def 1.1).\n\n2. The paper as a whole is intriguing. It builds upon existing works and, although the individual contributions may not be highly significant, the ideas are interconnected effectively, providing valuable insights into the behavior of one hidden layer neural networks.", "weaknesses": "1. The statements regarding the contribution in relation to prior work on simplicity bias (line 13, line 78, and the related work paragraph 1), such as [1], overlook certain aspects and do not accurately reflect the scope of the contribution.\n\n2. Definition 1.1 and its accompanying description require further clarification.\n\n3. Section 5, which presents empirical validation of LD-SB on real datasets, as well as the proposed ensemble method OrthoP, could be strengthened by including validation on additional datasets and conducting comparisons with more methods, respectively.\n\n4. The presentation of the paper should be refined. Some notation is non-standard or inconsistent, and the figures appear too small, which affects readability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their valuable comments and suggestions, and are happy to know that the reviewer finds our paper intriguing and insightful. The specific concerns are addressed below:\n\n1. **Contribution in relation to prior works on simplicity bias (SB)** - Thanks for pointing out this issue. The subtlety here is that there are two different notions of margin considered in different papers.\n\n    a. *Max input space margin* refers to maximum distance between the classification boundary and the data points in the input space. [1]  uses this notion.\n\n    b. *Max F1 margin* refers to minimizing L2 norm of the parameters. Both [2] and (Chizat and Bach, 2020) use this notion. (Chizat and Bach, 2020) shows that asymptotically, neural networks with infinite width converge to this classifier always, while [2] discusses early training dynamics and finite width networks.\n\n    Our key contribution is to show that these two notions can be very different. In fact, our IFM datasets have the property that the max input space margin classifiers are nonlinear while the max F1 margin classifier is linear. Consequently, neural networks which converge to the max F1 margin classifier suffer from simplicity bias in the context of max input space margin.\n\n2. **Regarding Definition 1.1** - Thanks for the catch. Yes, we will mention that we are assuming $f$ to have high accuracy.\n\n3. **Validation on additional datasets and comparison of OrthoP with other methods** - Regarding empirical validation, we have obtained results for comparison to an additional method - Evading-SB [3]. We have added this to the one page pdf. Our method performs either comparable or slightly better than this method. We will also move results for Imagenet and for MNIST-CIFAR from the appendix to the main paper.\n\n4. **Presentation** - Thanks for the detailed suggestions on presentation. We will improve the presentation and increase the size of figures as per your suggestions.\n\n[1] Shah, H., Tamuly, K., Raghunathan, A., Jain, P., and Netrapalli, P. The pitfalls of simplicity bias in neural networks. Advances in Neural Information Processing Systems, 33:9573–9585, 2020.\n\n[2] Lyu, K., Li, Z., Wang, R., and Arora, S. Gradient descent on two-layer nets: Margin maximization and simplicity bias. Advances in Neural Information Processing Systems, 34:12978–12991, 2021.\n\n[3] - Teney, D., Abbasnejad, E., Lucey, S., and van den Hengel, A. Evading the simplicity bias: Training a diverse set of models discovers solutions with superior ood generalization. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16761–16772, 2022"}, {"review_id": "E0zeyW5Ck9", "reviewer": "Reviewer_gAuS", "summary": "The main theoretical result in the paper is an application of the theorem of Chizat and Bach (COLT 2020), who showed implicit bias to margin maximisation with respect to the variation norm for logistic loss gradient flow on infinite-width one-hidden layer homogeneous networks.  The focus of the result is when the network activation is ReLU and the training data is given by an independent features model, in which only one feature is linearly separable and the data without that feature is separable (however not linearly); the result then establishes that there is a unique maximum-margin classifier, which only uses the linearly separable feature.  The other theoretical result in the paper is similar; in contrast to the first which applies for initialisation scales that result in the rich regime, it focuses on the lazy regime and on a more streamlined point-mass independent features model for the training data.  In the experimental part, a one-hidden layer ReLU network is trained on top of a pretrained Resnet-50 network with 2048 features, in both the rich and the lazy regimes, and on several datasets (four real: Imagenette, b-Imagenette, waterbirds-landbirds, and Imagenet; and one designed: MNIST-CIFAR).  The experimental results support the hypothesis of simplicity bias, namely that the training converges to a classifier that mostly uses only a small fraction of the features even though nearly as high accuracy is achievable using only the remaining features.  In the last part of the paper, the authors propose an ensemble obtained by adding a second network that is trained only on the features that are mostly not used by the first network, and report some experimental results that show performance superior to an ensemble of two independently trained networks, in particular when the features are corrupted by Gaussian noise.", "questions": "It seems to me that the literature on implicit bias of gradient descent is more related to the present paper than you suggest by reviewing it briefly and at the ends of Section 2 and Appendix C; after all, the work using which you prove the main theoretical result in the paper is from this literature (Chizat and Bach COLT 2020).  Related to this, you write in Appendix C that \"for non-linear networks, the work related to low-rank simplicity bias is rather sparse\", and in addition to Galanti and Poggio (discussed above in this review) only cite the empirical work of Huh et al.  You may want to refer to Phuong and Lampert ICLR 2021, Wang and Pilanci ICLR 2022, Sarussi et al. ICML 2021, Frei et al. arXiv:2210.07082 and ICLR 2023, Timor et al. ALT 2023, Boursier et al. NeurIPS 2022, and Chistikov et al. arXiv:2306.06479.\n\nIt would also be interesting to compare the present work with some of the most related literature on adversarial examples and adversarial robustness, such as Shamir et al. arXiv:2106.10151, Vardi et al. \"Gradient Methods Provably Converge to Non-Robust Networks\" NeurIPS 2022, Melamed et al. arXiv:2303.00783, Lv and Zhu ICLR 2022.\n\nMinor comment: something went wrong in footnote 1?\n", "limitations": "The theoretical results are for infinite-width networks.  In connection with their transfer to finite widths, another submission to NeurIPS 2023 is cited (Anonymous \"Feature-learning networks are consistent across widths at realistic scales\") and included in the supplementary materials, however that work is empirical rather than theoretical.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper addresses one of the greatest current challenges in machine learning research, namely understanding the bias of gradient-based algorithms for training neural networks.  The proofs of the theoretical results are provided in the appendix, which also contains additional material on the experiments and a further discussion of related work.  The setup for the experiments is realistic, using a pretrained Resnet-50 network and some challenging datasets.  The paper is well written.", "weaknesses": "The connections between the theoretical and the experimental results are a little loose, e.g. it is not clear to what extent the projection of the data by P is linearly separable; some depictions and measurements of this are in Figure 8 and Table 6 in the appendix, however they are not very conclusive.  The paper by Galanti and Poggio, cited in Section 2 and twice in Appendix C, is v1 of arXiv:2206.05794, where v2 and v3 have two additional authors and a different title; also a submission with the same title and authors as the latter was rejected by ICLR 2023 (https://openreview.net/forum?id=N7Tv4aZ4Cyx), so I am not sure this paper is sound, even its v3.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments. The specific concerns are addressed below:\n\n1. **Linear separability of the dataset** - Since the embeddings on real datasets are not linearly separable, we plotted the decision boundaries as well as linear classification accuracies to qualitatively/quantitatively measure closeness to linear decision boundaries. Qualitatively, the results suggest that the initial network has a simpler boundary. \n\n2. **More related works** - Thanks for pointing out the various related works. We will make sure to include these works in the next revision with a proper discussion about them.\n\n3. **Infinite width networks** - Yes, the anonymous cited work does an empirical analysis of whether the training dynamics of finite width networks currently being used in practice have started to converge to the infinite width limit, and finds reasonable agreements.\n"}, {"review_id": "VB78PZQs5b", "reviewer": "Reviewer_W8La", "summary": "The concept of simplicity bias (SB), that is, the behaviour for machine learning models only to learn the simplest features at hands when solving a given task, even when there are more complex but more robust features available, is central in this work. First, the authors aim toward a mathematical definition of a special case of it : low dimensional input dependence simplicity bias. The authors study SB in the special case of 1-hidden layer neural networks; from a theoretical perspective, with infinite width, and from a practical point of view, with finite width. The author propose an algorithm, OrthoP, an ensemble method for obtain diverse classifiers.", "questions": "1. How would you reconcile simplicity bias with the widely used technique of regularising in order to obtain simpler models and therefore being more robust? The two approaches seems contradictory. Especially since the use of weight decay was considered in the experimental section, which leads to simpler models.\n\n2. I might be wrong here, but it seems like for a projection matrix $P$, there exist an infinite number of orthogonal subspaces $P_{\\perp} \\in \\mathbb{R}^{d \\times d}$, and especially if rank(P) $\\neq$ d, for there are now more degrees of freedom for constructing $P_{\\perp}$ . Once a projection matrix $P$ is found, how is $P_{\\perp}$ found / computed?\n\n3. How LD-SB is a special case of SB (as stated on Line 50)? The definition of SB given in the abstract goes as follow : « […] [To] learn only the simplest features to solve a task at hand, even in the presence of other, more robust but more complex features ». Isn’t that what LD-SB aims to verify? When a predictor only uses a handful of features for rendering predictions while there are other that would be useful?\n\n4. Line 225 : « Though we assume axis aligned features, our results also hold for any rotation of the dataset. » Is there a proof of that somewhere in the article?\n\n5. In Theorem 4.1, what does $\\delta$ refers to? It is mentioned twice on Line 234, and nowhere else in the core of the paper. It might be explained in the proof, but since the proof is in supplementary material, it would be necessary to explain it in a few words.\n\n6. Is Theorem 4.2 an original result? If not, where does it come from?\n\n7. Line 299 : « In [the rich regime], we set P to be the subspace spanned by the top singular directions of the first layer weight matrix. » Why this choice of P? A justification would be necessary in the article.\n\n8. On Lines 309 to 311, an optimisation problem is described for obtaining P in the lazy regime. The optimisation of a d x d matrix, especially when d is high, can lead to an optimisation of tens of thousands of parameters simultaneously. How well did it converge? How much computing time was necessary in order to obtain an acceptable P matrix?\n\n9. When considering OrthoP in order to obtain a diverse ensemble predictor, is the ensemble necessarily constituted of two predictors? It would have been nice to witness a technique involving the adding of « uncorrelated » predictors until a certain criteria is met.\n\n10. I am not familiar with the metrics described in Section 4.4.1 and those used in Section 5; do they appear in the literature? Are they commonly used?", "limitations": "I feel like the authors discussed adequately the limitations of their work.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The related work section is thorough and discusses many topics, especially considering the work that has been made in Section C (supplementary material). I feel like such an aspect is truly important in a conference paper and is too oftenly neglected.\n\n2. The idea of having « independent », or orthogonal, representations of the input of a problem in order to obtain diverse classifiers is an interesting idea and is intuitive.", "weaknesses": "_Major_\n\n1. The main result of the article holds for a very specific case : on an IFM dataset (binary classification, with a few strong assumptions on it), which is tackled by a 1-hidden layer neural network, in the case of infinite width, with ReLU activation function, with cross-entropy loss. If the assumptions concerning the predictor are defended by the argument that this is respected when doing, for example, transfer learning, or fine-tuning a large model that has already been trained, not much argument is given toward the IFM assumption, apart from the fact that it « has a long history and is widely studied » (Lines 71-72).\n\n2. The metrics described in subsection 4.4.1. ($P_{\\perp}-LC$, $P-LC$, $P_{\\perp}-pC$ and $P-pC$) do not consider the correlation between examples $x_1$ and $x_2$, or the label values $y_1$ and $y_2$. That is a problem, because for example, for P-LC and P-pC on their current form, the goal would not be to obtain the highest values possible (as currently proposed), but to obtain pred($f(\\tilde{x})$) = pred($f(x_2)$) when $y_1$ = $y_2$, and pred($f(\\tilde{x})$) $\\neq$ pred($f(x_2)$) when $y_1 \\neq y_2$ for P-pC and to have , ||f($\\tilde{x}$) – f(x_2)|| / ||f(x_2)|| being low when $x_1$ and $x_2$ are similar, for example, for P-LC. Therefore, I think those metric are too simplistic and might not reflect what their purpose was intended to be.\n\n3. The article is dense. It seems like, in order to save place, many lines of relevant discussion have been removed. For examples, not a line of discussion regarding the results from Table 1, 2, 3 and 4 are given, while they constitute a big part of the empirical experiments. This leaves a feeling of lack in justifications, or the need for the reader to decide whether the reported results are interesting or not.\n\n4. The experimentations given in section 5 do not convince me. For example, Table 5 report mistake diversity when using OrthoP versus when using independent training of two models. It therefore is not surprising that OrthoP yields more diverse predictor when compared to a method that simply does not aim for that goal. The fact that no comparison is made with methods for obtaining diverse classifier in ensemble methods makes it that the reported results are hard to interpret (same goes for robustness experiments). Are the difference between OrthoP and the vanilla approach big or not? The metrics that are considered do not help in this sense. (The same goes for Table 1 and Table 2, but is worse in this last case for it is hard to interpret the results only on a single approach; should we compare P$_\\perp$-LC and P-LC together? And P$_\\perp$-pC and P-pC together?) The use of 5 different datasets satisfies me, but as for the robustness experiment, I do not feel like only looking at Gaussian perturbation of images yields a solid argument toward what’s defended.\n\n\n\n_Typos / Minor_\n\n1. The definition of LD-SB seems a bit arbitrary, for there are no criteria to determine how k should be small relatively to d, and there are no relationship between $\\epsilon_1$ and $\\epsilon_2$.\n\n2. In the definition of LD-SB, I would state clearly that model g and model f must have the same architecture, for it is implicit but crucial. In the same line of thoughts, it is mentioned that model g is trained on examples $(P_{\\perp}x,y)$, where $(x,y) \\~ \\mathcal{D}$, but it is not stated whether the accuracy (Acc(.)) refers to the training accuracy or the test accuracy.\n\n3. Line 72 : « IFM has a long history and is widely studied », but no citation toward recent or any other important work than Lewis (1998) is given, so the statement that IFM is currently widely studied is not supported.\n\n4. On Line 145 : « Consider instances $x \\in \\mathcal{R}^d$ […] » instead of $\\mathbb{R}^d$.\n\n5. On Lines 147-148, the model is defined as a function of weights and biases for the hidden layer, but only weights for the output layer. Was it on purpose?\n\n6. In subsection 3.2, the parameters from the hidden layer are sampled independently from the parameters from the output layer. But later in subsection 3.3, we are told that they are sampled dependently. Does that have something to do with the infinite width regime, or was that a mistake?\n\n7. At the beginning of Line 213 : missing end-of-sentence point.\n\n8. Line 306 : « […] it turns out that the \\textbf{effective} rank [...] ».\n\n9. Line 323 : no closed parenthesis in denominator\n\n10. Line 324 : « […] where we abuse notation by using f($x_i$) [...] »; no such notation is used.\n\n11. The only place where a reported accuracy is explicitly defined as being test accuracy is for Figure 4. It might be implicit that the same goes for Tables 3 and 4, but I would make it explicit.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback.\n\n1. **IFM assumption** - The IFM assumption is a simple and intuitive assumption that enables us to define features – in this model features are distinguished from each other through conditional independence. IFM played a key role in developing a thorough understanding of naive Bayes classifiers e.g., [1, 2] which has since been extended to more general data models [3, 4]. Similarly, given the close connection between features and simplicity bias (SB), we believe that IFM is a canonical setting to study SB. \n\n    *Linear separability*: Using linear evaluation protocol is the standard practice in many domains e.g. self-supervised learning domain. The pretrained representations obtained by self-supervised learning are clustered enough that they can be well-separated using a linear classifier [5]. Moreover, as demonstrated in Figure 8 and Table 6 in the Appendix, the pretrained representations on various datasets are close to being linearly separable.\n\n2. **Metrics in subsection 4.4.1** - We are not sure we fully understand the reviewer’s question. The main reasoning behind using these metrics is to verify the points specified in the definition of LD-SB (Definition 1.1), specifically the second point. In particular, $P_\\perp-PC$ exactly quantifies the value of $\\epsilon_1$ in Definition 1.1. Moreover, a very similar metric to these, randomized accuracy, was used for measuring the dependence of a classifier on a subset of coordinates in [6]. The mistake diversity used in Table 5 is a general metric for checking functional similarity of neural networks [7] .\n\n3. **Discussion of results** - We would like to point out that in Section 4.4.1, we had already discussed the relevance of various metrics being displayed in Tables 1, 2, 3 and 4. However, in the updated version, we will make sure to have more discussion regarding the results in their respective sections.\n\n4. **Experiments in Section 5** - Based on the reviewer’s suggestion, we conducted comparison experiments with another diverse training method  - Evading SB [8]. The results are attached in the one page pdf. Our method generally performs comparable to or better than this method in terms of class conditioned logit correlation or mistake diversity. In terms of gaussian robustness, the two methods are comparable as well, with our method slightly outperforming on ImageNet. Regarding gaussian robustness, we consider it as a metric that estimates the input space margin of the model.\n\n5. **Simplicity Bias and Regularization** - Our paper indeed shows that in some cases, simplicity bias does not lead to robust classifiers. This was first pointed out explicitly in [6], and we are demonstrating it in a more general setting. The use of weight decay in the experimental section is to replicate practical setup and demonstrate the adverse effects of simplicity bias in the setting. Also, note that F1 max margin classifier is the global minimizer of a limiting case of $\\ell_2$ regularization, where the weight decay coefficient tends to zero [9]. Our results prove that F1 max margin classifier suffers from the adverse effects of simplicity bias.\n\n6. **Construction of $P_{\\perp}$** - This is not the case. For a given projection matrix $P$,  $P_{\\perp} = I - P$.\n\n7. **LD-SB vs SB** - LD-SB provides a specific way of instantiating SB. Its not necessary that SB is a linear subspace of the input features, it could be a lower dimensional non-linear manifold. However, for 1 hidden layer networks, because of the low dimensional structure of the weights itself, we can identify the simple classifier within a low dimensional linear subspace of the inputs.\n\n8. **Proof for rotation of dataset** - We will add this proof to the appendix. It follows as we can rotate the weight matrix in response to any rotation of the inputs.\n\n9. **Theorem 4.1** - Delta represents the dirac delta distribution. Thanks for the catch.\n\n10. **Theorem 4.2** - Yes, theorem 4.2 is an original result. Its proof is provided in Appendix A.2\n\n11. **Choice of P** -  We have chosen this $P$ as the weight matrix itself is low rank. One justification is provided by Theorem 4.1, which shows that for a linearly separable dataset within IFM, this indeed represents the direction of the linearly separable coordinate. Thus projecting out this direction will force the model to rely on non-linearly separable coordinates.\n\n12. **Optimization of P in lazy regime** - If the projection is done over a $k$ dimensional subspace (where k is generally in the range of 1-10 as mentioned in Table 1), then the optimization is done over $d \\times k$ basis matrix, which can be used to create the projection matrix. This generally converged well, and was completed within a minute.\n\nWe hope that we addressed the reviewers questions and concerns and would appreciate it if the reviewer would update their score accordingly.\n\n[1] An empirical study of the naive Bayes classifier by I Rish -- IJCAI 2001 workshop on empirical methods in artificial intelligence, 2001\n\n[2] Naive (Bayes) at forty: The independence assumption in information retrieval DD Lewis - ECML, 1998\n\n[3] An improvement to naive bayes for text classification W Zhang, F Gao - Procedia Engineering, 2011\n\n[4] A non-parametric mixture of Gaussian naive Bayes classifiers based on local independent features AH Jahromi, M Taheri - AISP 2017\n\n[5] Chen et al. A simple framework for contrastive learning of visual representations. In ICML, pp. 1597–1607. PMLR, 2020a.\n\n[6] Shah et al. - The pitfalls of simplicity bias in neural networks. Advances in NeurIPS, 33:9573–9585, 2020.\n\n[7] Klabunde at al. 2023 - Similarity of Neural Network Models, arXiv:2305.06329\n\n[8] Teney et al. 2021 - Evading the Simplicity Bias: Training a Diverse Set of Models Discovers Solutions with Superior OOD Generalization, CVPR, 2022\n\n[9] Lenaic Chizat 2020 - Gradient descent for wide two-layer neural networks – II: Generalization and implicit bias"}]

Author rebuttal:
We thank the reviewers for their insightful comments. Based on the suggestions of the reviewers, we conducted a comparison experiment with another diverse training method - Evading SB [1]. The results are attached in the one page pdf and summarized below:

1. Our method generally performs comparable or slightly better than Evading-SB in terms of class conditioned logit correlation and mistake diversity.
2. In terms of gaussian robustness, both methods are comparable, with our method slightly outperforming on ImageNet and Evading-SB slightly outperforming on MNIST-CIFAR. 

[1] Teney et al. - Evading the Simplicity Bias: Training a Diverse Set of Models Discovers Solutions with Superior OOD Generalization, CVPR, 2022

Meta-review: In this paper, the authors introduce a definition of simplicity bias of a model (in terms of its being a function of a low-dimensional projection of input data.  The authors rigorously prove, in an infinite-width 1-hidden-layer neural network regime, that when data is linearly separable in a single dimension, the model will find that dimension, even if the resulting networks is less robust than other potential models.  The authors illustrate their notion of simplicity bias on networks trained on real datasets such as ImageNet.  Finally, the authors introduce an ensemble method that yields models that are more robust to Gaussian noise.  The main strength of the paper is that it provides an advance to the theory of simplicity bias, which is a central challenge of the deep learning field.   There were some concerns that the theory and experimental sections were not as tightly related as they could be, but given the difficulty of theoretical advances in this area, those differences are to be expected.
Decision: Accept (poster)