Original: S IMPLICITY BIAS IN 1 - HIDDEN LAYER NEURAL NETWORKS
Revised: S IMPLICITY BIAS IN 1 - HIDDEN LAYER NEURAL NETWORKS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Recent works Shah et al.
Revised: Recent works Shah et al.

Original: ; Chen et al.
Revised: ; Chen et al.

Original: (2021) have demonstrated that neural networks exhibit extreme simplicity bias (SB).
Revised: (2021) have demonstrated that neural networks exhibit extreme simplicity bias (SB).

Original: 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.
Revised: 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.

Original: Due to the lack of a general and rigorous definition of features , these works showcase SB on semi-synthetic datasets such as ColorMNIST, MNIST-CIFAR where defining features is relatively easier.
Revised: Due to the lack of a general and rigorous definition of features , these works showcase SB on semi-synthetic datasets such as ColorMNIST, MNIST-CIFAR where defining features is relatively easier.

Original: In this work, we rigorously define as well as thoroughly establish SB for one hidden layer neural networks.
Revised: In this work, we rigorously define as well as thoroughly establish SB for one hidden layer neural networks.

Original: 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 Imagenette and Waterbirds-Landbirds indeed depend on a low dimensional projection of the inputs, thereby demonstrating
Revised: 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 Imagenette and Waterbirds-Landbirds indeed depend on a low dimensional projection of the inputs, thereby demonstrating

Original: 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.
Revised: 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.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: It is well known that neural networks ( are vulnerable to distribution shifts as well as to adversarial examples (Szegedy et al., 2014;
Revised: It is well known that neural networks ( are vulnerable to distribution shifts as well as to adversarial examples (Szegedy et al., 2014;

Original: Hendrycks et al., 2021).
Revised: Hendrycks et al., 2021).

Original: A recent line of work (Geirhos et al., 2018; Shah et al., 2020;
Revised: A recent line of work (Geirhos et al., 2018; Shah et al., 2020;

Original: Geirhos et al., 2020) proposes that Simplicity
Revised: Geirhos et al., 2020) proposes that Simplicity

Original: Bias (SB) – aka shortcut learning – i.e., the tendency of neural networks (NNs) to learn only the simplest features over other useful but more complex features, is a key reason behind this non-robustness.
Revised: Bias (SB) – aka shortcut learning – i.e., the tendency of neural networks (NNs) to learn only the simplest features over other useful but more complex features, is a key reason behind this non-robustness.

Original: The argument is roughly as follows: for example, in the classification of swans vs bears, as illustrated in Figure 1, there are many features such as background, color of the animal, shape of the animal etc.
Revised: The argument is roughly as follows: for example, in the classification of swans vs bears, as illustrated in Figure 1, there are many features such as background, color of the animal, shape of the animal etc.

Original: that can be used for classification.
Revised: that can be used for classification.

Original: However using only one or few of them can lead to models that are not robust to specific distribution shifts, while using all the features can lead to more robust models.
Revised: However using only one or few of them can lead to models that are not robust to specific distribution shifts, while using all the features can lead to more robust models.

Original: Several recent works have demonstrated SB on a variety of semi-real constructed datasets (Geirhos et al., 2018; Shah et al., 2020; Chen et al., 2021), and have hypothesized SB to be the key reason for
Revised: Several recent works have demonstrated SB on a variety of semi-real constructed datasets (Geirhos et al., 2018; Shah et al., 2020; Chen et al., 2021), and have hypothesized SB to be the key reason for

Original: NN’s brittleness to distribution shifts (Shah et al., 2020).
Revised: NN’s brittleness to distribution shifts (Shah et al., 2020).

Original: However, such observations are still only for specific semi-real datasets, and a general method that can identify SB on a given dataset and a given model is still missing in literature.
Revised: However, such observations are still only for specific semi-real datasets, and a general method that can identify SB on a given dataset and a given model is still missing in literature.

Original: Such a method would be useful not only to estimate the robustness of a model but could also help in designing more robust models.
Revised: Such a method would be useful not only to estimate the robustness of a model but could also help in designing more robust models.

Original: A key challenge in designing such a general method to identify (and potentially fix) SB is that the notion of feature itself is vague and lacks a rigorous definition.
Revised: A key challenge in designing such a general method to identify (and potentially fix) SB is that the notion of feature itself is vague and lacks a rigorous definition.

Original: Existing works like Geirhos et al.
Revised: Existing works like Geirhos et al.

Original: (2018); Shah et al.
Revised: (2018); Shah et al.

Original: ; Chen et al.
Revised: ; Chen et al.

Original: (2021) avoid this challenge of vague feature definition by using carefully designed datasets (e.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.
Revised: (2021) avoid this challenge of vague feature definition by using carefully designed datasets (e.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.

Original: Contributions : One of the main contributions of this work is to provide a precise definition of a particular simplicity bias – LD-SB– of 1 -hidden layer neural networks .
Revised: Contributions : One of the main contributions of this work is to provide a precise definition of a particular simplicity bias – LD-SB– of 1 -hidden layer neural networks .

Original: In particular, we characterize
Revised: In particular, we characterize

Original: SB as low dimensional input dependence of the model.
Revised: SB as low dimensional input dependence of the model.

Original: Concretely,
Revised: Concretely,

Original: Definition 1.1 (LD-SB) .
Revised: Definition 1.1 (LD-SB) .

Original: A model f : R d → R c with inputs x ∈ R d and outputs f ( x ) ∈ R c (e.g., logits for c classes), trained on a distribution ( x, y ) ∼ D satisfies LD-SB if there exists a projection matrix P ∈ R d × d satisfying the following:
Revised: A model f : R d → R c with inputs x ∈ R d and outputs f ( x ) ∈ R c (e.g., logits for c classes), trained on a distribution ( x, y ) ∼ D satisfies LD-SB if there exists a projection matrix P ∈ R d × d satisfying the following:

Original: • rank ( P ) =
Revised: • rank ( P ) =

Original: k ≪ d , • f
Revised: k ≪ d , • f

Original: ( Px 1 + P ⊥ x 2 ) ≈ f ( x 1 ) ∀ ( x 1 , y 1 ) , ( x 2 , y 2 ) ∼ D , and • An independent model g trained on ( P ⊥ x, y ) where ( x, y ) ∼ D achieves high accuracy.
Revised: ( Px 1 + P ⊥ x 2 ) ≈ f ( x 1 ) ∀ ( x 1 , y 1 ) , ( x 2 , y 2 ) ∼ D , and • An independent model g trained on ( P ⊥ x, y ) where ( x, y ) ∼ D achieves high accuracy.

Original: Here P ⊥ is the projection matrix onto the subspace orthogonal to P .
Revised: Here P ⊥ is the projection matrix onto the subspace orthogonal to P .

Original: In words, LD-SB says that there exists a small k -dimensional subspace (given by the projection matrix P ) in the input space R d , which is the only thing that the model f considers in labeling any input point x .
Revised: In words, LD-SB says that there exists a small k -dimensional subspace (given by the projection matrix P ) in the input space R d , which is the only thing that the model f considers in labeling any input point x .

Original: In particular, if we mix two data points x 1 and x 2 by using the projection of x 1 onto
Revised: In particular, if we mix two data points x 1 and x 2 by using the projection of x 1 onto

Original: P and the projection of x 2 onto the orthogonal subspace P ⊥ , the output of f on this mixed point
Revised: P and the projection of x 2 onto the orthogonal subspace P ⊥ , the output of f on this mixed point

Original: Px 1 + P ⊥ x 2 is the same as that on x 1 .
Revised: Px 1 + P ⊥ x 2 is the same as that on x 1 .

Original: This would have been fine if the subspace P ⊥ does not contain any feature useful for classification.
Revised: This would have been fine if the subspace P ⊥ does not contain any feature useful for classification.

Original: However, the third bullet point says that P ⊥ indeed contains features that are useful for classification since an independent model g trained on ( P ⊥ x, y ) achieves high accuracy.
Revised: However, the third bullet point says that P ⊥ indeed contains features that are useful for classification since an independent model g trained on ( P ⊥ x, y ) achieves high accuracy.

Original: Furthermore, theoretically, we demonstrate LD-SB of 1 -hidden layer NNs 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.
Revised: Furthermore, theoretically, we demonstrate LD-SB of 1 -hidden layer NNs 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.

Original: IFM has a long history and is widely studied, especially in the context of naive-Bayes classifiers Lewis (1998).
Revised: IFM has a long history and is widely studied, especially in the context of naive-Bayes classifiers Lewis (1998).

Original: 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.
Revised: 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.

Original: Empirically, we demonstrate LD-SB on three real world datasets: binary and multiclass version of
Revised: Empirically, we demonstrate LD-SB on three real world datasets: binary and multiclass version of

Original: Imagenette (FastAI, 2021) as well as waterbirds-landbirds (Sagawa et al., 2020a) dataset.
Revised: Imagenette (FastAI, 2021) as well as waterbirds-landbirds (Sagawa et al., 2020a) dataset.

Original: Compared to the results in Shah et al.
Revised: Compared to the results in Shah et al.

Original: (2020), our results (i) theoretically show LD-SB in a fairly general setting and (ii) empirically show LD-SB on real datasets.
Revised: (2020), our results (i) theoretically show LD-SB in a fairly general setting and (ii) empirically show LD-SB on real datasets.

Original: Finally, building upon these insights, we propose a simple ensemble method – OrthoP – that sequentially constructs NNs by projecting out principle input data directions that are used by previous
Revised: Finally, building upon these insights, we propose a simple ensemble method – OrthoP – that sequentially constructs NNs by projecting out principle input data directions that are used by previous

Original: NNs.
Revised: NNs.

Original: We demonstrate that this method can lead to significantly more robust ensembles for realworld datasets in presence of simple distribution shifts like Gaussian noise.
Revised: We demonstrate that this method can lead to significantly more robust ensembles for realworld datasets in presence of simple distribution shifts like Gaussian noise.

Original: Why only 1 -hidden layer networks?
Revised: Why only 1 -hidden layer networks?

Original: : One might wonder why the results in this paper are restricted to 1 -hidden layer networks and why they are interesting.
Revised: : One might wonder why the results in this paper are restricted to 1 -hidden layer networks and why they are interesting.

Original: We present two reasons.
Revised: We present two reasons.

Original: From a theoretical standpoint, prior works have thoroughly characterized the training dynamics of infinite width 1 -hidden layer networks under different initialization schemes (Chizat et al., 2019) and have also identified the limit points of gradient descent for such networks (Chizat &
Revised: From a theoretical standpoint, prior works have thoroughly characterized the training dynamics of infinite width 1 -hidden layer networks under different initialization schemes (Chizat et al., 2019) and have also identified the limit points of gradient descent for such networks (Chizat &

Original: Bach, 2020).
Revised: Bach, 2020).

Original: Our results crucially build upon these prior works.
Revised: Our results crucially build upon these prior works.

Original: On the other hand,
Revised: On the other hand,

Original: we do not have such a clear understanding of the dynamics of deeper networks.
Revised: we do not have such a clear understanding of the dynamics of deeper networks.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: However, finetuning on target datasets might not utilize all the features in the pretrained model.
Revised: However, finetuning on target datasets might not utilize all the features in the pretrained model.

Original: Consequently, approaches that can train robust finetuning heads (such as a 1 -hidden layer network on top) can be quite effective.
Revised: Consequently, approaches that can train robust finetuning heads (such as a 1 -hidden layer network on top) can be quite effective.

Original: Extending our results to deeper networks and to other architectures is an exciting direction of research from both theoretical and practical points of view.
Revised: Extending our results to deeper networks and to other architectures is an exciting direction of research from both theoretical and practical points of view.

Original: Paper organization :
Revised: Paper organization :

Original: This paper is organized as follows.
Revised: This paper is organized as follows.

Original: Section 2 presents related work.
Revised: Section 2 presents related work.

Original: Sectionpresents preliminaries.
Revised: Sectionpresents preliminaries.

Original: Our main results on LD-SB are presented in Section 4.
Revised: Our main results on LD-SB are presented in Section 4.

Original: Section 5 presents results on training diverse classifiers.
Revised: Section 5 presents results on training diverse classifiers.

Original: We conclude in Section 6.
Revised: We conclude in Section 6.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Simplicity Bias : Subsequent to Shah et al.
Revised: Simplicity Bias : Subsequent to Shah et al.

Original: (2020), there have been several papers investigating the presence/absence of SB in various networks as well as reasons behind SB Scimeca et al.
Revised: (2020), there have been several papers investigating the presence/absence of SB in various networks as well as reasons behind SB Scimeca et al.

Original: Of these, Huh et al.
Revised: Of these, Huh et al.

Original: and Galanti & Poggio (2022) are the most closely related works to ours.
Revised: and Galanti & Poggio (2022) are the most closely related works to ours.

Original: Huh et al.
Revised: Huh et al.

Original: (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.
Revised: (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.

Original: Galanti & Poggio (2022) provide a theoretical intuition behind the relation between various hyperparameters (such as learning rate, batch size etc.)
Revised: Galanti & Poggio (2022) provide a theoretical intuition behind the relation between various hyperparameters (such as learning rate, batch size etc.)

Original: and rank of learnt weight matrices, and demonstrate it empirically.
Revised: and rank of learnt weight matrices, and demonstrate it empirically.

Original: In contrast, we prove LD-SB theoretically on the IFM model as well as empirically validate this on real datasets.
Revised: In contrast, we prove LD-SB theoretically on the IFM model as well as empirically validate this on real datasets.

Revised: Moreover, we also show how to use LD-SB to train a second diverse model and combine it to obtain a robust ensemble.

Original: Pezeshki et al.
Revised: Pezeshki et al.

Original: (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.
Revised: (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.

Original: In contrast, our results are shown for any dataset in the IFM model in the rich regime of training.
Revised: In contrast, our results are shown for any dataset in the IFM model in the rich regime of training.

Original: Finally Lyu et al.
Revised: Finally Lyu et al.

Original: (2021) considers anti-symmetric datasets and show that single hidden layer input homogeneous networks (i.e., without bias parameters) converge to linear classifiers.
Revised: (2021) consider anti-symmetric datasets and show that single hidden layer input homogeneous networks (i.e., without bias parameters) converge to linear classifiers.

Original: However, such networks have strictly weaker expressive power compared to those with bias parameters.
Revised: However, such networks have strictly weaker expressive power compared to those with bias parameters.

Original: Learning diverse classifiers : There have been several works that attempt to learn diverse classifiers.
Revised: Learning diverse classifiers : There have been several works that attempt to learn diverse classifiers.

Original: 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 al., 2022).
Revised: 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 al., 2022).

Original: Xu et al.
Revised: Xu et al.

Original: (2022) proposes a way to 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.
Revised: (2022) proposes a way to 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.

Original: Lee et al.
Revised: Lee et al.

Original: (2022) learns diverse classifiers by enforcing diversity on unlabeled target data.
Revised: (2022) learns diverse classifiers by enforcing diversity on unlabeled target data.

Original: Spurious correlations : There has been a large body of work which identifies the reasons for spurious correlations in NNs (Sagawa et al., 2020b) as well as proposing algorithmic fixes in different settings (Liu et al., 2021; Chen et al., 2020).
Revised: Spurious correlations : There has been a large body of work which identifies the reasons for spurious correlations in NNs (Sagawa et al., 2020b) as well as proposing algorithmic fixes in different settings (Liu et al., 2021; Chen et al., 2020).

Original: Implicit bias of gradient descent : There is also a large body of work understanding the implicit bias of gradient descent dynamics.
Revised: Implicit bias of gradient descent : There is also a large body of work understanding the implicit bias of gradient descent dynamics.

Original: Most of these works are for standard linear (Ji & Telgarsky, 2019) or deep linear networks (Soudry et al., 2018; Gunasekar et al., 2018).
Revised: Most of these works are for standard linear (Ji & Telgarsky, 2019) or deep linear networks (Soudry et al., 2018; Gunasekar et al., 2018).

Original: 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.
Revised: 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.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: In this section, we provide the notation and background on infinite width max-margin classifiers that is required to interpret the results of this paper.
Revised: In this section, we provide the notation and background on infinite width max-margin classifiers that is required to interpret the results of this paper.

Original: 3. B ASIC NOTIONS
Revised: 3. B ASIC NOTIONS

Original: 1-hidden layer neural networks and loss function.
Revised: 1-hidden layer neural networks and loss function.

Original: Consider instances x ∈ R d and labels y ∈ {± 1 } jointly distributed as D .
Revised: Consider instances x ∈ R d and labels y ∈ {± 1 } jointly distributed as D .

Original: A 1-hidden layer neural network model for predicting the label for a given instance x , is defined by parameters ( ¯ w ∈ R m × d , ¯ b ∈ R m , ¯ a ∈ R m ) .
Revised: A 1-hidden layer neural network model for predicting the label for a given instance x , is defined by parameters ( ¯ w ∈ R m × d , ¯ b ∈ R m , ¯ a ∈ R m ) .

Original: For a fixed activation function ϕ , given input instance x , the model is given as f (( ¯ w, ¯ b, ¯ a ) , x ) := ⟨ ¯ a, ϕ
Revised: For a fixed activation function ϕ , given input instance x , the model is given as f (( ¯ w, ¯ b, ¯ a ) , x ) := ⟨ ¯ a, ϕ

Original: ( ¯ wx + ¯ b ) ⟩ , where ϕ ( · ) is applied elementwise.
Revised: ( ¯ wx + ¯ b ) ⟩ , where ϕ ( · ) is applied elementwise.

Original: The cross entropy loss L for a given model f , input x and label y is given as L ( f ( x ) , y ) def = log(1 + exp( − yf (( ¯ w, ¯ b, ¯ a ) , x ))) .
Revised: The cross entropy loss L for a given model f , input x and label y is given as L ( f ( x ) , y ) def = log(1 + exp( − yf (( ¯ w, ¯ b, ¯ a ) , x ))) .

Original: Margin.
Revised: Margin.

Original: For data distribution D , the margin of a model f ( x ) is given as min ( x,y )

Original: ∼D yf ( x ) .

Original: Notation.
Revised: Notation.

Original: Here is some useful notation that we will use repeatedly.
Revised: Here is some useful notation that we will use repeatedly.

Original: For a matrix A , A ( i, . ) denotes the i th row of A .
Revised: For a matrix A , A ( i, . ) denotes the i th row of A .

Original: For any k ∈ N , S k − 1 denotes the surface of the unit norm Euclidean sphere in dimension k .
Revised: For any k ∈ N , S k − 1 denotes the surface of the unit norm Euclidean sphere in dimension k .

Original: 3. I NITIALIZATIONS
Revised: 3. I NITIALIZATIONS

Original: The gradient descent dynamics of the network depends strongly on the scale of initialization.
Revised: The gradient descent dynamics of the network depends strongly on the scale of initialization.

Original: In this work, we primarily consider rich regime initialization.
Revised: In this work, we primarily consider rich regime initialization.

Original: Rich regime.
Revised: Rich regime.

Original: In rich regime initialization, for any i ∈ [ m ] , the parameters ( ¯ w ( i, . ) , ¯ b ( i ) ) of the first layer are sampled from a uniform distribution on S d .
Revised: In rich regime initialization, for any i ∈ [ m ] , the parameters ( ¯ w ( i, . ) , ¯ b ( i ) ) of the first layer are sampled from a uniform distribution on S d .

Original: Each ¯ a ( i ) is sampled from Unif {− 1 , 1 } , and the output of the network is scaled down by 1 m (Chizat & Bach, 2020).
Revised: Each ¯ a ( i ) is sampled from Unif {− 1 , 1 } , and the output of the network is scaled down by 1 m (Chizat & Bach, 2020).

Original: This is roughly equivalent to Xavier initialization Glorot & Bengio (2010), where the weight parameters in both the layers are initialized approximately as N (0 , 2 m ) when m ≫ d .
Revised: This is roughly equivalent to Xavier initialization Glorot & Bengio (2010), where the weight parameters in both the layers are initialized approximately as N (0 , 2 m ) when m ≫ d .

Original: In addition, we also present some results for the lazy regime initialization described below.
Revised: In addition, we also present some results for the lazy regime initialization described below.

Original: Lazy regime.
Revised: Lazy regime.

Original: In the lazy regime, the weight parameters in the first layer are initialized with N (0 , 1 d ) , those of second layer are initialized with N (0 , 1 m ) and the biases are initialized to 

Revised: In the lazy regime, the weight parameters in the first layer are initialized with N (0 , 1 d ) , those of second layer are initialized with N (0 , 1 m ) and the biases are initialized to

Original: (Bietti & Mairal, 2019; Lee et al., 2019).
Revised: (Bietti & Mairal, 2019; Lee et al., 2019).

Original: This is approximately equivalent to Kaiming initialization (He et al., 2015).
Revised: This is approximately equivalent to Kaiming initialization (He et al., 2015).

Original: 3. T HE INFINITE WIDTH LIMIT
Revised: 3. T HE INFINITE WIDTH LIMIT

Original: For 1-hidden layer neural networks with ReLU activation in the infinite width limit i.e., as m → ∞ ,
Revised: For 1-hidden layer neural networks with ReLU activation in the infinite width limit i.e., as m → ∞ ,

Original: Jacot et al.
Revised: Jacot et al.

Original: ; Chizat et al.
Revised: ; Chizat et al.

Original: ; Chizat & Bach (2020) gave interesting characterizations of the trained model.
Revised: ; Chizat & Bach (2020) gave interesting characterizations of the trained model.

Original: As mentioned above, the training process of these models falls into one of two regimes depending on the scale of initialization (Chizat et al., 2019):
Revised: As mentioned above, the training process of these models falls into one of two regimes depending on the scale of initialization (Chizat et al., 2019):

Original: Rich regime.
Revised: Rich regime.

Original: In the infinite width limit, the neural network parameters can be thought of as a distribution ν over triples ( w, b, a ) ∈ S
Revised: In the infinite width limit, the neural network parameters can be thought of as a distribution ν over triples ( w, b, a ) ∈ S

Original: d +1 where w ∈ R d , b, a ∈ R .
Revised: d +1 where w ∈ R d , b, a ∈ R .

Original: Under the rich regime initialization, the function f computed by the model can be expressed as
Revised: Under the rich regime initialization, the function f computed by the model can be expressed as

Original: [Equation]
Revised: [Equation]

Original: Chizat & Bach (2020) showed that the training process with rich initialization can be thought of as gradient flow on the Wasserstein-2 space and gave the following characterization of the trained model under the cross entropy
Revised: Chizat & Bach (2020) showed that the training process with rich initialization can be thought of as gradient flow on the Wasserstein-2 space and gave the following characterization of the trained model under the cross entropy

Original: loss E ( x,y ) ∼D
Revised: loss E ( x,y ) ∼D

Original: [ L ( ν, ( x, y ))] .
Revised: [ L ( ν, ( x, y ))] .

Original: Theorem 3.1.
Revised: Theorem 3.1.

Original: (Chizat & Bach, 2020)
Revised: (Chizat & Bach, 2020)

Original: Under rich initialization in the infinite width limit with cross entropy loss, if gradient flow on 1-hidden layer NN with ReLU activation converges, it converges to a maximum margin classifier ν ∗ given as
Revised: Under rich initialization in the infinite width limit with cross entropy loss, if gradient flow on 1-hidden layer NN with ReLU activation converges, it converges to a maximum margin classifier ν ∗ given as

Original: [Equation]
Revised: [Equation]

Original: where P ( S d +1 ) denotes the space of distributions over S d +1 .
Revised: where P ( S d +1 ) denotes the space of distributions over S d +1 .

Original: This training regime is known as the ‘rich’ regime since it learns data dependent features ⟨ w, ·⟩ .
Revised: This training regime is known as the ‘rich’ regime since it learns data dependent features ⟨ w, ·⟩ .

Original: Lazy regime.
Revised: Lazy regime.

Original: Jacot et al.
Revised: Jacot et al.

Original: (2018) showed that in the infinite width limit, the neural network behaves like a kernel machine.
Revised: (2018) showed that in the infinite width limit, the neural network behaves like a kernel machine.

Original: This kernel is popularly known as the Neural Tangent Kernel(NTK), and is given by K ( x, x ′ )
Revised: This kernel is popularly known as the Neural Tangent Kernel(NTK), and is given by K ( x, x ′ )

Original: = (cid:68) ∂f ( x ) ∂W , ∂f ( x ′ )
Revised: = (cid:68) ∂f ( x ) ∂W , ∂f ( x ′ )

Original: ∂W (cid:69) , where W denotes the set of all trainable weight parameters.
Revised: ∂W (cid:69) , where W denotes the set of all trainable weight parameters.

Original: This initialization regime is called ’lazy’ regime since the weights do not change much from initialization, and the NTK remains almost constant, i.e, the network does not learn data dependent features.
Revised: This initialization regime is called ’lazy’ regime since the weights do not change much from initialization, and the NTK remains almost constant, i.e, the network does not learn data dependent features.

Original: We will use the following characterization of the NTK for 1-hidden layer neural networks.
Revised: We will use the following characterization of the NTK for 1-hidden layer neural networks.

Original: Theorem 3.2.
Revised: Theorem 3.2.

Original: Bietti & Mairal (2019)
Revised: Bietti & Mairal (2019)

Original: Under lazy regime initialization in the infinite width limit, the
Revised: Under lazy regime initialization in the infinite width limit, the

Original: NTK for 1-hidden layer neural networks with ReLU activation i.e., ϕ ( u ) = max( u, 0) , is given as
Revised: NTK for 1-hidden layer neural networks with ReLU activation i.e., ϕ ( u ) = max( u, 0) , is given as

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Revised: 1 Feature  1 Feature  1 Feature

Original: Lazy regime for binary classification.
Revised: Lazy regime for binary classification.

Original: Soudry et al.
Revised: Soudry et al.

Original: (2018) showed that for linearly separable datasets, gradient descent for linear predictors on logistic loss converges to the max-margin support vector machine (SVM) classifier.
Revised: (2018) showed that for linearly separable datasets, gradient descent for linear predictors on logistic loss converges to the max-margin support vector machine (SVM) classifier.

Original: This implies that, any sufficiently wide neural network, when1 Feature  1 Feature  1 Feature trained for a finite time in the lazy regime on a dataset that is separable by the finite-width induced
Revised: This implies that, any sufficiently wide neural network, when trained for a finite time in the lazy regime on a dataset that is separable by the finite-width induced

Original: NTK, will tend towards the L 2 max-margin-classifier given by
Revised: NTK, will tend towards the L 2 max-margin-classifier given by

Original: [Equation]
Revised: [Equation]

Original: (3) where H represents the Reproducing Kernel Hilbert Space (RKHS) associated with the finite width kernel (Chizat, 2020).
Revised: where H represents the Reproducing Kernel Hilbert Space (RKHS) associated with the finite width kernel (Chizat, 2020).

Original: With increasing width, this kernel tends towards the infinite-width NTK (which is universal (Ji et al., 2020)).
Revised: With increasing width, this kernel tends towards the infinite-width NTK (which is universal (Ji et al., 2020)).

Original: Therefore, in lazy regime, we will focus on the L 2 maxmargin-classifier induced by the infinite-width NTK.
Revised: Therefore, in lazy regime, we will focus on the L 2 maxmargin-classifier induced by the infinite-width NTK.

Original:  C HARACTERIZATION OF SB IN 1 - HIDDEN LAYER NEURAL NETWORKS
Revised:  C HARACTERIZATION OF SB IN 1 - HIDDEN LAYER NEURAL NETWORKS

Original: In this section, we first theoretically characterize the SB exhibited by gradient descent on linearly separable datasets in the independent features model (IFM) .
Revised: In this section, we first theoretically characterize the SB exhibited by gradient descent on linearly separable datasets in the independent features model (IFM) .

Original: The main result, stated in Theorem 4.1, is that for binary classification of inputs in R d , even if there is a single coordinate in which the data is linearly separable, gradient descent dynamics will learn a model that relies solely on this coordinate, even when there are an arbitrarily large number d − 1 of coordinates in which the data is separable, but by a non-linear classifier.
Revised: The main result, stated in Theorem 4.1, is that for binary classification of inputs in R d , even if there is a single coordinate in which the data is linearly separable, gradient descent dynamics will learn a model that relies solely on this coordinate, even when there are an arbitrarily large number d − 1 of coordinates in which the data is separable, but by a non-linear classifier.

Original: In other words, the simplicity bias of these networks is characterized by low dimensional input dependence , which we denote by LD-SB.
Revised: In other words, the simplicity bias of these networks is characterized by low dimensional input dependence , which we denote by LD-SB.

Original: We then experimentally verify that NNs trained on some real datasets do indeed satisfy LD-SB.
Revised: We then experimentally verify that NNs trained on some real datasets do indeed satisfy LD-SB.

Original: 4. D ATASET
Revised: 4. D ATASET

Original: We consider datasets in the independent features model (IFM), where the joint distribution over ( x, y ) satisfies p 
Revised: We consider datasets in the independent features model (IFM), where the joint distribution over ( x, y ) satisfies p ( x, y ) =

Original: [Equation]
Revised: r ( y ) (cid:81)

Original: d i =1 q i ( x i | y ) , i.e, the features are distributed independently conditioned on the label y .
Revised: d i = 1 q i ( x i | y ) , i.e, the features are distributed independently conditioned on the label y .

Original: Here r ( y ) is a distribution over {− 1 , +1 } and q i ( x i | y ) denotes the conditional distribution of i
Revised: Here r ( y ) is a distribution over {− 1 , +1 } and q i ( x i | y ) denotes the conditional distribution of i

Original: th -coordinate x i given y .
Revised: th -coordinate x i given y .

Original: IFM is widely studied in literature, particularly in the context of naive-Bayes classifiers Lewis (1998).
Revised: IFM is widely studied in literature, particularly in the context of naive-Bayes classifiers Lewis (1998).

Original: 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 at least one other with a non-linear boundary.
Revised: 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 at least one other with a non-linear boundary.

Original: See Figure 2 for an illustrative example.
Revised: See Figure 2 for an illustrative example.

Original: • One of the coordinates (say, the 1 st coordinate WLOG) is separable by a linear decision boundary with margin γ (see Figure 2), i.e, ∃ γ > 0 , such that γ ∈ Supp ( q 1 ( x 1 | y
Revised: • One of the coordinates (say, the 1 st coordinate WLOG) is separable by a linear decision boundary with margin γ (see Figure 2), i.e, ∃ γ > 0 , such that γ ∈ Supp ( q 1 ( x 1 | y

Original: = +1)) ⊆
Revised: = +1)) ⊆

Original: [ γ, ∞ ) and − γ ∈ Supp ( q 1 ( x 1 | y = − 1)) ⊆
Revised: [ γ, ∞ ) and − γ ∈ Supp ( q 1 ( x 1 | y = − 1)) ⊆

Original: ( −∞ , − γ ] , where Supp ( · ) denotes the support of a distribution.
Revised: ( −∞ , − γ ] , where Supp ( · ) denotes the support of a distribution.

Original: None of the other coordinates is linearly separable.
Revised: None of the other coordinates is linearly separable.

Original: More precisely, for all the other coordinates i ∈
Revised: More precisely, for all the other coordinates i ∈

Original: [ d ] \ { 1 } , 0 ∈ Supp ( q i ( x i | y = − 1)) and {− 1 , +1 } ⊆
Revised: [ d ] \ { 1 } , 0 ∈ Supp ( q i ( x i | y = − 1)) and {− 1 , +1 } ⊆

Original: Supp ( q i ( x i | y = +1)) .
Revised: Supp ( q i ( x i | y = +1)) .

Original: • The dataset can be perfectly classified even without using the linear coordinate.
Revised: • The dataset can be perfectly classified even without using the linear coordinate.

Original: This means, ∃ i ̸
Revised: This means, ∃ i ̸

Original: = 1 , such that q i ( x i | y ) has disjoint support for y
Revised: = 1 , such that q i ( x i | y ) has disjoint support for y

Original: = +

Revised: = +

Original: and y
Revised: and y

Original: =
Revised: =

Original: − 1 .
Revised: − 1 .

Original: Though we assume axis aligned features, our results also hold for any rotation of the dataset.
Revised: Though we assume axis aligned features, our results also hold for any rotation of the dataset.

Original: While our results hold in the general IFM setting, in comparison, current results for SB e.g., Shah et al.
Revised: While our results hold in the general IFM setting, in comparison, current results for SB e.g., Shah et al.

Original: (2020), are obtained for very specialized datasets within IFM, and do not apply to IFM in general.
Revised: (2020), are obtained for very specialized datasets within IFM, and do not apply to IFM in general.

Original: 4. M AIN RESULT
Revised: 4. M AIN RESULT

Original: Our main result states that, for rich initialization (Section 3.2), NNs demonstrate LD-SB for any
Revised: Our main result states that, for rich initialization (Section 3.2), NNs demonstrate LD-SB for any

Original: IFM dataset satisfying the above conditions.
Revised: IFM dataset satisfying the above conditions.

Original: Its proof appears in Appendix A.1.
Revised: Its proof appears in Appendix A.1.

Original: Theorem 4.1.
Revised: Theorem 4.1.

Original: For any dataset in the IFM model satisfying the above conditions and γ ≥ 1 , if gradient flow for 1-hidden layer FCN under rich initialization in the infinite width limit with cross entropy loss converges, it converges to ν ∗ = 0 .
Revised: For any dataset in the IFM model satisfying the above conditions and γ ≥ 1 , if gradient flow for 1-hidden layer FCN under rich initialization in the infinite width limit with cross entropy loss converges, it converges to ν ∗ = 0 .

Original: 5 δ θ 1 + 0 .
Revised: 5 δ θ 1 + 0 .

Original: 5 δ θ 2 on S d +1 , where θ 
Revised: 5 δ θ 2 on S d +1 , where θ 1 = (

Original: [Equation]
Revised: [Equation]

Original: [1 , 0 , · · · , 0] denotes first standard basis vector.
Revised: [1 , 0 , · · · , 0] denotes first standard basis vector.

Original: This implies f ( ν ∗ , Px 1 + P ⊥ x 2 ) =
Revised: This implies f ( ν ∗ , Px 1 + P ⊥ x 2 ) =

Original: f ( ν ∗ , x 1 ) ∀ ( x 1 , y 1 ) , ( x 2 , y 2 ) ∼ D , where P represents the (rank-1) projection matrix on first coordinate.
Revised: f ( ν ∗ , x 1 ) ∀ ( x 1 , y 1 ) , ( x 2 , y 2 ) ∼ D , where P represents the (rank-1) projection matrix on first coordinate.

Original: Moreover, since at least one of the coordinates { 2 , . . .
Revised: Moreover, since at least one of the coordinates { 2 , . . .

Original: , d } has disjoint support for q
Revised: , d } has disjoint support for q

Original: i
Revised: i

Original: ( x i | y
Revised: ( x i | y

Original: = +1)
Revised: = +1)

Original: and
Revised: and

Original: q i ( x i | y = − 1) , P ⊥ ( x ) can still perfectly classify the given dataset, thereby implying LD-SB.
Revised: q i ( x i | y = − 1) , P ⊥ ( x ) can still perfectly classify the given dataset, thereby implying LD-SB.

Original: It is well known that the rich regime is more relevant for the practical performance of NNs since it allows for feature learning, while lazy regime does not (Chizat et al., 2019).
Revised: It is well known that the rich regime is more relevant for the practical performance of NNs since it allows for feature learning, while lazy regime does not (Chizat et al., 2019).

Original: Nevertheless, in the next section, we present theoretical evidence that LD-SB holds even in the lazy regime, by considering a much more specialized dataset within IFM.
Revised: Nevertheless, in the next section, we present theoretical evidence that LD-SB holds even in the lazy regime, by considering a much more specialized dataset within IFM.

Original: 4. L AZY REGIME
Revised: 4. L AZY REGIME

Original: In this regime, we will work with the following dataset within the IFM family:
Revised: In this regime, we will work with the following dataset within the IFM family:

Original: For y ∈ {± 1 } we generate ( x, y ) ∈ D as[Equation] , x i [Equation]
Revised: For y ∈ {± 1 } we generate ( x, y ) ∈ D as

Original: (cid:26)

Original: [Equation]
Revised: [Equation]

Original: for y

Original: [Equation] for [Equation] Although the dataset above is a 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 are not.
Revised: Although the dataset above is a 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 are not.

Original: For this dataset, we provide the characterization of max-margin NTK (asdefined in Eqn.(3)) below:
Revised: For this dataset, we provide the characterization of max-margin NTK (as in Eqn.

Original: Theorem 4.2.
Revised: Theorem 4.2.

Original: For sufficiently small ϵ > 0 , there exists an absolute constant N such that for all d > N and γ ∈ [7 , ϵ √ d ) , the L 2 max-margin classifier for joint training of both the layers of 1-hidden layer FCN in the NTK regime on the dataset D , i.e., any f satisfying Eqn.
Revised: For sufficiently small ϵ > 0 , there exists an absolute constant N such that for all d > N and γ ∈ [7 , ϵ √ d ) , the L 2 max-margin classifier for joint training of both the layers of 1-hidden layer FCN in the NTK regime on the dataset D , i.e., any f satisfying Eqn.

Original: (3) satisfies:
Revised: (3) satisfies:

Original: [Equation]
Revised: [Equation]

Original: where P represents the projection matrix on the first coordinate and pred ( f ( x )) represents the predicted label by the model f on x .
Revised: where P represents the projection matrix on the first coordinate and pred ( f ( x )) represents the predicted label by the model f on x .

Original: The above theorem shows that the prediction on a mixed example Px 1 + P ⊥ x 2 is the same as that on x 1 , thus establishing LD-SB.
Revised: The above theorem shows that the prediction on a mixed example Px 1 + P ⊥ x 2 is the same as that on x 1 , thus establishing LD-SB.

Original: The proof for this theorem is provided in Appendix A.2.
Revised: The proof for this theorem is provided in Appendix A.2.

Original: 4. E MPIRICAL VERIFICATION
Revised: 4. E MPIRICAL VERIFICATION

Original: In this section, we will present empirical results demonstrating LD-SB on 3 real datasets: Imagenette (FastAI, 2021), a binary version of Imagenette (b-Imagenette) and waterbirds-landbirds (Sagawa et al., 2020a) as well as one designed dataset MNIST-CIFAR (Shah et al., 2020).
Revised: In this section, we will present empirical results demonstrating LD-SB on 3 real datasets: Imagenette (FastAI, 2021), a binary version of Imagenette (b-Imagenette) and waterbirds-landbirds (Sagawa et al., 2020a) as well as one designed dataset MNIST-CIFAR (Shah et al., 2020).

Original: More details about the datasets can be found in Appendix B.1.
Revised: More details about the datasets can be found in Appendix B.1.

Original: 4.4. E XPERIMENTAL SETUP
Revised: 4.4. E XPERIMENTAL SETUP

Original: We take Imagenet pretrained Resnet-50 models, with 2048 features, for feature extraction and train a 1 -hidden layer fully connected network, with ReLU nonlinearity, and 100 hidden units, for classification on each of these datasets.
Revised: We take Imagenet pretrained Resnet-50 models, with 2048 features, for feature extraction and train a 1 -hidden layer fully connected network, with ReLU nonlinearity, and 100 hidden units, for classification on each of these datasets.

Original: During the finetuning process, we freeze the backbone Resnetmodel and train only the 1 -hidden layer head (more details in Appendix B.1) .
Revised: During the finetuning process, we freeze the backbone Resnetmodel and train only the 1 -hidden layer head (more details in Appendix B.1) .

Original: Demonstrating LD-SB :
Revised: Demonstrating LD-SB :

Original: Given a model f ( · ) , we establish its low dimensional SB by identifying a small dimensional subspace, identified by its projection matrix P , such that if we mix inputs x 1 and
Revised: Given a model f ( · ) , we establish its low dimensional SB by identifying a small dimensional subspace, identified by its projection matrix P , such that if we mix inputs x 1 and

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: x ) is always close to the model’s output on x 1 i.e., f ( x 1 ) .
Revised: x ) is always close to the model’s output on x 1 i.e., f ( x 1 ) .

Original: We measure closeness in four metrics: (1) P ⊥ -randomized accuracy ( P ⊥ -RA): accuracy on the dataset ( Px 1 + P ⊥ x 2 , y 1 ) where ( x 1 , y 1 ) and ( x 2 , y 2 ) are sampled iid from the dataset, (2) P -randomized accuracy ( P -RA): accuracy on the dataset ( Px 1 + P ⊥ x 2 , y 2 ) , (3) P ⊥ logit change ( P ⊥ -LC): relative change wrt logits of x 1 i.e.,
Revised: We measure closeness in four metrics: (1) P ⊥ -randomized accuracy ( P ⊥ -RA): accuracy on the dataset ( Px 1 + P ⊥ x 2 , y 1 ) where ( x 1 , y 1 ) and ( x 2 , y 2 ) are sampled iid from the dataset, (2) P -randomized accuracy ( P -RA): accuracy on the dataset ( Px 1 + P ⊥ x 2 , y 2 ) , (3) P ⊥ logit change ( P ⊥ -LC): relative change wrt logits of x 1 i.e.,

Original: ∥ f
Revised: ∥ f

Original: ( (cid:101) x ) − f ( x 1 ) ∥ / ∥
Revised: ( (cid:101) x ) − f ( x 1 ) ∥ / ∥

Original: f
Revised: f

Original: ( x 1 ) ∥ , and (4) P logit change ( P -LC): relative change wrt logits of x 2 i.e.,
Revised: ( x 1 ) ∥ , and (4) P logit change ( P -LC): relative change wrt logits of x 2 i.e.,

Original: ∥ f
Revised: ∥ f

Original: ( (cid:101) x ) − f ( x 2 ) ∥ / ∥
Revised: ( (cid:101) x ) − f ( x 2 ) ∥ / ∥

Original: f
Revised: f

Original: ( x 2 ) ∥ .
Revised: ( x 2 ) ∥ .

Original: As described in Sections 4.2 and 4.3, the training of 1 -hidden layer neural networks might follow different trajectories depending on the scale of initialization.
Revised: As described in Sections 4.2 and 4.3, the training of 1 -hidden layer neural networks might follow different trajectories depending on the scale of initialization.

Original: So, the subspace projection matrix P will be obtained in different ways for rich vs lazy regimes.
Revised: So, the subspace projection matrix P will be obtained in different ways for rich vs lazy regimes.

Original: For rich regime, we will empirically show that the first layer weights have a low rank structure as per Theorem 4.1 while for lazy regime, we will show that though first layer weights do not exhibit low rank structure, the model still has low dimensional dependence on the input as per Theorem 4.2.
Revised: For rich regime, we will empirically show that the first layer weights have a low rank structure as per Theorem 4.1 while for lazy regime, we will show that though first layer weights do not exhibit low rank structure, the model still has low dimensional dependence on the input as per Theorem 4.2.

Original: 4.4. R ICH REGIME
Revised: 4.4. R ICH REGIME

Original: Theorem 4.1 suggests that asymptotically, the first layer weight matrix will be low rank.
Revised: Theorem 4.1 suggests that asymptotically, the first layer weight matrix will be low rank.

Original: However, since we train only for a finite amount of time, the weight matrix will only be approximately low rank.
Revised: However, since we train only for a finite amount of time, the weight matrix will only be approximately low rank.

Original: To quantify this, we use the notion of effective rank Roy & Vetterli (2007) to measure the rank of the first layer weight matrix.
Revised: To quantify this, we use the notion of effective rank Roy & Vetterli (2007) to measure the rank of the first layer weight matrix.

Original: Definition 4.3.
Revised: Definition 4.3.

Original: Given a matrix M , its effective rank is defined as:
Revised: Given a matrix M , its effective rank is defined as:

Original: [Equation]
Revised: [Equation]

Original: where σ i ( M ) denotes the i th singular value of M and σ i ( M ) 2 def =
Revised: where σ i ( M ) denotes the i th singular value of M and σ i ( M ) 2 def =

Original: [Equation]
Revised: [Equation]

Original: One way to interpret the effective rank is that it is the exponential of von-Neumann entropy Petz (2001) of the matrix
Revised: One way to interpret the effective rank is that it is the exponential of von-Neumann entropy Petz (2001) of the matrix

Original: MM ⊤ Tr ( MM ⊤ ) , where Tr ( · ) denotes the trace of a matrix.
Revised: MM ⊤ Tr ( MM ⊤ ) , where Tr ( · ) denotes the trace of a matrix.

Original: For illustration, the effective rank of a projection matrix onto k dimensions equals k .
Revised: For illustration, the effective rank of a projection matrix onto k dimensions equals k .

Original: Figure 3a shows the evolution of the effective rank through training on the four datasets.
Revised: Figure 3a shows the evolution of the effective rank through training on the four datasets.

Original: We observe that the effective rank of the weight matrix decreases drastically towards the end of training.
Revised: We observe that the effective rank of the weight matrix decreases drastically towards the end of training.

Original: To confirm that this indeed leads to LD-SB, we set P to be the subspace spanned by the top singular
Revised: To confirm that this indeed leads to LD-SB, we set P to be the subspace spanned by the top singular

Original: [Table] [Table]
Revised: [Table] [Table]

Original: directions of the first layer weight matrix and compute P and P ⊥ randomized accuracies as well as the relative logit change.
Revised: directions of the first layer weight matrix and compute P and P ⊥ randomized accuracies as well as the relative logit change.

Original: The results, presented in Table 1 confirm LD-SB in the rich regime on these datasets.
Revised: The results, presented in Table 1 confirm LD-SB in the rich regime on these datasets.

Original: 4.4. L AZY REGIME
Revised: 4.4. L AZY REGIME

Original: For the lazy regime, it turns out that the rank of first layer weight matrix remains high throughout training, as shown in Figure 3b.
Revised: For the lazy regime, it turns out that the rank of first layer weight matrix remains high throughout training, as shown in Figure 3b.

Original: However, we are able to find a low dimensional projection matrix P which satisfies the conclusions in Theorem 4.2 as the solution to an optimization problem.
Revised: However, we are able to find a low dimensional projection matrix P satisfying the conditions of LD-SB (as stated in Def 1.1) as the solution to an optimization problem.

Original: More concretely, given a pretrained model f and a rank r , we obtain a projection matrix P solving:
Revised: More concretely, given a pretrained model f and a rank r , we obtain a projection matrix P solving:

Original: [Equation]
Revised: [Equation]

Original: where U [ L ] represents a uniform distribution over all the L labels, ( x 1 , y 1 ) , · · · , ( x n , y n ) are training examples and L ( · , · ) is the cross entropy loss.
Revised: where U [ L ] represents a uniform distribution over all the L labels, ( x 1 , y 1 ) , · · · , ( x n , y n ) are training examples and L ( · , · ) is the cross entropy loss.

Original: We reiterate that the optimization is only over P , while the model parameters f are unchanged.
Revised: We reiterate that the optimization is only over P , while the model parameters f are unchanged.

Original: In words, the above function ensures that the neural network produces correct predictions along P and uninformative predictions along P ⊥ .
Revised: In words, the above function ensures that the neural network produces correct predictions along P and uninformative predictions along P ⊥ .

Original: Tablepresents the results for P ⊥ and P -RA as well as LC.
Revised: Tablepresents the results for P ⊥ and P -RA as well as LC.

Original: As can be seen, even in this case, we are able to find small rank projection matrices demonstrating LD-SB.
Revised: As can be seen, even in this case, we are able to find small rank projection matrices demonstrating LD-SB.

Original:  T RAINING DIVERSE CLASSIFIERS USING OrthoP
Revised:  T RAINING DIVERSE CLASSIFIERS USING OrthoP

Original: Motivated by our results on low dimensional SB, in this section, we present a natural way to train diverse models, so that an ensemble of such models could mitigate SB.
Revised: Motivated by our results on low dimensional SB, in this section, we present a natural way to train diverse models, so that an ensemble of such models could mitigate SB.

Original: More concretely, given an initial model f with a low dimensional projection P that captures its input dependence, we train another model f proj by projecting the input through P ⊥ i.e., instead of using dataset ( x i , y i ) for training, we use ( P ⊥ x i , y i ) for training the second model (denoted by f proj ).
Revised: More concretely, given an initial model f with a low dimensional projection P that captures its input dependence, we train another model f proj by projecting the input through P ⊥ i.e., instead of using dataset ( x i , y i ) for training, we use ( P ⊥ x i , y i ) for training the second model (denoted by f proj ).

Original: We refer to this training procedure as OrthoP for orthogonal projection .
Revised: We refer to this training procedure as OrthoP for orthogonal projection .

Original: Given any two models f and ˜ f , we evaluate their diversity using two metrics.
Revised: Given any two models f and ˜ f , we evaluate their diversity using two metrics.

Original: The first is mistake diversity: Mist-Div (cid:16) f, ˜ f (cid:17) def = 1 − | [Equation]
Revised: The first is mistake diversity: Mist-Div (cid:16) f, ˜ f (cid:17) def = 1 − | { i : f ( x i ) ̸

Original: [Equation] [Equation]
Revised: = y i & ˜ f ( x i ) ̸

Original: | min [Equation] [Equation] i 
Revised: = y i }| min( |{ i : f ( x i ) ̸

Original: [Equation]

Original: | , [Equation] { i : ˜ 
Revised: = y i }| , |{ i : ˜ f

Original: [Equation]

Original:  x i ) ̸
Revised: ( x i ) ̸

Original: = y i }| , where we abuse notation by using
Revised: = y i }| , where we abuse notation by using

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: f ( x i ) (resp.
Revised: f ( x i ) (resp.

Original: ˜
Revised: ˜

Original: f
Revised: f

Original: ( x i ) ) to denote the class predicted by f (resp ˜ f ) on x i .
Revised: ( x i ) ) to denote the class predicted by f (resp ˜ f ) on x i .

Original: Higher Mist-Div (cid:16) f, ˜ f (cid:17) means that there is very little overlap in the mistakes of f and ˜ f .
Revised: Higher Mist-Div (cid:16) f, ˜ f (cid:17) means that there is very little overlap in the mistakes of f and ˜ f .

Original: The second is class conditioned logit correlation i.e., correlation between outputs of f and ˜ f , conditioned on the class.
Revised: The second is class conditioned logit correlation i.e., correlation between outputs of f and ˜ f , conditioned on the class.

Original: More concretely,
Revised: More concretely,

Original: CC-LogitCorr (cid:16) f, ˜ f (cid:17) =
Revised: CC-LogitCorr (cid:16) f, ˜ f (cid:17) =

Original: (cid:80) y ∈Y Corr ( [ f ( x i )] , [ ˜ f ( x i )]: y i =
Revised: (cid:80) y ∈Y Corr ( [ f ( x i )] , [ ˜ f ( x i )]: y i =

Original: y ) |Y| , where corr ([ f ( x i )] , [ ˜ f
Revised: y ) |Y| , where corr ([ f ( x i )] , [ ˜ f

Original: ( x i )] : y i =
Revised: ( x i )] : y i =

Original: y ) represents the empirical correlation between the logits of f and ˜ f on the data points where the true label is y .
Revised: y ) represents the empirical correlation between the logits of f and ˜ f on the data points where the true label is y .

Original: Table 3 compares the diversity of two independently trained models ( f and f ind ) with that of two sequentially trained models ( f and f proj ) as above.
Revised: Table 3 compares the diversity of two independently trained models ( f and f ind ) with that of two sequentially trained models ( f and f proj ) as above.

Original: The results demonstrate that f and f proj are more diverse compared to f and f ind .
Revised: The results demonstrate that f and f proj are more diverse compared to f and f ind .

Original: Figure 4 shows the decision boundary of f and f proj on 2 -dimensional subspace spanned by top two singular vectors of the weight matrix.
Revised: Figure 4 shows the decision boundary of f and f proj on 2 -dimensional subspace spanned by top two singular vectors of the weight matrix.

Original: We observe that the decision boundary of the second model is more non-linear compared to that of the first model.
Revised: We observe that the decision boundary of the second model is more non-linear compared to that of the first model.

Original: Finally, Figure 5 shows the variation of test accuracy with the strength of gaussian noise added to the pretrained representations of the dataset.
Revised: Finally, Figure 5 shows the variation of test accuracy with the strength of gaussian noise added to the pretrained representations of the dataset.

Original: We can see that an ensemble of f and f proj is much more robust as compared to an ensemble of f and f ind .
Revised: We can see that an ensemble of f and f proj is much more robust as compared to an ensemble of f and f ind (where an ensemble is obtained by averaging the logits of the two models).

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: In this work, we characterize the simplicity bias exhibited by one hidden layer neural networks in terms of the low-dimensional input dependence of the model.
Revised: In this work, we characterize the simplicity bias exhibited by one hidden layer neural networks in terms of the low-dimensional input dependence of the model.

Original: We provide a theoretical proof for linearly separable datasets, and validate it empirically on real datasets.
Revised: We provide a theoretical proof for linearly separable datasets, and validate it empirically on real datasets.

Original: Based on this characterization, we also propose a simple way to train diverse models and show that it leads to models with significantly better Gaussian noise robustness.
Revised: Based on this characterization, we also propose a simple way to train diverse models and show that it leads to models with significantly better Gaussian noise robustness.

Original: This work is an initial step towards rigorously defining simplicity bias or shortcut learning of neural networks, which is one of the major challenges to their real-world deployment (Geirhos et al., 2020).
Revised: This work is an initial step towards rigorously defining simplicity bias or shortcut learning of neural networks, which is one of the major challenges to their real-world deployment (Geirhos et al., 2020).

Original: Providing a similar characterization for deeper nets and other architectures is an important research direction, which, in our opinion, requires a better understanding of the training dynamics and limit points of gradient descent on these networks.
Revised: Providing a similar characterization for deeper nets and other architectures is an important research direction, which, in our opinion, requires a better understanding of the training dynamics and limit points of gradient descent on these networks.

Original: R A P ROOFS FOR RICH AND LAZY REGIME A. R ICH REGIME
Revised: A P ROOFS FOR RICH AND LAZY REGIME A. R ICH REGIME

Original: We restate Theorem 4.1 below and prove it.
Revised: We restate Theorem 4.1 below and prove it.

Original: Theorem A.1.
Revised: Theorem A.1.

Original: For any dataset in IFM model satisfying the conditions in Section 4.1, γ ≥and f ( ν, x ) as in Eqn.
Revised: For any dataset in IFM model satisfying the conditions in Section 4.1, γ ≥and f ( ν, x ) as in Eqn.

Original: (1), the distribution ν ∗ = 0 .
Revised: (1), the distribution ν ∗ = 0 .

Original: 5 δ θ 1 + 0 .
Revised: 5 δ θ 1 + 0 .

Original: 5 δ θ 2 on S d +1 is the unique max-margin classifier satisfying Eqn.
Revised: 5 δ θ 2 on S d +1 is the unique max-margin classifier satisfying Eqn.

Original: (2), where θ 1 = (
Revised: (2), where θ 1 = (

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: √ 2) and e 1 def =
Revised: √ 2) and e 1 def =

Original: [1 , 0 , · · · , 0] denotes first standard basis vector.
Revised: [1 , 0 , · · · , 0] denotes first standard basis vector.

Original: In particular, this implies that if gradient flow for 1-hidden layer FCN under rich initialization in the infinite width limit with cross entropy loss converges, it converges to ν ∗ satisfying f ( ν ∗ , Px 1 + P ⊥ x 2 ) =
Revised: In particular, this implies that if gradient flow for 1-hidden layer FCN under rich initialization in the infinite width limit with cross entropy loss converges, it converges to ν ∗ satisfying f ( ν ∗ , Px 1 + P ⊥ x 2 ) =

Original: f ( ν ∗ , x 1 ) ∀ ( x 1 , y 1 ) , ( x 2 , y 2 ) ∈ D , where P represents the (rank-1) projection matrix on the first coordinate.
Revised: f ( ν ∗ , x 1 ) ∀ ( x 1 , y 1 ) , ( x 2 , y 2 ) ∈ D , where P represents the (rank-1) projection matrix on the first coordinate.

Original: Proof of Theorem A.1: (Chizat & Bach, 2020) showed the following primal-dual characterization of maximum margin classifiers in eqn.
Revised: Proof of Theorem A.1: (Chizat & Bach, 2020) showed the following primal-dual characterization of maximum margin classifiers in eqn.

Original: Lemma A.2.
Revised: Lemma A.2.

Original: (Chizat & Bach, 2020)
Revised: (Chizat & Bach, 2020)

Original: ν ∗ satisfies eqn.
Revised: ν ∗ satisfies eqn.

Original: (2) if there exists a data distribution p ∗ such that the following two complementary slackness conditions hold:
Revised: (2) if there exists a data distribution p ∗ such that the following two complementary slackness conditions hold:

Original: [Equation]
Revised: [Equation]

Original: The plan is to construct a distribution p ∗ that satisfies the conditions of the above Lemma.
Revised: The plan is to construct a distribution p ∗ that satisfies the conditions of the above Lemma.

Original: Uniqueness.
Revised: Uniqueness.

Original: Note further that for a fixed p ∗ , E ( x,y ) ∼ p ∗ yf ( ν, x ) is an upper bound for the margin min ( x,y )
Revised: Note further that for a fixed p ∗ , E ( x,y ) ∼ p ∗ yf ( ν, x ) is an upper bound for the margin min ( x,y )

Original: ∼D yf ( ν, x ) of any classifier ν .
Revised: ∼D yf ( ν, x ) of any classifier ν .

Original: Hence, for uniqueness, it suffices to show that δ θ 1 , δ θ 2 are the unique maximizers of the objective on the RHS of eqn.
Revised: Hence, for uniqueness, it suffices to show that δ θ 1 , δ θ 2 are the unique maximizers of the objective on the RHS of eqn.

Original: (4) and that the unique maximum margin convex combination of δ θ 1 , δ θ 2 over D is ν ∗ .
Revised: (4) and that the unique maximum margin convex combination of δ θ 1 , δ θ 2 over D is ν ∗ .

Original: We first describe the support D of p ∗ .
Revised: We first describe the support D of p ∗ .

Original: For y ∈ {± 1 } we generate ( x, y ) ∈ D as
Revised: For y ∈ {± 1 } we generate ( x, y ) ∈ D as

Original: [Equation]
Revised: [Equation]

Original: Now for ( x, y ) ∈ D , define
Revised: Now for ( x, y ) ∈ D , define

Original: [Equation]
Revised: [Equation]

Original: Note that p ∗ is supported on 2 d − 1 positive instances and one negative instance.
Revised: Note that p ∗ is supported on 2 d − 1 positive instances and one negative instance.

Original: We begin by showing eqn.
Revised: We begin by showing eqn.

Original: Claim A.3.
Revised: Claim A.3.

Original: p ∗ as in eqn.
Revised: p ∗ as in eqn.

Original: (6) satisfies eqn.
Revised: (6) satisfies eqn.

Original: Further, the unique maximum margin convex combination of δ θ 1 , δ θ 2 is ν ∗ .
Revised: Further, the unique maximum margin convex combination of δ θ 1 , δ θ 2 is ν ∗ .

Original: Proof.
Revised: Proof.

Original: Let us find the minimizers ( x, y ) ∼ D of yf ( ν, x ) =
Revised: Let us find the minimizers ( x, y ) ∼ D of yf ( ν, x ) =

Original: y E ( w,b,a ) ∼ ν ∗ [ a ( ϕ ( ⟨ w, x ⟩ + b ))] for any ν = λδ θ 1 + (1 − λ ) δ θ 2 , 0 ≤ λ ≤ 1 .
Revised: y E ( w,b,a ) ∼ ν ∗ [ a ( ϕ ( ⟨ w, x ⟩ + b ))] for any ν = λδ θ 1 + (1 − λ ) δ θ 2 , 0 ≤ λ ≤ 1 .

Original: yf ( ν, x ) for ( x, y ) with y = − 1 (denoting x 1 by − α 1 , where α 1 ≥ γ ) is
Revised: yf ( ν, x ) for ( x, y ) with y = − 1 (denoting x 1 by − α 1 , where α 1 ≥ γ ) is

Original: [Equation]
Revised: [Equation]

Original: [Equation][Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: and for ( x, y ) with y = 1 (denoting x 1 by α 2 , where α 2 ≥ γ ) is
Revised: and for ( x, y ) with y = 1 (denoting x 1 by α 2 , where α 2 ≥ γ ) is

Original: [Equation]
Revised: [Equation]

Original: As γ ≥ 1 , the expressions above equal λ √ γα 1 +1and (1 − λ ) √ γα 

Revised: As γ ≥ 1 , the expressions above equal λ √ γα 1 +1and (1 − λ ) √ γα

Original: +1respectively, and hence are minimized at α
Revised: +1respectively, and hence are minimized at α

Original: 1 = α 2 = γ .
Revised: 1 = α 2 = γ .

Original: Hence, the margin of ν is min( λ, 1 − λ ) √ 1+ γ which is uniquely maximized at λ = 1 / 2 .
Revised: Hence, the margin of ν is min( λ, 1 − λ ) √ 1+ γ which is uniquely maximized at λ = 1 / 2 .

Original: Further for λ = 1 / 2 , all points in D have the same value of yf ( ν, x ) .
Revised: Further for λ = 1 / 2 , all points in D have the same value of yf ( ν, x ) .

Original: In the rest of the proof we show eqn.
Revised: In the rest of the proof we show eqn.

Original: (4), Let us denote by g ( w, b, a ) :
Revised: (4), Let us denote by g ( w, b, a ) :

Original: = E ( x,y ) ∼ p ∗
Revised: = E ( x,y ) ∼ p ∗

Original: y
Revised: y

Original: [ a ( ϕ ( ⟨ w, x ⟩ + b ))] .
Revised: [ a ( ϕ ( ⟨ w, x ⟩ + b ))] .

Original: We show that δ θ 1 , δ θ 2 are the only maximizers of g ( w, b, a ) over S d +1 .
Revised: We show that δ θ 1 , δ θ 2 are the only maximizers of g ( w, b, a ) over S d +1 .

Original: We first find g ( θ 1 ) , g ( θ 2 ) .
Revised: We first find g ( θ 1 ) , g ( θ 2 ) .

Original: [Equation]
Revised: [Equation]

Original: where the first term is because w 2 , w 3 , . .
Revised: where the first term is because w 2 , w 3 , . .

Original: . , w d are zero for θ 1 .
Revised: . , w d are zero for θ 1 .

Original: Similarly, g ( θ 2 ) = √ γ 2 +  .
Revised: Similarly, g ( θ 2 ) = √ γ 2 +  .

Original: We now show that g ( w, a, b ) < √ γ 2 +  for ( w, a, b ) / ∈ { θ 1 , θ 2 } .
Revised: We now show that g ( w, a, b ) < √ γ 2 +  for ( w, a, b ) / ∈ { θ 1 , θ 2 } .

Original: We begin by showing the following simple but useful claim.
Revised: We begin by showing the following simple but useful claim.

Original: Claim A.4.
Revised: Claim A.4.

Original: All maximizers of g ( w, b, a ) over S d +1 satisfy | a | = 1 / √ 2 .
Revised: All maximizers of g ( w, b, a ) over S d +1 satisfy | a | = 1 / √ 2 .

Original: Proof.
Revised: Proof.

Original: The proof essentially follows from the 1 − homogeneity of the ReLU function ϕ and separability of g ( w, b, a ) .
Revised: The proof essentially follows from the 1 − homogeneity of the ReLU function ϕ and separability of g ( w, b, a ) .

Original: Note that g ( w, b, a ) =
Revised: Note that g ( w, b, a ) =

Original: (cid:112) ∥ w ∥ 2 + b 2 a · g ( w ′ , b ′ , 1) where ∥ w ′ ∥
Revised: (cid:112) ∥ w ∥ 2 + b 2 a · g ( w ′ , b ′ , 1) where ∥ w ′ ∥

Original: + b 2 = 1 .
Revised: + b 2 = 1 .

Original: Maximizing g ( w, b, a ) is equivalent to maximizing g ( w ′ , b ′ , 1) over S d and a (cid:112) ∥
Revised: Maximizing g ( w, b, a ) is equivalent to maximizing g ( w ′ , b ′ , 1) over S d and a (cid:112) ∥

Original: w ∥ 2 + b 2 over S d + respectively.
Revised: w ∥ 2 + b 2 over S d + respectively.

Original: The second of these has its unique maximum at | a | = 1 / √ 2 , completing the proof.
Revised: The second of these has its unique maximum at | a | = 1 / √ 2 , completing the proof.

Original: Now express g ( w, b, a ) as
Revised: Now express g ( w, b, a ) as

Original: [Equation]
Revised: [Equation]

Original: where σ i are independent Rademacher random variables.
Revised: where σ i are independent Rademacher random variables.

Original: We have two cases on a :
Revised: We have two cases on a :

Original: Case 1: a = 1 / √ 2 .
Revised: Case 1: a = 1 / √ 2 .

Original: By eqn.
Revised: By eqn.

Original: (7) we have
Revised: (7) we have

Original: [Equation]
Revised: [Equation]

Original: To simplify the above, define the random variable X
Revised: To simplify the above, define the random variable X

Original: = (cid:80) d i =2 σ i w i and denote γw 1 + b by α .
Revised: = (cid:80) d i =2 σ i w i and denote γw 1 + b by α .

Original: Note that | α | = | γw 1 + b | ≤ (cid:113) γ 2 + which follows from ∥ w ∥ 2 + b 2 = 1 / 2 .
Revised: Note that | α | = | γw 1 + b | ≤ (cid:113) γ 2 + which follows from ∥ w ∥ 2 + b 2 = 1 / 2 .

Original: The expectation in the last expression above becomes
Revised: The expectation in the last expression above becomes

Original: [Equation]
Revised: [Equation]

Original: where the last equality follows from symmetry of X .
Revised: where the last equality follows from symmetry of X .

Original: Note that Var ( X ) =
Revised: Note that Var ( X ) =

Original: (cid:80) di
Revised: (cid:80) di

Original: =2 w 2 i which is at most 12 − α 2 1+ γ 2 (using γw 1 + b = α and ∥ w ∥ 2 + b 2 = 1 / 2 ).
Revised: =2 w 2 i which is at most 12 − α 2 1+ γ 2 (using γw 1 + b = α and ∥ w ∥ 2 + b 2 = 1 / 2 ).

Original: Using A.5 to upper bound E [ X 1 { X ≥ α } ] we have
Revised: Using A.5 to upper bound E [ X 1 { X ≥ α } ] we have

Original: [Equation]
Revised: [Equation]

Original: We can check that the RHS of the above has its unique maximizer at α = (cid:113) 1+ γ 2for | α | ≤ (cid:113) 1+ γ 2 2 .
Revised: We can check that the RHS of the above has its unique maximizer at α = (cid:113) 1+ γ 2for | α | ≤ (cid:113) 1+ γ 2 2 .

Original: Hence g ( w, b, a ) ≤ √ 1+ γ in this case.
Revised: Hence g ( w, b, a ) ≤ √ 1+ γ in this case.

Original: We are now done since any ( w 1 , b ) satisfying γw 1 + b = (cid:113) 1+ γ 2and w 21 + b 2 ≤ 1 / 2 has b = 1 √ 2(1+ γ 2 ) .
Revised: We are now done since any ( w 1 , b ) satisfying γw 1 + b = (cid:113) 1+ γ 2and w 21 + b 2 ≤ 1 / 2 has b = 1 √ 2(1+ γ 2 ) .

Original: Case 2: a = − 1 / √ 2 .
Revised: Case 2: a = − 1 / √ 2 .

Original: Using eqn.
Revised: Using eqn.

Original: (7) we have g ( w, b, − 1 / √ 2) ≤ ϕ ( b − γw 1 ) / 2 √ 2 which for b 2 + w 21 ≤ 1 / 2 attains its unique maximum (cid:113) γ 2 + at b = 1 √ 2(1+ γ 2 ) .
Revised: (7) we have g ( w, b, − 1 / √ 2) ≤ ϕ ( b − γw 1 ) / 2 √ 2 which for b 2 + w 21 ≤ 1 / 2 attains its unique maximum (cid:113) γ 2 + at b = 1 √ 2(1+ γ 2 ) .

Original: Finally, note that the weights of the trained network ( w, b, a ) are sampled from ν ∗ .
Revised: Finally, note that the weights of the trained network ( w, b, a ) are sampled from ν ∗ .

Original: Hence, the final claim in the theorem about f ( ν ∗ , Px 1 + P ⊥ x 2 ) follows since the distribution of w only has a support on e 1 and − e 1 .
Revised: Hence, the final claim in the theorem about f ( ν ∗ , Px 1 + P ⊥ x 2 ) follows since the distribution of w only has a support on e 1 and − e 1 .

Original: A.1. A UXILIARY LEMMAS FOR RICH REGIME
Revised: A.1. A UXILIARY LEMMAS FOR RICH REGIME

Original: Lemma A.5.
Revised: Lemma A.5.

Original: For any symmetric discrete random variable X with bounded variance, for α > 0 ,
Revised: For any symmetric discrete random variable X with bounded variance, for α > 0 ,

Original: E
Revised: E

Original: [Equation]
Revised: [Equation]

Original: ≤
Revised: ≤

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: [Equation]
Revised: [Equation]

Original: where the last inequality is by Cauchy-Schwartz.
Revised: where the last inequality is by Cauchy-Schwartz.

Original: Also by Chebyshev’s inequality, p ( | X | ≥ α ) ≤
Revised: Also by Chebyshev’s inequality, p ( | X | ≥ α ) ≤

Original: V ar ( X ) / 2 α 2 .
Revised: V ar ( X ) / 2 α 2 .

Original: Combining this with eqn.
Revised: Combining this with eqn.

Original: (8) and using symmetry of X and non-negativity of α gives the required lemma.
Revised: (8) and using symmetry of X and non-negativity of α gives the required lemma.

Original: A. L AZY REGIME
Revised: A. L AZY REGIME

Original: Theorem 4.2 is a corollary of the following more general theorem.
Revised: Theorem 4.2 is a corollary of the following more general theorem.

Original: Theorem A.6.
Revised: Theorem A.6.

Original: Consider a point x ∈ D .
Revised: Consider a point x ∈ D .

Original: For sufficiently small ϵ > 0 , there exist an absolute constant N such that for all d > N, γ < ϵ √ d and γ ≥ 7 , for the joint training of both the layers of 1-hidden layer FCN in the NTK regime, the prediction of any point of the form ( ζ, x 2: d ) satisfies the following:
Revised: For sufficiently small ϵ > 0 , there exist an absolute constant N such that for all d > N, γ < ϵ √ d and γ ≥ 7 , for the joint training of both the layers of 1-hidden layer FCN in the NTK regime, the prediction of any point of the form ( ζ, x 2: d ) satisfies the following:

Original: For ζ ≥ 0 .
Revised: For ζ ≥ 0 .

Original: 73 , the prediction is positive.
Revised: 73 , the prediction is positive.

Original: For ζ ≤ − 0 .
Revised: For ζ ≤ − 0 .

Original: 95 γ , the prediction is negative.
Revised: 95 γ , the prediction is negative.

Original: The above theorem establishes that perturbing x 1 by O ( γ ) changes pred ( f ( x )) for x ∈ D (whereas a classifier exists that achieves a margin of Ω( √ d ) on D , as D has margin 1 for coordinates { 2 · · · d } ).
Revised: The above theorem establishes that perturbing x 1 by O ( γ ) changes pred ( f ( x )) for x ∈ D (whereas a classifier exists that achieves a margin of Ω( √ d ) on D , as D has margin 1 for coordinates { 2 · · · d } ).

Original: As γ = o ( d ) , this shows that the learned model is adversarially vulnerable.
Revised: As γ = o ( d ) , this shows that the learned model is adversarially vulnerable.

Original: Proof of Theorem A.6.
Revised: Proof of Theorem A.6.

Original: The idea of the proof is to obtain an explicit expression for f ( x ) by applying standard kernel max-margin SVM theory to the NTK kernel 3.2.
Revised: The idea of the proof is to obtain an explicit expression for f ( x ) by applying standard kernel max-margin SVM theory to the NTK kernel 3.2.

Original: We begin with some preliminaries.
Revised: We begin with some preliminaries.

Original: We will refer to the first coordinate of the instance as the ’linear’ coordinate, and to the rest as ’non-linear’ coordinates.
Revised: We will refer to the first coordinate of the instance as the ’linear’ coordinate, and to the rest as ’non-linear’ coordinates.

Original: Also, henceforth we append an extra coordinate with value 1 to all our instances (corresponding to bias term) - as is standard for working with unbiased SVM without loss of generality.
Revised: Also, henceforth we append an extra coordinate with value 1 to all our instances (corresponding to bias term) - as is standard for working with unbiased SVM without loss of generality.

Original: Explicit expression for f .
Revised: Explicit expression for f .

Original: Using representer theorem for max margin kernel SVM, we know that f can be expressed as
Revised: Using representer theorem for max margin kernel SVM, we know that f can be expressed as

Original: [Equation]
Revised: [Equation]

Original: for some λ t ≥ 0 (that are known as Lagrange multipliers ).
Revised: for some λ t ≥ 0 (that are known as Lagrange multipliers ).

Original: Further by KKT conditions, a function possessing such a representation (that correctly classifies D ) has maximum margin if y ( t ) f ( x ( t ) )
Revised: Further by KKT conditions, a function possessing such a representation (that correctly classifies D ) has maximum margin if y ( t ) f ( x ( t ) )

Original: = 1 whenever λ t > 0 (training points t satisfying λ t > 0 are called support vectors ).
Revised: = 1 whenever λ t > 0 (training points t satisfying λ t > 0 are called support vectors ).

Original: We begin with a useful claim.
Revised: We begin with a useful claim.

Original: Claim A.7.
Revised: Claim A.7.

Original: The max margin kernel SVM for D with the NTK kernel has all points in D as support vectors.
Revised: The max margin kernel SVM for D with the NTK kernel has all points in D as support vectors.

Original: Proof.
Revised: Proof.

Original: By the above discussion, it suffices to show that the (unique) solution α ∈ R | D | to Kα = y satisfies sign ( α i ) =
Revised: By the above discussion, it suffices to show that the (unique) solution α ∈ R | D | to Kα = y satisfies sign ( α i ) =

Original: y ( i ) for all i , where K is the | D | × | D | Gram matrix with ( i, j ) th entry
Revised: y ( i ) for all i , where K is the | D | × | D | Gram matrix with ( i, j ) th entry

Original: K ( x ( i ) , x ( j ) ) and y i =
Revised: K ( x ( i ) , x ( j ) ) and y i =

Original: y ( i ) (the Lagrange multipliers λ
Revised: y ( i ) (the Lagrange multipliers λ

Original: i are then given by y i α i ).
Revised: i are then given by y i α i ).

Original: Structure of Gram matrix.
Revised: Structure of Gram matrix.

Original: Order D so that the positive instances appear first.
Revised: Order D so that the positive instances appear first.

Original: Then the Gram matrix K has a block structure of the form (cid:18) B
Revised: Then the Gram matrix K has a block structure of the form (cid:18) B

Original: C
Revised: C

Original: C T
Revised: C T

Original: R (cid:19) where B ∈ R 2 d − 1 × 2 d − 1 and R ∈ R are the
Revised: R (cid:19) where B ∈ R 2 d − 1 × 2 d − 1 and R ∈ R are the

Original: Gram matrices for the positive and negative instances respectively, and C ∈ R 2 d − 1 × 1 represents the
Revised: Gram matrices for the positive and negative instances respectively, and C ∈ R 2 d − 1 × 1 represents the

Original: K ( x ( i ) , x ( | D | ) ) values for i < | D | .
Revised: K ( x ( i ) , x ( | D | ) ) values for i < | D | .

Original: Recall that for the NTK kernel, K ( x ( i ) , x ( j ) ) has the form ∥ x ( i ) ∥∥ x ( j ) ∥ κ ( ⟨ x ( i ) , x ( j ) ⟩ ) .
Revised: Recall that for the NTK kernel, K ( x ( i ) , x ( j ) ) has the form ∥ x ( i ) ∥∥ x ( j ) ∥ κ ( ⟨ x ( i ) , x ( j ) ⟩ ) .

Original: Note all the positive instances have the same norm (denoted by ρ 1 = (cid:112) d + γ 2 ) and the inner product between two positive instances depends only on the number i of non-matching non-linear coordinates (denoted by β i for 0 ≤ i ≤ d − 1 ).
Revised: Note all the positive instances have the same norm (denoted by ρ 1 = (cid:112) d + γ 2 ) and the inner product between two positive instances depends only on the number i of non-matching non-linear coordinates (denoted by β i for 0 ≤ i ≤ d − 1 ).

Original: Hence, the rows of B are permutations of each other, with the entry ρ 21 β i appearing (cid:0)
Revised: Hence, the rows of B are permutations of each other, with the entry ρ 21 β i appearing (cid:0)

Original: d − 1 i (cid:1) times.
Revised: d − 1 i (cid:1) times.

Original: Similarly, the entries in C are all equal and are denoted by ρ 

Revised: Similarly, the entries in C are all equal and are denoted by ρ

Original: ρ 2 β d where β d denotes κ ( x ( t ) , x | D | ) for any t < | D | and ρ 2 = ∥ x | D | ∥ =
Revised: ρ 2 β d where β d denotes κ ( x ( t ) , x | D | ) for any t < | D | and ρ 2 = ∥ x | D | ∥ =

Original: (cid:112) 1 + γ 2 .
Revised: (cid:112) 1 + γ 2 .

Original: The only entry in R is ρ 22 κ (1) .
Revised: The only entry in R is ρ 22 κ (1) .

Original: In particular, β i = κ (cid:18)
Revised: In particular, β i = κ (cid:18)

Original: [Equation]
Revised: [Equation]

Original: γ
Revised: γ

Original: [Equation]
Revised: [Equation]

Original: d + γ(cid:19) for i ∈ [ | D | − 1] , and β d = κ
Revised: d + γ(cid:19) for i ∈ [ | D | − 1] , and β d = κ

Original: [Equation]
Revised: [Equation]

Original: Now we are ready to solve Kα =
Revised: Now we are ready to solve Kα =

Original: y .
Revised: y .

Original: By symmetry in the structure of K, α looks like [ a, a, ......, b ] , where the first | D | − 1 entries are the same.
Revised: By symmetry in the structure of K, α looks like [ a, a, ......, b ] , where the first | D | − 1 entries are the same.

Original: Expanding Kα =
Revised: Expanding Kα =

Original: y , we get two equations given by
Revised: y , we get two equations given by

Original: [Equation]
Revised: [Equation]

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: Solving, we get
Revised: Solving, we get

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: We now show that a > 0 and b < 0 .
Revised: We now show that a > 0 and b < 0 .

Original: Note that for sufficiently large d , β d can be made arbitrarily close to κ (0)
Revised: Note that for sufficiently large d , β d can be made arbitrarily close to κ (0)

Original: =
Revised: =

Original: 1 /π (since κ is smooth around 0 ).
Revised: 1 /π (since κ is smooth around 0 ).

Original: Hence, a > 0 implies b < 0 .
Revised: Hence, a > 0 implies b < 0 .

Original: We in fact give the following estimate for a :
Revised: We in fact give the following estimate for a :

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: For the lower bound on ξ , write
Revised: For the lower bound on ξ , write

Original: [Equation]
Revised: [Equation]

Original: where for the first inequality we used convexity of κ and for the second inequality we used β d/ 

Revised: where for the first inequality we used convexity of κ and for the second inequality we used β d/

Original: = κ (0) + O (1 /d ) , β d = κ (0)
Revised: = κ (0) + O (1 /d ) , β d = κ (0)

Original: + O (1 / √ d ) .
Revised: + O (1 / √ d ) .

Original: For the upper bound on ξ , write
Revised: For the upper bound on ξ , write

Original: [Equation]
Revised: [Equation]

Original: where for the second inequality we used κ ( u ) ≤ 1 + u (which holds by convexity and κ ( − 1) = 0 , κ (1) = 2 ).
Revised: where for the second inequality we used κ ( u ) ≤ 1 + u (which holds by convexity and κ ( − 1) = 0 , κ (1) = 2 ).

Original: Now we analyze predicted labels for points of the form ( ζ, x 2: d +1 ) where x ∈ D .
Revised: Now we analyze predicted labels for points of the form ( ζ, x 2: d +1 ) where x ∈ D .

Original: We make two cases depending on the label of x .
Revised: We make two cases depending on the label of x .

Original: Predicted label for point ( ζ, x ( t ) 2: d +1 ) where x ( t ) ∈ D has positive label
Revised: Predicted label for point ( ζ, x ( t ) 2: d +1 ) where x ( t ) ∈ D has positive label

Original: Our point (denoted by x ) has the form ( ζ, ζ 1 , ζ 2 , . .
Revised: Our point (denoted by x ) has the form ( ζ, ζ 1 , ζ 2 , . .

Original: , ζ d , 1) where ζ i ∈ ± 1 .
Revised: , ζ d , 1) where ζ i ∈ ± 1 .

Original: 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 .
Revised: 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 .

Original: Explicit form for f .
Revised: Explicit form for f .

Original: Let τ i def = ⟨ x, x ′ ⟩ / ( ∥ x ∥∥ x ′ ∥ ) for a positive instance x ′ ∈ D , where x and x ′ have exactly i non-matching non-linear coordinates (for 0 ≤ i ≤ d − 1 ).
Revised: Let τ i def = ⟨ x, x ′ ⟩ / ( ∥ x ∥∥ x ′ ∥ ) for a positive instance x ′ ∈ D , where x and x ′ have exactly i non-matching non-linear coordinates (for 0 ≤ i ≤ d − 1 ).

Original: Similarly denote by τ d the quantity ⟨ x, x | D | ⟩ / ( ∥ x
Revised: Similarly denote by τ d the quantity ⟨ x, x | D | ⟩ / ( ∥ x

Original: ∥∥ x | D | ∥ ) .
Revised: ∥∥ x | D | ∥ ) .

Original: In particular, τ i = (cid:18) d − 2 i + γζ ρ 1 ∥ x ∥ (cid:19) and τ d = (cid:18) 1 − γζ ρ 2 ∥ x ∥ (cid:19) .
Revised: In particular, τ i = (cid:18) d − 2 i + γζ ρ 1 ∥ x ∥ (cid:19) and τ d = (cid:18) 1 − γζ ρ 2 ∥ x ∥ (cid:19) .

Original: By the above discussion, we have f ( x ) =
Revised: By the above discussion, we have f ( x ) =

Original: a   | D |− 1 (cid:88) t =
Revised: a   | D |− 1 (cid:88) t =

Original: K ( x, x ( t ) ) 
Revised: K ( x, x ( t ) ) 

Original:  + bK ( x, x | D | ) = aρ 1 ∥ x ∥ (cid:32) d − 1 (cid:88) i = (cid:18)
Revised:  + bK ( x, x | D | ) = aρ 1 ∥ x ∥ (cid:32) d − 1 (cid:88) i = (cid:18)

Original: d −i (cid:19) κ ( τ i ) (cid:33) + bρ 2 ∥ x
Revised: d −i (cid:19) κ ( τ i ) (cid:33) + bρ 2 ∥ x

Original: ∥ κ ( τ d ) .
Revised: ∥ κ ( τ d ) .

Original: Substituting b and denoting f ( x ) / ∥ x ∥ by g ( ζ ) we get g ( ζ ) = aρ(cid:34)
Revised: Substituting b and denoting f ( x ) / ∥ x ∥ by g ( ζ ) we get g ( ζ ) = aρ(cid:34)

Original: d − 1 (cid:88) i = (cid:18)
Revised: d − 1 (cid:88) i = (cid:18)

Original: d −i (cid:19) κ ( τ i ( ζ )) − 2 d − 1 β d κ (1) κ ( τ d
Revised: d −i (cid:19) κ ( τ i ( ζ )) − 2 d − 1 β d κ (1) κ ( τ d

Original: [Equation]
Revised: [Equation]

Original: Now try to expand g ( ζ ) using the Taylor series around ζ = γ (note that g ( γ ) = 1 /ρ 1 ).
Revised: Now try to expand g ( ζ ) using the Taylor series around ζ = γ (note that g ( γ ) = 1 /ρ 1 ).

Original: Note that κ ′ can however be unbounded around − 1 and 1 .
Revised: Note that κ ′ can however be unbounded around − 1 and 1 .

Original: To get around this, write g = h + q , where h has bounded first and second derivative, and q has lower order than h for ζ of interest.
Revised: To get around this, write g = h + q , where h has bounded first and second derivative, and q has lower order than h for ζ of interest.

Original: In particular,
Revised: In particular,

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: Observe that q ( ζ ) =
Revised: Observe that q ( ζ ) =

Original: o ( c d ) for c < 1 using the estimate eqn.
Revised: o ( c d ) for c < 1 using the estimate eqn.

Original: (9) for a and concentration for sums of independent Bernoullis.
Revised: (9) for a and concentration for sums of independent Bernoullis.

Original: By Taylor’s theorem,
Revised: By Taylor’s theorem,

Original: [Equation]
Revised: [Equation]

Original: for some θ ∈ [ γ, ζ ] , where h ( γ ) ≈ 1 / √ d .
Revised: for some θ ∈ [ γ, ζ ] , where h ( γ ) ≈ 1 / √ d .

Original: It will turn out that | h ′ ( γ ) | = Θ(1 / √ d ) , | h ′′ ( ζ ) |
Revised: It will turn out that | h ′ ( γ ) | = Θ(1 / √ d ) , | h ′′ ( ζ ) |

Original: = o (1 / √ d ) .
Revised: = o (1 / √ d ) .

Original: This will allow us to complete the proof using the linear approximation of g ( ζ ) by neglecting the second order term and q ( ζ ) .
Revised: This will allow us to complete the proof using the linear approximation of g ( ζ ) by neglecting the second order term and q ( ζ ) .

Original: We now compute h ′ , h ′′ , treating ∥ x ∥ =
Revised: We now compute h ′ , h ′′ , treating ∥ x ∥ =

Original: (cid:112) d + ζ 2 as a constant for exposition (the proof works without this approximation or the reader may think of γ as o ( √ d ) ).
Revised: (cid:112) d + ζ 2 as a constant for exposition (the proof works without this approximation or the reader may think of γ as o ( √ d ) ).

Original: Using τ ′
Revised: Using τ ′

Original: i ( ζ ) ≈
Revised: i ( ζ ) ≈

Original: [Equation]
Revised: [Equation]

Original: Plugging ∥ x
Revised: Plugging ∥ x

Original: ∥ ≈ ρ 1 ≈ √ d and substituting a from eqn.
Revised: ∥ ≈ ρ 1 ≈ √ d and substituting a from eqn.

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: which substituted in eqn
Revised: which substituted in eqn

Original: (11) with τ d ( ζ ) ≈ 0 , β d
Revised: (11) with τ d ( ζ ) ≈ 0 , β d

Original: ≈ κ (0) , κ ′ ( τ d ( ζ ))
Revised: ≈ κ (0) , κ ′ ( τ d ( ζ ))

Original: ≈ κ ′ (0) gives
Revised: ≈ κ ′ (0) gives

Original: [Equation]
Revised: [Equation]

Original: Hence, g ( ζ ) > 0 whenever the coefficient of 1 / √ d above is bounded above zero, and a similar condition holds for g ( ζ ) < 0 .
Revised: Hence, g ( ζ ) > 0 whenever the coefficient of 1 / √ d above is bounded above zero, and a similar condition holds for g ( ζ ) < 0 .

Original: Using the estimates of ξ from eqn.
Revised: Using the estimates of ξ from eqn.

Original: (9) and κ ′ (0) = 1 , κ (0)
Revised: (9) and κ ′ (0) = 1 , κ (0)

Original: = 

Revised: =

Original: /π, κ (1) = 2 , ρ 22 = 1 + γ 2 in the above gives that g ( ζ ) > 0 for ζ > − 0 .
Revised: /π, κ (1) = 2 , ρ 22 = 1 + γ 2 in the above gives that g ( ζ ) > 0 for ζ > − 0 .

Original: 68 γ − 1 .
Revised: 68 γ − 1 .

Original: 68 /γ
Revised: 68 /γ

Original: and g ( ζ ) < 0 for ζ < − 0 .
Revised: and g ( ζ ) < 0 for ζ < − 0 .

Original: 905 γ − 1 .
Revised: 905 γ − 1 .

Original: 905 /γ .
Revised: 905 /γ .

Original: Predicted label for point ( ζ, x ( t ) 2: d +1 )
Revised: Predicted label for point ( ζ, x ( t ) 2: d +1 )

Original: where x ( t ) ∈
Revised: where x ( t ) ∈

Original: D has negative label
Revised: D has negative label

Original: Following the same plan, write our point (denoted by x ) as ( ζ, 0 , . .
Revised: Following the same plan, write our point (denoted by x ) as ( ζ, 0 , . .

Original: Explicit form for f .
Revised: Explicit form for f .

Original: Begin by finding τ
Revised: Begin by finding τ

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

Original: (cid:19) and τ d
Revised: (cid:19) and τ d

Original: [Equation]
Revised: [Equation]

Original: eqn.
Revised: eqn.

Original: (10) now gives
Revised: (10) now gives

Original: [Equation]
Revised: [Equation]

Original: Expanding κ ( τ 0 ( ζ )) using Taylor series around ζ = − 1 /γ ,
Revised: Expanding κ ( τ 0 ( ζ )) using Taylor series around ζ = − 1 /γ ,

Original: [Equation]
Revised: [Equation]

Original: for some θ ∈ [ − 1 , 1] .
Revised: for some θ ∈ [ − 1 , 1] .

Original: For large d , τ 0 ( θ ) ≈ 0 and τ ′ 

Revised: For large d , τ 0 ( θ ) ≈ 0 and τ ′

Original: ( θ ) =
Revised: ( θ ) =

Original: O (1 / √ d ) .
Revised: O (1 / √ d ) .

Original: Hence we have
Revised: Hence we have

Original: [Equation]
Revised: [Equation]

Original: As before g ( ζ ) > 0 whenever the coefficient of 1 /ρ 2 above is bounded above zero which happens for ζ ≥ 0 .
Revised: As before g ( ζ ) > 0 whenever the coefficient of 1 /ρ 2 above is bounded above zero which happens for ζ ≥ 0 .

Original: 73 (for γ ≥ 3 ).
Revised: 73 (for γ ≥ 3 ).

Original: Similarly, g ( ζ ) < 0 for ζ ≤ 0 .
Revised: Similarly, g ( ζ ) < 0 for ζ ≤ 0 .

Original: B E XPERIMENTS
Revised: B E XPERIMENTS

Original: In this section, we provide experimental details, including hyperparameter tuning setup and some additional experiments.
Revised: In this section, we provide experimental details, including hyperparameter tuning setup and some additional experiments.

Original: B. D ETAILS ON THE EXPERIMENTAL SETTING
Revised: B. D ETAILS ON THE EXPERIMENTAL SETTING

Original: We will first describe the four datasets that have been used in this work.
Revised: We will first describe the four datasets that have been used in this work.

Original: Imagenette (FastAI, 2021):
Revised: Imagenette (FastAI, 2021):

Original: This is a subset of 10 classes of Imagenet, that are comparatively easier to classify.
Revised: This is a subset of 10 classes of Imagenet, that are comparatively easier to classify.

Original: b-Imagenette : This is a binarized version of Imagenette, where only a subset of two classes (tench and English springer) is used.
Revised: b-Imagenette : This is a binarized version of Imagenette, where only a subset of two classes (tench and English springer) is used.

Original: Waterbirds-Landbirds (Sagawa et al., 2020a):
Revised: Waterbirds-Landbirds (Sagawa et al., 2020a):

Original: This is a majority-minority group dataset, consisting of waterbirds on water and land background, as well as landbirds on land and water background.
Revised: This is a majority-minority group dataset, consisting of waterbirds on water and land background, as well as landbirds on land and water background.

Original: This dataset serves as a baseline for checking the dependence of model on the spurious background feature when predicting the bird class, as most of the training examples have waterbirds on water and landbirds on land background.
Revised: This dataset serves as a baseline for checking the dependence of model on the spurious background feature when predicting the bird class, as most of the training examples have waterbirds on water and landbirds on land background.

Original: MNIST-CIFAR (Shah et al., 2020):
Revised: MNIST-CIFAR (Shah et al., 2020):

Original: This is a collage dataset, created by concatenating MNIST and CIFAR images along an axis.
Revised: This is a collage dataset, created by concatenating MNIST and CIFAR images along an axis.

Original: This is a synthetic dataset for evaluating the simplicity bias of a trained model.
Revised: This is a synthetic dataset for evaluating the simplicity bias of a trained model.

Original: Setup Throughout the paper, we work with the pretrained representations of the above datasets, obtained by using an Imagenet pretrained Resnet 50.
Revised: Setup Throughout the paper, we work with the pretrained representations of the above datasets, obtained by using an Imagenet pretrained Resnet 50.

Original: We finetune a 1-hidden layer FCN (hidden dimension - 100 ) on top of these representations (keeping the backbone fixed) using SGD with a momentum of 0.9.
Revised: We finetune a 1-hidden layer FCN (hidden dimension - 100 ) on top of these representations (keeping the backbone fixed) using SGD with a momentum of 0.9.

Original: Every model is trained for 20000 steps (large enough for convergence) with a warmup and cosine decay learning rate scheduler.
Revised: Every model is trained for 20000 steps (large enough for convergence) with a warmup and cosine decay learning rate scheduler.

Original: For each of the runs, we tune the batch size, learning rate and weight decay using validation accuracy.
Revised: For each of the runs, we tune the batch size, learning rate and weight decay using validation accuracy.

Original: Below are the hyperparameter tuning details:
Revised: Below are the hyperparameter tuning details:

Original: • Batch size ∈ { 128 , 256 } • Learning rate: – Rich regime: ∈ { 0 . 5 , 1 . 0 } (as learning rate in rich regime needs to scale up with the hidden dimension) –
Revised: • Batch size ∈ { 128 , 256 } • Learning rate: – Rich regime: ∈ { 0 . 5 , 1 . 0 } (as learning rate in rich regime needs to scale up with the hidden dimension) –

Original: Lazy regime: ∈ { 0 . 01 , 0 . 05 } • Weight decay: ∈ { 0 ,
Revised: Lazy regime: ∈ { 0 . 01 , 0 . 05 } • Weight decay: ∈ { 0 ,

Original: [Equation] 
Revised: [Equation] −

Original: The final numbers reported are averaged across 3 independent runs with the selected hyperparameters.
Revised: The final numbers reported are averaged across 3 independent runs with the selected hyperparameters.

Original: Evaluation For Imagenette, b-Imagenette and MNIST-CIFAR, we report the standard test accuracy in all the experiments.
Revised: Evaluation For Imagenette, b-Imagenette and MNIST-CIFAR, we report the standard test accuracy in all the experiments.

Original: For waterbirds, we report train-adjusted test accuracy, as reported in
Revised: For waterbirds, we report train-adjusted test accuracy, as reported in

Original: Sagawa et al.
Revised: Sagawa et al.

Original: (2020a).
Revised: (2020a).

Original: 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.
Revised: 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.

Original: B. A DDITIONAL EXPERIMENTAL RESULTS
Revised: B. A DDITIONAL EXPERIMENTAL RESULTS

Original: In this section, we present a few additional experimental results.
Revised: In this section, we present a few additional experimental results.

Revised: Accuracy of f proj

Revised: In Table 4, we show the test accuracy of f proj .

Revised: As can be seen, even after projecting out the principal components used by f , f proj attains significantly high accuracy.

Revised: Note that, in these experiments, model 1 was kept fixed and the accuracy of f proj is averaged across 3 runs.

Original: Singular value decay .
Revised: Singular value decay .

Original: In Figure 6, we provide the singular value decay of the weight matrix for the first model trained in rich regime.
Revised: In Figure 6, we provide the singular value decay of the weight matrix for the first model trained in rich regime.

Original: As can be seen, the top few singular values capture most of the Frobenius norm of the matrix. 
Revised: As can be seen, the top few singular values capture most of the Frobenius norm of the matrix. [Table]

Original: MNIST-CIFAR
Revised: MNIST-CIFAR

Original: In Figure 7, we show that an ensemble of f and f proj has better gaussian robustness than an ensemble of f and f ind on MNIST-CIFAR dataset.
Revised: In Figure 7, we show that an ensemble of f and f proj has better gaussian robustness than an ensemble of f and f ind on MNIST-CIFAR dataset.

Revised: Quantitative measurement of non-linearity of decision boundary In this section, we report a quantitative measure of non-linearity of the decision boundary along the top two singular vectors for f and f proj .

Revised: Basically, we fit a linear classifier to the decision boundary and report its accuracy.

Revised: As shown in Table 5, the test accuracy obtained by the linear classifier for f proj is less than f .

Revised: Variation of LD-SB with depth In Figure 8 and 9, we show the evolution of effective rank of weight matrices for depth-2 and 3 ReLU networks.

Revised: As can be seen, the rank still decreases with training, however the effect is less pronounced for the initial layers.

Original:  [Figure]
Revised: Note that the initialization used in these runs was the feature learning initialization as proposed in Yang & Hu (2021). [Figure] [Figure] [Figure]

Revised: v e r a n k

Revised: MNIST-CIFARImagenetteWaterbirdsb-Imagenette

Original: [Figure]
Revised: [Figure]

Revised: C E XTENDED R ELATED W ORKS

Revised: In this section, we provide an extensive literature survey of various topics that the paper is based on.

Revised: Low rank Simplicity Bias in Linear Networks Multiple works have established low rank simplicity bias for gradient descent on linear networks, both for squared loss as well as cross-entropy loss.

Revised: For squared loss, Gunasekar et al.

Revised: (2017) conjectured that the network is biased towards finding minimum nuclear norm solutions for two-layer linear networks.

Revised: Arora et al.

Revised: (2019) refuted the conjecture and instead argued that the network is biased towards finding low rank solutions.

Revised: Razin &

Revised: Cohen (2020) provided empirical support to 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.

Revised: Li et al.

Revised: (2021) established that for small enough initialization, gradient flow on linear networks follows greedy low-rank learning trajectory.

Revised: For binary classification on linearly separable data, Ji & Telgarsky (2019) showed that the weight matrices of a linear network eventually become rank-1 as training progresses.

Revised: Low rank Simplicity Bias in Non-Linear Networks For non-linear networks, the work related to low-rank simplicity bias is rather sparse.

Revised: Two of the most notable works are Huh et al.

Revised: (2021) and

Revised: Galanti & Poggio (2022).

Revised: Huh et al.

Revised: (2021) empirically established that the rank of the embeddings learnt by a neural network with ReLU activations goes down as training progresses.

Revised: Galanti &

Revised: 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.

Revised: In contrast to these works, forlayer 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.

Revised: 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 an OOD point.

Revised: Other works also implicitly rely on the low-rank nature of the embeddings.

Revised: Ndiour et al.

Revised: (2020) use PCA on the learnt features, and only model the likelihood along the small subspace spanned by the top few directions.

Revised: Wang et al.

Revised: (2022) 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 it with logit information to detect OOD datapoints.

Revised: 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.

Revised: Other Simplicity Bias There have been many works exploring the nature of simplicity bias in neural networks, both empirically and theoretically.

Revised: Kalimeris et al.

Revised: (2019) empirically demonstrated that SGD on neural networks gradually learns functions of increasing complexity.

Revised: Rahaman et al.

Revised: (2018) empirically demonstrated that neural networks tend to learn lower frequency functions first.

Revised: Ronen et al.

Revised: (2019) theoretically established that in NTK regime, the convergence rate depends on the eigenvalues of the kernel spectrum.

Revised: Hacohen et al.

Revised: (2020) showed that neural networks always learn train and test examples almost in the same order, irrespective of the architecture.

Revised: Pezeshki et al.

Revised: (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.

Revised: In contrast, our results are shown for any dataset in the IFM model in the rich regime of training.

Revised: Lyu et al.

Revised: (2021) consider antisymmetric datasets and show that single hidden layer input homogeneous networks (i.e., without bias parameters) converge to linear classifiers.

Revised: However, such networks have strictly weaker expressive power compared to those with bias parameters.

Revised: Hacohen & Weinshall (2022) showed that for deep linear networks, in NTK regime, they learn the higher principal components of the input data first.

Revised: Most of the previous works used simplicity bias as a reason behind better generalization of neural nets.

Revised: However, Shah et al.

Revised: (2020) showed that extreme simplicity bias could also lead to worse OOD performance.

Revised: Learning diverse classifiers : There have been several works that attempt to learn diverse classifiers.

Revised: 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 al., 2022).

Revised: Xu et al.

Revised: (2022) proposes a way to 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.

Revised: Lee et al.

Revised: (2022) learns diverse classifiers by enforcing diversity on unlabeled target data.

Revised: Spurious correlations : There has been a large body of work which identifies the reasons for spurious correlations in NNs (Sagawa et al., 2020b) as well as proposing algorithmic fixes in different settings (Liu et al., 2021; Chen et al., 2020).

Revised: Implicit bias of gradient descent : There is also a large body of work understanding the implicit bias of gradient descent dynamics.

Revised: Most of these works are for standard linear (Ji & Telgarsky, 2019) or deep linear networks (Soudry et al., 2018; Gunasekar et al., 2018).

Revised: 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.