A Finite-Particle Convergence Rate for Stein Variational Gradient Descent
Abstract
We provide a first finite-particle convergence rate for Stein variational gradient descent (SVGD).
Specifically, whenever the target distribution satisfies Talagrand’s inequality, SVGD with n particles and an appropriate step size sequence drives the kernel Stein discrepancy to zero at an order 1 / √ log log n rate.
We suspect that the dependence on n can be improved, and we hope that our explicit, non-asymptotic proof strategy will serve as a template for future refinements
1. Introduction
Stein variational gradient descent [SVGD, 11] is an algorithm for approximating a target distribution P with a collection of n particles.
The algorithm iteratively evolves the particle locations to approach the target distribution through optimization in the space of probability measures.
It has demonstrated encouraging results for a wide variety of inferential tasks, including approximate inference [11, 18, 19], generative modeling [17], and reinforcement learning [8, 13].
Despite the popularity of SVGD, relatively little is known about its convergence.
A first analysis by Liu [10, Thm. 3.3] showed that an infinite-particle analogue of SVGD converges to P in kernel
Stein discrepancy [KSD, 2, 6, 12].
Follow-up work by Korba et al.
[9], Salim et al.
[14], Sun et al.
[15] sharpened this result with path-independent constants, weaker smoothness conditions, and explicit rates of convergence.
In addition, Duncan et al.
[4] analyzed the continuous-time limit of the infinite-particle analogue to provide conditions for exponential convergence.
However, each of these analyses applies only to the infinite-particle analogue of SVGD and not to the finite-particle algorithm used in practice.
To bridge this gap, Liu [10, Thm. 3.2] showed that n -particle SVGD converges to its infiniteparticle analogue in bounded-Lipschitz distance but only under boundedness assumptions violated by most common uses of SVGD.
To provide a more broadly applicable proof of convergence,
Gorham et al.
[7, Thm. 7] showed that n -particle SVGD converges to its infinite-particle analogue in 1 -Wasserstein distance under assumptions commonly satisfied in SVGD applications.
However, both convergence results are asymptotic, providing neither explicit error bounds nor rates of convergence.
In this work, we derive the first explicit error bounds and rates of convergence for finite-particle
SVGD.
To achieve this, we first bound the 1 -Wasserstein error between each round of finite- and infinite-particle SVGD in Theorem 1.
We next bound the KSD in terms of the 1 -Wasserstein distance in Theorem 2.
Finally, Theorem 3 combines our results with the established infinite-particle
KSD analysis of SVGD to arrive at an explicit KSD error bound for n -particle SVGD.
2. Notation and Assumptions
Throughout, we fix a reproducing kernel k : R d × R d → R , let H denote its induced reproducing kernel Hilbert space (RKHS), and let H d ≜ (cid:78) di =1 H denote the product RKHS [1].
For probability measures µ and ν , we let W 1 ( µ, ν ) ≜ inf X ∼ µ,Z ∼ ν E [ ∥ Z − X ∥ 2 ] denote the 1 -Wasserstein distance between them and introduce the shorthand m µ ≜ E µ [ ∥·∥ ] 2 , m µ,P ≜
E X ∼ µ ⊥⊥ Z ∼ P [ ∥
X − Z ∥ 2 ] , and
M µ,P ≜ E X ∼ µ ⊥⊥ Z ∼ P [ ∥
X − Z ∥ 22 ] .
We allow each of these quantities to take on the value ∞ when the random variables are not suitably integrable.
We further define the Kullback-Leibler divergence as KL ( µ ∥ ν ) ≜ E µ [log( dµdν )] when µ is absolutely continuous with respect to ν and as ∞ otherwise.
Our analysis will make use of the following standard assumptions on the SVGD kernel and target distribution.
Assumption 1 (Bounded kernel derivatives) sup x ∈ R d (cid:0) D
Ix D Iy k ( x, y ) | y
=
x (cid:1) ≤ κ 2 for any multiindex
I =
( I 1 , I 2 , . .
, I d ) with | I | ≜ (cid:80) di
=
I i ≤ 2 .
Here, D I is the differential operator defined as
[Equation]
Under Assumption 1, we could show that for any f ∈ H ,
[Equation]
Assumption 2 (Decaying kernel derivatives)
There exists γ > 0 such that
[Equation]
Assumption 3 (Lipschitz score function)
The target distribution P has a differentiable density p on R d and
an L -Lipschitz score function s p ≜ ∇ log p , i.e., ∥ s p ( x ) − s p ( y ) ∥ 2 ≤ L ∥ x − y ∥ 2 for any x ̸ = y .
We let x ∗ be any point that satisfies s p ( x ∗ ) = 0 .
In line with Salim et al.
[14], we also assume the distribution P satisfies Talagrand’s inequality
T1 [16, Def. 22.1].
We note that this assumption is only necessary for leveraging the infinite-particle convergence result of Salim et al.
In case it does not hold, one can alternatively resort to the path-dependent assumptions in Korba et al.
[9], Liu [10] .
Assumption 4 (Talagrand’s inequality T1 [16, Def. 22.1])
There exists λ > 0 such that for all µ with E µ [ ∥·∥ 2 ] < ∞
[Equation]
Finally we make use of the following notation specific to the SVGD algorithm.
Definition 1 The Langevin Stein operator [5] for a target measure P supported on R d is defined by
[Equation]
for any vector-valued function g : R d → R d .
Definition
The vector-valued Langevin Stein operator [12] for a target measure P on R d is defined by
[Equation]
with components
[Equation]
for any function h : R d → R .
Definition
The SVGD transport map [11] for a kernel k , target measure P , approximating measure µ , and step size ϵ is given by
[Equation]
Definition
The SVGD pushforward Φ ϵ ( µ ) represents the distribution of T µ,ϵ ( X ) when X ∼ µ .
Definition
The Langevin kernel Stein discrepancy [KSD, 2, 6, 12] with base kernel k , target measure P , and approximating measure µ is given by
[Equation]
3. Discretization Error of SVGD
Our first main result concerns the discretization error of SVGD and shows that n -particle SVGD does not stray too far from its infinite-particle limit when its step sizes are chosen appropriately.
Theorem 1 (Discretization error of SVGD) Suppose Assumptions 1, 2, and 3 hold.
Let µ n r , µ ∞ r denote the approximation given by n -particle and infinite-number-of-particle SVGD with step size ϵ r at step r .
And let b r ≜ (cid:80) rs =
ϵ s .
Then
[Equation]
n 0 , µ ∞ 0 ) exp( b r ( A + B exp( Cb r ))) , for A
= ( c 1 + c 2 )(1 + m P ) , B = c
1 m µ n 0 ,P + c 2 m µ ∞ 0 ,P , and C = κ 2 (3 L + d ) .
Here c 1 = max( √ dκ 2 L, √ dκ 2 L
∥ x ∗ ∥ 2 + dκ 2 ) , c 2 = κ 2 L + dκ 2 + ( γ + √ dκ 2 )(1 + ∥ x ∗ ∥ 2 ) .
The proof of Theorem 1 in Appendix A relies on two lemmas.
The first, due to Gorham et al.
[7], shows that the one-step SVGD mapping Φ ϵ is pseudo-Lipschitz with respect to the 1 -Wasserstein distance whenever the score function ∇ log p and kernel k fulfill a commonly-satisfied pseudoLipschitz condition.
Lemma 1 (Wasserstein pseudo-Lipschitzness of SVGD [7, Lem. 12]) Suppose that, for some constants c 1 , c 2 > 0 , the following pseudo-Lipschitz bounds hold:
[Equation]
Then,
[Equation]
where c µ,ν =
c 1 (1 + m µ ) + c 2 (1 + m ν ) .
The second lemma, proved in Appendix B, controls first and second absolute moment growth under SVGD using P as a reference.
Lemma 2 (SVGD moment growth) Suppose Assumptions 1 and 3 hold.
Let µ r denote the approximation given by SVGD at step r .
Then,
[Equation]
for C = κ 2 (3 L + d ) .
4. From Wasserstein to KSD
The above infinite-particle SVGD convergence rate is computed with respect to KSD.
To relate it to the Wasserstein discretization error bound in Theorem 1, we establish the following KSDWasserstein bound.
Theorem 2 (KSD-Wasserstein bound)
Suppose Assumptions 1 and 3 hold.
Let µ nr , µ ∞ r denote the approximation by n particle and infinitely many particle SVGD at iteration r .
Then,
[Equation]
Proof Using the result of Lemma [Equation] Cauchy-Schwarz, Eq.
(1) of Lemma [Equation] and min( a, b ) 2 ≤ ab [Equation]
we have
[Equation]
Taking the infinimum over the joint distribution of ( X, Y ) leads to
[Equation]
The proof of Theorem 2 uses the following lemma, a generalization of Gorham and Mackey [6,
Lem.
18] that we prove in Appendix C, which shows that the mean discrepancy between two T P transformed functions is controlled by the 1 -Wasserstein distance.
Lemma 3 (Single function Stein discrepancy bound) Suppose Assumptions 1 and 3 hold.
Let µ, ν be two probability measures on R d and g ∈ H d satisfying ∥ g ∥ 2 H d ≜ (cid:80) di =1 ∥
g i ∥ 2 H ≤ 1 .
Then,
[Equation]
5. A Finite-particle Convergence Rate for SVGD
To establish our main SVGD convergence result, we combine Theorems 1 and 2 with the following infinite-particle SVGD convergence rate due to Salim et al.
[14] (see also Korba et al.
[9], Liu [10]).
Lemma 4 (Infinite-particle SVGD descent lemma [14, Theorem 3.2]) Suppose Assumptions 1, 3, and 4 hold.
For any constant α > 1 and ϵ r < min (cid:16) 2 κ 2 ( L + α 2 ) , R (cid:17) , where R ≜ ( α − 1)(1 +
L E [Equation] ∥· [-− x ∗ ∥ 2 ] + 2 L (cid:112) 2 KL ( µ ∞ 
-][Equation]
∥ 2 ]
[Equation]
L (cid:112)
[Equation] ( µ [Equation] 0 ∥ P ) /λ ) , SVGD with infinite number of particles decreases KL and drives KSD to zero:
[Equation]
By summing (2) over r = 0 , . . .
, t , we have the following corollary:
Corollary 1 Let c ( ϵ ) ≜
ϵ (cid:16) 1 − [Equation] [Equation]
If [Equation] r ≤ min (cid:16) 1 κ
[Equation]
+ α
[Equation]
(cid:17) , then we have ϵ r 2 ≤ c ( ϵ r ) <
ϵ r and thereby
[Equation]
Finally, we arrive at our main result that shows finite-particle SVGD with appropriately chosen step size sequences drives KSD to zero as W 1 ( µ n 0 , µ ∞ 0 ) → 0 .
Theorem 3 (KSD error of finite-particle SVGD) Suppose Assumptions 1, 2, 3, and 4 hold.
Define π r ≜ c [Equation] r =0 c ( ϵ r ) and w 0 ≜ W 1 ( µ n 0 , µ ∞ 0 ) .
Let b t
=
(cid:80) tr =0 ϵ t .
Then,
[Equation]
where a t = κ ( d + L ) w 0 exp( b t − 1 ( A + B [Equation] A + B exp
[Equation]
Cb t − 1 )) / 2) b t − 1 )
Proof Let w r ≜
W 1 ( µ nr , µ ∞ r ) .
From Theorem 2, we know
[Equation]
Plugging in Theorem 1’s result
[Equation]
we have KSD ( µ nr , P ) − a r ≤ KSD ( µ ∞ r , P ) , where
[Equation]
Taking square on both sides and the weighted sum over r = 0 , . .
, t − 1 with weights π r ≜ c ( ϵ r ) (cid:80) t −r =
c ( ϵ r ) , we obtain
[Equation]
[Equation]
where the last inequality follows from Corollary 1.
Moreover, by Cauchy-Schwarz,
[Equation]
Combining (4) and (5), we have
[Equation]
We finish the proof by noticing that (cid:80)
[Equation]
The following corollary, proved in Appendix D, provides an explicit rate of SVGD by choosing the step size sequence and t such that the resulting b t − 1 balance the three terms on the right hand side of (3).
Corollary 2 (A finite-particle convergence rate for SVGD)
Define the shorthand ψ B,C ( x, y, β ) ≜ 1 C log (cid:16) 1 B ( 1 β log 1 x − y ) (cid:17) .
If max 0 ≤ r<t ϵ r ≤ min (cid:16) 1 κ 2 ( L + α 2 ) , R (cid:17) and b t − 1 = (cid:80) t − 1 r =0 ϵ r = min (cid:16) ψ B,C (cid:0) w 0 (cid:113) log log 1 w 0 , A, β 1 (cid:1) , ψ B,C (cid:0) w 0 log log 1 w 0 , A + 2 C + ϵ 0 C 2 , β 2 (cid:1)(cid:17) , for β 1 ≜ max (cid:16) 1 , ψ B,C (cid:0) w 0 (cid:113) log log 1 w 0 , A, 1 (cid:1)(cid:17) and β 2 ≜ max (cid:16) 1 , ψ B,C (cid:0) w 0 log log 1 w 0 , A + 2 C + ϵ 0 C 2 , 1 (cid:1)(cid:17) , then, min
[Equation]
[Equation]
If, in addition, µ n 0 = 1 n {+(cid:80) ni =
+} δ x i for x i i.i.d.
∼ µ ∞ 0 , then min
[Equation]
KSD
[Equation]
Appendix A. Proof of Theorem 1: Discretization error of SVGD
We first show that the pseudo-Lipshitz conditions in Lemma 1 holds given our assumptions.
According to Assumption 1, we have
[Equation]
Recall that x ∗ satisfies s p ( x ∗ ) = 0 by Assumption 3.
Then
[Equation]
by triangle inequality
[Equation]
the definition of ∥·∥ op
[Equation]
and Cauchy-Schwartz,
[Equation]
Letting c 1 = max( √ dκ 2 L
[Equation]
√ dκ 2 L ∥ x ∗ ∥ 2 + dκ 2 ) and taking supremum over z proves the first pseudo-Lipschitz condition.
Similarly
[Equation]
we have
[Equation]
where we used
[Equation]
The last inequality is due to
[Equation]
According to Assumption 2, there exists γ > 0 such that ∥∇ x k ( x, z ) ∥ 2 ≤ γ ∥ x − z
∥ − 1 2 .
Therefore, for ∥ x − z ∥ 2 ≥ 1 ,
[Equation]
And for ∥ x − z ∥ 2 < 1 ,
[Equation]
Combining (7) and (8),
[Equation]
Plugging (9) back into (6), we can show the second pseudo Lipschitzness condition holds for c 2 = κ 2 L + dκ
[Equation]
+
[Equation]
+
[Equation]
Now that the pseudo Lipschitz conditions hold, by repeated application of Lemma 1 and the inequality (1 + x ) ≤ e x , we have
[Equation]
for C s = c 1 (1 + m µ ns ) + c 2 (1 + m µ ∞ s ) .
Using the result from Lemma 2, we have
[Equation]
for A = ( c 1 + c 2 )(1 + m P ) and B = c 1 m µ n 0 ,P + c 2 m µ ∞ 0 ,P .
Therefore
[Equation]
Plugging this back into (10) proves the result.
Appendix B. Proof of Lemma 2: SVGD moment growth
From Assumption 1 we know
[Equation]
Let µ be any probability measure.
Using Jensen’s inequality and the facts that E P [ s p ( · )]
= 0 and
E Z ∼ P [( A P k ( · , x ))
( Z )]
= 0 , we have ∥
[Equation]
Let m
[Equation]
r ,
[Equation]
Z ∼ P [ ∥
X − Z ∥ 22 ] , and C = κ 2 (3 L + d ) .
Then,
[Equation]
where the last step used Jensen’s inequality m 2 µ r
[Equation]
P ≤ M µ r ,P .
Then
[Equation]
The final inequality uses the relation 1 + x ≤ e x .
Appendix C. Proof of Lemma
From Assumption 1 we know ∥ g ( x ) ∥ 22 ≤ (cid:80) di
=1 κ 2 ∥ g i ∥
2 H ≤ κ 2 , ∥∇ g ( x ) ∥ 2 op ≤ ∥∇ g ( x ) ∥ 2 F = (cid:80) d
i =1 (cid:80) d j =1 |∇ x i g j ( x ) | 2 ≤ (cid:80) d i =1 (cid:80) d j =1 κ 2 ∥ g j ∥ 2 H ≤ dκ 2 , ∥∇ ( ∇ · g ( x ))
∥
22 = (cid:80) di
=1 (cid:16)(cid:80)
dj
=1 ∇ x i ∇ x j g j ( x ) (cid:17) 2 ≤ d (cid:80) di
=1 (cid:80) dj
=1 |∇ x i ∇ x j g j ( x ) | 2 ≤ d 2 κ 2 .
Suppose X, Y, Z are distributed so that ( X, Y ) is a 1 -Wasserstein optimal coupling of ( µ, ν ) and Z is independent of ( X, Y ) .
Since E P [ s p ( · )] = 0 , repeated use of Cauchy-Schwarz gives
[Equation]
Appendix D. Proof of Corollary 2: A finite-particle convergence rate for SVGD
We first state and prove a useful lemma.
Lemma 5 Suppose f is a monotonic decreasing function.
Let x = f ( β ) for β = max(1 , f (1)) .
Then x ≤ f ( x ) .
Proof If β = 1 ≥ f (1) , then x = f ( β ) =
f (1) ≤ 1 .
Because f is monotonic decreasing, f ( x ) ≥ f (1) = x .
Otherwise, β = f (1) > 1 and hence x
= f ( β ) <
f
(1) =
β .
Again, due to monotonicity, f ( x ) > f ( β ) = x .
Combining the two cases finishes the proof.
Now we are ready to prove Corollary 2.
Let b (1) t − 1 = ψ B,C ( w 0 (cid:113) log log 1 w 0 , A, β 1 ) and b (2) t − 1 = ψ B,C ( w 0 log log 1 w 0 , A + 2 C + ϵ 0 C 2 , β 2 ) .
According to Lemma 5,
[Equation]
b
[Equation]
Rearranging the terms, we have
[Equation]
Similarly, we have b (2) t − 1 ≤ ψ B,C ( w 0 log log 1 w 0 , A + 2 C + ϵ 0 C 2 , b (2) t − 1 ) and √ w 0 exp
[Equation]
b
[Equation]
Since b t − 1 = min( b (1) t − 1 , b (2) t − 1 ) , we can substitute b t − 1 for b (1) t − 1 and b (2) t − 1 in (11) and (12), respectively.
Also, from the definition of b t − 1 we know
[Equation]
Plugging them into (3) proves the first claim.
The second claim follows from the Θ( n − 1 /d ) convergence rate of an empirical measure to its population measure in 1 -Wasserstein distance [3].