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 T 1 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 probability distribution P on R d with a collection of n particles.
Given an initial particle approximation µ n 0 = 1 n (cid:80) ni =
δ x i with locations x i ∈ R d , SVGD (Algorithm 1) iteratively evolves the particle locations to provide a more faithful approximation of the target P by performing optimization in the space of probability measures.
SVGD 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].
Algorithm 1 Stein Variational Gradient Descent [11]: SVGD ( µ 0 , r )
Input: Target density p , kernel k , step sizes ( ϵ s ) s ≥ 0 , approximating distribution µ 0 , rounds r for s = 0 , · · · , r − 1 do
Let µ s +1 be the distribution of X s + ϵ s (cid:82)
k
( x, X s ) ∇ log p ( x )+ ∇ x k ( x, X s ) dµ s ( x ) for X s ∼ µ s .
Output: Updated approximation µ r of the target P
Despite the popularity of SVGD, relatively little is known about its approximation quality.
A first analysis by Liu [10, Thm. 3.3] showed that continuous SVGD —that is, Algorithm 1 initialized with a continuous distribution µ ∞ 0 in place of the discrete particle approximation µ n 0 —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 continuoustime limit of continuous SVGD to provide conditions for exponential convergence.
However, each of these analyses applies only to continuous 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 continuous
SVGD in bounded-Lipschitz distance but only under boundedness assumptions violated by most applications of SVGD.
To provide a more broadly applicable proof of convergence, Gorham et al.
[7, Thm. 7] showed that n -particle SVGD converges to continuous SVGD 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-particle and continuous 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 KSD analysis of continuous SVGD to arrive at an explicit KSD error bound for n -particle SVGD.
2. Notation and Assumptions
Throughout, we fix a target distribution P in the set P of probability measures on R d , a nonnegative step size sequence ( ϵ s ) s ≥ 0 , and a reproducing kernel k : R d × R d → R with reproducing kernel Hilbert space (RKHS) H and product RKHS H d ≜ (cid:78)
di
=1 H [1].
For all µ, ν ∈ P , 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 ν , denoted by µ ≪ ν , and as ∞ otherwise.
Our analysis will make use of the following standard assumptions on the SVGD kernel and target distribution.
Assumption 1 (Lipschitz, mean-zero score function)
The target distribution P ∈ P has a differentiable density p with an L -Lipschitz score function s p ≜ ∇ log p , i.e., ∥ s p ( x ) − s p ( y ) ∥ 2 ≤
L ∥ x − y ∥ 2 for all x, y ∈ R d .
Moreover, E P [ s p ] = 0 and s p ( x ⋆ ) = 0 for some x ⋆ ∈ R d .
Assumption 2 (Bounded kernel derivatives)
For any multi-index I
=
( I 1 , I 2 , . .
, I d ) with | I | ≜ (cid:80) di
=
I i ≤ 2 , we have sup x ∈ R d (cid:0)
D Ix D Iy k ( x, y ) | y
=
x (cid:1) ≤ κ 2 .
Here, D I is the differential operator defined as D Ix =
[Equation]
Assumption 3 (Decaying kernel derivatives)
There exists γ > 0 such that
[Equation]
In line with Salim et al.
[14], we also assume the distribution P satisfies Talagrand’s T 1 inequality [16, Def. 22.1].
We note that this assumption is only necessary for leveraging the continuous
SVGD 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 T 1 inequality [16, Def. 22.1])
There exists λ > 0 such that for all µ with E µ [ ∥·∥
[Equation]
Finally we make use of the following notation specific to the SVGD algorithm.
Definition 1 (Stein operator) For any vector-valued function g : R d → R d , the Langevin Stein operator [5] for P satisfying Assumption 1 is defined by
[Equation]
Definition 2 (Vector-valued Stein operator)
For any function h : R d → R , the vector-valued
Langevin Stein operator [12] for P satisfying Assumption 1 is defined by
[Equation]
with components
[Equation]
Definition 3 (SVGD transport map and pushforward)
The SVGD transport map [11] for a kernel k , target P satisfying Assumption 1, approximating µ ∈ P , and step size ϵ is given by
[Equation]
Moreover, the SVGD pushforward Φ ϵ ( µ ) represents the distribution of T µ,ϵ ( X ) when X ∼ µ .
Definition 4 (Kernel Stein discrepancy) For measures µ, ν ∈ P and target P satisfying Assumption 1 we define the Langevin kernel Stein discrepancy [KSD, 2, 6, 12] with base kernel k by
[Equation]
3. Wasserstein 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 continuous SVGD limit when its step sizes are chosen appropriately.
Theorem 1 (Wasserstein discretization error of SVGD) Suppose Assumptions 1, 2, and 3 hold.
For any µ n 0 , µ ∞ 0 ∈ P , the Algorithm 1 outputs µ n r = SVGD ( µ n 0 , r ) and µ ∞ r = SVGD ( µ ∞ 0 , r ) satisfy
W 1 ( µ nr , µ ∞ r ) ≤ W 1 ( µ n 0 , µ ∞ 0 ) exp( b r − 1 ( A + B exp
( Cb r − 1 ))) for b r − 1 ≜ (cid:80) r − 1 s =
ϵ s , 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 ) and 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 pushforward Φ ϵ (Definition 3) is pseudo-Lipschitz with respect to the 1 -Wasserstein distance whenever the score function ∇ log p and kernel k fulfill a commonlysatisfied pseudo-Lipschitz condition.
Lemma 1 (Wasserstein pseudo-Lipschitzness of SVGD [7, Lem. 12])
For P ∈ P with differentiable density p , suppose that the following pseudo-Lipschitz bounds hold
[Equation]
for some constants c 1 , c 2 > 0 .
Then, for any µ, ν ∈ P ,
[Equation]
where Φ ϵ is the one-step SVGD pushforward (Definition 3) and c µ,ν
=
c 1 (1 + m µ ) + c 2 (1 + m ν ) .
The second lemma, proved in Appendix B, controls the growth of the first and second absolute moments under SVGD.
Lemma 2 (SVGD moment growth) Suppose Assumptions 1 and 2 hold, and let C = κ 2 (3 L + d ) .
Then the SVGD output µ r of Algorithm 1 with b r − 1 ≜ (cid:80) r − 1 s =
ϵ s satisfies
[Equation]
4. KSD Discretization Error of SVGD
Our next result translates the Wasserstein error bounds of Theorem 1 into KSD error bounds.
Theorem 2 (KSD discretization error of SVGD) Suppose Assumptions 1 and 2 hold.
For any µ n 0 , µ ∞ 0 ∈ P , the Algorithm 1 outputs µ nr = SVGD ( µ n 0 , r ) and µ ∞ r = SVGD ( µ ∞ 0 , r ) satisfy
[Equation]
for w 0 ,n ≜ W 1 ( µ n 0 , µ ∞ 0 ) and A, B, C defined as in Theorem 1.
Our proof of Theorem 2 relies on the following lemma, proved Appendix C, that shows that
the
KSD is controlled by the 1 -Wasserstein distance.
Lemma 3 (KSD-Wasserstein bound)
Suppose Assumptions 1 and 2 hold.
For any µ, ν ∈ P ,
KSD
[Equation]
Proof of Theorem
The result follows directly from Lemma 3, Lemma 2, and Theorem 1.
■
5. A Finite-particle Convergence Rate for SVGD
To establish our main SVGD convergence result, we combine Theorems 1 and 2 with the following descent lemma for continuous SVGD error due to Salim et al.
[14] which shows that continuous
SVGD decreases the KL divergence to P and drives the KSD to P to zero.
Lemma 4 (Continuous SVGD descent lemma [14, Thm. 3.2]) Suppose Assumptions 1, 2, andhold, and consider the outputs
µ ∞ r = SVGD ( µ ∞ 0 , r ) and µ ∞ r +1 =
SVGD ( µ ∞ 0 , r + 1) of Algorithmwith µ ∞ 0 ≪ P .
If max 0 ≤ s ≤ r ϵ s < R α, 2 for some α > 1 and
R α,p
[Equation]
then
[Equation]
By summing the result (1) over r = 0 , . . .
, t , we obtain the following corollary.
Corollary 1 Under the assumptions and notation of Lemma 4, suppose max 0 ≤ r ≤ t ϵ r ≤ R α, 1 for some α > 1 , and let π r ≜
[Equation]
ϵ r ) for c ( ϵ ) ≜ ϵ (cid:16) 1 − κ 2 ( L + α 2 )ϵ (cid:17) .
Since ϵ 2 ≤ c ( ϵ ) <
ϵ , we have
[Equation]
Finally, we arrive at our main result that bounds the approximation error of n -particle SVGD in terms of the chosen step size sequence and the initial discretization error W 1 ( µ n 0 , µ ∞ 0 ) .
Theorem 3 (KSD error of finite-particle SVGD) Suppose Assumptions 1, 2, 3, and 4 hold, fix any µ ∞
≪ P and µ n 0 ∈ P , and let w 0 ,n ≜ W 1 ( µ n 0 , µ ∞ 0 ) .
If max 0 ≤ r ≤ t ϵ r ≤ R α, 1 for α > 1 and
R α, 1 defined in Lemma 4, then the Algorithm 1 outputs µ nr = SVGD ( µ n 0 , r ) satisfy
[Equation]
for b t ≜ (cid:80) tr =0 ϵ r , π r as defined in Lemma 4, ( A, B, C ) as defined in Theorem 1, and
[Equation]
Proof By the triangle inequality and Theorem 2 we have
[Equation]
for each r .
Therefore
[Equation]
where the last inequality follows from Corollary
[Equation]
Moreover, by Jensen’s inequality,
[Equation]
Combining (3) and (4), we have
[Equation]
We finish the proof by noticing that (cid:80) tr =0 π r a r − 1 ≤ max 0 ≤ r ≤ t a r − 1 = a t − 1 ≤ a t .
The following corollary, proved in Appendix D, provides an explicit SVGD convergence rate by choosing the step size sequence to balance the three terms on the right-hand side of (2).
Corollary 2 (A finite-particle convergence rate for SVGD) Instantiate the notation and assumptions of Theorem 3, and define the shorthand ψ B,C ( x, y, β ) ≜ 1 C log( 1 B ( 1 β log 1 x − y )) .
If b t = (cid:80) tr =0 ϵ r = min (cid:16) ψ B,C (cid:0) w 0 ,n (cid:113) log log 1 w 0 ,n , A, β 1 (cid:1) , ψ B,C (cid:0) w 0 ,n log log 1 w 0 ,n , A + 2 C + ϵ 0 C 2 , β 2 (cid:1)(cid:17) , for β 1 ≜ max (cid:16) 1 , ψ B,C (cid:0) w 0 ,n (cid:113) log log 1 w 0 ,n , A, 1 (cid:1)(cid:17) and β 2 ≜ max (cid:16) 1 , ψ B,C (cid:0) w 0 ,n log log 1 w 0 ,n , A + 2 C + ϵ 0 C
[Equation]
1 (cid:1)(cid:17) , then,
[Equation]
If, in addition, µ n 0 = 1 n (cid:80) ni =
δ x i for x i i.i.d.
∼ µ ∞ 0 , then min 0 ≤ r ≤ t KSD ( µ nr , P ) =
O (cid:16) 1 √ log log n (cid:17) .
Appendix A. Proof of Theorem 1: Wasserstein discretization error of SVGD
We first show that the pseudo-Lipshitz conditions in Lemma 1 holds given our assumptions.
According to Assumption 2, we have
[Equation]
Recall that x ⋆ satisfies s p ( x ⋆ ) = 0 by Assumption 1.
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
[Equation]
there exists γ > 0 such that ∥∇ x k ( x, z ) ∥ 2 ≤ γ ∥ x − z ∥ − 1 2 .
Therefore, for ∥ x − z ∥ 2 ≥ 1 ,
[Equation]
A
And for ∥ x − z ∥ 2 < 1 ,
[Equation]
Combining (6) and (7),
[Equation]
Plugging (8) back into (5), we can show the second pseudo Lipschitzness condition holds for c 2 = κ 2 L +
[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 (9) proves the result.
Appendix B. Proof of Lemma 2: SVGD moment growth
From Assumption 2 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]
µ,
[Equation]
Let m
[Equation] [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
[Equation]
P .
Then
[Equation]
The final inequality uses the relation 1 + x ≤ e x .
Appendix C. Proof of Lemma 3: KSD-Wasserstein bound
Our proof generalizes that of Gorham and Mackey [6, Lem. 18].
Consider any g ∈ H d satisfying ∥ g ∥ 2 H d ≜ (cid:80) di =1 ∥ g i ∥ 2 H ≤ 1 .
From Assumption 2 we know ∥ g (
[Equation]
∥ 22 ≤
[Equation]
=
[Equation]
κ
[Equation]
g
[Equation]
≤ κ
[Equation]
, (10) ∥∇ g ( x ) ∥ 2 op ≤ ∥∇ g ( x ) ∥ 2 F = (cid:80) di
=1 (cid:80) dj
=1 |∇ x i g j ( x ) | 2 ≤ (cid:80) di
=1 (cid:80) dj
=1 κ 2 ∥ g j ∥ 2 H ≤ dκ 2 , and (11) ∥∇ ( ∇ · 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 s p is L -Lipschitz with E P [ s p ] =
(Assumption 1), g is bounded (10), and g and ∇ · g are Lipschitz (11), repeated use of Cauchy-Schwarz gives
[Equation]
Since our choice of g was arbitrary, the first advertised result now follows from the definition of KSD (Definition 4).
The second claim then follows from Cauchy-Schwarz and the inequality min( a, b ) 2 ≤ ab for a, b ≥ 0 , since
[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 = ψ B,C ( w 0 ,n (cid:113) log log 1 w 0 ,n , A, β 1 ) and b (2) t = ψ B,C ( w 0 ,n log log 1 w 0 ,n , A + 2 C + ϵ 0 C 2 , β 2 ) .
According to Lemma 5,
[Equation]
Rearranging the terms, we have
[Equation]
[Equation]
Similarly, we have b (2) t ≤
[Equation] [Equation]
[Equation]
log log
[Equation]
0 ,n , A +
[Equation]
+ ϵ
[Equation]
, b (2) t ) and
[Equation]
Since b t = min( b (1) t , b (2) t ) , we can substitute b t for b (1) t and b (2) t in (12) and (13), respectively.
Also, from the definition of b t we know 1 b t =
O (cid:18) 1 log log
[Equation]
w
[Equation]
n (cid:19) .
Plugging them into (2) proves the first claim.
The second claim follows from the O ( n − 1 /d ) convergence rate of an empirical measure to its population measure in 1 -Wasserstein
[Table]