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 is sub-Gaussian with a Lipschitz score, 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 [9, 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, 7, 15].
KSD convergence is also known to imply weak convergence under various conditions on the target P and the SVGD kernel k [1, 3, 7, 10].
Followup work by Korba et al.
[11], 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.
[5] 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.
[8, 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, 7, 10].
In addition, Korba et al.
[11] 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 [7, 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 [6] 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, 7, 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.
[8], 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 [8, 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 ∈ {
[Equation]
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 ≤ ϵ t ≜ R α, 1 for some α > 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 − 1 ≜ (cid:80) t − 1 r =
ϵ r , and a t − 1 ≜ κ ( d
[Equation]
) w
[Equation]
n exp
[Equation]
3) +
[Equation]
Proof By the triangle inequality and Theorem 2 we have
[Equation]
for each r .
Therefore
[Equation]
where the last inequality follows from Corollary 1.
Moreover, by Jensen’s inequality,
[Equation]
Combining (4) and (5)
[Equation]
we have
[Equation]
We finish the proof by noticing that (cid:80)
[Equation]
The following corollary, proved in Appendix D, provides an explicit SVGD convergence bound and rate by choosing the step size sum to balance the 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, let ( ¯ w 0 ,n , ¯ A, ¯ B, ¯ C ) be any upper bounds on ( w 0 ,n , A, B, C ) respectively, and define the growth functions
[Equation]
and
[Equation]
¯ B, ¯ C ( x, y, β ) ≜ 1¯ C log( 1¯ B max( ¯ B, 1 β log 1 x − y )) .
If the step size sum b
[Equation]
ψ ¯
[Equation]
¯
[Equation]
¯ w
[Equation]
n ϕ ( ¯ w
[Equation]
n )
, ¯ A + 2 ¯ C, β
[Equation]
(cid:1)(cid:17) , for β
[Equation]
max
[Equation] [Equation]
B,
[Equation]
C (cid:0)
[Equation]
,n
[Equation] [Equation]
[Equation]
and
[Equation]
then min
[Equation]
if
[Equation]
otherwise .
[Equation] [Equation]
[Equation] [Equation]
If, in addition, µ n 0 = 1 n (cid:80) ni =
δ x i for x i i.i.d.
∼ µ ∞ 0 with M µ ∞ 0 ≜ E µ ∞ 0 [ ∥·∥ 22 ] < ∞ , then
[Equation]
with probability at least 1 − cδ for a universal constant c > 0 .
Hence, with this choice of ¯ w 0 ,n ,
[Equation]
with probability at least 1 − cδ .
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 (10) and (11),
[Equation]
Plugging (12) back into (9)
[Equation]
we can show the second pseudo-Lipschitzness condition holds for c
[Equation]
Now that the pseudo-Lipschitzness conditions hold
[Equation]
by repeated application of Lemma 1 and the inequality (1 + x ) ≤ e x , we have
[Equation]
) for D s = c 1 (1 + m µ ns ) + c 2 (1 + m µ ∞ s ) .
Using the result from Lemma
[Equation]
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 (13) 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 (14)
[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 (15) and (16) 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 [7, 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 (17), and g and ∇ · g are Lipschitz (18), repeated use of Cauchy-Schwarz gives
E
[Equation]
ν [ 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 begin by establishing a lower bound on b t − 1 .
Let
[Equation]
β
[Equation]
and b (2) t − 1 = ψ ¯ B, ¯ C ( ¯ w 0 ,n ϕ ( ¯ w 0 ,n ) , ¯ A + 2 ¯ C, β 2 ) so that b t − 1 = min( b (1) t − 1 , b (2) t − 1 ) .
Since β 1 , β 2 , ϕ ( ¯ w 0 ,n ) ≥ 1 , we have β 1 = max(1 , 1 ¯ C log( 1 ¯ B (log ¯ w 0 ,n √ ϕ ( ¯ w 0 ,n ) − ¯ A )))
≤ max(1 , 1¯ C log( 1¯ B (log ¯ w 0 ,n √ ϕ ( ¯ w 0 ,n ) )))
≤ max(1 , 1¯ C log( 1¯ B (log 1¯ w 0 ,n ))) and β 2 = max(1 , 1¯ C log( 1¯ B (log 1 ¯ w 0 ,n ϕ ( ¯ w 0 ,n ) − ¯
A − 2 ¯ C ))) ≤ max(1 , 1¯ C log( 1¯ B (log 1 ¯ w 0 ,n ϕ ( ¯ w 0 ,n ) )))
≤
max
(1 , 1¯ C log( 1¯ B (log 1¯ w 0 ,n )
Hence, ϕ ( ¯ w 0 ,n ) ≥ 1 implies that
[Equation]
and
[Equation]
We divide the remainder of our proof into four parts.
First we prove each of the two cases in the generic KSD bound (6) in Appendices D.1 and D.2.
Next we show in Appendix D.3 that these two cases yield the generic convergence rate (7).
Finally, we prove the high probability upper estimate (8) for w 0 ,n under i.i.d.
initialization in Appendix D.4.
D.1. Case b t − 1 =
In this case, the error bound (6) follows directly from Theorem 3.
D.2. Case b t − 1 >
We first state and prove a useful lemma.
Lemma 5 Suppose x = f ( β ) for a non-increasing function f : R → R and β = max(1 , f (1)) .
Then x ≤ f ( x ) .
Proof If β = 1 ≥ f (1) , then x = f ( β ) =
f (1) ≤ 1 .
Because f is non-increasing, f ( x ) ≥ f (1) = x .
Otherwise, β = f (1) > 1 , and hence x = f ( β ) ≤ f (1) =
β since f is non-increasing.
Since x ≤ β and f is non-increasing, we further have f ( x ) ≥ f ( β ) =
x as advertised.
Since ψ ¯ B, ¯ C is non-increasing in its third argument, Lemma 5 implies that
[Equation]
Rearranging the terms and noting that
[Equation]
since b (1) t − 1 ≥ b t − 1 > 0 , we have
[Equation]
Similarly, we have b (2) t − 1 ≤ ψ ¯ B, ¯ C ( ¯ w 0 ,n log log 1¯ w 0 ,n , ¯ A + 2 ¯ C, b (2) t − 1 ) and
[Equation]
INITE
Since b t = min( b (1) t − 1 , b (2) t − 1 ) , the inequalities (20) and (21) are also satisfied when b t is substituted for b (1) t − 1 and b (2) t − 1 .
Since the error term a t (3) is non-decreasing in each of ( w 0 ,n , A, B, C ) , we have
[Equation]
Since b t − 1 = min( b (1) t − 1 , b (2) t − 1 ) , the claim (6) follows from this estimate, the lower bounds (19), and
Theorem 3.
D.3. Generic convergence rate
The generic convergence rate (7) holds as, by the lower bounds (19), b t − 1 = min( b (1) t − 1 , b (2) t − 1 ) >
whenever
[Equation]
and
[Equation]
a condition which occurs whenever ¯ w 0 ,n is sufficiently small since the right-hand side of each inequality converges to zero as ¯ w 0 ,n → 0 .
D.4. Initializing with i.i.d. particles
We begin by restating an expected Wasserstein bound due to Lei [12].
Lemma 6 (Lei [12, Thm. 3.1]) Suppose µ n 0 = 1 n (cid:80) ni
=
δ x i for x i i.i.d.
∼ µ ∞ 0 with M µ ∞ 0 ≜ E µ ∞ 0 [ ∥·∥ 22 ] < ∞ .
Then, for a universal constant c > 0 ,
E [
[Equation]
Together, Lemma 6 and Markov’s inequality imply that
[Equation]
with probability at least 1 − cδ , proving the high probability upper estimate (8).