Technical Summary: Randomized Midpoint Method for Log-Concave Sampling under Constraints

Problem Statement

The paper addresses the challenge of sampling from log-concave probability distributions supported on convex and compact sets $K \subset \mathbb{R}^p$. The target distribution $\nu$ is defined by a density $\nu(\theta) \propto e^{-U(\theta)}$, where the potential $U(\theta) = f(\theta) + \ell_K(\theta)$. Here, $f$ is a smooth potential function, and $\ell_K$ is the indicator function of the set $K$, which introduces non-differentiability and discontinuities at the boundary.

Existing diffusion-based sampling methods, such as Langevin Monte Carlo (LMC) and Kinetic Langevin Monte Carlo (KLMC), typically rely on the smoothness of the potential function. In constrained settings, the indicator function $\ell_K$ violates these smoothness assumptions, rendering standard discretization analyses invalid. While previous works have attempted to handle this via regularization (smoothing the indicator), existing theoretical guarantees often suffer from loose convergence rates or rely on specific projection schemes (e.g., Euclidean projection only).

Methodology

1. Generalized Smooth Approximation Framework

The authors propose a unified framework to approximate the non-smooth indicator $\ell_K$ with a smooth surrogate $\ell_{K,\lambda}$. Instead of relying solely on the Moreau envelope (Euclidean projection), the framework accommodates a broader class of projection operators:

• Metric-based Projections: Including Euclidean and Bregman projections. The distance function is defined as $d_K(\theta) = \|\theta - P_K(\theta)\|_2^2$, where $P_K$ is the projection operator.
• Scaling-based Projections: Specifically, the Gauge projection, where $d_K(\theta) = (g_K(\theta) - 1)^2$ and $g_K$ is the Minkowski functional (Gauge function) of $K$.

The surrogate potential is defined as $U_\lambda(\theta) = f(\theta) + \frac{1}{2\lambda^2}d_K(\theta)$, yielding a smooth target distribution $\nu_\lambda$. Under specific assumptions on the geometry of $K$ and the smoothness of $f$, the authors establish that $U_\lambda$ is strongly convex and smooth, enabling the application of standard Langevin algorithms.

2. Approximation Error Analysis (Wasserstein Bounds)

A core contribution is the rigorous analysis of the Wasserstein-$q$ distance ($W_q$) between the surrogate $\nu_\lambda$ and the true target $\nu$.

• Upper Bounds: The authors derive tight non-asymptotic upper bounds for $W_q(\nu, \nu_\lambda)$. The convergence rate depends on the dimension $p$ and the order $q$. Notably, they introduce a novel proof technique using a carefully constructed transport map in polar coordinates that aligns the radial components of $\nu_\lambda$ and $\nu$. This yields rates of $O(\lambda)$ for $1/p + 1/q > 1$ and $O(\lambda^{1/p + 1/q})$ for $1/p + 1/q < 1$, which are sharper than previous results (e.g., improving upon the $O(\lambda^{1/4})$ rate for $W_2$ in prior work).
• Lower Bounds: The paper derives the first lower bounds for $W_q(\nu, \nu_\lambda)$, demonstrating that the derived upper bounds are near-optimal in order.

3. Randomized Midpoint Discretization

Building on the smooth surrogate, the paper applies the Randomized Midpoint Method (originally proposed for unconstrained settings) to both overdamped and kinetic Langevin diffusions.

• CRLMC (Constrained Randomized Midpoint LMC): Applies the midpoint scheme to the overdamped Langevin diffusion on $\nu_\lambda$.
• CRKLMC (Constrained Randomized Midpoint KLMC): Applies the scheme to the kinetic Langevin diffusion, incorporating momentum and friction.

The authors also revisit the standard Euler discretization methods (CLMC and CKLMC) within this framework to provide improved baseline guarantees.

Key Contributions and Results

1. Unified Theoretical Framework

The paper unifies the treatment of various projection operators (Euclidean, Bregman, Gauge) under a single analytical umbrella. This generalizes previous work that was often restricted to Euclidean projections.

2. Improved Convergence Rates

By combining the tighter approximation bounds for $\nu_\lambda$ with the randomized midpoint discretization, the authors establish new iteration complexity bounds for achieving $\epsilon$-accuracy in Wasserstein distance:

• CRLMC: Achieves $\tilde{O}(\epsilon^{-10/3})$ for $W_1$ and $\tilde{O}(\epsilon^{-6 + 16/(3p)})$ for $W_2$.
• CRKLMC: Achieves $\tilde{O}(\epsilon^{-8/3})$ for $W_1$ and $\tilde{O}(\epsilon^{-5 + 14/(3p)})$ for $W_2$.
• CLMC/CKLMC (Euler): The refined analysis also improves the bounds for standard Euler-based constrained samplers, reducing the complexity from $\tilde{O}(\epsilon^{-18})$ to $\tilde{O}(\epsilon^{-6 + 4/p})$ for CLMC in $W_2$.

The results demonstrate that the randomized midpoint scheme provides a clear acceleration over standard Euler discretization in constrained settings.

3. Near-Optimality

The derivation of complementary lower bounds confirms that the upper bounds on the approximation error are near-optimal, validating the tightness of the proposed analysis.

4. Empirical Validation

Numerical experiments on synthetic data (2D scatter plots, Bayesian Lasso regression) and real-world data (UCI Adult dataset for fair logistic regression) support the theoretical findings. The experiments show that CRKLMC converges fastest, followed by CRLMC, with both outperforming their Euler-based counterparts (CKLMC and CLMC).

Significance and Claims

The paper claims to advance the theoretical understanding of constrained diffusion-based sampling by:

1. Resolving Non-Smoothness: Providing a systematic way to handle hard constraints via smooth surrogates that accommodate diverse projection geometries.
2. Tightening Bounds: Offering significantly sharper convergence rates than existing literature (specifically improving upon [BDMP17] and [GHZ24]) by utilizing a refined optimal transport argument.
3. Validating Midpoint Methods: Establishing that the randomized midpoint method, known for its benefits in unconstrained sampling, retains its advantages (reduced bias and better condition number dependence) when applied to constrained log-concave sampling.
4. Practical Guidance: Providing explicit constants and parameter dependencies that offer a principled basis for selecting step sizes and tuning parameters ($\lambda$) in constrained sampling algorithms.

The authors note that while their current analysis assumes strong convexity and smoothness for the sampling error analysis, the approximation error analysis relies on weaker assumptions, suggesting potential for future extension to non-convex settings.