Faster Sampling via Stochastic Gradient Proximal Sampler

TL;DR

This paper introduces stochastic variants of the proximal sampler, achieving faster convergence and improved gradient complexities for sampling from complex distributions, validated by empirical results.

Contribution

It develops a general framework for stochastic proximal samplers, providing convergence guarantees and practical algorithms based on SGLD and MALA with superior theoretical performance.

Findings

Achieves $ ilde{O}(d\, ext{epsilon}^{-2})$ TV error sampling complexity.

Outperforms existing methods by at least $ ilde{O}(d^{1/3})$ in gradient complexity.

Empirical results demonstrate efficiency on synthetic high-dimensional data.

Abstract

Stochastic gradients have been widely integrated into Langevin-based methods to improve their scalability and efficiency in solving large-scale sampling problems. However, the proximal sampler, which exhibits much faster convergence than Langevin-based algorithms in the deterministic setting Lee et al. (2021), has yet to be explored in its stochastic variants. In this paper, we study the Stochastic Proximal Samplers (SPS) for sampling from non-log-concave distributions. We first establish a general framework for implementing stochastic proximal samplers and establish the convergence theory accordingly. We show that the convergence to the target distribution can be guaranteed as long as the second moment of the algorithm trajectory is bounded and restricted Gaussian oracles can be well approximated. We then provide two implementable variants based on Stochastic gradient Langevin dynamics (SGLD) and Metropolis-adjusted Langevin algorithm (MALA), giving rise to SPS-SGLD and SPS-MALA. We further show that SPS-SGLD and SPS-MALA can achieve $\epsilon$-sampling error in total variation (TV) distance within $\tilde{\mathcal{O}}(d\epsilon^{-2})$ and $\tilde{\mathcal{O}}(d^{1/2}\epsilon^{-2})$ gradient complexities, which outperform the best-known result by at least an $\tilde{\mathcal{O}}(d^{1/3})$ factor. This enhancement in performance is corroborated by our empirical studies on synthetic data with various dimensions, demonstrating the efficiency of our proposed algorithm.