PaperlayerView exact revision on arXiv

Source-linked AI summary

Shifted Composition IV: Toward Ballistic Acceleration for Log-Concave Sampling

Jason M. Altschuler, Sinho Chewi, Matthew S. Zhang

arXiv:2506.23062v3math.PRcs.DSmath.APmath.NAmath.ST

TL;DR

The paper addresses whether accelerated sampling is possible for general non-quadratic κ-conditioned potentials, where prior work had not established any sublinear dependence on κ. It develops a KL local-error framework for discretized underdamped Langevin dynamics and obtains Õ(κ^5/6) dependence, while improving dimension dependence to d^1/3 remains a potential direction.

  • Problem

    Prior to this paper, it was unknown whether any sublinear rate κ^(1−δ) was possible for sampling with general non-quadratic κ-conditioned potentials.

  • Method

    The paper extends a shifted-chain-rule and auxiliary-process KL local-error framework to analyze discretizations of degenerate underdamped Langevin dynamics.

  • Results

    Õ(κ^5/6) iteration complexity for arbitrary κ-conditioned potentials establishes sublinear condition-number dependence, while the result scales as O(κ^5/6d^5/3/ε^2/3) and improves momentum-coordinate dependence by a factor of h^2.

  • Takeaways & Limitations

    The result answers affirmatively whether accelerated dependence on κ is possible for general potentials and opens the doorway to ballistic acceleration.

  • Takeaways & Limitations

    The low-friction setting is excluded because underdamped Langevin dynamics is non-contractive there, and the analysis relies on contractivity.

Abstract

from arXiv · show

Acceleration is a celebrated cornerstone of convex optimization, enabling gradient-based algorithms to converge sublinearly in the condition number. A major open question is whether an analogous acceleration phenomenon is possible for log-concave sampling. Underdamped Langevin dynamics (ULD) has long been conjectured to be the natural candidate for acceleration, but a central challenge is that its degeneracy necessitates the development of new analysis approaches, e.g., the theory of hypocoercivity. Although recent breakthroughs established ballistic acceleration for the (continuous-time) ULD diffusion via space-time Poincare inequalities, (discrete-time) algorithmic results remain entirely open: the discretization error of existing analysis techniques dominates any continuous-time acceleration. In this paper, we give a new coupling-based local error framework for analyzing ULD and its numerical discretizations in KL divergence. This extends the framework in Shifted Composition III from uniformly elliptic diffusions to degenerate diffusions, and shares its virtues: the framework is user-friendly, applies to sophisticated discretization schemes, and does not require contractivity. Applying this framework to the randomized midpoint discretization of ULD establishes the first ballistic acceleration result for log-concave sampling (i.e., sublinear dependence on the condition number). Along the way, we also obtain the first $d^{1/3}$ iteration complexity guarantee for sampling to constant total variation error in dimension $d$.

1 Introduction

Underdamped Langevin dynamics were conjectured to accelerate sampling, but continuous-time acceleration had not translated to algorithms because discretization errors overwhelmed the gain. This paper develops a KL-based coupling framework and applies it to randomized midpoint ULD, obtaining sublinear condition-number dependence and improved dimension guarantees.

  • Motivation: Space-time Poincare inequalities had established ballistic acceleration for continuous-time ULD, but discrete-time acceleration remained open because existing analyses made discretization error dominate.Underdamped dynamics were motivated by conjectured √κ dependence and faster mixing from non-reversible momentum lifts.
  • KL local error framework: The paper introduces a coupling-based KL local-error framework that extends Shifted Composition III from elliptic to degenerate diffusions and accumulates discretization errors linearly without requiring contractivity.It supports sophisticated schemes such as randomized midpoint discretization and is designed for both continuous-time and discrete-time analysis.
  • Applications to sampling: O(κ5/6d5/3/ε2/3) gradient evaluations yield an ε2-close KL sample for strongly log-concave, log-smooth targets, establishing sublinear dependence on κ for arbitrary κ-conditioned potentials.The result uses RM–ULMC initialized with log χ2(µ0 ∥π) =̃ O(d).
  • Applications to sampling: The framework produces O(κd1/3/ε2/3) KL complexity, matching prior W2 dependence while strengthening the guarantee to KL and therefore total variation.This rate is stated for the strongly convex case and is distinct from the accelerated result's dimension dependence.
  • Applications to sampling: For constant total variation error, the method achieves Õ(d1/3) dimension dependence, improving on prior Õ(d1/2) and Õ(d5/12) guarantees.The comparison is made against prior results [FYC23] [AC24b] and [KN24].
  • Techniques: The analysis extends a shifted-composition approach based on an auxiliary process and Girsanov-type energy control to compare exact and discretized degenerate diffusions.This construction is suited to accumulating numerical errors while retaining KL guarantees.

2 Preliminaries

The preliminaries introduce Wasserstein and Rényi divergences, then develop the shifted chain rule used to compare processes evolving under different kernels.

  • Wasserstein distance is defined through couplings, and the analysis may use a twisted norm because ULD contracts in that norm rather than the standard Euclidean norm.
  • Rényi divergence generalizes KL divergence through an order parameter q > 1, with q = 1 interpreted as KL and q = 2 related to chi-squared divergence.
  • The data processing inequality states that applying a Markov kernel cannot increase Rényi divergence.
  • The shifted chain rule bounds KL divergence between output laws using an auxiliary random variable X′ and a coupling-dependent conditional KL term.
  • Introducing X′ generalizes the standard chain rule and enables analysis of two Markov processes that use different update kernels.

3 Harnack inequalities for the underdamped Langevin diffusion

This section establishes reverse transport inequalities for ULD, dual to parabolic Harnack inequalities, under Hessian bounds covering strongly convex, weakly convex, and semi-convex potentials. The bounds capture sharp short-time behavior and distinguish exponential, polynomial, and absent long-time decay across friction and convexity regimes.

  • Main result: Theorem 3.2 establishes reverse transport, or Harnack, inequalities for ULD under α-convexity and β-smoothness assumptions.The result applies for all T > 0 and separates high- and low-friction regimes through the parameter ω.
  • Short-time behavior: For short times, the bounds exhibit T^-3 scaling in position and T^-1 scaling in momentum, reflecting the degenerate hypoelliptic structure of ULD.These scalings are sharp, including for quadratic potentials.
  • Long-time behavior: For large T, strongly convex, weakly convex, and semi-convex or low-friction settings yield respectively exponential decay, polynomial 1/T decay, and no decay.The long-time behavior changes continuously with α at the weak-convexity limit.
  • Proof strategy: The coupling constructs an auxiliary process that interpolates between two ULD trajectories and reaches the target at time T, with Girsanov’s theorem converting its shift energy into the inequality.Time-dependent twisted coordinates and effective friction enable contraction even in non-convex or otherwise non-contractive settings covered by the analysis.
  • Shift construction: The method uses time-dependent shifts selected to make the Girsanov KL bound small, and a simplified zero-potential example shows the resulting bound can be exact and sharp.The shifts are chosen explicitly rather than obtained from a closed-form optimization in the full ULD setting.

4 KL local error framework for underdamped Langevin discretizations

The paper develops a KL-divergence local error framework for discretizations of degenerate ULD, extending prior work beyond elliptic diffusions. Its master bound combines initialization, local position and momentum errors, and cross-regularity without requiring a single discretization kernel to satisfy the latter directly.

  • The framework extends the shifted-composition analysis from elliptic to degenerate diffusions and targets KL discretization error.
  • The master theorem bounds final KL error using initialization distance, local errors, and a cross-regularity term.The local-error contribution includes both position and momentum coordinates, with position errors assumed smaller by a factor of h for the analyzed schemes.
  • Cross-regularity is handled by replacing the final numerical step with ULMC, for which the required bound is easier to establish.This avoids proving cross-regularity for an arbitrary discretization kernel.
  • The framework distinguishes strongly convex, weakly convex, and semi-convex regimes, with continuous behavior as the convexity parameter approaches zero.The strongly convex bound is optimized for long times, while the weakly convex bound captures shorter-time behavior.
  • The analysis constructs an auxiliary shifted process that follows ULD while being coupled to the numerical algorithm through integrated shifts.The proof uses time-varying twisted coordinates and a contractive distance recursion to control the auxiliary distances.
  • The short-time coupling analysis exploits that initial momentum discrepancies affect the bound by a factor h^2 less than initial position discrepancies.

5 Application to sampling

The framework is applied first to ULMC and then to randomized midpoint ULD discretization. The randomized midpoint method yields sublinear condition-number dependence and the first d^1/3 total-variation dimension dependence at constant accuracy.

  • ULMC: ULMC provides a simpler warm-up application and yields a convergence guarantee under weak log-convexity, while its strongly convex rate matches known results.
  • Randomized midpoint discretization: Theorem 5.6 establishes the randomized midpoint method as the paper’s main application for convex sampling.It analyzes RM–ULMC together with a final ULMC step under strong and weak convexity.
  • Randomized midpoint discretization: RM–ULMC’s double-midpoint implementation removes the second term appearing in earlier rates of the form Õ(κ(d/ε^2)^1/3 + κ^7/6(d/ε^2)^1/6).
  • Randomized midpoint discretization: The first d^1/3 iteration-complexity guarantee for constant total variation error is obtained in the strongly convex, constant-condition-number regime.The result improves the prior d^5/12 dependence in that regime.
  • Randomized midpoint discretization: The strongly convex RM–ULMC rate κd^1/3/ε^2/3 is established in KL divergence and therefore also in total variation via Pinsker’s inequality.The κd^1/3/ε^2/3 dependence had previously been established in W2, while the KL guarantee is new here.
  • Randomized midpoint discretization: For weakly convex targets, the RM–ULMC rate is new and compares favorably with the rates listed for alternative algorithms.
  • Space-time Poincaré inequality: Under a Poincaré inequality, RM–ULMC achieves κ^5/6 dependence up to logarithmic factors, giving the first discretization bound with o(κ) dependence.For standard strongly log-concave initialization the rate is κ^5/6d^5/3/ε^2/3, while a warm start would reduce the dimension factor to d^1/3.

A.1 Proof of Example 3.7

The proof constructs an auxiliary process whose momentum drift is controlled so that its position and momentum marginals satisfy endpoint interpolation constraints. Optimizing this control yields the KL bound used in the continuous-time analysis.

  • The auxiliary process may add drift only to momentum, preserving absolute continuity of its path measure relative to ULD.
  • The control is chosen by minimizing the Girsanov KL bound subject to position and momentum interpolation constraints.
  • The optimal shift separates naturally into position and momentum components and produces the stated KL-divergence bound.
  • The resulting endpoint distributions are Gaussian, allowing the bound to be verified directly from the KL formula for Gaussians.

A.2 Proof of Lemma 3.14

The proof verifies the modified shift bounds across friction and convexity regimes. It then establishes that the auxiliary process reaches the target endpoint and recovers the long-time decay behavior required by the lemma.

  • The proof checks the required lower bounds separately in high-friction and low-friction regimes, including positive, zero, and negative convexity.
  • Strongly convex case: For strongly convex dynamics, the long-time bound decays exponentially as exp(−c|ω|T).
  • Strongly convex case: When the time horizon is shorter than |ω|^-1, the proof uses the weakly convex bound before switching to the long-time strongly convex estimate.
  • The endpoint limit is justified by showing the auxiliary distance vanishes as the remaining time δ approaches zero.

B.1 Proof of Lemma 4.6

The proof reduces Lemma 4.6 to controlling the error quantity ε_t, chiefly by bounding φ_t, N_t, and M_t plus skew terms under step-size and parameter restrictions.

  • The proof reduces Lemma 4.6 to showing ε_t ≤ 1/2, using earlier Lemmas 3.8 and 4.5.
  • For sufficiently small c0/A, Lemma B.1 controls φ_t and makes ∥φ_t − I∥op at most 1/2.
  • Under h ≲ γ^-1 ∧ γ/β and sufficiently small c0/A, Lemma B.2 bounds the perturbation term by ∥N_t∥op ≪ η_t.
  • The proof of Lemma B.3 separately controls the remaining terms, using M_t ⪰ 0 to restrict attention to upper eigenvalues.
  • When γ ≳ η_t, the resulting overall bound can be written β/γ_t + η_t.

B.2 Proof of Lemma 4.8

Lemma 4.8 bounds the coupling constant by combining the earlier error representation with Lemmas B.1 and 4.6, then applying Lemma B.3.

  • The expression from Lemma 4.5, valid for t ∈ [t−, t+], is combined with Lemmas B.1 and 4.6 to control the relevant terms.
  • Applying Lemma B.3 yields a coupling constant of O(βh/γ_t− + hη_p_t+).

B.3 Integral estimates

The integral estimates bound the main KL-error terms by splitting time into parameter-dependent regimes and treating convexity and friction cases separately.

  • Integral estimates: The auxiliary distance bounds distinguish strongly convex high-friction, weakly convex high-friction, and semi-convex low-friction regimes.
  • Regime decomposition: The analysis splits integrals over [0,T0], [T0,T1], and [T1,t] in the convex high-friction regime, and over [0,T0] and [T0,t] in the semi-convex low-friction regime.
  • Error bound: Across the considered regimes, the terminal squared error is bounded by h^2(Ē_w + Ē_s)^2.
  • Strongly convex, high friction: In the strongly convex high-friction case, the bound scales as 1/(αh^2)(Ē_w)^2 plus [1/(β^1/2h) log(1/|ω|h)](Ē_s)^2.
  • Weakly convex, high friction: In the weakly convex high-friction case, the bound scales as T/(β^1/2h^2)(Ē_w)^2 plus a logarithmic Ē_s term.
  • Theorem 4.1: Theorem 4.1 combines Lemmas 4.10 and B.4–B.5, while the remaining initial-distance contribution is compared against C(α,β,γ,T).

C.1 Recursive error control

The recursive error-control analysis transfers stationary error bounds to algorithm iterates under Lipschitz assumptions, then combines them with Theorem 4.1 for convex and non-convex settings.

  • Recursive error control: The framework controls iterate-based errors through auxiliary-process errors when the errors satisfy Assumption C.1’s Lipschitz conditions.
  • Change of measure: Change of measure converts expectations under the auxiliary process into expectations at stationarity, using Lemma C.3 under bounded Hessian assumptions.
  • Recursive bounds: Lemma C.4 provides recursive control of the error quantities, including bounds of the form (·)^2 ≲ d + B^2.
  • Convex case: For convex applications, the analysis checks Assumption C.1 and local-error bounds before applying Theorem 4.1 through Lemma C.2 over T = Nh.
  • Convex case: When β^2h^3/γ ≲ 1, the resulting step-size conditions simplify to h ≲ 1/β^1/2 for γ ≍ √β and h ≲ 1/(β^1/2κ^1/6) for γ ≍ √α.
  • Non-convex case: In the non-convex case, the analysis avoids the nonvanishing C(α,β,γ,T) term by initializing the auxiliary ULD process at the same distribution as the algorithm.
  • Final combination: The discretization bound is combined with convergence results ν_N → π for ULD.

C.2 Proofs for Subsection 5.1

The proofs verify the Lipschitz-error assumption for ULMC in strongly and weakly convex settings, then derive case-specific step-size and time-horizon conditions for the KL guarantee.

  • ULMC error verification: ULMC satisfies Assumption C.1 in both strongly and weakly convex cases under high friction, with step-size constraints h ≲1/(β1/2κ) and h ≲1/(βT), respectively.These conditions follow by controlling the relevant strong and weak error terms.
  • Strongly convex case: In the strongly convex case, the proof uses h ≲1/(β1/2κ1/2) to control the theorem bound before selecting T and h to make the error terms small.The strong-error contribution is ignored using the Jensen bound ¯Ew ⩽¯Es.
  • Weakly convex case: T ≍β1/2W 2/ε2 makes the first weakly convex error term at most ε2, after which h is chosen to control the second term.The proof obtains this by combining Theorem 4.1 with Lemmas C.2 and C.6.

C.3 Proofs for Subsection 5.2

This section establishes local movement and midpoint bounds for RM–ULMC, explains the double-midpoint design, and derives step-size conditions ensuring the required Lipschitz-error assumption.

  • Movement bounds: The ULD movement bound controls the process over [0,h] under h ≲1/β1/2 and γ ≲√β, providing the input for the midpoint estimate.The resulting gradient-deviation bound is independent of the choice of u.
  • Midpoint coupling: The midpoint analysis synchronously couples RM–ULMC and ULD with the same Brownian motion and bounds their discrepancy for every u ∈[0,1].The bound also applies when u is sampled from the distribution specified in (5.1).
  • Lipschitz-error verification: RM–ULMC satisfies Assumption C.1 in all cases when h obeys regime-dependent constraints derived from the strong and weak local errors.The proof compares one-step strong and weak errors and then applies the midpoint and movement lemmas.
  • Double midpoint design: The double-midpoint implementation chooses the random midpoint law to satisfy weak-error matching while using c(u)=(1−e−γh)/γ for strong-error control.The weak-error condition requires averaging the force at the second midpoint under the law in (5.1).
  • Step-size regimes: The resulting constraints include h ≲1/(β1/2κ1/3) in the strongly convex high-friction case and h ≲1/(β2/3T1/3) ∧ 1/(β3/4T1/2) in the weakly convex case.Additional conditions arise from neglecting or controlling terms involving β2h/γ2 and the accumulated horizon nh.

C.3.1 Log-concave case

The log-concave-case proofs apply the RM–ULMC error bounds across strongly, semi-, and weakly convex regimes, then convert KL guarantees into total variation or chi-squared guarantees.

  • Strongly convex case: The strongly convex proof combines Theorem 4.1 with Lemmas C.2 and C.10 under h ≲1/(β1/2κ), yielding the stated KL bound after choosing T and h.The resulting accuracy condition includes the restriction induced by the step-size constraint.
  • Strongly convex conclusion: In the strongly convex case, KL(µ,π) ⩽ε2 holds for T ≳|ω|−1 log(αW 2/ε2), with the ε condition also constrained by h ≲1/(β1/2κ).This conclusion follows after applying the step-size restriction in the preceding bound.
  • Weakly convex case: T ≍β1/2W 2/ε2 makes the first error term at most ε2, while h is chosen to control the second term subject to h ≲1/(β3/4T1/2).The proof retains only leading-order terms in 1/ε for simplicity.
  • Weakly convex conclusion: For n>N in the weakly convex regime, the proof combines Theorem 4.1 with Lemmas C.2 and C.10 using h ≲γ1/3/(β5/6T1/3).The final bound is obtained from (C.3).
  • Divergence conversion: The final sampling guarantees select T from continuous-time KL or chi-squared convergence, transfer the discrete bound, and conclude via Pinsker’s inequality or the weak triangle inequality.The KL route uses Theorem 5.7 after Lemma 5.8, while the chi-squared route uses Lemma 5.10.
Loading 2506.23062v3…