Source-linked AI summary
Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem
Andre Wibisono
TL;DR
The paper asks how to obtain an implementable, consistent discretization of Langevin dynamics despite ULA’s persistent discrete-time bias. It analyzes sampling as composite optimization in measure space, proposes SLA as a symmetrized alternative, and shows SLA is consistent for Gaussian targets while ULA is not.
Problem
ULA is implementable but biased for every fixed positive step size, while unbiased Forward-Backward discretization is generally not implementable because the heat flow’s forward and backward methods are unavailable.
Method
The paper analyzes Langevin dynamics as composite optimization of relative entropy and proposes symmetrizing the Forward-Flow discretization to obtain SLA.
Results
SLA is consistent for Gaussian target measures, whereas ULA remains biased; SLA also converges exponentially fast in the Gaussian case.
Takeaways & Limitations
Symmetrization can reduce or remove Langevin discretization bias, but the method requires a proximal gradient step in measure space and may involve numerical computation each iteration.
Takeaways & Limitations
The unbiased Forward-Backward approach is generally unavailable because neither the forward nor backward method for the heat flow is implementable, with explicit solvability limited to Gaussian initial data.
Abstract
from arXiv · showhide
We study sampling as optimization in the space of measures. We focus on gradient flow-based optimization with the Langevin dynamics as a case study. We investigate the source of the bias of the unadjusted Langevin algorithm (ULA) in discrete time, and consider how to remove or reduce the bias. We point out the difficulty is that the heat flow is exactly solvable, but neither its forward nor backward method is implementable in general, except for Gaussian data. We propose the symmetrized Langevin algorithm (SLA), which should have a smaller bias than ULA, at the price of implementing a proximal gradient step in space. We show SLA is in fact consistent for Gaussian target measure, whereas ULA is not. We also illustrate various algorithms explicitly for Gaussian target measure, including gradient descent, proximal gradient, and Forward-Backward, and show they are all consistent.
1. Introduction
The paper frames Langevin sampling as optimization in measure space and asks whether an implementable discretization can remain consistent while converging exponentially under LSI. It explains ULA’s bias through composite optimization and proposes symmetrization to reduce that bias.
- Langevin dynamics is the Wasserstein gradient flow of relative entropy and attracts initial distributions to the stationary target measure.
- The paper asks whether an implementable Langevin discretization can be consistent and converge exponentially fast under the logarithmic Sobolev inequality.
- ULA remains asymptotically biased for every fixed positive step size, even with a Gaussian target measure.
- The bias is attributed to ULA’s Forward-Flow discretization of relative entropy, a composite objective combining negative entropy and expected function value.
- Forward-Backward is consistent in principle, but its backward heat-flow method is generally not implementable, except for Gaussian initial data.
- The proposed symmetrized Forward-Flow method yields SLA, whose order-2 construction should reduce bias while requiring a proximal gradient step in measure space.
2. Sampling as optimization in the space of measures
The paper formulates sampling as optimization of relative entropy over measures and analyzes Langevin discretizations through composite optimization. It identifies ULA’s bias and proposes symmetrization, obtaining consistency for Gaussian targets while exposing implementation limits in general.
- Sampling as optimization: Sampling can be recast as minimizing relative entropy over probability measures, whose unique minimizer is the target distribution.The Wasserstein gradient flow of relative entropy is the Fokker–Planck equation, implemented in space by Langevin dynamics.
- Langevin dynamics and ULA: Under the logarithmic Sobolev inequality, continuous-time Langevin dynamics converges exponentially fast, but ULA has a nonzero asymptotic bias for every fixed step size.ULA converges to the Langevin dynamics as the step size tends to zero, yet even Gaussian targets exhibit bias at positive step size.
- ULA as composite optimization: ULA is the Forward-Flow discretization of relative entropy’s composite decomposition, and its bias arises because the flow step is not the adjoint of the forward step.The two components are gradient descent for expected function value and exact gradient flow for negative entropy.
- Forward-Backward method: Forward-Backward optimization preserves the minimizer through an adjoint backward method and is therefore consistent, with exponential convergence under suitable domination and smoothness assumptions.For sampling, the backward component corresponds to a proximal gradient step in measure space.
- Symmetrized Langevin algorithm: Symmetrizing Forward-Flow with its adjoint yields SLA, a second-order method whose Gaussian-target iteration is consistent and converges exponentially fast.For general targets, SLA requires solving a proximal gradient step numerically and may still have bias of order 2.
- Limitations and higher-order methods: Higher-order bias reduction is constrained because negative step sizes would require backward heat flow, which is generally non-implementable and mathematically ill-posed.The paper notes that this obstruction may potentially be circumvented for entropy-specific algorithms.
3. Optimization in the space of measures
The paper formulates sampling as optimization over probability measures, identifying Langevin dynamics with relative-entropy gradient flow. It develops measure-space optimization methods and examines their convergence and implementability, including Gaussian cases.
- Optimization takes place on probability measures with finite second moments, equipped with the Wasserstein metric.
- Expected-value functionals can be implemented from spatial samples, while negative-entropy gradient flow is the heat flow driven by Brownian motion.
- Relative entropy combines expected function value and negative entropy, is nonnegative, and is minimized uniquely by the target measure.
- Under the logarithmic Sobolev inequality, continuous-time Langevin dynamics converges exponentially fast; this includes strongly log-concave and some multimodal measures.
- ULA's forward-flow discretization is difficult to iterate because the next density is unavailable, whereas the corresponding backward method is generally implicit and analytically unsolvable.
- For Gaussian target measures with Gaussian initial data, the forward and backward methods preserve Gaussianity and are consistent.
4. Langevin dynamics as composite optimization in the space of measures
The paper treats relative entropy as a composite optimization objective and applies forward-backward or backward-forward compositions to its potential and entropy components. These unbiased constructions are generally not implementable, but become consistent for Gaussian data.
- Relative entropy is optimized as a sum of expected function value and negative entropy using composite optimization.
- Forward-backward optimization applies the forward method to expected function value and the backward method to negative entropy.
- Forward-backward is consistent because the backward method is adjoint to the forward method, and it can converge exponentially under gradient domination and smoothness assumptions.
- The forward-backward and backward-forward algorithms are not generally implementable because the required heat-flow method cannot be solved explicitly.
- For Gaussian targets and Gaussian initial data, both forward-backward and backward-forward preserve Gaussianity and are consistent.
5. Discussion and future work
The paper identifies ULA's bias and the non-implementability of unbiased composite methods as central obstacles to Langevin sampling. It proposes a symmetrized alternative and points to more advanced optimization-based sampling methods.
- ULA is implementable but biased, while unbiased forward-backward algorithms are not implementable because heat-flow forward and backward methods are unavailable in general.
- The symmetrized Langevin algorithm is proposed to reduce ULA's bias, at the cost of implementing a proximal gradient step in space.
- The paper suggests applying acceleration and related optimization techniques to sampling, including underdamped Langevin dynamics as a second-order analogue.
A.1. Proof of Lemma 1 (Contraction of ULA)
The proof analyzes ULA contraction by synchronously coupling two chains with the same Gaussian noise and bounding their squared-distance evolution under strong convexity and smoothness.
- Synchronous coupling runs two ULA chains with the same Gaussian noise.
- Subtracting the coupled updates and taking squared norms yields a one-step distance bound.
- The bound uses strong convexity and smoothness of f to control the gradient-difference term.
- When the relevant step-size condition makes the second term nonpositive, it can be dropped before taking expectations.
- Optimal coupling and the Wasserstein-distance definition convert the expected coupled distance bound into the desired Wasserstein claim.
A.2. Proof of Lemma 2 (Bias of ULA)
The proof invokes the noiseless case of a prior theorem.
- The result follows from the noiseless case δ = σ = 0 of Dalalyan and Karagulyan (2017, Theorem 4).
A.3. Proof of Lemma 3 (Contraction of SLA)
This section reviews integrators, adjoints, and symmetrization before applying these ideas to SLA and Gaussian-mixture examples.
- For Gaussian mixtures, SLA requires solving a one-dimensional equation for the projected variable before recovering the full update.
- B.2. The forward method: Forward and backward methods are first-order integrators, with the backward method equal to the forward method’s adjoint.
- Symmetrizing basic methods produces second-order integrators, including the trapezoid rule and implicit midpoint rule.
B.5. On order of error and order of bias
Under strong convexity, discretization order controls the asymptotic bias of an integrator, while contraction yields convergence to a biased limit.
- The discretization error and bias have the same order under strong convexity.
- Strong convexity makes the continuous flow contract exponentially and supports convergence of the discretized iteration to a biased limit.
- An order-p integrator has bias O(ǫ^p) under the stated monotonicity condition.
- Gradient domination yields exponential convergence of gradient flow in function value, while discrete gradient descent additionally requires smoothness.
- Proximal gradient avoids the smoothness requirement but requires solving an implicit update.
C.5. Symmetrized forward method
The symmetrized forward method combines gradient descent and proximal gradient as a Forward-Backward scheme for a self-decomposition of the objective.
- The symmetrized forward method composes gradient descent with the proximal gradient method.
- This composition is the Forward-Backward algorithm applied to the self-decomposition 2f = f + f.
- Under gradient domination and smoothness, the method converges exponentially in function value.
- Its convergence rate is slightly better than gradient descent under the stated conditions.
Appendix D. A brief review of the Wasserstein metric in the space of measures
The Wasserstein framework treats probability measures as an optimization space, where expected-value and entropy functionals admit gradient-based methods. It also characterizes heat flow and proximal operations, establishing convergence and implementability properties under stated assumptions.
- Wasserstein geometry: The Wasserstein metric W2 defines an infinite-dimensional Riemannian structure on probability measures with finite second moments.It is defined through an infimum over couplings of random variables with prescribed marginals.
- Expected-value optimization: For expected-value functionals F(ρ)=Eρ[f], the measure-space gradient, gradient domination, and strong convexity correspond to analogous properties of f.Strong convexity gives a unique point-mass minimizer, while gradient domination yields exponential convergence.
- Expected-value optimization: Gradient flows and discretizations for expected-value functionals can be implemented through corresponding algorithms applied directly in space.Under gradient domination, convergence is exponential; under strong convexity, coevolving solutions contract exponentially.
- Proximal operations: The proximal step for f is represented in measure space by a pushforward of the corresponding proximal map.The construction follows from the optimality condition for a Wasserstein-regularized expected-value functional.
- Negative entropy and heat flow: Negative entropy generates the heat equation, whose exact flow is implemented by adding independent Gaussian noise and is convex in measure space.Its gradient norm is Fisher information, and its entropy evolution satisfies De Bruijn’s identity.
- Negative entropy and heat flow: Negative entropy is convex but generally not strictly convex, has geodesic strong convexity under mean preservation, and has an unbounded Hessian above.Along heat flow, Gaussian data yield logarithmic entropy decrease, while variance grows linearly.
E.2.3. FORWARD METHOD FOR HEAT FLOW
The forward method for heat flow uses the score of the current density, but its iteration is generally not continuable because the next density’s score is unavailable. Gaussian data permit an explicit iteration whose variance grows too quickly.
- Forward method: The forward heat-flow method advances a measure using its score field ∇log ρ.Under K-log-semiconcavity, the step is subject to a step-size restriction.
- Forward method: The method generally cannot continue beyond one step because the analytic score of the updated density is unknown.Knowing the current score allows one update, but not the score required for the next update.
- Gaussian case: For Gaussian initial data, the forward method preserves Gaussianity but increases variance faster than the exact heat flow.The covariance update satisfies Σk+1=Σk(I+ǫΣk^-1)^2, exceeding the heat-flow growth at t=ǫk.
E.2.4. BACKWARD METHOD FOR HEAT FLOW
The backward heat-flow method is implicit and generally not analytically solvable, though Gaussian cases are explicit and grow more slowly than the heat flow. Composite optimization explains why adjoint forward-backward pairings avoid bias.
- Backward method: The backward method requires an implicit update and is generally not solvable analytically.For Gaussian initial data it is explicit, with a step-size condition tied to the smallest covariance eigenvalue.
- Gaussian case: For Gaussian data, backward heat flow preserves Gaussianity and produces variance growth slower than the exact heat flow.This contrasts with the forward method, whose variance grows faster.
- Composite optimization: Composite optimization alternates forward gradient descent on one term with a backward proximal step on the other.The minimizer is a fixed point because the backward step inverts the negative forward step.
- Composite optimization: Pairing adjoint methods preserves the true minimizer, whereas other pairings can converge to a biased point.Both Forward-Backward and Backward-Forward are unbiased in the quadratic example, while the other tested combinations are biased.
- Gaussian OU case: Forward-Backward and Backward-Forward Langevin schemes are consistent for Gaussian OU targets.Their means converge exponentially to the target mean, and the target covariance is the unique fixed point.