Accelerating Convergence of Stein Variational Gradient Descent via Deep Unfolding

Stein variational gradient descent (SVGD) is a prominent particle-based variational inference method used for sampling a target distribution. In this paper, we propose two novel trainable algorithms based on SVGD: deep-unfolded SVGD (DUSVGD) and Chebyshev-step based DUSVGD (C-DUSVGD). DUSVGD incorporates a deep-learning technique called deep unfolding into SVGD by embedding trainable parameters. C-DUSVGD combines DUSVGD with Chebyshev steps, which reduce the number of trainable parameters. DUSVGD and C-DUSVGD have T and 2 trainable parameters in T iterations to enhance the flexibility of SVGD, respectively. These parameters are learned to accelerate the convergence speed. To evaluate the proposed trainable SVGD algorithms, we conducted numerical simulations of three tasks: sampling a one-dimensional Gaussian mixture, performing Bayesian logistic regression, and learning Bayesian neural networks. The results show that our proposed algorithms exhibit faster convergence than the conventional variants of SVGD. For instance, in sampling a one-dimensional Gaussian mixture distribution, the proposed algorithms achieved more than 30% faster convergence compared to existing algorithms, indicating that DUSVGD and C-DUSVGD are promising trainable algorithms for fast sampling.