SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Moreover, it is worth noting that Assumption 2.2 implies that the data is w.h.p. linearly separable (see Lemma 2.1 below, and a proof in Appendix B). However, in this work we consider learning using overparameterized ReLU networks, and it is not obvious a priori that gradient methods do not harmfully overfit in this case. Indeed, it has been shown that ReLU networks trained by gradient descent can interpolate training data and fail to generalize well in some distributional settings [Kou+23].
Lemma 2.1. Let  $\mathbf{u} = \sum_{q\in [k]}y^{(q)}\pmb{\mu}^{(q)}$ . Then, with probability at least  $1 - 2d^{1 - \ln (d) / 2} = 1 - o_d(1)$  over  $(\mathbf{x},y)\sim \mathcal{D}_{clusters}$ , we have  $y = \mathrm{sign}(\mathbf{u}^{\top}\mathbf{x})$ .
3 Generalization
In this section, we show that under our assumptions on the distribution  $\mathcal{D}_{\mathrm{clusters}}$ , gradient flow does not harmfully overfit. Namely, even if the learned network is highly overparameterized, the implicit bias of gradient flow guarantees convergence to a solution that generalizes well. Moreover, we show that the sample complexity is optimal. The main result of this section is stated in the following theorem:
Theorem 3.1. Let  $\epsilon, \delta \in (0,1)$ . Let  $\mathcal{S} = \{(\mathbf{x}_i, y_i)\}_{i=1}^n \subseteq \mathbb{R}^d \times \{-1,1\}$  be a training set drawn i.i.d. from the distribution  $\mathcal{D}_{\text{clusters}}$ , where  $n \geq k \ln^2(d)$ . Let  $\mathcal{N}_{\theta}$  be a depth-2 ReLU network such that  $\pmb{\theta} = [\mathbf{w}_1, \dots, \mathbf{w}_m, \mathbf{b}, \mathbf{v}]$  is a KKT point of Problem (2). Provided  $d$  is sufficiently large such that  $\delta^{-1} \leq \frac{1}{3} d^{\ln(d)-1}$  and  $n \leq \min \left\{\sqrt{\frac{\delta}{3}} \cdot e^{d/32}, \frac{\sqrt{\delta}}{3} \cdot d^{\ln(d)/4}, \frac{\epsilon}{4} \cdot d^{\ln(d)/2}\right\}$ , then with probability at least  $1-\delta$  over  $\mathcal{S}$ , we have
$$
\Pr_ {(\mathbf {x}, y) \sim \mathcal {D} _ {\text {c l u s t e r s}}} \left[ y \mathcal {N} _ {\boldsymbol {\theta}} (\mathbf {x}) \leq 0 \right] \leq \epsilon .
$$
The sample complexity requirement in Theorem 3.1 is  $n = \tilde{\Omega}(k)$ . Essentially, it requires that the dataset  $S$  will include at least one example from each cluster. Clearly, any learning algorithm cannot perform well on unseen clusters. Hence the sample complexity requirement in the theorem is tight (up to log factors).
The assumptions in Theorem 3.1 include upper bounds on  $\delta^{-1}$  and  $n$ . Note that the expressions in these upper bounds are super-polynomial in  $d$ , and in particular if  $n, \delta^{-1}, \epsilon^{-1} = \mathrm{poly}(d)$ , then these assumptions hold for a sufficiently large  $d$ . Admittedly, enforcing an upper bound on the training dataset's size is uncommon in generalization results. However, if  $n$  is exponential in  $d$ , it is not hard to see that there will be clusters which have both positive and negative examples within radius  $\sigma$  of the cluster center, essentially introducing a form of label noise to the problem. Since KKT points of Problem (2) interpolate the training data, this would imply that the network has interpolated training data with label noise—in other words, it has 'overfit' to noise. Understanding the generalization behavior of interpolating neural networks in the presence of label noise is a very technically challenging problem for which much is unknown, especially if one seeks to understand this by only relying upon the properties of KKT conditions for margin maximization. It is noteworthy that all existing non-vacuous generalization bounds for interpolating nonlinear neural networks in the presence of label noise require  $n < d$  [FCB22; Cao+22; XG23; Fre+23a; Kou+23].
Combining Theorem 3.1 with Theorem 2.1, we conclude that w.h.p. over a training dataset of size  $n \geq k\ln^2(d)$  (and under some additional mild requirements), if gradient flow reaches empirical loss smaller than  $\frac{1}{n}$ , then it converges in direction to a neural network that generalizes well. This result is width-independent, thus, it holds irrespective of the network width. Specifically, even if the network is highly overparameterized, the implicit bias of gradient flow prevents harmful overfitting. Moreover, the result does not depend directly on the initialization of gradient flow. That is, it holds whenever gradient flow reaches small empirical loss after some finite time. Thus, by relying on the KKT conditions of the max-margin problem instead of analyzing the full gradient flow trajectory, we can prove generalization without the need to prove convergence.
3.1 Proof idea
The proof of Theorem 3.1 is given in Appendix A. Here we discuss the high-level approach. Let  $\pmb{\theta} = [\mathbf{w}_1, \dots, \mathbf{w}_m, \mathbf{b}, \mathbf{v}]$  be a KKT point of Problem (2). Thus, we have  $\mathcal{N}_{\pmb{\theta}}(\mathbf{x}) = \sum_{j \in [m]} v_j \phi(\mathbf{w}_j^\top \mathbf{x} + b_j)$ . Since  $\pmb{\theta}$  satisfies the KKT conditions of Problem (2), then there are  $\lambda_1, \dots, \lambda_n$  such that for every  $j \in [m]$  we have
$$
\mathbf {w} _ {j} = \sum_ {i \in [ n ]} \lambda_ {i} \nabla_ {\mathbf {w} _ {j}} \left(y _ {i} \mathcal {N} _ {\boldsymbol {\theta}} (\mathbf {x} _ {i})\right) = \sum_ {i \in [ n ]} \lambda_ {i} y _ {i} v _ {j} \phi_ {i, j} ^ {\prime} \mathbf {x} _ {i}, \tag {3}
$$
where  $\phi_{i,j}^{\prime}$  is a subgradient of  $\phi$  at  $\mathbf{w}_j^\top \mathbf{x}_i + b_j$ , i.e., if  $\mathbf{w}_j^\top \mathbf{x}_i + b_j \neq 0$  then  $\phi_{i,j}^{\prime} = \mathbb{1}[\mathbf{w}_j^\top \mathbf{x}_i + b_j \geq 0]$ , and otherwise  $\phi_{i,j}^{\prime}$  is some value in  $[0,1]$ . Also we have  $\lambda_i \geq 0$  for all  $i$ , and  $\lambda_i = 0$  if  $y_i \mathcal{N}_{\pmb{\theta}}(\mathbf{x}_i) \neq 1$ . Likewise, we have
$$
b _ {j} = \sum_ {i \in [ n ]} \lambda_ {i} \nabla_ {b _ {j}} \left(y _ {i} \mathcal {N} _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {i}\right)\right) = \sum_ {i \in [ n ]} \lambda_ {i} y _ {i} v _ {j} \phi_ {i, j} ^ {\prime}. \tag {4}
$$
In the proof, using a careful analysis of Eq. (3) and (4) we show that w.h.p.  $\mathcal{N}_{\theta}$  classifies correctly a fresh example. More precisely, the main argument can be described as follows. We denote  $J \coloneqq [m]$ ,  $J_{+} \coloneqq \{j \in J : v_{j} > 0\}$ , and  $J_{-} \coloneqq \{j \in J : v_{j} < 0\}$ . Moreover, we denote  $I \coloneqq [n]$  and  $Q \coloneqq [k]$ . For  $q \in Q$  we denote  $I^{(q)} = \{i \in I : \mathbf{x}_{i}$  is in cluster  $q\}$ . Consider the network's output for an input  $\mathbf{x}$  from cluster  $r \in Q$  with  $y^{(r)} = 1$ . Since  $\mathcal{N}_{\theta}(\mathbf{x}) = \sum_{j \in J_{+}} v_{j} \phi(\mathbf{w}_{j}^{\top} \mathbf{x} + b_{j}) + \sum_{j \in J_{-}} v_{j} \phi(\mathbf{w}_{j}^{\top} \mathbf{x} + b_{j})$  and  $\phi(z) \geq z$ , we have
$$
\mathcal {N} _ {\boldsymbol {\theta}} (\mathbf {x}) \geq \sum_ {j \in J _ {+}} v _ {j} \left(\mathbf {w} _ {j} ^ {\top} \mathbf {x} + b _ {j}\right) + \sum_ {j \in J _ {-}} v _ {j} \phi \left(\mathbf {w} _ {j} ^ {\top} \mathbf {x} + b _ {j}\right). \tag {5}
$$
This suggests the following possibility: if we can ensure that  $\sum_{j\in J_{+}}v_{j}(\mathbf{w}_{j}^{\top}\mathbf{x} + b_{j})$  is large and positive while  $\sum_{j\in J_{-}}v_{j}\phi (\mathbf{w}_{j}^{\top}\mathbf{x} + b_{j})$  is not too negative, then the network will accurately classify the example  $\mathbf{x}$ . Using Eq. (3) and (4) and that  $y^{(r)} = 1$  (so  $y_{i} = 1$  for  $i\in I^{(r)})$ , the first term in the above decomposition is equal to
$$
\begin{array}{l} \sum_ {j \in J _ {+}} v _ {j} \left(\mathbf {w} _ {j} ^ {\top} \mathbf {x} + b _ {j}\right) = \sum_ {j \in J _ {+}} v _ {j} \left[ \sum_ {i \in I} \lambda_ {i} y _ {i} v _ {j} \phi_ {i, j} ^ {\prime} \left(\mathbf {x} _ {i} ^ {\top} \mathbf {x} + 1\right) \right] \\ = \sum_ {j \in J _ {+}} \left[ \left(\sum_ {i \in I ^ {(r)}} \lambda_ {i} v _ {j} ^ {2} \phi_ {i, j} ^ {\prime} (\mathbf {x} _ {i} ^ {\top} \mathbf {x} + 1)\right) + \sum_ {q \in Q \setminus \{r \}} \sum_ {i \in I ^ {(q)}} \lambda_ {i} y _ {i} v _ {j} ^ {2} \phi_ {i, j} ^ {\prime} (\mathbf {x} _ {i} ^ {\top} \mathbf {x} + 1) \right] \\ \geq \left(\sum_ {i \in I ^ {(r)}} \sum_ {j \in J _ {+}} \lambda_ {i} v _ {j} ^ {2} \phi_ {i, j} ^ {\prime} (\mathbf {x} _ {i} ^ {\top} \mathbf {x} + 1)\right) - \sum_ {q \in Q \setminus \{r \}} \sum_ {i \in I ^ {(q)}} \sum_ {j \in J _ {+}} \lambda_ {i} v _ {j} ^ {2} \phi_ {i, j} ^ {\prime} | \mathbf {x} _ {i} ^ {\top} \mathbf {x} + 1 |. \\ \end{array}
$$
Since  $\mathbf{x}$  comes from cluster  $r$  and the clusters are nearly orthogonal, the pairwise correlations  $\mathbf{x}_i^\top \mathbf{x}$  will be large and positive when  $i \in I^{(r)}$  but will be small in magnitude when  $i \in I^{(q)}$  for  $q \neq r$ . Thus, we can hope that this term will be large and positive if we can show that the quantity  $\sum_{i \in I^{(r)}} \sum_{j \in J^+} \lambda_i v_j^2 \phi_{i,j}'$  is not too small relative to the quantity  $\sum_{q \in Q \setminus \{r\}} \sum_{i \in I^{(q)}} \sum_{j \in J_+} \lambda_i v_j^2 \phi_{i,j}'$ . By similar arguments, in order to show the second term in Eq. (5) is not too negative, we need to understand how the quantity  $\sum_{i \in I^{(q)}} \sum_{j \in J_-} \lambda_i v_j^2 \phi_{i,j}'$  varies across different clusters  $q \in Q$ . Hence, in the proof we analyze how the quantities  $\sum_{i \in I^{(q)}} \sum_{j \in J_+} \lambda_i v_j^2 \phi_{i,j}'$ ,  $\sum_{i \in I^{(q)}} \sum_{j \in J_-} \lambda_i v_j^2 \phi_{i,j}'$  relate to each other for different clusters  $q \in Q$ , and show that these quantities are all of the same order. Then, we conclude that w.h.p.  $\mathbf{x}$  is classified correctly.
4 Robustness
We begin by introducing the definition of  $R(\cdot)$ -robustness.
Definition 4.1. Given some function  $R(\cdot)$ , we say that a neural network  $\mathcal{N}_{\theta}$  is  $R(d)$ -robust w.r.t. a distribution  $\mathcal{D}_{\mathbf{x}}$  over  $\mathbb{R}^d$  if for every  $r = o(R(d))$ , with probability  $1 - o_d(1)$  over  $\mathbf{x} \sim \mathcal{D}_{\mathbf{x}}$ , for every  $\mathbf{x}' \in \mathbb{R}^d$  with  $\| \mathbf{x} - \mathbf{x}' \| \leq r$  we have  $\mathrm{sign}(\mathcal{N}_{\theta}(\mathbf{x}')) = \mathrm{sign}(\mathcal{N}_{\theta}(\mathbf{x}))$ .
Thus, a neural net  $\mathcal{N}_{\theta}$  is  $R(d)$ -robust if changing the label of an example cannot be done with a perturbation of size  $o(R(d))$ . Note that we consider here  $\ell_2$  perturbations.
For the distribution  $\mathcal{D}_{\mathrm{clusters}}$  under consideration, it is straightforward to show that classifiers cannot be  $R(d)$ -robust if  $R(d) = \omega (\sqrt{d})$ : since the distance between examples in different clusters is w.h.p.
$\mathcal{O}(\sqrt{d})$ , it is clearly possible to flip the sign of an example with a perturbation of size  $\mathcal{O}(\sqrt{d})$ . In particular, the best we can hope for is  $\sqrt{d}$ -robustness. In the following theorem, we show that there exist two-layer ReLU networks which can both achieve small test error and the optimal level of  $\sqrt{d}$ -robustness.
Theorem 4.1. For every  $r \geq k$ , there exists a depth-2 ReLU network  $\mathcal{N} : \mathbb{R}^d \to \mathbb{R}$  of width  $r$  such that for  $(\mathbf{x}, y) \sim \mathcal{D}_{\text{clusters}}$ , with probability at least  $1 - d^{-\omega_d(1)}$  we have  $y\mathcal{N}(\mathbf{x}) \geq 1$ , and flipping the sign of the output requires a perturbation of size larger than  $\frac{\sqrt{d}}{8}$  (for a sufficiently large  $d$ ). Thus,  $\mathcal{N}$  classifies the data correctly w.h.p., and it is  $\sqrt{d}$ -robust w.r.t.  $\mathcal{D}_{\mathbf{x}}$ .
Thus, we see that  $\sqrt{d}$ -robust networks exist. In the following theorem, we show that the implicit bias of gradient flow constrains the level of robustness of trained networks whenever the number of clusters  $k$  is large.
Theorem 4.2. Let  $\epsilon, \delta \in (0,1)$ . Let  $\mathcal{S} = \{(\mathbf{x}_i, y_i)\}_{i=1}^n \subseteq \mathbb{R}^d \times \{-1,1\}$  be a training set drawn i.i.d. from the distribution  $\mathcal{D}_{\text{clusters}}$ , where  $n \geq k \ln^2(d)$ . We denote  $Q_+ = \{q \in [k] : y^{(q)} = 1\}$  and  $Q_- = \{q \in [k] : y^{(q)} = -1\}$ , and assume that  $\min \left\{\frac{|Q_+|}{k}, \frac{|Q_-|}{k}\right\} \geq c$  for some  $c > 0$ . Let  $\mathcal{N}_{\boldsymbol{\theta}}$  be a depth-2 ReLU network such that  $\boldsymbol{\theta} = [\mathbf{w}_1, \ldots, \mathbf{w}_m, \mathbf{b}, \mathbf{v}]$  is a KKT point of Problem (2). Provided  $d$  is sufficiently large such that  $\delta^{-1} \leq \frac{1}{3} d^{\ln(d) - 1}$  and  $n \leq \min \left\{\sqrt{\frac{\delta}{3}} \cdot e^{d/32}, \frac{\sqrt{\delta}}{3} \cdot d^{\ln(d)/4}, \frac{\epsilon}{4} \cdot d^{\ln(d)/2}\right\}$ , with probability at least  $1 - \delta$  over  $\mathcal{S}$ , there is a vector  $\mathbf{z} = \eta \cdot \sum_{j \in [k]} y^{(j)} \boldsymbol{\mu}^{(j)}$  with  $\eta > 0$  and  $\|\mathbf{z}\| \leq O\left(\sqrt{\frac{d}{c^2k}}\right)$ , such that
$$
\operatorname * {P r} _ {(\mathbf {x}, y) \sim \mathcal {D} _ {\text {c l u s t e r s}}} \left[ \operatorname {s i g n} \left(\mathcal {N} _ {\boldsymbol {\theta}} (\mathbf {x})\right) \neq \operatorname {s i g n} \left(\mathcal {N} _ {\boldsymbol {\theta}} (\mathbf {x} - y \mathbf {z})\right) \right] \geq 1 - \epsilon .
$$
Note that the expressions in the upper bounds on  $n$  and  $\delta^{-1}$  are super-polynomial in  $d$ , and hence these requirements are mild (e.g., they hold for a sufficiently large  $d$  when  $n, \delta^{-1}, \epsilon^{-1} = \mathrm{poly}(d)$ ). As we mentioned in the discussion following Theorem 3.1, we believe removing the requirement for an upper bound on  $n$  would be highly nontrivial.
Theorem 4.2 implies that if  $c^2 k = \omega_d(1)$ , then w.h.p. over the training dataset, every KKT point of Problem (2) is not  $\sqrt{d}$ -robust. Specifically, if  $c$  is constant, namely, at least a constant fraction of the clusters have positive labels and a constant fraction of the clusters have negative labels, then the network is not  $\sqrt{d}$ -robust if  $k = \omega_d(1)$ . Recall that by Theorem 3.1, we also have w.h.p. that every KKT point generalizes well. Overall, combining Theorems 2.1, 3.1, 4.1, and 4.2, we conclude that for  $c^2 k = \omega_d(1)$ , w.h.p. over a training dataset of size  $n \geq k\ln^2(d)$ , if gradient flow reaches empirical loss smaller than  $\frac{1}{n}$ , then it converges in direction to a neural network that generalizes well but is not  $\sqrt{d}$ -robust, even though there exist  $\sqrt{d}$ -robust networks that generalize well. Thus, in our setting, there is bias towards solutions that generalize well but are non-robust.
Example 2. Consider the setting from the first item of Example 1. Thus, the cluster means satisfy  $|\langle \pmb{\mu}^{(i)},\pmb{\mu}^{(j)}\rangle| = \tilde{\mathcal{O}}(\sqrt{d})$  for every  $i \neq j$ , and we have  $\sigma = 1$  and  $k = \tilde{\Theta}(\sqrt{d})$ . Suppose that  $c = \Theta(1)$ , namely, there is at least a constant fraction of clusters with each label  $\{-1,1\}$ . Then, the adversarial perturbation  $\mathbf{z}$  from Theorem 4.2 satisfies  $\| \mathbf{z}\| = \mathcal{O}\left(\sqrt{\frac{d}{k}}\right) = \tilde{\mathcal{O}}\left(d^{1/4}\right) = o(\sqrt{d})$ .
Similarly to our discussion after Theorem 3.1, we note that Theorem 4.2 is width-independent, i.e., it holds irrespective of the network width. It implies that we cannot hope to obtain a robust solution by choosing an appropriate width for the trained network. As we discussed in the related work section, Bubeck, Li, and Nagaraj [BLN21] and Bubeck and Sellke [BS21] considered the expressive power of neural networks, and showed that overparameterization might be necessary for robustness. By Theorem 4.2, even when the network is overparameterized, the implicit bias of the optimization method can prevent convergence to robust solutions. Moreover, our result does not depend directly on the initialization of gradient flow. Recall that by Theorem 2.1 if gradient flow reaches small empirical loss then it converges in direction to a KKT point of Problem (2). Hence our result holds whenever gradient flow reaches a small empirical loss.
Note that in Theorem 4.2, the adversarial perturbation does not depend on the input (up to sign). It corresponds to the well-known empirical phenomenon of universal adversarial perturbations, where one can find a single perturbation that simultaneously flips the label of many inputs (cf. [Moo+17;
Zha+21]). Moreover, the same perturbation applies to all depth-2 networks to which gradient flow might converge (i.e., all KKT points). It corresponds to the well-known empirical phenomenon of transferability in adversarial examples, where one can find perturbations that simultaneously flip the labels of many different trained networks (cf. [Liu+17; AM18]).
It is worth noting that Theorems 3.1 and 4.2 demonstrate that trained neural networks exhibit different properties than the 1-nearest-neighbour learning rule, irrespective of the number of parameters in the network. For example, consider the case where  $\sigma = \frac{1}{\sqrt{d}}$ , namely, the examples of each cluster are concentrated within a ball of radius  $O(1)$  around its mean. Then, the distance between every pair of points from the same cluster is  $O(1)$ , and the distance between points from different clusters is  $\Omega(\sqrt{d})$ . In this setting, both the 1-nearest-neighbour classifier and the trained neural network will classify a fresh example correctly w.h.p., but in the 1-nearest-neighbour classifier flipping the output's sign will require a perturbation of size  $\Omega(\sqrt{d})$ , while in the neural network a much smaller perturbation will suffice.
Finally, we remark that in the limit  $\sigma \to 0$ , we get a distribution supported on  $\pmb{\mu}^{(1)},\dots,\pmb{\mu}^{(k)}$ . Then, a training dataset of size  $n\geq k\ln^2 (d)$  will contain w.h.p. all examples in the support, and hence robustness w.r.t. test data is equivalent to robustness w.r.t. the training data. In this case, we recover the results of Vardi, Yehudai, and Shamir [VYS22] which characterized the non-robustness of KKT points of ReLU networks trained on nearly orthogonal training data. In particular, our Theorem 4.2 is a strict generalization of their Theorem 4.1.
4.1 Proof ideas
Here we discuss the main ideas in the proofs of Theorem 4.1 and 4.2 (see Appendices C and D for the formal proofs).
The proof of Theorem 4.1 follows by the following simple construction. The robust network includes  $k$  neurons, each corresponding to a single cluster. That is, we have  $\mathcal{N}(\mathbf{x}) = \sum_{j=1}^{k} v_j \sigma(\mathbf{w}^\top \mathbf{x} + b_j)$ , where  $v_j = y^{(j)}$ ,  $\mathbf{w}_j = \frac{4\mu^{(j)}}{d}$ , and  $b_j = -2$ . Note that the  $j$ -th neuron points at the direction of the  $j$ -th cluster and has a negative bias term, such that the neuron is active on points from the  $j$ -th cluster, and inactive on points from the other clusters. Then, given a fresh example  $(\mathbf{x}, y) \sim \mathcal{D}_{\mathrm{clusters}}$ , we show that the network classifies it correctly w.h.p. with margin at least 1. Also, there is w.h.p exactly one neuron that is active on  $\mathbf{x}$ , and hence the gradient of the network w.r.t. the input is affected only by this neuron and is of size  $\mathcal{O}(1 / \sqrt{d})$ . Therefore, we need a perturbation of size  $\Omega(\sqrt{d})$  in order to flip the output's sign.
The intuition for Theorem 4.2 can be described as follows. Recall that in our construction of a robust network above, an example  $(\mathbf{x},y)\sim \mathcal{D}_{\mathrm{clusters}}$  is w.h.p. in an active region of exactly one neuron, and hence in the neighborhood of  $\mathbf{x}$  the output of the network is sensitive only to perturbations in the direction of that neuron. Now, consider the linear model  $\mathbf{x}\mapsto \mathbf{w}^{\top}\mathbf{x}$ , where  $\mathbf{w} = \sum_{q = 1}^{k}\frac{1}{d} y^{(q)}\pmb{\mu}^{(q)}$ . It is not hard to verify that for  $(\mathbf{x},y)\sim \mathcal{D}_{\mathrm{clusters}}$  we have w.h.p. that  $0 < y\mathbf{w}^{\top}\mathbf{x}\leq \mathcal{O}(1)$ . Moreover, the gradient of this linear predictor is of size  $\| \mathbf{w}\| = \Omega (\sqrt{k / d})$ . Hence, we can flip the output's sign with a perturbation of size  $\mathcal{O}(\sqrt{d / k})$ . Thus, the linear classifier is non-robust if  $k = \omega_d(1)$ . Intuitively, the difference between our robust ReLU network and the non-robust linear classifier is the fact that in the neighborhood of  $\mathbf{x}$  the robust network is sensitive only to perturbations in the direction of one cluster, while the linear classifier is sensitive to perturbations in the directions of all  $k$  clusters. In the proof, we analyze ReLU networks which are KKT points of Problem (2), and show that although these ReLU networks are non-linear, they are still sensitive to perturbations in the directions of all  $k$  clusters, similarly to the above linear classifier. The formal proof follows by a careful analysis of the KKT conditions of Problem (2), given in Eq. (3) and (4).
We remark that in the proof of Theorem 4.2 we use some technical ideas from Vardi, Yehudai, and Shamir [VYS22]. However, there are significant differences between the two settings. For example, they assume that the training data are nearly orthogonal, which only holds when the dimension is large relative to the number of samples; thus, it is unclear whether the existence of small adversarial perturbations in their setting is due to the high-dimensionality of the data or if a similar phenomenon exists in the more common  $n > d$  setting. At a more technical level, their proof relies on showing that
in a KKT point all inputs must lie exactly on the margin, while in our setting they are not guaranteed to lie exactly on the margin.
In this paper, we considered clustered data, and showed that gradient flow in two-layer ReLU networks does not harmfully overfit, but also hinders robustness. Our results follow by analyzing the KKT points of the max-margin problem in parameter space. In our distributional setting, the clusters are well-separated, and hence there exist robust classifiers, which allows us to consider the effect of the implicit bias of gradient flow on both generalization and robustness. Understanding generalization and robustness in additional data distributions and neural network architectures is a challenging but important question. As a possible next step, it would be interesting to study whether the approach used in this paper can be extended to the following data distributions:
First, our assumption on the data distribution (Assumption 2.2) implies that the number of clusters cannot be too large, and as a result the data is linearly separable (Lemma 2.1). We conjecture that our results hold even for a significantly larger number of clusters, such that the data is not linearly separable.
Second, it would be interesting to understand whether our generalization result holds for linearly separable data distributions that are not clustered. That is, given a distribution that is linearly separable with some margin  $\gamma > 0$  and a training dataset that is large enough to allow learning with a max-margin linear classifier, are there KKT points of the max-margin problem for two-layer ReLU networks that do not generalize well? In other words, do ReLU networks that satisfy the KKT conditions generalize at least as well as max-margin linear classifiers?
SF, GV, PB, and NS acknowledge the support of the NSF and the Simons Foundation for the Collaboration on the Theoretical Foundations of Deep Learning through awards DMS-2031883 and #814639, and of the NSF through grant DMS-2023505.