Random-Projection Tree Stein Variational Gradient Descent

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.