Fractional Underdamped Langevin Dynamics: Retargeting SGD with Momentum under Heavy-Tailed Gradient Noise
Umut Şimşekli, Lingjiong Zhu, Yee Whye Teh, Mert Gürbüzbalaban
Introduction
Gradient-based optimization algorithms have been the de facto choice in deep learning for solving the optimization problems of the form:
where denotes a random subset drawn from the set of data points with for all .
where is still called the velocity. The connection between this system and (2) becomes clearer, if we discretize this system by using the Euler scheme with step-size :
While the Gaussianity assumption can be accurate in certain settings such as small networks (Martin & Mahoney, 2019; Panigrahi et al., 2019), recently it has been empirically demonstrated that in several deep learning setups, the stochastic gradient noise can exhibit a heavy-tailed behavior (Şimşekli et al., 2019a; Zhang et al., 2019b)In two recent studies, Gürbüzbalaban et al. (2020) and Hodgkinson & Mahoney (2020) have shown that the stationary distribution of the stochastic gradient descent (SGD) algorithm can be indeed a heavy-tailed distribution depending on the choice of the step-size and the batch-size. On the other hand, in another recent study, Şimşekli et al. (2020) have provided generalization bounds for a general class of SDEs, including heavy- and light-tailed ones.. While the Gaussianity assumption would not be appropriate in this case since the conventional CLT would not hold anymore, nevertheless we can invoke the generalized CLT, which states that the asymptotic distribution of will be a symmetric -stable distribution (); a class of distributions that are commonly used in the statistical physics literature as an approximation to heavy-tailed random variables (Sliusarenko et al., 2013; Dubkov et al., 2008). As we will define in more detail in the next section, in the core of , lies the parameter , which determines the heaviness of the tail of the distribution. The tails get heavier as gets smaller, the case reduces to the Gaussian random variables. This is illustrated in Figure 1.
Şimşekli et al. (2019b, a) empirically illustrated that, in deep neural networks, the statistical structure of can be better captured by using an -stable distribution. With the assumption of being distributed, the choice of Brownian motion will be no longer appropriate and should be replaced with an -stable Lévy motion, which motivates the following Lévy-driven SDE:
A more striking property of (7) was very recently revealed in a statistical physics study (Capała & Dybiec, 2019), where the authors numerically illustrated that, even when has a single minimum, the invariant measure of (7) can exhibit multiple maxima, none of which coincides with the minimum of . A similar property has been formally proven in the overdamped dynamics with Cauchy noise (i.e., and ) by Sliusarenko et al. (2013). Since the process (7) would spend more time around the modes of its invariant measure (i.e., the high probability region), in an optimization context (i.e., for larger ) the sample paths would concentrate around these modes, which might be arbitrarily distant from the optima of . In other words, the heavy-tails of the gradient noise could result in an undesirable bias, which would be still present even when the step-size is taken to be arbitrarily small. As we will detail in Section 3, informally, this phenomenon stems from the fact that the heavy-tailed noise leads to aggressive updates on , which are then directly transmitted to due to the dynamics. Unless ‘tamed’, these updates create an hurling effect on and drift it away from the modes of the “potential” that is sought to be minimized.
Contributions: In this study, we develop a fractional underdamped Langevin dynamics whose invariant distribution is guaranteed to be in the form of the Boltzmann-Gibbs measure, hence its optima exactly match the optima of . We first prove a general theorem which holds for any kinetic energy function, which is not necessarily the Gaussian kinetic energy. However, it turns out that some components of the dynamics might not admit an analytical form for an arbitrary choice of the kinetic energy. Then we identify two choices of kinetic energies, where all the terms in dynamics can be written in an analytical form or accurately computable. We also analyze the Euler discretization of (14) and identify sufficient conditions for ensuring weak convergence of the ergodic averages computed over the iterates.
We observe that the discretization of the proposed dynamics has interesting algorithmic similarities with natural gradient descent (Amari, 1998) and gradient clipping (Pascanu et al., 2013), which we believe bring further theoretical understanding for their role in deep learning. Finally, we support our theory with experiments conducted on both synthetic settings and neural networks.
Technical Background & Related Work
Lévy motions are stochastic processes with independent and stationary increments. Their successive displacements are random and independent, and statistically identical over different time intervals of the same length, and can be viewed as the continuous-time analogue of random walks. The best known and most important examples are the Poisson process, Brownian motion, the Cauchy process and more generally stable processes. Lévy motions are prototypes of Markov processes and of semimartingales, and concern many aspects of probability theory. We refer to (Bertoin, 1996) for a survey on the theory of Lévy motions.
In general, Lévy motions are heavy-tailed, which make it appropriate to model natural phenomena with possibly large variations, that often occurs in statistical physics (Eliazar & Klafter, 2003), signal processing (Kuruoglu, 1999), and finance (Mandelbrot, 1997).
where the drift is defined as follows:
Here, and denotes the fractional Riesz derivative (Riesz, 1949):
The important property of the process (8) is that it admits an invariant distribution whose density is proportional to (Nguyen et al., 2019). It is easy to show that, when , the drift reduces to , hence we recover the classical overdamped dynamics:
Since the fractional Riesz derivative is costly to compute, Şimşekli (2017) proposed an approximation of based on the alternative definition of given in (Ortigueira, 2006), such that:
From a pure Monte Carlo perspective, Ye & Zhu (2018) extended the fractional overdamped dynamics (8) to higher-order dynamics and proposed the so-called fractional Hamiltonian dynamics (FHD), given as follows:
where , and . They showed that the invariant measure of the process has a density proportional to , i.e., the Boltzmann-Gibbs measure. Similar to the overdamped case (8), the Riesz derivatives do not admit an analytical form in general. Hence they approximated them by using the same approximation given in (12), which yields the SDE given in (7) (up to a scaling factor). This observation also confirms that the heavy-tailed noise requires an adjustment in the dynamics, otherwise the induced bias might drive the dynamics away from the minima of (Capała & Dybiec, 2019).
Fractional Underdamped Langevin Dynamics
In this section, we develop the fractional underdamped Langevin dynamics (FULD), which is expressed by the following SDE:
Let has the following form:
One of the main features of FULD is that the fractional Riesz derivatives only appears in the drift , which only depends on . This is highly in contrast with FHD (13), where the Riesz derivatives are taken over both and , which is the source of intractability. Moreover, FULD enjoys the freedom to choose different kinetic energy functions . In the sequel, we will investigate two options for , such that the drift can be analytically obtained.
In classical overdamped Langevin dynamics and Hamiltonian dynamics, the default choice of kinetic energy is the Gaussian kinetic energy, which corresponds to taking (Neal, 2010; Livingstone et al., 2019; Dalalyan & Riou-Durand, 2020). With this choice, the fractional dynamics becomes:
In the next result, we will show that in this case, the drift admits an analytical solution.
Let . Then, for any ,
where is the gamma function and is the Kummer confluent hypergeometric function. In particular, when , we have .
We observe that the fractional dynamics (16) strictly extends the underdamped Langevin dynamics (6) as .
Even though this aggressive behavior of can be beneficial for the continuous-time system, it is unfortunately clear that its Euler-Maruyama discretization will not yield a practical algorithm due to the same behavior. Indeed, we would need the function to be Lipschitz continuous in order to guarantee the algorithmic stability of its discretization (Kloeden & Platen, 1999); however, if we consider the integral form of (cf. (Abramowitz & Stegun, 1972)), we observe that the function
is clearly not Lipschitz continuous in . Therefore, we conclude that FULD with the Gaussian kinetic energy is mostly of theoretical interest.
2 Alpha-stable kinetic energy
The dynamics with the Gaussian kinetic energy requires a very strong drift mainly because we force the dynamics to make sure that the invariant distribution of to be a Gaussian. Since the Gaussian distribution has light-tails, it cannot tolerate samples with large magnitudes, hence requires a large dissipation to make sure does not take large values.
In order to avoid such an explosive drift that potentially degrades practicality, next we explore heavy-tailed kinetic energies, which would allow the components of to take large values, while still making sure that the drift in (15) admits an analytical form.
Let be the probability density function of . Choose in (15), where for any . Then,
It now follows from Theorem 3 that the FULD with -stable kinetic energy reduces to the following SDE:
It can be easily verified that for , as , hence, the SDE (19) also reduces to the classical underdamped Langevin dynamics (6).
While this choice of results in an analytically available , unfortunately the function itself admits a closed-form analytical formula only when or , due to the properties of the densities. Nevertheless, as is based on one-dimensional densities, it can be very accurately computed by using the recent methods developed in (Ament & O’Neil, 2018). On the other hand, in the next section, we will show that is Lipschitz continuous for all , which implies that under standard regularity conditions on , the Boltzmann-Gibbs measure is the unique invariant measure of (19).
We visually inspect the behavior of in Figure 2 for dimension one. We observe that, as soon as , takes a very smooth form. Besides, for small the function behaves like a linear function and when goes to infinity, it vanishes. This behavior can be interpreted as follows: since can take larger values due to the heavy tails of the kinetic energy, in order to be able target the correct distribution, the dynamics compensates the potential bursts in by passing it through the asymptotically vanishing .
3 Euler discretization and weak convergence analysis
As visually hinted in Figure 2, the function has strong regularity, which makes (19) to be potentially beneficial for practical implementations. Indeed, it is easy to verify that is Lipschitz continuous for and , and in our next result, we show that this observation is true for any admissible , which is a desired property when discretizing continuous-time dynamics.
For , the map is Lipschitz continuous, hence is also Lipschitz continuous.
Accordingly we consider the following Euler-Maruyama discretization for (19):
where is the cumulative sum of the step-size sequence.
We note that Langevin-based algorithms have been used in the literature to obtain global convergence guarantees for non-convex optimization, see e.g. (Raginsky et al., 2017; Xu et al., 2018; Gao et al., 2018b; Zou et al., 2019; Nguyen et al., 2019). In particular, Nguyen et al. (2019) used an overdamped fractional Langevin dynamics for non-convex optimizations. The proposed model in our paper can also be used to study the non-convex optimizations and we expect that our underdamped dynamics may have improved theoretical guarantees compared to (Nguyen et al., 2019).
We now present the assumptions that imply our results.
The step-size sequence is non-increasing and satisfies and .
These are common assumptions ensuring that the SDE is simulated with infinite time-horizon and the process is not explosive (Panloup, 2008; Şimşekli, 2017). We can now establish the weak convergence of (21) and present it as a corollary to Theorem 1, Proposition 1, and (Panloup, 2008) (Theorem 2).
Assume that the gradient is Lipschitz continuous and has linear growth i.e., there exists such that for all . Furthermore, assume that Assumptions 1 and 2 hold for some . If the test function then
4 Connections to existing approaches
Let us now consider gradient-clipping, a heuristic approach for eliminating the problem of ‘exploding gradients’, which often appear in training neural networks (Pascanu et al., 2013; Zhang et al., 2019a). Very recently, Zhang et al. (2019b) empirically illustrated that such explosions stem from heavy-tailed gradients and formally proved that gradient clipping indeed improves convergence rates under heavy-tailed perturbations. We notice that, the behavior of (22) is reminiscent of gradient clipping: due to the vanishing behavior of for , as the components of gets larger in magnitude, the update applied on gets smaller. The behavior becomes more prominent in (23). On the other hand, (22) is more aggressive in the sense that the updates can get arbitrarily small as the value of decreases as opposed to being ‘clipped’ with a threshold.
Numerical Study
In this section, we will illustrate our theory on several experiments which are conducted in both synthetic and real-data settingsWe provide our implementation in https://github.com/umutsimsekli/fuld.. We note that, as expected, FULD with Gaussian kinetic energy did not yield a numerically stable discretization due to the explosive behavior of . Hence, in this section, we only focus on FULD with kinetic energy and from now on we will simply refer to FULD with kinetic energy as FULD.
We first consider a one-dimensional synthetic setting, similar to the one considered in (Capała & Dybiec, 2019). We consider a quartic potential function with a quadratic component, . We then simulate the ‘uncorrected dynamics’ (UD) given in (7) and FULD (19) by using the Euler-Maruyama discretization to compare their behavior for different . For , we used the software given in (Ament & O’Neil, 2018) for computing .
Figure 3 illustrates the distribution of the samples generated by simulating the two dynamics. In this setup, we set , , with number of iterations . We observe that, for , FULD very accurately captures the form of the distribution, whereas UD exhibits a visible bias and the shape of its resulting distribution is slightly distorted. Nevertheless, since the perturbations are close to a Gaussian in this case (i.e., is close to ), the difference is not substantial and can be tolerable in an optimization context. However, this behavior becomes much more emphasized when we use a heavier-tailed driving process: when , we observe that the target distribution of UD becomes distant from the Gibbs measure , and more importantly its modes no longer match the minima of ; agreeing with the observations presented in (Capała & Dybiec, 2019)We note that the overdamped dynamics with the uncorrected drift exhibits a similar behavior to the one of the uncorrected underdamped dynamics with sufficiently large . . On the other hand, thanks to the correction brought by , FULD still captures the target distribution very accurately, even when the driving force is Cauchy.
On the other hand, in our experiments we observed that, for small values of , UD can quickly become numerically unstable and even diverge for slightly larger step-sizes, whereas this problem never occurred for FULD. This outcome also stems from the fact that UD does not have any mechanism to compensate the potential large updates originating from the heavy-tailed perturbations. To illustrate this observation more clearly, in Figure 4 we illustrate the iterates which were used for producing Figure 3. We observe that, while the iterates of UD are well-behaved for , the magnitude range of the iterates gets quite large when is set to . On the other hand, for both values of , FULD iterates are always kept in a reasonable range, thanks to the clipping-like effect of .
2 Neural networks
In our next set of experiments, we evaluate our theory on neural networks. In particular, we apply the iterative scheme given in (22) as an optimization algorithm for training neural networks, and compare its behavior with classical SGDm defined in (2). In this setting, we do not add any explicit noise, all the stochasticity comes from the potentially heavy-tailed stochastic gradient noise (3) under the assumption that the noise can be well-modeled by using an vector (see Section 3.4 for the explicit assumption).
We consider a fully-connected network for a classification task on the MNIST and CIFAR10 datasets, with different depths (i.e. number of layers) and widths (i.e. number of neurons per layer). For each depth-width pair, we train two neural networks by using SGDm (2) and our modified version (22), and compare their final train/test accuracies and loss values. We use the conventional train-test split of the datasets: for MNIST we have K training and K test samples, and for CIFAR10 these numbers are K and K, respectively. We use the cross entropy loss (also referred to as the ‘negative-log-likelihood’).
We note that the modified scheme (22) reduces to (2) when , since . Hence in this section, we will refer to SGDm as the special case of (22) with . On the other hand, in these experiments, directly computing becomes impractical for , since the algorithms given in (Ament & O’Neil, 2018) become prohibitively slow with the increased dimension . However, since is based on the derivatives of the one-dimensional densities (see Theorem 3), for , we first precomputed the values of over a fine grid of ; then, during the SGDm recursion, we approximated by linearly interpolating the values of that are precomputed over this grid. We expect that, if the stochastic gradient noise can be well-approximated by using an distribution, then the modified dynamics should exhibit an improved performance since it would eliminate the potential bias brought by the heavy-tailed noise.
In these experiments, we set , for MNIST, and for CIFAR10. We run the algorithms for iterations Since the scale of the gradient noise is proportional to (see (20)), in this setup, a fixed implicitly determines . . We measure the accuracy and the loss at every 100th iteration and we report the average of the last two measurements. Figures 5 and 6 show the results obtained on the MNIST dataset. We observe that, in most of the cases, setting yields a better performance in terms both training and testing accuracies/losses. This difference becomes more visible when the width is set to : the accuracy difference between the algorithms reaches . We obtain a similar result on the CIFAR10 dataset, as illustrated in Figures 7 and 8. In most of the cases performs better, with the maximum accuracy difference being , implying the gradient noise can be approximated by an random variable.
We observed a similar behavior when the width was set to . However, when we set the width to we did not perceive a significant difference in terms of the performance of the algorithms. On the other hand, when the width was set to , resulted in a slightly better performance, which would be an indication that the Gaussian approximation is closer. The corresponding figures are provided in the supplementary document.
Conclusion and Future Directions
We considered the continuous-time variant of SGDm, known as the underdamped Langevin dynamics (ULD), and developed theory for the case where the gradient noise can be well-approximated by a heavy-tailed -stable random vector. As opposed to naïvely replacing the driving stochastic force in ULD, which correspondonds to running SGDm with heavy-tailed gradient noise, the dynamics that we developed exactly target the Boltzmann-Gibbs distribution, and hence do not introduce an implicit bias. We further established the weak convergence of the Euler-Maruyama discretization and illustrated interesting connections between the discretized algorithm and existing approaches commonly used in practice. We supported our theory with experiments on a synthetic setting and fully connected neural networks.
Our framework opens up interesting future directions. Our current modeling strategy requires a state-independent, isotropic noise assumption, which would not accurately reflect the reality. While anisotropic noise can be incorporated to our framework by using the approach of Ye & Zhu (2018), state-dependent noise introduces challenging technical difficulties. Similarly, it has been illustrated that the tail-index can depend on the state and different components of the noise can have a different (Şimşekli et al., 2019a). Incorporating such state dependencies would be an important direction of future research. Finally, it has been shown that the heavy-tailed perturbations yield shorter escape times (Nguyen et al., 2019) in the overdamped dynamics, and extending such results to the underdamped case is still an open problem.
Acknowledgments
We thank Jingzhao Zhang for fruitful discussions. The contribution of Umut Şimşekli to this work is partly supported by the French National Research Agency (ANR) as a part of the FBIMATRIX (ANR-16-CE23-0014) project, and by the industrial chair Data science & Artificial Intelligence from Télécom Paris. Lingjiong Zhu is grateful to the support from Simons Foundation Collaboration Grant. Mert Gürbüzbalaban acknowledges support from the grants NSF DMS-1723085 and NSF CCF-1814888.
References
Proof of Theorem 1
Let denote the probability density of . Then it satisfies the fractional Fokker-Planck equation (see Proposition 1 and Section 7 in (Schertzer et al., 2001)):
where we used the property (Proposition 1 in (Şimşekli, 2017)) and the semi-group property of the Riesz derivative .
Therefore, it follows from (24), (25) and (26) that we have
Proof of Theorem 2
Recall the definition of Fourier transform and its inverse:
Notice that the Fourier transform of is itself, i.e. , and moreover, , and therefore,
By the Taylor expansion of sine function, we get
where we used the identity , for any given . Moreover, for any given , we have the identity:
where is the Kummer confluent hypergeometric function. Therefore, we conclude that
By the identity , we get
In particular, when , by applying the identity
Proof of Theorem 3
Let be the probability density function of the symmetric -stable distribution such that
Proof of Proposition 1
It is straightforward to verify that the result holds for the cases and . Assume or . Let be the unit symmetric -stable random variable defined by its characteristic function
By taking inverse Fourier transformation, its density can be expressed as
Writing , we compute
where we used the fact that and are even functions of , whereas is an odd function of . If we define,
In particular, since and this implies that
It is also well-known that a symmetric stable random variable has a decay in its density satisfying when is large. In fact, Wintner (1941) derived a large- expansion for when and . This expansion is equivalent to
(see eqn. (11) from (Montroll & Bendler, 1984)) where it can be seen from the Stirling’s approximation of the gamma function and the ratio test that the series converges absolutely. A similar absolutely convergent series sum (with exactly the same leading term) is also available in the literature for which says that
(see eqn. (3.58) from (Montroll & West, 1979)). By differentiating the series sum for with respect to , we can express and as a series sum. After a straightforward computation, we obtain
which implies from (40) that as . This shows that is bounded on the interval . On the other hand, is an even function and therefore is an even function satisfying . We conclude that is bounded on the real line. This completes the proof. ∎
Proof of Corollary 1
By Proposition 1, we know that is Lipschitz and by our hypthesis is also Lipschitz and has linear growth. Then the process (19) admits a unique invariant measure (cf. (Schertzer et al., 2001) Section 9), which is given by Theorem 1. The rest of the proof follows from (Panloup, 2008) (Theorem 2). ∎
Alternative forms of the drift function c𝑐c with the Gaussian kinetic energy
For some special values of , we can get alternative formulas for , .
(1) . Using the identity , where is the modified Bessel function of the first kind, we get
(2) . Using the identity , we get
Visual Illustrations
On the other hand, we visualize the conformal Hamiltonian field generated by this dynamics in Figure 10 for . The figure shows that conformal Hamiltonian generated by the dynamics with has a very slow concentration behavior towards the minimum at the origin, whereas this behavior is alleviated when where the field concentrates faster.
Additional Experimental Results
In this section, we provide the additional experimental results that were mentioned in the main document for width , , and .