Splitting Regularized Wasserstein Proximal Algorithms for Nonsmooth Sampling Problems

TL;DR

This paper introduces a splitting-based Wasserstein proximal algorithm for efficient sampling from nonsmooth distributions, with theoretical convergence guarantees and successful high-dimensional applications.

Contribution

It develops a novel splitting-based sampling method using regularized Wasserstein proximal operators, applicable to nonsmooth distributions like Laplace priors, with convergence analysis and practical demonstrations.

Findings

Converges to target distributions in Rényi divergence and Wasserstein-2 distance.

Effective in high-dimensional nonsmooth sampling tasks.

Demonstrates robustness and efficiency in various applications.

Abstract

Sampling from nonsmooth target probability distributions is essential in various applications, including the Bayesian Lasso. We propose a splitting-based sampling algorithm for the time-implicit discretization of the probability flow for the Fokker-Planck equation, where the score function, defined as the gradient logarithm of the current probability density function, is approximated by the regularized Wasserstein proximal. When the prior distribution is the Laplace prior, our algorithm is explicitly formulated as a deterministic interacting particle system, incorporating softmax operators and shrinkage operations to efficiently compute the gradient drift vector field and the score function. We verify the convergence towards target distributions regarding R\'enyi divergences and Wasserstein-2 distance under suitable conditions. Numerical experiments in high-dimensional nonsmooth sampling problems, such as sampling from mixed Gaussian and Laplace distributions, logistic regressions, image restoration with $L_1$-TV regularization, and Bayesian neural networks, demonstrate the efficiency and robust performance of the proposed method.