Regularized Stein Variational Gradient Flow
Ye He, Krishnakumar Balasubramanian, Bharath K. Sriperumbudur, Jianfeng Lu
TL;DR
This work tackles the mismatch between deterministic Stein variational gradient flows and the true Wasserstein gradient flow by introducing Regularized SVGF (R-SVGF), which injects a regularization via $\left((1-\nu)\mathcal{T}_{k,\rho}+\nu I\right)^{-1}$ to interpolate between the Stein and Wasserstein dynamics. The authors develop a mean-field PDE for $\rho_t$, establish the existence and uniqueness of weak solutions, and prove stability and convergence results in Fisher information and KL divergence, including under a log-Sobolev inequality. They also analyze a time-discretized version and present a practical Regularized SVGD algorithm with a concrete particle-update scheme and computational considerations, complemented by synthetic numerical evidence of improved performance. The framework provides a principled path to closely approximate the WGF while retaining a controllable, implementable discretization, with explicit rates and dependencies on the regularization parameter $\nu$, kernel spectral data, and functional inequalities.
Abstract
The Stein Variational Gradient Descent (SVGD) algorithm is a deterministic particle method for sampling. However, a mean-field analysis reveals that the gradient flow corresponding to the SVGD algorithm (i.e., the Stein Variational Gradient Flow) only provides a constant-order approximation to the Wasserstein Gradient Flow corresponding to the KL-divergence minimization. In this work, we propose the Regularized Stein Variational Gradient Flow, which interpolates between the Stein Variational Gradient Flow and the Wasserstein Gradient Flow. We establish various theoretical properties of the Regularized Stein Variational Gradient Flow (and its time-discretization) including convergence to equilibrium, existence and uniqueness of weak solutions, and stability of the solutions. We provide preliminary numerical evidence of the improved performance offered by the regularization.
