SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 7: Mistake diversity and class conditioned logit correlation of models trained sequentially after projecting out important features of the first model (Mist-Div  $(f, f_{\mathrm{proj}})$  and CC-LogitCorr  $(f, f_{\mathrm{proj}})$  resp.) and using the method mentioned in Evading-SB ((Mist-Div  $(f, f_{\mathrm{esb}})$  and CC-LogitCorr  $(f, f_{\mathrm{esb}})$  resp.). Numbers have been bolded if they are better than the other even after taking one standard deviation into account.
Dataset	Mist-Div (f, fesb) (↑)	Mist-Div (f, fproj) (↑)	CC-LogitCorr (f, fesb) (↓)	CC-LogitCorr (f, fproj) (↓)
B-Imagenette	17.86 ± 0.0	21.15 ± 1.57	90.61 ± 0.48	90.86 ± 1.08
Imagenette	11.23 ± 0.24	11.44 ± 0.65	92.48 ± 0.68	91 ± 0.59
Waterbirds	13.6 ± 0.36	14.53 ± 0.48	95.41 ± 0.41	93.81 ± 0.48
MNIST-CIFAR	0.0 ± 0.0	5.56 ± 7.89	76.19 ± 2.38	78.74 ± 2.28
Imagenet	4.95 ± 0.11	12.31 ± 0.16	98.83 ± 0.0	92.52 ± 0.01
C Extended Related Works
In this section, we provide an extensive literature survey of various topics that the paper is based on.
Low rank Simplicity Bias in Linear Networks Multiple works have established low rank simplicity bias for gradient descent on linear networks, both for squared loss as well as cross-entropy loss. For squared loss, Gunasekar et al. (2017) conjectured that the network is biased towards finding minimum nuclear norm solutions for two-layer linear networks. Arora et al. (2019) refuted the conjecture and instead argued that the network is biased towards finding low rank solutions. Razin & Cohen (2020) provided empirical support 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. Li et al. (2021) established that for small enough initialization, gradient flow on linear networks follows greedy low-rank learning trajectory. For binary classification on linearly separable data, Ji & Telgarsky (2019) showed that the weight matrices of a linear network eventually become rank-1 as training progresses.
Low rank Simplicity Bias in Non-Linear Networks There have been a few works establishing low-rank simplicity bias for non-linear networks. Phuong & Lampert (2021) showed that, for orthogonally separable datasets, the rank of first layer weights for two-layer ReLU networks eventually tends to two under gradient flow on logistic loss. Frei et al. (2023) extend this result to nearly-orthogonal datasets and Leaky-ReLU networks. Timor et al. (2023) showed that, maximum margin ReLU networks of sufficiently large depth implicitly also have low rank. Huh et al. (2021) empirically established that the rank of the embeddings learnt by a neural network with ReLU activations goes down as training progresses. Galanti & Poggio (2022) provided an intuition behind the relation between the rank of the weight matrices and various hyperparameter such as batch size, weight decay etc. In contrast to these works, for 1 layer nets, we theoretically and empirically establish that the network depends on an extremely low dimensional projection of the input, and this bias can be utilized to develop a robust classifier.
Adversarial Robustness Some of the recent works have also shown that gradient flow on shallow networks generally leads to non-robust classifiers. Vardi et al. (2022) showed that for nearly orthogonal datasets, gradient flow on 2-layer ReLU networks converges to non-robust networks, even when robust classifiers are possible for the dataset. Melamed et al. (2023) show similar results for datasets which lie on a low-dimensional linear subspace.
Relation to OOD Many recent works in OOD detection (Cook et al., 2020; Zaeemzadeh et al., 2021) explicitly create low-rank embeddings so that it is easier to discriminate them for an OOD point. Other works also implicitly rely on the low-rank nature of the embeddings. Ndiour et al. (2020) use PCA on the learnt features, and only model the likelihood along the small subspace spanned by the top few directions. Wang et al. (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. While there have been works implicitly utilizing the low rank property of embeddings, we note that our paper (i) demonstrates low rank property of the weights, rather than that of embeddings, and (ii) shows that it is a consequence of SB.
Other Simplicity Bias There have been many works exploring the nature of simplicity bias in neural networks, both empirically and theoretically. Kalimeris et al. (2019) empirically demonstrated that SGD on neural networks gradually learns functions of increasing complexity. Rahaman et al. (2018) empirically demonstrated that neural networks tend to learn lower frequency functions first. Ronen et al. (2019) theoretically established that in NTK regime, the convergence rate depends on the eigenvalues of the kernel spectrum. Hacohen et al. (2020) showed that neural networks always learn train and test examples almost in the same order, irrespective of the architecture. Pezeshki et al. (2021) proposes that gradient starvation at the beginning of training is a potential reason for SB in the lazy/NTK regime but the conditions are hard to interpret. In contrast, our results are shown for any dataset in the IFM model in the rich regime of training. Lyu et al. (2021) consider antisymmetric datasets and show that single hidden layer input homogeneous networks (i.e., without bias parameters) converge to linear classifiers. However, such networks have strictly weaker expressive power compared to those with bias parameters. Hacohen & Weinshall (2022) showed that for deep linear networks, in NTK regime, they learn the higher principal components of the input data first.
Most of the previous works used simplicity bias as a reason behind better generalization of neural nets. However, Shah et al. (2020) showed that extreme simplicity bias could also lead to worse OOD performance.
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) 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. Lee et al. (2022) learns diverse classifiers by enforcing diversity on unlabeled target data.
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., 2020b).
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.
D More discussion on the extension of results to deep nets
Extending our theoretical results to deep nets is a very exciting and challenging research direction. For shallow as well as deep nets, even in the mean field regime of training, results regarding convergence to global minima have been established (Chizat & Bach, 2018; Fang et al., 2021). However, to the best of our knowledge, only for 1-hidden layer FCN (Chizat & Bach, 2020), a precise characterization of the global minima to which gradient flow converges has been established. Understanding this implicit bias of gradient flow is still an open problem for deep nets, which we think is essential for extension of our results to deep nets.
E Convergence to  $\mathcal{F}_1$  -max-margin classifier for ReLU networks
In this section, we will provide a brief background on Wasserstein gradient flow and state the precise result of Chizat & Bach (2020) regarding the asymptotic convergence point of gradient flow on ReLU networks. We will follow the notation of Chizat & Bach (2020) for ease of the reader. In this entire section, we will consider that a neural network is parameterized by a probability measure  $\mu$  on the neurons and is given by
$$
h (\mu , x) = \int \phi (w, x) d \mu (w)
$$
where  $\phi (w,x) = b(a^{\top}(x,1))_{+}$  (+ denotes the positive component, i.e the ReLU activation) with  $w = (a,b)\in \mathbb{R}^{d + 2}$ .
E.1 Wasserstein gradient flow
Gradient flow can be defined for many functions  $f$  over a general metric space  $\mathcal{X}$ . For a given step size  $\eta$ , define
$$
x _ {k + 1} \in \arg \min  f (x) + \frac {1}{2 \eta} d (x, x _ {k}) ^ {2}
$$
where  $d$  is the metric associated with  $\mathcal{X}$ . With appropriate interpolation schemes (Santambrogio, 2016), this curve converges to the gradient flow curve as step size tends to 0.
Wasserstein metric on the space of probability measures is defined as
$$
W _ {p} (\nu_ {1}, \nu_ {2}) = \inf _ {\gamma \in \Gamma_ {\nu_ {1}, \nu_ {2}}} \left[ \int \| x - y \| ^ {p} d \gamma \right] ^ {1 / p}
$$
where  $\nu_{1},\nu_{2}$  are two probability measures and  $\gamma$  is a coupling between them (i.e. marginals of  $\gamma$  are  $\nu_{1}$  and  $\nu_{2}$ ). Here, we will be particularly concerned with the case  $p = 2$ . For two discrete measures
$\nu_{1} = \frac{1}{m}\sum \delta_{x_{i}}$  and  $\nu_{2} = \frac{1}{m}\sum \delta_{y_{j}}$ , their Wasserstein distance is defined as
$$
W _ {2} ^ {2} (\nu_ {1}, \nu_ {2}) = \frac {1}{m} \min \| x _ {i} - y _ {\sigma (i)} \| ^ {2}
$$
over all permutations  $\sigma : \{1,..,m\} \to \{1,..,m\}$ . Notice that if we are considering a small neighborhood of  $\nu_{1}$ , then the mapping would remain the same within that small neighborhood. Thus, Wasserstein gradient flow on the discrete measure would be the same as gradient flow on the particles as the notion of distance is the same locally. This intuition leads to the proof that gradient flow on a 2-layer neural net converges to Wasserstein gradient flow on the probability measure  $\mu$  as width tends to infinity. This is made formal in Chizat & Bach (2018).
E.2 Asymptotic convergence point of gradient flow on ReLU networks
A neural network is parameterized by a probability measure  $\mu$  on the neurons and is given by
where  $\phi (w,x) = b(a^{\top}(x,1))_{+}$ $(+$  denotes the positive component, i.e the ReLU activation) with  $w = (a,b)\in \mathbb{R}^{d + 2}$ . As the network is 2-homogeneous, a projection of the measure  $\mu$  on the unit sphere can be defined. The projection operator  $(\Pi_2)$  on the sphere for a measure  $\mu$  is defined such that for any continuous function  $\varphi$  on the sphere,
$$
\int_ {\mathbb {S} ^ {d + 1}} \varphi (\theta) d [ \Pi_ {2} (\mu) ] (\theta) = \int_ {\mathbb {R} ^ {d + 2}} \| w \| ^ {2} \varphi (w / \| w \|) d \mu (w)
$$
Now, let  $\rho$  denote the input distribution on the input space  $\mathcal{X}$  and let the labeling function  $y:\mathcal{X}\to \mathcal{Y}$  be deterministic. Then, consider the population objective given by
$$
F (\mu) = - \log \left[ \int_ {\mathcal {X}} \exp (- y (x) h (\mu , x)) d \rho (x) \right]
$$
Note that log doesn't affect the direction of the gradients, thus, the trajectory of gradient flow on this loss is the same as on exponential loss. Also, let the population smooth margin be given by
$$
S (f) = - \log \left(\int_ {\mathcal {X}} \exp (- f (x)) d \rho (x)\right)
$$
For this particular case,  $f(x) = y(x)h(\mu, x)$ . Denote  $y(x) \cdot h(\mu, x)$  by  $\hat{h}(\mu)$ .
Theorem E.1. Suppose that  $\rho$  has bounded density and bounded support, and labeling function  $y$  is continuous, then there exists a Wasserstein gradient flow  $(\mu_t)$  on  $F$  with  $\mu_0 = \mathcal{U}(\mathbb{S}^d) \otimes \mathcal{U}\{-1,1\}$ , i.e., input (resp. output) weights uniformly distributed on the sphere (resp. on  $\{-1,1\}$ ). If  $\nabla S(\hat{h}(\mu_t))$  converges weakly in  $\mathcal{P}(\mathcal{X})$ , if  $\bar{\nu}_t = \Pi_2(\mu_t) / ([\Pi_2(\mu_t)](\mathbb{S}^{d+1}))$  converges weakly in  $\mathcal{P}(\mathbb{S}^{d+1})$  and  $F'(\mu_t)$  converges in  $C_{loc}^1$  to  $F'$  that satisfies the Morse-Sard property, then  $h(\bar{\nu}_{\infty},.)$  is a maximizer for  $\max_{\|f\|_{\mathcal{F}_1} \leq 1} \min_{x \in \mathcal{X}} y(x)f(x)$ .
where  $\mathcal{P}(\mathcal{X})$  denotes the space of probability distributions on  $\mathcal{X}$  and  $[\Pi_2(\mu_t)](\mathbb{S}^{d + 1})$  denotes the total mass of the measure  $\Pi_2(\mu_t)$  on  $\mathbb{S}^{d + 1}$ .
To parse the theorem, note that
$$
\nabla S (f) = \frac {\exp (- f (x)) d \rho (x)}{\int_ {\mathcal {X}} \exp (- f (x ^ {\prime})) d \rho (x ^ {\prime})}
$$
Thus,  $\nabla S(f)$  convergence means that the exponentiated normalized margins converge. Also,  $\bar{\nu}_t$  is similar to the directional convergence of weights, however, in this case, weights are replaced by directions in  $\mathbb{S}^{d + 1}$ . For explanation of the Morse-Sard property and the metric  $C_{loc}^{1}$ , please refer to Appendix H of Chizat & Bach (2020).