Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Gradient-adjusted underdamped Langevin dynamics for sampling

Xinzhe Zuo, Stanley Osher, Wuchen Li

TL;DR

A class of stochastic differential equations (SDEs) called gradient-adjusted underdamped Langevin dynamics (GAUL), which add stochastic perturbations in primal-dual damping dynamics and Hessian-driven damping dynamics from optimization, and it is proved that GAUL admits the correct stationary distribution, whose marginal is the target distribution.

Abstract

Sampling from a target distribution is a fundamental problem. Traditional Markov chain Monte Carlo (MCMC) algorithms, such as the unadjusted Langevin algorithm (ULA), derived from the overdamped Langevin dynamics, have been extensively studied. From an optimization perspective, the Kolmogorov forward equation of the overdamped Langevin dynamics can be treated as the gradient flow of the relative entropy in the space of probability densities embedded with Wassrstein-2 metrics. Several efforts have also been devoted to including momentum-based methods, such as underdamped Langevin dynamics for faster convergence of sampling algorithms. Recent advances in optimizations have demonstrated the effectiveness of primal-dual damping and Hessian-driven damping dynamics for achieving faster convergence in solving optimization problems. Motivated by these developments, we introduce a class of stochastic differential equations (SDEs) called gradient-adjusted underdamped Langevin dynamics (GAUL), which add stochastic perturbations in primal-dual damping dynamics and Hessian-driven damping dynamics from optimization. We prove that GAUL admits the correct stationary distribution, whose marginal is the target distribution. The proposed method outperforms overdamped and underdamped Langevin dynamics regarding convergence speed in the total variation distance for Gaussian target distributions. Moreover, using the Euler-Maruyama discretization, we show that the mixing time towards a biased target distribution only depends on the square root of the condition number of the target covariance matrix. Numerical experiments for non-Gaussian target distributions, such as Bayesian regression problems and Bayesian neural networks, further illustrate the advantages of our approach.

Gradient-adjusted underdamped Langevin dynamics for sampling

TL;DR

A class of stochastic differential equations (SDEs) called gradient-adjusted underdamped Langevin dynamics (GAUL), which add stochastic perturbations in primal-dual damping dynamics and Hessian-driven damping dynamics from optimization, and it is proved that GAUL admits the correct stationary distribution, whose marginal is the target distribution.

Abstract

Sampling from a target distribution is a fundamental problem. Traditional Markov chain Monte Carlo (MCMC) algorithms, such as the unadjusted Langevin algorithm (ULA), derived from the overdamped Langevin dynamics, have been extensively studied. From an optimization perspective, the Kolmogorov forward equation of the overdamped Langevin dynamics can be treated as the gradient flow of the relative entropy in the space of probability densities embedded with Wassrstein-2 metrics. Several efforts have also been devoted to including momentum-based methods, such as underdamped Langevin dynamics for faster convergence of sampling algorithms. Recent advances in optimizations have demonstrated the effectiveness of primal-dual damping and Hessian-driven damping dynamics for achieving faster convergence in solving optimization problems. Motivated by these developments, we introduce a class of stochastic differential equations (SDEs) called gradient-adjusted underdamped Langevin dynamics (GAUL), which add stochastic perturbations in primal-dual damping dynamics and Hessian-driven damping dynamics from optimization. We prove that GAUL admits the correct stationary distribution, whose marginal is the target distribution. The proposed method outperforms overdamped and underdamped Langevin dynamics regarding convergence speed in the total variation distance for Gaussian target distributions. Moreover, using the Euler-Maruyama discretization, we show that the mixing time towards a biased target distribution only depends on the square root of the condition number of the target covariance matrix. Numerical experiments for non-Gaussian target distributions, such as Bayesian regression problems and Bayesian neural networks, further illustrate the advantages of our approach.

Paper Structure

This paper contains 25 sections, 33 theorems, 166 equations, 8 figures, 1 table.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Gradient descent, unadjusted Langevin algorithms, and optimal transport gradient flows
  4. Nesterov's ODEs and underdamped Langevin dynamics
  5. Primal-dual damping dynamics and Hessian driven damping dynamics
  6. Gradient-adjusted underdamped Langevin dynamics
  7. Analysis of GAUL on quadratic potential functions
  8. Problem set-up
  9. Continuous time analysis
  10. Discrete time analysis
  11. Numerical experiment
  12. Gaussian examples
  13. One dimension
  14. 20 dimensions
  15. Mixture of Gaussian
  16. ...and 10 more sections

Key Result

Theorem 1.1

Suppose that $f:\mathbb R^d \to \mathbb R^d$ is given by $f({\bm{x}})=\frac{1}{2}{\bm{x}}^T \Lambda {\bm{x}}$ with a symmetric positive definite matrix $\Lambda\in\mathbb{R}^{d\times d}$ with eigenvalues $s_1 \geq s_2 \geq \ldots \geq s_d >0$. Let $\kappa=s_1/s_d$ be the condition number of matrix $

Figures (8)

  • Figure 1: Convergence and density comparisons of three methods. (a) and (c): KL divergence between the sample and the target distribution, which is a one-dimensional Gaussian with zero mean and variance 0.01 (a), 100 (c). 'ol' represents overdamped Langevin dynamics; 'ul' represents underdamped Langevin dynamics. x-axis represents time and y-axis is in $\log_{10}$ scale. (b) and (d): density comparison at the end of the experiment between the three methods and the true density.
  • Figure 2: Convergence and scatter plots. (a)--(d): $h=0.005$. (e)--(h): $h=0.05$. (a) and (e): KL divergence between the sample and target distribution. The x-axis represents time and the y-axis is in $\log_{10}$ scale. Rest panels: scatter plot of the three methods at the end of the experiment for different step sizes. Contours of the true density are also provided for comparisons. In (g) there are no scatter points shown as 'ul' does not converge for this choice of $h$.
  • Figure 3: Convergence and scatter plots. (a): KL divergence between the sample and target distribution, which is a mixture of two unit variance Gaussians located at $(1/2,1/2)$ and $(-1/2,-1/2)$. x-axis represents time and y-axis is in $\log_{10}$ scale. (b)--(d): scatter plot of the three methods a the end of the experiment. Contour of the true density is also provided for comparison.
  • Figure 4: Convergence and scatter plots for mixture of Guassians centered at $(3,3)$ and $(-3,-3)$.
  • Figure 5: Convergence and scatter plots for the quadratic cosine example.
  • ...and 3 more figures

Theorems & Definitions (66)

  • Theorem 1.1: Informal
  • Proposition 2.1: feng2024fisher Proposition 1
  • Proposition 2.2
  • Remark 2.3
  • Remark 2.4
  • Remark 3.1
  • Proposition 3.2
  • Corollary 3.3
  • Theorem 3.4
  • proof
  • ...and 56 more