Implicit Bias of Gradient Descent for Wide Two-layer Neural Networks Trained with the Logistic Loss

Introduction

Artificial neural networks are successfully used in a variety of difficult supervised classification tasks, but the mechanisms behind their performance remain unclear. The situation is particularly intriguing when the number of parameters of these models exceeds by far the number of input data points and they are trained with gradient-based methods until zero training error, without any explicit regularization. In this case, the training algorithm induces an implicit bias: among the many classifiers which overfit on the training set, it selects a specific one which often turns out to perform well on the test set. In this paper, we study the implicit bias of wide neural networks with two layers (i.e., with a single hidden-layer) trained with gradient descent on the logistic loss, or any loss with an exponential tail. Our analysis lies at the intersection of two lines of research that study (i) the implicit bias of gradient methods, and (ii) the training dynamics of wide neural networks.

Dynamics of infinitely-wide neural networks.

This fine characterization is made possible by looking at the infinite width limit of two-layer neural networks. This strategy has been used in several works to obtain insights on their statistical properties (Bengio et al., 2006; Bach, 2017a) or training behavior (Nitanda and Suzuki, 2017; Rotskoff and Vanden-Eijnden, 2018; Chizat and Bach, 2018; Mei et al., 2018; Sirignano and Spiliopoulos, 2019), which can be described by a Wasserstein gradient flow (Ambrosio et al., 2008). In particular, Chizat and Bach (2018) show that if the loss is convex, if the initialization is “diverse enough”, and if the gradient flow of the objective converges, then its limit is a global minimizer. This result does not apply in our context because the gradient flow diverges, which turns out to be beneficial for the analysis of the implicit bias that we propose.

A general drawback of those mean-field analyses is that they are mostly non-quantitative, both in terms of number of neurons and number of iterations. While some works have shown quantitative results by modifying the dynamics (Mei et al., 2019; Wei et al., 2019; Chizat, 2019), we do not take this path in order to stay close to the way neural networks are used in practice and because our numerical experiments suggest that those modifications are not necessary to obtain a good practical behavior. Finally, we stress that our analysis does not take place in the lazy training regime (Chizat et al., 2019) which consists of training dynamics that can be analyzed in a perturbative regime around the initialization (see, e.g., Li and Liang, 2018; Jacot et al., 2018; Du et al., 2019). Lazy training is another kind of implicit bias that amounts to training a linear model and does not lead to adaptivity results as those shown in Section 6 (see Figure 3 for an illustration in our context).

1 Organization and contributions

After preliminaries on wide neural networks in Section 2, we make the following contributions :

In Section 3, we show that for a class of two-layer neural networks and for losses with an exponential tail, the classifier learnt by the non-convex gradient flow is a max-margin classifier for a certain functional norm known as the variation norm.

When fixing the “directions” of the neurons (Section 4), or when only training the output layer (Section 5), we show that the dynamics implicitly performs online mirror ascent on a sequence of smooth-margin objectives and thus naturally maximizes the margin. This leads to convergence guarantees in $O(\log(t)/\sqrt{t})$ in situations where no rate was previously known.

In Section 6, we study the margins of those classifiers and prove dimension-independent generalization bounds for classification in presence of hidden linear structures.

We perform numerical experiments in Section 7 for two-layer ReLU neural networks which confirm the statistical efficiency of this implicit bias in a high-dimensional setting.

In summary, we show that training two-layer ReLU neural networks implicitly solves a problem with strong statistical benefits. We stress however that the runtime of the algorithm is still unknown.

2 Notation

Preliminaries on infinitely wide two-layer networks

Here are examples of models which satisfy (A1):

2 Parameterizing with a measure

Finite width networks as in Eq. (1) are recovered when $\mu$ is a discrete measure with $m$ atoms.

3 Max-margins and functional norms

In this paper we deal with two notions of norms, that in turn define two types of max-margin classifiers. We refer to Bach (2017a) for a more detailed presentation.

RKHS norm.

This is a separable kernel support vector machine problem.

Statistical and computational properties.

In Section 6, we will show that the margin $\gamma_{1}$ can be large even in high dimension when the dataset has hidden low dimensional structure, which leads to strong generalization guarantees, which is a priori not true for $\gamma_{2}$. While $\mathcal{F}_{2}$-max-margin classifiers can be found with convex optimization techniques (such as training only the output layer, as shown in Section 5), it is not clear a priori how to find $\mathcal{F}_{1}$-max-margin classifiers. In the next section, we show that training an over-parameterized two-layer neural network precisely does that.

4 Training dynamics in the infinite width limit

The family $(\phi(\cdot,x_{i}))_{i\in[n]}$ is linearly independent and for $i\in[n]$, the function $\phi(\cdot,x_{i})$ is differentiable with a Lipschitz-continuous gradient and subanalytic (i.e., its graph is locally the linear projection of a bounded semianalytic set).

Gradient flow of the smooth-margin objective.

Up to the gradient sign, this gradient flow is an approximation of gradient descent (Gautschi, 1997; Scieur et al., 2017) and stochastic gradient descent (SGD) (Kushner and Yin, 2003, Thm. 2.1) with small step sizes.Although Theorem 3.1 below could be extended to discrete time analysis, this would be of little interest since the result is so far purely qualitative. In simpler settings, we study discrete time dynamics in Sections 4 and 5. Classical results guarantee that under Assumption (A2-3), this gradient flow is uniquely well defined.

Wasserstein gradient flow.

It can be directly checked that when $\mu_{0}$ is discrete, we recover the training dynamics defined in Eq. (7). In this case, $w_{j}(t)=X(t,w_{j}(0))$ is the position (in parameter space) at time $t$ of the hidden unit initialized with parameters $w_{j}(0)$. The following theorem shows that Wasserstein gradient flows characterize the training dynamics of infinitely wide two-layer neural networks. It is an application of Chizat and Bach (2018, Thm. 2.6), see details in Appendix C (hereafter, by convergence in $\mathcal{P}_{2}$, we mean weak convergence and convergence of the second moments (Ambrosio et al., 2008)).

This limit can be made quantitative using the geodesic convexity estimates of Chizat and Bach (2018), and the stability results of Ambrosio et al. (2008, Thm. 11.2.1) but with an exponential dependency in time. In the different setting of the square loss, error estimates for SGD have been derived by Mei et al. (2018, 2019). This limit dynamics covers, but is not limited to, the lazy training dynamics studied by Li and Liang (2018); Jacot et al. (2018); Du et al. (2019) which here corresponds to a short time analysis when the initialization has a large variance (see Figure 3 in Section 7).

Main result: implicit bias of gradient flow

We are now in position to state the main theorem of this paper, which characterizes the implicit bias of training infinitely wide two-layer neural networks with a loss with an exponential tail.

Under (A1-3), assume that $\Pi_{2}(\mu_{0})$ has full support on $\SS^{p-1}$. If $\nabla S(\hat{h}(\mu_{t}))$ converges and $\bar{\nu}_{t}=\Pi_{2}(\mu_{t})/([\Pi_{2}(\mu_{t})](\SS^{p-1}))$ converges weakly to some $\bar{\nu}_{\infty}$, then this limit $\bar{\nu}_{\infty}$ is a maximizer for the $\mathcal{F}_{1}$-max-margin problem in Eq. (4).

The strength of this result is that the limit $\bar{\nu}_{\infty}$ of a non-convex dynamics is a global minimizer of Eq. (4). Its proof relies, among other things, on a compatibility between the optimality conditions and the gradient flow dynamics, which is specific to the $2$-homogeneous case.

It is an open question to prove that $\nabla S(\hat{h}(\mu_{t}))$ and $\bar{\nu}_{t}$ converge is this setting. Note that the unnormalized measure $\nu_{t}$ does not converge, so the global convergence result from Chizat and Bach (2018) (which has a similar assumption regarding the existence of a limit) does not apply.

Unlike in the convex case (Soudry et al., 2018), the dynamics does not completely forget where it started from. For instance, when initialized with a Dirac measure, the Wasserstein gradient flow can only converge to a Dirac measure, which is typically not a global minimizer.

Together, Theorems 2.2 and 3.1 give asymptotic guarantees for training finite width neural networks.

There exists $(\mu_{t})_{t\geq 0}$ a Wasserstein gradient flow of the objective Eq. (9) with $\mu_{0}=\mathcal{U}(\SS^{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})](\SS^{p-1}))$ converges weakly in $\mathcal{P}(\SS^{p-1})$ and if (*) $F^{\prime}_{\mu_{t}}$ converges in $\mathcal{C}^{1}_{loc}$ to some $F^{\prime}$ that satisfies the Morse-Sard property (see details in Appendix H), then $h(\bar{\nu}_{\infty},\cdot)$ is a maximizer for $\max_{\|f\|_{\mathcal{F}_{1}}\leq 1}\min_{x\in\mathcal{X}}y(x)f(x).$

Insights on the convergence rate and choice of step-size

While making Corollary 3.2 quantitative in terms of number of neurons and the number of iterations is left as an open question, it is of practical importance to better understand the effect of the choice of step-size. In this section, we look at a simplified dynamics where the direction of each parameter $w_{j}(t)$ is fixed after initialization and only its magnitude evolves. A complete discrete-time analysis is possible in this case, using tools from convex analysis.

Let $a_{j}(t)=r_{j}(t)^{2}/m$ for $j\in[m]$, $\beta(t)=\|a(t)\|_{1}$ and $\bar{a}(t)=a(t)/\beta(t)$. For the step-sizes $\eta(t)=1/(16\|z\|_{\infty}\sqrt{t+1})$ and a uniform initialization $r(0)\propto\mathbf{1}$, it holds

where $\gamma_{1}^{(m)}:=\max_{a\in\Delta^{m-1}}\min_{i\in[n]}z_{i}^{\top}a$ and $B:=\sum_{s=0}^{\infty}\frac{1}{\beta(s)\sqrt{s+1}}<\infty$ when $\gamma_{1}^{(m)}>0$.

In the proof of Lemma E.3, it can be seen that our bound on $B$ grows to $\infty$ as $\gamma_{1}^{(m)}$ goes to zero.

To prove Proposition 4.1, we consider the family of smooth-margin functions

and we show that $\bar{a}(t)$ approximately follows online mirror ascent for the sequence of concave functions $G_{\beta(t)}$ in the simplex $\Delta^{m-1}$ with step-sizes $\eta(t)$. It then only remains to apply classical bounds for mirror descent and use the fact that $|\min_{i\in[n]}z_{i}^{\top}a-G_{\beta}(a)|\leq\log(n)/\beta$. This algorithm thus implicitly performs online optimization on the regularization path. It is also analogous to smoothing techniques in non-smooth optimization (Nesterov, 2005).

Continuous limit.

Using the notations from Section 2, the dynamics $\sum_{j=1}^{m}\bar{a}_{j}(t)\delta_{\theta_{j}}$ solves

When $\frac{1}{m}\sum_{j=1}^{m}\delta_{\theta_{j}}$ converges to the uniform measure on the sphere, we thus recover the same implicit bias as in Theorem 3.1 and $\gamma_{1}^{(m)}\to\gamma_{1}$ (note that the logarithmic dependency in $m$ in Proposition 4.1 could be removed with a slightly finer analysis as done in Chizat (2019)). While functions in $\mathcal{F}_{1}$ may be well-approximated with a small number of neurons (Bach, 2017a; Jones, 1992), this is not anymore true if the positions $\{\theta_{j}\}_{j\in[m]}$ of those neurons are fixed a priori (see Barron (1993) for exponential lower bounds in a similar setting). In Theorem 3.1, positions are allowed to vary during training: this makes its setting more challenging but also much more relevant.

Training only the output layer

Let $a(t)=r(t)/m$, $\beta(t)=\max\{1,\max_{0\leq s\leq t}\sqrt{m}\|a(t)\|_{2}\}$ and $\bar{a}(t)=a(t)/\beta(t)$. Assume $\gamma_{2}^{(m)}:=\max_{\sqrt{m}\|a\|_{2}\leq 1}\min_{i\in[n]}z_{i}^{\top}a>0$. For the step-sizes $\eta(t)=\beta(t){\sqrt{2}}/(\|z\|_{\infty}\sqrt{t+1})$ and initialization $r(0)=0$, it holds

Random features for kernel max-margin classifier.

Using the notations from Section 2, the dynamics $(r_{j})_{j\in[m]}$ converges to a solution to

Dimension independent generalization bounds

In this section, we give arguments showing the favorable statistical properties of the bias exhibited in Theorem 3.1 for ReLU networks. We propose to measure the complexity of the dataset $S_{n}=(x_{i},y_{i})_{i=1}^{n}$ with the following projected interclass distance defined, for $r\in[d]$, as

For each dimension $r$, it looks for the $r$-dimensional subspace which maximizes the distance between the two classes. Interclass distance often appears in the statistical analysis of classification problems (see, e.g., Li and Liang, 2018) often complemented with “clustered data” assumptions. Our definition is designed to capture the fact that if $\Delta_{r}\approx\Delta_{d}$ for $r\ll d$, then there is a hidden structure which can be exploited for statistical efficiency.

We first lower-bound the margins $\gamma_{1}$ and $\gamma_{2}$ in terms of $\Delta_{r}(S_{n})$ and then apply margin-based generalization bounds (Koltchinskii and Panchenko, 2002) and bounds on the Rademacher complexity of the unit ball of $\mathcal{F}_{1}$ and $\mathcal{F}_{2}$.

Numerical experiments

In this section, we consider a large ReLU network with $m=1000$ hidden units, and compare the implicit bias and statistical performances of training both layers – which leads to a max margin classifier in $\mathcal{F}_{1}$ – versus the output layer – which leads to max margin classifier in $\mathcal{F}_{2}$. The experiments are reproducible with the Julia code that can be found online\urlhttps://github.com/lchizat/2020-implicit-bias-wide-2NN.

Low dimensional illustrations.

Figure 1 illustrates the differences in the implicit biases when $d=2$. It represents a sampled training set and the resulting decision boundary between the two classes for $4$ examples. The $\mathcal{F}_{1}$- max-margin classifier is non-smooth and piecewise affine, which comes from the fact that the mass constraint in Eq. (4) favors sparse solutions. In contrast, the max-margin classifier in $\mathcal{F}_{2}$ has a smooth decision boundary, which is typical of learning in a RKHS.

Performance.

Two implicit biases in one dynamics.

In Figure 3, we illustrate for $d=2$ a case where two different kinds of implicit biases show up in a single dynamics ($t$ is the number of iterations with a constant step-size). We initialize the ReLU network with a large variance ($\mathcal{N}(0,40^{2})$). The model is at first in the lazy regime (Chizat et al., 2019) and follows closely the dynamics of its linearization around initialization, which converges to the max-margin classifier for the tangent kernel (Jacot et al., 2018). It then converges to the $\mathcal{F}_{1}$-max-margin classifier as suggested by Theorem 3.1. In order to observe this intermediate implicit bias, one needs an initial step-size inversely proportional to the scale of the initialization (Chizat et al., 2019).

Conclusion

We have shown that for wide two-layer ReLU neural networks, training both layers or only the output layer leads to very different implicit biases. When training both layers, the classifier converges to a max-margin classifier for a non-Hilbertian norm, which enjoys favorable statistical properties. Interestingly, this problem does not seem to be directly solvable with known convex methods in high dimension. Proving complexity guarantees for this non-convex gradient flow is an important open question for future work. In particular, even for infinite width, continuous time dynamics as in Theorem 3.1, it is still unknown whether a convergence rate can be given under reasonable conditions.

Acknowledgements

Part of this work was carried through while the first author was visiting the Chair of Statistical Field Theory at the École Polytechnique Fédérale de Lausanne (EPFL), Switzerland. This work was funded in part by the French government under management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference ANR-19-P3IA-0001 (PRAIRIE 3IA Institute). We also acknowledge support the European Research Council (grant SEQUOIA 724063).

References

Appendix A Organization of the appendix

In Appendix B, we prove the equivalence between our definition of the variation norm in Section 2.3 with the one that is used in the literature on convex neural networks.

In Appendix C, we discuss properties of the Wasserstein gradient flow and justify Theorem 2.2.

In Appendix D, we prove our main theorem Theorem 3.1 and its corollary.

In Appendix E, we prove Proposition 4.1 on the convergence rate with fixed “positions”.

In Appendix F, we prove Proposition 5.1 on the convergence rate when training the output layer.

In Appendix G, we prove Theorem 6.1 on the margins and generalization performance.

In Appendiz H, we prove Theorem 3.3 which covers the case of ReLU networks.

Appendix B Equivalence of two variation norms

Let us introduce, for ReLU networks, the variation norm introduced in Section 2.3 and the different definition from the literature (Bengio et al., 2006; Bach, 2017a). We will show that they are equal up to a factor $2$. This is a known result (Neyshabur et al., 2014), and we provide here a natural proof using the measure theoretic formalism, for the sake of completeness. We stress that the analogous equivalence would fail for the RKHS norms, i.e., such a modification of the feature function could lead to different functional spaces.

An interesting consequence of this result is that empirical risk minimization with the commonly used weight decay regularization and the total variation regularization used by Bengio et al. (2006); Bach (2017a) (the path-norm) are equivalent.

As for the total variation norm $\|\Pi(\nu)\|:=|\Pi(\nu)|(\SS^{p-2})$ of $\Pi(\nu)$, it can be bounded as follows. In the definition of $\Pi(\nu)$, we may restrict the integral over $\{c>0\}$, which defines a measure $\Pi_{+}(\nu)\in\mathcal{M}_{+}(\SS^{p-2})$. Similarly restricting the integral over $\{c<0\}$ an taking the opposite gives another measure $\Pi_{-}(\nu)\in\mathcal{M}_{+}(\SS^{p-2})$. It holds $\Pi(\nu)=\Pi_{+}(\nu)-\Pi_{-}(\nu)$ and thus $\|\Pi(\nu)\|\leq\|\Pi_{+}(\nu)\|+\|\Pi_{-}(\nu)\|$. Moreover, by integrating against $\varphi=1$, it holds

since $c\|a\|\leq(\|a\|^{2}+c^{2})/2=1/2$ for $(a,c)\in\SS^{p-1}$. Using a similar bound for $\|\Pi_{-}(\nu)\|$, we get that

Finally, tracking the equality cases, it holds $\|\Pi(\nu)\|=\frac{1}{2}\nu(\SS^{p-1})$ if and only if $\nu$ is concentrated on the set given in Proposition B.1, which is the intersection of the sphere with the set of points satisfying $2|c|\|a\|=\|a\|^{2}+|c|^{2}$.

Conversely, let $\nu\in\mathcal{M}(\SS^{p-2})$ and consider its Jordan decomposition $\nu=\nu_{+}-\nu_{-}$ into two nonnegative measures, which is such that $\|\nu\|=\|\nu_{+}\|+\|\nu_{-}\|$. We define two maps $T^{+},T^{-}:\SS^{p-2}\to\SS^{p-1}$ as $T^{+}(a)=(a,1)/\sqrt{2}$ and $T^{-}(a)=(a,-1)/\sqrt{2}$. Now, define the linear map $T:\mathcal{M}(\SS^{p-2})\to\mathcal{M}_{+}(\mathcal{S}^{p-1})$ as

Since pushforwards preserve the mass of nonnegative measures, it holds $\|T(\nu)\|=2(\|\nu_{+}\|+\|\nu_{-}\|)=2\|\nu\|$. Moreover, using the definition of pushforward measures, we have

Appendix C Details on Wasserstein gradient flows

(Divergence form) It can be shown that a Wasserstein gradient flow as defined in Definition 2.1 satisfies, in the sense of distributions, the following partial differential equation (Ambrosio et al., 2008)

(Projected representation) If we look at the projected trajectory $\nu_{t}=\Pi_{2}(\mu_{t})$, it can be shown that it solves the following dynamic, which is known as Wasserstein-Fisher-Rao or Hellinger-Kantorovich gradient flow of the functional $J:\mathcal{M}_{+}(\SS^{p-1})$ satisfying $J(\Pi_{2}(\mu))=F(\mu)$. In equation,

where $J^{\prime}_{\nu}(\theta)=\sum_{i=1}^{n}\nabla_{i}S(\hat{h}(\nu))y_{i}\phi(\theta,x_{i})$ is defined on the sphere, see Chizat (2019). Note that there is also a Lagrangian representation for this projected dynamics (Maniglia, 2007) .

(Renormalized dynamics) It can also be seen with a direct computation that the normalized dynamics $\bar{\nu}_{t}=\nu_{t}/\|\nu_{t}\|$ satisfies the following equation

When the driving potential is $J^{\prime}_{\bar{\nu}_{t}}$ (instead of $J^{\prime}_{\nu_{t}}$), this dynamics is known as the spherical Wasserstein-Fisher-Rao or spherical Hellinger Kantorovich gradient flow (Kondratyev and Vorotnikov, 2019) and was considered by Rotskoff et al. (2019) for neural networks training.

C.2 Proof of Theorem 2.2

We just need to prove that the assumptions of Chizat and Bach (2018, Theorem 2.6) are satisfied and justify that non-compactly supported initialization are also allowed.

under Assumption (A3), the function $\Phi$ is differentiable with a locally Lipschitz continuous gradient. Moreover its gradient has at most a linear growth by $2$-homogeneity (this verifies assumptions from Chizat and Bach (2018, Assumptions 2.1-(iii)-(c))).

Removing the compact support assumption.

Appendix D Appendix to Section 3: main theorem

Let us restate Theorem 3.1 and give its proof. We recall that $\nu_{t}:=\Pi_{2}(\mu_{t})$ and $\bar{\nu}_{t}:=\nu_{t}/(\nu_{t}(\SS^{p-1}))$.

Step 1: mass grows unbounded. In a first step, we prove that $\nu_{t}(\SS^{p-1})\to\infty$. Assume that $J^{\prime}_{\infty}$ is not constant (the other case will be considered later), and let $v\in{]0,M/8[}$ be such that $M-v$ is a regular value of $J^{\prime}_{\infty}$, i.e., be such that $\|\nabla J^{\prime}_{\infty}\|$ does not vanish on the $M-v$ level-set of $J^{\prime}_{\infty}$. Such a $v$ is guaranteed to exist thanks to the fact that $\phi(\cdot,x_{i})$ is subanalytic (which implies that $J^{\prime}_{\infty}$, which is a finite sum of such $\phi(\cdot,x_{i})$, is also subanalytic) and that the sphere is a subanalytic set, and then applying (Bolte et al., 2006, Thm. 14). Note that such admissible $v$ are dense in the range of $J^{\prime}_{\infty}$, which will be useful in Step. 3. Let $K_{v}=(J^{\prime}_{\infty})^{-1}([M-v,M])\subset\SS^{p-1}$ be the corresponding super-level set. By the regular value theorem, the boundary $\partial K_{v}$ of $K_{v}$ is a differentiable orientable compact submanifold of $\SS^{p-1}$ and is orthogonal to $\nabla J^{\prime}_{\infty}$. By construction, it holds for all $\theta\in K_{v}$, $J^{\prime}_{\infty}(\theta)\geq M-v$ and, for some $u>0$, by the regular value property, $\nabla J^{\prime}_{\infty}(\theta)\cdot\vec{n}_{\theta}\geq u$ for all $\theta\in\partial K_{v}$ where $\vec{n}_{\theta}$ is the unit normal vector to $\partial K_{v}$ at $\theta$ pointing inwards. Since $J^{\prime}_{\nu_{t}}$ converges in $\mathcal{C}^{1}(\SS^{p-1})$ towards $J^{\prime}_{\infty}$, there exists $t_{0}>0$ such that for all $t\geq t_{0}$, $\|J^{\prime}_{\nu_{t}}-J^{\prime}_{\infty}\|_{\mathcal{C}^{1}(\SS^{p-1})}\leq\min\{v,u/2\}$ and thus

It follows by Grönwall’s lemma that $\nu_{t}(K_{v})\geq\exp(4(M-2v)t)\nu_{t_{0}}(K_{v})$ for $t\geq t_{0}$. On the other hand, $\nu_{t_{0}}$ has full support on $\SS^{p-1}$ since it can be written as the pushforward of a rescaled version of $\nu_{0}$ by a diffeomorphism, see Maniglia (2007, Eq. (1.3)) (this is the only place where the assumption on the support is needed). Thus $\nu_{t_{0}}(K_{v})>0$ and it follows that $\nu_{t}(\SS^{p-1})\to\infty$. To deal with the case where $J^{\prime}_{\infty}$ is constant and is equal to $M>0$, we can directly take $K=\SS^{p-1}$ to show that $\nu_{t}(\SS^{p-1})\to\infty$. In the rest of the proof, we show that $(\bar{\nu}_{\infty},p(\infty))$ satisfy the optimality conditions of Eq. (4) given by Proposition D.5, which we refer to as the complementary slackness conditions.

Step 2: complementary slackness (I). We first show that $\min_{i}\hat{h}_{i}\to\infty$. Using the property of gradient flows and previously established estimates, we have for $T>0$,

Step 3: complementary slackness (II). We now show that $\bar{\nu}_{\infty}$ is concentrated on $(J^{\prime}_{\infty})^{-1}(M)$, where $\bar{\nu}_{t}=\nu_{t}/\|\nu_{t}\|\in\mathcal{P}(\SS^{p-1})$ is the normalized path and $\bar{\nu}_{\infty}$ its limit. This is immediate if $J^{\prime}_{\infty}$ is constant. Otherwise, assuming that $M-4v$ is also a regular value of $J^{\prime}_{\infty}$ and taking a potentially smaller $u$ (which can always be achieved by perturbing $v$ if needed, since regular values are dense as mentioned in Step. 1), it holds for $t\geq t_{0}$,

using the fact that no mass enters into $\SS^{p-1}\setminus K_{4v}$ due to the divergence term in Eq. (11) for $t\geq t_{0}$. Comparing the rate of growth of the mass in $K_{v}$ and in $\SS^{p-1}\setminus K_{4v}$, we get that $\bar{\nu}_{\infty}(\SS^{p-1}\setminus K_{4v})\leq\lim_{t\to\infty}\bar{\nu}_{t}(\SS^{p-1}\setminus K_{4v})=0$ since $\SS^{p-1}\setminus K_{4v}$ is open and by the properties of weak convergence of measures (Portmanteau Theorem). Since this holds for $v$ arbitrarily close to , it follows that $\bar{\nu}_{\infty}$ is concentrated on $(J^{\prime}_{\infty})^{-1}(M)$.

Step 4: conclusion. We have proved the two complementary slackness properties, so by Proposition D.5, the pair $(\bar{\nu}_{\infty},p(\infty))$ satisfies the optimality conditions, which concludes the proof.

D.2 Proof of Corollary 3.2

On the other hand, direct computations using the fact that $\nabla S(u)\in\Delta^{n-1}$ leads to

Combining both gives the total derivative

Moreover, by Lemma D.7, we have $|S_{\|\nu_{t}\|}(u)-\min_{i}u|\leq\epsilon/3$ for $\|\nu_{t}\|$ large enough, uniformly for $u$ in a compact set. Hence for all $m\geq m_{0}$,

It remains to show that for all $C>0$, there exists $m_{0}$ such that if $m\geq m_{0}$ then $\lim\inf_{t\to\infty}\|\nu_{t,m}\|\geq C$, so that we deduce from the above

Since $\epsilon$ is arbitrary, it would follow that $\lim_{m\to\infty}\lim_{t\to\infty}\min_{i\in[n]}\hat{h}(\bar{\nu}_{\infty,t})\geq\gamma_{1}$ which is our claim.

To see this, it is sufficient to notice that since $F(\nu_{t,m})$ is increasing in $t$ for all $m$, and $S(u)-\log(n)\leq\min_{i}u$, we have that for all $C>0$, there exists $t_{0},m_{0}$ such that $\min_{i}\hat{h}(\nu_{t,m})[i]\geq C$ for all $t\geq t_{0}$ and $m\geq m_{0}$.

D.3 Intermediate results

This section contains intermediate results used in the proof of Theorem 3.1.

By minimax duality (Sion, 1958), we can rewrite Eq (4) as the minimax problem

and it admits (at least) a saddle point $(\nu^{\star},p^{\star})$. Moreover, the optimality conditions are necessary and sufficient for the right-hand side to equal the left-hand side.

Let us now prove some useful properties of the function

which is a soft-min function under assumption (A2), as shown in the next lemma.

As a consequence, $\exp\left(-\beta\left(S_{\beta}(\bar{u})-m\right)\right)\to\frac{1}{n}\#\arg\min_{i}\bar{u}_{i}\in{]0,1]}$. Thus, $S_{\beta}(\bar{u})\to m$.

Under Assumption (A2), let $u(t)$ be a sequence such that $S(u(t))$ is lower bounded and $\nabla S(u(t))$ converges. Then $\lim_{t\to\infty}\nabla S(u(t))\neq 0$.

We analyze separately two cases, whether there is $i_{0}\in[n]$ such that $u_{i_{0}}(t)$ is upper bounded or not. In the first case, we have

The next lemma is adapted from Gunasekar et al. (2018a, Lemma 8) and exploits the fact that the gradient of $S$ is a soft-argmax.

Taking any $\gamma\in{]\min_{i}\bar{u}_{i}(\infty),\bar{u}_{i_{0}}(\infty)[}$, we have

Appendix E Appendix for Section 4

Let us recall Proposition 4.1 and prove it.

Let $a_{j}(t)=r_{j}(t)^{2}/m$ for $j\in[m]$, $\beta(t)=\|a(t)\|_{1}$ and $\bar{a}(t)=a(t)/\beta(t)$. For the step-sizes $\eta(t)=1/(16\|z\|_{\infty}\sqrt{t+1})$ and a uniform initialization $r(0)\propto\mathbf{1}$, it holds

where $\gamma_{1}^{(m)}:=\max_{a\in\Delta^{m-1}}\min_{i\in[n]}z_{i}^{\top}a$ and the last sum is uniformly bounded in $t$ as soon as $\gamma_{1}^{(m)}>0$.

Let us first prove that the normalized dynamics $\bar{a}(t)$ satisfies the (perturbed) online mirror ascent recursion (on the simplex, with the entropy mirror map):

We mention that it is perturbed because for the plain online mirror ascent, the multiplicative term in the first line would be $\exp(4\eta(t)\nabla G_{\beta(t)}(\bar{a}(t))$, so we have second order corrections in $\eta(t)$.

Since $\nabla_{j}F(r(t))=(2/m)r_{j}(t)\nabla G_{1}(a(t))$, it follows

Now using the fact that $\nabla G_{1}(a)=\nabla G_{\|a\|_{1}}(a/\|a\|_{1})$, it follows

which are the online mirror ascent updates, with a second-order error term $e(t)_{j}=\log(1+4\eta(s)g_{j}(t)+4\eta(t)^{2}g_{j}(t)^{2})-4\eta(t)g_{j}(t)$. Notice that $\forall t$, $\|g(t)\|_{\infty}\leq\|z\|_{\infty}$ so if we assume that $\eta(t)\leq 1/(16\|z\|_{\infty})$ we have $\eta(t)\|g(t)\|_{\infty}\leq 1/16$. Using the inequality $|\log(1+u)-u|\leq u^{2}$ for $|u|\leq 1/2$, we get by applying it with $u_{j}=4\eta(t)g_{j}(t)+4\eta(t)^{2}g_{j}(t)^{2}$,

We now follow the usual proof of mirror ascent from Bubeck (2015, Thm. 4.2) (or Beck and Teboulle (2003) for the variable step-size case) and including this error term leads to, for all $\bar{a}^{*}\in\Delta^{m-1}$,

We get a telescopic sum and using the concavity of each $G_{\beta}$,

With our choice of initialization, $D(\bar{a}^{*},\bar{a}(0))\leq\log(m)$. Let us choose $\eta(t)=\tau/\sqrt{t+1}$. Using the inequalities

In particular, with the choice $\tau=1/(16\|z\|_{\infty})$, we get

where we used $\sqrt{t}/4\leq\sqrt{t+1}-1$ to simplify the expression. Finally, using inequality (14), we have

In the next result, we show that the norm of the iterates grows to $+\infty$ and that $\sum_{s=0}^{t-1}\frac{1}{\beta(s)\sqrt{s+1}}$ is finite. For simplicity, we do not track the constants.

Under the assumptions of Proposition 4.1, we have that $\beta(t)\to\infty$ and $\sum_{s=0}^{t-1}\frac{1}{\beta(s)\sqrt{s+1}}$ is bounded uniformly in $t$.

where $C$ only depend on $\|z\|_{\infty}$. By summing we get

where we only track the dependency in $t$ and $\alpha$. On the other hand, using inequality (14),

Taking for instance $\alpha=\sqrt{t}$ shows that $\sum_{s=0}^{t-1}\frac{\beta(s)}{\sqrt{s+1}}\gtrsim t-\log(t)$ and thus $\beta(t)\to\infty$. By Lemma E.5 then $\beta(t)$ is increasing for $t\geq t_{0}$ and grows to $\infty$ at a super-polynomial rate and the conclusion follows.

We now prove the asymptotic rate of growth of the norm of the iterates.

If $\eta(t)\asymp 1/\sqrt{t+1}$ and $\beta(t)\to\infty$, then $\beta(t)$ is increasing for $t$ large enough and

The result follows since $\sum_{s=0}^{t-1}(t+1)^{-1/2}\geq 2(\sqrt{t-1}-1)$.

Appendix F Appendix to Section 5

Let us recall Proposition 5.1 and prove it.

Let $a(t)=r(t)/m$, $\beta(t)=\max\{1,\max_{0\leq s\leq t}\sqrt{m}\|a(t)\|_{2}\}$ and $\bar{a}(t)=a(t)/\beta(t)$ and assume that $\gamma_{2}^{(m)}$ is positive. For the step-sizes $\eta(t)=\beta(t){\sqrt{2}}/(\|z\|_{\infty}\sqrt{t+1})$ and initialization $r(0)=0$, it holds

Using the fact that $a(t)=r(t)/m$ and $m\nabla F(r)=\nabla G(r/m)$ we have

Finally, since $\beta(t)\max\{1,\sqrt{m}\|b(t+1)\|_{2}\}=\max\{\beta(t),\sqrt{m}\|a(t+1)\|_{2}\}=\beta(t+1)$ it follows that $\bar{a}(t+1)=b(t+1)/\max\{1,\sqrt{m}\|b(t+1)\|_{2}\}$. Thus $\bar{a}(t)$ follows the iterations $\bar{a}(0)=0$ and

Now using the bound of Eq. (14), it follows

since $\nabla G(a(t))\in\Delta^{m-1}$. It follows that $\beta(t)\geq\sqrt{m}\|a(t)\|_{2}\geq m\gamma_{2}^{(m)}\sum_{s=0}^{s-1}\eta(s)$. Using the fact that $\beta(t)\geq C\sum_{s=0}^{s-1}\beta(s)/\sqrt{s+1}$ with $C=\gamma_{2}^{(m)}\sqrt{2}/\|z\|_{\infty}$ and the bound $\sum_{s=0}^{s-1}1/\sqrt{s+1}\geq 2(\sqrt{t-1}-1)$, it follows that $\beta(t)\geq 2C\sqrt{t+1}/\sqrt{6}$ and thus

Plugging into the previous bound gives the conclusion. Note that we did not attempt to make the lower bound on $\beta(t)$ tight.

Appendix G Appendix to Section 6: generalization bounds

With the notations of Section 6, let us first lower bound the margins in $\mathcal{F}_{1}$ and in $\mathcal{F}_{2}$.

Assume that $\|x_{i}\|_{2}\leq R$ for $i\in[n]$. For any $\epsilon\in(0,1)$ and $r\in[d]$, there exists $C(r),C_{\epsilon}(r)>0$ such that

This function is $4/\Delta_{r}(S_{n})$-Lipschitz continuous, satisfies $\|f\|_{\infty}\leq 2$ and $y_{i}f(x_{i})=2$ for all $i\in[n]$. Let us first consider the case $r=d$. Using the approximation results of Lipschitz functions in $\mathcal{F}_{2}$ from Bach (2017a, Prop. 6), we know that if $N>0$ is larger than a constant independent of $\Delta_{d}(S_{n})$ and satisfies

where $\eta=\max\{2,4R/\Delta_{d}(S_{n})\}=4R/\Delta_{d}(S_{n})$, then there exists $\hat{f}$ such that $\|\hat{f}\|_{\mathcal{F}_{2}}\leq N$ and $\sup_{\|x\|_{2}\leq R}|\hat{f}(x)-f_{d}(x)|\leq 1$. Since $\hat{f}/N$ is feasible for the $\mathcal{F}_{2}$-max-margin problem Eq. (5), this shows that $\gamma_{2}\leq 1/N$ and it remains to estimate how large $N$ must be. In the next computations, the dimension dependent constant $C(d)$ might change from line to line. Using the bound $\log(u)\leq C(\epsilon,d)u^{\epsilon/(d+1)}$ for $\epsilon>0$, we obtain the stronger condition on $N$:

This is a direct application of the margin-based generalization bounds of Theorem G.5, using that for any $f\in\mathcal{F}_{1}$ with $\|f\|_{\mathcal{F}_{1}}\leq 1$,

Assume that $\forall f\in\mathcal{F}$ we have $\sup_{x}|f(x)|\leq C$. Then, with probability at least $1-\delta$ over the sample, for all margins $\gamma>0$ and all $f\in\mathcal{F}$ we have

Appendix H Proof for the ReLU case

In this appendix, we detail how to rigorously cover the case of ReLU networks, i.e. models as in Eq. (2) with a feature function of the form

We assume that the input distribution has the following properties:

The assumption on $\rho$ excludes discrete measures. Also, the continuity assumption on $y$ means that the sets where $y=+1$ and where $y=-1$ are disconnected, and this implies that they are at a positive distance from each other since the level-sets $\{y=1\}$ and $\{y=-1\}$ are both compact and with empty intersection. This distribution $\rho$ could be for instance the population distribution of the input, or also could be obtained by taking the expectation over small perturbations of the input training set, which is a well-known smoothing technique (e.g., Duchi et al. (2012)). This leads to the definition of the population smooth-margin

defined for $f\in\mathcal{C}(\mathcal{X})$, where $\mathcal{C}(\mathcal{X})$ is the space of continuous and real-valued functions on $\mathcal{X}$ endowed with the supremum norm. Let us give some facts about this function.

Let $x^{*},x_{0}\in\mathcal{X}$ be such that $x^{*}\in\arg\min_{x\in\mathcal{X}}f_{\infty}$ and $x_{0}\notin\arg\min f_{\infty}$. By continuity of $f_{\infty}$ and uniform convergence, there exists $\epsilon,\delta,t_{0}>0$ such that $\forall t\geq t_{0}$, it holds $f_{t}(x)\geq f(x^{*})+\epsilon$ for all $x\in B_{\delta}(x_{0})$ and $f_{t}(x)\leq f(x^{*})+\epsilon/2$ for all $x\in B_{\delta}(x^{*})$, where $B_{\delta}(x)$ is the open ball of radius $\delta$ centered at $x$. It follows

This shows that $x_{0}\notin\operatorname{spt}(\lambda)$. Since this is true for any $x_{0}\notin\arg\min f_{\infty}$, it follows that $\operatorname{spt}(\lambda)\subset\arg\min_{x\in\mathcal{X}}f_{\infty}(x)$.

Our next step is to gather some regularity properties of $F^{\prime}_{\mu}$. Let us consider the sets

We endow $\SS_{\pm}$ with its Riemannian geometry inherited from the sphere (it is a disconnected manifold with two connected components $\SS_{+}$ and $\SS_{-}$) and let us denote $J^{\prime}_{\mu}$ the restriction of $F^{\prime}_{\mu}$ to $\SS_{\pm}$.

Let us now show that $\nabla F^{\prime}_{\mu}$ is Lipschitz continuous on $\SS_{\pm}$, which implies the same for $\nabla J^{\prime}_{\mu}$. Since this is immediate for the component $\nabla_{b}$ so let us focus on $\nabla_{a}$. For $(a,b),(\bar{a},b)\in\SS_{+}$, defining $H_{a,\bar{a}}=\{x\in\mathcal{X}\;;\;(a^{\top}x)_{+}^{0}\neq(\bar{a}^{\top}x)_{+}^{0}\}$, we have

Existence of a Wasserstein gradient flow.

After these preliminaries, we are in position to show the existence of Wasserstein gradient flows for certain initializations. In this section, we will not attempt to prove uniqueness of the dynamics.

So far, we have proved that for all $m$, there exists a Wasserstein gradient flow $\mu_{t,m}$ starting from $\mu_{0,m}$, defined on $[0,\infty[$ and such that $\operatorname{spt}(\mu_{t,m})\subset D$ for $t\geq 0$. The rest of the proof follows that in (Chizat and Bach, 2018, Thm. 2.6) where we extract a (weak) limit curve $\mu_{t}$ by compactness and show that it satisfies the Wasserstein gradient flow equation Eq. (8). The only technical point is Step.(iii) in that proof, which is taken care of by the last claim of Lemma H.5. Note that here we do not prove uniqueness of the Wasserstein gradient flow, only its existence.

Implicit bias for ReLU networks.

Let us restate Theorem 3.3 in a slightly more general form and making Assumption (*) explicit. We recall that $\nu_{t}:=\Pi_{2}(\mu_{t})$ and $\bar{\nu}_{t}:=\nu_{t}/(\nu_{t}(\SS^{p-1}))$.

if $\bar{\nu}_{t}=\Pi_{2}(\mu_{t})/([\Pi_{2}(\mu_{t})](\SS^{p-1}))$ converges weakly to some $\bar{\nu}_{\infty}\in\mathcal{P}(\SS^{p-1})$,

if $\nabla S[\hat{h}(\mu_{t})]$ converges weakly to some $\lambda_{\infty}\in\mathcal{P}(\mathcal{X})$, and

then $h(\bar{\nu}_{\infty},\cdot)$ is a maximizer for the $\mathcal{F}_{1}$-max-margin problem $\max_{\|f\|_{\mathcal{F}_{1}}\leq 1}\min_{x\in\operatorname{spt}(\rho)}y(x)f(x).$

Before proving this theorem, let us discuss its assumptions. First, the assumption on the initialization given here is satisfied by $\mu_{0}$ given in Theorem 3.3 (which is an example given for the sake of concreteness). The conditions in the two first bullets are similar to those of Theorem 3.1 and just require the uniqueness of limits of some sequences that live in compact spaces. In Assumption (*), the fact that the Morse-Sard property holds for $J^{\prime}_{\infty}$ was already required in Chizat and Bach (2018) and it is an open question to guarantee that this property holds in this context (where $J^{\prime}_{\infty}$ is potentially an infinite sum of subanalytic functions, instead of a finite sum as in Theorem 3.3). Finally, the most undesirable assumption is perhaps the convergence of $J^{\prime}_{\mu_{t}}$ to $J^{\prime}_{\infty}$ in $\mathcal{C}^{1}(\SS_{\pm})$. The fact that it converges in $\mathcal{C}(\SS_{\pm})$ can be shown a priori, so the assumption is really on the uniform convergence of the gradient. In particular, it requires $J^{\prime}_{\infty}$ to be continuously differentiable, which is for instance not true if $\lambda_{\infty}$ is a discrete measure.

Step 1. The argument of Theorem 3.1 goes through if we replace $\SS^{p-1}$ by the set $\SS_{\pm}$, in particular thanks to Assumption (*). Lemma H.5 and the positive $2$-homogeneity guarantee that the restriction of the flow $X$ to $D$ is a diffeomorphism.

Step 2. Since we focus on the exponential loss, we can directly apply Lemma H.3, which gives $\operatorname{spt}(\lambda_{\infty})\subset\arg\min_{x\in\operatorname{spt}(\rho)}y(x)h(\bar{\nu}_{\infty},x)$.

Steps 3. The argument of Theorem 3.1 goes through, again replacing $\SS^{p-1}$ by the set $\SS_{\pm}$.

Steps 4. We conclude with the optimality conditions in Proposition H.13 and the structure of the minimizers given in Proposition B.1.

The maximization problem (17) admits global maximizers $\nu^{\star}\in\mathcal{P}(\SS^{p-1})$. Moreover, a measure $\nu^{\star}\in\mathcal{P}(\SS^{p-1})$ is a global maximizer of (4) if and only if $\nu^{\star}\in\mathcal{P}(\SS^{p-1})$ and there exists $\lambda^{\star}\in\mathcal{P}(\mathcal{X})$ such that