Fast Non-Log-Concave Sampling under Nonconvex Equality and Inequality Constraints with Landing
Kijung Jeon, Michael Muehlebach, Molei Tao
TL;DR
This work tackles sampling from distributions constrained to nonlinear sets defined by smooth equalities and inequalities. It introduces OLLA, a projection-free overdamped Langevin framework that relies on a tangent-space projection and a deterministic landing term to enforce constraints, proving exponential convergence in 2-Wasserstein distance to the constrained target ρ_Σ ∝ e^{-f} dσ_Σ. It further proposes OLLA–H, an Euler–Maruyama discretization using Hutchinson trace estimation to approximate Itô–Stratonovich corrections, achieving favorable computational costs while preserving sampling accuracy. Across synthetic 2D, high-dimensional, molecular, and Bayesian logistic regression tasks, OLLA and especially OLLA–H match or exceed the performance of projection-based baselines with significantly reduced runtime, demonstrating practical impact for constrained probabilistic inference and physical simulations. The results offer a scalable, theoretically grounded approach to constrained sampling in nonconvex settings, with clear pathways to broader applicability and further theoretical sharpening of the discrete-time guarantees.
Abstract
Sampling from constrained statistical distributions is a fundamental task in various fields including Bayesian statistics, computational chemistry, and statistical physics. This article considers the cases where the constrained distribution is described by an unconstrained density, as well as additional equality and/or inequality constraints, which often make the constraint set nonconvex. Existing methods for nonconvex constraint set $Σ\subset \mathbb{R}^d$ defined by equality or inequality constraints commonly rely on costly projection steps. Moreover, they cannot handle equality and inequality constraints simultaneously as each method only specialized in one case. In addition, rigorous and quantitative convergence guarantee is often lacking. In this paper, we introduce Overdamped Langevin with LAnding (OLLA), a new framework that can design overdamped Langevin dynamics accommodating both equality and inequality constraints. The proposed dynamics also deterministically corrects trajectories along the normal direction of the constraint surface, thus obviating the need for explicit projections. We show that, under suitable regularity conditions on the target density and $Σ$, OLLA converges exponentially fast in $W_2$ distance to the constrained target density $ρ_Σ(x) \propto \exp(-f(x))dσ_Σ$. Lastly, through experiments, we demonstrate the efficiency of OLLA compared to projection-based constrained Langevin algorithms and their slack variable variants, highlighting its favorable computational cost and reasonable empirical mixing.