On Learning Rates and Schrödinger Operators

Introduction

Gradient-based optimization has been the workhorse algorithm powering recent developments in statistical machine learning. Many of these developments involve solving nonconvex optimization problems, which raises new challenges for theoreticians, given that classical theory has often been restricted to the convex setting.

A particular focus in machine learning is the class of gradient-based methods referred to as stochastic gradient descent (SGD), given its desirable runtime properties, and its desirable statistical performance in a wide range of nonconvex problems. Consider the minimization of a (nonconvex) function $f$ defined in terms of an expectation:

where the expectation is over the randomness embodied in $\zeta$. A simple example of this is empirical risk minimization, where the loss function,

is averaged over $n$ data points, where the datapoint-specific losses, $f_{i}(x)$, are indexed by $i$ and where $x$ denotes a parameter. When $n$ is large, it is computationally prohibitive to compute the full gradient of the objective function, and SGD provides a compelling alternative. SGD is a gradient-based update based on a (noisy) gradient evaluated from a single data point or a mini-batch:

where the set $\mathcal{B}$ of size $B$ is sampled uniformly from the $n$ data points and therefore the noise term $\xi$ has mean zero. Starting from an initial point $x_{0}$, SGD updates the iterates according to

where $\xi_{k}$ denotes the noise term at the $k$th iteration. Note that the step size $s>0$, also known as the learning rate, can either be constant or vary with the iteration [Bot10].

The learning rate plays an essential role in determining the performance of SGD and many of the practical variants of SGD [Ben12].Note that the mini-batch size as another parameter can be, to some extent, incorporated into the learning rate. See discussion later in this section. The overall effect of the learning rate can be complex. In convex optimization problems, theoretical analysis can explain many aspects of this complexity, but in the nonconvex setting the effect of the learning rate is yet more complex and theory is lacking [Zei12, KB14]. As a numerical illustration of this complexity, Figure 1 plots the error of SGD with a piecewise constant learning rate in the training of a neural network on the CIFAR-10 dataset. With a constant learning rate, SGD quickly reaches a plateau in terms of training error, and whenever the learning rate decreases, the plateau decreases as well, thereby yielding better optimization performance. This illustration exemplifies the idea of learning rate decay, a technique that is used in training deep neural networks (see, e.g., [HZRS16, BCN18, SS19]). Despite its popularity and the empirical evidence of its success, however, the literature stops short of providing a general and quantitative approach to understanding how the learning rate impacts the performance of SGD and its variants in the nonconvex setting [YLWJ19, LWM19]. Accordingly, strategies for setting learning rate decay schedules are generally adhoc and empirical.

In the current paper we provide theoretical insight into the dependence of SGD on the learning rate in nonconvex optimization. Our approach builds on a recent line of work in which optimization algorithms are studied via the analysis of their behavior in continuous-time limits [SBC16, Jor18, SDJS18]. Specifically, in the case of SGD, we study stochastic differential equations (SDEs) as surrogates for discrete stochastic optimization methods (see, e.g., [KY03, LTE17, KB17, COO+18, DJ19]). The construction is roughly as follows. Taking a small but nonzero learning rate $s$, let $t_{k}=ks$ denote a time step and define $x_{k}=X_{s}(t_{k})$ for some sufficiently smooth curve $X_{s}(t)$. Applying a Taylor expansion in powers of $s$, we obtain:

Plugging the last two displays into (1.1), we get

Retaining both $O(1)$ and $O(\sqrt{s})$ terms but ignoring smaller terms, we obtain a learning-rate-dependent stochastic differential equation (lr-dependent SDE) that approximates the discrete-time SGD algorithm:

Intuitively, a larger learning rate $s$ gives rise to more stochasticity in the lr-dependent SDE (1.2), and vice versa. Accordingly, the learning rate must have a substantial impact on the dynamics of SGD in its continuous-time formulation. In stark contrast, this parameter plays a fundamentally different role in gradient descent (GD) and stochastic gradient Langevin dynamics (SGLD) when one considers their limiting differential equations. In particular, consider GD:

which can be modeled by the following ordinary differential equation (ODE):

and the SGLD algorithm, which adds Gaussian noise $\xi_{k}$ to the GD iterates:

These differential equations are derived in the same way as (1.2), namely by the Taylor expansion and retaining $O(1)$ and $O(\sqrt{s})$ terms.The coefficients of the $O(\sqrt{s})$ terms turn out to be zero in both differential equations. See more discussion in Appendix A.1 and particularly Figure 12 therein. While the SDE for modeling SGD sets the square root of the learning rate to be its diffusion coefficient, both the GD and SGLD counterparts are completely free of this parameter. This distinction between SGD and the other two methods is reflected in their different numerical performance as revealed in Figure 2. The right plot of this figure shows that the behaviors of both GD and SGLD in the time $t=ks$ scale are almost invariant in terms of optimization error with respect to the learning rate. In striking contrast, the stationary optimization error of SGD decreases significantly as the learning rate decays. As a consequence of this distinction, GD and SGLD do not exhibit the phenomenon that is shown in Figure 1.

The discussion thus far suggests that one may examine the effect of the learning rate in SGD using the lr-dependent SDE (1.2). In particular, this SDE distinguishes SGD from GD and SGLD. Accordingly, in the current paper we study the lr-dependent SDE, and make the following contributions.

Linear convergence to stationarity. We show that, for a large class of (nonconvex) objectives, the continuous-time formulation of SGD converges to its stationary distribution at a linear rate.Roughly speaking, stationarity refers to the distribution of $X_{s}(t)$ in the limit $t\rightarrow\infty$. See a more precise definition in Section 3. In particular, we prove that the solution $X_{s}(t)$ to the lr-dependent SDE obeys

This convergence result sheds light on why SGD performs so well in many practical nonconvex problems. In particular, note that while GD can be trapped in a saddle point or a local minimum, SGD can efficiently escape saddle points, provided that the linear rate $\lambda_{s}$ is not too small (this is the case if $s$ is sufficiently large; see the second contribution). This superiority of SGD in the nonconvex setting must be attributed to the noise in the gradient and this implication is consistent with earlier work showing that stochasticity in gradients significantly accelerates the escape of saddle points for gradient-based methods [JGN+17, LSJR16].

Distinctions between convexity and nonconvexity. The first contribution stops short of saying anything about how $\lambda_{s}$ depends on the learning rate $s$ and the geometry of the objective $f$. Such an analysis is fundamental to an explanation of the differing effects of the learning rate in deep learning (nonconvex optimization) and convex optimization. In the current paper we show that if the objective $f$ is a nonconvex function and satisfies certain regularity conditions, we have:We write $a_{m}\asymp b_{m}$ if there exist positive constants $c$ and $c^{\prime}$ such that $cb_{m}\leq a_{m}\leq c^{\prime}b_{m}$ for all $m$.

for a certain value $H_{f}>0$ that only depends on $f$. This expression for $\lambda_{s}$ enables a concrete interpretation of the effect of learning rate in Figure 1. In brief, in the nonconvex setting, $\lambda_{s}$ decreases to zero quickly as the learning rate $s$ tends to zero. As a consequence, with a large learning rate $s$ at the beginning, SGD converges rapidly to stationarity and the rate becomes smaller as the learning rate decreases.

For comparison, $\lambda_{s}$ is equal to $\mu$ if $f$ is $\mu$-strongly convex for $\mu>0$, regardless of the learning rate $s$. As such, the convergence behaviors of SGD are necessarily different between convex and nonconvex objectives. To appreciate this implication, we refer to Figure 3. Note that all four plots show that a larger learning rate gives rise to a larger stationary risk, as predicted by the monotonically increasing nature of $\epsilon$ with respect to $s$ in (1.3). The most salient part of this figure is, however, shown in the right panel. Specifically, the right panel, which uses time $t$ as the $x$-axis, shows that in the (strongly) convex setting the linear rate of the convergence is roughly the same between the two choices of learning rate, which is consistent with the result that $\lambda_{s}$ is constant in the case of a strongly convex objective. In the nonconvex case (bottom right), however, the rate of convergence is more rapid with the larger learning rate $s=0.1$, which is implied by the fact that $\lambda_{0.1}>\lambda_{0.05}$. In stark contrast, the two plots in the left panel, which use the number $k$ of iterations for the $x$-axis, are observed to have a larger rate of linear convergence with a larger learning rate. This is because in the $k$ scale the rate $\lambda_{s}s$ of linear convergence always increases as $s$ increases no matter if the objective is convex or nonconvex.

The mathematical tools that we bring to bear in analyzing the lr-dependent SDE (1.2) are as follows. We establish the linear convergence via a Poincaré-type inequality that is due to Villani [Vil09]. The asymptotic expression for the rate $\lambda_{s}$ is proved by making use of the spectral theory of the Schrödinger operator or, more concretely, the Witten-Laplacian associated with the Fokker–Planck–Smoluchowski equation that governs the lr-dependent SDE. We believe that these tools will prove to be useful in theoretical analyses of other stochastic approximation methods.

2 Related work

Recent years have witnessed a surge of research devoted to explanations of the effectiveness of deep neural networks, with a particular focus on understanding how the learning rate affects the behavior of stochastic optimization. In [SKYL17, KMN+16], the authors uncovered various tradeoffs linking the learning rate and the mini-batch size. Moreover, [JKA+17, JKB+18] related the learning rate to the generalization performance of neural networks in the early phase of training. This connection has been further strengthened by the demonstration that learning rate decay encourages SGD to learn features of increasing complexity [LWM19, YLWJ19]. From a topological perspective, [DDC19] establish connections between the learning rate and the sharpness of local minima. Empirically, deep learning models work well with non-decaying schedules such as cyclical learning rates [LH16, Smi17] (see also the review [Sun19]), with recent theoretical justification [LA19].

In a different direction, there has been a flurry of activity in using dynamical systems to analyze discrete optimization methods. For example, [SBC16, WWJ16, SDJS18] derived ODEs for modeling Nesterov’s accelerated gradient methods and used the ODEs to understand the acceleration phenomenon (see the review [Jor18]). In the stochastic setting, this approach has been recently pursued by various authors [COO+18, CS18, MHB16, LSJR16, CH19, LTE17] to establish various properties of stochastic optimization. As a notable advantage, the continuous-time perspective allows us to work without assumptions on the boundedness of the domain and gradients, as opposed to older analyses of SGD (see, for example, [HRB08]).

Our work is motivated in part by the recent progress on Langevin dynamics, in particular in nonconvex settings [Vil09, Pav14, HKN04, BGK05]. In relating to Langevin dynamics, $s$ in the lr-dependent SDE can be thought of as the temperature parameter and, under certain conditions, this SDE has a stationary distribution given by the Gibbs measure, which is proportional to $\exp(-2f/s)$. Of particular relevance to the present paper from this perspective is a line of work that has considered the optimization properties of SGLD and analyzed its convergence rates [Hwa80, RRT17, ZLC17]. Compared to these results, however, the present paper is distinct in that our analysis provides a more concise and sharp delineation of the convergence rate based on geometric properties of the objective function.

3 Organization

The remainder of the paper is structured as follows. In Section 2 we introduce basic assumptions and techniques employed throughout this paper. Next, Section 3 develops our main theorems. In Section 4, we use the results of Section 3 to offer insights into the benefit of taking a larger initial learning rate followed by a sequence of decreasing learning rates in training neural networks. Section 5 formally proves the linear convergence (1.3) and Section 6 further specifies the rate of convergence (1.4). Technical details of the proofs are deferred to the appendices. We conclude the paper in Section 7 with a few directions for future research.

Preliminaries

A function $f$ is said to be confining if it is infinitely differentiable and satisfies $\lim_{\|x\|\rightarrow+\infty}f(x)=+\infty$ and $\exp(-2f/s)$ is integrable for all $s>0$:

Next, we need to show that the lr-dependent SDE (1.2) with an arbitrary learning rate $s>0$ admits a unique global solution under mild conditions on the objective $f$. We will show in Section 3.3 that the solution to this SDE approximates the SGD iterates well. The formal description is shown rigorously in Proposition 3.5. Recall that the lr-dependent SDE (1.2) is

with the boundary condition $\rho_{s}(0,\cdot)=\rho$. Here, $\Delta\equiv\nabla\cdot\nabla$ is the Laplacian. For completeness, in Appendix A.2 we derive this Fokker–Planck–Smoluchowski equation from the lr-dependent SDE (1.2) by Itô’s formula. If the objective $f$ satisfies the confining condition, then this equation admits a unique invariant Gibbs distribution that takes the form

For any confining function $f$ and any initial probability density $\rho\in L^{2}(\mu_{s}^{-1})$, the lr-dependent SDE (1.2) admits a weak solution whose probability density in $C^{1}\left([0,+\infty),L^{2}(\mu_{s}^{-1})\right)$ is the unique solution to the Fokker–Planck–Smoluchowski equation (2.1).

The proof of 2.2 is shown in Appendix A.4. For more information, 5.2 in Section 5 shows that the probability density $\rho_{s}(t,\cdot)$ converges to the Gibbs distribution as $t\rightarrow\infty$.

Finally, we need a condition that is due to Villani for the development of our main results in the next section.

A confining function $f$ is said to satisfy the Villani condition if $\|\nabla f(x)\|^{2}/s-\Delta f(x)\rightarrow+\infty$ as $\|x\|\rightarrow+\infty$ for all $s>0$.

This condition amounts to saying that the gradient has a sufficiently large squared norm compared with the Laplacian of the function. Strictly speaking, some loss functions used for training neural networks might not satisfy this condition. However, the Villani condition does not look as stringent as it appears since the SGD iterates in the training process are bounded and this condition is essentially concerned with the function at infinity.

Main Results

In this section, we state our main results. In brief, in Section 3.1 we show linear convergence to stationarity for SGD in its continuous formulation, the lr-dependent SDE. In Section 3.2, we derive a quantitative expression of the rate of linear convergence and study the difference in the behavior of SGD in the convex and nonconvex settings. This distinction is further elaborated in Section 3.3 by carrying over the continuous-time convergence guarantees to the discrete case. Finally, Section 3.4 offers an exposition of the theoretical results in the univariate case. Proofs of the results presented in this section are deferred to Section 5 and Section 6.

Let $f$ satisfy both the confining condition and the Villani condition. Then there exists $\lambda_{s}>0$ for any learning rate $s>0$ such that the expected excess risk satisfies

for all $t\geq 0$. Here $\epsilon(s)=\epsilon(s;f)\geq 0$ is a strictly increasing function of $s$ depending only on the objective function $f$, and $D(s)=D(s;f,\rho)\geq 0$ depends only on $s,f$, and the initial distribution $\rho$.

Briefly, the proof of this theorem is based on the following decomposition of the excess risk:

Under the assumptions of 1, there exists $\lambda_{s}>0$ for any learning rate $s$ such that

for all $t\geq 0$, where the constant $C(s)>0$ depends only on $s$ and $f$, and where

measures the gap between the initialization and the stationary distribution.

Loosely speaking, it takes $O(1/\lambda_{s})$ time to converge to stationarity. In relating to 1, $D(s)$ can be set to $C(s)\left\|\rho-\mu_{s}\right\|_{\mu_{s}^{-1}}$. Notably, the proof of 3.1 shall reveal that $C(s)$ increases as $s$ increases.

Under the assumptions of 1, the excess risk at stationarity, $\epsilon(s)$, is a strictly increasing function of $s$. Moreover, for any $S>0$, there exists a constant $A$ that depends only on $S$ and $f$ and satisfies

The two propositions are proved in Section 5. The proof of 1 is a direct consequence of 3.1 and 3.2. More precisely, the two propositions taken together give

Taken together, these results offer insights into the phenomena observed in Figure 1. In particular, 3.1 states that, from the continuous-time perspective, the risk of SGD with a constant learning rate applied to a (nonconvex) objective function converges to stationarity at a linear rate. Moreover, 3.2 demonstrates that the excess risk at stationarity decreases as the learning rate $s$ tends to zero. This is in agreement with the numerical experiments illustrated in Figures 1, 2, and 3. For comparison, this property is not observed in GD and SGLD.

The following result gives the iteration complexity of SGD in its continuous-time formulation.

Under the assumptions of 3.2, for any $\epsilon>0$, if the learning rate $s\leq\min\{\epsilon/(2A),S\}$ and $t\geq\frac{1}{\lambda_{s}}\log\frac{2C(s)\left\|\rho-\mu_{s}\right\|_{\mu_{s}^{-1}}}{\epsilon}$, then

2 The rate of linear convergence

We now turn to the key issue of understanding how the linear rate $\lambda_{s}$ depends on the learning rate. In this subsection, we show that for certain objective functions, $\lambda_{s}$ admits a simple expression that allows us to interpret how the convergence rate depends on the learning rate.

We begin by considering a strongly convex function. Recall the definition of strong convexity: for $\mu>0$, a function $f$ is $\mu$-strongly convex if

In addition to the assumptions of 1, assume that the objective $f$ is a $\mu$-strongly convex function. Then, $\lambda_{s}$ in (3.1) satisfies $\lambda_{s}=\mu$.

We turn to the more challenging setting where $f$ is nonconvex. Let us refer to the objective $f$ as a Morse function if its Hessian has full rank at any critical point $x$ (that is, $\nabla f(x)=0$).See Section 6.2 for a discussion of Morse functions. Note that (infinitely differentiable) strongly convex functions are Morse functions.

In addition to the assumptions of Theorem 1, assume that the objective $f$ is a Morse function and has at least two local minima.We call $x$ a local minimum of $f$ if $\nabla f(x)=0$ and the Hessian $\nabla^{2}f(x)$ is positive definite. By convention, in this paper a global minimum is also considered a local minimum. Then the constant $\lambda_{s}$ in (3.1) satisfies

for $0<s\leq s_{0}$, where $s_{0}>0,\alpha>0$, and $H_{f}>0$ are constants that all depend only on $f$.

The proof of this result relies on tools in the spectral theory of Schrödinger operators and is deferred to Section 6. From now on, we call $\lambda_{s}$ in (3.1) the exponential decay constant. To obviate any confusion, $o(s)$ in 2 stands for a quantity that tends to zero as $s\rightarrow 0$, and the precise expression for $H_{f}$ shall be given in Section 6, with a simple example provided in Section 3.4. To leverage 2 for understanding the phenomena discussed in Section 1, however, it suffices to recognize the fact that $H_{f}$ is completely determined by $f$. Moreover, we remark that while 1 shows that $\lambda_{s}$ exists for any learning rate, the present theorem assumes a bounded learning rate.

The key implication of this result is that the rate of convergence is highly contingent upon the learning rate $s$: the exponential decay constant increases as the learning rate $s$ increases. Accordingly, the linear convergence to stationarity established in Section 3.1 is faster if $s$ is larger, and, by recognizing the exponential dependence of $\lambda_{s}$ on $s$, the convergence would be very slow if the learning rate $s$ is very small. For example, if $H_{f}=0.05$, setting $s=0.1$ and $s=0.001$ gives

Moreover,as we will see clearly in Section 6, $\lambda_{s}$ is completely determined by the geometry of $f$. In particular, it does not depend on the probability distribution of the initial point or the dimension $d$ given that the constant $H_{f}$ has no direct dependence on the dimension $d$. For comparison, the linear rate in the nonconvex case is shown by 2 to depend on the learning rate $s$, while the linear rate of convergence stays constant regardless of $s$ if the objective is strongly convex. This fundamental distinction between the convex and nonconvex settings enables an interpretation of the observation brought up in Figure 1, in particular the right panel of Figure 3. More precisely, with time $t$ being the $x$-axis, SGD with a larger learning rate leads to a faster convergence rate in the nonconvex setting, while for the (strongly) convex setting the convergence rate is independent of the learning rate. For further in-depth discussion of the implications of 2 (see Section 4).

3 Discretization

In this subsection, we carry over the results developed from the continuous perspective to the discrete regime. In addition to assuming that the objective function $f$ satisfies the Villani condition, satisfies the confining condition, and is a Morse function, we also now assume $f$ to be $L$-smooth; that is, $f$ has $L$-Lipschitz continuous gradients in the sense that $\left\|\nabla f(x)-\nabla f(y)\right\|\leq L\|x-y\|$ for all $x,y$. Moreover, we restrict the learning rate $s$ to be no larger than $1/L$. The following proposition is the key theoretical tool that allows translation to the discrete regime.

We note that there exists a sharp bound on $B(T)$ in [BT96]. For completeness, we also remark that the convergence can be strengthened to the strong sense:

This result has appeared in [Mil75, Tal82, PT85, Tal84, KP92] and we provide a self-contained proof in Appendix B.3.

We now state the main result of this subsection.

In addition to the assumptions of 1, assume that $f$ is $L$-smooth. Then, the following two conclusions hold:

For any $T>0$, the iterates of SGD with learning rate $0<s\leq 1/L$ satisfy

for all $k\leq T/s$, where $\lambda_{s}$ is the exponential decay constant in (3.1), $A$ as in 3.2 depends only on $1/L$ and $f$, $C=C_{1/L}$ is as in 3.1, and $B(T)$ depends only on the time horizon $T$ and the Lipschitz constant $L$.

If $f$ is a Morse function with at least two local minima, with $\lambda_{s}$ appearing in (3.4) being given by (3.3), and if $f$ is $\mu$-strongly convex then $\lambda_{s}=\mu$.

3 follows as a direct consequence of 1 and 3.5. Note that the second part of 3 is simply a restatement of 2 and 3.4. As earlier in the continuous-time formulation, we also mention that the dimension parameter $d$ is not an essential parameter for characterizing the rate of linear convergence. In relating to Figure 3, note that its left panel with $k$ being the $x$-axis shows a faster linear convergence of SGD when using a larger learning rate, regardless of convexity or nonconvexity of the objective. This is because the linear rate $s\lambda_{s}$ in (3.4) is always an increasing function of $s$ even for the strongly convex case, where $\lambda_{s}$ itself is constant.

4 A one-dimensional example

In this section we provide some intuition for the theoretical results presented in the preceding subsections. Our priority is to provide intuition rather than rigor. Consider the simple example of $f$ presented in Figure 4, which has a global minimum $x^{\star}$, a local minimum $x^{\bullet}$, and a local maximum $x^{\circ}$.We can also regard $x^{\circ}$ as a saddle point in the sense that the Hessian at this point has one negative eigenvalue. See Section 6.2 for more discussion. We use this toy example to gain insight into the expression (3.3) for the exponential decay constant $\lambda_{s}$; deferring the rigorous derivation of this number in the general case to Section 6.

From (3.1) it suggests that the lr-dependent SDE (1.2) takes about $O(1/\lambda_{s})$ time to achieve approximate stationarity. Intuitively, for the specific function in Figure 4, the bottleneck in achieving stationarity is to pass through the local maximum $x^{\circ}$. Now, we show that it takes about $O(1/\lambda_{s})$ time to pass $x^{\circ}$ from the local minimum $x^{\bullet}$. For simplicity, write

where $g(x)=f(x^{\bullet})$ stays constant if $x\leq x^{\circ}-\nu$ for a very small positive $\nu$ and $\theta>0$. Accordingly, the lr-dependent SDE (1.2) is reduced to the Ornstein–Uhlenbeck process,

before hitting $x^{\circ}$. Denote by $\tau_{x^{\circ}}$ the first time the Ornstein–Uhlenbeck process hits $x^{\circ}$. It is well known that the hitting time obeys

In passing, we remark that the discussion above can be made rigorous by invoking the theory of the Kramers escape rate, which shows that for this univariate case the hitting time satisfies

See, for example, [FW12, Pav14]. Furthermore, we demonstrate the view from the theory of viscosity solution and singular perturbation in Section B.4.

Why Learning Rate Decay?

As a widely used technique for training neural networks, learning rate decay refers to taking a large learning rate initially and then progressively reducing it during the training process. This technique has been observed to be highly effective especially in the minimization of nonconvex objective functions using stochastic optimization methods, with a very recent strand of theoretical effort toward understanding its benefits [YLWJ19, LWM19]. In this section, we offer a new and crisp explanation by leveraging the results in Section 3. To highlight the intuition, we primarily work with the continuous-time formulation of SGD.

For purposes of illustration, Figure 5 presents numerical examples for this technique where the learning rate is set to $0.1$ or $0.05$. This figure clearly demonstrates that SGD with a larger learning rate converges much faster to the global minimum than SGD with a smaller learning rate. This comparison reveals that a large learning rate would render SGD able to quickly explore the landscape of the objective function and efficiently escape bad local minima. On the other hand, a larger learning rate would prevent SGD iterates from concentrating around a global minimum, leading to substantial suboptimality. This is clearly illustrated in Figure 6. As suggested by the heuristic work on learning rate decay, we see that it is important to decrease the learning rate to achieve better optimization performance whenever the iterates arrive near a local minimum of the objective function.

In contrast, the sharp dependence of $k_{s}$ on the learning rate $s$ is not seen for strongly convex functions, because $\lambda_{s}=\mu$ stays constant as the learning rate $s$ varies. Following the preceding example, we have

In addition to taking a large $s$, an alternative way to make $T_{s}^{\delta}$ small is to have an initial distribution $\rho$ that is close to the stationary distribution $\mu_{s}$. This can be achieved by using the technique of learning rate decay. More precisely, taking a larger learning rate $s_{1}$ for a while, at the end the distribution of the iterates is approximately the stationary distribution $\mu_{s_{1}}$, which serves as the initial distribution for SGD with a smaller learning rate $s_{2}$ in the second phase. Taking $\rho\approx\mu_{s_{1}}$, the factor $\left\|\rho-\mu_{s}\right\|_{\mu_{s}^{-1}}$ in (4.1) for the second phase of learning rate decay is approximately

Both $\mu_{s_{1}}$ and $\mu_{s_{2}}$ are decreasing functions of $f$ and, therefore, have the same modes. As a consequence, the integral of $\mu_{s_{1}}^{2}/\mu_{s_{2}}$ is small by appeal to the rearrangement inequality, thereby leading to fast convergence of SGD with learning rate $s_{2}$ to the stationary risk $\epsilon(s_{2})$. In contrast, $\|\rho-\mu_{s_{2}}\|_{\mu_{s_{2}}^{-1}}$ would be much larger for a general random initialization $\rho$. Put simply, SGD with learning rate $s_{2}$ cannot achieve a risk of approximately $\epsilon(s_{2})$ given the same number of iterations without the warm-up stage using learning rate $s_{1}$. See Figure 8 for an illustration.

Proof of the Linear Convergence

In this section, we prove 3.1 and 3.2, leading to a complete proof of 1.

To better appreciate the linear convergence of the lr-dependent SDE (1.2), as established in 3.1, we start by showing the convergence to stationarity without a rate. In fact, this intermediate result constitutes a necessary step in the proof of 3.1.

The following result shows that the solution to the lr-dependent SDE converges to stationarity in terms of the dynamics of its probability densities over time.

Let $f$ satisfy the confining condition and denote the initial distribution as $\rho\in L^{2}(\mu_{s}^{-1})$. Then, the unique solution $\rho_{s}(t,\cdot)\in C^{1}\left([0,+\infty),L^{2}(\mu_{s}^{-1})\right)$ to the Fokker–Planck–Smoluchowski equation (2.1) converges in $L^{2}(\mu_{s}^{-1})$ to the Gibbs invariant distribution $\mu_{s}$, which is specified by (2.2).

Note that the existence and uniqueness of $\rho_{s}(t,\cdot)$ is ensured by 2.2. The convergence guarantee on $\rho_{s}(t,\cdot)$ in 5.2 relies heavily on the following lemma (5.3). This preparatory lemma introduces the transformation

with the initial distribution $h_{s}(0,\cdot)=\rho\mu_{s}^{-1}\in L^{2}(\mu_{s})$. The linear operator

has a crucial property, as stated in the following lemma. Its proof is postponed to Appendix C.2.

The linear operator $\mathscr{L}_{s}$ in (5.2) is self-adjoint and nonpositive in $L^{2}(\mu_{s})$. Explicitly, for any $g_{1},g_{2}$, this operator obeys

where the last equality is due to (5.1). Next, we proceed by making use of 5.3:

Thus, $\left\|\rho_{s}(t,\cdot)-\mu_{s}\right\|_{\mu_{s}^{-1}}^{2}$ is a strictly decreasing function, decreasing asymptotically towards the equilibrium state

This equality holds, however, only if $h_{s}(t,\cdot)$ is constant. Because both $\rho_{s}(t,\cdot)$ and $\mu_{s}$ are probability densities, this case must imply that $h_{s}(t,\cdot)\equiv 1$; that is, $\rho_{s}(t,\cdot)\equiv\mu_{s}$. Therefore, $\rho_{s}(t,\cdot)\in C^{1}\left([0,+\infty),L^{2}(\mu_{s}^{-1})\right)$ converges to the Gibbs invariant distribution $\mu_{s}$ in $L^{2}(\mu_{s}^{-1})$. ∎

We turn towards the proof of linear convergence. We first state a lemma which serves as a fundamental tool for us to prove a linear rate of convergence for 3.1.

for any $h$ such that the integrals above are well-defined.

For completeness, we provide a proof of this Poincaré-type inequality in Section C.3. For comparison, the usual Poincaré inequality is put into use for a bounded domain, as opposed to the entire Euclidean space as in 5.4. In addition, while the constant in the Poincaré inequality in general depends on the dimension (see, for example, [Eva10, Theorem 1, Chapter 5.8]), $\lambda_{s}$ in 5.4 is completely determined by geometric properties of the objective $f$. See details in Section 6.

Importantly, 5.4 allows us to obtain the following lemma, from which the proof of 3.1 follows readily. The proof of this lemma is given at the end of this subsection.

Under the assumptions of 3.1, $\rho_{s}(t,\cdot)$ converges to the Gibbs invariant distribution $\mu_{s}$ in $L^{2}(\mu_{s}^{-1})$ at the rate

where the first inequality applies the Cauchy-Schwarz inequality and

We conclude this subsection with the proof of 5.5.

Integrating both sides yields (5.4), as desired.

2 Proof of 3.2

Next, we turn to the proof of 3.2. We first state a technical lemma, deferring its proof to Section C.4.

Under the assumptions of 3.2, the excess risk at stationarity $\epsilon(s)$ satisfies

Using 5.6, we now finish the proof of 3.2.

Letting $g=f-f^{\star}$, we write the excess risk at stationarity as

for $s>0$. Consequently, we have proven that the excess risk $\epsilon(s)$ at stationarity is a strictly increasing function of $s\in[0,+\infty)$.

Geometrizing the Exponential Decay Constant

Having established the linear convergence to stationarity for the lr-dependent SDE, we now offer a quantitative characterization of the exponential decay constant $\lambda_{s}$ for a class of nonconvex objective functions. This is crucial for us to obtain a clear understanding of the dynamics of SGD and especially its dependence on the learning rate in the nonconvex setting.

We begin by deriving a relationship between the lr-dependent SDE (1.2) and a Schrödinger operator. Recall that the probability density $\rho_{s}(t,\cdot)$ of the SDE solution is assumed to be in $L^{2}(\mu_{s}^{-1})$. Consider the transformation

This transformation allows us to equivalently write the Fokker–Planck–Smoluchowski equation (2.1) as

is positive for sufficiently large $\|x\|$ due to the Villani condition.

Now, we collect some basic facts concerning the spectrum of the Schrödinger operator $-s\Delta+V_{s}$. First, it is a positive semidefinite operator, as shown below. Recognizing the uniqueness of the Gibbs distribution (2.2), it is not hard to show that $\sqrt{\mu_{s}}$ is the unique eigenfunction of $-s\Delta+V_{s}$ with a corresponding eigenvalue of zero. Using this fact, from the proof of 5.5, we get

Assume that $V$ is continuous, and $V(x)\rightarrow+\infty$ as $\|x\|\rightarrow+\infty$. Then the operator $-\Delta+V$ has a purely discrete spectrum.

A crucial fact from this representation is that the exponential decay constant $\lambda_{s}$ in Theorem 5.5 can be set to

As such, we can take $\lambda_{s}=\frac{1}{2}\zeta_{s,1}$ in the proof of 5.5.

As a consequence of this discussion, we seek to study the Fokker–Planck–Smoluchowski equation (2.1) by analyzing the spectrum of the linear Schrödinger operator (6.1), especially its smallest positive eigenvalue $\delta_{s,1}$. To facilitate the analysis, a crucial observation is that this Schrödinger operator is equivalent to the Witten-Laplacian,

The spectrum of the Witten-Laplacian has been the subject of a large literature [HN05, BGK05, Nie04, AK99], and in the next subsection, we exploit this literature to derive a closed-from expression for the first positive eigenvalue of the Witten-Laplacian, thereby obtaining the dependence of the exponential decay constant on the learning rate for a certain class of nonconvex objective functions [HHS11, Mic19].

2 The spectrum of the Witten-Laplacian: nonconvex Morse functions

We proceed by imposing the mild condition on the objective function that its first-order and second-order derivatives cannot be both degenerate anywhere. Put differently, the objective function is a Morse function. This allows us to use the theory of Morse functions to provide a geometric interpretation of the spectrum of the Witten-Laplacian.

Let $x$ be an index-$1$ saddle point and $r>0$ be sufficiently small. If $C_{1}(x,r)$ and $C_{2}(x,r)$ are contained in two different (maximal) connected components of the sublevel set $\mathcal{K}_{f(x)}$, we call $x$ an index-$1$ separating saddle point.

The remainder of this section aims to relate index-1 separating saddle points to the convergence rate of the lr-dependent SDE. For ease of reading, the remainder of the paper uses $x^{\circ}$ to denote an index-1 separating saddle point and writes $\mathcal{X}^{\circ}$ for the set of all these points. To give a geometric interpretation of 6.2, let $x_{1}^{\bullet}$ and $x_{2}^{\bullet}$ denote local minima in the two maximal connected components of $\mathcal{K}_{f(x^{\circ})}$, respectively. Intuitively speaking, the index-1 separating saddle point $x^{\circ}$ is the bottleneck of any path connecting the two local minima. More precisely, along a path connecting $x_{1}^{\bullet}$ and $x_{2}^{\bullet}$, by definition the function $f$ must attain a value that is at least as large as $f(x^{\circ})$. In this regard, the function value at $x^{\circ}$ plays a fundamental role in determining how long it takes for the lr-dependent SDE initialized at $x_{1}^{\bullet}$ to arrive at $x_{2}^{\bullet}$. See an illustration in Figure 9.

As is assumed in this section, $f$ is a Morse function and satisfies both the confining and the Villani conditions; in this case, it can be shown that the number of the critical points of $f$ is finite. Thus, denote by $n^{\circ}$ the number of index-$1$ separating saddle points of $f$ and let $n^{\bullet}$ denote the number of local minima.

To describe the labeling procedure, consider the set of the objective values at index-$1$ separating saddle points $\mathcal{V}=\{f(x^{\circ}):x^{\circ}\in\mathcal{X}^{\circ}\}$. This is a finite set and we use $I$ to denote the cardinality of this set. Write $\mathcal{V}=\{\nu_{1},\ldots,\nu_{I}\}$ and sort these values as

where by convention $\nu_{0}=+\infty$ corresponds to a fictive saddle point at infinity.

Next, we follow [HHS11] and define a type of connected components of sublevel set.

With the preparatory notions above in place, we describe the following procedure for labeling index-1 separating saddle points and local minima [HHS11]. See Figure 10 for an illustration of this process.

Let $\mathcal{X}^{\bullet}_{0}$ denote the singleton set $\{x^{\star}\}$.

Let $E^{1}_{j}$ for $j=1,\ldots,m_{1}$ be the critical components of the sublevel set $\mathcal{K}_{\nu_{1}}$. Note that $E^{1}_{1}\cup\cdots\cup E^{1}_{m_{1}}$ is a (proper) subset of $\mathcal{K}_{\nu_{1}}$. Without loss of generality, assume $x^{\star}\in E^{1}_{m_{1}}$. Then, we select $x_{1,j_{1}}^{\bullet}$ as

Define $\mathcal{X}^{\bullet}_{1}:=\{x_{1,1}^{\bullet},\ldots,x_{1,m_{1}-1}^{\bullet}\}$.

For $i=2,\ldots,I$, let $E^{i}_{j}$ for $j=1,\ldots,m_{i}$ be the critical components of the sublevel set $\mathcal{K}_{\nu_{i}}$. Without loss of generality, we assume that the critical components are ordered such that there exists an integer $k_{i}\leq m_{i}$ satisfying

for any $j=k_{i}+1,\ldots,m_{i}$. Set $x^{\bullet}_{i,j}$ to

for $j=1,\ldots,k_{i}$. Define $\mathcal{X}_{i}^{\bullet}:=\{x_{i,1}^{\bullet},\ldots,x_{i,k_{i}}^{\bullet}\}$.

To make the labeling process above valid, however, we need to impose the following assumption on the objective. This assumption is generic in the sense that it should be satisfied by a generic Morse function.

For every critical component $E^{i}_{j}$ selected in the labeling process above, where $i=0,1,\ldots,I$, we assume that

The minimum $x^{\bullet}_{i,j}$ of $f$ in any critical component $E^{i}_{j}$ is unique.

If $E_{j}^{i}\cap\mathcal{X}^{\circ}\neq\varnothing$, there exists a unique $x_{i,j}^{\circ}\in E_{j}^{i}\cap\mathcal{X}^{\circ}$ such that $f(x_{i,j}^{\circ})=\max\limits_{x\in E_{j}^{i}\cap\mathcal{X}^{\circ}}f(x)$. In particular, $E_{j}^{i}\cap\mathcal{K}_{f(x_{i,j}^{\circ})}$ is the union of two distinct critical components.

The first condition in this assumption requires that there exists a unique minimum of the objective $f$ in every critical component $E^{i}_{j}$. In particular, the global minimum $x^{\star}$ is unique under this assumption. In addition, the second condition requires that among all index-$1$ separating saddle points in $E^{i}_{j}$, if any, $f$ attains the maximum at exactly one of these points.

Under 6.4, the above labeling process includes all the local minima of $f$. Moreover, it reveals a remarkable result: there exists a bijection between the set of local minima and the set of index-$1$ separating saddle points (including the fictive one) $\mathcal{X}^{\circ}\cup\{\infty\}$. As shown in the labeling process, for any local minimum $x_{i,j}^{\bullet}$, we can relate it to the index-1 separating saddle point at which $f$ attains the maximum in the critical component $E_{j}^{i}$. See Figure 10 for an illustrative example. Interestingly, this shows that the number of local minima is always larger than the number of index-1 separating saddle points by one; that is, $n^{\circ}=n^{\bullet}-1$.

where $f(x^{\circ}_{0})-f(x^{\bullet}_{0})=f(\infty)-f(x^{\star})=+\infty$. A detailed description of this bijection is given in [HHS11, Proposition 5.2].

Under 6.4 and the assumptions of 2, there exists $s_{0}>0$ such that for any $s\in(0,s_{0}]$, the first $n^{\bullet}-1$ smallest positive eigenvalues of the Witten-Laplacian $\Delta^{s}_{f}$ associated with $f$ satisfy

Using 6.5 in conjunction with the simple relationship between the exponential decay constant and the spectrum of the Schrödinger operator/Witten-Laplacian (6.2), it is a stone’s throw to prove 2 when $f$ is generic. First, we give the definition of the Morse saddle barrier.

Let $f$ satisfy the assumptions of 2. We call $H_{f}=f(x_{1}^{\circ})-f(x_{1}^{\bullet})$ the Morse saddle barrier of $f$.

By 6.5, we can set the exponential decay constant to

in 2. Taking $\alpha=\frac{1}{2}\frac{|\eta_{d}(x_{1}^{\circ})|}{2\pi}\left(\frac{\det(\nabla^{2}f(x_{1}^{\bullet}))}{-\det(\nabla^{2}f(x_{1}^{\circ}))}\right)^{\frac{1}{2}}$ in (3.3), we complete the proof when $f$ falls into the generic case.

However, the generic assumption for the labeling process is complex, leading to the lack of a geometric interpretation of the objective function required for the labeling process. To gain further insight, we present a simplifying assumption that is a special case of 6.4. This simplification is due to [Nie04].

The following result follows immediately from 6.5.

Under 6.7 and the assumptions of 2, 6.5 holds. Therefore, 2 holds in this case.

We say that a Morse function is degenerate if it satisfies the assumptions of 2 but not 6.4. To violate the generic assumption, for example, we can change the objective value $f(x_{3,1}^{\bullet})$ to $f(x_{1,1}^{\bullet})$ or change $f(x_{3,2}^{\bullet})$ to $f(x_{2,3}^{\bullet})$ in Figure 10. In this situation, the first condition in 6.4 is not satisfied. Alternatively, if the objective value at $x_{3,1}^{\circ}$ is changed to $f(x_{2,1}^{\circ})$, the second condition in 6.4 is not met. Figure 11 presents an example of a degenerate Morse function.

In particular, the number of local minima remains larger than that of index-1 separating saddle points by one in this case. The following result extends 6.5 to the degenerate case, which is adapted from Theorem 2.8 of [Mic19].

Assume that the assumptions of 2 are satisfied but not 6.4. Then, there exists $s_{0}>0$ such that for any $s\in(0,s_{0}]$, the first $n^{\bullet}-1$ smallest positive eigenvalues of the Witten-Laplacian $\Delta^{s}_{f}$ associated with $f$ satisfy

Discussion

In this paper, we have presented a theoretical perspective on the convergence of SGD in nonconvex optimization as a function of the learning rate. Introducing the notion of an lr-dependent SDE, we have leveraged modern tools for the study of diffusions, in particular the spectral theory of diffusion operators, to analyze the dynamics of SGD in a continuous-time model. Specifically, we have shown that the solution to the SDE converges linearly to stationarity under certain regularity conditions and we have presented a concise expression for the linear rate of convergence with transparent dependence on the learning rate for nonconvex Morse functions. Our results show that the linear rate is a constant in the strongly convex case, whereas it decreases rapidly as the learning rate decreases in the nonconvex setting. We have thus uncovered a fundamental distinction between convex and nonconvex problems. As one implication, we note that noise in the gradients plays a more determinative role in stochastic optimization with nonconvex objectives as opposed to convex objectives. We also note that our results provide a justification for the use of a large initial learning rate in training neural networks.

From a practical standpoint, our work offers several promising avenues for future research in deep learning. First, a seemingly straightforward direction is to extend our SDE-based analysis to various learning rate schedules used in practice in training deep neural networks, such as diminishing learning rate and cyclical learning rates [BCN18, Smi17]. More broadly, it is of great interest to use SDEs to study and improve on practical variants of SGD, including RMSProp and Adam [TH12, KB14]. Second, our results would likely to be useful in guiding the choice of hyperparameters of deep neural networks from an optimization viewpoint. For instance, recognizing the essence of the exponential decay constant $\lambda_{s}$ in determining the convergence rate of SGD, how to choose the neural network architecture and the loss function so as to get a small value of the Morse saddle barrier $H_{f}$? Finally, we wonder if the lr-dependent SDE might give insights into generalization properties of neural networks such as local elasticity [HS20] and implicit regularization [ZBH+16, GLSS18].

We would like to thank Zhuang Liu and Yu Sun for helpful conversations about practical experience in deep learning. This work was supported in part by NSF through CAREER DMS-1847415, CCF-1763314, and CCF-1934876, and the Wharton Dean’s Research Fund. We also recognize support from the Mathematical Data Science program of the Office of Naval Research under grant number N00014-18-1-2764.

References

Appendix A Technical Details for Sections 1 and 2

Figure 12 presents a diagram that shows approximating surrogates for GD, SGD, and SGLD at multiple scales. In the case of SGD, for example, the inclusion of only $O(1)$ terms leads to the ODE $\dot{X}=-\nabla f(X)$, whereas the inclusion of up to $O(\sqrt{s})$ terms leads to the lr-dependent SDE (1.2). For GD and SGLD, $O(\sqrt{s})$ terms are not found in the expansion as in the derivation of (1.2). The $O(\sqrt{s})$-approximation, therefore, leads to the same differential equation as the $O(1)$-approximation for both GD and SGLD.

A.2 Derivation of the Fokker–Planck–Smoluchowski equation

To derive the lr-dependent Fokker–Planck–Smoluchowski equation (2.1), we first state the following lemma.

Recognizing the invariance of translation of time and letting $u_{s}(t-t^{\prime},x)=v_{s}(t^{\prime},x)$, we can reduce (A.3) to the following backward Fokker–Planck–Smoluchowski equation:

Next, from the Chapman–Kolmogorov equation, we get

and by switching the order of the integration, we obtain

Making use of the backward Fokker–Planck–Smoluchowski equation (A.4) and switching the order of integration (A.5), we get

Hence, we derive the forward Fokker–Planck–Smoluchowski equation at $t=0$ for an arbitrary smooth function $u_{s}(0,x)=g(t,x)$. Noting that $t=0$ can be replaced by any time $t$, we complete the derivation of the Fokker–Planck–Smoluchowski equation.

A.3 The uniqueness of Gibbs invariant distribution

We begin by proving that the probability density $\mu_{s}$ is an invariant distribution of (2.1). Plugging

We now proceed to show that the probability density $\mu_{s}$ is unique. To derive a contradiction, we assume that there exists another distribution $\vartheta_{s}$ satisfying the Fokker–Planck–Smoluchowski equation:

Write $\varpi_{s}=\vartheta_{s}\mu_{s}^{-1}$ and recall the operator $\mathscr{L}_{s}$ defined in Section 5.1. We can rewrite (A.8) as

A.4 Proof of 2.2

Thus, the solution to the partial differential equation (6.1) is

Recognizing the transformation $\psi_{s}(t,\cdot)=\rho_{s}(t,\cdot)/\sqrt{\mu_{s}}$, we recover

Appendix B Technical Details for Section 3

Here, we prove 3.4 using the Bakry–Emery theorem, which is a Poincaré-type inequality for $\mu$-strongly convex functions. As a direct consequence of this theorem, the exponential decay constant for strongly convex objectives does not depend on the learning rate $s$ and the ambient dimension $d$.

B.1 serves as the main technical tool in the proof of 3.4. Its proof is in Section B.1.1. Now, we prove the following result using B.1.

Under the same assumptions as in 3.4, $\rho_{s}(t,\cdot)$ converges to the Gibbs distribution $\mu_{s}$ in $L^{2}(\mu_{s}^{-1})$ at the rate

Integrating both sides yields (B.1), as desired. ∎

Leveraging B.2, we proceed to complete the proof of 3.4.

where the first inequality applies the Cauchy-Schwarz inequality and

We introduce two operators $\Gamma_{s}$ and $\Gamma_{s,2}$ that are built on top of the linear operator $\mathscr{L}_{s}$ defined in (5.2). For any $g_{1},g_{2}\in L^{2}(\mu_{s})$, let

A simple relationship between the two operators is described in the following lemma.

Under the same assumptions as in B.1, for any $g\in L^{2}(\mu_{s})$ we have

Then, the operator $\Gamma_{s}$ must satisfy

we obtain that the operator $\Gamma_{s,2}$ satisfies

where $\mathbf{Tr}$ is the standard trace of a squared matrix. Recognizing that the objective $f$ is $\mu$-strongly convex, a comparison between (B.4) and (B.5) completes the proof.

Recall that $h_{s}(t,\cdot)\in L^{2}(\mu_{s})$ is the solution to the partial differential equation (5.1), with the initial condition $h_{s}(0,\cdot)=h$. Define

The following lemma considers the derivatives of $\Lambda_{1,s}(t)$.

Under the same assumptions as in B.1, we have

Since $h_{s}(t,\cdot)\in L^{2}(\mu_{s})$ is the solution to the partial differential equation (5.1), we get

Furthermore, by the definition of $\Gamma_{s,2}$ and integration by parts, we haveSee the calculation in [BGL13].

From 5.3, we know that the linear operator $\mathscr{L}_{s}$ is self-adjoint. Then, we obtain the second derivative as

Using B.3 and B.4, we obtain the following inequality:

From the definition of $\Lambda_{1,s}(t)$, we have

where the second term on the right-hand side follows from 5.2 and

By 5.2, we get $h_{s}(\infty,\cdot)\equiv 1$, which together with (B.7) gives

The final equality follows from (B.4). Integrating both sides of the inequality (B.8), we have

Similarly, the proof of B.5 will be based on the following log-Sobolev type inequality.

By the Csiszár–Kullback inequality, we have

Finally, using B.6, we obtain the estimate for the derivative of the entropy (B.12) as

with $h_{s}(0,x)=h^{2}(x)$. We now proceed to find the derivatives of $\Lambda_{2,s}(t)$ with respect to time $t$, which is formulated as the following lemma.

Under the same assumptions as in B.6, we have

Recall the linear operator $\mathscr{L}_{s}$ defined in (5.2). Using integration by parts, we have

Since $h_{s}(t,\cdot)$ is the solution to the partial differential equation (5.1), we obtain the first derivative as

Furthermore, recognizing the definitions of $\Gamma_{s,2}$, $\Gamma_{s}$ and $\mathscr{L}_{s}$, we have

Then the first equality above by the definitions can be written as

Finally, we obtain the second derivative as

Next, we complete the proof of Lemma B.6.

Since $\dot{\Lambda}_{2,s}(t)$ is positive, we get

By the definition of $\Lambda_{2,s}(t)$, we have

where the second term in the right-hand side follows from 5.2 and

By 5.2, we know that $h_{s}(\infty,\cdot)\equiv 1$. Plugging it into (B.14), we get

where the last equality follows from (B.4). Integrating both sides of the inequality (B.15), we have

B.3 Proof of Proposition 3.5

By 2.2, let $\rho_{s}(t,\cdot)\in C^{1}([0,+\infty),L^{2}(\mu_{s}^{-1}))$ denote the unique transition probability density of the solution to the lr-dependent SDE. Taking an expectation, we get

Hence, the uniqueness has been proved. Using the Cauchy–Schwarz inequality and Theorem 5.5, we obtain:

B.4 Connection with vanishing viscosity

Taking $s=0$, the zero-viscosity steady-state equation of the Fokker–Planck–Smoluchowski equation (2.1) reads

A solution to this zero-viscosity steady-state equation takes the form

where $x_{i}$’s are critical points of the objective $f$. As is clear, the solution is not unique. However, we have shown previously that the invariant distribution $\mu_{s}$ is unique and converges to

in the sense of distribution, which is a special case of (B.17). Clearly, when there exists more than one critical point, $\mu_{s\rightarrow 0}(x)$ is different from $\mu_{0}(x)$ in general. In contrast, $\mu_{s\rightarrow 0}(x)$ and $\mu_{0}(x)$ must be the same for (strictly) convex functions. In light of this comparison, the correspondences between the case $s>0$ and the case $s=0$ are fundamentally different in nonconvex and convex problems.

Next, we consider the rate of convergence in the convex setting. Let

where $\theta>0$. Plugging into the Fokker-Planck-Smoluchowski equation (2.1), we have

Furthermore, using the following inequality

The phenomenon presented above is called singular perturbation. It appears in mathematical models of boundary layer phenomena [CM90, Chapter 2.2, Example 1 and Example 2], WKB theory for Schrödinger equations [Gas07, Supplement 4A], KAM theory for circle diffeomorphisms [Arn12, Chapter 2, Section 11] and that for Hamilton systems [Arn13, Appendix 8]. Moreover, the singular perturbation phenomenon shows that there exists a fundamental distinction between the $O(1)$-approximating ODE for SGD and the lr-dependent SDE (1.2). In particular, the learning rate $s\rightarrow 0$ in the Fokker–Planck–Smoluchowski equation (2.1) corresponds to vanishing viscosity. The vanishing viscosity phenomenon was originally observed in fluid mechanics [CM90, KCD08], particularly in the degeneration of the Navier–Stokes equation to the Euler equation [CF99]. As a milestone, the vanishing viscosity method has been used to study the Hamilton–Jacobi equation [CL83, Eva80, CEL84]. In fact, the Fokker–Planck–Smoluchowski equation (2.1) and its stationary equation are a form of Hamilton–Jacobi equation with a viscosity term, for which the Hamiltonian is

The Hamiltonian (B.21) is different from the classical case [Lio82, CS04, Eva10], which is generally nonlinear in $\nabla\rho$ (cf. Burger’s equation). Although the Hamiltonian depends linearly on $\rho$ and $\nabla\rho$, the coefficients depend on $\Delta f$ and $\nabla f$. Hence, it is not reasonable to apply directly the well-established theory of Hamilton–Jacobi equations [CL83, Eva80, CEL84, Lio82, CS04, Eva10] to the Fokker–Planck–Smoluchowski equation (2.1) and its stationary equation. Furthermore, for the aforementioned example, which proves the $O(\sqrt{s})$ convergence for the Fokker–Planck–Smoluchowski equation with the quadratic potential $f(x)=\frac{\theta}{2}x^{2}$, is also a viscosity solution to the Hamilton–Jacobi equation [CL83], since the Hamiltonian (B.21) for the quadratic potential degenerates to

where $f(x)=x^{T}Ax$ and $A$ is positive definite and symmetric. Thus, we remark that the general theory of viscosity solutions to Hamilton–Jacobi equations cannot be used directly to prove the theorems in the main body of this paper.

In closing, we present several open problems.

Consider the stationary solution $\mu_{s}(x)$ to the Fokker–Planck–Smoluchowski equation (2.1). For a convex or strongly convex objective $f$ with Lipschitz gradients, can we quantify the rate of convergence? Does the rate of convergence remain $O(\sqrt{s})$?

Let $T>0$ be fixed and consider the solution $\rho_{s}(t,x)$ to the Fokker–Planck–Smoluchowski equation (2.1) in $[0,T]$. For a convex or strongly convex objective $f$ with Lipschitz gradients, does the solution to the Fokker–Planck–Smoluchowski equation (2.1) converge to the solution to its zero-viscosity equation (B.20)? Is the rate of convergence still $O(\sqrt{s})$?

Consider the solution $\rho_{s}(t,x)$ to the Fokker–Planck–Smoluchowski equation (2.1) in $[0,+\infty)$. For a convex or strongly convex objective $f$ with Lipschitz gradients, does the global solution to the Cauchy problem of the Fokker–Planck–Smoluchowski equation (2.1) converge to the solution of its zero-viscosity equation (B.20)? Is the rate of convergence still $O(\sqrt{s})$?

Appendix C Technical Details for Section 5

From the Cauchy–Schwarz inequality, we get

C.2 Proof of Lemma 5.3

Recall that the linear operator $\mathscr{L}_{s}$ in (5.2) is defined as

Therefore, $\mathscr{L}_{s}$ is self-adjoint in $L^{2}(\mu_{s})$ and is non-positive.

C.3 Proof of Lemma 5.4

we can obtain the following key inequality [DS01]

Let $R_{0,s}>0$ be large enough such that $V_{s}(x)>0$ for $\|x\|\geq R_{0,s}$. For $R_{s}>R_{0,s}$, we can define $\epsilon_{s}$ as

Then $\epsilon(R_{s})\rightarrow 0$ as $R_{s}\rightarrow\infty$. Furthermore, we assume the $R_{s}$ is large enough such that

From the key inequality (C.2), we obtain that

Using the Poincaré inequality in a bounded domain [Eva10, Theorem 1, Chapter 5.8], we get

where $C(R_{s})$ is a constant depending on $R_{s}$. Furthermore, using the inequality (C.4), we obtain

Making use of the mean-zero property of $h$, we have

Taking (C.5) and (C.8) together, we obtain

C.4 Proof of 5.6

For convenience, we introduce a shorthand:

Next, we assume that $\zeta_{k}(x)=x^{k}e^{-x^{\alpha}}$, where $\alpha<1$ is a fixed positive constant and $k=1,2$. The facts that $\zeta_{k}(0)=0$ and $\lim\limits_{x\rightarrow+\infty}\zeta_{k}(x)=0$ give