Random-Projection Tree Stein Variational Gradient Descent

Shishuo Guo, Xiaoyuan Cheng, Xing Liu, Ziang Niu, Zonghao Chen, Xu Liu, Zhuo Sun

Keywords: Stein Variational Gradient Descent

Abstract: Vanilla SVGD is known to have a computational cost of $\mathcal{O}(N^2)$ per iteration. In this work, we introduce Random-Projection Tree Stein Variational Gradient Descent (RP-SVGD) to alleviate this computational challenge. This is achieved by restricting kernel interactions to spatially proximal particles clustered via a random projection tree, further reducing to a cost of $\mathcal{O}(N d(\log (N/C)+C))$, where $C$ is a hyperparameter representing the maximum leaf node capacity of the spanning tree. To establish theoretical validity, we introduce a smoothed RP-Tree kernel for any fixed tree realization and prove that it belongs to the Stein class of the target distribution, thereby ensuring the generation of valid gradient flows. In addition, by considering the effective kernel as the expectation of random tree partitions, we verify that it preserves regularity along with other necessary conditions, which guarantee convergence under the approximate gradient flow framework. Extensive experimental results demonstrate that RP-SVGD tends to have competitive performance and significant speedups across various tasks.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 47