Technical Summary: Time-Inhomogeneous Preconditioned Langevin Dynamics

Problem Statement

Sampling from Gibbs distributions of the form $p(x) \propto \exp(-\Psi(x))$ is a fundamental task in machine learning and Bayesian inverse problems. Standard approaches rely on discretizing the overdamped Langevin stochastic differential equation (SDE). However, Langevin sampling faces two distinct challenges:

1. Global Mode Coverage: Difficulty in traversing between disjoint modes in multi-modal distributions.
2. Local Mode Exploration: Slow convergence within modes due to ill-conditioned local geometry (diverse curvature) of the potential $\Psi$.

Existing preconditioning strategies attempt to address these by incorporating problem-specific information (e.g., sample covariance or local curvature). However, a trade-off exists: global preconditioners (like covariance) aid mode coverage but fail to adapt to local geometry, while local preconditioners (like inverse Hessian approximations) improve local exploration but may struggle with global connectivity. No prior method simultaneously resolves both challenges within a single framework.

Methodology

The authors propose TIPreL (Time-Inhomogeneous Preconditioned Langevin), an algorithm that utilizes a preconditioner $B(t, x)$ which varies with both time $t$ and position $x$. The dynamics are governed by the SDE:
$$dX_t = [-B(t, X_t)\nabla\Psi(X_t) + \text{div } B(t, X_t)] dt + \sqrt{2}B^{1/2}(t, X_t) dW_t$$
where the $\text{div } B$ term is necessary to maintain the target distribution $\pi$ as the stationary distribution.

Key Design Choices

• Time-Dependent Interpolation: The preconditioner is constructed as a time-dependent convex combination of a global preconditioner (e.g., estimated covariance) and a local, position-dependent preconditioner (e.g., an approximation of the inverse Hessian).
  $$B(t, x) = \lambda_t B_{\text{global}} + (1 - \lambda_t) B_{\text{local}}(x)$$
  where $\lambda_t \to 1$ as $t \to \infty$. This design aims to leverage global information for initial exploration and switch to local curvature information for fine-tuning near modes.
• Tamed Euler Discretization: Since the position-dependence of $B(t, x)$ leads to a drift term that is only locally Lipschitz (violating global Lipschitz assumptions required for standard Euler-Maruyama stability), the authors employ a "tamed" discretization scheme:
  $$Y_{t_{k+1}} = Y_{t_k} + h_k \hat{b}_{h_k}(t_k, Y_{t_k}) + \sqrt{2h_k}B^{1/2}(t_k, Y_{t_k})Z_k$$
  where the drift is tamed via $\hat{b}_h(t, x) = \frac{b(t, x)}{1 + h|b(t, x)|}$.

Theoretical Contributions

The paper establishes rigorous convergence guarantees for TIPreL under general conditions that extend the state of the art:

1. Convergence in Continuous Time: The authors prove exponential convergence in the Kullback-Leibler (KL) divergence for the continuous-time dynamics, assuming the target satisfies a Preconditioned Log-Sobolev Inequality (PLSI).
2. Discretization Error Bounds: They derive bounds on the discretization error for the tamed Euler scheme. Crucially, these results hold for:
   • Locally Lipschitz drifts (rather than globally Lipschitz).
   • Time- and space-dependent diffusion coefficients.
   • These conditions have not been covered by prior work on preconditioned Langevin dynamics.
3. Wasserstein-2 Convergence: The discrete scheme is proven to converge to the target distribution in the Wasserstein-2 distance ($W_2$) as the step size vanishes or the number of steps increases, provided the step sizes are square-summable but not summable.

Experimental Results

The authors evaluate TIPreL against competing schemes (constant scalar, global covariance, inverse expected Fisher information, and static curvature-aware preconditioning) on two tasks:

1. 2D Rosenbrock Potential: A benchmark known for its ill-conditioned, banana-shaped valley.

   • Findings: The time-dependent interpolated preconditioner and the static curvature-aware preconditioner demonstrated faster convergence in $W_2$ distance and lower bias compared to global methods. The interpolated method showed superior performance even in this unimodal setting, attributed to better initial exploration followed by adaptive step-size scaling.
   • Observables: The proposed methods provided faithful estimates of the mean and cosine-wave statistics, whereas global methods exhibited bias or divergence.

2. Bayesian Logistic Regression: A higher-dimensional problem ($d=13$) with an anisotropic Gaussian prior, creating an ill-conditioned posterior.

   • Findings: The interpolated preconditioner consistently exhibited faster convergence to the posterior mean and lower marginal Wasserstein-2 distances compared to global preconditioners.
   • Robustness: Under smaller step sizes, global preconditioners (covariance and Fisher) showed problematic behavior (increasing distance after initial dips), while the interpolated method maintained stability.

Significance and Claims

The paper claims that TIPreL effectively bridges the gap between global mode coverage and local mode exploration by unifying them within a time-inhomogeneous framework. The significance of the work lies in:

• Algorithmic Innovation: Introducing a time-dependent interpolation strategy that dynamically shifts from global to local preconditioning.
• Theoretical Advancement: Providing the first convergence analysis for Langevin sampling with time- and space-dependent diffusion and locally Lipschitz drifts using a tamed Euler discretization.
• Empirical Validation: Demonstrating that combining global and local curvature information yields superior sampling efficiency in both low-dimensional ill-posed problems and higher-dimensional Bayesian inference tasks.

The authors note that a limitation of the approach is the computational cost of calculating the divergence term $\text{div } B$, which they suggest could be addressed in future work via approximate computation or learned preconditioners.