A proximal gradient algorithm for composite log-concave sampling

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 ${\alpha }_f$-convex $f$ and a restricted Gaussian oracle (RGO) for ${\alpha }_g$-convex $g$, where ${\alpha }_g\ge 0$. 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$-convex and $f$ is $\beta$-smooth, our sampler achieves $\epsilon$ error in total variation distance using $\tilde{\mathcal{O}}(\kappa \sqrt{d}{\log }^4(1/\epsilon ))$ oracle calls where $\kappa :=1\vee \beta /\alpha$, with $0\le {\alpha }_g\le \beta$ and $\alpha :={\alpha }_f+{\alpha }_g>0$, 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é or log-Sobolev inequality, and (2) $f$ is non-smooth but Lipschitz.