Source-linked AI summary
Underdamped Langevin MCMC: A non-asymptotic analysis
Xiang Cheng, Niladri S. Chatterji, Peter L. Bartlett, Michael I. Jordan
TL;DR
The paper asks whether underdamped Langevin dynamics can yield faster non-asymptotic sampling for smooth, strongly log-concave targets than overdamped Langevin methods. It analyzes an MCMC algorithm obtained by discretizing the diffusion and proves improved 2-Wasserstein convergence rates, while identifying remaining condition-number and distributional-scope questions.
Problem
The paper studies non-asymptotic sampling for smooth, strongly log-concave target distributions and the quantitative basis for underdamped dynamics outperforming overdamped Langevin MCMC.
Method
The paper analyzes an MCMC algorithm based on discretizing underdamped Langevin diffusion, including bounds controlling discretization error and kinetic energy.
Results
The method achieves ε error in 2-Wasserstein distance in O(√d/ε) steps, improving over overdamped Langevin MCMC's O(d/ε^2) rate under the same assumptions.
Takeaways & Limitations
Underdamped, second-order dynamics provide quantitative support for the faster convergence empirically associated with Hamiltonian MCMC and connect to acceleration methods such as Nesterov acceleration and heavy ball.
Takeaways & Limitations
The analysis leaves open improving condition-number dependence from κ^2 to κ and extending efficient sampling guarantees beyond log-concave distributions.
Abstract
from arXiv · showhide
We study the underdamped Langevin diffusion when the log of the target distribution is smooth and strongly concave. We present a MCMC algorithm based on its discretization and show that it achieves $\varepsilon$ error (in 2-Wasserstein distance) in $\mathcal{O}(\sqrt{d}/\varepsilon)$ steps. This is a significant improvement over the best known rate for overdamped Langevin MCMC, which is $\mathcal{O}(d/\varepsilon^2)$ steps under the same smoothness/concavity assumptions. The underdamped Langevin MCMC scheme can be viewed as a version of Hamiltonian Monte Carlo (HMC) which has been observed to outperform overdamped Langevin MCMC methods in a number of application areas. We provide quantitative rates that support this empirical wisdom.
1 Introduction
The paper studies underdamped Langevin diffusion and an implementable discretization for sampling strongly log-concave targets, providing non-asymptotic convergence guarantees and a quantitative account of its acceleration over overdamped Langevin methods.
- 1 Introduction: The discretized underdamped Langevin diffusion provides an algorithmic way to sample from p*(x) ∝ e^-f(x) without knowing the normalization constant.The continuous process has invariant distribution proportional to exp(-(f(x) + ||v||^2/2u)), whose x-marginal is proportional to exp(-f(x)).
- 1 Introduction: The paper proves convergence of both the continuous diffusion and its discretization to the invariant distribution for log-smooth, strongly log-concave targets.These results yield explicit sampling rates for Algorithm 1.
- 1 Introduction: Underdamped Langevin diffusion includes a Hamiltonian component, and its discretization can be viewed as a form of Hamiltonian MCMC.The method has a physical interpretation as particle dynamics under a force field and drag.
- 1 Introduction: Hamiltonian MCMC has been empirically observed to converge faster than standard Langevin MCMC, the discretization of overdamped Langevin diffusion.The paper aims to provide a non-asymptotic quantitative explanation for this observation.
- 1.1 Related Work: The paper addresses a gap in prior theory: no prior polynomial-in-dimension convergence result was known for HMC under log-smooth or strongly log-concave targets.Earlier related work either treated broader function classes with dimension-exponential scaling or analyzed overdamped Langevin methods.
- 1 Introduction: The analysis studies an exact diffusion, its discrete update, and the Wasserstein framework used to compare their distributions.The discrete update replaces the current position in the velocity drift with the initial position of the step.
2 Results
The paper analyzes a discretized underdamped Langevin algorithm and gives convergence guarantees in 2-Wasserstein distance, including a noisy-gradient setting. Its stated runtime dependence improves over overdamped Langevin diffusion under the corresponding assumptions.
- Main algorithm: Algorithm 1 discretizes the underdamped Langevin diffusion using a Gaussian transition obtained by integrating the diffusion for step size δ.The analyzed choice uses γ = 2 and u = 1/L.
- Main result: Theorem 1 provides a convergence guarantee for Algorithm 1 from a point-mass initialization at (x(0), 0).The theorem parameterizes the guarantee by the initial distance to the optimum and the number of iterations.
- Main result: The resulting dependence on dimension and accuracy is reported as a significant improvement over the runtime of overdamped Langevin diffusion.The comparison is made with the corresponding overdamped runtime cited in [4].
- Main result: A time-varying step size can remove the logarithmic factor in the stated bound.The authors defer this refinement to Theorem 14 and do not optimize constants in either theorem.
- Result with stochastic gradients: With noisy gradients, the paper assumes an additive independent noise vector with zero mean and analyzes the resulting stochastic-gradient algorithm.The noisy-gradient dynamics are formulated through a separate SDE and Algorithm 2.
- Result with stochastic gradients: When gradient-noise variance is large, the noisy-gradient rate recovers the overdamped Langevin rate and requires ˜O(σ^2κ^2d/ε^2) steps for ε accuracy in W2.This comparison is stated explicitly in Remark 4.
3 Convergence of the Continuous-Time Process
The continuous-time analysis establishes contraction of the underdamped Langevin diffusion through synchronous coupling. This contraction yields exponential convergence to the stationary distribution in W2.
- Contraction proof: Theorem 5 chooses u = 1/L and γ = 2 and establishes a coupling between the evolved distributions from arbitrary initial points.Here L is the smoothness parameter of f.
- Convergence consequence: The continuous-time process converges exponentially to the stationary distribution in W2.Corollary 7 transfers the contraction through the map g(x, v) = (x, x + v).
- Contraction proof: The proof uses synchronous coupling, sharing Brownian motion between two processes so their Brownian terms cancel.The coupled difference variables are the position and velocity differences.
- Contraction proof: Strong convexity and smoothness bound the Hessian eigenvalues by 0 < m ≤ Λ_j ≤ L, enabling the matrix-based contraction argument.The proof diagonalizes the Hessian and analyzes the resulting characteristic equations.
- Contraction proof: The contraction estimate is converted into a convergence rate using Grönwall’s inequality.The argument works with a quadratic objective and its associated matrix inequalities.
4 Discretization Analysis
The discretization analysis bounds the W2 error between the continuous and discrete underdamped Langevin processes over one step. The result relies on a bounded kinetic-energy assumption and synchronous coupling.
- Discretization bound: Theorem 9 bounds W2(Φδp0, ˜Φδp0) between the continuous-time and discrete-time processes initialized from the same distribution.The analysis assumes δ ≤ 1 and uses u = 1/L and γ = 2.
- Proof strategy: The proof couples the two processes through the same initial distribution and common Brownian motion.Velocity and position errors are bounded separately.
- Assumptions: The discretization estimates use L-smoothness of f and a uniform upper bound on the continuous-time kinetic energy.The kinetic-energy bound is supplied by Lemma 12.
- Discretization bound: The resulting upper bound is obtained for δ smaller than 1, and taking square roots gives the desired Wasserstein estimate.The position estimate is combined with the preceding velocity estimate.
5 Proof of Theorem 1
The proof of Theorem 1 combines continuous-time contraction with the one-step discretization bound. A geometric-series argument then controls the accumulated error over n iterations.
- Combining bounds: The proof starts from the continuous-time convergence estimate and applies it at each iterate of the discretized chain.This uses the bound from Corollary 7.
- Combining bounds: Theorem 9 and the sandwich inequality bound the discrepancy introduced by discretization.The triangle inequality for W2 is then used to combine contraction and discretization terms.
- Iteration argument: Defining η = e^(-δ/2κ), repeated application of the one-step inequality produces a geometric series over n iterations.The proof bounds the resulting terms using the upper bound from the sandwich inequality.
- Final guarantee: The step size is selected so that the contraction and discretization contributions are each at most ε/2, yielding total W2 error ε.The proof uses δ/κ < 1 and a kinetic-energy bound.
6 Conclusion
The paper establishes convergence guarantees for underdamped Langevin diffusion and its discretization, while connecting the method to Hamiltonian MCMC and second-order optimization dynamics.
- The algorithm provides 2-Wasserstein convergence guarantees for sampling strongly log-concave distributions using underdamped Langevin diffusion.The analysis covers both the continuous-time process and its discretization.
- An open question is whether the condition-number dependence can improve from κ^2 to κ, alongside extensions to non-log-concave sampling and broader second-order Langevin equations.The paper also notes that lower bounds are largely unknown in the MCMC field.
- The continuous and discrete processes admit explicit Gaussian representations whose coordinates can be sampled independently in time linear in d.The Gaussian moments are obtained from integral representations of the processes.
B Controlling the Kinetic Energy
The kinetic-energy analysis bounds the velocity contribution under specified initialization and uses that bound to control discretization error and establish convergence estimates.
- The kinetic-energy bound is introduced to control discretization error at each step.
- Under Dirac initialization at (x(0), 0), the analysis assumes a bounded initial distance from the optimum.This initialization and distance condition are used in the kinetic-energy argument.
- The kinetic-energy estimate is EK = 26(d/m + D2).
- The proof bounds kinetic energy for an arbitrary distribution by comparing it with the optimum distribution through an optimal coupling.Young’s inequality is used repeatedly in the resulting inequalities.
- The resulting bounds are propagated across iterations by induction and Wasserstein-distance contraction estimates.
C Varying Step Size
A varying-step-size schedule reduces the logarithmic factor in the convergence guarantee by progressively halving the step size and doubling the iteration count.
- An adaptive step size removes the log factor appearing in Theorem 1.
- The initial step size and iteration count are selected under Dirac initialization with initial distance to optimum below ε0.
- The method uses ℓ epochs, with each new step size halved and its iteration count doubled.The epoch count is chosen as ℓ = ⌈log(ε0/ε)/log(2)⌉.
- The total computational cost is the sum of the iteration counts across all epochs.
- After each epoch, the Wasserstein error is reduced geometrically, with the final guarantee W2(p(n), p*) ≤ ε0/2^ℓ < ε.
D Analysis with Stochastic Gradients
The stochastic-gradient extension replaces exact gradients in the underdamped Langevin discretization while retaining Gaussian update structure derived from the stochastic diffusion.
- The paper states an underdamped Langevin MCMC algorithm using stochastic gradients under the assumptions from Section 2.2.1.
- The update noise conditioned on the current state has a Gaussian distribution with computable conditional mean and covariance.
- The update distribution is obtained by integrating the discrete stochastic underdamped Langevin diffusion for time δ with γ = 2 and u = 1/L.
- The stochastic-gradient derivation mirrors the exact-gradient calculation after replacing ∇f(·) with stochastic gradients ∇̂f(·).The resulting transition is represented by p(i+1) = Φ̂δp(i).
D.1 Discretization Analysis
The section analyzes discretization errors by comparing discrete processes with and without noisy gradients, then combines these bounds to establish the theorem. The proof controls multiple error terms at level ε/3 and selects δ accordingly.
- Discretization Analysis: The analysis bounds the discretization error between the discrete process without gradient noise and the process with noisy gradients.Both processes are considered from the same initial distribution, and Lemma 16 formalizes their comparison.
- Discretization Analysis: The noisy-gradient difference is analyzed using a zero-mean variance-bounded random variable independent of the Brownian motion.The variance is bounded by σ^2d, and optimal couplings with the target distribution are used in the argument.
- D.1 Discretization Analysis: The proof combines discretization bounds, the sandwich inequality, the triangle inequality for W2, and a recursive-sequence bound.These ingredients connect intermediate process comparisons to the final theorem.
- D.1 Discretization Analysis: The error terms are controlled separately at level ε/3, with δ chosen under conditions involving ε, κ, and an upper bound on E K.The argument uses 1 − e^−δ/2κ ≥ δ/4κ when δ/κ < 1.
- D.1 Discretization Analysis: The stated parameter choices complete the claim.The section concludes after establishing the required bounds.
E Technical Results
The technical-results section collects external theorems and standard lemmas used in the paper’s proofs. These results cover prior bounds, block-matrix determinants, and recursive inequalities.
- E Technical Results: Theorem 17 is imported from reference [4] and is used in the proof of Lemma 12.It is stated for all t ≥ 0 and x ∈ R^d.
- E Technical Results: A standard block-matrix determinant lemma is applied in the proof of Theorem 5.The stated condition is that the square matrices C and D commute.
- E Technical Results: Lemma 19 from reference [7] provides a bound for a non-negative sequence satisfying a recursive inequality.The lemma assumes non-negative numbers A, B, and C, with A ∈ {0, 1}, and applies for all integers k ≥ 0.