Stochastic gradient descent performs variational inference, converges to limit cycles for deep networks

Introduction

where $H(\rho)$ is the entropy of the distribution $\rho$ and $\eta$ and $\mathscr{b}$ are the learning rate and batch-size, respectively. The potential $\Phi(x)$, which we characterize explicitly, is related but not necessarily equal to $f(x)$. It is only a function of the architecture and the dataset. This implies that SGD implicitly performs variational inference with a uniform prior, albeit of a different loss than the one used to compute back-propagation gradients.

We next prove that the implicit potential $\Phi(x)$ is equal to our chosen loss $f(x)$ if and only if the noise in mini-batch gradients is isotropic. This condition, however, is not satisfied for deep networks. Empirically, we find gradient noise to be highly non-isotropic with the rank of its covariance matrix being about $1\%$ of its dimension. Thus, SGD on deep networks implicitly discovers locations where $\nabla\Phi(x)=0$, these are not the locations where $\nabla f(x)=0$. This is our second main result: the most likely locations of SGD are not the local minima, nor the saddle points, of the original loss. The deviation of these critical points, which we compute explicitly scales linearly with $\eta/\mathscr{b}$ and is typically large in practice.

When mini-batch noise is non-isotropic, SGD does not even converge in the classical sense. We prove that, instead of undergoing Brownian motion in the vicinity of a critical point, trajectories have a deterministic component that causes SGD to traverse closed loops in the weight space. We detect such loops using a Fourier analysis of SGD trajectories. We also show through an example that SGD with non-isotropic noise can even converge to stable limit cycles around saddle points.

Background on continuous-time SGD

Stochastic gradient descent performs the following updates while training a network $x_{k+1}=x_{k}-\eta\ \nabla f_{\mathscr{b}}(x_{k})$ where $\eta$ is the learning rate and $\nabla f_{\mathscr{b}}(x_{k})$ is the average gradient over a mini-batch $\mathscr{b}$,

Note that $D(x)$ is independent of the learning rate $\eta$ and the batch-size $\mathscr{b}$. It only depends on the weights $x$, architecture and loss defined by $f(x)$, and the dataset. We will often discuss two cases: isotropic diffusion when $D(x)$ is a scalar multiple of identity, independent of $x$, and non-isotropic diffusion, when $D(x)$ is a general function of the weights $x$.

We now construct a stochastic differential equation (SDE) for the discrete-time SGD updates.

The continuous-time limit of SGD is given by

We refer to Li et al., 2017b (, Thm. 1) for the proof of the convergence of discrete SGD to Eq. 3. Note that $\beta^{-1}$ completely captures the magnitude of noise in SGD that depends only upon the learning rate $\eta$ and the mini-batch size $\mathscr{b}$.

SGD performs variational inference

which is also the steady-state solution of

as can be verified by direct substitution in Eq. FP.

The above observation is very useful because it suggests that, if $\nabla f(x)$ can be written in terms of the diffusion matrix and a gradient term $\nabla\Phi(x)$, the steady-state distribution of this SDE is easily obtained. We exploit this observation to rewrite $\nabla f(x)$ in terms a term $D\ \nabla\Phi$ that gives rise to the above steady-state, the spatial derivative of the diffusion matrix, and the remainder:

interpreted as the part of $\nabla f(x)$ that cannot be written as $D\ \Phi^{\prime}(x)$ for some $\Phi^{\prime}$. We now make an important assumption on $j(x)$ which has its origins in thermodynamics.

This leads us to the main result of this section.

decreases monotonically along the trajectories of the Fokker-Planck equation Eq. FP and converges to its minimum, which is zero, at steady-state. Moreover, we also have an energetic-entropic split

Theorem 5, proven in Section F.1, shows that SGD implicitly minimizes a combination of two terms: an “energetic” term, and an “entropic” term. The first is the average potential over a distribution $\rho$. The steady-state of SGD in Eq. 6 is such that it places most of its probability mass in regions of the parameter space with small values of $\Phi$. The second shows that SGD has an implicit bias towards solutions that maximize the entropy of the distribution $\rho$.

Note that the energetic term in Eq. 11 has potential $\Phi(x)$, instead of $f(x)$. This is an important fact and the crux of this paper.

If the diffusion matrix $D(x)$ is isotropic, i.e., a constant multiple of the identity, the implicit potential is the original loss itself

This is proven in Section F.2. The definition in Eq. 8 shows that $j\neq 0$ when $D(x)$ is non-isotropic. This results in a deterministic component in the SGD dynamics which does not affect the functional $F(\rho)$, hence $j(x)$ is called a “conservative force.” The following lemma is proven in Section F.3.

The force $j(x)$ does not decrease $F(\rho)$ in Eq. 11 and introduces a deterministic component in SGD given by

The condition $\nabla\cdot j(x)=0$ in 4 implies that most likely trajectories of SGD traverse closed trajectories in weight space.

Theorem 5 applies for a general $D(x)$ and it is equivalent to the celebrated JKO functional (Jordan et al.,, 1997) in optimal transportation (Santambrogio,, 2015; Villani,, 2008) if the diffusion matrix is isotropic. Appendix D provides a brief overview using the heat equation as an example.

If $D(x)=I$, trajectories of the Fokker-Planck equation Eq. FP are gradient flow in the Wasserstein metric of the functional

2 Connection to Bayesian inference

Note the absence of any prior in Eq. 11. On the other hand, the evidence lower bound (Kingma and Welling,, 2013) for the dataset $\Xi$ is,

where $H(q,p)$ is the cross-entropy of the estimated steady-state and the variational prior. The implicit loss function of SGD in Eq. 11 therefore corresponds to a uniform prior $p(x\,|\,\Xi)$. In other words, we have shown that SGD itself performs variational optimization with a uniform prior. Note that this prior is well-defined by our hypothesis of $x\in\Omega$ for some compact $\Omega$.

It is important to note that SGD implicitly minimizes a potential $\Phi(x)$ instead of the original loss $f(x)$ in ELBO. We prove in Section 5 that this potential is quite different from $f(x)$ if the diffusion matrix $D$ is non-isotropic, in particular, with respect to its critical points.

The functional Eq. 11 is equivalent to the information bottleneck principle in representation learning (Tishby et al.,, 1999). Minimizing this functional, explicitly, has been shown to lead to invariant representations (Achille and Soatto,, 2017). Theorem 5 shows that SGD implicitly contains this bottleneck and therefore begets these properties, naturally.

3 Practical implications

We will show in Section 5 that the potential $\Phi(x)$ does not depend on the optimization process, it is only a function of the dataset and the architecture. The effect of two important parameters, the learning rate $\eta$ and the mini-batch size $\mathscr{b}$ therefore completely determines the strength of the entropic regularization term. If $\beta^{-1}\to 0$, the implicit regularization of SGD goes to zero. This implies that

is a good tenet for regularization of SGD.

In order to maintain the entropic regularization, the learning rate $\eta$ needs to scale linearly with the batch-size $\mathscr{b}$. This prediction, based on Theorem 5, fits very well with empirical evidence wherein one obtains good generalization performance only with small mini-batches in deep networks (Keskar et al.,, 2016), or via such linear scaling (Goyal et al.,, 2017).

The diffusion matrix for the case when mini-batches are sampled with replacement is very close to Eq. 2, see Section A.2. However, the corresponding inverse temperature is

The extra factor of $\left(1-\frac{\mathscr{b}}{N}\right)$ reduces the entropic regularization in Eq. 11, as $\mathscr{b}\to N$, the inverse temperature $\beta^{\prime}\to\infty$. As a consequence, for the same learning rate $\eta$ and batch-size $\mathscr{b}$, Theorem 5 predicts that sampling with replacement has better regularization than sampling without replacement. This effect is particularly pronounced at large batch-sizes.

Empirical characterization of SGD dynamics

Section 4.1 shows that the diffusion matrix $D(x)$ for modern deep networks is highly non-isotropic with a very low rank. We also analyze trajectories of SGD and detect periodic components using a frequency analysis in Section 4.2; this validates the prediction of Lemma 7.

We consider three networks for these experiments: a convolutional network called small-lenet, a two-layer fully-connected network on MNIST (LeCun et al.,, 1998) and a smaller version of the All-CNN-C architecture of Springenberg et al., (2014) on the CIFAR-10 and CIFAR-100 datasets (Krizhevsky,, 2009); see Appendix E for more details.

Figs. 1 and 2 show the eigenspectrumthresholded at $\lambda_{\max}\times d\times\textrm{machine-precision}$. This formula is widely used, for instance, in numpy. of the diffusion matrix. In all cases, it has a large fraction of almost-zero eigenvalues with a very small rank that ranges between $0.3\%$ - $2\%$. Moreover, non-zero eigenvalues are spread across a vast range with a large variance.

We have plotted the eigenspectra of the diffusion matrix in Fig. 1 and Fig. 2 at three different instants, $20\%$, $40\%$ and $100\%$ training completion; they are almost indistinguishable. This implies that the variance of the mini-batch gradients in deep networks can be considered a constant, highly non-isotropic matrix.

The eigenspectra in Fig. 2 for CIFAR-10 and CIFAR-100 have much larger eigenvalues and standard-deviation than those in Fig. 1, this is expected because the images in the CIFAR datasets have more variety than those in MNIST. Similarly, while CIFAR-100 has qualitatively similar images as CIFAR-10, it has $10\times$ more classes and as a result, it is a much harder dataset. This correlates well with the fact that both the mean and standard-deviation of the eigenvalues in Fig. 2(b) are much higher than those in Fig. 2(a). Input augmentation increases the diversity of mini-batch gradients. This is seen in Fig. 2(c) where the standard-deviation of the eigenvalues is much higher as compared to Fig. 2(a).

Remark 14 shows that the mean of the eigenspectrum is large if the dataset is diverse. Based on this, we propose that the inverse temperature $\beta$ should scale linearly with the mean of the eigenvalues of $D$:

where $d$ is the number of weights. This keeps the noise in SGD constant in magnitude for different values of the learning rate $\eta$, mini-batch size $\mathscr{b}$, architectures, and datasets. Note that other hyper-parameters which affect stochasticity such as dropout probability are implicit inside $D$.

Compare the eigenspectra in Figs. 1(a) and 1(b) with those in Figs. 2(a) and 2(c). The former pair shows that small-lenet which is a much better network than small-fc also has a much larger rank, i.e., the number of non-zero eigenvalues ($D(x)$ is symmetric). The second pair shows that for the same dataset, data-augmentation creates a larger variance in the eigenspectrum. This suggests that both the quantities, viz., rank of the diffusion matrix and the variance of the eigenspectrum, inform the performance of a given architecture on the dataset. Note that as discussed in Remark 15, the mean of the eigenvalues can be controlled using the learning rate $\eta$ and the batch-size $\mathscr{b}$.

This observation is useful for automated architecture search where we can use the quantity

to estimate the efficacy of a given architecture, possibly, without even training, since $D$ does not depend on the weights much. This task currently requires enormous amounts of computational power (Zoph and Le,, 2016; Baker et al.,, 2016; Brock et al.,, 2017).

2 Analysis of long-term trajectories

We train a smaller version of small-fc on $7\times 7$ down-sampled MNIST images for $10^{5}$ epochs and store snapshots of the weights after each epoch to get a long trajectory in the weight space. We discard the first $10^{3}$ epochs of training (“burnin”) to ensure that SGD has reached the steady-state. The learning rate is fixed to $10^{-3}$ after this, up to $10^{5}$ epochs.

The auto-correlation (AC) in Fig. 3(b) should be compared with the AC for Brownian motion which decays to zero very quickly and stays within the red confidence bands ($99\%$). Our iterates are significantly correlated with each other even at very large lags. This further indicates that trajectories of SGD do not perform Brownian motion.

Fig. 3(c) shows that the full-gradient computed over the entire dataset (without burnin) does not decrease much with respect to the number of epochs. While it is expected to have a non-zero gradient norm because SGD only converges to a neighborhood of a critical point for non-zero learning rates, the magnitude of this gradient norm is quite large. This magnitude drops only by about a factor of $3$ over the next $10^{5}$ epochs. The presence of a non-zero $j(x)$ also explains this, it causes SGD to be away from critical points, this phenomenon is made precise in Theorem 22. Let us note that a similar plot is also seen in Shwartz-Ziv and Tishby, (2017) for the per-layer gradient magnitude.

SGD for deep networks is out-of-equilibrium

This section now gives an explicit formula for the potential $\Phi(x)$. We also discuss implications of this for generalization in Section 5.3.

The fundamental difficulty in obtaining an explicit expression for $\Phi$ is that even if the diffusion matrix $D(x)$ is full-rank, there need not exist a function $\Phi(x)$ such that $\nabla\Phi(x)=D^{-1}(x)\ \nabla f(x)$ at all $x\in\Omega$. We therefore split the analysis into two cases:

a local analysis near any critical point $\nabla f(x)=0$ where we linearize $\nabla f(x)=Fx$ and $\nabla\Phi(x)=Ux$ to compute $U=G^{-1}\ F$ for some $G$, and

the general case where $\nabla\Phi(x)$ cannot be written as a local rotation and scaling of $\nabla f(x)$.

Let us introduce these cases with an example from Noh and Lee, (2015).

The matrix $F$ in Eq. 15 can be uniquely decomposed into

$D$ and $Q$ are the symmetric and anti-symmetric parts of a matrix $G$ with $GF^{\top}-FG^{\top}=0$, to get $\Phi(x)=\frac{1}{2}x^{\top}Ux$.

The above lemma is a classical result if the critical point is a local minimum, i.e., if the loss is locally convex near $x=0$; this case has also been explored in machine learning before (Mandt et al.,, 2016). We refer to Kwon et al., (2005) for the proof that linearizes around any critical point.

We see from Lemma 20 that, near a critical point,

up to the first order. This suggests that the effect of $j(x)$ is to rotate the gradient field and move the critical points, also seen in Fig. 4(b). Note that $\nabla\cdot D=0$ and $\nabla\cdot Q=0$ in the linearized analysis.

2 General case

We next give the general expression for the deviation of the critical points $\nabla\Phi$ from those of the original loss $\nabla f$.

The main result of the section now follows. It exploits the A-type interpretation to compute the difference between the most likely locations of SGD which are given by the critical points of the potential $\Phi(x)$ and those of the original loss $f(x)$.

is equivalent to the A-type SDE (Ao et al.,, 2007; Shi et al.,, 2012)

The anti-symmetric matrix $Q(x)$ and the potential $\Phi(x)$ can be explicitly computed in terms of the gradient $\nabla f(x)$ and the diffusion matrix $D(x)$. The potential $\Phi(x)$ does not depend on $\beta$.

See Section F.4 for the proof. It exploits the fact that the the Ito SDE Eq. 3 and the A-type SDE Eq. 18 should have the same Fokker-Planck equations because they have the same steady-state distributions.

Theorem 22 presents a picture that is completely consistent with Lemma 20. If $j(x)=0$ and $Q(x)=0$, or if $Q$ is a constant like the linear case in Lemma 20, the divergence of $Q(x)$ in Eq. 19 is zero.

The presence of a $Q(x)$ with non-zero divergence is the consequence of a non-isotropic $D(x)$ and it persists even if $D$ is constant and independent of weights $x$. So long as $D$ is not isotropic, as we discussed in the beginning of Section 5, there need not exist a function $\Phi(x)$ such that $\nabla\Phi(x)=D^{-1}\ \nabla f(x)$ at all $x$. This is also seen in our experiments, the diffusion matrix is almost constant with respect to weights for deep networks, but consequences of out-of-equilibrium behavior are still seen in Section 4.2.

The effect predicted by Eq. 19 becomes more pronounced if $\beta^{-1}=\frac{\eta}{2\mathscr{b}}$ is large. In other words, small batch-sizes or high learning rates cause SGD to be drastically out-of-equilibrium. Theorem 5 also shows that as $\beta^{-1}\to 0$, the implicit entropic regularization in SGD vanishes. Observe that these are exactly the conditions under which we typically obtain good generalization performance for deep networks (Keskar et al.,, 2016; Goyal et al.,, 2017). This suggests that non-equilibrium behavior in SGD is crucial to obtain good generalization performance, especially for high-dimensional models such as deep networks where such effects are expected to be more pronounced.

3 Generalization

It was found that solutions of discrete learning problems that generalize well belong to dense clusters in the weight space (Baldassi et al.,, 2015; 2016). Such dense clusters are exponentially fewer compared to isolated solutions. To exploit these observations, the authors proposed a loss called “local entropy” that is out-of-equilibrium by construction and can find these well-generalizable solutions easily. This idea has also been successful in deep learning where Chaudhari et al., (2016) modified SGD to seek solutions in “wide minima” with low curvature to obtain improvements in generalization performance as well as convergence rate (Chaudhari et al., 2017a, ).

Local entropy is a smoothed version of the original loss given by

Actively constructing out-of-equilibrium behavior leads to good generalization in practice. Our evidence that SGD on deep networks itself possesses out-of-equilibrium behavior then indicates that SGD for deep networks generalizes well because of such behavior.

Related work

The idea that SGD is related to variational inference has been seen in machine learning before (Duvenaud et al.,, 2016; Mandt et al.,, 2016) under assumptions such as quadratic steady-states; for instance, see Mandt et al., (2017) for methods to approximate steady-states using SGD. Our results here are very different, we would instead like to understand properties of SGD itself. Indeed, in full generality, SGD performs variational inference using a new potential $\Phi$ that it implicitly constructs given an architecture and a dataset.

It is widely believed that SGD is an implicit regularizer, see Zhang et al., (2016); Neyshabur et al., (2017); Shwartz-Ziv and Tishby, (2017) among others. This belief stems from its remarkable empirical performance. Our results show that such intuition is very well-placed. Thanks to the special architecture of deep networks where gradient noise is highly non-isotropic, SGD helps itself to a potential $\Phi$ with properties that lead to both generalization and acceleration.

SGD and noise:

Noise is often added in SGD to improve its behavior around saddle points for non-convex losses, see Lee et al., (2016); Anandkumar and Ge, (2016); Ge et al., (2015). It is also quite indispensable for training deep networks (Hinton and Van Camp,, 1993; Srivastava et al.,, 2014; Kingma et al.,, 2015; Gulcehre et al.,, 2016; Achille and Soatto,, 2017). There is however a disconnect between these two directions due to the fact that while adding external gradient noise helps in theory, it works poorly in practice (Neelakantan et al.,, 2015; Chaudhari and Soatto,, 2015). Instead, “noise tied to the architecture” works better, e.g., dropout, or small mini-batches. Our results close this gap and show that SGD crucially leverages the highly degenerate noise induced by the architecture.

Gradient diversity:

Yin et al., (2017) construct a scalar measure of the gradient diversity given by $\sum_{k}\lVert\nabla f_{k}(x)\rVert/\lVert\nabla f(x)\rVert$, and analyze its effect on the maximum allowed batch-size in the context of distributed optimization.

Markov Chain Monte Carlo:

MCMC methods that sample from a negative log-likelihood $\Phi(x)$ have employed the idea of designing a force $j=\nabla\Phi-\nabla f$ to accelerate convergence, see Ma et al., (2015) for a thorough survey, or Pavliotis, (2016); Kaiser et al., (2017) for a rigorous treatment. We instead compute the potential $\Phi$ given $\nabla f$ and $D$, which necessitates the use of techniques from physics. In fact, our results show that since $j\neq 0$ for deep networks due to non-isotropic gradient noise, very simple algorithms such as SGLD by Welling and Teh, (2011) also benefit from the acceleration that their sophisticated counterparts aim for (Ding et al.,, 2014; Chen et al.,, 2016).

Discussion

The continuous-time point-of-view used in this paper gives access to general principles that govern SGD, such analyses are increasingly becoming popular (Wibisono et al.,, 2016; Chaudhari et al., 2017b, ). However, in practice, deep networks are trained for only a few epochs with discrete-time updates. Closing this gap is an important future direction. A promising avenue towards this is that for typical conditions in practice such as small mini-batches or large learning rates, SGD converges to the steady-state distribution quickly (Raginsky et al.,, 2017).

Acknowledgments

PC would like to thank Adam Oberman for introducing him to the JKO functional. The authors would also like to thank Alhussein Fawzi for numerous discussions during the conception of this paper and his contribution to its improvement.

References

Appendix A Diffusion matrix D​(x)𝐷𝑥D(x)

In this section we denote $g_{k}:=\nabla f_{k}(x)$ and $g:=\nabla f(x)=\frac{1}{N}\ \sum_{k=1}^{N}\ g_{k}$. Although we drop the dependence of $g_{k}$ on $x$ to keep the notation clear, we emphasize that the diffusion matrix $D$ depends on the weights $x$.

Let $i_{1},\dots,i_{\mathscr{b}}$ be $\mathscr{b}$ iid random variables in $\left\{1,2,\ldots,N\right\}$. We would like to compute

Note that we have that for any $j\neq k$, the random vectors $g_{i_{j}}$ and $g_{i_{k}}$ are independent. We therefore have

and assimilate the factor of $\mathscr{b}^{-1}$ in the inverse temperature $\beta$.

A.2 Without replacement

Let us define an indicator random variable ${\bf 1}_{i\in\mathscr{b}}$ that denotes if an example $i$ was sampled in batch $\mathscr{b}$. We can show that

Similar to Li et al., 2017a , we can now compute

and assimilate the factor of $\mathscr{b}^{-1}\left(1-\frac{\mathscr{b}}{N}\right)$ that depends on the batch-size in the inverse temperature $\beta$.

Appendix B Discussion on 4

Let $F(\rho)$ be as defined in Eq. 11. In non-equilibrium thermodynamics, it is assumed that the local entropy production is a product of the force $-\nabla\left(\frac{\delta F}{\delta\rho}\right)$ from Eq. A8 and the probability current $-J(x,t)$ from Eq. FP. This assumption in this form was first introduced by Prigogine, (1955) based on the works of Onsager, 1931a ; Onsager, 1931b . See Frank, (2005, Sec. 4.5) for a mathematical treatment and Jaynes, (1980) for further discussion. The rate of entropy $(S_{i})$ increase is given by

This can now be written using Eq. A8 again as

The first term in the above expression is non-negative, in order to ensure that $\frac{dS_{i}}{dt}\geq 0$, we require

where the second equality again follows by integration by parts. It can be shown (Frank,, 2005, Sec. 4.5.5) that the condition in 4, viz., $\nabla\cdot j(x)=0$, is sufficient to make the above integral vanish and therefore for the entropy generation to be non-negative.

Appendix C Some properties of the force j𝑗j

The Fokker-Planck equation Eq. FP can be written in terms of the probability current as

Using the definition of $j(x)$ in Eq. 8, we have detailed balance when

Appendix D Heat equation as a gradient flow

As first discovered in the works of Jordan, Kinderleherer and Otto (Jordan et al.,, 1998; Otto,, 2001), certain partial differential equations can be seen as coming from a variational principle, i.e., they perform steepest descent with respect to functionals of their state distribution. Section 3 is a generalization of this idea, we give a short overview here with the heat equation. The heat equation

can be written as the steepest descent for the Dirichlet energy functional

However, the same PDE can also be seen as the gradient flow of the negative Shannon entropy in the Wasserstein metric (Santambrogio,, 2017; 2015),

More precisely, the sequence of iterated minimization problems

converges to trajectories of the heat equation as $\tau\to 0$. This equivalence is extremely powerful because it allows us to interpret, and modify, the functional $-H(\rho)$ that PDEs such as the heat equation implicitly minimize.

This equivalence is also quite natural, the heat equation describes the probability density of pure Brownian motion: $dx=\sqrt{2}\ dW(t)$. The Wasserstein point-of-view suggests that Brownian motion maximizes the entropy of its state distribution, while the Dirichlet functional suggests that it minimizes the total-variation of its density. These are equivalent. While the latter has been used extensively in image processing, our paper suggests that the entropic regularization point-of-view is very useful to understand SGD for machine learning.

Appendix E Experimental setup

We consider the following three networks on the MNIST (LeCun et al.,, 1998) and the CIFAR-10 and CIFAR-100 datasets (Krizhevsky,, 2009).

small-lenet: a smaller version of LeNet (LeCun et al.,, 1998) on MNIST with batch-normalization and dropout ($0.1$) after both convolutional layers of $8$ and $16$ output channels, respectively. The fully-connected layer has $128$ hidden units. This network has $13,338$ weights and reaches about $0.75\%$ training and validation error.

small-fc: a fully-connected network with two-layers, batch-normalization and rectified linear units that takes $7\times 7$ down-sampled images of MNIST as input and has $64$ hidden units. Experiments in Section 4.2 use a smaller version of this network with $16$ hidden units and $5$ output classes ($30,000$ input images); this is called tiny-fc.

small-allcnn: this a smaller version of the fully-convolutional network for CIFAR-10 and CIFAR-100 introduced by Springenberg et al., (2014) with batch-normalization and $12,24$ output channels in the first and second block respectively. It has $26,982$ weights and reaches about $11\%$ and $17\%$ training and validation errors, respectively.

We train the above networks with SGD with appropriate learning rate annealing and Nesterov’s momentum set to $0.9$. We do not use any data-augmentation and pre-process data using global contrast normalization with ZCA for CIFAR-10 and CIFAR-100.

Appendix F Proofs

which helps us write the Fokker-Planck equation Eq. FP as

As we show in Appendix B, the first term above is zero due to 4. Under suitable boundary condition on the Fokker-Planck equation which ensures that no probability mass flows across the boundary of the domain $\partial\Omega$, after an integration by parts, the second term can be written as

In the above expression, $A:B$ denotes the matrix dot product $A:B=\sum_{ij}\ A_{ij}B_{ij}$. The final inequality with the quadratic form holds because $D(x)\succeq 0$ is a covariance matrix. Moreover, we have from Eq. A7 that

F.2 Lemma 6

F.3 Lemma 7

from Eqs. 8 and FP can be split into two operators

We first note that $L_{A}$ does not affect $F(\rho)$ in Theorem 5. For solutions of $\rho_{t}=L_{A}\ \rho$, we have

by 4. The dynamics of the anti-symmetric operator is thus completely deterministic and conserves $F(\rho)$. In fact, the equation Eq. A10 is known as the Liouville equation (Frank,, 2005) and describes the density of a completely deterministic dynamics given by

F.4 Theorem 22

All the matrices below depend on the weights $x$; we suppress this to keep the notation clear. Our original SDE is given by

We will transform the original SDE into a new SDE

where $S$ and $A$ are the symmetric and anti-symmetric parts of $G^{-1}$,

Since the two SDEs above are equal to each other, both the deterministic and the stochastic terms have to match. This gives

Using the above expression, we can now give an explicit, although formal, expression for the potential:

where $\Gamma:\to\Omega$ is any curve such that $\Gamma(1)=x$ and $\Gamma(0)=x(0)$ which is the initial condition of the dynamics in Eq. 3. Note that $\Phi(x)$ does not depend on $\beta$ because $G(x)$ does not depend on $\beta$.

We now write the modified SDE Eq. A12 as a second-order Langevin system after introducing a velocity variable $p$ with $q\triangleq x$ and mass $m$:

The key idea in Yin and Ao, (2006) is to compute the Fokker-Planck equation of the system above and take its zero-mass limit. The steady-state distribution of this equation, which is also known as the Klein-Kramer’s equation, is

where the position and momentum variables are decoupled. The zero-mass limit is given by

We now exploit the fact that $Q$ is defined to be an anti-symmetric matrix. Note that $\sum_{i,j}\partial_{i}\partial_{j}\left(Q_{ij}\rho\right)=0$ because $Q$ is anti-symmetric. Rewrite the third term on the last step ($*$) as

Observe that the brown terms are equal. Moving the blue terms together and matching the drift terms in the two Fokker-Planck equations then gives