SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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). 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.
4.3 Lazy regime
In this regime, we will work with the following dataset within the IFM family:
For  $y \in \{\pm 1\}$  we generate  $(x, y) \in D$  as
$$
\mathbf {x} _ {1} = \gamma y,   \forall i \in 2,.., d, \mathbf {x} _ {i} = \left\{ \begin{array}{c c c} \pm 1 & \text {f o r} & y = 1 \\ 0 & \text {f o r} & y = - 1 \end{array} \right.
$$
4Using linear probe for classifying pretrained representations is a standard practice in self-supervised learning (Chen et al., 2020a; Grill et al., 2020).
In Shah et al. (2020), the theoretical results were obtained mainly for linear and one non-linear coordinate
Table 1: Demonstration of LD-SB in the rich regime: This table presents  $P_{\perp}$  and  $P$  logit as well as prediction changes on the five datasets. These results confirm that projection of input  $x$  onto the subspace spanned by  $P$  essentially determines the model's prediction on  $x$ .  $\uparrow$  (resp.  $\downarrow$ ) indicates that LD-SB implies a large (resp. small) value.
Dataset	rank (P)	P⊥-LC (↓)	P-LC (↑)	P⊥-pC (↓)	P-pC (↑)
b-Imagenette	1	28.57 ± 0.26	92.13 ± 0.24	6.35 ± 0.06	47.02±0.24
Imagenette	10	33.64 ± 1.21	106.29 ± 0.53	12.04±0.29	89.88±0.08
Waterbirds	3	25.24 ± 1.03	102.35 ± 0.19	6.78 ± 0.15	35.96±0.02
MNIST-CIFAR	1	38.97 ± 0.76	101.98 ± 0.31	5.41 ± 0.55	45.15±0.44
Imagenet	150	15.78 ± 0.05	132.05 ± 0.06	13.05±0.03	99.76±0.01
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. For this dataset, we provide the characterization of max-margin NTK (as in Eqn. (3)):
Theorem 4.2. There exists  $\delta_0 > 0$  such that for every  $\delta < \delta_0$ , there exists an absolute constant  $N$  such that for all  $d > N$  and  $\gamma \in [7, \delta \sqrt{d})$ , the  $\mathcal{L}_2$  max-margin classifier for joint training of both the layers of 1-hidden layer FCN with ReLU activation in the NTK regime on the dataset  $D$ , i.e., any  $f$  satisfying Eqn. (3) satisfies:
$$
\operatorname {p r e d} \left(f \left(P x ^ {(1)} + P _ {\perp} x ^ {(2)}\right)\right) = \operatorname {p r e d} \left(f \left(x ^ {(1)}\right)\right) \quad \forall \left(x ^ {(1)}, y ^ {(1)}\right), \left(x ^ {(2)}, y ^ {(2)}\right) \in D
$$
where  $P$  represents the projection matrix on the first coordinate and  $\operatorname{pred}(f(x))$  represents the predicted label by the model  $f$  on  $x$ .
The proof of this theorem is presented in Appendix A.2. The above theorem shows that the prediction on a mixed example  $P x^{(1)} + P_{\perp}x^{(2)}$  is the same as that on  $x^{(1)}$  (i.e.,  $\epsilon_1 = 0$  in Definition 1.1). Furthermore, since there exists at least one coordinate  $i \neq 1$  which can be used to perfectly classify the dataset, we have that Definition 1.1 is satisfied with  $\epsilon_2 = 0$ , thus establishing LD-SB.
4.4 Empirical verification
In this section, we will present empirical results demonstrating LD-SB on 4 real datasets: Imagenette (FastAI, 2021), a binary version of Imagenette (b-Imagenette), waterbirds-landbirds (Sagawa et al., 2020a) and Imagenet (Deng et al., 2009) as well as one designed dataset MNIST-CIFAR (Shah et al., 2020). More details about the datasets can be found in Appendix B.1.
4.4.1 Experimental setup
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. During finetuning, we freeze the backbone Resnet-50 model and train only the 1-hidden layer head (details in Appendix B.1).
Demonstrating LD-SB: Given a model  $f(\cdot)$ , 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  $x_{2}$  as  $Px_{1} + P_{\perp}x_{2}$ , the model's output on the mixed input  $\widetilde{x} \stackrel{\mathrm{def}}{=} Px_{1} + P_{\perp}x_{2}$ ,  $f(\widetilde{x})$  is always close to the model's output on  $x_{1}$  i.e.,  $f(x_{1})$ . We measure closeness in four metrics: (1)  $P_{\perp}$  logit change  $(P_{\perp}$ -LC): relative change of logits wrt  $x_{1}$  i.e.,  $\|f(\widetilde{x}) - f(x_{1})\| / \|f(x_{1})\|$ , (2)  $P$  logit change (P-LC): relative change wrt logits of  $x_{2}$  i.e.,  $\|f(\widetilde{x}) - f(x_{2})\| / \|f(x_{2})\|$ , (3)  $P_{\perp}$ -prediction change  $(P_{\perp}$ -pC):  $\mathbb{P}[\mathrm{pred}(f(\widetilde{x})) \neq \mathrm{pred}(f(x_{1}))]$ , and (4)  $P$ -prediction change  $(P$ -pC):  $\mathbb{P}[\mathrm{pred}(f(\widetilde{x})) \neq \mathrm{pred}(f(x_{2}))]$ . The quantities rank  $(P)$  and  $P_{\perp}$ -pC correspond to  $k$  and  $\epsilon_{1}$  in Definition 1.1 respectively. To demonstrate that the subspace  $P_{\perp}$  has features that are useful for prediction, we also train a new model  $f_{\mathrm{proj}}$  as follows. Given the initial model  $f$  and the corresponding projection matrix  $P$ , we then train another model  $f_{\mathrm{proj}}$  by projecting the input through  $P_{\perp}$  i.e., instead of using dataset  $(x^{(i)}, y^{(i)})$  for training, we use  $(P_{\perp}x^{(i)}, y^{(i)})$  for training the second model (denoted by  $f_{\mathrm{proj}}$ ). We refer to this training procedure as OrthoP for orthogonal projection. The quantity  $|\mathrm{Acc}(f) - \mathrm{Acc}(f_{\mathrm{proj}})|$  corresponds to  $\epsilon_{2}$  in Definition 1.1. We now describe how we identify  $P$  in rich and lazy regimes.
Table 2: Demonstration of LD-SB in lazy regime: This table presents  $P_{\perp}$  and  $P$  logit as well as prediction changes on the five datasets. These results confirm that the projection of input  $x$  onto the subspace spanned by  $P$  essentially determines the model's prediction on  $x$ .
Dataset	rank (P)	P⊥-LC (↓)	P-LC (↑)	P⊥-pC (↓)	P-pC (↑)
b-Imagenette	1	36.94±1.01	138.41±1.62	5.5 ± 1.13	47.7 ± 1.55
Imagenette	15	55.99±3.86	133.86±5.42	11.25±0.36	89.75±0.15
Waterbirds	6	36.89±5.18	105.41±7.06	20.74±0.64	45.96±0.69
MNIST-CIFAR	2	24.9 ± 0.61	141.12±1.86	0.53 ± 0.24	49.83±0.78
Imagenet	200	32.74±0.02	132.47±0.04	18.2 ± 0.16	99.74±0.01
(a) Rich regime
(b) Lazy regime
Figure 3: Evolution of effective rank of first layer weight matrices in rich and lazy regimes.
4.4.2 Rich regime
Theorem 4.1 suggests that asymptotically, the first layer weight matrix will be low rank. However, since we train only for a finite amount of time, the weight matrix will only be approximately low rank. To quantify this, we use the notion of effective rank Roy & Vetterli (2007).
Definition 4.3. Given a matrix  $M$ , its effective rank is defined as  $e^{-\sum_{i}\overline{\sigma_{i}(M)^{2}}\log \overline{\sigma_{i}(M)^{2}}}$  where  $\sigma_{i}(M)$  denotes the  $i^{\mathrm{th}}$  singular value of  $M$  and  $\overline{\sigma_i(M)^2}\stackrel {\mathrm{def}}{=}\frac{\sigma_i(M)^2}{\sum_i\sigma_i(M)^2}$ .
One way to interpret the effective rank is that it is the exponential of von-Neumann entropy Petz (2001) of the matrix  $\frac{MM^{\top}}{\operatorname{Tr}(MM^{\top})}$ , where  $\operatorname{Tr}(\cdot)$  denotes the trace of a matrix. For illustration, the effective rank of a projection matrix onto  $k$  dimensions equals  $k$ .
Figure 3a shows the evolution of the effective rank through training on the four datasets. We observe that the effective rank of the weight matrix decreases drastically towards the end of training. In this case, we set  $P$  to be the subspace spanned by the top singular directions of the first layer weight matrix. Table 1 presents the results for  $P_{\perp}$  and  $P$ -LC as well as pC, while Table 3 presents  $\mathrm{Acc}(f)$  and  $\mathrm{Acc}(f_{\mathrm{proj}})$ . These results establish LD-SB in the rich regime.
Table 3: Accuracy of  ${f}_{\text{proj }}$  in rich regime
Dataset	Acc(f)	Acc(fproj)
b-Imagenette	93.35	91.35 ± 0.32
Imagenette	79.67	71.93 ± 0.12
Waterbirds	90.29	89.92 ± 0.08
MNIST-CIFAR	99.69	98.95 ± 0.02
Imagenet	72.02	69.63 ± 0.08
Table 4: Accuracy of  ${f}_{\text{proj }}$  in lazy regime
Dataset	Acc(f)	Acc(fproj)
b-Imagenette	93.09	91.77 ± 0.34
Imagenette	80.31	77.34 ± 0.21
Waterbirds	90.4	89.5 ± 0.18
MNIST-CIFAR	99.74	98.54 ± 0.00
Imagenet	72.6	72.07 ± 0.08
Table 5: Mistake diversity and class conditioned logit correlation of models trained independently (Mist-Div  $(f,f_{\mathrm{ind}})$  and CC-LogitCorr  $(f,f_{\mathrm{ind}})$  resp.) vs 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.). The results demonstrate that  $f$  and  $f_{\mathrm{proj}}$  are more diverse compared to  $f$  and  $f_{\mathrm{ind}}$ .
Dataset	Mist-Div (f, find) (↑)	Mist-Div (f, fproj) (↑)	CC-LogitCorr (f, find) (↓)	CC-LogitCorr (f, fproj) (↓)
B-Imagenette	3.87 ± 1.54	21.15 ± 1.57	99.88 ± 0.01	90.86 ± 1.08
Imagenette	6.6 ± 0.46	11.44 ± 0.65	99.31 ± 0.12	91 ± 0.59
Waterbirds	2.9 ± 0.52	14.53 ± 0.48	99.66 ± 0.04	93.81 ± 0.48
MNIST-CIFAR	0.0 ± 0.0	5.56 ± 7.89	99.76 ± 0.17	78.74 ± 2.28
Imagenet	6.97 ± 0.06	12.31 ± 0.16	99.5 ± 0.0	92.52 ± 0.01
4.4.3 Lazy regime
For the lazy regime, it turns out that the rank of first layer weight matrix remains high throughout training, as shown in Figure 3b. However, we are able to find a low dimensional projection matrix  $P$  satisfying the conditions of LD-SB (as stated in Def 1.1) as the solution to an optimization problem. More concretely, given a pretrained model  $f$  and a rank  $r$ , we obtain a projection matrix  $P$  solving:
$$
\min _ {P} \frac {1}{n} \sum_ {i = 1} ^ {n} \left(\mathcal {L} \left(f (P x ^ {(i)}), y ^ {(i)}\right) + \lambda \mathcal {L} \left(f (P ^ {\perp} x ^ {(i)}), \mathcal {U} [ L ]\right)\right)
$$
where  $\mathcal{U}[L]$  represents a uniform distribution over all the  $L$  labels,  $(x^{(1)},y^{(1)}),\dots ,(x^{(n)},y^{(n)})$  are training examples and  $\mathcal{L}(\cdot ,\cdot)$  is the cross entropy loss. We reiterate that the optimization is only over  $P$ , while the model parameters  $f$  are unchanged. In words, the above function ensures that the neural network produces correct predictions along  $P$  and uninformative predictions along  $P_{\perp}$ . Table 2 presents the results for  $P_{\perp}$  and  $P$ -LC as well as pC, while Table 4 presents  $\mathrm{Acc}(f)$  and  $\mathrm{Acc}(f_{\mathrm{proj}})$ . These results again establish LD-SB in the lazy regime.
5 Training diverse classifiers using OrthoP
Our results above motivate a natural strategy to construct diverse ensembles i.e., use  $f$  and  $f_{\mathrm{proj}}$  instead of two independently trained models. In this section, we provide two natural diversity metrics and empirically demonstrate that OrthoP leads to diverse models in practice. We also demonstrate that an ensemble of  $f$  and  $f_{\mathrm{proj}}$  has higher robustness to Gaussian noise compared to an ensemble of independently trained models.
Diversity Metrics: Given any two models  $f$  and  $\tilde{f}$ , we empirically evaluate their diversity using two metrics. The first is mistake diversity: Mist-Div  $\left(f, \tilde{f}\right) \stackrel{\mathrm{def}}{=} 1 - \frac{|\{i : f(\mathbf{x}^{(i)}) \neq y^{(i)} \& \tilde{f}(\mathbf{x}^{(i)}) \neq y^{(i)}\}|}{\min(|\{i : f(\mathbf{x}^{(i)}) \neq y^{(i)}\}|, |\{i : \tilde{f}(\mathbf{x}^{(i)}) \neq y^{(i)}\}|}$ , where we abuse notation by using  $f(x_{i})$  (resp.  $\tilde{f}(x_{i})$ ) to denote the class predicted by  $f$  (resp.  $\tilde{f}$ ) on  $x_{i}$ . Higher Mist-Div  $\left(f, \tilde{f}\right)$  means that there is very little overlap in the mistakes of  $f$  and  $\tilde{f}$ . The second is class conditioned logit correlation i.e., correlation between outputs of  $f$  and  $\tilde{f}$ , conditioned on the class. More concretely, CC-LogitCorr  $\left(f, \tilde{f}\right) = \frac{\sum_{y \in \mathcal{Y}} \operatorname{Corr}\left([f(\mathbf{x}_{i})], [\tilde{f}(\mathbf{x}_{i})]: y_{i} = y\right)}{|\mathcal{Y}|}$ , where  $\operatorname{corr}\left([f(\mathbf{x}_{i})], [\tilde{f}(\mathbf{x}_{i})]: y_{i} = y\right)$  represents the empirical correlation between the logits of  $f$  and  $\tilde{f}$  on the data points where the true label is  $y$ . Table 5 compares the diversity of two independently trained models  $(f$  and  $f_{\mathrm{ind}})$  with that of two sequentially trained models  $(f$  and  $f_{\mathrm{proj}})$ . The results demonstrate that  $f$  and  $f_{\mathrm{proj}}$  are more diverse compared to  $f$  and  $f_{\mathrm{ind}}$ . We have also compared OrthoP to another diverse training method Evading-SB (Teney et al., 2022). The results are provided in Table 7 in the Appendix. As can be seen, our results are either better or comparable to Evading-SB.
Ensembling: Figure 4 shows the variation of test accuracy with the strength of gaussian noise added to the pretrained representations of the dataset. Here, an ensemble is obtained by weighted averaging of the logits of multiple models, trained either independently ( $f_{\mathrm{ind}}$ ) or using OrthoP ( $f_{\mathrm{proj}}$ ). Moreover, we also compare our method to another diversity training method ( $f_{\mathrm{esb}}$ ) termed as
(a) Imagenette
Figure 4: Variation of test accuracy vs standard deviation of Gaussian noise added to the pretrained representations of the dataset. Here model1, ensemble-ind, ensemble-proj and ensemble-esb refer to the original model  $f$ , ensemble of  $f$  and independently trained model  $f_{\mathrm{ind}}$ , ensemble of  $f$  and  $f_{\mathrm{proj}}$  trained using OrthoP, and ensemble of  $f$  and  $f_{\mathrm{esb}}$  obtained using Evading-SB method respectively.
(b) Waterbirds-Landbirds
(c) Imagenet
Evading-SB (Teney et al., 2022). We can see that, an ensemble of  $f$  and  $f_{\mathrm{proj}}$  is much more robust as compared to an ensemble of  $f$  and  $f_{\mathrm{ind}}$ , and generally comparable to and ensemble of  $f$  and  $f_{\mathrm{esb}}$ .
6 Conclusion: Summary, Limitations and Future Directions
In this work, we propose a rigorous definition of simplicity bias, which is believed to be a key reason for their brittleness (Shah et al., 2020). In particular, we prove that 1-hidden layer networks suffer from low dimensional input dependence (LD-SB), and empirically verify this phenomenon on several real world datasets. We also propose a novel approach – OrthoP– to train diverse models, and demonstrate that an ensemble consisting of such diverse models is more robust to Gaussian noise. Extending these insights to deeper models or in the finite width setting are interesting directions for future work.
We acknowledge support from Simons Investigator Fellowship, NSF grant DMS-2134157, DARPA grant W911NF2010021, and DOE grant DE-SC0022199.
A Proofs for rich and lazy regime
A.1 Rich regime
We restate Theorem 4.1 below and prove it.
Theorem A.1. For any dataset in IFM model satisfying the conditions in Section 4.1,  $\gamma \geq 1$  and  $f(\nu, x)$  as in Eqn. (1), the distribution  $\nu^{*} = 0.5\delta_{\theta_{1}} + 0.5\delta_{\theta_{2}}$  on  $\mathcal{S}^{d+1}$  is the unique max-margin classifier satisfying Eqn. (2), where  $\theta_{1} = \left(\frac{\gamma}{\sqrt{2(1 + \gamma^{2})}}\mathbf{e}_{1}, \frac{1}{\sqrt{2(1 + \gamma^{2})}}, 1 / \sqrt{2}\right), \theta_{2} = \left(-\frac{\gamma}{\sqrt{2(1 + \gamma^{2})}}\mathbf{e}_{1}, \frac{1}{\sqrt{2(1 + \gamma^{2})}}, -1 / \sqrt{2}\right)$  and  $\mathbf{e}_{1} \stackrel{\text{def}}{=} [1, 0, \dots, 0]$  denotes first standard basis vector. In particular, this implies that if gradient flow for 1-hidden layer FCN with ReLU activation under rich initialization in the infinite width limit with cross entropy loss converges, and satisfies the technical conditions in Theorem E.1, then it converges to  $\nu^{*}$  satisfying  $f(\nu^{*}, Px_{1} + P_{\perp}x_{2}) = f(\nu^{*}, x_{1})\forall (x_{1}, y_{1}), (x_{2}, y_{2}) \in D$ , where  $P$  represents the (rank-1) projection matrix on the first coordinate.
Proof of Theorem A.1: The proof relies on showing that  $\nu^{*}$  is a max-margin classifier as in Theorem 3.1. To this end, we employ a primal-dual characterization of max-margin classifiers and construct a dual certificate that proves the optimality of margin of  $\nu^{*}$ . Chizat & Bach (2020) showed the following primal-dual characterization of maximum margin classifiers in eqn. (2):
Lemma A.2. (Chizat & Bach, 2020)  $\nu^{*}$  satisfies eqn. (2) if there exists a data distribution  $p^*$  such that the following two complementary slackness conditions hold:
$$
\operatorname {S u p p} \left(\nu^ {*}\right) \subseteq \underset {(w, b, a) \in \mathbb {S} ^ {d + 1}} {\arg \max } \mathbb {E} _ {(x, y) \sim p ^ {*}} y [ a (\phi (\langle w, x \rangle + b)) ] \quad \text {a n d} \tag {4}
$$
$$
\operatorname {S u p p} \left(p ^ {*}\right) \subseteq \underset {(x, y) \sim \mathcal {D}} {\arg \min } y \mathbb {E} _ {(w, b, a) \sim \nu^ {*}} [ a (\phi (\langle w, x \rangle + b)) ]. \tag {5}
$$
The plan is to construct a distribution  $p^*$  that satisfies the conditions of the above Lemma.
Uniqueness. Note further that for a fixed  $p^*$ ,  $\mathbb{E}_{(x,y) \sim p^*} yf(\nu, x)$  is an upper bound for the margin  $\min_{(x,y) \sim \mathcal{D}} yf(\nu, x)$  of any classifier  $\nu$ . Hence, for uniqueness, it suffices to show that  $\delta_{\theta_1}, \delta_{\theta_2}$  are the unique maximizers of the objective on the RHS of eqn. (4) and that the unique maximum margin convex combination of  $\delta_{\theta_1}, \delta_{\theta_2}$  over  $\mathcal{D}$  is  $\nu^*$ .
We first describe the support  $D$  of  $p^*$ . For  $y \in \{\pm 1\}$  we generate  $(x, y) \in D$  as
$$
\mathbf {x} _ {1} = \gamma y
$$
$$
\forall i \in 2,.., d, \mathbf {x} _ {i} = \left\{ \begin{array}{c c c} \pm 1 & \text {f o r} & y = 1 \\ 0 & \text {f o r} & y = - 1 \end{array} \right.
$$
Now for  $(x,y)\in D$  , define
$$
p ^ {*} (x, y) = \left\{ \begin{array}{l l} 0. 5 & \text {f o r} \quad y = 1 \\ 0. 5 ^ {d} & \text {f o r} \quad y = - 1 \end{array} \right. \tag {6}
$$
Note that  $p^*$  is supported on  $2^{d-1}$  positive instances and one negative instance. We begin by showing eqn. (5).
Claim A.3.  $p^*$  as in eqn. (6) satisfies eqn. (5). Further, the unique maximum margin convex combination of  $\delta_{\theta_1}, \delta_{\theta_2}$  is  $\nu^*$ .
Proof. Let us find the minimizers  $(x,y)\sim \mathcal{D}$  of  $yf(\nu ,x) = y\mathbb{E}_{(w,b,a)\sim \nu^{*}}[a(\phi (\langle w,x\rangle +b))]$  for any  $\nu = \lambda \delta_{\theta_1} + (1 - \lambda)\delta_{\theta_2},0\leq \lambda \leq 1$
$yf(\nu ,x)$  for  $(x,y)$  with  $y = -1$  (denoting  $x_{1}$  by  $-\alpha_{1}$ , where  $\alpha_{1}\geq \gamma$ ) is
$$
\begin{array}{l} y f (\nu , x) = - 1 \Big [ \lambda * \phi \left(\frac {\gamma}{\sqrt {2 (1 + \gamma^ {2})}} \mathbf {e} _ {1} ^ {\top} (- \alpha_ {1} \mathbf {e} _ {1}) + \frac {1}{\sqrt {2 (1 + \gamma^ {2})}}\right) * \frac {1}{\sqrt {2 (1 + \gamma^ {2})}} \Big ] \\ + (1 - \lambda) * \phi \left(- \frac {\gamma}{\sqrt {2 (1 + \gamma^ {2})}} \mathbf {e} _ {1} ^ {\top} (- \alpha_ {1} \mathbf {e} _ {1}) + \frac {1}{\sqrt {2 (1 + \gamma^ {2})}}\right) * \frac {- 1}{\sqrt {2}} ], \\ \end{array}
$$
and for  $(x,y)$  with  $y = 1$  (denoting  $x_{1}$  by  $\alpha_{2}$ , where  $\alpha_{2} \geq \gamma$ ) is
$$
\begin{array}{l} y f (\nu , x) = 1 \left[ \lambda * \phi \left(\frac {\gamma}{\sqrt {2 (1 + \gamma^ {2})}} \mathbf {e} _ {1} ^ {\top} (\alpha_ {2} \mathbf {e} _ {1}) + \frac {1}{\sqrt {2 (1 + \gamma^ {2})}}\right) * \frac {1}{\sqrt {2}} \right. \\ + (1 - \lambda) * \phi \left(- \frac {\gamma}{\sqrt {2 (1 + \gamma^ {2})}} \mathbf {e} _ {1} ^ {\top} (\alpha_ {2} \mathbf {e} _ {1}) + \frac {1}{\sqrt {2 (1 + \gamma^ {2})}}\right) * \left. \frac {- 1}{\sqrt {2}} \right]. \\ \end{array}
$$
As  $\gamma \geq 1$ , the expressions above equal  $\lambda \frac{\sqrt{\gamma\alpha_1 + 1}}{2}$  and  $(1 - \lambda)\frac{\sqrt{\gamma\alpha_2 + 1}}{2}$  respectively, and hence are minimized at  $\alpha_{1} = \alpha_{2} = \gamma$ . Hence, the margin of  $\nu$  is  $\min (\lambda ,1 - \lambda)\frac{\sqrt{1 + \gamma^2}}{2}$  which is uniquely maximized at  $\lambda = 1 / 2$ . Further for  $\lambda = 1 / 2$ , all points in  $D$  have the same value of  $yf(\nu ,x)$ .
In the rest of the proof we show eqn. (4). Let us denote by  $g(w, b, a) \coloneqq \mathbb{E}_{(x, y) \sim p^*} y[a(\phi(\langle w, x \rangle + b))]$ . We show that  $\delta_{\theta_1}, \delta_{\theta_2}$  are the only maximizers of  $g(w, b, a)$  over  $\mathbb{S}^{d+1}$ .
We first find  $g(\theta_1), g(\theta_2)$ .
$$
\begin{array}{l} g (\theta_ {1}) = \operatorname * {P r} (y = 1) \cdot 1 \cdot \frac {1}{\sqrt {2}} \cdot \phi \left(\frac {\gamma}{\sqrt {2 (1 + \gamma^ {2})}} \mathbf {e} _ {1} ^ {T} (\gamma \mathbf {e} _ {1}) + \frac {1}{\sqrt {2 (1 + \gamma^ {2})}}\right) \\ + \operatorname * {P r} (y = - 1) \cdot - 1 \cdot \frac {1}{\sqrt {2}} \cdot \phi \left(\frac {\gamma}{\sqrt {2 (1 + \gamma^ {2})}} \mathbf {e} _ {1} ^ {T} (- \gamma \mathbf {e} _ {1}) + \frac {1}{\sqrt {2 (1 + \gamma^ {2})}}\right) = \frac {\sqrt {\gamma^ {2} + 1}}{4}, \\ \end{array}
$$
where the first term is because  $w_{2}, w_{3}, \ldots, w_{d}$  are zero for  $\theta_{1}$ . Similarly,  $g(\theta_{2}) = \frac{\sqrt{\gamma^{2} + 1}}{4}$ . We now show that  $g(w, a, b) < \frac{\sqrt{\gamma^{2} + 1}}{4}$  for  $(w, a, b) \notin \{\theta_{1}, \theta_{2}\}$ .
We begin by showing the following simple but useful claim.
Claim A.4. All maximizers of  $g(w, b, a)$  over  $\mathbb{S}^{d + 1}$  satisfy  $|a| = 1 / \sqrt{2}$ .
Proof. The proof essentially follows from the 1-homogeneity of the ReLU function  $\phi$  and separability of  $g(w,b,a)$ . Note that  $g(w,b,a) = \sqrt{\|w\|^2 + b^2} a \cdot g(w',b',1)$  where  $\| w' \|^2 + b^2 = 1$ . Maximizing  $g(w,b,a)$  is equivalent to maximizing  $g(w',b',1)$  over  $\mathbb{S}^d$  and  $a \sqrt{\|w\|^2 + b^2}$  over  $\mathbb{S}^{d+1}$  respectively. The second of these has its unique maximum at  $|a| = 1/\sqrt{2}$ , completing the proof.
Now express  $g(w, b, a)$  as
$$
\begin{array}{l} g (w, b, a) = a (\Pr (y = 1) \mathbb {E} [ \phi (w ^ {T} x + b) | y = 1 ] - \Pr (y = - 1) \mathbb {E} [ \phi (w ^ {T} x + b) | y = - 1 ]) \\ = \frac {a}{2} \left(\mathbb {E} _ {\sigma} \left[ \phi \left(\gamma w _ {1} + b + \sum_ {i = 2} ^ {d} \sigma_ {i} w _ {i}\right) \right] - \phi (b - \gamma w _ {1})\right), \tag {7} \\ \end{array}
$$
where  $\sigma_{i}$  are independent Rademacher random variables. We have two cases on  $a$ :
Case 1:  $a = 1 / \sqrt{2}$ . By eqn. (7) we have
$$
g (w, b, 1 / \sqrt {2}) \leq \frac {1}{2 \sqrt {2}} \mathbb {E} _ {\sigma} \left[ \phi (\gamma w _ {1} + b + \sum_ {i = 2} ^ {d} \sigma_ {i} w _ {i}) \right].
$$
To simplify the above, define the random variable  $X = \sum_{i=2}^{d} \sigma_i w_i$  and denote  $\gamma w_1 + b$  by  $\alpha$ . Note that  $|\alpha| = |\gamma w_1 + b| \leq \sqrt{\frac{\gamma^2 + 1}{2}}$  which follows from  $\|w\|^2 + b^2 = 1/2$ . The expectation in the last expression above becomes
$$
\begin{array}{l} \mathbb {E} [ \phi (X + \alpha) ] = \mathbb {E} [ (X + \alpha) \mathbf {1} \{X + \alpha \geq 0 \} ] = \mathbb {E} [ X \mathbf {1} \{X \geq - \alpha \} ] + \alpha \Pr (X \geq - \alpha) \\ = \mathbb {E} [ X \mathbf {1} \{X \geq \alpha \} ] + \alpha (1 - \Pr (X \geq \alpha)) \leq \mathbb {E} [ X \mathbf {1} \{X \geq \alpha \} ] + \alpha , \\ \end{array}
$$
where the last equality follows from symmetry of  $X$ . Note that  $\operatorname{Var}(X) = \sum_{i=2}^{d} w_i^2$  which is at most  $\frac{1}{2} - \frac{\alpha^2}{1 + \gamma^2}$  (using  $\gamma w_1 + b = \alpha$  and  $\|w\|^2 + b^2 = 1/2$ ). Using A.5 to upper bound  $\mathbb{E}[X \mathbf{1}\{X \geq \alpha\}]$  we have
$$
\mathbb {E} [ \phi (X + \alpha) ] \leq \alpha + \sqrt {\frac {1}{2} \min \left(\frac {1}{2} , \frac {\frac {1}{2} - \frac {\alpha^ {2}}{1 + \gamma^ {2}}}{2 \alpha^ {2}}\right) \left(\frac {1}{2} - \frac {\alpha^ {2}}{1 + \gamma^ {2}}\right)}.
$$