A proximal gradient algorithm for composite log-concave sampling

TL;DR

This paper introduces a proximal gradient sampling algorithm for composite log-concave distributions, achieving near-optimal convergence rates under certain smoothness and convexity conditions.

Contribution

It develops a novel sampling method combining gradient evaluations and a restricted Gaussian oracle for efficient sampling from composite distributions.

Findings

Achieves ( ilde{ ext{O}}(\u03ba \u221a d \u2227  (1/\u03b5))) convergence for strongly convex, smooth functions.

Extends to non-log-concave distributions satisfying Poincare9 or log-Sobolev inequalities.

Handles non-smooth Lipschitz functions with maintained convergence guarantees.

Abstract

We propose an algorithm to sample from composite log-concave distributions over $\mathbb{R}^d$, i.e., densities of the form $\pi\propto e^{-f-g}$, assuming access to gradient evaluations of $f$ and a restricted Gaussian oracle (RGO) for $g$. The latter requirement means that we can easily sample from the density $\text{RGO}_{g,h,y}(x) \propto \exp(-g(x) -\frac{1}{2h}||y-x||^2)$, which is the sampling analogue of the proximal operator for $g$. If $f + g$ is $\alpha$-strongly convex and $f$ is $\beta$-smooth, our sampler achieves $\varepsilon$ error in total variation distance in $\widetilde{\mathcal O}(\kappa \sqrt d \log^4(1/\varepsilon))$ iterations where $\kappa := \beta/\alpha$, which matches prior state-of-the-art results for the case $g=0$. We further extend our results to cases where (1) $\pi$ is non-log-concave but satisfies a Poincar\'e or log-Sobolev inequality, and (2) $f$ is non-smooth but Lipschitz.