Fast Non-Log-Concave Sampling under Nonconvex Equality and Inequality Constraints with Landing

AI summary

Overview

Research area: Stochastic sampling / Bayesian inference under constraints, combining stochastic differential equations, Riemannian geometry, and nonconvex optimization.

Technical level: Advanced. The paper builds on Itô calculus, Riemannian manifolds, Wasserstein distances, and log-Sobolev inequalities, though its central ideas are geometric intuition that can be stated simply.

Scope in one sentence: The paper introduces OLLA, a projection-free Langevin sampler that handles nonconvex equality and inequality constraints simultaneously and provably converges exponentially fast to the constrained target distribution.

What This Paper Is About

Sampling from probability distributions restricted to a feasible region — such as a curved surface defined by equations and inequalities — is a core task in Bayesian statistics, chemistry, and physics, but almost all existing Langevin-based samplers rely on projections, which are expensive and require the feasible set to be convex. When the constraint set is nonconvex, projections may be intractable and existing methods often come with no quantitative convergence guarantee. This paper designs a sampler that never projects: instead, it splits motion into a tangential part that explores the constraint surface and a normal "landing" drift that pulls the trajectory back onto the feasible set at an exponential rate.

Key Contributions

1. A unified SDE for mixed constraints. OLLA enforces equality constraints h(x) = 0 and inequality constraints g(x) ≤ 0 in a single stochastic differential equation, constructed by solving two least-squares problems: the diffusion matrix is chosen closest to the unconstrained noise √2·I subject to being orthogonal to all active constraint gradients, and the drift is chosen closest to −∇f subject to the landing equations. It recovers classical equality-only constrained Langevin dynamics as a special case and needs no slack variables or projections.

2. Non-asymptotic exponential convergence guarantees. The authors prove that the continuous-time OLLA converges to the constrained target ρ_Σ ∝ exp(−f) dσ_Σ exponentially in 2-Wasserstein distance, in all three regimes: equality-only, inequality-only, and mixed. Rates depend on the landing rate α and the log-Sobolev constant λ_LSI of the constrained manifold, with the requirement α > 2λ_LSI.

3. A cheap, scalable discretization (OLLA-H). The Euler–Maruyama scheme requires traces of the form Tr(Π∇²h_i) and Tr(Π∇²g_j). The authors replace exact Hessians with a Hutchinson trace estimator using N random probe vectors and Hessian-vector products, reducing per-iteration cost from O(d · grad-cost) to O(N · grad-cost). Only N = 5 probes suffice for low-variance estimates in their experiments.

4. Empirical validation against established baselines. OLLA and OLLA-H are compared to CLangevin, CHMC, and CGHMC on 2D synthetic manifolds and on a controllable high-dimensional stress test, showing competitive sampling accuracy, low constraint violation, and favorable scaling in both dimension and number of constraints.

Main Findings

• Exponential decay of constraint violation. The landing mechanism makes equality residuals decay as h_i(X_t) = h_i(X_0)e^{−αt} exactly, while an initially violated inequality g_j > 0 reaches the boundary in finite time t = (1/α) ln((g_j(X_0)+ε)/ε) and stays feasible afterward. The boundary repulsion parameter ε drives the normal probability flux at the boundary to −αε ρ_t, which vanishes after the landing time.

• Landing plus manifold mixing explains convergence. The analysis splits the dynamics into two effects: fast normal contraction toward the constraint set (rate α), and tangential Langevin mixing on the landed manifold governed by λ_LSI. A key step is showing that the projected process Y_t = π(X_t) behaves like a Langevin diffusion on Σ with perturbation terms of size O(e^{−αt}), which the log-Sobolev inequality can absorb.

• Infusible initialization. Unlike projection-based methods, OLLA can start from an infeasible point and land within finite time — the theory accounts for this explicitly through the cutoff time t_cut.

• Mean curvature = Itô–Stratonovich correction. The correction term ℋ, which arises when converting the Stratonovich SDE to Itô form, is exactly the mean curvature of the constraint intersection manifold, and equals div Π + Π ∇ ln((det G)^{1/2}). This connects OLLA to the orthogonal-direction variational KL samplers of Zhang et al., though OLLA targets the Riemannian volume-weighted density e^{−f} dσ_Σ rather than the coarea measure e^{−f} δ_Σ(dx).

• Inequality constraints add a genuinely new mechanism. Where prior work on inequality constraints (e.g., Zhang et al.) uses purely deterministic normal drift with no tangential noise, OLLA projects noise and drift tangentially even during the landing phase, preserving exploration on the evolving manifold and improving mixing.

• Hyperparameter behavior matches theory. Increasing α improves both W₂² and E|h| (e.g., α from 1 to 200 drops W₂² from 0.363 to 0.121 and E|h| from 0.682 to 0.008), consistent with faster landing and contraction. Increasing ε mainly reduces inequality violation (E[max g⁺] falls from 0.082 at ε = 0.1 to 0.019 at ε = 10) without materially changing W₂², consistent with ε affecting only the finite landing time, not the asymptotic rate. Both parameters destabilize the discretization if set too large.

• Dimension scaling favors OLLA-H. With m = l = 5 fixed and d growing from 10 to 700, OLLA-H's CPU time per effective sample stays roughly flat near 0.05 s/sample, while exact OLLA grows linearly to about 1.1 s/sample at d = 700. CLangevin shows severe bias at high dimension (estimated P(x₁ > 0) collapsing to about 0.2 instead of 0.5), whereas OLLA, OLLA-H, CHMC, and CGHMC agree.

• Constraint count is the harder axis. When the number of equality constraints grows large, OLLA's advantage over equality-specialized baselines narrows, because the largest stable α in the discretized algorithm must be reduced to control discretization error.

Methodology in Plain English

Start with ordinary overdamped Langevin dynamics, whose drift is the negative gradient of the target energy and whose noise is isotropic. Now imagine you are only allowed to move along directions that respect the constraints. The authors formulate this as two small optimization problems: pick the noise matrix closest to the standard isotropic one that is annihilated by every active constraint gradient, and pick the drift vector closest to the usual gradient drift that additionally satisfies one linear equation per constraint.

Those equations are chosen deliberately: applying Itô's lemma to each constraint function shows that if the drift satisfies ∇h_iᵀq + ½Tr(∇²h_i QQᵀ) + αh_i = 0, then h_i decays exponentially at rate α. For inequality constraints, the same trick is applied to g_j + ε, so the constraint hits zero in finite time and then stays negative forever. Solving these least-squares problems in closed form yields a single SDE: tangential gradient drift, an α-scaled normal "landing" drift, a mean-curvature correction, and noise living entirely in the tangent space.

For the convergence proof, the authors observe that once the trajectory is close enough to the feasible set, it can be projected onto it. The projected process approximates a well-behaved Langevin diffusion on the manifold, with error terms bounded by O(e^{−αt}). As long as the landing rate exceeds twice the log-Sobolev constant of the target on the manifold, contraction wins and the Kullback–Leibler divergence — and hence the Wasserstein distance — decays exponentially. The inequality case additionally requires showing the boundary probability flux vanishes after landing; the mixed case requires extra assumptions controlling how the boundary of the inequality region moves as the equality residual shrinks.

For implementation, they discretize with Euler–Maruyama and replace the expensive trace terms with Hutchinson's estimator, which uses random vectors and Hessian-vector products. This is the whole trick behind OLLA-H's dimension-independent per-sample cost.

Why This Matters

Research impact. Constraints are ubiquitous in modern machine learning — safety, fairness, physical consistency, structured priors — yet constrained sampling theory has largely been limited to convex or equality-only settings. This paper provides the first projection-free framework that handles both constraint types together with non-asymptotic guarantees, and it bridges two previously separate literatures: projection-free constrained optimization (which inspired the landing mechanism) and constrained Langevin sampling.

Real-world applications:

• Bayesian inference with structured priors, where parameters must satisfy physical or logical relationships such as conservation laws, normalization conditions, or simplicial constraints.
• Molecular dynamics and computational chemistry, where bond lengths, angles, and steric exclusion correspond directly to equality and inequality constraints on atomic configurations.
• Safe and fair generative modeling, where a generator must produce samples from a region defined by complex nonlinear requirements rather than an unconstrained latent distribution.
• Imitation learning and constrained reinforcement learning, where the feasible action set is implicitly defined and projections onto it are impractical.

Industry relevance. Sampling is the computational bottleneck in uncertainty quantification, simulation-based inference, and probabilistic modeling pipelines. A method whose per-iteration cost does not scale with the ambient dimension, that tolerates infeasible initialization, and that avoids the inner optimization loop of projection, is directly useful for high-dimensional industrial models where projection is the dominant cost.

Future Directions

1. Closing the constraint-count gap. OLLA's advantage over equality-specialized baselines shrinks when the number of equality constraints is large, because numerical stability caps the usable landing rate α. Better discretizations or adaptive α schedules could restore the advantage.

2. Relaxing the geometric assumptions. The convergence theory requires bounded initial constraint violation, a tubular neighborhood with recovery properties, and (in the mixed case) assumptions on how the inequality boundary moves during landing. Extending to degenerate constraints that violate LICQ is a natural open problem.

3. Sharpening the discretization analysis. Convergence guarantees are proven for the continuous-time SDE; a finite-step analysis incorporating the Hutchinson estimator's variance and Euler–Maruyama bias would give end-to-end complexity bounds.

4. Scaling to nonconvex, non-LSI targets and real applications. The theory assumes the constrained target satisfies a log-Sobolev inequality, which can fail for multimodal targets. Demonstrating OLLA on molecular dynamics or large Bayesian models would test whether the synthetic scaling results transfer.

Target Audience

Researchers and graduate students in machine learning, statistics, and applied mathematics working on sampling algorithms, Bayesian inference with constraints, and stochastic differential equations. It is most valuable to readers already comfortable with Langevin dynamics and manifold calculus who want to understand how to sample under nonconvex constraints without projections. Practitioners in computational chemistry, simulation-based inference, and safety-constrained generative modeling will also benefit from the algorithm and its Hutchinson-accelerated implementation, even if they skip the convergence proofs.

Authors’ 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.