Improved sampling algorithms and functional inequalities for non-log-concave distributions
Yuchen He, Zhehan Lei, Jianan Shao, Chihao Zhang
TL;DR
The paper tackles the problem of sampling from μ ∝ e^{-V} with access to V and ∇V under minimal smoothness, showing an exponential gap between the classic (1) and a stronger (1*) assumption along the OU flow. It introduces a variant of the proximal (restricted Gaussian) sampler guided by stochastic localization and OU analysis, achieving polynomial-time sampling in d and 1/ε when L is constant and moment bounds hold. The authors establish explicit Poincaré and modified log-Sobolev inequalities under (1*) and stronger moments, and derive mLSI bounds for mixtures of log-concave components via concatenation across SL phases. In addition to improved sampling guarantees, the work provides structural functional-inequality results for μ, including bounds for mixtures and a clarified link between SL and OU processes, with practical RGO-based sampling and late-initialization techniques.
Abstract
We study the problem of sampling from a distribution $μ$ with density $\propto e^{-V}$ for some potential function $V:\mathbb R^d\to \mathbb R$ with query access to $V$ and $\nabla V$. We start with the following standard assumptions: (1) $V$ is $L$-smooth. (2) The second moment $\mathbf{E}_{X\sim μ}[\|X\|^2]\leq M$. Recently, He and Zhang (COLT'25) showed that the query complexity of this problem is at least $\left(\frac{LM}{dε}\right)^{Ω(d)}$ where $ε$ is the desired accuracy in total variation distance, and the Poincaré constant can be unbounded. Meanwhile, another common assumption in the study of diffusion based samplers (see e.g., the work of Chen, Chewi, Li, Li, Salim and Zhang (ICLR'23)) strengthens (1) to the following: (1*) The potential function of *every* distribution along the Ornstein-Uhlenbeck process starting from $μ$ is $L$-smooth. We show that under the assumptions (1*) and (2), the query complexity of sampling from $μ$ can be $\mathrm{poly}(L,d)\cdot \left(\frac{Ld+M}{ε^2}\right)^{\mathcal{O}(L+1)}$, which is polynomial in $d$ and $\frac{1}ε$ when $L=\mathcal{O}(1)$ and $M=\mathrm{poly}(d)$. This improves the algorithm with quasi-polynomial query complexity developed by Huang et al. (COLT'24). Our results imply that the seemingly moderate strengthening from (1) to (1*) yields an exponential gap in the query complexity. Furthermore, we show that together with the assumption (1*) and the stronger moment assumption that $\|X\|$ is $λ$-sub-Gaussian for $X\simμ$, the Poincaré constant of $μ$ is at most $\mathcal{O}(λ)^{2(L+1)}$. We also establish a modified log-Sobolev inequality for $μ$ under these conditions. As an application of our technique, we obtain a new estimate of the modified log-Sobolev constant for a specific class of mixtures of strongly log-concave distributions.