SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Depen Morwani *
dmorwani@g.harvard.edu
Jatin Batra
Tata Institute of Fundamental Research (TIFR)
jatin.batra@tifr.res.in
Prateek Jain<sup>1</sup>
Praneeth Netrapalli
pnetrapalli@google.com
Recent works (Shah et al., 2020; Chen et al., 2021) have demonstrated that neural networks exhibit extreme simplicity bias (SB). That is, they learn only the simplest features to solve a task at hand, even in the presence of other, more robust but more complex features. Due to the lack of a general and rigorous definition of features, these works showcase SB on semi-synthetic datasets such as Color-MNIST, MNIST-CIFAR where defining features is relatively easier.
In this work, we rigorously define as well as thoroughly establish SB for one hidden layer neural networks. More concretely, (i) we define SB as the network essentially being a function of a low dimensional projection of the inputs (ii) theoretically, in the infinite width regime, we show that when the data is linearly separable, the network primarily depends on only the linearly separable (1-dimensional) subspace even in the presence of an arbitrarily large number of other, more complex features which could have led to a significantly more robust classifier, (iii) empirically, we show that models trained on real datasets such as Imagenet and Waterbirds-Landbirds indeed depend on a low dimensional projection of the inputs, thereby demonstrating SB on these datasets, iv) finally, we present a natural ensemble approach that encourages diversity in models by training successive models on features not used by earlier models, and demonstrate that it yields models that are significantly more robust to Gaussian noise.
It is well known that neural networks (NNs) are vulnerable to distribution shifts as well as to adversarial examples (Szegedy et al., 2014; Hendrycks et al., 2021). A recent line of work (Geirhos et al., 2018; Shah et al., 2020; Geirhos et al., 2020) proposes that Simplicity Bias (SB) (or 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 non-robustness of the trained networks. The argument is roughly as follows: for example, in the classification of swans vs bears, as illustrated in Figure 1, there are many features such as background, color of the animal, shape of the animal etc. that can be used for classification. However using only one or few of them can lead to models that are not robust to specific distribution shifts, while using all the features can lead to more robust models.
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
* part of the work done while at Google Research, India 1 Alphabetical Ordering
2Image source: Wikipedia swa, bea.
for NN's brittleness to distribution shifts (Shah et al., 2020). However, such observations are still only for specific semi-real datasets, and a general method that can identify SB on a given dataset and a given model is still missing in literature. Such a method would be useful not only to estimate the robustness of a model but could also help in designing more robust models.
A key challenge in designing such a general method to identify (and potentially fix) SB is that the notion of feature itself is vague and lacks a rigorous definition. Existing works Geirhos et al. (2018); Shah et al. (2020); Chen et al. (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.
Figure 1: Classification of swans vs bears. There are several features such as background, color of the animal, shape of the animal etc., each of which is sufficient for classification but using all of them will lead to a more robust model.
Contributions: Our first contribution is to provide a precise definition of a particular simplicity
bias - LD-SB- referring to low dimensional input dependence of the model.
Definition 1.1 (LD-SB). A model  $f: \mathbb{R}^d \to \mathbb{R}^c$  with inputs  $x \in \mathbb{R}^d$  and outputs  $f(x) \in \mathbb{R}^c$  (e.g., logits for  $c$  classes), trained on a distribution  $(x, y) \sim \mathcal{D}$  satisfies LD-SB if there exists a projection matrix  $P \in \mathbb{R}^{d \times d}$  satisfying:
-  $\operatorname{rank}(P) = k \ll d$ ,
-  $\mathbb{P}[\operatorname{pred}(f(Px^{(1)} + P_{\perp}x^{(2)})) = \operatorname{pred}(f(x^{(1)}))]\geq 1 - \epsilon_1$  for  $(x^{(1)},y^{(1)}),(x^{(2)},y^{(2)})\sim \mathcal{D}$ , where  $\operatorname{pred}(f(x))$  represents the predicted label for  $x$ ,
- An independent model  $g$  trained on  $(P_{\perp}x,y)$  where  $(x,y)\sim \mathcal{D}$  satisfies  $|\mathrm{Acc}(g) - \mathrm{Acc}(f)|\leq \epsilon_2$
for some small  $\epsilon_{1}$  and  $\epsilon_{2}$ . Here  $P_{\perp}$  is the projection matrix onto the subspace orthogonal to  $P$ , and  $\mathrm{Acc}(f)$  represents the accuracy of  $f$ .
In words, LD-SB says that there exists a small  $k$ -dimensional subspace (given by the projection matrix  $P$ ) in the input space  $\mathbb{R}^d$ , which is the only thing that the model  $f$  considers in labeling any input point  $x$ . In particular, if we mix two data points  $x_1$  and  $x_2$  by using the projection of  $x_1$  onto  $P$  and the projection of  $x_2$  onto the orthogonal subspace  $P_{\perp}$ , the output of  $f$  on this mixed point  $Px_1 + P_{\perp}x_2$  is the same as that on  $x_1$ . This would have been fine if the subspace  $P_{\perp}$  does not contain any feature useful for classification. However, the third bullet point says that  $P_{\perp}$  indeed contains features that are useful for classification since an independent model  $g$  trained on  $(P_{\perp}x,y)$  achieves high accuracy.
Theoretically, we demonstrate LD-SB of 1-hidden layer NNs in the infinite width limit for a fairly general class of distributions called independent features model (IFM), where the features (i.e., coordinates) are distributed independently conditioned on the label. IFM has a long history and is widely studied, especially in the context of naive-Bayes classifiers Lewis (1998). For IFM, we show that as long as there is even a single feature in which the data is linearly separable, NNs trained using SGD will learn models that rely almost exclusively on this linearly separable feature, even when there are an arbitrarily large number of features in which the data is separable but with a non-linear boundary. Empirically, we demonstrate LD-SB on three real world datasets: binary and multiclass version of Imagenette (FastAI, 2021), waterbirds-landbirds (Sagawa et al., 2020a) as well as the ImageNet (Deng et al., 2009) dataset. Compared to the results in Shah et al. (2020), our results (i) theoretically show LD-SB in a fairly general setting and (ii) empirically show LD-SB on real datasets.
Finally, building upon these insights, we propose a simple ensemble method - OrthoP - that sequentially constructs NNs by projecting out the input data directions that are used by previous NNs. We demonstrate that this method can lead to significantly more robust ensembles for real-world datasets in presence of simple distribution shifts like Gaussian noise.
Why study 1-hidden layer networks in the infinite width regime?
1. From a practical standpoint, the dominant paradigm in machine learning right now is to pretrain large models on large amounts of data and then finetune on small target datasets. Given the large and diverse pretraining data seen by these models, it has been observed that they do learn rich features (Rosenfeld et al., 2022; Nasery et al., 2022). However, finetuning on target datasets might not utilize all the features in the pretrained model. Consequently, approaches that can train robust finetuning heads (such as a 1-hidden layer network on top) can be quite effective.
2. From a theoretical standpoint, there have been several works that analyze training dynamics of finite width networks (Ding et al., 2022), and show convergence to global minima on the training data. However, these results do not identify which among the many global minima, the training dynamics converge to, which is crucial in determining the nature of SB of the converged model. Such a precise characterization of the final convergence point is known only for infinite width 1-hidden layer networks (Chizat et al., 2019; Chizat & Bach, 2020).
3. While our theoretical analysis works in the setting of infinite width networks, our extensive experiments on several large scale datasets suggest that the results continue to hold even for finite width networks. Furthermore, Vyas et al. (2023) show that the behavior of neural networks remains consistent with width in the feature learning regime.
To summarize, this paper characterizes the nature of SB in 1-hidden layer networks, and also proposes a novel ensemble training approach, called OrthoP, which leads to more robust ensembles. While the theoretical results are in the infinite width regime, empirical results on several real world datasets show that these results continue to hold even for finite width networks.
Paper organization: This paper is organized as follows. Section 2 presents related work. Section 3 presents preliminaries. Our main results on LD-SB are presented in Section 4. Section 5 presents results on training diverse classifiers. We conclude in Section 6.
Simplicity Bias: Subsequent to Shah et al. (2020), there have been several papers investigating the presence/absence of SB in various networks as well as reasons behind SB (Scimeca et al., 2021). Of these, Huh et al. (2021) is the most closely related work to ours. Huh et al. (2021) empirically observe that on certain synthetic datasets, the embeddings of NNs both at initialization as well as after training have a low rank structure. In contrast, we prove LD-SB theoretically on the IFM model as well as empirically validate this on real datasets. Furthermore, our results show that while the network weights exhibit low rank structure in the rich regime (see Section 3.2 for definition), the manifestation of LD-SB is far more subtle in lazy regime. Moreover, we also show how to use LD-SB to train a second diverse model and combine it to obtain a robust ensemble. Galanti & Poggio (2022) provide a theoretical intuition behind the relation between various hyperparameters (such as learning rate, batch size etc.) and rank of learnt weight matrices, and demonstrate it empirically. Pezeshki et al. (2021) propose that gradient starvation at the beginning of training is a potential reason for SB in the lazy/NTK regime but the conditions are hard to interpret. In contrast, our results are shown for any dataset in the IFM model in the rich regime of training. Finally, Lyu et al. (2021) consider anti-symmetric datasets and show that single hidden layer input homogeneous networks (i.e., without bias parameters) converge to linear classifiers. However, our results hold for general datasets and do not require input homogeneity.
Learning diverse classifiers: There have been several works that attempt to learn diverse classifiers. Most works try to learn such models by ensuring that the input gradients of these models do not align (Ross & Doshi-Velez, 2018; Teney et al., 2022). Xu et al. (2022) propose 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. Lee et al. (2022) learn diverse classifiers by enforcing diversity on unlabeled target data.
Spurious correlations: There has been a large body of work which identifies 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., 2020b). Simplicity bias seems to be one of the primary reasons behind learning spurious correlations within NNs (Shah et al., 2020).
Implicit bias of gradient descent: There is also a large body of work understanding the implicit bias of gradient descent dynamics. Most of these works are for standard linear (Ji & Telgarsky, 2019) or deep linear networks (Soudry et al., 2018; Gunasekar et al., 2018). For nonlinear neural networks,
one of the well-known results is for the case of 1-hidden layer neural networks with homogeneous activation functions (Chizat & Bach, 2020), which we crucially use in our proofs. More related works are provided in Appendix.
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.
3.1 Basic notions
1-hidden layer neural networks and loss function. Consider instances  $x \in \mathcal{R}^d$  and labels  $y \in \{\pm 1\}$  jointly distributed as  $\mathcal{D}$ . A 1-hidden layer neural network model (or fully connected network (FCN)) for predicting the label for a given instance  $x$ , is defined by parameters  $(\bar{w} \in \mathbb{R}^{m \times d}, \bar{b} \in \mathbb{R}^m, \bar{a} \in \mathbb{R}^m)$ . For a fixed activation function  $\phi$ , given input instance  $x$ , the model is given as  $f((\bar{w}, \bar{b}, \bar{a}), x) := \langle \bar{a}, \phi(\bar{w}x + \bar{b}) \rangle$ , where  $\phi(\cdot)$  is applied elementwise. The cross entropy loss  $\mathcal{L}$  for a given model  $f$ , input  $x$  and label  $y$  is given as  $\mathcal{L}(f(x), y) \stackrel{\mathrm{def}}{=} \log(1 + \exp(-yf((\bar{w}, \bar{b}, \bar{a}), x)))$ .
Margin. For data distribution  $\mathcal{D}$ , the margin of a model  $f(x)$  is given as  $\min_{(x,y)\sim \mathcal{D}}yf(x)$ .
Notation. Here is some useful notation that we will use repeatedly. For a matrix  $A$ ,  $A(i,.)$  denotes the  $i$ th row of  $A$ . For any  $k \in \mathbb{N}$ ,  $\mathbb{S}^{k-1}$  denotes the surface of the unit norm Euclidean sphere in dimension  $k$ .
3.2 Initializations
The gradient descent dynamics of the network depends strongly on the scale of initialization. In this work, we primarily consider rich regime initialization.
Rich regime. In rich regime initialization, for any  $i \in [m]$ , the parameters  $(\bar{w}(i, .), \bar{b}(i))$  of the first layer are sampled from a uniform distribution on  $\mathbb{S}^d$ . Each  $\bar{a}(i)$  is sampled from Unif\{-1, 1\}, and the output of the network is scaled down by  $\frac{1}{m}$  (Chizat & Bach, 2020). This is roughly equivalent in scale to Xavier initialization Glorot & Bengio (2010), where the weight parameters in both the layers are initialized approximately as  $\mathcal{N}(0, \frac{2}{m})$  when  $m \gg d$ .
In addition, we also present some results for the lazy regime initialization described below.
Lazy regime. In the lazy regime, the weight parameters in the first layer are initialized with  $\mathcal{N}(0,\frac{1}{d})$ , those of second layer are initialized with  $\mathcal{N}(0,\frac{1}{m})$  and the biases are initialized to 0 (Bietti & Mairal, 2019; Lee et al., 2019). This is approximately equivalent in scale to Kaiming initialization (He et al., 2015).
3.3 Infinite Width Case
For 1-hidden layer neural networks with ReLU activation in the infinite width limit i.e., as  $m\to \infty$  Jacot et al. (2018); Chizat et al. (2019); Chizat & Bach (2020) gave interesting characterizations of the trained model. As mentioned above, the training process of these models falls into one of two regimes depending on the scale of initialization (Chizat et al., 2019):
Rich regime. In the infinite width limit, the neural network parameters can be thought of as a distribution  $\nu$  over triples  $(w,b,a)\in \mathbb{S}^{d + 1}$  where  $w\in \mathbb{R}^d,b,a\in \mathbb{R}$ . Under the rich regime initialization, the function  $f$  computed by the model can be expressed as
$$
f (\nu , x) = \mathbb {E} _ {(w, b, a) \sim \nu} [ a (\phi (\langle w, x \rangle + b) ]. \tag {1}
$$
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  $^3$  of the trained model under the cross entropy loss  $\mathbb{E}_{(x,y)\sim \mathcal{D}}[\mathcal{L}(\nu ,(x,y))]$ .
<sup>3</sup>Theorem 3.1 is an informal version of Chizat & Bach 2020, Theorem 5. For exact result, refer Theorem E.1 in Appendix E
Theorem 3.1. (Informal)(Chizat & Bach, 2020) Under rich initialization in the infinite width limit with cross entropy loss, if gradient flow on 1-hidden layer NN with ReLU activation converges, it converges to a maximum margin classifier  $\nu^{*}$  given as
$$
\nu^ {*} = \underset {\nu \in \mathcal {P} (\mathbb {S} ^ {d + 1})} {\arg \max } \underset {(x, y) \sim \mathcal {D}} {\min } y f (\nu , x), \tag {2}
$$
where  $\mathcal{P}(\mathbb{S}^{d + 1})$  denotes the space of distributions over  $\mathbb{S}^{d + 1}$ .
This training regime is known as the 'rich' regime since it learns data dependent features  $\langle w,\cdot \rangle$
Lazy regime. Jacot et al. (2018) showed that in the infinite width limit, the neural network behaves like a kernel machine. This kernel is popularly known as the Neural Tangent Kernel(NTK), and is given by  $K(x,x^{\prime}) = \left\langle \frac{\partial f(x)}{\partial W},\frac{\partial f(x^{\prime})}{\partial W}\right\rangle$ , where  $W$  denotes the set of all trainable weight parameters. This initialization regime is called 'lazy' regime since the weights do not change much from initialization, and the NTK remains almost constant, i.e., the network does not learn data dependent features. We will use the following characterization of the NTK for 1-hidden layer neural networks.
Theorem 3.2. Bietti & Mairal (2019) Under lazy regime initialization in the infinite width limit, the NTK for 1-hidden layer neural networks with ReLU activation i.e.,  $\phi(u) = \max(u, 0)$ , is given as
$$
K (x, x ^ {\prime}) = \| x \| \| x ^ {\prime} \| \kappa \left(\frac {\langle x , x ^ {\prime} \rangle}{\| x \| \| x ^ {\prime} \|}\right), w h e r e \kappa (u) = \frac {1}{\pi} (2 u (\pi - c o s ^ {- 1} (u)) + \sqrt {1 - u ^ {2}}).
$$
Lazy regime for binary classification. Soudry et al. (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. 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 NTK, will tend towards the  $\mathcal{L}_2$  max-margin-classifier given by
$$
\underset {f \in \mathcal {H}} {\arg \min } \| f \| _ {\mathcal {H}} \text {s . t .} y f (x) \geq 1 \forall (x, y) \sim \mathcal {D}, \tag {3}
$$
where  $\mathcal{H}$  represents the Reproducing Kernel Hilbert Space (RKHS) associated with the finite width kernel (Chizat, 2020). With increasing width, this kernel tends towards the infinite-width NTK (which is universal (Ji et al., 2020)). Therefore, in lazy regime, we will focus on the  $\mathcal{L}_2$  max-margin-classifier induced by the infinite-width NTK.
4 Characterization of SB in 1-hidden layer neural networks
In this section, we first theoretically characterize the SB exhibited by gradient descent on linearly separable datasets in the independent features model (IFM). The main result, stated in Theorem 4.1, is that for binary classification of inputs in  $\mathbb{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. In other words, the simplicity bias of these networks is characterized by low dimensional input dependence, which we denote by LD-SB. We then experimentally verify that NNs trained on some real datasets do indeed satisfy LD-SB.
We consider datasets in the independent features model (IFM), where the joint distribution over  $(x,y)$  satisfies  $p(x,y) = r(y)\prod_{i = 1}^{d}q_{i}(x_{i}|y)$ , i.e., the features are distributed independently conditioned on the label  $y$ . Here  $r(y)$  is a distribution over  $\{-1, + 1\}$  and  $q_{i}(x_{i}|y)$  denotes the conditional distribution of  $i^{\mathrm{th}}$ -coordinate  $x_{i}$  given  $y$ . IFM is widely studied in literature, particularly in the context of naive-Bayes classifiers Lewis (1998). We make the following assumptions which posit that there are at least two features of differing complexity for classification: one with a linear boundary and at least one other with a non-linear boundary. See Figure 2 for an illustrative example.
Figure 2: Illustration of an IFM dataset. Given a class  $\pm 1$  represented by blue and red respectively, each coordinate value is drawn independently from the corresponding distribution. Shown above are the supports of distributions on three different coordinates for an illustrative IFM dataset, for positive and negative labels.
- One of the coordinates (say, the  $1^{\mathrm{st}}$  coordinate WLOG) is separable by a linear decision boundary with margin  $\gamma$  (see Figure 2), i.e.,  $\exists \gamma > 0$ , such that  $\gamma \in \operatorname{Supp}(q_1(x_1 | y = +1)) \subseteq [\gamma, \infty)$  and  $-\gamma \in \operatorname{Supp}(q_1(x_1 | y = -1)) \subseteq (-\infty, -\gamma]$ , where  $\operatorname{Supp}(\cdot)$  denotes the support of a distribution.
- None of the other coordinates is linearly separable. More precisely, for all the other coordinates  $i \in [d] \setminus \{1\}$ ,  $0 \in \operatorname{Supp}(q_i(x_i | y = -1))$  and  $\{-1, +1\} \subseteq \operatorname{Supp}(q_i(x_i | y = +1))$ .
- The dataset can be perfectly classified even without using the linear coordinate. This means,  $\exists i \neq 1$ , such that  $q_{i}(x_{i}|y)$  has disjoint support for  $y = +1$  and  $y = -1$ .
Though we assume axis aligned features, our results also hold for any rotation of the dataset. While our results hold in the general IFM setting, in comparison, current results for SB e.g., Shah et al. (2020), are obtained for very specialized datasets within  $\mathrm{IFM}^5$ , and do not apply to IFM in general.
4.2 Main result
Our main result states that, for rich initialization (Section 3.2), NNs demonstrate LD-SB for any IFM dataset satisfying the above conditions. Its proof appears in Appendix A.1.
Theorem 4.1. For any dataset in the IFM model with bounded density and bounded support, satisfying the above conditions and  $\gamma \geq 1$ , and for 1-hidden layer networks with ReLU activation in the infinite width limit (i.e., Eqn. (1)), there is a unique max margin classifier  $\nu^{*}$  (i.e., satisfying Eqn. (2)). This  $\nu^{*}$  is given by:  $\nu^{*} = 0.5\delta_{\theta_{1}} + 0.5\delta_{\theta_{2}}$  on  $S^{d+1}$ , where  $\delta$  represents the dirac-delta distribution,  $\theta_{1} = (\frac{\gamma}{\sqrt{2(1 + \gamma^{2})}}\mathbf{e}_{1}, \frac{1}{\sqrt{2(1 + \gamma^{2})}}, 1 / \sqrt{2}), \theta_{2} = (-\frac{\gamma}{\sqrt{2(1 + \gamma^{2})}}\mathbf{e}_{1}, \frac{1}{\sqrt{2(1 + \gamma^{2})}}, -1 / \sqrt{2})$  and  $\mathbf{e}_{1} \stackrel{\text{def}}{=} [1,0,\dots,0]$  denotes first standard basis vector. This implies  $f(\nu^{*}, Px^{(1)} + P_{\perp}x^{(2)}) = f(\nu^{*}, x^{(1)}) \forall (x^{(1)}, y^{(1)}), (x^{(2)}, y^{(2)}) \sim \mathcal{D}$ , where  $P$  represents the (rank-1) projection matrix on first coordinate.
Together with Theorem 3.1, this implies that if gradient flow converges, it converges to  $\nu^{*}$  given above. Since  $P$  is a rank-1 matrix and  $f(\nu^{*}, Px^{(1)} + P_{\perp}x^{(2)}) = f(\nu^{*}, x^{(1)})$ ,  $\nu^{*}$  satisfies the first two conditions of LD-SB (Definition 1.1) with  $k = 1$  and  $\epsilon_1 = 0$ . Moreover, since at least one of the coordinates  $\{2, \ldots, d\}$  has disjoint support for  $q_i(x_i|y = +1)$  and  $q_i(x_i|y = -1)$ ,  $P_{\perp}(x)$  can still perfectly classify the given dataset, thereby implying the third condition of LD-SB with  $\epsilon_2 = 0$ .