Title: Provably Fast Finite Particle Variants of SVGD via Virtual Particle Stochastic Approximation
Venue: NeurIPS 2023
Accepted: yes
Keywords: Stein Variational Gradient Descent

Variational Inference

Sampling
Abstract: Stein Variational Gradient Descent (SVGD) is a popular particle-based variational inference algorithm with impressive empirical performance across various domains. Although the population (i.e, infinite-particle) limit dynamics of SVGD is well characterized, its behavior in the finite-particle regime is far less understood. To this end, our work introduces the notion of *virtual particles* to develop novel stochastic approximations of population-limit SVGD dynamics in the space of probability measures, that are exactly realizable using finite particles. As a result, we design two computationally efficient variants of SVGD, namely VP-SVGD and GB-SVGD, with provably fast finite-particle convergence rates. Our algorithms can be viewed as specific random-batch approximations of SVGD, which are computationally more efficient than ordinary SVGD. We show that the $n$ particles output by VP-SVGD and GB-SVGD, run for $T$ steps with batch-size $K$, are at-least as good as i.i.d samples from a distribution whose Kernel Stein Discrepancy to the target is at most $O(\tfrac{d^{1/3}}{(KT)^{1/6}})$ under standard assumptions. Our results also hold under a mild growth condition on the potential function, which is much weaker than the isoperimetric (e.g. Poincare Inequality) or information-transport conditions (e.g. Talagrand's Inequality $\mathsf{T}_1$) generally considered in prior works. As a corollary, we analyze the convergence of the empirical measure (of the particles output by VP-SVGD and GB-SVGD) to the target distribution and demonstrate a **double exponential improvement** over the best known finite-particle analysis of SVGD. Beyond this, our results present the **first known oracle complexities for this setting with polynomial dimension dependence**, thereby completely eliminating the curse of dimensionality exhibited by previously known finite-particle rates.

Provably Fast Finite Particle Variants of SVGD via Virtual Particle Stochastic Approximation

Abstract

Stein Variational Gradient Descent (SVGD) is a popular particle-based variational inference algorithm with impressive empirical performance across various domains. Although the population (i.e, infinite-particle) limit dynamics of SVGD is well characterized, its behavior in the finite-particle regime is far less understood. To this end, our work introduces the notion of virtual particles to develop novel stochastic approximations of population-limit SVGD dynamics in the space of probability measures, that are exactly realizable using finite particles. As a result, we design two computationally efficient variants of SVGD, namely VP-SVGD and GB-SVGD, with provably fast finite-particle convergence rates. Our algorithms can be viewed as specific random-batch approximations of SVGD, which are computationally more efficient than ordinary SVGD. We show that the n particles output by VP-SVGD and GB-SVGD, run for T steps with batch-size K , are atleast as good as i.i.d samples from a distribution whose Kernel Stein Discrepancy to the target is at most O ( d 1 / 3 / ( KT ) 1 / 6 ) under standard assumptions. Our results also hold under a mild growth condition on the potential function, which is much weaker than the isoperimetric (e.g. Poincare Inequality) or information-transport conditions (e.g. Talagrand's Inequality T 1 ) generally considered in prior works. As a corollary, we analyze the convergence of the empirical measure (of the particles output by VP-SVGD and GB-SVGD) to the target distribution and demonstrate a double exponential improvement over the best known finite-particle analysis of SVGD. Beyond this, our results present the first known oracle complexities for this setting with polynomial dimension dependence , thereby completely eliminating the curse of dimensionality exhibited by previously known finite-particle rates.

1 Introduction

Sampling from a distribution over R d whose density π ⋆ ( x ) ∝ exp( -F ( x )) is known only upto a normalizing constant, is a fundamental problem in machine learning [52, 20], statistics [41, 36], theoretical computer science [28, 17] and statistical physics [39, 15]. A popular approach to this is the Stein Variational Gradient Descent (SVGD) algorithm introduced by Liu and Wang [31], which uses a positive definite kernel k to evolve a set of n interacting particles ( x ( i ) t ) i ∈ [ n ] ,t ∈ N as follows:

SVGD exhibits remarkable empirical performance in a variety of Bayesian inference, generative modeling and reinforcement learning tasks [31, 51, 54, 22, 50, 34, 45, 19] and usually converges rapidly to the target density while using only a few particles, often outperforming Markov Chain Monte Carlo (MCMC) methods. However, in contrast to its wide practical applicability, theoretical

analysis of the behavior SVGD is a relatively unexplored problem. Prior works on the analysis of SVGD [26, 14, 30, 42, 7] mainly consider the population limit, where the number of particles n →∞ . These works assume that the initial distribution of the (infinite number of) particles has a finite KL divergence to the target π ⋆ and subsequently, interpret the dynamics of population-limit SVGD as projected gradient descent updates for the KL divergence on the space of probability measures, equipped with the Wasserstein geometry. Under appropriate assumptions on the target density, one can then use the theory of Wasserstein Gradient Flows to establish non-asymptotic (in time) convergence of population-limit SVGD to π ⋆ in the Kernel Stein Discrepancy (KSD) metric.

While the framework of Wasserstein Gradient Flows suffices to explain the behavior of SVGD in the population limit, the same techniques are insufficient to effectively analyze SVGD in the finite-particle regime. This is primarily due to the fact that the empirical measure ˆ µ ( n ) of a finite number of particles does not admit a density with respect to the Lebesgue measure, and thus, its KL divergence to the target is always infinite (i.e. KL ( ˆ µ ( n ) || π ⋆ ) = ∞ ). In such a setting, a direct analysis of the dynamics of finite-particle SVGD becomes prohibitively difficult due to complex inter-particle dependencies. To the best of our knowledge, the pioneering work of Shi and Mackey [43] is the only result that obtains an explicit convergence rate of finite-particle SVGD by tracking the deviation between the law of n -particle SVGD and that of its population-limit. To this end, Shi and Mackey [43] show that for subgaussian target densities, the empirical measure of n -particle SVGD converges to π ⋆ at a rate of O ( √ poly ( d ) log log n Θ(1 /d ) ) in KSD 1 . The obtained convergence rate is quite slow and fails to adequately explain the impressive practical performance of SVGD.

Our work takes a starkly different approach to this problem and deliberately deviates from tracking population-limit SVGD using a finite number of particles. Instead, we directly analyze the dynamics of KL divergence along a carefully constructed trajectory in the space of distributions. To this end, our proposed algorithm, Virtual Particle SVGD (VP-SVGD) devises an unbiased stochastic approximation (in the space of measures) of the population-limit dynamics of SVGD . We achieve this by considering additional particles called virtual particles 2 which evolve in time but aren't part of the output (i.e. real particles ). These virtual particles are used only to compute information about the current population-level distribution of the real particles, and enable exact implementation of our stochastic approximation to population-limit SVGD, while using only a finite number of particles.

Our analysis is similar in spirit to non-asymptotic analyses of stochastic gradient descent (SGD) that generally do not attempt to track gradient descent (analogous to population-limit SVGD in this case), but instead directly track the evolution of the objective function along the SGD trajectory using appropriate stochastic descent lemmas [24, 21, 11]. The key feature of our proposed stochastic approximation is the fact that it can be implemented using only a finite number of particles. This allows us to design faster variants of SVGD with provably fast finite-particle convergence.

1.1 Contributions

VP-SVGD and GB-SVGD We propose two variants of SVGD that enjoy provably fast finite-particle convergence guarantees to the target distribution: Virtual Particle SVGD (VP-SVGD in Algorithm 1) and Global Batch SVGD (GB-SVGD in Algorithm 2). VP-SVGD is a conceptually elegant stochastic approximation (in the space of probability measures) of population-limit SVGD, and GB-SVGD is a practically efficient version of SVGD which achieves good empirical performance. Our analysis of GB-SVGD builds upon that of VP-SVGD. When the potential F is smooth and satisfies a quadratic growth condition (which holds under subgaussianity of π ⋆ , a common assumption in prior works [42, 43]), we show that the n particles output by T steps of our algorithms, run with batch-size K , are at least as good as i.i.d draws from a distribution whose Kernel Stein Discrepancy to π ⋆ is at most O ( d 1 / 3 / ( KT ) 1 / 6 ) . Our results also hold under a mild subquadratic growth condition for F , which is much weaker than isoperimetric (e.g. Poincare Inequality) or information-transport (e.g. Talagrand's Inequality T 1 ) assumptions generally considered in the sampling literature [47, 42, 43, 8, 2].

State-of-the-art Finite Particle Guarantees As corollaries of the above result, we establish that VP-SVGD and GB-SVGD exhibit the best known finite-particle guarantees in the literature which significantly outperform that of prior works . Our results are summarized in Table 1. In particular,

under subgaussianity of the target distribution π ⋆ , we show that the empirical measure of the n particles output by VP-SVGD converges to π ⋆ in KSD at a rate of O (( d / n ) 1 / 4 +( d / n ) 1 / 2 ) . Similarly, the empirical measure of the n output particles of GB-SVGD converges to π ⋆ at a KSD rate of O ( d 1 / 3 / n 1 / 12 +( d / n ) 1 / 2 ) . Both these results represent a double exponential improvement over the O ( poly ( d ) √ log log n Θ( 1 /d ) ) KSD rate of n -particle SVGD obtained by Shi and Mackey [43], which, to our knowledge, is the best known prior result for SVGD in the finite particle regime. When benchmarked in terms of gradient oracle complexity, i.e., the number of evaluations of ∇ F required by an algorithm to achieve KSD π ⋆ ( ·|| π ⋆ ) ≤ ϵ , we demonstrate that for subgaussian π ⋆ , the oracle complexity of VP-SVGD is O ( d 4 / ϵ 12 ) while that of GB-SVGD is O ( d 6 / ϵ 18 ) . To the best of our knowledge, our result presents the first known oracle complexity guarantee with polynomial dimension dependence , and consequently, does not suffer from a curse of dimensionality unlike prior works. Furthermore, as discused above, the conditions under which our result holds is far weaker than subgaussianity of π ⋆ , and as such, includes sub-exponential targets and beyond. In particular, our guarantees for sub-exponential target distributions are (to the best of our knowledge) the first of its kind.

Computational Benefits: VP-SVGD and GB-SVGD can be viewed as specific random batch approximations of SVGD. Our experiments (Section 8) show that GB-SVGD obtains similar performance as SVGD but requires fewer computations. In this context, a different kind of random batch method that divides the particles into random subsets of interacting particles, has been proposed by Li et al. [29]. However, the objective in Li et al. [29] is to approximate finite-particle SVGD dynamics using the random batch method, instead of analyzing convergence of the random batch method itself. Beyond this, their guarantees also suffer from an exponential dependence on the time T . As explained below, their approach is also conceptually different from our method since we use the same random batch to evolve every particle, allowing us to interpret this as a stochastic approximation in the space of distributions instead of in the path space.

1.2 Technical Challenges

We resolve the following important conceptual challenges, which may be of independent interest.

Stochastic Approximation in the Space of Probability Measures Stochastic approximations are widely used in optimization, control and sampling [27, 52, 23]. In the context of sampling, stochastic approximations are generally implemented in path space, e.g., Stochastic Gradient Langevin Dynamics (SGLD) [52] takes a random batch approximation of the drift term via the update x t +1 = x t -η K ∑ K -1 j =0 ∇ f ( x t , ξ j ) + √ 2 ηϵ t , ϵ t ∼ N (0 , I ) where E [ f ( x t , ξ j ) | x t ] = F ( x t ) . Such stochastic approximations are then analyzed using the theory of stochastic processes over R d [12, 40, 55, 25]. However, when viewed in the space of probability measures (i.e, µ t = Law( x t ) ), the time-evolution

of these algorithms is deterministic. In contrast, our approach designs stochastic approximations in the space of probability measures . In particular, the time-evolution of the law of any particle in VP-SVGD and GB-SVGD are a stochastic approximation of the dynamics of population-limit SVGD. We ensure that requires only a finite number of particles for exact implementation.

Tracking KL Divergence in the Finite-Particle Regime The population limit ( n →∞ ) ensures that the initial empirical distribution ( µ 0 ) of SVGD admits a density (w.r.t the Lebesgue measure). When the KL divergence between µ 0 and π ⋆ is finite, prior works on population-limit SVGD analyze the time-evolution of the KL divergence to π ⋆ . However, this approach cannot be directly used to analyze finite-particle SVGD since the empirical distribution of a finite number of particles does not admit a density, and thus its KL divergence to π ⋆ is infinite. Our analysis of VP-SVGD and GB-SVGD circumvents this obstacle by considering the dynamics of an infinite number of particles, whose empirical measure then admits a density. However, the careful design ensures that the dynamics of n of these particles can be computed exactly, using only a finite total number of (real + virtual) particles. When conditioned on the virtual particles, these particles are i.i.d. and their conditional law is close to the target distribution with high probability.

2 Notation and Problem Setup

We use ∥ · ∥ , ⟨· , ·⟩ to denote the Euclidean norm and inner product over R d respectively. All other norms and inner products are subscripted to indicate their underlying space. P 2 ( R d ) denotes the space of probability measures on R d with finite second moment, with the Wasserstein-2 metric denoted as W 2 ( µ, ν ) for µ, ν ∈ P 2 ( R d ) . For any two probability measures µ, ν , we denote their KL divergence as KL ( µ || ν ) . For any function f : X → Y and any probability measure µ over X , we let f # µ denote the law of f ( x ) : x ∼ µ . Given a sigma algebra F over some space Ω , and a measurable space X , µ ( · ; · ) : F ×X → R + is a probability kernel if for every x ∈ X , µ ( · ; x ) is a measure over F and for every A ∈ F , the map x → µ ( A ; x ) is measurable. We make use of probability measures µ ( · ; x ) where x is a random element of some appropriate space X , resulting in random probability measures. We use [ m ] and ( m ) to denote the sets { 1 , . . . , m } and { 0 , . . . , m -1 } respectively. For any finite set A , S A denotes the group of all permutations of A . We use the O notation to characterize the dependence of our rates on the number of iterations T , dimension d and batch-size K , suppressing numerical and problem-dependent constants. We use ≲ to denote ≤ upto universal constants.

We fix a symmetric positive definite reproducing kernel k : R d × R d → R and let the corresponding reproducing kernel Hilbert space (RKHS) [44] be denoted as H 0 . We denote the product RKHS as H = ∏ d i =1 H 0 , equipped with the standard inner product for product spaces. We assume k is differentiable in both its arguments and let ∇ 2 k ( x , y ) denote the gradient of k ( · , · ) with respect to the second argument. For any µ ∈ P 2 ( R d ) , we assume H ⊂ L 2 ( µ ) and the inclusion map i µ : H → L 2 ( µ ) is continuous. We use P µ : L 2 ( µ ) →H to denote the adjoint of i µ , i.e., the unique operator which satisfies ⟨ f, i µ g ⟩ L 2 ( µ ) = ⟨ P µ f, g ⟩ H for any f ∈ L 2 ( µ ) , g ∈ H . Carmeli et al. [6] shows that P µ can be expressed as a kernel convolution, i.e., ( P µ f )( x ) = ∫ k ( x , y ) f ( y ) d µ ( y )

We define the function h : R d × R d → R as h ( x , y ) = k ( x , y ) ∇ F ( y ) -∇ 2 k ( x , y ) and h µ ∈ H as h µ = P µ ( ∇ x log( d µ d π ⋆ ( x ))) for any µ ∈ P 2 ( R d ) . Integration by parts shows that h µ ( x ) = ∫ h ( x , y ) d µ ( y ) . The convergence metric we use is the Kernel Stein Discrepancy (KSD) metric, which is widely used for comparing probability distributions [32, 9] and analyzing SVGD [42, 26].

Definition 1 (Kernel Stein Discrepancy) . Define the Langevin Stein Operator of π ⋆ acting on any differentiable g : R d → R d :

For any two probability measures µ, ν , the Kernel Stein Discrepancy between µ and ν (with respect to π ⋆ ), denoted as KSD π ⋆ ( µ || ν ) is defined as

Using integration by parts (see Chwialkowski et al. [9]), it follows that KSD π ⋆ ( µ || ν ) = ∥ h µ -h ν ∥ H

Organization: We review population-limit SVGD in Section 3 and derive VP-SVGD and GBSVGD in Section 4. We state our technical assumptions in Section 5, main results in Section 6 and provide a proof sketch in Section 7. We present empirical evaluation in Section 8.

3 Background on Population-Limit SVGD

We briefly introduce the analysis of population-limit SVGD using the theory of Wasserstein Gradient Flows and refer the readers to Korba et al. [26] and Salim et al. [42] for a detailed treatment.

The space P 2 ( R d ) equipped with the 2-Wasserstein metric W 2 is known as the Wasserstein space, which admits the following Riemannian structure : For any µ ∈ P 2 ( R d ) , the tangent space T µ P 2 ( R d ) can be identified with the Hilbert space L 2 ( µ ) . We can then define differentiable functionals L : P 2 ( R d ) → R and compute their Wasserstein gradients, denoted as ∇ W 2 L . Note that the target π ⋆ is the unique minimizer over P 2 ( R d ) for the functional L [ µ ] = KL ( µ || π ⋆ ) . The Wasserstein Gradient of L [ µ ] is ∇ W 2 L [ µ ] = ∇ x log( d µ d π ⋆ ( x )) [1]. This powerful machinery has served as a backbone for the analysis of algorithms such as LMC [53, 3, 2] and population-limit SVGD [14, 26, 42, 45, 7].

The updates of population-limit SVGD can be viewed as Projected Gradient Descent in the Wasserstein space. Recall from Section 2 that the function h µ ( x ) = P µ ( ∇ log( d µ d π ⋆ ))( x ) = ∫ h ( x , y ) d µ ( y ) . Let ˆ µ n t denote the empirical measures of the SVGD particles ( x ( i ) t ) i ∈ [ n ] at timestep t . We note that the SVGD updates in (1) can be recast as ˆ µ n t +1 = ( I -γh ˆ µ n t ) # ˆ µ t n . In the limit of infinite particles n →∞ , suppose the empirical measure ˆ µ n t converges to the population measure ¯ µ t . In this population limit, the updates of SVGD can be expressed as,

Recall from Section 2 that P ¯ µ t : L 2 (¯ µ t ) →H is the Hilbert adjoint of i ¯ µ t . Since H ⊂ L 2 (¯ µ t ) , the updates of SVGD in the population limit can be seen as Projected Wasserstein Gradient Descent for L [ µ ] = KL ( µ || π ⋆ ) , with the Wasserstein Gradient at each step being projected onto the RKHS H . Assuming KL (¯ µ 0 || π ⋆ ) < ∞ , convergence of population limit SVGD is then established by tracking the evolution of KL (¯ µ t || π ⋆ ) under appropriate structural assumptions (such as subgaussianity) on π ⋆ .

4 Algorithm and Intuition

In this section, we derive VP-SVGD (Algorithm 1), and build upon it to obtain GB-SVGD. Consider a countably infinite collection of particles x ( l ) 0 ∈ R d , l ∈ N ∪ { 0 } , sampled i.i.d from a measure µ 0 , having a density w.r.t. the Lebesgue measure. By the strong law of large numbers, the empirical measure of x ( l ) 0 is almost surely equal to µ 0 (see Dudley [13, Theorem 11.4.1]). Let batch size K ∈ N denote the batch size, and F t denote the filtration F t , t ≥ 0 as F t = σ ( { x ( l ) 0 | l ≤ Kt -1 } ) , ∀ t ∈ N , with F 0 being the trivial σ algebra. For ease of exposition, we discuss the case of K = 1 in this section and present a complete derivation for arbitrary K ≥ 1 in Appendix C. Recall from Section 3 that the updates of population-limit SVGD in P 2 ( R d ) can be expressed as follows:

We aim to design a stochastic approximation in P 2 ( R d ) for the updates (2), such that it admits a finite-particle realization. To this end, we propose the following dynamics in R d

Now, for each time-step t , we focus on the time evolution of the particles ( x ( l ) t ) l ≥ t (called the lower triangular evolution ). From (3), we observe that for any t ∈ N and l ≥ t , x ( l ) t depends only on x (0) 0 , . . . , x ( t -1) 0 , x ( l ) 0 . Hence, there exists a deterministic, measurable function H t such that:

Since x (0) 0 , . . . , x ( t -1) 0 , x ( l ) 0 i.i.d. ∼ µ 0 , we conclude from (4) that ( x ( l ) t ) l ≥ t are i.i.d when conditioned on x (0) 0 , . . . , x ( t -1) 0 . To this end, we define the random measure µ t |F t as the law of x ( t ) t conditioned on F t , i.e., µ t |F t is a probability kernel µ t ( · ; x (0) 0 , . . . , x ( t -1) 0 ) , where µ 0 |F 0 := µ 0 . By the strong law of large numbers, µ t |F t is equal to the empirical measure of ( x ( l ) t ) l ≥ t conditioned on F t . We will use µ t |F t and µ t ( · ; x (0) 0 , . . . , x ( t -1) 0 ) interchangeably.

Define the random function g t : R d → R d as g t ( x ) := h ( x , x ( t ) t ) . From (4), we note that g t is F t +1 measurable. From (3), we infer that the particles satisfy the following relation:

Recall that x ( s ) t +1 | x (0) 0 , . . . , x ( t ) 0 ∼ µ t +1 |F t +1 for any s ≥ t +1 . Furthermore, from Equation (4), we note that for s ≥ t +1 , x ( s ) t depends only on x (0) 0 , . . . , x ( t -1) 0 and x ( s ) 0 . Hence, we conclude that Law( x ( s ) t | x (0) 0 , . . . , x ( t ) 0 ) = Law( x ( s ) t | x (0) 0 , . . . , x ( t -1) 0 ) = µ t |F t . With this insight, the dynamics of the lower-triangular evolution in P 2 ( R d ) that the following holds almost surely:

x ( t ) t |F t ∼ µ t |F t implies E [ g t ( x ) |F t ] = h µ t |F t ( x ) . Thus lower triangular dynamics (5) is a stochastic approximation in P 2 ( R d ) to the population limit of SVGD (2). Setting the batch size to general K and tracking the evolution of the first KT + n particles, we obtain VP-SVGD (Algorithm 1).

Algorithm 1 Virtual Particle SVGD ( VP-SVGD )

Input : Number of steps T , number of output particles n , batch size K , Initial positions x (0) 0 , . . . , x ( n + KT -1) 0 i.i.d. ∼ µ 0 , Kernel k , step size γ .

1: for t ∈ { 0 , . . . , T -1 } do

2: for s ∈ { 0 , . . . , KT + n -1 } do

4: end for

5: end for

6: Draw S uniformly at random from { 0 , . . . , T -1 }

Virtual Particles In Algorithm 1, ( x ( l ) t ) KT ≤ l ≤ KT + n -1 are the real particles which constitute the output. ( x ( l ) t ) l<KT are virtual particles which propagate information about the probability measure µ t |F t to enable computation of g t , an unbiased estimate of the projected Wasserstein gradient h µ t |F t .

VP-SVGD as SVGD Without Replacement VP-SVGD is a without-replacement random-batch approximation of SVGD (1), where a different batch is used across timesteps, but the same batch is across particles given a fixed timestep. With i.i.d. initialization, picking the 'virtual particles' in a fixed order or from a random permutation does not change the evolution of the real particles. With this insight, we design GB-SVGD (Algorithm 2) where we consider n particles and output n particles (instead of wasting KT particles as 'virtual particles') via a random-batch approximation of SVGD.

Algorithm 2 Global Batch SVGD ( GB-SVGD )

Input : # of time steps T , # of particles n , x (0) 0 , . . . , x ( n -1) 0 i.i.d. ∼ µ 0 , Kernel k , step size γ , Batch size K , Sampling method ∈ { with replacement , without replacement }

1: for t ∈ { 0 , . . . , T -1 } do

2: K t ← random subset of [ n ] of size K (via. sampling method)

3: for s ∈ { 0 , . . . , n -1 } do

4: x ( s ) t +1 = x ( s ) t -γ K ∑ r ∈K t [ k ( x ( s ) t , x ( r ) t ) ∇ F ( x ( r ) t ) -∇ 2 k ( x ( s ) t , x ( r ) t )]

5: end for

6: end for

7: Draw S uniformly at random from { 0 , 1 , . . . , T -1 }

8: Output (¯ y (0) , . . . , ¯ y ( n -1) ) = ( x (0) S , . . . , x ( n -1) S )

In Algorithm 2, with replacement sampling means selecting a batch of K particles i.i.d. from the uniform distribution over [ n ] . Without replacement sampling means fixing a random permutation σ over { 0 , . . . , n -1 } and selecting the batches in the order specified by the permutation.

5 Assumptions

In this section, we discuss the key assumptions required for our analysis of VP-SVGD and GB-SVGD. Our first assumption is smoothness of F , which is standard in optimization and sampling.

Assumption 1 (L-Smoothness) . ∇ F exists and is L Lipschitz. Moreover ∥∇ F (0) ∥ ≤ √ L .

√

It is easy find a point such that ∥∇ F ( x ∗ ) ∥ ≤ L (e.g., using Θ(1) gradient descent steps [37]) and center the initialization at µ 0 at x ∗ . For clarity, we take x ∗ = 0 without loss of generality. We now impose the following growth condition on F .

Assumption 2 (Growth Condition) . There exist α, d 1 , d 2 > 0 such that

Note that Assumption 1 ensures α ≤ 2 . Assumption 2 is essentially a tail decay assumption on the target density π ⋆ ( x ) ∝ e -F ( x ) . In fact, as we shall show in Appendix B, Assumption 2 ensures that the tails of π ⋆ decay as ∝ e -∥ x ∥ α . Consequently, Assumption 2 holds with α = 2 when π ⋆ is subgaussian and with α = 1 when π ⋆ is subexponential. Subgaussianity is equivalent to π ⋆ satisfying the T 1 inequality [5, 49], commonly assumed in prior works on SVGD [42, 43]. We also note that subexponentiality is implied when π ⋆ satisfies the Poincare Inequality [4, Section 4], which is considered a mild condition in the sampling literature [47, 8, 2, 12, 7]. This makes Assumption 1 significantly weaker than the isoperimetric or information-transport assumptions considered in prior works.

Next, we impose a mild assumption on the RKHS of the kernel k , which has been used by several prior works [42, 26, 45, 43].

Assumption 3 (Bounded RKHS Norm) . For any y ∈ R d , k ( · , y ) satisfies ∥ k ( · , y ) ∥ H 0 ≤ B . Furthermore, ∇ 2 k ( · , y ) ∈ H and ∥∇ 2 k ( · , y ) ∥ H ≤ B

Assumption 3 ensures that the adjoint operator P µ , used in Sections 2 and 3, is well-defined. We also make the following assumptions on the kernel k , which is satisfied by a large class of standard kernels such as Radial Basis Function kernels and Matérn kernels of order ≥ 3 / 2 .

Assumption 4 (Kernel Decay) . The kernel k satisfies the following for constants A 1 , A 2 , A 3 > 0 .

Finally, we make the following mild assumption on the initialization.

Assumption 5 (Initialization) . The initial distribution µ 0 is such that KL ( µ 0 || π ⋆ ) < ∞ . Furthermore, µ 0 is supported in B ( R ) , the ℓ 2 ball of radius R

Since prior works usually assume Gaussian initialization [42, 47], Assumption 5 may seem slightly non-standard. However, this is not a drawback. In fact, whenever R = Θ( √ d + polylog ( n / δ )) , Gaussian initialization can be made indistinguishable from Uniform ( B ( R )) initialization, with probability at least 1 -δ , via a coupling argument. To this end, we impose Assumption 5 for ease of exposition and our results can be extended to consider Gaussian initialization. In Appendix B we show that taking R = √ d / L and µ 0 = Uniform ( B ( R )) suffices to ensure KL ( µ 0 || π ⋆ ) = O ( d ) .

6 Results

6.1 VP-SVGD

Our first result, proved in Appendix C, shows that the law of the real particles of VP-SVGD , when conditioned on the virtual particles, is close to π ⋆ in KSD. As a consequence, it shows that the particles output by VP-SVGD are i.i.d. samples from a random probability measure ¯ µ ( · ; x (0) 0 , . . . , x ( KT -1) 0 , S ) which is close to π ⋆ in KSD. We also present a high-probability version of this result in Appendix C.

Theorem 1 ( Convergence of VP-SVGD ) . Let µ t be as defined in Section 4. Let Assumptions 1 2, 3, 4, and 5 be satisfied and let γ ≤ min { 1 / 2 A 1 L , 1 / (4+ L ) B } . There exist ( ζ i ) 0 ≤ i ≤ 3 depending polynomially

on A 1 , A 2 , A 3 , B, L, d 1 , d 2 for any fixed α ∈ (0 , 2] , such that whenever γξ ≤ 1 2 B , with ξ = ζ 0 + ζ 1 ( γT ) 1 / α + ζ 2 ( γ 2 T ) 1 / α + ζ 3 R 2 / α , the following holds:

Define the probability kernel ¯ µ ( · ; · ) as follows: For any x τ ∈ R d , τ ∈ ( KT ) and s ∈ ( T ) , ¯ µ ( · ; x 0 , . . . , x KT -1 , s ) := µ s ( · ; x 0 , . . . , x Ks -1 ) and ¯ µ ( · ; x 0 , . . . , x KT -1 , s = 0) := µ 0 ( · ) . Conditioned on x (0) τ = x τ , S = s for every τ ∈ ( KT ) , the outputs y (0) , . . . , y ( n -1) of VP-SVGD are i.i.d samples from ¯ µ ( · ; x 0 , . . . , x KT -1 , s ) . Furthermore,

Convergence Rates Taking µ 0 = Uniform ( B ( R )) with R = √ d / L ensures KL ( µ 0 |F 0 || π ⋆ ) = O ( d ) (see Appendix B). Under this setting, choosing γ = O ( ( Kd ) η T 1 -η ) ensures that E [ KSD 2 π ⋆ (¯ µ || π ⋆ )] = O ( d 1 -η ( KT ) η ) where η = α 2(1+ α ) . Thus, for α = 2 , (i.e, sub-Gaussian π ⋆ ), KSD 2 = O ( d 2 / 3 ( KT ) 1 / 3 ) . For α = 1 (i.e, sub-Exponential π ⋆ ), the rate (in squared KSD) becomes O ( d 3 / 4 ( KT ) 1 / 4 ) . To the best of our knowledge, our convergence guarantee for sub-exponential π ⋆ is the first of its kind.

Comparison with Prior Works Salim et al. [42] analyzes population-limit SVGD for subgaussian π ⋆ , obtaining KSD 2 = O ( d 3 / 2 / T ) rate. We note that population-limit SVGD (which requires infinite particles) is not implementable whereas VP-SVGD is an implementable algorithm whose outputs are conditionally i.i.d samples from a distribution with guaranteed convergence to π ⋆ .

6.2 GB-SVGD

We now use VP-SVGD as the basis to analyze GB-SVGD. Assume n > KT . Then, with probability at least 1 -K 2 T 2 / n (for with-replacement sampling) and 1 (for without-replacement sampling), the random batches K t in GB-SVGD (Algorithm 2) are disjoint and contain distinct elements. When conditioned on this event E , we note that the n -KT particles that were not included in any random batch K t evolve exactly like the n real particles of VP-SVGD. With this insight, we show that, conditioned on E , the outputs of VP-SVGD and GB-SVGD can be coupled such that the first n -KT particles of both the algorithms are exactly equal. This allows us to derive the following squared KSD bound between the empirical measures of the outputs of VP-SVGD and GB-SVGD. The proof of this result is presented in Appendix D.

Theorem 2 ( KSD Bounds for GB-SVGD ) . Let n > KT and let Y = ( y (0) , . . . , y ( n -1) ) and ¯ Y = (¯ y (0) , . . . , ¯ y ( n -1) ) denote the outputs of VP-SVGD and GB-SVGD respectively. Moreover, let ˆ µ ( n ) = 1 n ∑ n -1 i =0 δ y ( i ) and ˆ ν ( n ) = 1 n ∑ n -1 i =0 δ ¯ y ( i ) denote their respective empirical measures. Under the assumptions and parameter settings of Theorem 1, there exists a coupling of Y and ¯ Y such that the following holds:

6.3 Convergence of the Empirical Measure to the Target

As a corollary of Theorem 1 and Theorem 2, we show that the empirical measure of the output of VP-SVGD and GB-SVGD rapidly converges to π ⋆ in KSD. We refer to Appendix E for the full statement and proof.

Corollary 1 ( VP-SVGD and GB-SVGD: Fast Finite Particle Rates ) . Let the assumptions and parameter settings of Theorem 1 be satisfied. Let ˆ µ ( n ) be the empirical measure of the n particles output by VP-SVGD, run with run with KT = d α 2+ α , R = √ d / L and appropriately chosen γ . Then:

Similarly, let ˆ ν ( n ) be the empirical measure of the output of GB-SVGD without replacement, run with KT = √ n , R = √ d / L and appropriately chosen γ . Then, the following holds:

Comparison to Prior Work For subgaussian π ⋆ (i.e. α = 2 ), VP-SVGD has a finite-particle rate of E [ KSD π ⋆ (ˆ µ ( n ) || π ⋆ )] = O (( d / n ) 1 / 4 +( d / n ) 1 / 2 ) while that of GB-SVGD is E [ KSD π ⋆ (ˆ ν ( n ) || π ⋆ )] = O ( d 1 / 3 / n 1 / 12 + ( d / n ) 1 / 2 ) . Both these rates are a double exponential improvement over the ˜ O ( poly ( d ) √ log log n Θ( 1 /d ) ) KSD rate obtained by Shi and Mackey [43] for SVGD with subgaussian π ⋆ .

For subexponential π ⋆ (i.e. α = 1 ) the KSD rate of VP-SVGD is O ( d 1 / 3 n 1 / 6 + d n 1 / 2 ) while that of GB-SVGD is O ( d 3 / 8 n 1 / 16 + d n 1 / 2 ) . To our knowledge, both these results are the first of their kind.

Oracle Complexity As illustrated in Section E.3, for subgaussian π ⋆ , the oracle complexity of VP-SVGD to achieve ϵ -convergence in KSD is O ( d 4 / ϵ 12 ) and that of GB-SVGD is O ( d 6 / ϵ 18 ) . To our knowledge, these results are the first known oracle complexities for this problem with polynomial dimension dependence , and significantly improve upon the O ( poly ( d ) ϵ 2 e Θ( de poly ( d ) /ϵ 2 ) ) oracle complexity of SVGD as implied by Shi and Mackey [43]. For subexponential π ⋆ , the oracle complexity of VP-SVGD is O ( d 6 / ϵ 16 ) and that of GB-SVGD is O ( d 9 / ϵ 24 ) .

7 Proof Sketch

We now present a sketch of our analysis. As shown in Section 4, the particles ( x ( l ) t ) l ≥ Kt are i.i.d conditioned on the filtration F t , and the random measure µ t |F t is the law of ( x ( Kt ) t ) conditioned on x (0) 0 , . . . , x ( Kt -1) 0 . Moreover, from equation (5), we know that µ t |F t is a stochastic approximation of population limit SVGD dynamics, i.e., µ t +1 |F t +1 = ( I -γg t ) # µ t |F t . Lemma 1 (similar to Salim et al. [42, Proposition 3.1] and Korba et al. [26, Proposition 5]) shows that under appropriate conditions, the KL divergence between µ t |F t and π ⋆ satisfies a (stochastic) descent lemma . Hence µ t |F t admits a density and KL ( µ t |F t || π ⋆ ) is almost surely finite.

Lemma 1 (Descent Lemma for µ t |F t ) . Let Assumptions 1, 3 and 5 be satisfied and let β > 1 be an arbitrary constant. On the event γ ∥ g t ∥ H ≤ β -1 βB , the following holds almost surely

Lemma 1 is analogous to the noisy descent lemma which is used in the analysis of SGD for smooth functions. Notice that E [ g t |F t ] = h µ t |F t (when interpreted as a Gelfand-Pettis integral [46], as discussed in Appendix B and Appendix C) and hence in expectation, the KL divergence decreases in time. In order to apply Lemma 1, we establish an almost-sure bound on ∥ g t ∥ H below.

Lemma 2. Let Assumptions 1, 2, 3, 4 and 5 hold. For γ ≤ 1 / 2 A 1 L , the following holds almost surely,

where ζ 0 , ζ 1 , ζ 2 and ζ 3 which depend polynomially on A 1 , A 2 , A 3 , B, d 1 , d 2 , L for any fixed α .

Let K = 1 for clarity. To prove Lemma 2, we first note via smoothness of F ( · ) and Assumption 3 that ∥ g t ∥ H ≤ C 0 ∥ x ( t ) t ∥ + C 1 , and then bound ∥ x ( t ) t ∥ . Now, g s ( x ) = k ( x , x ( s ) s ) ∇ F ( x ( s ) s ) -∇ 2 k ( x , x ( s ) s ) . When ∥ x ( s ) s -x ∥ is large, ∥ g s ( x ) ∥ is small due to decay assumptions on the kernel (Assumption 4) implying that the particle does not move much. When x ( s ) s ≈ x , we have g s ( x ) ≈ k ( x , x ( s ) s ) ∇ F ( x ) -∇ 2 k ( x , x ( s ) s ) and k ( x , x ( s ) s ) ≥ 0 . This is approximately a gradient descent update on F ( · ) along with a bounded term ∇ 2 k ( x , x ( s ) s ) . Thus, the value of F ( x ( l ) t ) cannot grow too large after T iterations. By Assumption 2, F ( x ( l ) t ) being small implies that ∥ x ( l ) t ∥ is small.

Equipped with Lemma 2, we set the step-size γ to ensure that the descent lemma (Lemma 1) always holds. The remainder of the proof involves unrolling through Lemma 1 by taking iterated expectations on both sides. To this end we control 〈 h µ t |F t , g t 〉 H and ∥ g t ∥ 2 H in expectation, in Lemma 3.

Lemma 3. Let Assumptions 1,2,3,4,5 hold and ξ be as defined in Lemma 2. Then, for γ ≤ 1 / 2 A 1 L ,

8 Experiments

We compare the performance of GB-SVGD and SVGD. We take n = 100 and use the Laplace kernel with h = 1 for both. We pick the stepsize γ by a grid search for each algorithm. Additional details are presented in Appendix G. We observe that SVGD takes fewer iterations to converge, but the compute time for GB-SVGD is lower. This is similar to the typical behavior of stochastic optimization algorithms like SGD.

SVGD

0.25

GB-SVGD (K=10)

ance

0.20

istar

0.15

D

MMD

0.10

0.05

0.00

0.000

0.0020.0040.006(

0.0080.010

0.012

WallClockTime(seconds)

SVGD

0.25

GB-SVGD (K=10)

ance

0.20

istar

0.15

D

MMD

0.10

0.05

0.00

0

5

10

15

20

25

30

Iterations

(a) MMD vs Compute Time

(b) MMD vs Iterations

(a) Accuracy vs Compute Time

(b) Accuracy vs Iterations

0.75

ccuracy

0.74

0.73

0.72

SVGD

GB-SVGD(K=40)

0.71

20

40

60

WallClockTime(Seconds)

0.75

0.74

cura

0.73

0.72

SVGD

0.71

GB-SVGD(K=40)

2500

5000

7500

10000

12500

15000

Iterations

Sampling from Isotropic Gaussian (Figure 1): As a sanity check, we set π ⋆ = N (0 , I ) with d = 5 . We pick K = 10 for GB-SVGD. The metric of convergence is MMD with respect to the empirical measure of 1000 i.i.d. sampled Gaussians.

Bayesian Logistic Regression (Figure 2) We consider the Covertype dataset which contains ∼ 580 , 000 data points with d = 54 . We consider the same priors suggested in Gershman et al. [16] and implemented in Liu and Wang [31]. We take K = 40 for GB-SVGD. For both VP-SVGD and GB-SVGD, we use AdaGrad with momentum to set the step-sizes as per Liu and Wang [31]

9 Conclusion

We develop two computationally efficient variants of SVGD with provably fast convergence guarantees in the finite-particle regime, and present a wide range of improvements over prior work. A promising avenue of future work could be to establish convergence guarantees for SVGD with general non-logconcave targets, as was considered in recent works on LMC and SGLD [2, 12]. Other important avenues include establishing minimax lower bounds for SVGD and related particle-based variational inference algorithms. Beyond this, we also conjecture that the rates of GB-SVGD can be improved even in the regime n ≪ KT . However, we believe this requires new analytic tools.

References

L. Ambrosio, N. Gigli, and G. Savaré. Gradient flows: in metric spaces and in the space of probability measures . Springer Science & Business Media, 2005.

K. Balasubramanian, S. Chewi, M. A. Erdogdu, A. Salim, and S. Zhang. Towards a theory of non-log-concave sampling:first-order stationarity guarantees for langevin monte carlo. In P.-L. Loh and M. Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory , volume 178 of Proceedings of Machine Learning Research , pages 2896-2923. PMLR, 02-05 Jul 2022. URL https://proceedings.mlr.press/v178/balasubramanian22a.html .

E. Bernton. Langevin monte carlo and jko splitting. In Conference On Learning Theory , pages 1777-1798. PMLR, 2018.

S. Bobkov and M. Ledoux. Poincaré's inequalities and talagrand's concentration phenomenon for the exponential distribution. Probability Theory and Related Fields , 107:383-400, 1997.

S. G. Bobkov and F. Götze. Exponential integrability and transportation cost related to logarithmic sobolev inequalities. Journal of Functional Analysis , 163(1):1-28, 1999.

C. Carmeli, E. De Vito, A. Toigo, and V. Umanitá. Vector valued reproducing kernel hilbert spaces and universality. Analysis and Applications , 8(01):19-61, 2010.

S. Chewi, T. Le Gouic, C. Lu, T. Maunu, and P. Rigollet. Svgd as a kernelized wasserstein gradient flow of the chi-squared divergence. Advances in Neural Information Processing Systems , 33:2098-2109, 2020.

S. Chewi, M. A. Erdogdu, M. Li, R. Shen, and S. Zhang. Analysis of langevin monte carlo from poincare to log-sobolev. In P.-L. Loh and M. Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory , volume 178 of Proceedings of Machine Learning Research , pages 1-2. PMLR, 02-05 Jul 2022. URL https://proceedings.mlr.press/v178/chewi22a. html .

K. Chwialkowski, H. Strathmann, and A. Gretton. A kernel test of goodness of fit. In International conference on machine learning , pages 2606-2615. PMLR, 2016.

J. B. Conway. A course in functional analysis , volume 96. Springer, 2019.

A. Das, B. Schölkopf, and M. Muehlebach. Sampling without replacement leads to faster rates in finite-sum minimax optimization. Advances in Neural Information Processing Systems , 2022.

A. Das, D. Nagaraj, and A. Raj. Utilising the clt structure in stochastic gradient based sampling: Improved analysis and faster algorithms. In Conference on Learning Theory , 2023.

R. M. Dudley. Real analysis and probability . CRC Press, 2018.

A. Duncan, N. Nüsken, and L. Szpruch. On the geometry of stein variational gradient descent. arXiv preprint arXiv:1912.00894 , 2019.

A. El Alaoui, A. Montanari, and M. Sellke. Sampling from the sherrington-kirkpatrick gibbs measure via algorithmic stochastic localization. In 2022 IEEE 63rd Annual Symposium on Foundations of Computer Science (FOCS) , pages 323-334. IEEE, 2022.

S. Gershman, M. D. Hoffman, and D. M. Blei. Nonparametric variational inference. In Proceedings of the 29th International Conference on International Conference on Machine Learning , 2012.

S. Gopi, Y. T. Lee, and D. Liu. Private convex optimization via exponential mechanism. In Conference on Learning Theory , pages 1948-1989. PMLR, 2022.

J. Gorham, A. Raj, and L. Mackey. Stochastic stein discrepancies. Advances in Neural Information Processing Systems , 33:17931-17942, 2020.

T. Haarnoja, H. Tang, P. Abbeel, and S. Levine. Reinforcement learning with deep energy-based policies. In International conference on machine learning , pages 1352-1361. PMLR, 2017.

J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

P. Jain, D. M. Nagaraj, and P. Netrapalli. Making the last iterate of sgd information theoretically optimal. SIAM Journal on Optimization , 31(2):1108-1130, 2021.

P. Jaini, L. Holdijk, and M. Welling. Learning equivariant energy based models with equivariant stein variational gradient descent. Advances in Neural Information Processing Systems , 34: 16727-16737, 2021.

S. Jin, L. Li, and J.-G. Liu. Random batch methods (rbm) for interacting particle systems. Journal of Computational Physics , 400:108877, 2020.

A. Khaled and P. Richtárik. Better theory for sgd in the nonconvex world. arXiv preprint arXiv:2002.03329 , 2020.

Y. Kinoshita and T. Suzuki. Improved convergence rate of stochastic gradient langevin dynamics with variance reduction and its application to optimization. arXiv preprint arXiv:2203.16217 , 2022.

A. Korba, A. Salim, M. Arbel, G. Luise, and A. Gretton. A non-asymptotic analysis for stein variational gradient descent. Advances in Neural Information Processing Systems , 33: 4672-4682, 2020.

H. J. Kushner and D. S. Clark. Stochastic approximation methods for constrained and unconstrained systems , volume 26. Springer Science & Business Media, 2012.

Y. T. Lee and S. S. Vempala. The manifold joys of sampling (invited talk). In 49th International Colloquium on Automata, Languages, and Programming (ICALP 2022) . Schloss DagstuhlLeibniz-Zentrum für Informatik, 2022.

L. Li, Y. Li, J.-G. Liu, Z. Liu, and J. Lu. A stochastic version of stein variational gradient descent for efficient sampling. Communications in Applied Mathematics and Computational Science , 15(1):37-63, 2020.

Q. Liu. Stein variational gradient descent as gradient flow. Advances in neural information processing systems , 30, 2017.

Q. Liu and D. Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in neural information processing systems , 29, 2016.

Q. Liu, J. Lee, and M. Jordan. A kernelized stein discrepancy for goodness-of-fit tests. In International conference on machine learning , pages 276-284. PMLR, 2016.

T. Liu, P. Ghosal, K. Balasubramanian, and N. S. Pillai. Towards understanding the dynamics of gaussian-stein variational gradient descent, 2023.

Y. Liu, P. Ramachandran, Q. Liu, and J. Peng. Stein variational policy gradient. arXiv preprint arXiv:1704.02399 , 2017.

J. Lu, Y. Lu, and J. Nolen. Scaling limit of the stein variational gradient descent: The mean field regime. SIAM Journal on Mathematical Analysis , 51(2):648-671, 2019.

R. M. Neal et al. Mcmc using hamiltonian dynamics. Handbook of markov chain monte carlo , 2(11):2, 2011.

Y. Nesterov. Introductory lectures on convex programming volume i: Basic course. Lecture notes , 3(4):5, 1998.

N. Nüsken and D. Renger. Stein variational gradient descent: many-particle and long-time asymptotics. arXiv preprint arXiv:2102.12956 , 2021.

G. Parisi. Correlation functions and computer simulations. Nuclear Physics B , 180(3):378-384, 1981.

M. Raginsky, A. Rakhlin, and M. Telgarsky. Non-convex learning via stochastic gradient langevin dynamics: a nonasymptotic analysis. In Conference on Learning Theory , pages 1674-1703. PMLR, 2017.

G. O. Roberts and R. L. Tweedie. Exponential convergence of langevin distributions and their discrete approximations. Bernoulli , pages 341-363, 1996.

A. Salim, L. Sun, and P. Richtarik. A convergence theory for svgd in the population limit under talagrand's inequality t1. In International Conference on Machine Learning , pages 19139-19152. PMLR, 2022.

J. Shi and L. Mackey. A finite-particle convergence rate for stein variational gradient descent. arXiv preprint arXiv:2211.09721 , 2022.

I. Steinwart and A. Christmann. Support vector machines . Springer Science & Business Media, 2008.

L. Sun, A. Karagulyan, and P. Richtarik. Convergence of stein variational gradient descent under a weaker smoothness condition. In International Conference on Artificial Intelligence and Statistics , pages 3693-3717. PMLR, 2023.

M. Talagrand. Pettis integral and measure theory . American Mathematical Soc., 1984.

S. Vempala and A. Wibisono. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. Advances in neural information processing systems , 32, 2019.

R. Vershynin. High-dimensional probability: An introduction with applications in data science , volume 47. Cambridge university press, 2018.

C. Villani et al. Optimal transport: old and new , volume 338. Springer, 2009.

D. Wang and Q. Liu. Learning to draw samples: With application to amortized mle for generative adversarial learning. arXiv preprint arXiv:1611.01722 , 2016.

D. Wang, Z. Zeng, and Q. Liu. Stein variational message passing for continuous graphical models. In International Conference on Machine Learning , pages 5219-5227. PMLR, 2018.

M. Welling and Y. W. Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11) , pages 681-688, 2011.

A. Wibisono. Sampling as optimization in the space of measures: The langevin dynamics as a composite optimization problem. In Conference on Learning Theory , pages 2093-3027. PMLR, 2018.

J. Zhuo, C. Liu, J. Shi, J. Zhu, N. Chen, and B. Zhang. Message passing stein variational gradient descent. In International Conference on Machine Learning , pages 6018-6027. PMLR, 2018.

D. Zou, P. Xu, and Q. Gu. Faster convergence of stochastic gradient langevin dynamics for non-log-concave sampling. In Uncertainty in Artificial Intelligence , pages 1152-1162. PMLR, 2021.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a implementable stochastic approximation algorithm based on Stein Variational Gradient Descent, which the idea is based on a stochastic approximation at the probability measure level.  The efficiency and convergence analysis of the proposed algorithm, i.e. VP-SVGD and GB-SVGD,  are presented. A better convergence rate for the proposed algorithm is proved given a weaker assumption for the potential function\n\nSTRENGTHS:\nThe authors propose a new algorithm, a stochastic approximation at the probability measure level according to the authors, to study the SVGD. The focus is to improve the convergence rate and make the algorithm more efficient. The idea of the proposed algorithm in (3) seems to be new. The intuition and design of the algorithm is interesting, and it seems to work as shown in the numerical experiments.\n\nWEAKNESSES:\nThe reviewer feels that the intuition and the theoretical results have some gaps. To be more precise, some of the arguments are only explained at the intuition level without rigorous proof. It might be possible the reviewer missed the key point of the proof, if this is the case, please correct me.\n\nQUESTIONS:\n1. What is exactly the mild assumption for potential function F, which is weaker then Poincare and Talagrand inequality? How significantly different for this condition, e.g. regarding the proof ?\n\n2. Section 3 is not well presented. The push-forward map $(\\cdot)_{\\sharp}$ is not defined, but only mentioned, which is not clear. But this is a key step in designing the algorithm. \n\n3. Section 4, the notation for filtration is not good, $\\mathcal F_t$ should be measurable up to time $t$, but here it is measurable up to time $t-1$, which results the notation $\\mu_t|\\mathcal F_t$ confusing. \n\n4. Up this this point,  updates in (2) is not rigorously derived, but only intuitively mentioned. Then the algorithm (3) is based on (2), which is not clear to me why (3) is a stochastic approximation? It seems to be deterministic given the previous time step particle values. The same questions for $\\mu_t|\\mathcal F_t$, the conditional measure based on updated in (3) does not seem to be random. If there is a better way to present this idea, please clarify. \n\n5. The random function $g_t(\\mathbf X)$ is not clear, where does the randomness come? \n\n6. the mini batch size comes in Algorithm 1. which does not show before in section 4. There is a gap between the Algorithm 1 & 2  and (2)(3)(4). \n\n7. What is exactly the role of countably infinite many particles, as in the algorithm only $KT+n-1$ are used? I seems that this assumption is important, if this can be explained more clear, it should help improving the presentation.\n\nLIMITATIONS:\nsee above", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis study analyzed the finite particle behavior of Stein Variational Gradient Descent (SVGD) in an asymptotic manner. For that purpose, the authors designed two computationally efficient variants of SVGD, namely VP-SVGD and GB-SVGD, and prove their convergence rates. In those methods, the authors introduced a new concept of \"virtual particles\" and develop a new stochastic approximation of the population-limited SVGD dynamics in the probability measure space. The proposed algorithm shows a significant improvement of the convergence rates under  looser conditions than the usual SVGD.\n\nSTRENGTHS:\n- The convergence rates in this paper are significantly better than existing methods.\n\n- The idea of virtual particles introduced for the analysis is very novel and very interesting in itself.\n\n- Numerical experiments, although simple,  suggest that the proposed method is indeed an improvement over existing methods.\n\nWEAKNESSES:\n- Notation in Sec 4 is very confusing, especially about $l, s, t$\n\n-  In the explanation of Sec 4, only the case of K=1 is clearly stated, and I had to spend a lot of time thinking about what happens when K>2. I think an additional explanation of K > 2 is needed in the Appendix.\n\n- The algorithm of GB-SVGD is difficult to understand. In particular, the third line mentions \"l\", but \"l\" does not appear at all in the fourth line. This may be because the \"l\" in the third line may be a mistake for \"s\".\n\n- I am not sure if the theoretical comparison of the convergence order with the existing analysis is justified since the VP-SVGD is no longer the original SVGD. The GB-SVGD is close to the existing SVGD, but it is also difficult to compare with the existing SVGD because of the approximation errors as shown in Theorem 2.\n\nQUESTIONS:\n- What is $\\nabla_2$ in Line 143 ? Partial derivative with respect to the second variable?\n\n- Is my following understanding correct? Intuitively, the motivation for introducing virtual particles is to ensure that the resulting particles are i.i.d, which is required to use Lemma 7 in the theoretical analysis.\n\n- Algorithm 1 requires a larger number of particles in proportion to the time step T than ordinary SVGD, and therefore the number of gradient calculations increases considerably. To achive $KSD(particles,\\pi^*)\\leq \\epsilon$, does VP-SVGD increase the order of magnitude of the number of the gradient calculation compared to standard SVGD?\n\n- I was wondering why S is taken randomly in Algorithm 1 line 6 and Algorithm 2 line 7, etc., and I thought it would greatly reduce the practical performance. But after reading the proof, I understand that this is necessary to convert the bounds of $\\frac{1}{T}$KSD^2(\\mu_t|\\pi^*)$ to KSD(\\bar{\\mu}|\\pi^*}) as shown in Theorem 1. Is this understanding correct?\n\n- In the numerical experiment, are random $S$ used as written in TAlgorithm 1 and 2? Since the objective function is convex, such random S may not have such a strong effect on the final performance. I wonder how harmful this random $S$ is when the objective distribution is non-log concave?\n\nLIMITATIONS:\nTheoretical limits are discussed in detail.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies Stein Variational Gradient Descent (SVGD) which is a popular variational inference algorithm that simulates an interacting particle system to approximately sample from a target distribution. The authors propose two variants of SVGD, Virtual Particle SVGD (VP-SVGD) and Global Batch SVGD (GB-SVGD), that differ in their implementation details and computational efficiency. Both variants use the so-called \"virtual particles\", which are additional particles that evolve in time but are not part of the output, to compute information about the current population-level distribution of the real particles and enable exact implementation of the stochastic approximation using only a finite number of particles. The use of virtual particles allows for faster variants of SVGD with provably fast finite-particle convergence.\n\nSTRENGTHS:\nThis paper introduced two variants of SVGD, VP-SVGD and GB-SVGD, that use virtual particles to enable the exact implementation of the stochastic approximation using only a finite number of particles. The authors present a non-asymptotic analysis of the convergence of VP-SVGD and GB-SVGD to the target distribution, which demonstrates a double exponential improvement over the best-known finite-particle analysis of SVGD. The paper proposes a new algorithm, VP-SVGD, that achieves provably fast finite-particle convergence and is computationally more efficient than ordinary SVGD.\n\nThe paper is well-organized and clearly written, with detailed explanations of key concepts and technical details that make it accessible to a wide range of readers. Literature review of related work on SVGD and its variants was conducted comprehensively, which helps to contextualize the proposed algorithms and their contributions.\n\nI went through the proofs of Theorem 1 and Corollary 1 or 2 in Section E in detail, and believe they are technically sound. The KSD metric satisfies the triangle inequality where the upper bound is in two additional parts. In (34), the KSD metric satisfies a triangle inequality, and the first term (probability kernel vs population limit) is polynomial with a polynomial decaying exponent that is independent of dimension $d$ ($d^{1/3}T^{-1/6}$ when subgaussian). The second term appears dominating, where the polynomial decaying exponent in Wasserstein-1 distance is $\\Theta(1/d)$, which is due to Lemma 8 by Lei [29].\n\nWEAKNESSES:\nDespite its significant contributions, this work assumes a certain level of familiarity with the concepts of variational inference and particle methods and is relatively specialized and may make it difficult for readers who are not already familiar with or not working on the topic of SVGD and its variants to follow.  Despite the technical contributions of proposed SVGD-style algorithms and their convergence properties, the authors did not provide a detailed discussion of the limitations of the proposed algorithms or potential avenues for future research, which may limit their impact and relevance in the long term it and make it less accessible to readers who are more interested in the broader implications of the work. Lastly, the work does not provide sufficient details on the proposed algorithms' empirical evaluations of real-world datasets, which limits the ability to assess their practical usefulness. These weaknesses did affect my current rating, but I reserve the possibility to raise the score if any of the above points can be resolved in the authors' rebuttal.\n\nQUESTIONS:\nI am not a leader in the topic of SVGD but had a chance to work on the topic at a senior educative level. Hence I recognize the importance of this topic. Some questions are as follows:\n\nHow do the VP-SVGD and GB-SVGD variants of SVGD differ from the original SVGD introduced by Liu and Wang [32]?\nMy understanding is that the key difference between VP-SVGD and GB-SVGD variants of SVGD is that VP-SVGD is a stochastic approximation of the exact SVGD dynamics, while GB-SVGD is a specific random-batch approximation of SVGD which enjoys computationally more efficient than ordinary SVGD. Is that correct?\n\nCan you explain more about the concept of virtual particles and how they are used in these algorithms? I understand that virtual particles are additional particles that evolve in time but are not part of the output (i.e., real particles) which are used only to compute information about the current population-level distribution of the real particles. Are they absolutely necessary --- are there any equivalent or similar implementations that do not necessarily require the introduction of virtual particles?\n\nIn the technical analysis of Corollary 2, is there any hope to make the decaying exponent in polynomial rate dimension-free or have a weaker dependence on dimension (perhaps using an alternative metric or developing new analytic tools to yield an improved bound)?\n\nLIMITATIONS:\nThe work is purely theoretical and admits no negative social impacts.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper propose two variants of stein variational gradient descent (SVD) named as Virtual Particle SVGD (VP-SVGD) and Global Batch SGVD (GB-SVGD) with provably fast finite-particle convergence rates to a target distribution. Introducing the notion of virtual particles, the authors show that VP-SVGD is an exactly implementable finite-particle stochastic approximation of population-limit SVGD dynamics in the space of probability measures. Viewing VP-SVGD as a specific random-batch SVGD algorithm, the authors argue that  the output particles from GB-SVGD after T time steps with batch-size K are as good as sampling iid from a distribution whose Kernel Stein Discrepancy (KST) to a sub-gaussian target distribution is at most $O(d^{1/3}/(KT)^{1/6})$. Furthermore, under mild growth conditions, the authors show that the empirical distribution of the n-output particles from both VP-SVGD and GB-SVGD converges to the target in KSD at a rate of $O(d^{2}/(n)^{\\Theta(1/d)})$. The authors note that this is a doubly exponential improvement over the state-of-the-art finite-particle analysis of SVGD under weaker assumptions. Finally, the authors provide numerical experiments to show fast convergence of GB-SVGD which is computationally more efficient than SVGD.\n\nSTRENGTHS:\n1)Originality: This paper is introduces two new variants of SVGD appealing to a novel construction involving countably many virtual particles.  Although the analysis involves infinitely many particles, a finite-number of particles exactly suffice for the actual algorithm VP-SVGD. This enables the authors to claim that VP-SVGD performs a finite-particle stochastic approximation of the population (infinite-particle) dynamics of SVGD over the space of probability distributions. Finally, they show VP-SVGD and the computationally efficient GB-SVGD behave very closely.\n\n2)Quality: I think this paper is technically quite solid. The theoretical claims are all sound and the assumptions are reasonable.\n\n3)Clarity: This paper is well written. It is easy to follow the main ideas including the problem statement, technical challenges, relevance to past works, and the contributions.\n\n4)Significance: It is quite significant to improve the SOTA result doubly exponentially under weaker assumptions. Besides, the virtual particle method is worthwhile studying by itself ( I am not aware of a similar method discussed somewhere else).\n\nWEAKNESSES:\nIt wasn't easy for me to parse some of the theoretical results, which can be expected as this paper is theory heavy.\n\nA table of past works comparing the results in terms of assumptions, method, and the convergence rate might be a helpful addition.\n\nQUESTIONS:\nI suspect there are a few typos in the text.\n\n1)bold x with varying subscript  in line 249 vs bold x with varying superscript in the equation after line 250. also the index runs to T but I would expect T-1.\n\n2)bold x instead of bold y in line 224\n\n3) In line 204, we are already given that $g_s(x) = ...$ and then in line 306 we are told that $g_s(x) \\approx ...$, Is this a typo ?\n\nLIMITATIONS:\nLimitations and future works are discussed.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-19"}