SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Spencer Frei*
sfrei@ucdavis.edu
Gal Vardi*
Peter L. Bartlett
UC Berkeley and Google DeepMind
peter@berkeley.edu
In this work, we study the implications of the implicit bias of gradient flow on generalization and adversarial robustness in ReLU networks. We focus on a setting where the data consists of clusters and the correlations between cluster means are small, and show that in two-layer ReLU networks gradient flow is biased towards solutions that generalize well, but are vulnerable to adversarial examples. Our results hold even in cases where the network is highly overparameterized. Despite the potential for harmful overfitting in such settings, we prove that the implicit bias of gradient flow prevents it. However, the implicit bias also leads to non-robust solutions (susceptible to small adversarial  $\ell_2$ -perturbations), even though robust networks that fit the data exist.
A central question in the theory of deep learning is how neural networks can generalize even when trained without any explicit regularization, and when there are more learnable parameters than training examples. In such optimization problems there are many solutions that label the training data correctly, and gradient descent seems to prefer solutions that generalize well [Zha+17]. Thus, it is believed that gradient descent induces an implicit bias towards solutions which enjoy favorable properties [Ney+17]. Characterizing this bias in various settings has been a subject of extensive research in recent years, but it is still not well understood when the implicit bias provably implies generalization in non-linear neural networks.
An additional intriguing phenomenon in deep learning is the abundance of adversarial examples in trained neural networks. In a seminal paper, Szegedy et al. [Sze+14] observed that deep networks are extremely vulnerable to adversarial examples, namely, very small perturbations to the inputs can significantly change the predictions. This phenomenon has attracted considerable interest, and various attacks (e.g., [GSS15; CW17; Pap+17; ACW18; CW18; Wu+20]) and defenses (e.g., [Pap+16; KGB17; Mad+18; WK18; CH20; WRK20]) were developed. However, the fundamental principles underlying the existence of adversarial examples are still unclear, and it is believed that for most tasks where trained neural networks suffer from a vulnerability to adversarial attacks, there should exist other neural networks which can be robust to such attacks. This is suggestive of the possible role of the optimization algorithms used to train neural networks in the existence of adversarial examples.
In this work, we study the implications of the implicit bias for generalization and robustness in ReLU networks, in a setting where the data consists of clusters (i.e., Gaussian mixture model) and the
correlations between cluster means are small. We show that in two-layer ReLU networks trained with the logistic loss or the exponential loss, gradient flow is biased towards solutions that generalize well, albeit they are non-robust. Our results are independent of the network width, and hence they hold even where the network has significantly more parameters than training examples. In such an overparameterized setting, one might expect harmful overfitting to occur, but we prove that the implicit bias of gradient flow prevents it. On the flip side, in our setting the distances between clusters are large, and thus one might hope that gradient flow will converge to a robust network. However, we show that the implicit bias leads to non-robust solutions.
Our results rely on known properties of the implicit bias in two-layer ReLU networks trained with the logistic or the exponential loss, which were shown by Lyu and Li [LL20] and Ji and Telgarsky [JT20]. They proved that if gradient flow in homogeneous models (which include two-layer ReLU networks) with such losses reaches a small training loss, then it converges (in direction) to a KKT point of the maximum-margin problem in parameter space. We show that in clustered data distributions, with high probability over the training dataset, every network that satisfies the KKT conditions of the maximum-margin problem generalizes well but is non-robust. Thus, instead of analyzing the trajectory of gradient flow directly in the complex setting of training two-layer ReLU networks, we demonstrate that investigating the KKT points is a powerful tool for understanding generalization and robustness. We emphasize that our results hold in the rich (i.e., feature learning) regime, namely, the neural network training does not lie in the kernel regime, and thus we provide guarantees which go beyond the analysis achieved using NTK-based results.
In a bit more detail, our main contributions are the following:
- Suppose that the data distribution consists of  $k$  clusters, and the training dataset is of size  $n \geq \tilde{\Omega}(k)$ . We show that with high probability over the size- $n$  dataset, if gradient flow achieves training loss smaller than  $\frac{1}{n}$  at some time  $t_0$ , then it converges in direction to a network that generalizes well (i.e., has a small test error). Thus, gradient-flow-trained networks cannot harmfully overfit even if the network is highly overparameterized. The sample complexity  $\tilde{\Omega}(k)$  in this result is optimal (up to log factors), since we cannot expect to perform well on unseen data using a training dataset that does not include at least one example from each cluster.
- In the same setting as above, we prove that gradient flow converges in direction to a non-robust network, even though there exist robust networks that classify the data correctly. Specifically, we consider data distributions on  $\mathbb{R}^d$  such that the distance between every pair of clusters is  $\Omega(\sqrt{d})$ , and we show that there exists a two-layer ReLU network where flipping the output sign of a test example requires w.h.p. an  $\ell_2$ -perturbation of size  $\Omega(\sqrt{d})$ , but gradient flow converges to a network where we can flip the output sign of a test example with an  $\ell_2$ -perturbation of size much smaller than  $\sqrt{d}$ . Moreover, the adversarial perturbation depends only on the data distribution, and not on the specific test example or trained neural network. Thus, the perturbation is both universal [Moo+17; Zha+21] and transferable [Liu+17; AM18]. We argue that clustered data distributions are a natural setting for analyzing the tendency of gradient methods to converge to non-robust solutions. Indeed, if positive and negative examples are not well-separated (i.e., the distances between points with opposite labels are small), then robust solutions do not exist. Thus, in order to understand the role of the optimization algorithm, we need a setting with sufficient separation between positive and negative examples.
The remainder of this paper is structured as follows: Below we discuss related work. In Section 2 we provide necessary notations and background, and introduce our setting and assumptions. In Sections 3 and 4 we state our main results on generalization and robustness (respectively), and provide some proof ideas, with all formal proofs deferred to the appendix. We conclude with a short discussion (Section 5).
Implicit bias in neural networks. The literature on implicit bias in neural networks has rapidly expanded in recent years, and cannot be reasonably surveyed here (see Vardi [Var22] for a survey). In what follows, we discuss results that apply to two-layer ReLU or leaky-ReLU networks trained with gradient flow in classification settings.
By Lyu and Li [LL20] and Ji and Telgarsky [JT20], homogeneous neural networks (and specifically two-layer ReLU networks, which are the focus of this paper) trained with exponentially-tailed classification losses converge in direction to a KKT point of the maximum-margin problem. Our analysis of the implicit bias relies on this result. We note that the aforementioned KKT point may not be a global optimum of the maximum-margin problem [VSS22]. Recently, Kunin et al. [Kun+22] extended this result by showing bias towards margin maximization in a broader family of networks called quasi-homogeneous.
Lyu et al. [Lyu+21], Sarussi, Brutzkus, and Globerson [SBG21], and Frei et al. [Fre+23b] studied implicit bias in two-layer leaky-ReLU networks with linearly-separable data, and proved that under some additional assumptions gradient flow converges to a linear classifier. Chizat and Bach [CB20] studied the dynamics of gradient flow on infinite-width homogeneous two-layer networks with exponentially-tailed losses, and showed bias towards margin maximization w.r.t. a certain function norm known as the variation norm. Phuong and Lampert [PL20] studied the implicit bias in two-layer ReLU networks trained on orthogonally separable data.
Safran, Vardi, and Lee [SVL22] proved implicit bias towards minimizing the number of linear regions in univariate two-layer ReLU networks, and used this result to obtain generalization bounds. Similarly to our work, they used the KKT conditions of the maximum-margin problem in parameter space to prove generalization in overparameterized networks. However, our setting is significantly different. Implications of the bias towards KKT points of the maximum-margin problem were also studied in Haim et al. [Hai+22], where they showed that this implicit bias can be used for reconstructing training data from trained ReLU networks.
Theoretical explanations for non-robustness in neural networks. Despite much research, the reasons for the abundance of adversarial examples in trained networks are still unclear [GSS15; FFF18; Sha+19; Sch+18; KH18; Bub+19; AL21; Wan+20; Sha+20; SMB21; Sin+21; Wan+22; DB22]. Below we discuss several prior theoretical works on this question.
In one line of work, it has been shown that small adversarial perturbations can be found for any fixed input in certain neural networks with random weights (drawn from the Gaussian distribution) [DS20; Bub+21; BBC21; MW22]. These works differ in the assumptions about the width and depth of the networks as well as the activation functions considered. However, since trained networks are non-random, these works are unable to capture the existence of adversarial examples in trained networks.
The result closest to ours was shown in Vardi, Yehudai, and Shamir [VYS22]. Similarly to our result, they used the KKT conditions of the maximum-margin problem in parameter space, in order to prove that gradient flow converges to non-robust two-layer ReLU networks under certain assumptions. More precisely, they considered a setting where the training dataset  $S$  consists of nearly-orthogonal points, and proved that every KKT point is non-robust w.r.t.  $S$ . Namely, for every two-layer network that satisfies the KKT conditions of the maximum-margin problem, and every point  $\mathbf{x}_i$  from  $S$ , it is possible to flip the output's sign with a small perturbation. Their result has two main limitations: (1) It considers robustness w.r.t. the training data, while the more common setting in the literature considers robustness w.r.t. test data, as it is often more crucial to avoid adversarial perturbations in test examples; (2) Since they assume near orthogonality of the training data, the size of the dataset  $S$  must be smaller than the input dimension. Thus, they considered a high dimensional setting. We note that high-dimensional settings often have a different generalization behavior than low-dimensional settings (e.g., overfitting can be benign in the high-dimensional setting, but harmful in a low-dimensional setting [KYS23]). Our result does not suffer from these limitations, since we consider robustness w.r.t. test data, and the size of our training dataset might be very large. In our results, we essentially require near orthogonality of the cluster means, as opposed to near orthogonality of the training dataset in their result.
Finally, in Bubeck, Li, and Nagaraj [BLN21] and Bubeck and Sellke [BS21], the authors proved (under certain assumptions) that overparameterization is necessary if one wants to interpolate training data using a neural network with a small Lipschitz constant. Namely, neural networks with a small number of parameters are not expressive enough to interpolate the training data while having a small
2They also give a version of their result, where instead of assuming this upper bound on the size of the dataset, they assume an upper bound on the number of points that attain the margin in the trained network, but it is not clear a priori when this assumption is likely to hold.
Lipschitz constant. These results suggest that overparameterization might be necessary for robustness. In this work, we show that even if the network is highly overparameterized, the implicit bias of the optimization method can prevent convergence to robust solutions.
We use bold-face letters to denote vectors, e.g.,  $\mathbf{x} = (x_{1},\ldots ,x_{d})$ . For  $\mathbf{x}\in \mathbb{R}^d$  we denote by  $\| \mathbf{x}\|$  the Euclidean norm. We denote by  $\mathbb{1}[\cdot ]$  the indicator function, for example  $\mathbb{1}[t\geq 5]$  equals 1 if  $t\geq 5$  and 0 otherwise. We denote  $\mathrm{sign}(z) = 1$  if  $z > 0$  and  $-1$  otherwise. For an integer  $d\geq 1$  we denote  $[d] = \{1,\dots ,d\}$ . For a set  $A$  we denote by  $\mathcal{U}(A)$  the uniform distribution over  $A$ . We denote by  $\mathsf{N}(\mu ,\sigma^2)$  the normal distribution with mean  $\mu \in \mathbb{R}$  and variance  $\sigma^2$ , and by  $\mathsf{N}(\boldsymbol {\mu},\boldsymbol {\Sigma})$  the multivariate normal distribution with mean  $\pmb{\mu}$  and covariance matrix  $\boldsymbol{\Sigma}$ . The identity matrix of size  $d$  is denoted by  $I_{d}$ . We use standard asymptotic notation  $\mathcal{O}(\cdot)$  and  $\Omega (\cdot)$  to hide constant factors, and  $\tilde{\mathcal{O}} (\cdot),\tilde{\Omega} (\cdot)$  to hide logarithmic factors. We use log for the logarithm with base 2 and in for the natural logarithm.
In this work, we consider depth-2 ReLU neural networks. The ReLU activation function is defined by  $\phi(z) = \max\{0, z\}$ . Formally, a depth-2 network  $\mathcal{N}_{\theta}$  of width  $m$  is parameterized by  $\theta = [\mathbf{w}_1, \ldots, \mathbf{w}_m, \mathbf{b}, \mathbf{v}]$  where  $\mathbf{w}_i \in \mathbb{R}^d$  for all  $i \in [m]$  and  $\mathbf{b}, \mathbf{v} \in \mathbb{R}^m$ , and for every input  $\mathbf{x} \in \mathbb{R}^d$  we have
$$
\mathcal {N} _ {\boldsymbol {\theta}} (\mathbf {x}) = \sum_ {j \in [ m ]} v _ {j} \phi \left(\mathbf {w} _ {j} ^ {\top} \mathbf {x} + b _ {j}\right).
$$
We sometimes view  $\theta$  as the vector obtained by concatenating the vectors  $\mathbf{w}_1, \ldots, \mathbf{w}_m$ ,  $\mathbf{b}$ ,  $\mathbf{v}$ . Thus,  $\|\theta\|$  denotes the  $\ell_2$  norm of the vector  $\theta$ . We note that in this work we train both layers of the ReLU network.
We denote  $\Phi (\pmb {\theta};\mathbf{x})\coloneqq \mathcal{N}_{\pmb{\theta}}(\mathbf{x})$  . We say that a network is homogeneous if there exists  $L > 0$  such that for every  $\alpha >0$  and  $\pmb{\theta}$  , x we have  $\Phi (\alpha \pmb {\theta};\mathbf{x}) = \alpha^L\Phi (\pmb {\theta};\mathbf{x})$  . Note that depth-2 ReLU networks as defined above are homogeneous (with  $L = 2$
We next define gradient flow and remind the reader of some recent results on the implicit bias of gradient flow in two-layer ReLU networks. Let  $\mathcal{S} = \{(x_i, y_i)\}_{i=1}^n \subseteq \mathbb{R}^d \times \{-1, 1\}$  be a binary classification training dataset. Let  $\Phi(\theta; \cdot): \mathbb{R}^d \to \mathbb{R}$  be a neural network parameterized by  $\theta$ . For a loss function  $\ell: \mathbb{R} \to \mathbb{R}$  the empirical loss of  $\Phi(\theta; \cdot)$  on the dataset  $S$  is
$$
\mathcal {L} (\boldsymbol {\theta}) := \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(y _ {i} \Phi \left(\boldsymbol {\theta}; \mathbf {x} _ {i}\right)\right). \tag {1}
$$
We focus on the exponential loss  $\ell(q) = e^{-q}$  and the logistic loss  $\ell(q) = \log(1 + e^{-q})$ .
We consider gradient flow on the objective given in Eq. (1). This setting captures the behavior of gradient descent with an infinitesimally small step size. Let  $\pmb{\theta}(t)$  be the trajectory of gradient flow. Starting from an initial point  $\pmb{\theta}(0)$ , the dynamics of  $\pmb{\theta}(t)$  is given by the differential equation  $\frac{d\pmb{\theta}(t)}{dt} \in -\partial^{\circ}\mathcal{L}(\pmb{\theta}(t))$ . Here,  $\partial^{\circ}$  denotes the Clarke subdifferential [Cla+08], which is a generalization of the derivative for non-differentiable functions.
We now remind the reader of a recent result concerning the implicit bias of gradient flow over the exponential and logistic losses for homogeneous neural networks. Note that since homogeneous networks satisfy  $\mathrm{sign}(\Phi (\alpha \pmb {\theta};\mathbf{x})) = \mathrm{sign}(\Phi (\pmb {\theta};\mathbf{x}))$  for any  $\alpha >0$ , the sign of the network output of homogeneous networks depends only on the direction of the parameters  $\pmb{\theta}$ . The following theorem provides a characterization of the implicit bias of gradient flow by showing that the trajectory of the weights  $\pmb{\theta}(t)$  converge in direction to a first-order stationary point of a particular constrained optimization problem, where  $\pmb{\theta}$  converges in direction to  $\tilde{\pmb{\theta}}$  means  $\lim_{t\to \infty}\frac{\pmb{\theta}(t)}{\|\pmb{\theta}(t)\|} = \frac{\tilde{\pmb{\theta}}}{\|\tilde{\pmb{\theta}}\|}$ . Note that since ReLU networks are non-smooth, the first-order stationarity conditions (i.e., the Karush-Kuhn-Tucker conditions, or KKT conditions for short) are defined using the Clarke subdifferential (see Lyu and Li [LL20] and Dutta et al. [Dut+13] for more details on the KKT conditions in non-smooth optimization problems).
Theorem 2.1 (Paraphrased from Lyu and Li [LL20] and Ji and Telgarsky [JT20]). Let  $\Phi(\pmb{\theta};\cdot)$  be a homogeneous ReLU neural network parameterized by  $\pmb{\theta}$ . Consider minimizing either the exponential or the logistic loss over a binary classification dataset  $\{(\mathbf{x}_i,y_i)\}_{i=1}^n$  using gradient flow. Assume that there exists time  $t_0$  such that  $\mathcal{L}(\pmb{\theta}(t_0)) < \frac{1}{n}$  (and thus  $y_i\Phi(\pmb{\theta}(t_0);\mathbf{x}_i) > 0$  for every  $\mathbf{x}_i$ ). Then, gradient flow converges in direction to a first-order stationary point (KKT point) of the following
maximum margin problem in parameter space:
$$
\min  _ {\boldsymbol {\theta}} \frac {1}{2} \| \boldsymbol {\theta} \| ^ {2} \quad \text {s . t .} \quad \forall i \in [ n ] y _ {i} \Phi (\boldsymbol {\theta}; \mathbf {x} _ {i}) \geq 1. \tag {2}
$$
Moreover,  $\mathcal{L}(\pmb {\theta}(t))\to 0$  and  $\| \pmb {\theta}(t)\| \to \infty$  as  $t\to \infty$
Theorem 2.1 gives a characterization of the implicit bias of gradient flow with the exponential and the logistic loss for homogeneous ReLU networks. Note that the theorem makes no assumption on the initialization, training data, or number of parameters in the network; the only requirement is that the network is homogeneous and that at some time point in the gradient flow trajectory, the network is able to achieve small training loss. The theorem shows that although there are many ways to configure the network parameters to achieve small training loss (via overparameterization), gradient flow only converges (in direction) to networks which satisfy the KKT conditions of Problem (2). It is important to note that satisfaction of the KKT conditions is not sufficient for global optimality of the constrained optimization problem [VSS22]. We further note that if the training data are sampled i.i.d. from a distribution with label noise (e.g., a class-conditional Gaussian mixture model, or a distribution where labels  $y_{i}$  are flipped to  $-y_{i}$  with some nonzero probability), networks which have parameters that are feasible w.r.t. the constraints of Problem (2) have overfit to noise, and understanding the generalization behavior of even globally optimal solutions to Problem (2) in this setting is the subject of significant research [Mon+19; CL21; Fre+23a].
Finally, we introduce the distributional setting that we consider. We consider a distribution  $\mathcal{D}_{\mathrm{clusters}}$  on  $\mathbb{R}^d\times \{-1,1\}$  that consists of  $k$  clusters with means  $\pmb{\mu}^{(1)},\dots,\pmb{\mu}^{(k)}\in \mathbb{R}^d$  and covariance  $\sigma^2 I_d$  (i.e., a Gaussian mixture model), such that the examples in the  $j$ -th cluster are labeled by  $y^{(j)}\in \{-1,1\}$ . More formally,  $(\mathbf{x},y)\sim \mathcal{D}_{\mathrm{clusters}}$  is generated as follows: we draw  $j\sim \mathcal{U}([k])$  and  $\mathbf{x}\sim \mathsf{N}(\pmb{\mu}^{(j)},\sigma^2 I_d)$  and set  $y = y^{(j)}$ . We assume that there exist  $i,j\in [k]$  with  $y^{(i)}\neq y^{(j)}$ . Moreover, we assume the following:
Assumption 2.2. We have:
$\left\| \pmb{\mu}^{(j)}\right\| = \sqrt{d}$  for all  $j\in [k]$
-  $0 < \sigma \leq 1$ .
-  $k\left(\max_{i \neq j} |\langle \pmb{\mu}^{(i)}, \pmb{\mu}^{(j)} \rangle| + 4\sigma \sqrt{d} \ln(d) + 1\right) \leq \frac{d - 4\sigma \sqrt{d} \ln(d) + 1}{10}$ .
Example 1. Below we provide simple examples of settings that satisfy the assumption:
- Suppose that the cluster means satisfy  $|\langle \pmb{\mu}^{(i)},\pmb{\mu}^{(j)}\rangle| = \tilde{\mathcal{O}}(\sqrt{d})$  for every  $i \neq j$ . This condition holds, e.g., if we choose each cluster mean i.i.d. from the uniform distribution on the sphere  $\sqrt{d} \cdot \mathbb{S}^{d - 1}$  (see, e.g., Vardi, Yehudai, and Shamir [VYS22, Lemma 3.1]). Let  $\sigma = 1$ , namely, each cluster has a radius of roughly  $\sqrt{d}$ . Then, the assumption can be satisfied by choosing  $k = \tilde{\mathcal{O}}(\sqrt{d})$ .
- Suppose that the cluster means are exactly orthogonal (i.e.,  $\langle \pmb{\mu}^{(i)},\pmb{\mu}^{(j)}\rangle = 0$  for all  $i\neq j$ ), and  $\sigma = 1 / \sqrt{d}$ . Then, the assumption can be satisfied by choosing  $k = \tilde{\mathcal{O}} (d)$ .
- If the number of clusters is  $k = \tilde{\mathcal{O}}(1)$ , then the assumption may hold even where  $\max_{i \neq j} |\langle \pmb{\mu}^{(i)}, \pmb{\mu}^{(j)} \rangle| = \tilde{\Theta}(d)$  (for any  $0 < \sigma \leq 1$ ).
A few remarks are in order. First, the assumption that  $\|\pmb{\mu}^{(j)}\|$  is exactly  $\sqrt{d}$  is for convenience, and we note that it may be relaxed (to have all cluster means approximately of the same norm) without affecting our results significantly. Note that in the case where  $\sigma = 1$ , the radius of each cluster is roughly of the same magnitude as the cluster mean. Second, we assume for convenience that the noise (i.e., the deviation from the cluster's mean) is drawn from a Gaussian distribution with covariance matrix  $\sigma^2 I_d$ . However, we note that this assumption can be generalized to any distribution  $\mathcal{D}_{\mathrm{noise}}$  such that for every unit vector  $\mathbf{e}$  the noise  $\xi \sim \mathcal{D}_{\mathrm{noise}}$  satisfies w.h.p. that  $\langle \xi, \mathbf{e} \rangle = \tilde{\mathcal{O}}(1)$  and  $\|\xi\| = \tilde{\mathcal{O}}(\sqrt{d})$ . This property holds, e.g., for a  $d$ -dimensional Gaussian distribution  $\mathsf{N}(0, \Sigma)$ , where  $\operatorname{tr}[\Sigma] = d$  and  $\|\Sigma\|_2 = O(1)$  (see Frei et al. [Fre+23b, Lemma 3.3]), and more generally for a class of sub-Gaussian distributions (see Hu et al. [Hu+20, Claim 3.1]). Third, note that the third part of Assumption 2.2 essentially requires that the number of clusters  $k$  cannot be too large
and the correlations between cluster means cannot be too large. Finally, we remark that when  $k$  is small, our results may be extended to the case where  $\sigma > 1$ . For example, if  $k = \tilde{\mathcal{O}}(1)$  and  $\max_{i \neq j} |\langle \pmb{\mu}^{(i)}, \pmb{\mu}^{(j)} \rangle| = \tilde{\mathcal{O}}(\sqrt{d})$ , our generalization result (Theorem 3.1) can be extended to the case where  $\sigma = \tilde{\mathcal{O}}(d^{1/8})$ . We preferred to avoid handling  $\sigma > 1$  in order to simplify the proofs.