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, 14] 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 [14, 22, 23], generative modeling [21], and reinforcement learning [10, 16].
Algorithm 1 Stein Variational Gradient Descent [14]: SVGD ( µ 0 , r )
Input: Target P with density p , kernel k , step sizes ( ϵ s ) s ≥ 0 , approximating measure µ 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 [13, 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, 4, 8, 15].
KSD convergence is also known to imply weak convergence under various conditions on the target P and the SVGD kernel k [1, 3, 8, 11].
Followup work by Korba et al.
[12], Salim et al.
[17], Sun et al.
[19] sharpened the result of Liu with path-independent constants, weaker smoothness conditions, and explicit rates of convergence.
In addition, Duncan et al.
[6] analyzed the continuous-time limit of continuous SVGD to provide
A 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 [13, 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.
[9, 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.
Korba et al.
Prop. 7] explicitly bounded the expected squared Wasserstein distance between n -particle and continuous SVGD but only under the assumption of bounded ∇ log p , an assumption that rules out all strongly log concave or dissipative distributions and all distributions for which the KSD is currently known to control weak convergence [1, 3, 8, 11].
In addition, Korba et al.
[12] do not provide a unified bound for the convergence of n -particle SVGD to P .
In this work, we derive a first unified convergence bound for finite-particle SVGD to its target.
To achieve this, we first bound the 1 -Wasserstein discretization error between finite-particle and continuous SVGD under assumptions commonly satisfied in SVGD applications and compatible with
KSD weak convergence control (see Theorem 1).
We next bound KSD in terms of 1 -Wasserstein distance and SVGD moment growth to explicitly control KSD discretization error 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 [2].
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 twice 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]
To leverage the continuous SVGD convergence rates of Salim et al.
[17], we additionally assume that the target P satisfies Talagrand’s T 1 inequality [20, Def. 22.1].
Remarkably, Villani [20,
Thm. 22.10] showed that Assumption 4 is equivalent to P being a sub-Gaussian distribution.
Hence, this mild assumption holds for all strongly log concave P [18, Def. 2.9], all P satisfying the log
Sobolev inequality [20, Thm. 22.17], and all distantly dissipative P for which KSD is known to control weak convergence [8, Def. 4].
Assumption 4 (Talagrand’s T 1 inequality [20, Def. 22.1]) E P [ ∥·∥ 2 ] < ∞ , and there exists λ >such that for all µ with E µ [ ∥·∥ 2 ] < ∞
[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 [7] 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 [15] for P satisfying Assumption 1 is defined by
[Equation]
with components
[Equation]
Definition 3 (SVGD transport map and pushforward)
The SVGD transport map [14] for a kernel k , target P satisfying Assumption 1, approximating µ ∈ P , and step size ϵ is given by
[Equation]
+ ϵ E X ∼ µ [( A P k ( · , x ))
( X )] .
Moreover, the SVGD pushforward Φ ϵ ( µ ) represents the distribution of T µ,ϵ ( X ) when X ∼ µ .
Definition 4 (Kernel Stein discrepancy)
For a target P satisfying Assumption 1 and measures µ, ν ∈ P , we define the Langevin kernel Stein discrepancy [KSD, 4, 8, 15] 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 µ nr = SVGD ( µ n 0 , r ) and µ ∞ r = SVGD ( µ ∞ 0 , r ) satisfy
[Equation]
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 + L ( γ + √ dκ 2 )(1 + ∥ x ⋆ ∥ 2 ) .
The proof of Theorem 1 in Appendix A relies on two lemmas.
The first, due to Gorham et al.
[9], 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 [9, 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 in 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 ,
[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.
[17] 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 [17, 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 ≜ min (cid:16)
[Equation]
for p ∈ { 1 , 2 } , 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 ≜ c ( ϵ r ) (cid:80) tr
=0 c ( ϵ 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 π r as defined in Lemma 4, ( A, B, C ) as defined in Theorem 1, b t ≜ (cid:80) t r =0 ϵ r , and a t ≜ κ ( d + L ) w 0 ,n exp( b t ( A + B exp( Cb t )))
+ d 1 / 4 κL (cid:112)
2 M µ ∞ 0 ,P w 0 ,n exp( b t (2 C + A + B exp( Cb t )) / 2) .
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)
[Equation]
we have
[Equation]
We finish the proof by noticing that (cid:80) tr =
[Equation]
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
[Equation]
and define the shorthand ψ B,C ( x, y, β ) ≜ 1 C log( 1 B ( 1 β log 1 x − y )) .
If
[Equation]
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, β 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, 1 (cid:1)(cid:17) , then, min 0 ≤ r ≤ t KSD P ( µ nr , P ) =
O (cid:16) 1 / (cid:113) log log 1 w 0 ,n (cid:17) .
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 P ( µ nr , P ) =
O p (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-Lipschitzness conditions of Lemma 1 hold given our assumptions.
According to Assumption
[Equation]
we have
[Equation]
Recall that x ⋆ satisfies s p ( x ⋆ ) = 0 by Assumption 1.
Then, by the 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-Lipschitzness condition.
Similarly
[Equation]
we have
[Equation]
where we used the Lipschitzness of s p from Assumption 1 and
[Equation]
The last inequality is due to
[Equation]
According to Assumption 3, 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 (6) and (7),
[Equation]
Plugging (8) back into (5), we can show the second pseudo-Lipschitzness condition holds for c 2 = κ
[Equation]
+
[Equation]
Now that the pseudo-Lipschitzness conditions hold, by repeated application of Lemma 1 and the inequality (1 + x ) ≤ e x , we have
[Equation]
(9) for D 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 ) , B = c 1 m µ n 0 ,P + c 2 m µ ∞ 0 ,P , and C = κ 2 (3 L + d ) .
Therefore
[Equation]
Plugging this back into (9) proves the result
[Equation]
Appendix B. Proof of Lemma 2: SVGD moment growth
From Assumption 2 we know
[Equation]
The last inequality implies
[Equation]
Let µ be any probability measure.
Using the above results, Jensen’s inequality and the facts that
E P [ s p ( · )]
= 0 and E Z ∼ P
[( A P k ( · , x ))( Z )]
= 0 , we have ∥ T µ,ϵ
[Equation] [Equation]
The last step used the definitions m µ,P ≜ E X ∼ µ ⊥⊥ Z ∼ P [ ∥
X − Z ∥ 2 ] and C
= κ 2 (3 L + d ) .
Then, applying triangle inequality and (10)
[Equation]
we have
[Equation]
where the second last step used Jensen’s inequality m 2 µ r ,P ≤ M µ r ,P .
Then
[Equation]
we repeatedly apply (11) and (12) together with the inequality 1 + x ≤ e x to get
[Equation]
Appendix C. Proof of Lemma 3: KSD-Wasserstein bound
Our proof generalizes that of Gorham and Mackey [8, 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
[Equation] [Equation]
[Equation]
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 (13), and g and ∇ · g are Lipschitz (14), repeated use of Cauchy-Schwarz gives
E µ
[Equation]
E ν
[ T P g ]
[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
[Equation]
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, β 2 ) .
According to Lemma 5,
[Equation]
Rearranging the terms, we have
[Equation]
Similarly, we have b (2) t
[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 (15) and (16), respectively.
Also, from the definition of b t we know 1 b t =
O
[Equation]
Plugging them into (2) proves the first claim.
The second claim follows from the O p ( n − 1 / max( d, 2) ) convergence rate of an empirical measure to its population measure in 1 -Wasserstein distance [5, Thm. 1].