O N THE INDUCTIVE BIAS OF STOCHASTIC GRADIENT DESCENT
Anonymous authors
Paper under double-blind review
A BSTRACT
Stochastic gradient descent (SGD) is widely believed to perform implicit regularization when used to train deep neural networks, but the precise manner in which this occurs has thus far been elusive.
We prove that SGD minimizes an average potential over the posterior distribution of weights along with an entropic regularization term.
This potential is however not the original loss function in general.
So SGD does perform variational inference, but for a different loss than the one used to compute the gradients.
Even more surprisingly, SGD does not even converge in the classical sense: we show that the most likely trajectories of SGD for deep networks do not behave like Brownian motion around critical points.
Instead, they resemble closed loops with deterministic components.
We prove that such “out-of-equilibrium” behavior is a consequence of the fact that the gradient noise in SGD is highly non-isotropic; the covariance matrix of mini-batch gradients has a rank as small as 1% of its dimension.
We provide extensive empirical validation of these claims, proven in the appendix.
I NTRODUCTION
Our ﬁrst result is to show precisely in what sense stochastic gradient descent (SGD) implicitly performs variational inference, as is often claimed informally in the literature.
For a loss function f
( x ) with weights x ∈ R d , if ρ ss is the steady-state distribution over the weights estimated by SGD,
[Equation]
; where H ( ρ ) is the entropy of the distribution ρ and η and (cid:98) are the learning rate and batch-size, respectively.
The potential Φ ( 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 gradients.
We next prove that the implicit potential Φ ( 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 satisﬁed for deep networks.
Empirically, we ﬁnd 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 ∇Φ ( x ) = 0, these are not the locations where ∇ 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, can be quite large in practice, and scales linearly with η / (cid:98) .
Indeed, 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.
These loops can be non-trivial if the deviation of Φ ( x ) from f ( x ) is large.
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.
B ACKGROUND ON CONTINUOUS - TIME SGD
Stochastic gradient descent performs the following updates while training a network
[Equation]
where η is the learning rate and ∇ f (cid:98) ( x k ) is the average back-propagation (back-prop) gradient over a mini-batch (cid:98) ,
[Equation]
We overload notation (cid:98) for both the set of examples in a mini-batch and its size.
We assume that weights belong to a compact subset Ω ⊂ R d , to ensure appropriate boundary conditions for the evolution of steady-state densities in SGD, although all our results hold without this assumption if the loss grows unbounded as (cid:107) x (cid:107) → ∞ , for instance, with weight decay as a regularizer.
Deﬁnition 1 ( Diffusion matrix D ( x ) ) .
If examples in a mini-batch are sampled with replacement, we show in Appendix A.1 that the variance of mini-batch gradients is var (cid:16) ∇ f (cid:98) ( x ) (cid:17) =
D ( x ) (cid:98) where
[Equation]
Note that D ( x ) is independent of the learning rate η and the batch-size (cid:98) .
It only depends on the weights x , architecture and loss deﬁned by f ( x ) , and the dataset.
We will often discuss two cases:
(i) isotropic diffusion when D ( x ) is a constant multiple of the identity, independent of x , and (ii) 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 updates in (SGD).
Lemma 2 ( Continuous-time SGD ) .
The continuous-time limit of SGD is given by
[Equation]
where W ( t ) is a Brownian motion and the parameter β is the inverse temperature deﬁned as
[Equation]
The steady-state density of the weights ρ ( x , t ) ∝ P (cid:0) x ( t ) =
x | x ( 0 ) (cid:1) , evolves according to the FokkerPlanck equation (Risken, 1996):
[Equation]
where the notation ∇ · v denotes the divergence ∇ · v =
∑ i ∂ x i v i ( x ) for any vector v ( x ) ∈ R d .
We refer to Li et al.
(2017b, Thm. 1) for the proof but note that (3) and (SGD) are equivalent in the sense that any function computed on trajectories of (3) differs by at most O ( dt ) from its value on trajectories of (SGD), the parameter η is identitiﬁed with the time discretization dt .
Note that β −completely captures the magnitude of noise in SGD that depends only upon the learning rate η and the mini-batch size (cid:98) .
Assumption 3 ( Steady-state distribution exists ) .
We assume that the steady-state distribution of the Fokker-Planck equation (FP) exists, this is denoted by ρ ss ( x ) and satisﬁes,
[Equation]
Note that the “probability current” J ss needs not be zero at steady-state, only its divergence does.
We will see that the case J ss = 0 corresponds to D ( x ) being isotropic while the general case ∇ · J ss =will correspond to a non-isotropic diffusion matrix.
SGD PERFORMS VARIATIONAL INFERENCE
In this section, we implicitly deﬁne the potential Φ ( x ) using the steady-state distribution ρ ss as
[Equation]
up to a constant.
The potential Φ ( x ) depends only on the full-gradient and the diffusion matrix, and will be made explicit in Section 5.
For now, we express ρ ss in terms of the potential using a normalizing constant Z ( β ) as
[Equation]
which is also the steady-state solution of
[Equation]
as can be veriﬁed by direct substitution in (FP).
The above observation is very useful because it suggests that, if ∇ f ( x ) can be written in terms of the diffusion matrix and a gradient term ∇Φ ( x ) , the steady-state distribution of this SDE is easily obtained.
We exploit this observation to rewrite ∇ f ( x ) in terms of two quantities, a term D ∇Φ that gives rise to the above steady-state and the remainder, a “conservative “force”
[Equation]
interpreted as the part of ∇ f ( x ) that cannot be written as D Φ (cid:48)
( x ) for some Φ (cid:48) .
We now make an important assumption on j ( x ) which has its origins in thermodynamics.
Assumption 4 ( Force j ( x ) is conservative ) .
We assume that
[Equation]
The Fokker-Planck equation (FP) typically models a physical system which exchanges energy with an external environment (Ottinger, 2005; Qian, 2014).
In our case, this physical system is the gradient descent part ∇ · ( ∇ f ρ ) while the interaction with the environment is the stochastic part β − 1 ∇ · ( D ∇ ρ ) .
The second law of thermodynamics states that the entropy of a system can never decrease; in Appendix B we show how the above assumption is sufﬁcient to satisfy the second law.
We also discuss some properties of j ( x ) in Appendix D that are a consequence of this.
The most important is that, out of the two forces j ( x ) and D ∇Φ , in (9), only D Φ ( x ) makes progress towards convergence to the steady-state ρ ss , whereas j ( x ) is always orthogonal to ∇ ρ ss .
This leads us to the main result of this section.
Theorem 5 ( SGD performs variational inference ) .
The functional
[Equation]
decreases monotonically along the trajectories of the Fokker-Planck equation (FP) and converges to its minimum, which is zero, at steady-state.
Moreover, we also have an energetic-entropic split
[Equation]
Theorem 5, proven in Appendix E.1 shows that SGD implicitly minimizes a combination of two terms: an “energetic” term, and an “entropic” term.
The ﬁrst is the average potential over a distribution ρ .
The steady-state of SGD in (7) is such that it places most of its probability mass in regions of the parameter space with small values of Φ .
The second shows that SGD has an implicit bias towards solutions that maximize the entropy of the distribution ρ .
Note that the energetic term in (12) has potential Φ ( x ) , instead of f ( x ) .
This is an important fact and the crux of this paper.
Lemma 6 ( Potential equals original loss iff isotropic diffusion ) .
If the diffusion matrix D ( x ) is isotropic, i.e., a constant multiple of the identity, the implicit potential is the original loss itself
[Equation]
This is proven in Appendix E.2.
The deﬁnition in (9) shows that j (cid:54) = 0 when D ( x ) is non-isotropic.
This results in a deterministic component in the SGD dynamics which does not affect the functional
F ( ρ ) , hence j ( x ) is called a “conservative force.”
The following lemma is proven in Appendix E.3.
Lemma 7 ( Deterministic dynamics ) .
The force j ( x ) does not decrease F ( ρ ) in (12) and introduces a deterministic component in SGD given by
[Equation]
[Equation]
The condition ∇ · j ( x ) = 0 in Assumption 4 implies that this dynamic makes SGD traverse closed trajectories in weight space.
3. C ONNECTION TO B AYESIAN INFERENCE
Note the absence of any prior in (12).
On the other hand, the evidence lower bound (Kingma and
Welling, 2013) for the dataset Ξ is,
[Equation]
where H ( q , p ) is the cross-entropy of the estimated steady-state and the variational prior.
The implicit loss function of SGD in (12) therefore corresponds to a uniform prior p ( x | Ξ ) .
In other words, we have shown that SGD itself performs variational optimization with a uniform prior.
Note that this prior is well-deﬁned because x ∈ Ω and Ω is compact.
It is important to note that SGD implicitly minimizes a potential Φ ( x ) instead of the original loss f ( x ) in (ELBO).
We prove in Section 5 that this potential is quite different from f ( x ) , in particular with respect to critical points.
Remark 8 ( Information in representations ) .
Working with ELBO in practice involves one or multiple steps of SGD to minimize the energetic term along with an analytically computable KLdivergence term which is often achieved using a factored Gaussian prior (Kingma and Welling, 2013).
As Theorem 5 shows, such an approach implicitly also enforces a uniform prior whose strength is determined by β − 1 and conﬂicts with the externally imposed Gaussian prior in ELBO.
This conﬂict, which fundamentally arises from using SGD to minimize the energetic term, has resulted in researchers modulating the strength of the Gaussian prior (Mandt et al., 2016) in ELBO using a pre-factor in the KL-divergence term.
This has also been shown to lead to invariant representations (Achille and Soatto, 2017) via the information bottleneck principle (Tishby et al., 1999).
3. P RACTICAL IMPLICATIONS
We will show in Section 5 that the potential Φ ( x ) does not depend upon the optimization process, it is only a function of the dataset and the architecture.
Therefore, the effect of two important parameters, the learning rate η and the mini-batch size (cid:98) completely determines the strength of the entropic regularization term.
If β − 1 → 0, the implicit regularization of SGD goes to zero.
This implies that
[Equation]
should not be small is a good tenet for regularization of SGD.
Remark 9 ( Learning rate should scale linearly with batch-size to generalize well ) .
In order to maintain the same level of entropic regularization, the learning rate η needs to scale linearly with the batch-size (cid:98) .
This theoretical prediction ﬁts 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).
Remark 10 ( Sampling with replacement is better regularized than without replacement ) .
The diffusion matrix for the case when mini-batches are sampled with replacement is almost the same as that of (2), see Appendix A.2.
However, the expression for the inverse temperature shows that
[Equation]
should not be small .
The extra factor of (cid:0) 1 − (cid:98) N (cid:1) in β (cid:48)− 1 reduces the entropic regularization in (12).
In fact, as (cid:98) → N , the inverse temperature β (cid:48) → ∞ .
As a consequence, for the same learning rate η and batch-size (cid:98) , Theorem 5 predicts that sampling with replacement has better regularization than sampling without replacement.
This effect is particularly pronounced at large batch-sizes.
E MPIRICAL CHARACTERIZATION OF SGD DYNAMICS
Section 4.1 shows that the diffusion matrix D ( x ) for modern deep networks is highly non-isotropic with very low rank.
We also analyze trajectories of SGD and detect periodic components using a frequency analysis in Section 4.2.
We consider three networks for these experiments: a convolutional network called small-lenet and a two-layer, fully-connected network named small-fc 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 CIFARdatasets (Krizhevsky, 2009), see Appendix C for more details.
We use networks with about 20 , 000 weights to keep the eigen-decomposition of D ( x ) ∈ R d × d tractable.
These networks however possess all the architectural intricacies such as convolutions, dropout, batch-normalization etc.
We evaluate D ( x ) using (2) with the network in evaluation mode.
4. H IGHLY NON - ISOTROPIC D ( x ) FOR DEEP NETWORKS
Figs. 1 and 2 show the eigenspectrum 1 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 .
Moreover, non-zero eigenvalues are spread across a vast range with a large variance.
[Figure]
Remark 11 ( Noise in SGD is largely independent of the weights ) .
The variance of noise in (3) is
[Equation]
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.
Remark 12 ( More non-isotropic diffusion if data is diverse ) .
The eigenspectra in Fig.
2 for
CIFAR-10 and CIFAR-100 have much larger eigenvalues and standard-deviation than those in Fig.
[Figure]
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 × 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.
2b are much higher than those in Fig.
2a.
Input augmentation increases the diversity of mini-batch gradients.
This is seen in Fig.
2c where the standard-deviation of the eigenvalues is much higher as compared to Fig.
2a.
4. A NALYSIS OF LONG - TERM TRAJECTORIES
We train a smaller version of small-fc on 7 × 7 down-sampled MNIST images (cf.
Appendix C)
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 ﬁrst 10 3 epochs of training (“burnin”) to ensure that SGD has reached the steady-state.
The learning rate is ﬁxed to 10 − 3 for the remaining 10 5 epochs.
[Figure]
Iterates of SGD after it reaches a neighborhood of a critical point (cid:107) ∇ f ( x k ) (cid:107) ≤ ε are expected to perform Brownian motion with variance var ( ∇ f (cid:98) ( x )) , the FFT in Fig.
3a would be ﬂat if this were so.
Instead, we see low-frequency modes in the trajectory that are indicators of a periodic dynamics of the force j ( x ) .
These modes are not sharp peaks in the FFT because j ( x ) can be a non-linear function of the weights thereby causing the modes to spread into all dimensions of x .
The FFT is dominated by jittery high-frequency modes on the right with a slight increasing trend; this suggests the presence of colored noise in SGD at high-frequencies.
The auto-correlation (AC) in Fig.
3b should be compared with the AC for Brownian motion which decays to zero very quickly and stays within the red conﬁdence bands (99%).
Our iterates are signiﬁcantly correlated with each other even at very large lags.
This further indicates that trajectories of SGD do not perform Brownian motion.
Fig.
3c 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 ofover 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 17.
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 implicit potential Φ ( x ) that SGD creates.
We also discuss implications of this for generalization in Section 5.3.
The fundamental difﬁculty in obtaining an explicit expression for Φ is that even if the diffusion matrix D ( x ) is full-rank, there need not exist a function Φ ( x )
such that ∇Φ ( x ) =
D − 1 ( x ) ∇ f ( x ) at all x ∈ Ω .
This precludes us from exploiting (8).
We therefore split the analysis into two cases: (i) a local analysis near any critical point ∇ f ( x ) = 0 where we linearize ∇ f ( x ) =
Fx and ∇Φ ( x )
=
Ux to compute U =
G − 1 F , and (ii) the general case which pertains to Fig.
4c where ∇Φ ( x ) cannot be written as a local rotation and scaling of ∇ f ( x ) .
We introduce these cases with an example from Noh and Lee (2015).
[Figure]
Example 13 ( Double-well potential with limit cycles ) .
Deﬁne Φ ( x ) =
( x 2 1 − 1 ) 2+ x 2 2 2 .
Instead of constructing a diffusion matrix D ( x ) , we will directly construct different gradients ∇ f ( x ) that lead to the same potential Φ ; these are equivalent but the later is much easier.
The dynamics is given by dx = − ∇ f ( x ) dt + √ 2 dW ( t ) where ∇ f ( x ) = − j ( x )+ ∇Φ ( x ) .
We pick j = λ e Φ J ss ( x ) for some parameter λ >
where J ss ( x ) =
e − ( x 21 + x 22 ) 2( − x 2 , x 1 ) .
Note that this satisﬁes (7) and does not change ρ ss = e − Φ .
Fig.
4 shows the gradient ﬁeld f ( x ) along with a discussion.
5. L INEARIZATION AROUND A CRITICAL POINT
Without loss of generality, let x = 0 be a critical point of f ( x ) .
This critical point can be a local minimum, maximum, or even a saddle point.
We linearize the gradient around the origin and deﬁne a ﬁxed matrix F ∈ R d × d (the Hessian) to be ∇ f ( x ) =
Fx .
Let D
= D ( 0 ) be the constant diffusion matrix matrix.
The dynamics in (3) can now be written as
[Equation]
Lemma 14 ( Linearization ) .
The matrix F in (15) can be uniquely decomposed into
[Equation]
[Equation]
D and Q are the symmetric and anti-symmetric parts of a matrix G with GF (cid:62) − FG (cid:62) = 0 , to get Φ ( x ) = 12 x (cid:62) 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.
for the proof that linearizes around any critical point.
Remark 15 ( Rotation of gradients ) .
We see from Lemma 14 that, locally near a critical point,
[Equation]
This suggests that the effect of j ( x ) is to rotate the local gradient ﬁeld, this is also seen in Fig.
4b.
5. G ENERAL CASE
The key idea in the general case is that all the expressions in Lemma 14 and (17) are kept the same but with the matrices now dependent on the weights x .
[Equation]
We have again decomposed the gradient ∇ f ( x ) into the part that comes from a potential ∇Φ ( x ) and the extra term j off ( x ) .
It is already clear from (18) that the critical points of f ( x ) are different from the most likely locations of ρ ss if j off ( x ) (cid:54) = 0.
Assumption 16 ( Steady-state does not change ) .
As done in Ao et al.
(2007), we will assume that the force j off ( x ) does not change the steady-state distribution, i.e., we still have ρ ss ∝
e
− β Φ .
This effectively amounts to assuming that the symmetric matrix D ( x ) and an anti-symmetric matrix Q ( x ) captures most of the gradient ∇ f ( x ) that results from a potential Φ ( x ) .
The assumption above is often used in non-equilibrium studies in physics and biology because it explains experimental data (Wang et al., 2008), its implications have been supported by perturbationbased calculations (Kwon and Ao, 2011) and mathematical analyses such as Tel et al. ; Qian (2015); Kaiser et al. (2017) do not invalidate it.
It also explains our experimental data from Section 4.2.
The main result of this section now follows.
Theorem 17 ( Most likely locations are not the critical points of the loss ) .
The most likely locations of SGD where ∇Φ ( x ) = 0 are different from the critical points of the original loss ∇ f ( x ) = 0 , More precisely, for j off deﬁned in (18),
[Equation]
(19) where the divergence operator is applied to the matrix Q ( x ) column-wise.
The matrix Q ( x ) and the potential Φ ( x ) can be explicitly computed in terms of the gradient ∇ f ( x ) and the diffusion matrix
D ( x ) .
In particular, Φ ( x ) does not depend on β .
See Appendix E.4 for the proof.
The time spent by a Markov chain at a state x is proportional to its steady-state distribution ρ ss ( x ) .
While it is easily seen that SGD does not converge in the Cauchy sense due to the stochasticity, it is very surprising that it may spend a signiﬁcant amount of time away from the critical points of the original loss.
If Q ( x ) has a large divergence, the set of states with ∇Φ ( x ) = 0 might be drastically different than those with ∇ f ( x ) = 0.
The example in Fig.
4c also demonstrates this effect.
If the diffusion matrix D ( x ) is such that the magnitude of the conservative force j ( x ) is large, the trajectories can be highly non-trivial; as Fig.
4 demonstrates, SGD may even converge around a saddle point.
This also closes the logical loop we began in Section 3 where we assumed the existence of ρ ss and deﬁned the potential Φ using it.
Lemma 14 and Theorem 17 show that both can be deﬁned uniquely in terms of the original quantities, i.e. , the gradient term ∇ f ( x ) and the diffusion matrix D ( x ) .
There is no ambiguity as to whether the potential Φ ( x ) results in the steady-state ρ ss ( x ) or vice-versa.
Remark 18 ( Consistent with the linear case ) .
Theorem 17 presents a picture that is completely consistent with Lemma 14.
If j ( x ) = 0 and Q ( x ) = 0, or if Q is a constant like the linear case in Lemma 14, the divergence of Q ( x ) in (19) is zero.
Remark 19 ( Out-of-equilibrium effect can be large even if D is constant ) .
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 Φ ( x ) such that ∇Φ ( x ) =
D − 1 ∇ 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.
Remark 20 ( Out-of-equilibrium effect increases with β − 1 ) .
The effect predicted by (19) becomes more pronounced if β − 1 = η 2 (cid:98) 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 β − 1 → 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.
5. G ENERALIZATION
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 ﬁnd these well-generalizable solutions easily.
This idea has also been successful in deep learning where Chaudhari et al.
(2016) modiﬁed 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 f γ ( x ) = − log (cid:16) G γ ∗ e − f ( x ) (cid:17) where
G γ is a Gaussian kernel of variance γ .
Even with an isotropic diffusion matrix, the steady-state distribution with f γ ( x ) as the loss function is ρ ss γ ( x ) ∝ e − β f γ ( x ) .
For large values of γ , the new loss makes the original local minima exponentially less likely.
In other words, local entropy does not rely on non-isotropic gradient noise to obtain out-of-equilibrium behavior, it gets it explicitly, by construction.
This is also seen in Fig.
4c: if SGD is drastically out-of-equilibrium, it converges around the “wide” saddle point region at the origin which has a small local entropy.
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.
R ELATED WORK
SGD, variational inference and implicit regularization: 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 Φ that it implicitly constructs given an architecture and a dataset.
It is widely believed that SGD is an implicit regularizer, see Zhang et al.
; 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 Φ 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.
; Anandkumar and Ge (2016);
Ge et al.
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, 2016).
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 ∑ k (cid:107) ∇ f k ( x ) (cid:107) / (cid:107) ∇ f ( x ) (cid:107) , 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 Φ ( x ) have employed the idea of designing a force j = ∇Φ − ∇ f to accelerate convergence, see Ma et al.
(2015) for a thorough survey.
We instead compute the potential Φ given ∇ f and D , which necessitates the use of techniques from physics.
In fact, our results show that very simple algorithms such as SGLD by Welling and Teh (2011) also beneﬁt from the acceleration that their sophisticated counterparts aim for (Ding et al., 2014; Chen et al., 2016), if they use deep networks.
D ISCUSSION
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)
A D IFFUSION MATRIX D ( x )
In this section we denote g k : =
∇ f k ( x ) and g : =
∇ f ( x ) =
1 N ∑ Nk = 1 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 .
A. W ITH REPLACEMENT
Let i 1 ,..., i (cid:98) be (cid:98) iid random variables in { 1 , 2 ,..., N } .
We would like to compute
[Equation]
Note that we have that for any j (cid:54)
=
k , the random vectors g i j and g i k are independent.
We therefore have
[Equation]
We use this to obtain
[Equation]
We will set
[Equation]
and assimilate the factor of (cid:98) − 1 in the inverse temperature β .
A. W ITHOUT REPLACEMENT
Let us deﬁne an indicator random variable 1 i ∈ (cid:98) that denotes if an example i was sampled in batch (cid:98) .
We can show that
[Equation]
and for i (cid:54) =
j ,
[Equation]
Similar to Li et al.
(2017a)
[Equation]
we can now compute
[Equation]
We will again set
D (
[Equation]
=
N
[Equation]
and assimilate the factor of (cid:98) − 1 (cid:0) 1 − (cid:98) N (cid:1) that depends on the batch-size in the inverse temperature β .
B D ISCUSSION ON A SSUMPTION
Let F ( ρ ) be as deﬁned in (12).
In non-equilibrium thermodynamics, it is assumed that if the system is not far from equilibrium, i.e., if (cid:107) j ( x ) (cid:107) is small, the local entropy production is a product of the force − ∇ (cid:16) δ F δρ (cid:17) from (A7) and the probability current − J ( x , t ) from (FP).
This assumption in this form was ﬁrst introduced by Prigogine (1955), see Frank (2005, Sec. 4.5) for a mathematical treatment.
The total entropy increase is given by
[Equation]
This can now be written using (A7) again as
[Equation]
The ﬁrst term in the above expression is non-negative, in order to ensure that dS i dt ≥ 0, we require
[Equation]
where the second equality again follows by integration by parts.
It can be shown (Frank, 2005, Sec.
that the condition in Assumption 4, viz., ∇ · j ( x ) = 0, is sufﬁcient to make the above integral vanish and therefore for the entropy generation to be non-negative.
C E XPERIMENTAL SETP
The three networks used in the experiments in Section 4 are as follows.
(i) small-lenet: a smaller version of LeNet (LeCun et al., 1998) on MNIST with batchnormalization 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 13338 weights and reaches about 0 .
75% training and validation error.
(ii) small-fc: a fully-connected network with two-layers, batch-normalization and rectiﬁed linear units that takes 7 × 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 andoutput classes (30 , 000 input images); this is called tiny-fc .
(iii) small-allcnn: this a smaller version of the fully-convolutional network for CIFAR-10/ introduced by Springenberg et al.
(2014) with batch-normalization and 12 , 24 output channels in the ﬁrst 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 .
We do not use any data-augmentation and pre-process data using global contrast normalization with ZCA for CIFAR-10 and CIFAR-100.
D S OME PROPERTIES OF THE FORCE j ( x )
The Fokker-Planck equation (FP) can be written in terms of the probability current as
[Equation]
Since we have ρ ss ∝ e − β Φ ( x ) , we also have that
[Equation]
which gives
[Equation]
In other words, the conservative force is non-zero only if detailed balance is broken, i.e., J ss (cid:54) = 0.
We also have
[Equation]
which shows using Assumption 4 and ρ ss ( x ) > 0 for all x ∈ Ω that j ( x ) is always orthognal to the gradient of the potential
[Equation]
Using the deﬁnition of j ( x ) in (9), we have detailed balance when
[Equation]
E P ROOFS E. T HEOREM
The KL-divergence is non-negative F ( ρ ) ≥ 0 with equality if and only if ρ = ρ ss .
The expression in (12) follows after writing
[Equation]
We now show that dF ( ρ ) dt ≤ 0 with equality only at ρ = ρ ss when F ( ρ ) reaches its minimum and the
Fokker-Planck equation achieves its steady-state.
The ﬁrst variation (Santambrogio, 2015) of F ( ρ ) computed from (12) is
[Equation]
which helps us write the Fokker-Planck equation (FP) as
[Equation]
A
[Equation]
Together, we can now write
[Equation]
As we show in Appendix B, the ﬁrst term above is zero due to Assumption 4.
Under suitable boundary condition on the Fokker-Planck equation which ensures that no probability mass ﬂows across the boundary of the domain ∂ Ω , after an integration by parts, the second term can be written as
[Equation]
In the above expression, A : B denotes the matrix dot product A : B =
∑ ij A ij B ij .
The ﬁnal inequality with the quadratic form holds because D ( x ) (cid:23) 0 is a covariance matrix.
Moreover, we have from (A6)
that
[Equation]
E. L EMMA
The forward implication can be checked by substituting ρ ss ( x ) ∝ e − c β f ( x ) in the Fokker-Planck equation (FP) while the reverse implication is true since otherwise (A4) would not hold.
E. L EMMA
The Fokker-Planck operator written as
[Equation]
from (9) and (FP) can be split into two operators
[Equation]
where the symmetric part is
[Equation]
and the anti-symmetric part is
[Equation]
We ﬁrst note that L A does not affect F ( ρ ) in Theorem 5.
For solutions of ρ t =
L A ρ , we have
[Equation]
by Assumption 4.
The dynamics of the anti-symmetric operator is thus completely deterministic and conserves F ( ρ ) .
In fact, the equation (A9) is known as the Liouville equation (Frank, 2005) and describes the density of a completely deterministic dynamics given by
[Equation]
where j ( x ) =
J ss / ρ ss = − ∇ f + D Φ ( x ) from Appendix D.
E. T HEOREM
All the matrices below depend on the weights x ; we suppress this to keep the notation clear.
Our original SDE is given by
[Equation]
The matrix G again satisﬁes the the same conditions as (16) that it is the symmetric part of D and ∇Φ ( x ) =
G − 1 ∇ f ( x ) is the gradient of a scalar function, i.e., its curl is zero:
[Equation]
The curl is an anti-symmetric matrix ( ∇ × v ) ij =
∂ i v j − ∂
j v i for any vector ﬁeld v .
Note that G is deﬁned uniquely using the two conditions in (A11).
The ﬁrst condition that D is the symmetric part of G adds d ( d + 1 ) / 2 constraints while the second condition on the curl adds another d ( d − 1 ) /constraints.
We will transform the original SDE into a new SDE
[Equation]
where S and A are the symmetric and anti-symmetric parts of G − 1 ,
[Equation]
Since the two SDEs above are equal to each other, both the deterministic and the stochastic terms have to match.
This gives
[Equation]
Using the above expression, we can now give an explicit, although formal expression for the potential as
[Equation]
where Γ : [ 0 , 1 ] → Ω is any curve such that Γ ( 1 ) = x and Γ ( 0 )
=
x ( 0 ) which is the initial condition of the dynamics in (3).
Note that Φ ( x ) does not depend on β because G ( x ) does not depend on β .
Now deﬁne the difference
[Equation]
This is really Assumption 16 in action, such a Q may not be anti-symmetric.
This is a technical trick, not uncommon in physics, and its validation comes from experimental evidence of its implications.
We now write the modiﬁed SDE (A12) as a second-order Langevin system after introducing a velocity variable p with q (cid:44) x and mass m :
[Equation]
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
[Equation]
where the position and momentum variables are decoupled.
The zero-mass limit is given by
[Equation]
We now exploit the fact that Q is deﬁned to be an anti-symmetric matrix.
Rewrite the third term on the last step ( ∗ ) as
[Equation]
We have assumed that the steady-state distribution of (3) exists and this also means that the potential Φ = − β − 1 log ρ ss exists.
Now consider the Fokker-Planck equation for the force decomposed like (18),
[Equation]
Since both (A15) and (A18) have the same steady-state
[Equation]
we can compare the two equations to get
[Equation]