Directional convergence and alignment in deep learning

Ziwei Ji, Matus Telgarsky

Introduction

Recent efforts to rigorously analyze the optimization of deep networks have yielded many exciting developments, for instance the neural tangent (Jacot et al., 2018; Du et al., 2018; Allen-Zhu et al., 2018; Zou et al., 2018) and mean-field perspectives (Mei et al., 2019; Chizat and Bach, 2018). In these works, it is shown that small training or even testing error are possible for wide networks.

The above theories, with finite width networks, usually require the weights to stay close to initialization in certain norms. By contrast, practitioners run their optimization methods as long as their computational budget allows (Shallue et al., 2018), and if the data can be perfectly classified, the parameters are guaranteed to diverge in norm to infinity (Lyu and Li, 2019). This raises a worry that the prediction surface can continually change during training; indeed, even on simple data, as in Figure 1, the prediction surface continues to change after perfect classification is achieved, and even with large width is not close to the maximum margin predictor from the neural tangent regime. If the prediction surface never stops changing, then the generalization behavior, adversarial stability, and other crucial properties of the predictor could also be unstable.

In this paper, we resolve this worry by guaranteeing stable convergence behavior of deep networks as training proceeds, despite this growth of weight vectors to infinity. Concretely:

Directional convergence: the parameters converge in direction, which suffices to guarantee convergence of many other relevant quantities, such as the prediction margins.

Alignment: when gradients exist, they converge in direction to the parameters, which implies various margin maximization results and saliency map convergence, to name a few.

\bm{+}”), trained via gradient descent. Figure 1(a) shows the prediction surface reached by freezing activations, which is also the prediction surface of the corresponding Neural Tangent Kernel (NTK) maximum margin predictor (Soudry et al., 2017). Figure 1(b) shows the same network, but now without frozen activations, at the first moment with perfect classification. Training this network much longer converges to Figure 1(c). We show that the network parameters WtW_{t} converge in direction, meaning the normalized iterates \nicefracWt∥Wt∥\nicefrac{{W_{t}}}{{\|W_{t}\|}} converge. Details are deferred to Section 3, but here is a brief overview.

Our networks are LL-positively homogeneous in the parameters, meaning scaling the parameters by c>0c>0 scales the predictions by cLc^{L}, and definable in some oo-minimal structure, a mild technical assumption which we will describe momentarily. Our networks can be arbitrarily deep with many common types of layers (e.g., linear, convolution, ReLU, and max-pooling layers), but homogeneity rules out some components such as skip connections and biases, which all satisfy definability.

Under these conditions, we prove the following result, without any other assumptions about the distribution of the parameters or the width of the network (cf. Theorem 3.1):

The curve swept by \nicefracWt∥Wt∥\nicefrac{{W_{t}}}{{\|W_{t}\|}} has finite length, and thus \nicefracWt∥Wt∥\nicefrac{{W_{t}}}{{\|W_{t}\|}} converges.

Our main corollary is that prediction margins converge (cf. Section 3), meaning convergence of the normalized per-example values \nicefracyiΦ(xi;Wt)∥Wt∥L\nicefrac{{y_{i}\Phi(x_{i};W_{t})}}{{\|W_{t}\|^{L}}}, where yiy_{i} is the label and Φ(xi;Wt)\Phi(x_{i};W_{t}) is the prediction on example xix_{i}. These quantities are central in the study of generalization of deep networks, and their stability also implies stability of many other useful quantities (Bartlett et al., 2017; Jiang et al., 2019, 2020). As an illustration of directional convergence and margin convergence, we plot the margin values for all examples in the standard cifar data against training iterations in Figure 2; these trajectories exhibit strong convergence behavior, both within our theory (a modified homogeneous AlexNet, as in Figure 2(a)), and outside of it (DenseNet, as in Figure 2(b)).

Directional convergence is often assumed throughout the literature (Gunasekar et al., 2018a; Chizat and Bach, 2020), but has only been established for linear predictors (Soudry et al., 2017). It is tricky to prove because it may still be false for highly smooth functions: for instance, the homogeneous Mexican Hat function satisfies all our assumptions except definability, and can be adjusted to have arbitrary order of continuous derivatives, but its gradient flow does not converge in direction, instead it spirals (Lyu and Li, 2019). To deal with similar pathologies in many branches of mathematics, the notion of functions definable in some o-minimal structure was developed: these are rich classes of functions built up to limit oscillations and other bad behavior. Using techniques from this literature, we build general tools, in particular unbounded nonsmooth Kurdyka-Łojasiewicz inequalities, which allows us to prove directional convergence, and may also be useful outside deep learning. More discussion on the o-minimal literature is given in Section 1.3, technical preliminaries are introduced in Section 2, and a proof overview is given in Section 3, with full details in the appendices.

2 Second result: gradient alignment

Our second contribution, in Section 4, is that if the network has locally Lipschitz gradients, then these gradients also converge, and are aligned to the gradient flow path (cf. Theorem 4.1).

The gradient flow path, and the gradient of the risk along the path, converge to the same direction.

As a practical consequence of this, recall the use of gradients within the interpretability literature, specifically in saliency maps (Adebayo et al., 2018): if gradients do not converge in direction then saliency maps can change regardless of the number of iterations used to produce them. As a theoretical consequence, directional convergence and alignment imply margin maximization in a variety of situations: this holds in the deep linear case, strengthening prior work (Gunasekar et al., 2018b; Ji and Telgarsky, 2018a), and in the 2-homogeneous network case, with an assumption taken from the infinite width setting (Chizat and Bach, 2020), but presented here with finite width.

3 Further related work

Our analysis is heavily inspired and influenced by the work of Lyu and Li (2019), who studied margin maximization of homogeneous networks, establishing monotonicity of a smoothed margin, a quantity we also use. However, they did not prove directional convergence but instead must use subsequences. Their work also left open alignment and global margin maximization.

A standard approach to resolve directional convergence and similar questions is to establish that the objective function in question is definable in some o-minimal structure, which as mentioned before, limits oscillations and other complicated behavior. This literature cannot be directly applied to our setting, owing to a combination of nonsmooth layers like the ReLU and max-pooling, and the exponential function used in the cross entropy loss, and as a result, our proofs need to rebuild many o-minimal results from the ground up.

In more detail, an important problem in the o-minimal literature is the gradient conjecture of René Thom: it asks when the existence of lim⁡t→∞Wt=z\lim_{t\to\infty}W_{t}=z further implies lim⁡t→∞\nicefrac(Wt−z)∥Wt−z∥\lim_{t\to\infty}\nicefrac{{(W_{t}-z)}}{{\|W_{t}-z\|}} exists, and was established in various definable scenarios by Kurdyka et al. (2000a, 2006) via related Kurdyka-Łojasiewicz inequalities (Kurdyka, 1998). The underlying proof ideas can also be used to analyze lim⁡t→∞\nicefracWt∥Wt∥\lim_{t\to\infty}\nicefrac{{W_{t}}}{{\|W_{t}\|}} when the weights go to infinity (Grandjean, 2007). However, the prior results require the objective function to be either real analytic, or definable in a “polynomially-bounded” o-minimal structure. The first case causes the aforementioned nonsmoothness issue, and excludes many common layers in deep learning such as the ReLU and max-pooling. The second case excludes the exponential function, and means the logistic and cross-entropy losses cannot be handled. To resolve these issues, we had to redo large portions of the o-minimality theory, such as the nonsmooth unbounded Kurdyka-Łojasiewicz inequalities that can handle the exponential/logistic loss, as presented in Section 3.

Alignment.

As discussed in Section 4, alignment implies the gradient flow reaches a stationary point of the limiting margin maximization objective, and therefore is related to various statements and results throughout the literature on implicit bias and margin maximization (Soudry et al., 2017; Ji and Telgarsky, 2018b). This stationary point perspective also appears in some nonlinear works, for instance in the aforementioned work on margins by Lyu and Li (2019), which showed that subsequences of the gradient flow converge to such stationary points; in addition to fully handling the gradient flow, the present work also differs in that alignment is in general a stronger notion, in that it is unclear how to prove alignment as a consequence of convergence to KKT points. Additionally, alignment can still hold when the objective function is not definable and directional convergence is false, for example on the homogeneous Mexican hat function, which cannot be handled by the approach in (Lyu and Li, 2019, Appendix J). As a final pointer to the literature, many implicit bias works explicitly assume directional convergence and some version of alignment (Gunasekar et al., 2018b; Chizat and Bach, 2020), but neither do these works indicate a possible proof, nor do they provide conclusive evidence.

4 Experimental overview

The experiments in Figures 1 and 2 are performed in as standard a way as possible to highlight that directional convergence is a reliable property; full details are in Appendix A. Briefly, Figure 1 uses synthetic data and vanilla gradient descent (no momentum, no weight decay, etc.) on a 10,000 node wide 2-layer squared ReLU network and its Neural Tangent Kernel classifier; by using the squared ReLU, both our directional convergence and our alignment results apply. Figure 2 uses standard cifar firstly with a modified homogeneous AlexNet and secondly with an unmodified DenseNet, respectively inside and outside our assumptions. SGD was used on cifar due to training set size, and seeing how directional convergence still seems to occur, suggests another open problem.

Preliminaries and assumptions

If ff is locally Lipschitz, it holds that ff is differentiable a.e. (Borwein and Lewis, 2000, Theorem 9.1.2). The Clarke subdifferential of ff at x∈Dx\in D is defined as

which is nonempty convex compact (Clarke, 1975), and if ff is continuously differentiable at xx, then ∂f(x)={∇f(x)}\partial f(x)=\{\nabla f(x)\}. Vectors in ∂f(x)\partial f(x) are called subgradients, and we let ∂ˉf(x)\bar{\partial}f(x) denote the unique minimum-norm subgradient:

In the following analysis, we use ∂ˉf\bar{\partial}f in many places that seem to call on ∇f\nabla f.

O-minimal structures and definable functions.

Many natural functions and operations are definable. First of all, definability of functions is stable under algebraic operations, composition, inverse, maximum and minimum, etc. Moreover, Wilkie (1996) proved that there exists an o-minimal structure where polynomials and the exponential function are definable. Consequently, definability allows many common layer types in deep learning, such as fully-connected/convolutional/ReLU/max-pooling layers, skip connections, the cross entropy loss, etc.; moreover, they can be composed arbitrarily As will be discussed later, what is still missing is the handling of the gradient flow on such functions.

The network model.

For any fixed xx, the prediction W↦Φ(x;W)W\mapsto\Phi(x;W) as a function of WW is locally Lipschitz, LL-positively homogeneous for some L>0L>0, and definable in some o-minimal structure including the exponential function.

As mentioned before, homogeneity means that Φ(x;cW)=cLΦ(x;W)\Phi(x;cW)=c^{L}\Phi(x;W) for any c≥0c\geq 0. This means, for instance, that linear, convolutional, ReLU, and max-pooling layers are permitted, but not skip connections and biases. Homogeneity is used heavily throughout the theoretical study of deep networks (Lyu and Li, 2019).

Gradient flow.

Our second assumption is on the initial risk, and appears in prior work (Lyu and Li, 2019).

Directional convergence

We now turn to stating our main result on directional convergence and sketching its analysis. As Sections 2 and 2 imply ∥Wt∥→∞\|W_{t}\|\to\infty (Lyu and Li, 2019), we study the normalized flow W~t:=Wt/∥Wt∥\widetilde{W}_{t}\mathrel{\mathop{\ordinarycolon}}=W_{t}/\|W_{t}\|, whose convergence is a formal way of studying the directional convergence of WtW_{t}. As mentioned before, directional convergence is false in general (Lyu and Li, 2019), but definability suffices to ensure it. Throughout, for general nonzero WW, we will use W~:=W/∥W∥\widetilde{W}\mathrel{\mathop{\ordinarycolon}}=W/\|W\|.

A direct consequence of Theorem 3.1 is the convergence of the margin distribution (i.e., normalized outputs). Due to homogeneity, for any nonzero WW, we have pi(W)/∥W∥L=pi(W~)p_{i}(W)/\|W\|^{L}=p_{i}(\widetilde{W}), and thus the next result follows from Theorem 3.1.

Next we give a proof sketch of Theorem 3.1; the full proofs of the Kurdyka-Łojasiewicz inequalities (Sections 3.1 and 3.1) are given in Section B.3, while the other proofs are given in Appendix C.

The smoothed margin introduced in (Lyu and Li, 2019) is crucial in our analysis: given W≠0W\neq 0, let

Given any function ff which is locally Lipschitz around a nonzero WW, let

For simplicity, in the discussion here we consider the case that all subgradients in Section 3.1 are nonzero, with the general case handled in the full proofs in the appendices. Then Section 3.1 implies

Given a locally Lipschitz definable function ff with an open domain D⊂{x | ∥x∥>1}D\subset\mathinner{\left\{x\,\middle|\,\|x\|>1\right\}}, for any c,η>0c,\eta>0, there exists ν>0\nu>0 and a definable desingularizing function Ψ\Psi on [0,ν)[0,\nu) such that

Given a locally Lipschitz definable function ff with an open domain D⊂{x | ∥x∥>1}D\subset\mathinner{\left\{x\,\middle|\,\|x\|>1\right\}}, for any λ>0\lambda>0, there exists ν>0\nu>0 and a definable desingularizing function Ψ\Psi on [0,ν)[0,\nu) such that

Alignment between the gradient flow path and gradients

Theorem 3.1 gave our directional convergence result, namely that the normalized iterate Wt/∥Wt∥W_{t}/\|W_{t}\| converges to some direction. Next we show and discuss our alignment result, that if all pip_{i} have locally Lipschitz gradients, then along the gradient flow path, −∇L(Wt)-\nabla\mathcal{L}(W_{t}) converges to the same direction as WtW_{t}.

Under Sections 2 and 2, if all pip_{i} further have locally Lipschitz gradients, then −∇L(Wt)-\nabla\mathcal{L}(W_{t}) and WtW_{t} converge to the same direction, meaning the angle between WtW_{t} and −∇L(Wt)-\nabla\mathcal{L}(W_{t}) converges to zero. If all pip_{i} are twice continuously differentiable, then the same result holds without the definability condition (cf. Section 2).

Below we first sketch the proof of Theorem 4.1, with full details in Appendix D, and then in Section 4.2 present a few global margin maximization consequences, which are proved in Appendix E.

Recall that lim⁡t→∞α(Wt)/∥Wt∥L=a\lim_{t\to\infty}\alpha(W_{t})/\|W_{t}\|^{L}=a. The first observation is that α(Wt)\alpha(W_{t}), the smoothed margin function, asymptotes to the exact margin min⁡1≤i≤npi(Wt)\min_{1\leq i\leq n}p_{i}(W_{t}) which is LL-positively homogeneous. Therefore α\alpha is asymptotically LL-positively homogeneous, and formally we can show

which can be viewed as an asymptotic version of Euler’s homogeneous function theorem (cf. Appendix C). Consequently, the inner product between ∇α(Wt)/∥Wt∥L−1\nabla\alpha(W_{t})/\|W_{t}\|^{L-1} and W~t\widetilde{W}_{t} converges.

Let θt\theta_{t} denote the angle between WtW_{t} and −∇L(Wt)-\nabla\mathcal{L}(W_{t}), which is also the angle between WtW_{t} and ∇α(Wt)\nabla\alpha(W_{t}), since ∇L(Wt)\nabla\mathcal{L}(W_{t}) and ∇α(Wt)\nabla\alpha(W_{t}) point to opposite directions by the chain rule. By (Lyu and Li, 2019, Corollary C.10), given any ϵ>0\epsilon>0, there exists a time tϵt_{\epsilon} such that θtϵ<ϵ\theta_{t_{\epsilon}}<\epsilon. The question is whether such a small angle can be maintained after tϵt_{\epsilon}. This is not obvious since, as mentioned above, the smoothed margin α(Wt)\alpha(W_{t}) asymptotes to the exact margin min⁡1≤i≤npi(Wt)\min_{1\leq i\leq n}p_{i}(W_{t}), which may be nondifferentiable even with smooth pip_{i}, due to nondifferentiability of the minimum. Consequently, the exact margin may have discontinuous Clarke subdifferentials, and since the smoothed margin asymptotes to it, it is unclear whether θt→0\theta_{t}\to 0. (This point was foreshadowed earlier, where it was pointed out that alignment is not a clear consequence of convergence to stationary points of the margin maximization objective.)

To handle this, the key to our analysis is the potential function J(W):= ⁣∥∇α(Wt)∥2/∥Wt∥2L−2\mathcal{J}(W)\mathrel{\mathop{\ordinarycolon}}=\mathinner{\!\left\lVert\nabla\alpha(W_{t})\right\rVert}^{2}/\|W_{t}\|^{2L-2}. Suppose at time tt, it holds that ⟨∇α(Wt)/∥Wt∥L−1,W~t⟩\left\langle\nabla\alpha(W_{t})/\|W_{t}\|^{L-1},\widetilde{W}_{t}\right\rangle is close to aLaL, and θt\theta_{t} is very small. If θt′\theta_{t^{\prime}} becomes large again at some t′>tt^{\prime}>t, it must follows that J(Wt′)\mathcal{J}(W_{t^{\prime}}) is much larger than J(Wt)\mathcal{J}(W_{t}). We prove that this is impossible, by showing that

and thus Theorem 4.1 follows. The proof of eq. 4.3 is motivated by the dual convergence analysis in (Ji and Telgarsky, 2019), and also uses the positive homogeneity of ∇pi\nabla p_{i} and ∇2pi\nabla^{2}p_{i} (which exist a.e.).

2 Main alignment consequence: margin maximization

A variety of (global) margin maximization results are immediate consequences of directional convergence and alignment. This subsection investigates two examples: deep linear networks, and shallow squared ReLU networks.

Thanks to directional convergence and alignment (cf. Theorems 3.1 and 4.1), the proof boils down to writing down the gradient expression for each layer and doing some algebra.

A more interesting example is a certain 2-homogeneous case, which despite its simplicity is a universal approximator; this setting was studied by Chizat and Bach (2020), who considered the infinite width case, and established margin maximization under assumptions of directional convergence and gradient convergence. Unfortunately, it is not clear if Theorems 3.1 and 4.1 can be applied to fill these assumptions, since they do not handle infinite width, and indeed it is not clear if infinite width networks or close relatives are definable in an o-minimal structure. Instead, here we consider the finite width case, albeit with an additional assumption.

Following (Chizat and Bach, 2020, S-ReLU), organize WtW_{t} into mm rows (wj(t))j=1m(w_{j}(t))_{j=1}^{m}, with normalizations θj(t):=wj(t)/∥wj(t)∥\theta_{j}(t)\mathrel{\mathop{\ordinarycolon}}=w_{j}(t)/\|w_{j}(t)\| where θj(t)=0\theta_{j}(t)=0 when ∥wj(t)∥=0\|w_{j}(t)\|=0, and consider

whereby pi(W)=∑jφij(wj)p_{i}(W)=\sum_{j}\varphi_{ij}(w_{j}), and Φ\Phi, pip_{i}, and φij\varphi_{ij} are all 2-homogeneous and definable. (The “(−1)j(-1)^{j}” may seem odd, but is an easy trick to get universal approximation without outer weights.)

(Global guarantee.) Suppose the covering condition: there exist t0t_{0} and ϵ>0\epsilon>0 with

The first part (the “local guarantee”) characterizes the limiting margin as the maximum margin of a linear problem obtained by taking the limiting directions (θˉj)j=1m(\bar{\theta}_{j})_{j=1}^{m} and treating the resulting φij(θˉj)\varphi_{ij}(\bar{\theta}_{j}) as features. The quality of this margin is bad if the limiting directions are bad, and therefore we secondly (the “global guarantee”) consider a case where our margin is nearly as good as the infinite width global max margin value as defined by (Chizat and Bach, 2020, eq. (5)); see discussion therein for a justification of this choice, and moreover calling it the globally maximal margin.

The covering condition deserves further discussion. In the infinite width setting, it holds for all ϵ>0\epsilon>0 assuming directional convergence (Chizat and Bach, 2020, Proof of Theorem D.1), but cannot hold in such generality here as we are dealing with finite width. Similar properties have appeared throughout the literature: Wei et al. (2018, Section 3) explicitly re-initialized network nodes to guarantee a good covering, and more generally (Ge et al., 2015) added noise to escape saddle points in general optimization problems.

Concluding remarks and open problems

In this paper, we established that the normalized parameter vectors \nicefracWt∥Wt∥\nicefrac{{W_{t}}}{{\|W_{t}\|}} converge, and that under an additional assumption of locally Lipschitz gradients, the gradients also converge and align with the parameters.

There are many promising avenues for future work based on these results. One basic line is to weaken our assumptions: dropping homogeneity to allow for DenseNet and ResNet, and analyzing finite-time methods like (stochastic) gradient descent, and moreover their rates of convergence. We also handled only the binary classification case, however our tools should directly allow for cross-entropy.

Another direction is into further global margin maximization results, beyond the simple networks in Section 4.2, and into related generalization consequences of directional convergence and alignment.

The authors thank Zhiyuan Li and Kaifeng Lyu for lively discussions during an early phase of the project. The authors are grateful for support from the NSF under grant IIS-1750051, and from NVIDIA via a GPU grant.

References

Appendix A Experimental setup

The goal of the experiments is to illustrate that directional convergence is a clear, reliable phenomenon. Below we detail the setup for the two types of experiments: contour plots in Figure 1, and margin plots in Figure 2 (with ResNet here in Figure 3).

Figure 2 used the standard cifar dataset in its 10 class configuration (Krizhevsky, 2009). There are 50,000 data points, each with 3072 dimensions, organized into 32×3232\times 32 images with 3 color channels.

Models.

A few simple models both inside and outside our technical assumptions were used. All code was implemented in PyTorch (Paszke et al., 2019).

Figure 1 worked with a style of 2-layer network which appears widely throughout theoretical investigations: specifically, there is first a wide linear layer (in our case, 10,00010,000 nodes), then a squared ReLU layer, and then a layer of random signs which is not trained. This squared ReLU network with one trainable layer is 2-homogeneous, and was chosen both to fit with the alignment guarantee in Theorem 4.1, and also to amplify differences with the NTK. Note that this simple architecture is still a universal approximator with non-convex training. Figures 1(b) and 1(c) trained this network, which can be written as x\mapsto\sum_{j}s_{j}\max\mathinner{\bigl{\{}0,\left\langle w_{j},x\right\rangle\bigr{\}}}^{2}, where sj∈±1s_{j}\in\pm 1 are fixed random signs and (wj)j=1m(w_{j})_{j=1}^{m} are the trainable parameters. Figure 1(a) trained the corresponding NTK (Jacot et al., 2018; Du et al., 2018; Allen-Zhu et al., 2018; Zou et al., 2018), meaning the linear predictor obtained by freezing the network activations, which thus has the form x↦∑jsj⟨vj,x⟩max⁡{0,⟨wj,x⟩}x\mapsto\sum_{j}s_{j}\left\langle v_{j},x\right\rangle\max\{0,\left\langle w_{j},x\right\rangle\}, where (wj)j=1m(w_{j})_{j=1}^{m} from before are now fixed, and only (vj)j=1m(v_{j})_{j=1}^{m} are trained.

Figure 2 used convolutional networks. Firstly, Figure 2(a) used “H-AlexNet”, which is based on a simplified version of the standard AlexNet (Krizhevsky et al., 2012) as presented in the PyTorch cifar tutorial (Paszke et al., 2019), but with biases disabled in order to give a homogeneous network. The network ultimately consists of ReLU layers, max-pooling layers, linear layers, and convolutional layers, and is 5-homogeneous. In particular, H-AlexNet satisfies all conditions we need for directional convergence.

The two models outside the assumptions were DenseNet (cf. Figure 2(b) and ResNet (cf. Figure 3), used unmodified from the PyTorch source, namely by invoking torchvision.models.densetnet121 and torchvision.models.resnet18 with argument num_classes=10.

Training.

To help reach such small risk, the main ideas were to rewrite the objective functions to be numerically stable, and secondly to scale the step size by \nicefrac1L(Wt−1)\nicefrac{{1}}{{\mathcal{L}(W_{t-1})}}, which incidentally is consistent with gradient flow on α\alpha with exponential loss, and is moreover an idea found across the margin literature, most notably as the step size used in AdaBoost (Freund and Schapire, 1997). This can lead to some numerical instability, so the step size was reduced if the norm of the induced update was too large, meaning the norm of the gradient times the step size was too large. A much more elaborate numerical scheme was reported by Lyu and Li (2019, Appendix L), but not used here.

One point worth highlighting is the role of SGD, which seems as though it should have introduced a great deal of noise into the plots, and after all is outside the assumptions of the paper (which requires gradient flow, let alone gradient descent). Though not depicted here, experiments in Figure 2 were also tried on subsampled data and full gradients, and Figure 1 was tried with SGD in place of GD; while gradient descent does result in smoother plots, the difference is small overall, leaving the rigorous analysis of directional convergence with SGD as a promising future direction.

Margin plots.

A few further words are in order for the margin plots in Figures 2 and 3.

While margins are well-motivated from generalization and other theoretical perspectives (Bartlett et al., 2017; Jiang et al., 2019, 2020), we also use margin plots as a visual surrogate for prediction surface contour plots from Figure 1, but now for high-dimensional data, even with high-dimensional outputs. In particular, Figures 2 and 3 track the prediction surface but restricted to the training set, showing, in a sense, the output trajectory for each data example. Since the output dimension is 10 classes, we convert this to a single real number via the usual multi-class margin (x,y)↦Φ(x;Wt)y−max⁡j≠yΦ(x;Wt)j(x,y)\mapsto\Phi(x;W_{t})_{y}-\max_{j\neq y}\Phi(x;W_{t})_{j}.

In the case of homogeneous networks, it is natural to normalize this quantity by ∥Wt∥L\|W_{t}\|^{L}; for the inhomogeneous cases DenseNet and ResNet, no such normalization is available. Therefore, for consistency, at each time tt, margins were normalized by the median nonnegative margin across all data.

To show the evolution of the margins most clearly, we sorted margins according to the final margin level, and used this fixed data ordering for all time; as a result, lines in the plot indeed correspond to trajectories of single examples. Moreover, we indexed time by the log of the inverse risk, namely ln⁡\nicefracnL(Wt)\ln\nicefrac{{n}}{{\mathcal{L}(W_{t})}} in our notation. While this may seem odd at first, importantly it washes out the effect of small step-sizes and other implementation choices; and crucially disallows an artificial depiction of directional convergence by choosing rapidly-vanishing step sizes.

Appendix B Results on o-minimal structures

S1\mathcal{S}_{1} is the collection of all finite unions of open intervals and points.

Sn\mathcal{S}_{n} is closed under finite union, finite intersection, and complement.

S\mathcal{S} is closed under Cartesian products: if A∈SmA\in\mathcal{S}_{m} and B∈SnB\in\mathcal{S}_{n}, then A×B∈Sm+nA\times B\in\mathcal{S}_{m+n}.

S\mathcal{S} is closed under projection Πn\Pi_{n} onto the first nn coordinates: if A∈Sn+1A\in\mathcal{S}_{n+1}, then Πn(A)∈Sn\Pi_{n}(A)\in\mathcal{S}_{n}.

A convenient way to construct definable sets and functions is to use first-order formulas:

If AA is a definable set, then “x∈Ax\in A” is a first-order formula.

If ϕ\phi and ψ\psi are first-order formulas, then ϕ∧ψ\phi\wedge\psi, ϕ∨ψ\phi\vee\psi, ¬ϕ\neg\phi and ϕ⇒ψ\phi\Rightarrow\psi are first-order formulas.

Given a first-order formula, the set of free variables which satisfy the formula is definable (Van den Dries and Miller, 1996, Appendix A). The following basic properties of definable sets and functions can then be shown (see (Van den Dries and Miller, 1996; Coste, 2000; Lê Loi, 2010)).

Any composition of definable functions is definable.

Any coordinate permutation of a definable set is definable. Consequently, if the inverse of a definable function exists, it is also definable.

The image and pre-image of a definable set by a definable function is definable. Particularly, given any real-valued definable function ff, all of f−1(0)f^{-1}(0), f−1((−∞,0))f^{-1}\mathinner{\left((-\infty,0)\right)} and f−1((0,∞))f^{-1}\mathinner{\left((0,\infty)\right)} are definable.

Any combination of finitely many definable functions with disjoint domains is definable. For example, the pointwise maximum and minimum of definable functions are definable.

The proofs are standard and omitted. To illustrate the idea, we give a proof of the following standard result on the infimum and supremum operation.

Given a definable set AA, the function dA(x):=inf⁡y∈A∥x−y∥d_{A}(x)\mathrel{\mathop{\ordinarycolon}}=\inf_{y\in A}\|x-y\| is definable, which implies the closure, interior and boundary of AA are definable.

The lower-semicontinuous envelope of a definable function is definable.

is definable, since it is given by the following first-order formula:

Let GfG_{f} denote the graph of ff, and GgG_{g} denote the graph of gg. We can just apply the main claim to the following definable set:

First, the Minkowski sum of two definable sets AA and BB is definable:

Then we can just apply the main claim to the Minkowski sum of the graphs of ff and gg.

Let GfG_{f} denote the graph of ff. If GfG_{f} is definable, then the epigraph is definable:

If the epigraph is definable, then GfG_{f} is definable due to the main claim.

We can just apply the main claim to the set

The closure of AA is just dA−1(0)d_{A}^{-1}(0). The interior of AA is the complement of dAc−1(0)d_{A^{c}}^{-1}(0). The boundary is the difference between the closure and interior.

The epigraph of the lower-semicontinuous envelope of ff is the closure of the epigraph of ff.

As another example, note that the types of networks under discussion are definable.

then all hjh_{j} are definable. It suffices if each output coordinate of gjg_{j} is the minimum or maximum over some finite set of polynomials, which allows for linear, convolutional, ReLU, max-pooling layers and skip connections.

The definability of hjh_{j} can be proved by induction using the fact that definability is preserved under composition. Next, note that the minimum and maximum of a finite set of polynomials is definable. Lastly, note that each output coordinate of linear and convolutional layers can be written as a polynomial of their input and the parameters; each output coordinate of a ReLU layer is the maximum of two polynomials; each output of a max-pooling layer is a maximum of polynomials. Skip connections are allowed by the definition of hjh_{j}. ∎

Below are some useful properties of definable functions.

Section B.1 and Theorem B.1 imply the following result which we need later.

Let z:=lim⁡s→∞γ(s)z\mathrel{\mathop{\ordinarycolon}}=\lim_{s\to\infty}\gamma(s). Since  ⁣∥z−γ(s)∥\mathinner{\!\left\lVert z-\gamma(s)\right\rVert} is definable, either it is for all large enough ss, or it is positive for all large enough ss. In the first case, since γ\gamma is C1C^{1}, it has finite length. In the second case, Theorem B.1 implies that there exists an interval [a,∞)[a,\infty) on which  ⁣∥z−γ(s)∥>0\mathinner{\!\left\lVert z-\gamma(s)\right\rVert}>0 and d ⁣⁡  ⁣∥z−γ(s)∥/d ⁣⁡s<0\operatorname{d\!}\,\mathinner{\!\left\lVert z-\gamma(s)\right\rVert}/\operatorname{d\!}s<0, and thus  ⁣∥γ′(s)∥>0\mathinner{\!\left\lVert\gamma^{\prime}(s)\right\rVert}>0. Let

The existence of the above limits is guaranteed by Section B.1. Note that ⟨u,v⟩\langle u,v\rangle is equal to

Since γ′(s)/ ⁣∥γ′(s)∥→v\gamma^{\prime}(s)/\mathinner{\!\left\lVert\gamma^{\prime}(s)\right\rVert}\to v, given any ϵ>0\epsilon>0, for large enough ss it holds that ⟨γ′(s)/ ⁣∥γ′(s)∥,v⟩≥1−ϵ\left\langle\gamma^{\prime}(s)/\mathinner{\!\left\lVert\gamma^{\prime}(s)\right\rVert},v\right\rangle\geq 1-\epsilon, and thus

which implies that u=vu=v. Since ϵ>0\epsilon>0 was arbitrary, then

which implies that γ\gamma has finite length. ∎

The following Curve Selection Lemma is crucial in proving the Kurdyka-Łojasiewicz inequalities.

We also need the following version at infinity, from (Némethi and Zaharia, 1992, Lemma 2) and (Kurdyka et al., 2000b, Lemma 3.4).

With this in hand, define a curve ρ0:[1,∞)→A\rho_{0}\mathrel{\mathop{\ordinarycolon}}[1,\infty)\to A as

B.2 Clarke subdifferentials

Here we prove the definability of Clarke subdifferential, and a chain rule along arcs which is crucial in our analysis.

is definable, since it is given by the following first-order formula:

and that ∂ˉf(x)\bar{\partial}f(x) denotes the unique minimum-norm subgradient. Similarly to the gradients, the following result holds for the Clarke subdifferentials.

is definable. Moreover, the function D∋x↦∂ˉf(x)D\ni x\mapsto\bar{\partial}f(x) is definable.

Let D′:={x∈D | ∇f(x) exists}D^{\prime}\mathrel{\mathop{\ordinarycolon}}=\mathinner{\left\{x\in D\,\middle|\,\nabla f(x)\textrm{ exists}\right\}}, which is definable. The set AA given by

is also definable. Now by Carathéodory’s Theorem, Γ\Gamma is given by

It then follows from Section B.1 that x↦ ⁣∥∂ˉf(x)∥x\mapsto\mathinner{\!\left\lVert\bar{\partial}f(x)\right\rVert} and x↦∂ˉf(x)x\mapsto\bar{\partial}f(x) are definable. ∎

The following chain rule is important in our analysis; it allows us to use ∂ˉf\bar{\partial}f in many places that seem to call on ∇f\nabla f. It is basically from (Davis et al., 2020, Theorem 5.8 and Lemma 5.2), though we detail how their proof handles our slight extension.

Moreover, for the gradient flow in eq. 2.1, it holds for a.e. t≥0t\geq 0 that d ⁣⁡Wt/d ⁣⁡t=−∂ˉL(Wt)\operatorname{d\!}W_{t}/\operatorname{d\!}t=-\bar{\partial}\mathcal{L}(W_{t}) and d ⁣⁡L(Wt)/d ⁣⁡t=− ⁣∥∂ˉL(Wt)∥2\operatorname{d\!}\mathcal{L}(W_{t})/\operatorname{d\!}t=-\mathinner{\!\left\lVert\bar{\partial}\mathcal{L}(W_{t})\right\rVert}^{2}.

B.3 Kurdyka-Łojasiewicz inequalities

We have the following result regarding the asymptotic Clarke critical values of a definable function, which is basically from (Bolte et al., 2007, Corollary 9).

To state the proof in a bit more detail, (Bolte et al., 2007, Corollary 9) shows that if ff is lower semi-continuous and f>−∞f>-\infty, then ff has finitely many asymptotic Clarke critical values. To get Section B.3, we just need to apply (Bolte et al., 2007, Corollary 9) to the lower semi-continuous envelopes of f∣f−1((0,∞))f|_{f^{-1}\mathinner{\left((0,\infty)\right)}} and −f∣f−1((−∞,0))-f|_{f^{-1}\mathinner{\left((-\infty,0)\right)}}.

The bounded setting.

Here we consider the case where the domain of ff is bounded. (Kurdyka, 1998, Theorem 1) gives a Kurdyka-Łojasiewicz inequality assuming ff is differentiable; below we extend it to the locally Lipschitz setting.

for any x∈f−1((0,ν))x\in f^{-1}\mathinner{\left((0,\nu)\right)}.

By Sections B.1 and B.2, ϕ\phi is definable. Section B.3 implies that there are only finitely many asymptotic Clarke critical values on (0,ϵ)(0,\epsilon), and thus there exists ϵ′∈(0,ϵ)\epsilon^{\prime}\in(0,\epsilon) such that on (0,ϵ′)(0,\epsilon^{\prime}) there is no asymptotic Clarke critical value and ϕ(z)>0\phi(z)>0.

Since ρ\rho is C1C^{1} on $,thereexists, there existsB>0suchthatsuch that\mathinner{\!\left\lVert\rho^{\prime}(s)\right\rVert}\leq Bonon$.

Since hh is definable, h(0)=0h(0)=0, and h(s)>0h(s)>0 on (0,1](0,1], Theorem B.1 implies that there exists a constant ω∈(0,1]\omega\in(0,1] such that h′(s)>0h^{\prime}(s)>0 on (0,ω)(0,\omega).

Section B.2 implies that for a.e. s∈(0,ω)s\in(0,\omega),

Since the left hand side of eq. B.2 is definable, it can actually be nonzero only for finitely many ss, and thus is equal to on some interval (0,μ)(0,\mu) where μ≤ω\mu\leq\omega.

Let ν=h(μ)\nu=h(\mu), the Inverse Function Theorem implies that Ψ:(0,ν)→(0,2Bμ)\Psi\mathrel{\mathop{\ordinarycolon}}(0,\nu)\to(0,2B\mu) given by Ψ(z):=2Bh−1(z)\Psi(z)\mathrel{\mathop{\ordinarycolon}}=2Bh^{-1}(z) is also C1C^{1} definable with a positive derivative, and lim⁡z→0Ψ(z)=0\lim_{z\to 0}\Psi(z)=0.

Now for any x∈f−1((0,ν))x\in f^{-1}\mathinner{\left((0,\nu)\right)}, let s=h−1(f(x))s=h^{-1}\mathinner{\left(f(x)\right)}, we have

The unbounded setting.

The unbounded setting is more complicated: to show directional convergence, we need two Kurdyka-Łojasiewicz inequalities (cf. Sections 3.1 and 3.1), depending on the relationship between the spherical and radial parts of ∂ˉf\bar{\partial}f.

In any o-minimal structure, Uϵ,c,ηU_{\epsilon,c,\eta} is definable if η\eta is rational. Now we prove Section 3.1, a Kurdyka-Łojasiewicz inequality on some Uν,c,ηU_{\nu,c,\eta}, using ideas from (Kurdyka et al., 2006, Proposition 6.3).

Since there is no asymptotic Clarke critical value on (0,ϵ′)(0,\epsilon^{\prime}), it holds that ϕ(z)>0\phi(z)>0.

Since f(Uϵ′,c,η)=(0,ϵ′)f(U_{\epsilon^{\prime},c,\eta})=(0,\epsilon^{\prime}) as above, there exists a sequence xix_{i} in AA such that f(xi)→0f(x_{i})\to 0. If the xix_{i} are bounded, then the claim follows from the proof of Section B.3 and D⊂{x | ∥x∥>1}D\subset\mathinner{\left\{x\,\middle|\,\|x\|>1\right\}}. If the xix_{i} are unbounded, then without loss of generality (e.g., by taking a subsequence) we can assume ∥xi∥→∞\|x_{i}\|\to\infty. Section B.1 asserts that there exists a C1C^{1} definable curve ρ:[a,∞)→A\rho\mathrel{\mathop{\ordinarycolon}}[a,\infty)\to A such that  ⁣∥ρ(s)∥=s\mathinner{\!\left\lVert\rho(s)\right\rVert}=s and lim⁡s→∞f(ρ(s))=0\lim_{s\to\infty}f\mathinner{\left(\rho(s)\right)}=0. Let h(s):=f(ρ(s))h(s)\mathrel{\mathop{\ordinarycolon}}=f\mathinner{\left(\rho(s)\right)}, and ρr′(s):=⟨ρ′(s),ρ(s)⟩ρ(s)/s2\rho_{r}^{\prime}(s)\mathrel{\mathop{\ordinarycolon}}=\left\langle\rho^{\prime}(s),\rho(s)\right\rangle\rho(s)/s^{2} denote the radial part of ρ′(s)\rho^{\prime}(s), and ρ⊥′(s):=ρ′(s)−ρr′(s)\rho_{\perp}^{\prime}(s)\mathrel{\mathop{\ordinarycolon}}=\rho^{\prime}(s)-\rho_{r}^{\prime}(s) denote the spherical part of ρ′(s)\rho^{\prime}(s).

Theorem B.1 implies that h′h^{\prime} is negative and continuous on some interval [ω,∞)[\omega,\infty).

As in the proof of Section B.3, it follows from Section B.2 that there exists μ≥ω\mu\geq\omega, such that

It holds that lim⁡z→0Ψ(z)=0\lim_{z\to 0}\Psi(z)=0. Moreover, for any x∈Uν,c,ηx\in U_{\nu,c,\eta}, let s=h−1(f(x))s=h^{-1}\mathinner{\left(f(x)\right)}, we have

Below we prove Section 3.1, a Kurdyka-Łojasiewicz inequality which is useful outside of Uν,c,ηU_{\nu,c,\eta}.

Note that ξλ−1=ξ1/λ\xi_{\lambda}^{-1}=\xi_{1/\lambda}. If y=ξλ(x)y=\xi_{\lambda}(x), then x=ξ1/λ(y)x=\xi_{1/\lambda}(y), which has the Jacobian

Note that gg is locally Lipschitz and definable with an open bounded domain. Therefore Section B.3 implies that there exists ν>0\nu>0 and a definable desingularizing function Ψ\Psi on [0,ν)[0,\nu) such that

for any y∈g−1((0,ν))y\in g^{-1}\mathinner{\left((0,\nu)\right)}. Let x=ξλ−1(y)x=\xi_{\lambda}^{-1}(y), it holds that gg is differentiable at yy if and only if ff is differentiable at xx, and by the definition of Clarke subdifferential,

which finishes the proof for rational λ\lambda. To handle real λ>0\lambda>0, we can apply the above result to any rational λ′∈(λ/2,λ)\lambda^{\prime}\in(\lambda/2,\lambda). ∎

Appendix C Omitted proofs from Section 3

We first give a generalization of Euler’s homogeneous function theorem, which can also be found in (Lyu and Li, 2019, Theorem B.2), but with an additional requirement of a chain rule.

Let D′D^{\prime} denote the set of xx where ff is differentiable. For any nonzero x∈D′x\in D^{\prime}, it holds that

Since ff is LL-positively homogeneous, f(x+δx)=(1+δ)Lf(x)f(x+\delta x)=(1+\delta)^{L}f(x), and thus

which implies ⟨x,∇f(x)⟩=Lf(x)\left\langle x,\nabla f(x)\right\rangle=Lf(x). This property trivially holds if 0∈D′0\in D^{\prime}.

Since ∂f(x)\partial f(x) consists of convex combinations of such x∗x^{*}, Appendix C holds. ∎

Next we prove a few technical lemmas. Recall the definitions of unnormalized and normalized smoothed margin: given W≠0W\neq 0, let

Additionally, given any function ff which is locally Lipschitz around a nonzero WW, let

denote the radial and spherical parts of ∂ˉf(W)\bar{\partial}f(W) respectively.

We first characterize the Clarke subdifferentials of α\alpha, the unnormalized smoothed margin.

Note that L\mathcal{L} is differentiable at WW if and only if α\alpha is differentiable at WW, and when both gradients exist, the chain rule and inverse function theorem together imply that

whereby the first claim follows from the definition of Clarke subdifferential. To prove the second claim, the chain rule for Clarke subdifferentials (Clarke, 1983, Theorem 2.3.9) implies that

and thus Appendix C ensures for any W∗∈∂α(W)W^{*}\in\partial\alpha(W),

By the definition of Clarke subdifferential, for any nonzero WW,

The first claim of Appendix C holds since for any W∈∂α(W)W\in\partial\alpha(W), by Appendix C,

The last technical result we need is that α\alpha and β\beta are close.

Note that α(W)=π(p(W))\alpha(W)=\pi\mathinner{\left(p(W)\right)} where p(W)=(p1(W),…,pn(W))p(W)=\mathinner{\left(p_{1}(W),\ldots,p_{n}(W)\right)}.

We want to show that ∇2π(v)⪯0\nabla^{2}\pi(v)\preceq 0, or equivalently

Note that for a,b>0a,b>0, we have ea+b−1>(ea−1)+(eb−1)e^{a+b}-1>(e^{a}-1)+(e^{b}-1), which implies

Using Appendix C, we can prove Appendix C.

For simplicity, let p:=(p1(W),…,pn(W))p\mathrel{\mathop{\ordinarycolon}}=\mathinner{\left(p_{1}(W),\ldots,p_{n}(W)\right)}. Recall that α(W)=π(p)\alpha(W)=\pi(p), and from the proof of Appendix C we know that

By the super-additivity of the function σ\sigma defined in eq. C.2, we know that

Let c=−ln⁡(exp⁡(ln⁡(2)/n)−1)≤ln⁡(n)−ln⁡ln⁡(2)c=-\ln\mathinner{\left(\exp\mathinner{\left(\ln(2)/n\right)}-1\right)}\leq\ln(n)-\ln\ln(2) and 1⃗\vec{1} denote the all-ones vector, we have π(c1⃗)=0\pi\mathinner{\left(c\vec{1}\right)}=0, and

Section B.2 implies that for a.e. t≥0t\geq 0,

First note that Section 2 implies that ∥W0∥>0\|W_{0}\|>0, and moreover Lyu and Li (2019, Lemma 5.1) proved that d ⁣⁡∥Wt∥/d ⁣⁡t>0\operatorname{d\!}\|W_{t}\|/\operatorname{d\!}t>0 for a.e. t≥0t\geq 0, and thus ∥Wt∥\|W_{t}\| is increasing and ∥Wt∥≥∥W0∥>0\|W_{t}\|\geq\|W_{0}\|>0.

Now consider W~t\widetilde{W}_{t} and ζt\zeta_{t}. Since WtW_{t} is an arc, and ∥Wt∥≥∥W0∥>0\|W_{t}\|\geq\|W_{0}\|>0, it follows that W~t\widetilde{W}_{t} is also an arc. Moreover, for a.e. t≥0t\geq 0,

Since W~t\widetilde{W}_{t} is an arc, d ⁣⁡W~t/d ⁣⁡t\operatorname{d\!}\widetilde{W}_{t}/\operatorname{d\!}t and  ⁣∥d ⁣⁡W~t/d ⁣⁡t∥\mathinner{\!\left\lVert\operatorname{d\!}\widetilde{W}_{t}/\operatorname{d\!}t\right\rVert} are both integrable, and by definition of the curve length,

Finally we prove the core Section 3.1, which directly implies Theorem 3.1.

Note that eq. C.8 is the opposite to eq. C.4. Appendices C and C implies that

By Appendix C, ∂ˉα(Wt)\bar{\partial}\alpha(W_{t}) is parallel to ∂ˉL(Wt)\bar{\partial}\mathcal{L}(W_{t}), therefore

Now Section 3.1 and eqs. C.11 and C.12 imply

for some constant c>0c>0. Section 3.1 then follows. ∎

Appendix D Omitted proofs from Section 4

We first give the following technical result.

If ∇f\nabla f is differentiable at a nonzero xx, then for any c>0c>0, it holds that

Moreover, there exists Kσ>0K_{\sigma}>0 such that for any ∥x∥=1\|x\|=1, if ∇2f(x)\nabla^{2}f(x) exists, then  ⁣∥∇2f(x)∥σ≤Kσ\mathinner{\!\left\lVert\nabla^{2}f(x)\right\rVert}_{\sigma}\leq K_{\sigma}.

which proves the claim. The homogeneity of ∇2f\nabla^{2}f when it exists can be proved in the same way.

To get KσK_{\sigma}, note that for any ∥x∥=1\|x\|=1, there exists an open neighborhood UxU_{x} of xx on which ∇f\nabla f is KxK_{x}-Lipschitz continuous, and thus the spectral norm of ∇2f\nabla^{2}f is bounded by KxK_{x} when it exists. All the UxU_{x} form an open cover of the compact unit sphere, and thus has a finite subcover, which implies the claim. ∎

Below we estimate various quantities using Appendix D.

where π\pi is defined in eq. C.3 and all partial derivatives are evaluated at p(W):=(p1(W),…,pn(W))p(W)\mathrel{\mathop{\ordinarycolon}}=(p_{1}(W),\ldots,p_{n}(W)). It is shown in the proof of Appendix C that  ⁣∥π(p)∥1≤2\mathinner{\!\left\lVert\pi(p)\right\rVert}_{1}\leq 2. Moreover, Appendix D implies that all  ⁣∥∇pi(W)∥/∥W∥L−1\mathinner{\!\left\lVert\nabla p_{i}(W)\right\rVert}/\|W\|^{L-1} are bounded. Consequently,  ⁣∥∇α(W)∥/∥W∥L−1\mathinner{\!\left\lVert\nabla\alpha(W)\right\rVert}/\|W\|^{L-1} is bounded. ∎

If all ∇pi\nabla p_{i} are locally Lipschitz, then J\mathcal{J} is also locally Lipschitz. We further have the following result.

for some constant K>0K>0, where θ\theta denotes the angle between WW and −∇L(W)-\nabla\mathcal{L}(W).

Below we fix an arbitrary W∈D′∩S0W\in D^{\prime}\cap S_{0}. All the partial derivatives below with respect to pip_{i} are evaluated at p(W):=(p1(W),…,pn(W))p(W)\mathrel{\mathop{\ordinarycolon}}=(p_{1}(W),\ldots,p_{n}(W)). Recall that

where π\pi is defined in eq. C.3. Since ∇pi\nabla p_{i} are also differentiable at WW, we have

Now for any W∈D′∩S0W\in D^{\prime}\cap S_{0}, we have (recall that W~=W/∥W∥\widetilde{W}=W/\|W\|)

Comparing eqs. D.1 and D, first note that

since π\pi is concave by Appendix C, and moreover

Let ∇rα(W)\nabla_{r}\alpha(W) and ∇⊥α(W)\nabla_{\perp}\alpha(W) denote the radial and spherical part of ∇α(W)\nabla\alpha(W), respectively. Let θ\theta denote the angle between WW and ∇α(W)\nabla\alpha(W). Appendices C and C imply that

and thus θ\theta is between and π/2\pi/2. Now Appendix D and the proof of Appendix C imply that

In addition, the proof of Appendix C shows that  ⁣∥π(p)∥1≤2\mathinner{\!\left\lVert\pi(p)\right\rVert}_{1}\leq 2, and Appendix D ensures that ∥∇2f∥σ\|\nabla^{2}f\|_{\sigma} has a uniform bound KσK_{\sigma} on the unit sphere, therefore

Combining appendices D, D.3, D, D and D gives

The following result helps us control θt\theta_{t}.

Under the same condition as Appendix D and Section 2, it holds that

Since β(Wt)/∥Wt∥L\beta(W_{t})/\|W_{t}\|^{L} is bounded due to Appendix D, the proof is finished. ∎

Fix an arbitrary ϵ∈(0,1)\epsilon\in(0,1), and let JtJ_{t} denote J(Wt)J(W_{t}). Recall that lim⁡t→∞α(Wt)/∥Wt∥L=a\lim_{t\to\infty}\alpha(W_{t})/\|W_{t}\|^{L}=a. Appendix C then implies lim⁡t→∞β(Wt)/∥Wt∥L=a\lim_{t\to\infty}\beta(W_{t})/\|W_{t}\|^{L}=a, and thus we can find t1t_{1} such that for any t>t1t>t_{1},

Moreover, Sections B.2, D and D imply that there exists t2t_{2} such that for any t′>t>t2t^{\prime}>t>t_{2},

(Lyu and Li, 2019, Corollary C.10) implies that there exists t3>max⁡{t1,t2}t_{3}>\max\{t_{1},t_{2}\} such that

We claim that δt<1+ϵ\delta_{t}<1+\epsilon for any t>t3t>t_{3}.

To see this, note that eqs. D.7 and D.9 imply

Moreover, using eq. D.8, for any t>t2t>t_{2},

Since ϵ\epsilon is arbitrary, we have lim⁡t→∞θt=0\lim_{t\to\infty}\theta_{t}=0.

If all pip_{i} are C2C^{2}, then the above proof holds without definability: it is only used in eq. D.8 to ensure the chain rule, which always holds for C2C^{2} functions. ∎

Appendix E Global margin maximization proofs for Section 4.2

First, a technical lemma regarding directional convergence and alignment properties inherited by these subsets of WtW_{t}. This will be used in both the deep linear case and in the 2-homogeneous case.

Suppose the conditions for Theorems 3.1 and 4.1 hold. Let (U1(t),…,Ur(t))(U_{1}(t),\ldots,U_{r}(t)) be any partition of WtW_{t}, and set sj(t):=∥Uj(t)∥L/∥Wt∥Ls_{j}(t)\mathrel{\mathop{\ordinarycolon}}=\|U_{j}(t)\|^{L}/\|W_{t}\|^{L}. Then s(t)s(t) converges to some sˉ\bar{s}, and for each jj,

First note that s(t)s(t) converges since Wt/∥Wt∥W_{t}/\|W_{t}\| converges, and alignment grants

By directional convergence (cf. Theorem 3.1), alignment (cf. Theorem 4.1), and Cauchy-Schwarz,

which starts and ends with −1-1 and is thus a chain of equalities. Applying eq. E.1 and he equality case of Cauchy-Schwarz to each jj with sˉj>0\bar{s}_{j}>0,

For the final claim, note Theorem 4.1 and eq. 4.2 imply that

Applying the preceding lemma to network layers, we handle the deep linear case as follows.

where (AL⋯Aj+1)T(A_{L}\cdots{}A_{j+1})^{\scriptscriptstyle\mathsf{T}} is a column vector, and (Aj−1⋯A1∇uL(W))T(A_{j-1}\cdots A_{1}\nabla_{u}\mathcal{L}(W))^{\scriptscriptstyle\mathsf{T}} is a row vector, and moreover ⟨Aj,∇AjL(W)⟩=⟨u,∇uL(W)⟩\left\langle A_{j},\nabla_{A_{j}}\mathcal{L}(W)\right\rangle=\left\langle u,\nabla_{u}\mathcal{L}(W)\right\rangle, where this last inner product does not depend on jj.

Applying the subset-alignment of Appendix E to layers (Aj,…,A1)(A_{j},\ldots,A_{1}) gives, for each jj,

whereby sˉj\bar{s}_{j} is independent of jj, which can only mean sˉj2/L=1/L>0\bar{s}_{j}^{2/L}=1/L>0 for all jj, but more importantly ∥Aj(t)∥→∞\|A_{j}(t)\|\to\infty for all jj. By Appendix E, this means all layers align with their gradients.

Next it is proved by induction from ALA_{L} to A1A_{1} that there exist unit vectors v0,…,vLv_{0},\ldots,v_{L} with vL=1v_{L}=1 and Aj/∥Aj∥=vjvj−1TA_{j}/\|A_{j}\|=v_{j}v_{j-1}^{\scriptscriptstyle\mathsf{T}}. The base case ALA_{L} holds immediately, since ALA_{L} is a row vector, meaning we can choose vL:=1v_{L}\mathrel{\mathop{\ordinarycolon}}=1 and vL−1:=ALT/∥AL∥v_{L-1}\mathrel{\mathop{\ordinarycolon}}=A_{L}^{\scriptscriptstyle\mathsf{T}}/\|A_{L}\| since ALA_{L} converges in direction. For the inductive step AjA_{j} with j<Lj<L, note

Since vjv_{j} is a fixed unit vector and since ∇AjL(W)\nabla_{A_{j}}\mathcal{L}(W) converges in direction, the row vector part of the above expression must also converge to some fixed unit vector vj−1Tv_{j-1}^{\scriptscriptstyle\mathsf{T}}, namely

Since AjA_{j} and −∇AjL(W)-\nabla_{A_{j}}\mathcal{L}(W) asymptotically align as above, then \nicefracAj∥Aj∥→vjvj−1T\nicefrac{{A_{j}}}{{\|A_{j}\|}}\to v_{j}v_{j-1}^{\scriptscriptstyle\mathsf{T}}.

Now consider v0v_{0} and uu, where it still needs to be shown that v0=u/∥u∥v_{0}=u/\|u\|. To this end, note

whereby u/∥u∥=v0u/\|u\|=v_{0}. By a similar calculation,

which means u/∥u∥u/\|u\| asymptotically satisfies the optimality conditions for the optimization problem

Before moving on to the 2-homogeneous case, we first produce another technical lemma, which we will use to control dual variables qi(t):=∂α/∂pi(Wt)q_{i}(t)\mathrel{\mathop{\ordinarycolon}}=\partial\alpha/\partial p_{i}(W_{t}), which also appear in Section 4.2.

With this in hand, we can handle the 2-homogeneous case.

Applying Appendix E to the per-node weights (w1,…,wm)(w_{1},\ldots,w_{m}), a limit sˉ\bar{s} exists and due to 2-homogeneity satisfies sˉ∈Δm\bar{s}\in\Delta_{m}. Whenever, sˉj>0\bar{s}_{j}>0, then

Consequently, this means that either sˉj>0\bar{s}_{j}>0 and lim⁡t→∞∑iqi(t)φij(θj(t))=a\lim_{t\to\infty}\sum_{i}q_{i}(t)\varphi_{ij}(\theta_{j}(t))=a, or else sˉj=0\bar{s}_{j}=0 and by the choice θˉj=0\bar{\theta}_{j}=0 then lim⁡t→∞∑iqi(t)φij(θj(t))=0\lim_{t\to\infty}\sum_{i}q_{i}(t)\varphi_{ij}(\theta_{j}(t))=0. In particular, this means sˉj>0\bar{s}_{j}>0 iff θˉj\bar{\theta}_{j} attains the maximal value aa, meaning sˉ\bar{s} satisfies the Sion primal optimality conditions for the saddle point problem over the fixed points (θˉ1,…,θˉm)(\bar{\theta}_{1},\ldots,\bar{\theta}_{m}) (Chizat and Bach, 2020, Proposition D.3).

Now consider the dual variables qi(t)=∂α/∂pi(Wt)q_{i}(t)=\partial\alpha/\partial p_{i}(W_{t}). By Appendix E, any accumulation point qˉ\bar{q} is an element of Δn\Delta_{n} and moreover is supported on those examples ii minimizing pi(W‾)p_{i}(\overline{W}), which means qˉ\bar{q} satisfies the Sion dual optimality conditions for the margin saddle point problem again over fixed points (θˉ1,…,θˉm)(\bar{\theta}_{1},\ldots,\bar{\theta}_{m}) (Chizat and Bach, 2020, Proposition D.3). Thus applying the Sion Theorem over discrete domain (θˉ1,…,θˉm)(\bar{\theta}_{1},\ldots,\bar{\theta}_{m}) to the primal-dual optimal pair (sˉ,qˉ)(\bar{s},\bar{q}) gives

and directional convergence of W~t\widetilde{W}_{t} combined with definition of qˉ\bar{q} gives

Since qˉ\bar{q} was an arbitrary accumulation point, it holds in general that

Next, for any q∈Δnq\in\Delta_{n} and s∈Δms\in\Delta_{m}, using the first part of the cover condition,