SGD C AN C ONVERGE TO L OCAL M AXIMA
Anonymous authors
Paper under double-blind review
A BSTRACT
Stochastic gradient descent (SGD) is widely used for the nonlinear, nonconvex problem of training deep neural networks, but its behavior remains poorly understood.
Many theoretical works have studied SGD, but they commonly rely on restrictive assumptions about the nature of its noise.
In this work, we construct example optimization problems illustrating that, when not in the regimes that the previous works often assume, SGD can exhibit many strange behaviors that the established wisdom of the ﬁeld cannot explain.
Our constructions show that (1)
SGD can converge to local maxima, (2) SGD might only escape saddle points arbitrarily slowly, (3) SGD can prefer sharp minima over ﬂat ones, and (4) AMSGrad can converge to local maxima.
We realize our most surprising results in a simple neural network-like construction, suggesting their relevance to deep learning.
I NTRODUCTION
SGD is the main optimization algorithm for training deep neural networks, and understanding SGD is widely regarded as a key step on the path to understanding deep learning (Bottou, 2012; Zhang et al., 2018;
Xing et al., 2018; Mori et al., 2021; Du et al., 2018; Allen-Zhu et al., 2018; Wojtowytsch, 2021a;b; Gower et al., 2021; Ziyin et al., 2021; Gurbuzbalaban et al., 2021; Zou et al., 2021; Li et al., 2021; Feng and Tu, 2021).
Despite its algorithmic simplicity (describable with only two lines of equations), SGD is hard to understand.
The difﬁculty is threefold: (1) SGD is discrete-time in nature, and discrete-time dynamics are typically much more complicated to solve than their continuous-time counterparts (May, 1976); (2) SGD noise is state-dependent (Ziyin et al., 2021; Hodgkinson and
Mahoney, 2020); and (3) the loss landscape can be nonlinear and non-convex, involving many local minima, saddle points, and degeneracies (the Hessian is not full-rank) in the landscape (Xing et al., 2018; Wu et al., 2017).
Each of these difﬁculties is so challenging that very few works attempt to deal with all of them simultaneously.
Most previous works on SGD in the context of deep learning are limited to the cases when the loss landscape is strongly convex (Ziyin et al., 2021; Liu et al., 2021;
Hoffer et al., 2017; Mandt et al., 2017), or when the noise is assumed to be Gaussian and time-independent (Zhu et al., 2019; Jastrzebski et al., 2017; Xie et al., 2021); for the works that tackle SGD in a non-convex setting, often strong conditions are assumed.
The reliance on strong assumptions regarding each of the challenges means that our present understanding of SGD for deep learning could be speculative.
This work aims to examine some commonly held presuppositions about
SGD and show that when all the three challenging factors are taken together, many counter-intuitive phenomena may arise.
In this work, we study the behavior of SGD in a landscape with non-convex loss and multi-minima.
Our aim is the direct opposite of that of many related works: instead of showing when SGD will converge, we set out to answer the question of when SGD does not.
In particular, the problem setting considers discrete-time SGD close to saddle points, where the noise is due to minibatch sampling, and the learning rate is held constant throughout training.
In the next section, we deﬁne the SGD algorithm and the necessary notations.
In Sec. 3, we discuss the related works.
A warmup example is provided in Sec. 4.
Sec.
5 introduces the main results.
Sec.
6 presents the numerical simulations, including a minimal example involving a neural network.
We also encourage the readers to examine the appendix.
Sec.
A presents additional experiments.
Sec.
B discusses additional theoretical concerns.
Sec.
C presents a continuous-time analysis of the problems in the main text and is relevant for discussing the unique discrete-time features of SGD.
Sec.
D presents all the proofs.
B ACKGROUND
Notation and Terminology .
We use λ to denote the learning rate.
w is the model parameter.
ˆ L ( w ; x ) , with the caret symbol ˆ ⋅ , denotes the sampled loss function, which is a function of the data point, and
[Figure]
L ∶= E x [ ˆ L ] is the true loss landscape averaged over the data distribution; for notational conciseness, we often hide the dependence of ˆ L on x when the context is clear.
The lowercase t denotes the time step of optimization.
In this work, the terms learning rate and step size are used interchangeably.
The phrases loss , objective , and potential are also used interchangeably.
Now, we introduce the minibatch SGD algorithm.
The objective of SGD can be deﬁned as a pair ( ˆ L,p ( x ))
such that ˆ L is a differentiable function and p ( x ) is a probability density.
We aim at ﬁnding the minimizer of the following differentiable objective:
[Equation]
where x is a data point drawn from distribution p ( x ) and w ∈ R D is the model parameter.
The gradient descent (GD) algorithm for L is deﬁned as the update rule w t = w t − 1 − λ ∇ w L ( w,x ) for some randomly initialized w 0 , where λ is the learning rate.
The following deﬁnition deﬁnes SGD.
Deﬁnition 1.
The minibatch SGD algorithm computes the update to the parameter w with the following equation: w t = w t − 1 − λ ˆ g t , where ˆ g t = ∇ ˆ L ( w t − 1 ; x t ) , and x t is drawn from p ( x )
such that x i is independent of x j for i ≠ j .
R ELATED W ORKS
Escaping Saddle Points .
Because neural networks are highly nonlinear, the landscape of neural networks is believed to have many local minima and saddle points (Du et al., 2017; Kleinberg et al., 2018; Reddi et al., 2018b).
It is thus crucial for any optimization algorithm to be able to overcome saddle points efﬁciently and not become stuck in a suboptimal minimum.
However, the majority of works on escaping saddle points assume no stochasticity or only artiﬁcially injected noise (Du et al., 2017; Jin et al., 2017;
Ge et al., 2015; Reddi et al., 2018b; Lee et al., 2016).
Some physics-inspired approaches take the stochasticity into account but often rely on continuous-time approximation and nontrivial assumptions regarding the stationary distribution of the parameter (Mori et al., 2021; Xie et al., 2021; Liu et al., 2021).
A few works that do consider the stochasticity in ﬁnite step-size SGD rely on strong assumptions about the noise, whose relevance to deep learning is unclear (see the next paragraph).
Role of Minibatch Noise .
One indispensable aspect of SGD is its special stochasticity due to minibatch sampling (Wu et al., 2020; Ziyin et al., 2021; Zhu et al., 2019; Hodgkinson and Mahoney, 2020), whose strength and structure are determined by the model architecture and the distribution of the data points.
In fact, it has been shown that the models trained with SGD signiﬁcantly outperform the models directly trained with GD (Hoffer et al., 2017), which suggests that minibatch noise in
SGD is beneﬁcial for training.
Understanding the role of minibatch noise is therefore of fundamental and practical value.
However, in the previous works, the focus is on analyzing the behavior of
SGD on the landscape speciﬁed by L ; the unique structures due to the noise in SGD are often treated in an oversimpliﬁed manner without referring to the actual models in deep learning or the actual data distributions.
For example, Daneshmand et al.
(2018) assumes that the noisy gradient is negatively correlated with the direction of the Hessian with the smallest eigenvalue, the validity of which is unclear when the Hessian is not full rank.
Kleinberg et al.
(2018) assumes that the noise in SGD makes the landscape one-point strongly convex, which implies that the loss landscape before the convolution is already close to a strongly convex landscape.
Another commonly assumed condition is that the loss-landscape satisﬁes the Lojasiewicz condition, which is non-convex but implies the unrealistic condition that there is no local minimum except the global minimum (Karimi et al., 2020; Wojtowytsch, 2021a; Vaswani et al., 2019), which is not sufﬁcient to understand the complicated dynamics that often happen in a setting with many minima.
In Sec.
6.3, we show that these assumptions do not hold even for the simplest kind of nonlinear neural network with two layers and one hidden neuron.
The central message of our work is that, without making model and task-dependent assumptions regarding the SGD noise, the behavior of the SGD algorithm may be arbitrary and independent of the underlying landscape.
Large Learning Rate Regime .
Recently, it is realized that networks trained at a large learning rate have a dramatically better performance than networks trained with a vanishing learning rate (lazy training)
(Chizat and Bach, 2018).
Lewkowycz et al.
(2020) shows that there is a qualitative difference between the lazy training regime and the large learning rate regime; the performance features two plateaus in testing accuracy in the two regimes, with the large learning rate regime performing much better.
However, almost no previous theory exists for understanding SGD at a large learning rate.
Our work is also relevant for understanding what happens at a large learning rate.
A W ARM - UP E XAMPLE
This section studies a special example of the main results of our work.
While the setting of this example is restricted, it captures the essential features and mechanisms of SGD that we will utilize to prove the main results.
Consider the following loss function:
[Equation]
where x,w ∈ R and x is drawn from an underlying distribution p ( x ) at every time-step.
We let E [ x ] = a and assume that Var [ x ] = σ 2 is ﬁnite.
This gives rise to a true, sample-averaged loss of L ( w ) = a 2 w 2 .
The stationary point of this loss function is w = 0 .
When a > 0 , the underlying (deterministic) landscape is a minimum; the a > 0 case is now well understood in the discrete-time limit and when the underlying noise is state-dependent (Liu et al., 2021).
When a < 0 , the point w =is a local maximum; one hopes that the underlying dynamics escape to inﬁnity, and our goal is to understand the behavior of SGD in this case.
The following proposition and proof show that for the seemingly benign data distribution p ( x ) (zero-mean and bounded) above, SGD cannot escape from the local maximum when the learning rate is set to be 1 .
Proposition 1.
Let λ = 1 and a < 0 and p ( x ) be the distribution such that x = 1 and x = − 1 + a with equal probability.
Then, w converges to w = 0 with probability 1 with the SGD algorithm, independent of the initialization.
Proof .
By the deﬁnition of the SGD algorithm, we have
[Equation]
Therefore, after t time steps, w t has at most 2 − t probability of being non-zero.
For inﬁnite t , w t =with probability 1 . ◻
This simple toy example illustrates a few out-of-expectation points.
First of all, it suggests that one cannot use the expected value of w to study the escape problem of SGD.
The expectation value of w t is
[Equation]
which is an exponentially fast escape from any initialization w 0 for a < 0 .
In fact, ln ( 1 − a ) is exactly the escape time scale of the GD algorithm.
This is counter-intuitive: SGD will converge to the local maximum with probability 1 , even if its expected value escapes the local maximum exponentially fast.
Alternatively, one might also hope to use the expected loss or the norm E w t [ L ( w t )] ∼ ∣∣ w t ∣∣ 22 as the metric of escape, but it is not hard to see that they suffer the same problem.
At time t , because w t =with probability 1 − 2 − t and w t = ( 2 − a ) t w 0 with probability 2 − t , we have E [ L ( w t )]
∝ w 20 e 2 t ln 2 − a √ 2 , which also predicts an exponentially fast escape.
Again, this fails to reﬂect the fact that SGD has only a vanishingly small probability of escaping the constructed local maximum.
Statistically speaking, the norms of w t fail to be good metrics to measure the escape rate because these metrics are not robust against outliers.
In our example, there is only 2 − t probability for SGD to escape the local minimum, yet the speed at which this outlier escapes the maximum outweighs the decay in its probability through time, thus contributing more to the expected norm than any other events.
This example also shows the difﬁculty and subtlety of dealing with the escape problem.
In essence, the results in the rest of this work are extensions of this special case to less restrictive settings.
M AIN R ESULTS
In this section, we present four main theoretical results: (1) SGD may converge to a local maximum; (2) SGD may escape a saddle point arbitrarily slowly; (3) SGD may prefer sharp minima over ﬂat ones; and (4) AMSGrad may converge to a local maximum.
Numerically, we show that these results are relevant for a minimal example involving a neural network, where the common assumptions regarding SGD are violated (Sec. 6.3).
5. SGD M AY C ONVERGE TO A L OCAL M AXIMUM
We show that there are cases where SGD may fail to escape saddle points, while GD can successfully escape; this contradicts the suggestions in Kleinberg et al.
(2018) that SGD always escapes a local minimum faster than GD.
It is appropriate to begin with the deﬁnition of a saddle point.
Deﬁnition 2.
w is said to be a stationary point of L ( w ) if ∇ w L ( w ) = 0 .
w is a minimum of L if, for all w ′ in a sufﬁciently small neighborhood of w , L ( w ′ ) ≥ L ( w ) .
w is said to be a saddle point if w is a stationary point but not a local minimum.
This deﬁnition of saddle points, which includes local maxima, agrees with the common deﬁnition in the literature (Daneshmand et al., 2018).
The following proposition generalizes our warmup example to a family of 1D problems in which SGD converges to a local maximum.
Proposition 2.
Let Eq.
2 be the loss function.
Let λ > 0 , and p ( x ) be such that µ ∶= E
[ ln ∣ 1 − λx ∣] and s 2 ∶= Var [ ln ∣ 1 − λx ∣] ; then SGD converges to the local maximum in probability if µ < 0 .
Proof Sketch .
We ﬁrst show that 1 √ t ln ∣ w t / w 0 ∣ obeys a normal distribution as t → ∞ .
This allows us to deduce the cumulative distribution function (c.d.f.) of w as t → ∞ .
The c.d.f. has an bifurcative dependence on the sign of µ ∶= E [ ln ∣ 1 − λx ∣] .
When µ < 0 , the distribution of w t converges to 0 in probability.
See Sec.
D.2 for the detailed proof.
◻
Corollary 1.
Let Eq.
2 be the loss function and p ( x ) = 12 δ
( x − 1 )
+ 12 δ ( x + 1 − a ) ; then SGD converges to the local maximum in probability if
[Equation]
Note that the convergence is independent of the actual shape of the the distribution p ( x ) , and so may apply to many kinds of distributions.
For simplicity, we consider the case when p ( x ) is a
Bernoulli variable.
There are three different regimes/phases for this simple escaping problem (also see Fig.
1-Right).
The escape regime λ > a −√
[Equation]
) is also present when a > 0 ; it means that in this regime, SGD will also escape w 0 even if it is a local minimum.
This regime is not present if we perform a continuous-time analysis (Sec. C), showing that this region is due to the instability of the discrete-time SGD algorithm.
This kind of escaping is undesirable because, in this regime, the local gradient cannot provide any guidance for minimizing the loss.
The region
[Equation]
( a − 1 ) is the trapped regime.
This is due to the special structure of the minibatch noise – a comparison with the continuous-time analysis suggests that SGD will not be trapped if the noise is not position dependent
(Sec. C.1).
The small learning rate regime λ <
a a − 1 is the successful escape regime because the SGD noise is of order λ 2 and becomes negligible at a small λ (Liu et al., 2021).
This regime corresponds to the lazy learning regime of training: with negligible noise and vanishing gradient, SGD is known to optimize neural networks well (Du et al., 2018).
This discussion also justiﬁes our later choice for placing an upper bound to λ and focus only on the second and third regime in Sec.
5.3 and 5.4.
It is interesting to compare this result with the previous results on the convergence of SGD to local minima in, for example, Ge et al.
In comparison, the SGD algorithm is not guaranteed to escape the saddle point in this example because two crucial assumptions of Ge et al.
(2015) is violated: (1) the gradient does not have a bounded ﬂuctuation around its mean (and so even parameters very far away from the saddle can be brought back close to the saddle point due to the noise), and (2) we do not inject additive noise to the gradient.
The standard theoretical guarantee of convergence in Ge et al.
(2015) thus does not apply to this example.
One naturally wonders what is the cause of the convergence to a local maximum.
We thus perform a continuous-time analysis in Section C of this example, and the continuous-time result shows that the convergence to local maxima occurs when the minibatch noise dominates the gradient, namely, when the signal to noise ratio becomes smaller than 1 .
This suggests the importance of minibatch noise in inﬂuencing the dynamics of SGD.
5. SGD M AY E SCAPE S ADDLE P OINTS A RBITRARILY S LOWLY
In the previous section, we have shown that one can adversarially construct a loss function for a ﬁxed learning rate such that the SGD algorithm converges to a local maximum.
This section shows that, even if one can tune the learning rate at will, there is a loss landscape where SGD can take an arbitrarily long time to escape, no matter how one chooses the learning rate.
We begin by deﬁning the escape rate.
Deﬁnition 3.
The asymptotic average escape rate of w at learning rate λ , γ ( λ ) , is deﬁned as γ ∶= lim t →∞ 1 t E w t
[ ln ∣ w t ∣] .
The optimal escape rate is deﬁned as γ ∗ ∶= sup λ γ ( λ ) .
Remark.
Note that the deﬁnition above can be written as
[Equation]
for a large τ .
This alternative deﬁnition provides a way to estimate the escape rate during a ﬁnite time.
Also, the readers familiar with the theory of dynamical systems will immediately notice this deﬁnition mirrors the deﬁnition of the Lyapunov exponent in the study of deterministic chaos (Schuster and Just, 2006; Lyapunov, 1992).
This deﬁnition can be seen as a generalization of the Lyapunov exponent to a probabilistic setting.
We note that this deﬁnition of escape rate differentiates this work from the standard literature on escaping saddles, where the focus is on the time scale for reaching a local minimum when saddle points exist (Ge et al., 2015; Daneshmand et al., 2018; Jin et al., 2017), while our deﬁnition focuses on the time scale of escaping a speciﬁed saddle point.
These two time scales might be vastly different.
For example, when the loss is the cross-entropy, the global minimum is never reachable, yet escaping saddle points should be important and happen at a reasonable time.
This deﬁnition is, therefore, more realistic and natural for studying escaping saddle points.
As discussed in the previous section, the undesired escape can be achieved by setting the learning rate to be arbitrarily large; however, this is not possible in practice because the learning rate of SGD is inherently associated with the stability of learning (Ziyin et al., 2021), and λ usually has to be much smaller than 1 .
Proposition 3.
Let the loss function be Eq.
(2) , a < 0 and p ( x ) ∼ 12 δ ( x − 1 ) + 12 δ ( x − a + 1 ) , where δ is the Dirac delta function.
Assuming that λ ≤ 1 , SGD algorithm escapes the saddle point of Eq.
(2) arbitrarily slowly, i.e., for any (cid:15) > 0 , there exists loss function such that sup λ γ ( λ ) < (cid:15) .
Speciﬁcally, γ ∗ =
ln 2 − a 2 √ 1 − a , and, for any (cid:15) > 0 , γ ∗ < (cid:15) if ∣ a ∣ ≤ 2 ∣− e (cid:15) √ e 2 (cid:15)
[Table]
Intuitively, one expects escaping to be easier as we increase the learning rate (Kleinberg et al., 2018), and that the escape rate should monotonically increase as we increase the learning rate.
This example conveys a surprising message that the escaping efﬁciency is not monotonically increasing as the learning rate increases.
For this example, the escape rate ﬁrst decreases and starts to increase only when λ is quite large.
This example also shows the subtlety in the escaping problem.
On the one hand, one needs to avoid a too large step size to avoid training instability.
On the other hand, one cannot use a too small learning rate because a small learning rate also makes optimization slow.
This interesting example suggests that there must be a tradeoff between learning speed and learning stability and that we may construct a general theory for ﬁnding the best learning rate that achieves the best tradeoff; however, this tradeoff problem is a sufﬁciently complex problem on its own and is beyond the scope of the present work.
5. SGD M AY P REFER S HARPER M INIMA
Modern deep neural networks are often overparametrized and capable of memorizing all the training data points quite easily.
This means that the traditional metrics such as Rademacher complexity cannot be used to guarantee the generalization capability of SGD (Zhang et al., 2017).
Nevertheless, neural networks are found to generalize very well.
This inspired the hypothesis that SGD, the main training algorithm of neural networks, contains some implicit regularization effect such that it biases the neural network towards simpler solutions (Neyshabur et al., 2017; Soudry et al., 2018).
One hypothesized mechanism of this regularization is that SGD selects ﬂat minima over the sharp ones (Hochreiter and Schmidhuber, 1997; Meng et al., 2020; Xie et al., 2021;
Liu et al., 2021; Mori et al., 2021; Smith and Le, 2018; Wojtowytsch, 2021b).
In this work, we show that this may not be the case because the dynamics and convergence of SGD crucially depend on the underlying mini-batch noise, while the sharpness of the landscape is independent of the noise.
Before introducing the results, let us ﬁrst deﬁne the sharpness of a local minimum.
Deﬁnition 4.
Let w = w ∗ be a local minimum of the loss function L ( w ) ; the sharpness s ( w ∗ ) of a local minimum w ∗ is deﬁned as s ( w ∗ ) ∶=
Tr [∇ 2 w L ( w ∗ )] .
We say that a minimum w ∗ 1 is sharper than w ∗ 2 if and only
if s ( w ∗ 1 ) > s ( w ∗ 2 ) .
In this deﬁnition, the sharpness is the trace of the Hessian of the quadratic approximation to the local minimum.
Other possible deﬁnitions such as the determinant of the Hessian are essentially similar to this deﬁnition (Dinh et al., 2017).
For an explicit construction, consider the following objective on a 2-dimensional landscape:
[Equation]
for a,b > 0 and p ( x ) = 12 δ ( x − 1 )
+ 12 δ ( x + 1 ) .
The diagonal terms of the Hessian of the loss function is given by
[Equation]
There are four local minima in total for this loss function.
These minima and their sharpnesses are
[Equation]
We see that for positive a , the minima at ( 0 , ±√ 1 + a 2 ) are sharper than the other two minima.
We show that SGD will converge to these sharper minima.
For this example, we assume that w 1 and ware bounded because, if initialized sufﬁciently far from the origin, the fourth-order term will cause divergence.
Proposition 4.
Let the loss function be Eq.
(7) , ∣ w 1 ∣ ≤ 1 and ∣ w 2 ∣ ≤ 1 .
For any λ < c for some constant c = O ( 1 ) , there exists some b such that if SGD converges in probability, it will converge to the sharper minimum at ( w 1 ,w 2 )
= ( 0 , ±√ 1 + a 2 ) .
Remark.
In this example, the noise is not full-rank, which is often the case in a deep learning setting for overparametrized networks (Wojtowytsch, 2021a).
For example, when weight decay is used, the
Hessian of the loss function should be full rank, while the rank of the noise covariance should be proportional to the inherent dimension of the data points, which is in general much smaller than the dimension of the Hessian in deep learning (Ansuini et al., 2019).
The proof is technical and given in the appendix Sec.
D.5.
In fact, it has been controversial whether ﬁnding a ﬂat minimum can help generalization.
For example, Dinh et al.
(2017) shows that, for every ﬂat minimum of a ReLU-based net, there exists a minimum that is arbitrarily sharper and generalizes as well.
However, this work does not rule out the possibility that conditioning by using gradient-based optimization, the performance of the sharper minima that gradient descent ﬁnds is worse than the performance of the ﬂatter minima.
If this assumption is valid, then biasing gradient descent towards ﬂatter minima can indeed help, and the stochasticity of SGD has been hypothesized to help in this regard.
However, our construction shows that SGD, on its own, may be incapable of helping SGD converge to a ﬂatter minimum.
At least some other assumptions need to be invoked to show that
SGD may help.
For example, the deﬁnition of the neural networks may endow these models with a special kind of Hessian and noise structure that, when combined with SGD, results in a miraculous generalization behavior.
However, no previous work has pursued this direction in sufﬁcient depth to our knowledge, and future studies in this direction may be fruitful.
Moreover, it is worth commenting that, in the particular example we studied, SGD will converge to the local minimum at ( 0 , ±√ 1 + a 2 ) independent of its sharpness due to the special noise we constructed.
This fact hints at the crucial role of noise in SGD optimization.
In fact, if we allow arbitrary spacedependent noise, it might be the case that SGD may converge to an arbitrary distribution that is independent of the underlying landscape; we brieﬂy discuss this possibility in the form of a conjecture in Sec.
B.1.
5. N ON -C ONVERGENCE OF A DAPTIVE G RADIENTS
[Equation]
Adam (Kingma and Ba, 2014) and its closely related variants such as RMSProp (Tieleman and Hinton, 2012) have been shown to converge to a local maximum even for some simple convex loss landscapes (Reddi et al., 2018a).
The proposed ﬁx, named AMSGrad, takes the maximum of all the previous preconditioners in Adam.
This section shows that AMSGrad may also converge to a local maximum even in simple non-convex settings.
Let x t ∼ p ( x ) and ˆ g t = ∇ ˆ L ( x t ,w t − 1 ) , the Adam algorithm is given by Eq.
(10)-(13), where v 0 = m
= 0 .
We have removed the numerical smoothing constant (cid:15) from the denominator, which causes no problem if w 0 is initialized away from 0 .
Here, β 1 is the momentum hyperparameter.
v t is called the preconditioner, and β 2 is the associated hyperparameter.
The standard value for β 1 is 0 .
9 and β 2 is 0 .
In our theory, we only consider the case when β 1 = 0 .
The experiment section shows that a similar problem exists when β 1 > 0 (with additional interesting ﬁndings).
Intuitively, this is easy to understand because AMSGrad behaves like SGD asymptotically, so we only have to wait for long enough, and the results in previous sections would apply.
The construction below follows this intuition.
Proposition 5.
Let w t ∈
[− 1 , 1 ] , and w 0 ≠ 0 .
For ﬁxed λ < 1 and the loss function in Eq.
(2) there exists a < 0 such that the AMSGrad algorithm converges in probability to a local maximum.
Remark.
The proof is given in Sec.
D.6.
It is interesting to note that the example we construct is similar to the original construction in Reddi et al.
(2018a) that shows that Adam converges to a maximum while AMSGrad succeeds in reaching the global minimum.
In their example (with some rescaling), the gradient is a random variable
[Equation]
such that the expected gradient is negative, and Adam is shown to converge to the direction opposite to the gradient descent (therefore to a maximum).
In our example, the gradient is (roughly)
[Equation]
Therefore, our example can be seen as a minimal generalization of the Reddi et al.
(2018a) example to a non-convex setting.
There is no reason to expect an algorithm that converges in a convex setting to converge in a general non-convex setting, and there is even less reason to expect convergence in a stochastic setting.
In fact, there has not been any convincing evidence that Adam fails badly in a practical setting (Choi et al., 2020).
E XPERIMENTAL D EMONSTRATIONS
We perform experiments to illustrate the examples studied in this work.
Due to space limitations, we illustrate the escape rate and the convergence to sharper minima example in Appendix Sec. A.
6. SGD C ONVERGES TO A L OCAL M AXIMUM
Numerical results are obtained for the setting described in section 5.1.
See Fig.
1-Left.
The loss landscape is deﬁned by L ( w )
=
aw 2 .
In this numerical example, we set λ = 0 .
8 and a = − 1 , and the histogram is plotted with 2000 independent runs.
We see that the distribution converges to the local maximum at w = 0 as the theory predicts.
For better qualitative understanding, we also plot the empirical phase diagram of this setting in Fig.
1-Right.
We perform numerical evaluations for different combinations of a and λ on a a − λ plane, and for each combination, we plot the percentage of w that escapes the central bin (white is 100% and dark blue is 0% ).
For completeness, we also plot the case when a > 0 , i.e., when the critical point at w = 0 is a local minimum.
The orange line shows where the phase transition is expected to happen according to Eq.
The fact that the orange line agrees exactly with the empirical phase transition line conﬁrms our theory.
We also numerically study a closely related fourth-order loss landscape, deﬁned as ˆ L ( w ) =
12 xw 2 + 14 w 4 , where the distribution p ( x ) is the same as before.
The expected loss L ( w ) has two local minima located as w = ±√ a .
We are interested in whether SGD can converge to these two points successfully; note that, without the fourth-order term, this loss is the same as the quadratic loss we studied.
See Fig.
1-Middle.
As before, a = − 1 and λ = 0 . 8 .
We see that the distribution of w also concentrates towards the local maximum after training, and no w is found to converge to the two local minimum even if a signiﬁcant proportion of w is initialized close to these two minima.
A phase diagram analysis is given in the Sec.
A.1.
6. N ON -C ONVERGENCE OF AMSG RAD
In this section, we illustrate the example of the convergence behavior of AMSGrad described in
Sec. 5.4; for reference, we also plot the behavior of Adam and gradient descent for this example.
See
Fig.
The experiment is the average over 2000 runs, each with 5 × 10 4 steps.
The uncertainty is reﬂected by the shaded region (almost invisibly small).
The loss function is the same as discussed in Sec.
5.4 with a = − 0 .
For illustration purposes, we set λ = 0 .
2 and β 2 = 0 .
999 for both Adam and AMSGrad.
When momentum is used, we set β 1 = 0 .
GD is run with a learning rate of 0 . 01 .
Without momentum, both Adam and AMSGrad converge to the local maximum with almost the same speed.
When momentum is added, AMSGrad still converges to the local maximum.
Curiously, Adam is no longer attracted to the local maximum but also remains away from the local minimum.
This suggests that Adam, with momentum, might have a better capability of escaping saddle points than
AMSGrad and might be preferable to AMSGrad in a non-convex setting.
6. A N EURAL N ETWORK E XAMPLE
We have now illustrated several unexpected SGD phenomena in artiﬁcial settings, inspiring the natural question of whether they might be important in practical cases.
Here we construct a neural-net-like optimization problem to show that our zoo of odd effects might indeed arise in tasks of interest.
The neural network function is given by f ( x ) =
w 2 σ ( w 1 x ) , where w 1 ,w 2 ∈ R .
σ is the nonlinearity, and we let σ ( x ) =
x 2 as a minimal example.
This is the simplest kind of nonlinear feedforward network one can construct ( 2 layer with a single hidden neuron).
We also pick a minimal dataset with a single data point: x = 1 with probability 1 and y ∈ {− 1 , 2 } , each with probability 0 .
5 , i.e., the label has a degree of inherent uncertainty, which is often the case in real problems.
We use mean squared error (MSE) as the loss function.
Therefore, the expected loss is L ( w 1 ,w 2 ) = 1 2 ( w 2 w 2 1 + 1 ) 2 + 1 2 ( w 2 w 2 1 − 2 ) 2 .
The global minimum L ∗ = 2 . 25 is degenerate, and is achieved when w 1 = 1 /√ 2 w 2 .
There is also a manifold of saddle points given by {( w 1 ,w 2 )∣ w 1 = 0 ,w 2 ≥ 0 } , all with L = 2 . 5 .
Despite the simplicity of this example, it contains many essential features in a realistic problem, including (1) inherent uncertainty in the data, (2) a nonlinear hidden layer, and (3) a degenerate minimum with zero eigenvalues in the Hessian.
Moreover, the fact that this example contains only two parameters makes it feasible to illustration.
[Figure]
More importantly
[Figure]
most previous works on the convergence or escaping behavior of SGD are inapplicable to this example.
One dominant assumption is the ρ -Hessian Lipschitz property (Jin et al., 2017;
Ge et al., 2015)
[Figure]
which does not hold for any deep neural network in general and for this example in particular.
Another common assumption is the PL condition (Karimi et al., 2020;
Wojtowytsch, 2021a; Vaswani et al., 2019)
[Figure]
which does not apply due to the existence of the saddle point.
One recent assumption is that the loss function is 1 -point convex for all points (Kleinberg et al., 2018)
[Figure]
which is also ruled out because ( 0 , 0 ) is not 1 -point strongly convex.
Another relevant assumption is the correlated negative curvature assumption, which also does not hold for ( 0 , 0 ) .
We show in Sec.
A.5 in detail why these conditions are violated.
Therefore, while this example is a highly limited special example, it is arguably closer to actual deep learning than the popular settings of escaping saddle points with SGD.
See Fig.
3 for the experiment with this example.
Here, w 1 is initialized uniformly in [− 1 , 1 ] ; w 2 is initialized uniformly in [ 0 , 1 ] to be closer to the global minimum and away from the saddle point (standard initialization such as initializing in [− 1 , 1 ] does not change the conclusion).
The left ﬁgure shows the stationary distribution of the model parameter at a small learning rate ( λ = 0 . 001 ).
All the parameters are located in the global-minimum valley as expected.
In contrast, when the learning rate is large ( λ = 0 . 1 ), the central ﬁgure shows that all models ( 1000 independent runs) converge to the saddle point at ( 0 , 0 ) .
The right ﬁgure investigates this change more systematically and shows the change in the average stationary training loss of the models as we increase the learning rate from 0 .
001 to 0 .
We see that for a small learning rate, the training loss is close to that of the global minima ( L = 2 . 25 ), and for a signiﬁcant range of large learning rates, the model invariably converges to a saddle point ( L = 2 . 5 ).
One additional interesting observation is that the model diverges at λ ≈ 0 .
11 , and there is almost no sign of such divergence (such as increased ﬂuctuation) when the learning rate is close to the divergent threshold.
This example shows the relevance of our results to the study of neural networks.
In fact, it has often been noticed that at convergence, many large-scale neural networks in real tasks exhibit negative eigenvalues in the Hessian (Alain et al., 2019; Granziol et al., 2019).
The existence of negative eigenvalues at a late time suggests either the possibility that the algorithm has yet to escape or that it has indeed been attracted to such saddle points; if the latter is true, then our work offers an explanation.
The emergence of negative eigenvalues at the end of training serves as indirect evidence that our result may be relevant for larger neural networks.
D ISCUSSION AND O PEN Q UESTIONS
We have shown that, when not in the restrictive regimes that previous works often assume, SGD can exhibit many unexpected phenomena, such as convergence to local maxima or saddle points.
These behaviors are relevant to understanding the role of SGD in deep learning because the loss landscapes of neural networks are full of saddle points (Wu et al., 2017) which one hopes SGD will avoid.
Our results suggest the following questions for future work: (1) Do the behaviors we realize in simple constructions also occur in realistic deep networks?
(2) Does convergence to saddle points help or hurt generalization?
(3) If it hurts, how can we modify SGD to better avoid saddle points?
We suspect changing learning rates, changing batch size, or injecting noise may help, but a convincing theoretical guarantee in a realistic setting is as yet lacking.
(4) If saddle points do not have worse generalization, our results motivate for understanding why.
This is especially relevant because we showed that using a large learning rate is more likely to converge to saddle points, while previous works have shown that using a large learning rate can improve generalization; combined, this may imply that certain saddle points may have intriguing but unknown regularization effects.
A A DDITIONAL E XPERIMENTS A. P HASE D IAGRAM OF THE FOURTH - ORDER L OSS F UNCTION
The fourth-order loss function is a more realistic loss function than the one we considered in Sec. 5.1.
We also performed one experiment in the main text (See Fig. 1).
Here, we plot its phase diagram.
We see that for this loss landscape also, there is some region such that for all λ , SGD converges to the local maximum.
This loss landscape is very difﬁcult to study in discrete time as it is known to lead to chaotic behavior at large step size (May, 1976).
Therefore, we alternatively try to understand this landscape using continuous approximation.
See Sec.
C.
[Figure]
A. E STIMATING THE E SCAPE R ATE
In the escape rate experiments, the empirical results are obtained from the proposed approximation in
Eq.
Here, we show that it is valid.
See Fig.
5, where we plot the estimated γ as a function of the training step.
We see that the estimated value converges within about 20 steps.
We, therefore, estimate the escape rates at time step 100 in the main text.
[Figure]
A. E SCAPE R ATE E XPERIMENTS
This section illustrates the slow escape problem studied in Sec. 5.2.
See Fig.
6-Right.
γ is calculated by averaging the ﬁrst 50 time steps across 2000 independent runs.
We see that, as our theory predicts, as ∣ a ∣ gets closer to 0 , the optimal escape rate decreases towards 0 .
This implies that there exists a landscape such that SGD is arbitrarily slow at learning, independent of parameter tuning.
One additional observation is that the optimal learning rate is neither too large nor too small, i.e., there seems to be a tradeoff between the escape speed and escape probability (and between the speed and stability).
See also the discussion at the end of Sec. 5.2.
[Figure] [Figure]
Proposition
A. C ONVERGENCE TO THE S HARPER M INIMUM
We use the same 2-dimensional loss landscape deﬁned in Sec.
The experiment is run withindependent simulations and learning rate λ = 0 .
See Fig.
7, where we overlap the underlying landscape with the empirical distribution (the heat map).
We see that, even though the initialization overlaps signiﬁcantly with the local ﬂatter minimum, all points converge to the sharper minimum, as the theorem predicts.
A. F AILURE OF S TANDARD A SSUMPTION FOR THE N EURAL N ETWORK E XAMPLE
In this section, we show that the standard assumptions in the previous literature do not hold for the example studied in Sec.
One common assumption is that the loss function is ρ -Hessian Lifschitz.
Deﬁnition 5.
A loss function L is said to be ρ -Hessian Lifschitz if, for some ρ > 0 ,
[Equation]
This certainly does not hold for the example we considered.
Explicitly, let w 2 = 0 , the Hessian H ( w ) of the loss function is
[Equation]
and
[Equation]
Let w 1 → ∞ violates the ρ -Hessian Lifschitz property.
Another common assumption for the non-convex setting is the PL condition.
Deﬁnition 6.
Let L ∗ be the value of L at the global minimum.
A loss function L is said to satisfy the PL (Lojasiewitz) condition if
[Equation]
(19) for some µ > 0 .
This also does not hold due to the existence of the saddle point.
Let ( w 1 ,w 2 ) =
The gradient is zero, but the right-hand side is greater than 0 .
Recently, the correlated negative curvature assumption has been proposed to study the escaping behavior of SGD.
Deﬁnition 7.
Let v w be the eigenvector of the minimum eigenvalue of the Hessian H ( w ) .
L satisﬁes the correlated negative curvature assumption if, for some γ > 0 ,
[Equation]
where x is the data point.
This also does not hold.
The point ( 0 , 0 ) violates this condition because ∇ L ( x, w ) = 0 for all x at ( w 1 ,w 2 )
= ( 0 , 0 ) .
In fact, it is exactly this point that violates this condition that the SGD converges to in this example.
Recently, the one point strongly convex assumption has been proposed to study the escaping behavior of SGD.
Deﬁnition 8.
A loss function is said to satisfy the c -one point strongly convex condition with respect to w ∗ for all gradient noise (cid:15) , c > 0 , deﬁne v = w − λ ∇ L ( w ) , we have
[Equation] [Equation]
where (cid:15) is the random noise caused by minibatch sampling.
This assumption is equivalent to that the loss landscape is strongly convex after convoluting with the noise.
This condition implies that there is only a single stationary point.
However, this is not the case, consider any point with w 1 = 0 and w 2 ≤ 0 .
These points have zero gradient for w 1 ≤ 0 and so (cid:15)
=
with probability 1 .
Then, these points all have the same loss after the convolution:
[Equation]
22) which either imply that the landscape is not convex or that there is more than 1 stationary point, and so the system is not one point strongly convex.
B A DDITIONAL T HEORETICAL C ONCERNS B. A C ONJECTURE ABOUT THE I MPORTANCE OF THE N OISE
The mathematical constructions in this work may all be related to the following conjecture: if the noise is arbitrary, SGD may converge to any stationary distribution .
In fact, we can prove this hypothesis in one dimension and when continuous approximation is adapted.
In continuous time,
SGD is equivalent to the following stochastic differential equation (SDE):
[Equation]
(23) where η ( t ) is a zero-mean noise with unit variance, and C ( w ) is the w − dependent variance of the noise.
Proposition 6.
Let w ∈ R , the loss function L ( w ) be an arbitrary differential function, and p ( w ) be an arbitrary distribution such that 0 < p ( w ) < ∞ and lim w →∞ wL ′ ( w ) p ( w ) < ∞ .
Then there exists
C ( w ) such that the stationary distribution of SGD is p ( w ) .
We prove this proposition below.
This proposition shows that if the noise is arbitrary (even if it is smooth and zero-mean), the stationary distribution of SGD can be independent of the landscape as long as it is sufﬁciently regular and smooth.
It is possible to extend this proposition to some (mildly) non-smooth distributions.
This result becomes more surprising since we can show that the converse is not true, i.e., for an arbitrary p ( w ) , there exists C ( w ) such that there is no L ( w ) such that p S ( w )
=
p ( w ) .
Proposition 7.
Let w ∈ R , p ( w ) be an arbitrary (sufﬁciently smooth) distribution.
Then there exist
C ( w ) such that such the stationary distribution p S ( w ) ≠ p ( w ) for any loss landscape L ( w ) .
Proof .
Simply let C ( w ) = 0 , and let p ( w ) be a Gaussian distribution with non-zero variance.
Then the SDE either converges to distributions with Dirac measures or does not converge at all for any
L ( w ) .
◻
The above two propositions combined have an important implication: the noise may be more crucial than the landscape itself in inﬂuencing the dynamics and the convergence of SGD.
This relates to the central message of this work.
To study the role of SGD in deep learning, one has to make realistic assumptions about its noise, and the structure of the noise needs to be a function of the model architecture and the actual data distribution.
Before the proof of the ﬁrst proposition, we state a more “constructive" version of the proposition.
The proof follows by solving the Fokker-Planck equation, whose validity has been rigorously studied (Ndao, 2016).
The Fokker-Planck equation and some relevant studies will be presented in Sec.
C.
Proposition 8.
Let w ∈ R , the loss function L ( w ) be an arbitrary differential function, and p ( w ) be an any distribution such that 0 < p ( w ) < ∞ and lim w →∞ wL ′ ( w ) p ( w ) < ∞ .
Then, if
[Equation]
the stationary distribution of SGD is p ( w ) .
Proof .
The stationary distribution p S ( w ) is given by the solution to the Fokker-Planck equation in (37).
Integrating over w , one ﬁnds
[Equation]
(25) where c 0 is the integration constant.
For the notion-wise simplicity, c 0 multiplied or added by a constant is also denoted as c 0 .
By the deﬁnition of a distribution on R , lim ∣ w ∣→∞ p S ( w ) = 0 .
The constant probability ﬂow J 0 going through the system is also 0 .
Applying J 0 = 0 and rearranging the equation, one ﬁnds
[Equation]
The statement of this proposition is therefore equivalent to the statement that we can ﬁnd C ( w ) ≥such that Eq.
(26) is satisﬁed.
By assumption, p S ( w ) > 0 everywhere; thus, the solution C ( w ) exists if
[Equation]
or, equivalently,
[Equation] [Equation]
i.e. the average value of L ′ ( w ) is upper bounded by c 0 .
This condition is satisﬁed, when p S ( w ) itself is ﬁnite everywhere, by tuning p S ( w ) such
that lim w →∞ wL ′ ( w ) p S ( w ) < ∞ .
Thus, the stationary distribution of w under SGD dynamics could be any distribution, as long as the probability distribution p S ( w ) satisﬁes
• p S ( w ) > 0 ; • p S ( w ) < ∞ ; • lim w →∞ wL ′ ( w ) p S ( w )
< ∞ .
This completes the proof.
◻
Remark.
The condition lim w →∞ wL ′ ( w ) p ( w ) < ∞ is equivalent to saying that the tail of p ( w ) has to decay sufﬁciently fast.
Since the loss landscape is usually polynomial, the aforementioned condition is usually satisﬁed in practice.
Depending on L ( w ) , the condition 0 < p ( w ) < ∞ could be relaxed such that 0 < p ( w ) < ∞ almost everywhere.
We also would like to comment that, even if the target distribution is not smooth and does not satisfy the condition in Proposition 8, it is still possible that we can approximate p ( x ) arbitrarily close.
For example, if the target p ( x ) is the (not necessarily smooth) limiting distribution of a sequence of smooth distributions in distance function f
([ p ]) , then we can still approximate p ( x ) arbitrarily well in the speciﬁed distance function f .
This allows us to extend the applicability of the above proposition to common non-smooth ones such as a delta distribution (see the example below) or a uniform distribution.
The following examples follow from direct applications of Proposition 6.
Example 1 .
If L ( w ) = 12 aw 2 and C ( w ) = w 2 , then, by Proposition 8 p S ( w ) =
p ( w ) ∝ w − 2 − 2 Saλ .
When p S ( w ) = p ( w ) ∝ w − 2 − 2 Saλ and L ( w ) = 12 aw 2 , according to (26),
[Equation]
when c 0 = 0 , B ( w ) =
w 2 , which is consistent with the results in C.4.
Notice that, in this example, p S ( w ) diverges at w =
and this distribution is deﬁned only as a limit, where Proposition 8 still holds.
Example 2 .
If L ( w ) = 12 w 2 , and C ( w ) = exp
[( w −
[Equation] [Equation]
dw } , then p ( w ) = 1 √ π exp [−( w − w 0 ) 2 ] .
It is straightforward to check that the expression of C ( w ) is derived using (26).
One can also check that p ( w ) is indeed the solution to the Fokker-Planck equation.
[Equation]
This example shows that, with proper noise, the stationary distribution and the landscape could even have different symmetry, and that it is possible to have a stationary distribution such that the expected gradient is non-vanishing.
[Figure]
The ﬁrst example has been closely studied in the main text and the be directly with numerical simulation
C C ONTINUOUS -T IME A PPROXIMATION WITH F OKKER -P LANCK E QUATION C. F OKKER -P LANCK E QUATION AND I TS S TATIONARY D ISTRIBUTION
In this section, we approach the relevant problems in this paper using continuous-time approximation.
In fact, many results we have obtained in discrete-time can also be derived in the continuous-time approximation.
The additional advantage of continuous-time approximation is that it gives us more ﬂexibility to deal with more complicated loss landscapes.
Recall that the SGD update takes the form (see Sec. 2)
[Equation]
when λ < 1 , the above equation may be approximated by a continuous-time Ornstein-Uhlenbeck process (Mandt et al., 2017)
[Equation]
where λ is the learning rate; we have also introduced S as the batch size.
dW ( t ) is a stochastic process satisfying
[Equation]
By the deﬁnition of the underlying discrete-time process, the term dw ( t ) depends only on w ( t ) and has no dependence on w ( t + dt ) .
Thus, the stochastic integration should be interpreted as Ito, leading to the following Fokker-Planck equation (Risken, 1996).
[Equation] [Equation]
where Tr [⋅] is the trace.
For simplicity, only one-dimensional version of the Fokker-Planck equation is considered, which is
[Equation]
[Equation]
The probability P [ w,t ∣ w ( 0 ) , 0 ] is normalized by deﬁnition, thus, the total probability is conserved.
Expressing the conservation of probability in a microscopic manner:
[Equation]
where J [ w,t ∣ w ( 0 ) , 0 ] is the probability ﬂow.
Comparing with the Fokker-Planck equation, the expression of J [ w,t ∣ w ( 0 ) , 0 ] is obtained:
[Equation]
Assuming that the SGD dynamics is ergodic, there is a unique stationary distribution for w when t → ∞ , and it satisﬁes
[Equation]
The dependence on the initial condition could be dropped due to ergodicity.
Integrating over w on both sides, one ﬁnds
[Equation] [Equation]
where J 0 is the integral constant, which could be interpreted as the stationary ﬂow of probability going from −∞ to ∞ .
When the loss function is convex, P (∞) = 0 , leading to a trivial ﬂow J 0 = 0 .
In this case, the stationary solution takes a simple form:
[Equation]
We will apply this equation to study the relevant problems in this work.
C. A G ENERAL C ASE
In this section, we consider (a slightly more general version of) the fourth-order potential we studied in the main text.
The average loss landscape in this model is
[Equation]
where b > 0 , guarantees that the w is bounded regardless the value of a .
In the limit of b → 0 , this function reduces to saddle point problem we studied in Sec. 5.1.
Besides the SGD noise, additive noise is also present in the dynamics.
The variance of the additive noise has no dependence on w .
When the additive noise is present, the update rule (equation of motion) is
[Equation]
where both η m ( t ) and η a ( t ) are both Ornstein-Uhlenbeck process, denoting multiplicative noise and additive noise respectively.
The w -dependent covariance of w is C ( w ) =
w 2 .
The SGD noise is thus a multiplicative noise 2 η m ( t ) and has a variance w ( t ) 2 S dt , while the additive noise η a ( t ) has a variance of σ 2 dt , where σ is a positive constant denoting the strength of the additive noise.
If the multiplicative noise is seen as a part of the loss landscape, (2) coincides with the 2nd-order term in the loss landscape.
There is no correlation between the additive noise and the SGD noise, i.e. E [ η m ( t ) η a ( t ′ )]
= 0 .
The additive noise can be seen as the artiﬁcially injected noise, a technique sometimes used for aiding the escape or for Bayesian learning purposes (Jin et al., 2017; Welling and
Teh, 2011).
The additive noise vanishes in the limit of σ → 0 , and the model becomes a normal SGD with 4th-order loss function.
It is always possible to deﬁne a noise η ′ ( t ) , whose contribution is equivalent to the contribution of both the SGD and the additive noise, and the equation of motion becomes
[Equation]
[Equation]
This transformed noise η ′ ( t ) thus has 0 mean and a variance of ( w ( t ) 2 S + σ
[Equation]
dt .
Deﬁne η (
[Equation]
) = η ′ ( t )/ √ w ( t )
S + σ 2 , the equation of motion becomes
[Equation]
Comparing with (30), one ﬁnds
[Equation]
The solution of the corresponding Fokker-Planck equation is
[Equation]
Further simpliﬁcation yields
[Equation]
The function P ( w ) is ﬁnite everywhere and decays exponentially to 0 when w → ∞ , regardless of the values of the parameters a,b,σ,S .
This indicates that ∫ ∞−∞ dw ( w 2 + Sσ 2 ) − 1 − Saλ + S
[Figure]
exp [− Sbλ w 2 ] has a well-deﬁned value.
As a consequence, there is no concentration of measure.
Thus, w ( t ) can be fall into any interval with ﬁnite probability after big enough t , indicating that w ( t ) can be arbitrarily far away from 0 .
In practice, this means that w ( t ) escapes from the saddle point at w =
regardless of the value of the parameters.
[Figure]
All the other parameters are identical to those in Fig.1.
This distribution has ﬁnite width.
The effect of additive noise is better understood when b = 0 .
In this case, which corresponds to the case shown in Fig.1, w neither concentrates at 0 , nor escapes to ±∞ .
Instead, it stays near 0 without converging to it.
The stationary distribution of w when b
= 0 , obtained by numerical simulation, is shown in Fig.9.
The other settings are identical to the ones in Fig.1.
The stationary distribution exists, but w does not concentrate at w = 0 in the stationary distribution.
First of all, this ﬁgure shows that the additive noise helps SGD to escape saddle point.
If the designed landscape is a part of a realistic landscape, having a broader distribution means having more chance of being attracted by another minimum.
However, this ﬁgure also shows that w would stay in the neighborhood of w = 0 until the noise is big enough.
Thus, having additive noise and long enough training time does not guarantee the efﬁciency of escape.
In the language of Bayesian inference, ln [ P ( w )] is the likelihood of w .
ˆ w , the most probable value of w , is deﬁned by the relation
[Equation]
which reads
[Equation]
The likelihood maximizer is
[Equation]
When a > − λS , the maximum likelihood parameter is w = 0 .
When a < − λS , w = ±√− λSb − ab equally likely.
This resembles the phase transition in Landau model and − λS is referred to as the critical a .
For the energy landscape, when a > 0 there is only one global minimum at w = 0 , and when a < 0 there are two global minima at w = ±√− ab .
Comparing with the maximum likelihood solutions, we see that the likelihood maximizer given by SGD is in fact a biased estimator of the underlying minima.
There is a bias term introduced by the SGD noise in both the critical a and the value of the most probable w .
This term vanishes when S → ∞ , i.e. when the SGD noise vanishes.
This indicates that, when the noise is state-dependent, there is no reason to expect SGD to be an unbiased or consistent estimator of the minimizer, as some works assume.
C. 4- TH O RDER P OTENTIAL WITH M ULTIPLICATIVE N OISE
In this subsection, we study a 4-th order loss landscape in SGD dynamics without additive noise.
The average loss function with 4th-order term is
[Equation]
47) the SGD update rule (i.e., the equation of motion) in this case is
[Equation]
48) where the deﬁnition of η is identical to the previous example.
The solution of corresponding stationary
Fokker-Planck equation is
[Equation]
When
[Equation]
P ( w ) diverges at w = 0 .
The normalization factor, i.e. integral ∫ +∞−∞ dww − 2 − 2 Saλ exp [− Sbλ w 2 ] , also diverges due to the divergence at w = 0 .
This solution could be seen as a limit of (44) when the additive noise vanishes.
As the strength of the additive noise σ → 0 , the normalization factor keep growing while the distribution remains normalized.
For w ≠ 0 , the function w − 2 − 2 Saλ exp [− Sbλ w 2 ] is always ﬁnite.
Thus, in the limit σ → 0 , P ( w ) approaches 0 everywhere except for at w = 0 .
It is straight forward to check that the normalization factor diverges slower
that P ( 0 ) as σ → 0 .
As a consequence, when σ = 0 , P ( w ) diverges at w = 0 and P ( w ) becomes a Dirac delta function.
This corroborates with the result in Sec.
5.3 that a small a leads to a concentration of measure towards the local maximum.
Compared with (44), one ﬁnds that the additive noise prevents the concentration of measure.
Thus, in practice, the additive noise helps SGD escape the local minimum or saddle point.
However, the existence of additive noise does not change the critical a and the position of the peaks.
C. Q UADRATIC P OTENTIAL
At last, we consider the continuous approximation of the model treated in the Sec. 5.1.
Let the loss function be
[Equation]
and we have
[Equation] [Equation]
The SGD update rule (i.e., the equation of motion) in 1d is
[Equation]
With the continuous approximation, the equation of motion becomes
[Equation] [Equation]
where η ( t ) ∼ N( 0 , √ dt ) .
By comparing with the 1d Fokker-Planck equation,
[Equation]
[Equation]
The stationary distribution of w is
[Equation]
The only difference between the solution in this case and the one in the 4th-order-loss case is the exponential factor, showing that the 4th-order potential does nothing but keeping w ( t ) bounded.
The function w − 2 − 2 Saλ diverges at w = 0 or w = ±∞
depending on the value of a .
However, the stationary distribution, being a limit of (44), is well deﬁned.
In the case that it diverges at w = 0 ,
P ( w ) becomes a delta function.
On the contrary, there are two peaks inﬁnitely far away from 0 when w − 2 − 2 Saλ diverges at w = ±∞ .
Being able to escape requires that the probability measure of w does not concentrate at 0 , corresponding to
[Equation]
This straight line agrees with the boundary of the lower branch of the phase diagram in Fig.
6 when a → 0 .
This condition of escaping is worth interpretation.
Note that a is the local curvature and reﬂects the strength of the gradient signal.
In contrast, the strength of the SGD noise is proportional to λ / S .
Also, this example raises an interesting question regarding the mechanism of SGD that causes a convergence to the local maximum.
From the perspective of the types of convergence, the SGD mechanism behind proposition 1 may be different from that of proposition 2.
However, from the perspective of continuous-time analysis, there is really the same mechanism that is governing the SGD dynamics behind proposition 1 and 2 (namely, the fact that the noise has dominated the gradient).
D D ELAYED P ROOFS D. T WO F REQUENTLY U SED L EMMAS
We ﬁrst prove the following lemma regarding the limiting distribution of ln ∣ w t ∣ .
Lemma 1.
Let the loss function be ˆ L ( w ) = 12 xw 2 , x ∼ p ( x ) such that Var [ x ] = σ 2 and E x
[ x ] = a < 0 and that p ( x ) is continuous in a δ -neighborhood of x = 1 , and w 0 ≠ 0 .
Then, for w t generated by
SGD after t time steps,
[Equation] [Equation]
where µ = E x [ ln ∣ 1 − λx ∣] and s 2 =
Var [ ln ∣ 1 − λx ∣] , where λ > 0 is the learning rate.
Proof .
After t steps of SGD, we have
[Equation]
This leads to
[Equation]
Now, since p ( 1 − x ) is continuous in the neighborhood of 1 by assumption, we can apply Proposition(appendix) to show that the ﬁrst and second moment of ln ∣ 1 − λx ∣ exists, whereby we can apply the central limit theorem to obtain
[Equation]
where s 2 =
Var [ ln ∣ 1 − λx ∣] .
This proves the lemma. ◻
Proposition 9.
E x ∼ p ( x ) [ ln 2 ∣ x ∣] is ﬁnite if the second moment of p ( x ) exists and if p ( x ) is continuous in the neighborhood of x = 0 .
Proof .
Since p ( x ) is continuous in the neighborhood of x = 0 , then there exists δ > 0 such that for all ∣ (cid:15) ∣ < δ , ∣ p ( δ ) − p ( x )∣ ≤ c for some c > 0 .
This means that we can divide the integral over p ( x ) into three regions:
[Equation]
The second term can be bounded as
[Equation]
The ﬁrst and the third term can also be bounded.
Since ln ∣ x ∣ is a convex function in the regions x ≥ δ and x ≤ δ respectively.
We can ﬁnd linear functions ax + b of x such that ax + b ≥ ln x for x ≥ δ .
This leads to
[Equation]
where we have used the notation E [ x 2 ] = µ 2 and E [ x 1 ] = µ , which by assumption is ﬁnite.
The case for x ≤ − δ is completely symmetric and can also be bounded by a 2 µ 2 + 2 abµ 1 + b 2 .
Therefore, we have shown that
[Equation]
This proves the proposition.
◻
Remark.
Note that, the continuity in the neighborhood of 0 is not a necessary condition.
For example, one can also prove that E x [ ln 2 ∣ x ∣] is ﬁnite if p ( x ) is bounded and its second moment exists.
D. P ROOF OF P ROPOSITION
Proof .
The case when w 0 = 0 is trivially true, we therefore consider the case w 0 ≠ 0 .
By Lemma 1, we have
[Equation]
where we have deﬁned z t = 1 √ t ( ln ∣ w t w 0 ∣ − tµ ) , µ = E [ ln ∣ 1 − λx ∣] and s
=
Var [ ln ∣ 1 − λx ∣] .
This means that
[Equation]
By deﬁnition, ∣ w t ∣ =
∣ w
∣ e √ tz + tµ .
Therefore, we have
[Equation]
(76) for all (cid:15) > 0 .
In other words, since µ and σ are t -independent, in the inﬁnite t limit, the sign of µ becomes crucial.
When µ > 0 , the limiting distribution diverges to inﬁnity; when µ < 0 , ∣ w t ∣ converges to 0 in probability.
This ﬁnishes the proof.
◻
D. P ROOF OF C OROLLARY
Proof .
By the deﬁnition of p ( x ) ,
[Equation]
while s 2 is indeed constant in time.
Therefore, the asymptotic distribution of w t is solely dependent on the sign of µ .
The above equation implies that µ ≥ 0 when
[Equation]
When λ ≤ 1 , the above equation is solved by
[Equation] [Equation]
When λ > 1 , the above equation solves to
[Equation]
only when the learning rate satisﬁes the above two conditions can it escape from the local maximum.
Conversely, SGD cannot escape the local minimum when
[Equation]
We are done.
◻
D. P ROOF OF P ROPOSITION
Proof .
By Lemma 1, we have that
[Equation]
where the second equality follows from the assumption that λ < 1 .
We differentiate with respect to λ to ﬁnd the critical escape rate:
[Equation]
Since µ is convex in λ , it follows that this critical escape rate is the maximum escape rate.
Plugging this into µ , we obtain that
[Equation]
i.e., the optimal escape rate can be made smaller than any (cid:15) if we set
[Equation] [Equation]
This ﬁnishes the proof.
◻
D. P ROOF OF P ROPOSITION
Proof .
It sufﬁces to show that w 1 converges to 0 with probability 1 because, if this is the case, the only possible local minimum to converge to is ( w 1 ,w 2 ) =
( 0 , ±√ 1 + a 2 ) .
The SGD dynamics is
[Equation]
We focus on the dynamics of w 1 .
By the deﬁnition of the noise x , we have that
[Equation]
probability
[Equation]
with probability
[Equation]
Equivalently,
[Equation]
Since 0 ≤ 4 λw 21 ,t + 12 λw 22 ,t ≤ 16 λ , we can deﬁne a new random variable r t :
[Equation]
By construction, r t ≥ ∣ w 1 ,t + 1 / w 1 ,t ∣ for all values of x t .
This implies that ∣ w 1 ,t / w 1 , 0 ∣ ≤ ∏ ti = 1 r t , and so,
[Equation]
i.e., if ∏ ti = 1 r t converges to 0 with probability 1 , w 1 ,t must also converge to zero with probability 1 .
We let b = 1 λ − 6 (note that this is the value of b such that r t is minimized for both cases), and we obtain that
[Equation]
The rest of the proof follows from applying the central limit theorem to 1 √ t ∑ ti =
ln r t as in Lemma 1.
The result is that ∏ t i = 1 r i converges to 0 with probabiliy 1 if
[Equation]
The above inequality solves to
[Equation]
which is of order O ( 1 ) as stated in the theorem statement.
This proves the proposition.
◻
D. P ROOF OF P ROPOSITION
Proof .
First we note that, since ∣ w t ∣
[Equation] [Equation]
By the deﬁnition of the AMSGrad algorithm, we have that
[Equation]
Since v 0 = 0 and ∣ ˆ g t ∣ =
∣ x t w t ∣ ≤ 1 + a , we have that
[Equation]
for all t , and so
[Equation]
Meanwhile, since ˆ v t is a monotonically increasing series and is upper bounded by 1 + a , it must converge to a constant 0 < c ≤ 1 + a .
Now, as before, we want to upper bound the random variable
[Equation]
where µ = E [ ln ∣ w t + 1 w t ∣] .
Since ˆ v t converges to c , there must exists a positive integer N ( (cid:15) ) such that for any (cid:15) > 0 , ˆ v N ≤ c − (cid:15) .
This allows us to divide the sum to two terms:
[Equation]
where we introduced 0 ≤ k t ≤ (cid:15) .
But the ﬁrst term is of order O ( 1 /√ t ) and converges to 0 , and so for sufﬁciently large t , the ﬁrst term is also smaller than arbitrary (cid:15) .
This means that
[Equation]
We can now consider the random variable
[Equation]
We now deﬁne a new random variable
[Equation]
where we deﬁned the constant m ∶= max (∣ 1 − λ /√ c ∣ , ∣ 1 − λ /√ c − (cid:15) ∣) .
One can easily check that r t ≥ ∣ 1 − λ √ c − k t x t ∣ .
This means that
[Equation]
i.e., if r t converges to 0 in probability, ∣ w t + 1 w 0 ∣ must also converge to 0 in probability.
Proceeding as in the proof of Proposition 2.
One ﬁnds that the condition for r t to converge in probability to 0 is
[Equation]
which is equivalent to
[Equation]
Since (cid:15) is arbitrary, we let (cid:15) → 0 and obtain
[Equation]
Denote λ /√ c as λ ′ , this condition solves to
[Equation]
Setting a such that the above condition is met, AMSGrad will converge to 0 in probability.
For λ < 1 ,
This completes the proof.
◻