Sampling with Mirrored Stein Operators
Abstract
Accurately approximating an unnormalized distribution with a discrete sample is a fundamental challenge in machine learning, probabilistic inference, and Bayesian inference.
Particle evolution methods like Stein variational gradient descent have found great success in approximating unconstrained distributions but break down for constrained targets.
We introduce a new family of particle evolution samplers suitable for constrained domains and non-Euclidean geometries.
They minimize the Kullback-Leibler (KL) divergence to constrained target distributions by evolving particles in a dual space deﬁned by a mirror map.
We derive these samplers from a new class of mirrored Stein operators and adaptive kernels developed in this work.
We establish the convergence of our new procedures under veriﬁable conditions on the target distribution.
Finally, we demonstrate that these new samplers yield accurate approximations to distributions on the simplex and deliver valid conﬁdence intervals in post-selection inference.
1. Background: Mirror Descent
Standard gradient descent can be viewed as optimizing a local quadratic approximation to the target function f : θ t +1 = argmin θ ∈ Θ ∇ f ( θ t ) (cid:62) θ + 12 (cid:15) t (cid:107) θ − θ t (cid:107) 22 .
When Θ ⊆ R d is constrained, it can be advantageous to replace (cid:107) · (cid:107) 2 with a function Ψ that reﬂects the geometry of a problem (Nemirovskij and Yudin, 1983; Beck and Teboulle, 2003): θ t +1 = argmin θ ∈ Θ ∇ f ( θ t ) (cid:62) θ + 1 (cid:15) t Ψ
( θ, θ t ) .
The mirror descent (MD) algorithm chooses Ψ to be the
Bregman divergence induced by a strongly convex, essentially smooth 1 function ψ : Θ →
R ∪ {∞} : Ψ( θ, θ (cid:48) )
= ψ ( θ ) − ψ ( θ (cid:48) )
− ∇ ψ ( θ [Equation]
) (cid:62) ( θ − θ (cid:48) ) .
The solution of the problem is
[Equation]
where ψ ∗ ( η ) (cid:44) sup θ ∈
Θ η (cid:62) θ − ψ ( θ ) is the convex conjugate of ψ and ∇ ψ is a bijection from Θ to dom( ψ ∗ ) with inverse map ( ∇ ψ ) − 1 = ∇ ψ ∗ .
We can view (1) as ﬁrst mapping θ t to η t by ∇ ψ , applying the update η t +1 =
η t − (cid:15) t ∇ f ( θ t ), and mapping back through θ t +1 = ∇ ψ ∗ ( η t +1 ).
We can view MD as a discretization of the continuous-time dynamics dη t = −∇ f
( θ t ) dt, θ t = ∇ ψ ∗ ( η t ).
It is equivalent to the Riemannian gradient ﬂow (see App.
B):
[Equation]
or equivalently,
[Equation]
where ∇ 2 ψ ( θ ) and ∇ 2 ψ ∗ ( η ) are Riemannian metric tensors.
In information geometry, the discretization of (2) is known as natural gradient descent (Amari, 1998).
There is considerable theoretical and practical evidence (Martens, 2014) showing that natural gradient works eﬃciently in learning.
[Equation]
2. Sampling with Mirrored Stein Operators
Particle evolution methods like Stein variational gradient descent (SVGD, Liu and Wang, 2016) apply deterministic updates to particles using operators based on Stein’s method (Stein, 1972; Gorham and Mackey, 2015) and reproducing kernels (Berlinet and Thomas-Agnan, 2011) to sequentially minimize Kullback-Leibler (KL) divergence.
Speciﬁcally, SVGD updates each particle in its approximation by applying θ t +
=
θ t + (cid:15)
t g t ( θ t ) for a chosen mapping g t : R d → R d .
It chooses the mapping g ∗ t that leads to the largest decrease in KL divergence to p in the limit as (cid:15) t → 0.
The following theorem summarizes their main ﬁndings.
Theorem 1 (Liu and Wang, 2016, Thm. 3.1) Suppose ( θ t ) t ≥ 0 satisﬁes dθ t =
g t ( θ t ) dt for bounded Lipschitz g t ∈ C 1 : R d → R d and that θ t has density
q t with E q t
(cid:107)∇ log q t ( θ ) (cid:107) 2 ] < ∞ .
If KL( q t (cid:107) p ) (cid:44) E q t [log( q t ( θ ) /p
( θ ))] exists then,
[Equation]
(3) where S p is the Langevin Stein operator (Gorham and Mackey, 2015) given by
[Equation]
For an unconstrained domain with E p [ (cid:107)∇ log p ( θ ) (cid:107) 2 ] < ∞ , Stein’s identity E p
[( S p g )
( θ )]
= 0 holds whenever g ∈ C 1 is bounded and Lipschitz by (Gorham et al., 2019, proof of Prop. 3), which ensures q t = p is a stationary point of the dynamics.
To improve its current particle approximation, SVGD ﬁnd the choice of g t that most quickly decreases KL( q t (cid:107) p ) at time t , i.e., they minimize ddt KL( q t (cid:107) p ) over a set G of candidate directions g t .
SVGD chooses g t in a reproducing kernel Hilbert space (RKHS, Berlinet and Thomas-Agnan, 2011) norm ball
B H d = { g : (cid:107) g (cid:107) H d ≤ 1 } , where H d is the product RKHS containing vector-valued functions with each component in the RKHS H of k .
Then the optimal g ∗ t ∈ B H d that minimizes (3) is
[Equation]
where we let K k ( θ, θ (cid:48) )
=
k ( θ, θ (cid:48) )
I , and S p K k ( · , θ ) denotes applying S p to each row of K k ( · , θ ).
2.1. Mirrored dynamics
SVGD encounters two diﬃculties when faced with a constrained support.
First, the SVGD updates can push θ t outside of its support Θ, rendering all future updates undeﬁned.
Second,
Stein’s identity E p
[( S p g )
( θ )]
= 0 often fails to hold for B H d (cf.
Ex. 1) if p is non-vanishing or explosive at the boundary.
When this occurs, SVGD need not converge to p as p is not a stationary point of its dynamics.
To ﬁx this, we consider the following mirrored dynamics θ t = ∇ ψ ∗ (
[Equation]
t ) for dη t =
g t ( θ t ) dt, or, equivalently,
[Equation]
where g t : Θ → R d now represents the update direction in η space.
The inverse mirror map ∇ ψ ∗ guarantees that θ t belongs to the constrained domain Θ. Since ψ is strongly convex and ∇ 2 ψ − 1 is bounded Lipschitz, from Thm.
1 it follows for any bounded Lipschitz g t that
[Equation]
where we introduce the mirrored Stein operator M p,ψ 2 :
[Equation]
Here ψ is as in Sec.
1 with ( ∇ 2 ψ ) − 1 diﬀerentiable and Lipschitz on Θ. The following result, proved in App.
L.1, shows that M p,ψ generates mean-zero functions under p whenever ∇ 2 ψ − 1 suitably cancels the growth of p at the boundary.
Proposition 2 Suppose that ∇ 2 ψ − 1 ∇ log p and ∇·∇ 2 ψ − 1 are p -integrable.
If lim r →∞ (cid:82) ∂ Θ r p ( θ ) (cid:107)∇ 2 ψ ( θ ) − 1 n r ( θ ) (cid:107) 2 dθ = 0 for Θ r (cid:44) { θ ∈ Θ : (cid:107) θ (cid:107) ∞ ≤ r } and n r ( θ ) the outward unit normal vectorto
∂
Θ r at θ , then E p [( M p,ψ g )
( θ )]
= 0 if g ∈ C 1 is bounded Lipschitz.
Example 1 (Dirichlet p , Negative entropy ψ )
When θ 1: d +1 ∼ Dir( α ) for α ∈ R d +1+ , even setting g ( θ ) = 1 in (4) need not cause E p [( S p g )
( θ )]
= 0 when ∃ j, α j ≤ 1 .
However, we show in App.
C that the conditions of Prop. 2 are met for any α if ψ ( θ ) = (cid:80) d +1 j =1 θ j log θ j .
We propose new deterministic sampling algorithms by seeking optimal directions g t that minimizes (6) over diﬀerent function classes.
Thm.
3 forms the basis of our analysis.
Theorem 3 (Optimal mirror updates in RKHS, proved in App.
L.2)
Let ( θ t ) t ≥follow the mirrored dynamics (5) .
Let H K denote the RKHS of a matrix-valued kernel
K : Θ × Θ → S d × d (Micchelli and Pontil, 2005).
Then, the optimal direction of g t that minimizes (6) in the norm ball B H K (cid:44) { g : (cid:107) g (cid:107) H K ≤ 1 } is
[Equation]
2.2. Mirrored Stein Variational Gradient Descent
We can simply choose the K to be
K k
( θ, θ (cid:48) )
=
k ( θ, θ (cid:48) )
I .
In this case, the update g ∗ q t ,K k ( · ) =
E q t [ M p,ψ K k ( · , θ )] is equivalent to running SVGD in the η space before mapping back to Θ.
Theorem 4 (Mirrored SVGD, proof is in App.
L.3)
In the setting of Thm. 3, let k ψ ( η, η (cid:48) )
=
k ( ∇ ψ ∗ ( η ) , ∇ ψ ∗ ( η (cid:48) )) , p H
[Equation]
η ) =
p
[Equation]
∇ ψ
[Equation]
η )) · | det ∇ 2 ψ ∗ ( η ) | denote the density of η = ∇ ψ ( θ ) when θ ∼ p , and q t,H denote the distribution of η t under the mirrored dynamics (5) .
If K k = kI ,
[Equation]
The proof is in App.
L.3.
By discretizing the dynamics dη t =
g ∗ q t ,K k ( θ t ) dt and initializing with any particle approximation q 0 = 1 n (cid:80) ni
=
δ θ i 0 , we obtain Mirrored SVGD (MSVGD) , our ﬁrst algorithm for sampling in constrained domains.
The details are summarized in Alg. 1.
When a single particle is used and the kernel satisﬁes ∇ k
( θ, θ ) = 0, the MSVGD update (9) reduces to gradient descent on − log p H ( η ).
This however is not MD: although MD operates in η space, it uses the gradient of − log p ( θ ) instead of − log p H ( η ).
Also note the modes of p H ( η ) need not match those of p ( θ ) (cf.
examples in App. F).
Since we are primarily interested in the θ -space density, it is natural to ask whether there exists a mirrored dynamics that reduces to ﬁnding the mode of p ( θ ) when n = 1.
In the next section, we give an answer to this question by designing an adaptive kernel that yields a MD-like update.
2.3. Stein Variational Mirror Descent
Our second sampling algorithm for constrained problems is called Stein Variational Mirror
Descent (SVMD) .
We start by introducing a new adaptive matrix-valued kernel.
Deﬁnition 5 (Kernels for SVMD)
Given a scalar-valued kernel k , consider the Mercer representation 3 k ( θ, θ (cid:48) )
= (cid:80)
i
[Equation]
λ t
[Equation]
u t
[Equation]
i
[Equation]
θ
[Equation]
u t
[Equation]
i ( θ (cid:48) ) w.r.t.
q t , where u t,i is an eigenfunction:
[Equation]
For k 1 / 2 t
[Equation]
) , we deﬁne the adaptive SVMD kernel at time t ,
[Equation]
By Thm. 3, the optimal update direction for the SVMD kernel ball is g ∗ q t ,K ψ,t =
E q t [ M p,ψ K ψ,t ( · , θ )] .
We obtain the SVMD algorithm (summarized in Alg. 1) by discretizing dη t =
g ∗ q t ,K ψ,t ( θ t ) dt and initializing with q 0 = 1 n (cid:80)
ni
=
δ θ i 0 .
Because of the discrete representation of q t , K ψ,t takes the form K ψ,t ( θ, θ (cid:48) )
= (cid:80) ni
=1 (cid:80) nj =1 λ 1 / 2 t,i λ 1 / 2 t,j u t,i ( θ ) u t,
j ( θ (cid:48) )
Γ t,ij , for Γ t,ij = 1 n (cid:80) n(cid:96)
=1 u t,i ( θ (cid:96)t )
u t,j ( θ (cid:96)t ) ∇ 2 ψ ( θ (cid:96)t ) .
λ t,j , u t,j ( θ it ) and its gradients can be computed by solving a matrix eigenvalue problem involving the particle set { θ it } ni =1 .
We give the details in App. H.
Notably, SVMD diﬀers from MSVGD only in its choice of kernel, but, whenever ∇ k ( θ, θ ) = 0, this change is suﬃcient to exactly recover MD when n = 1.
Proposition 6 (Single-particle SVMD is MD)
If n = 1 , then one step of SVMD becomes η t +1 =
η t + (cid:15) t ( k ( θ t , θ t ) ∇ log p ( θ t ) + ∇ k
( θ t , θ t )) , θ t +1 =
[Equation]
In App.
I, we establish convergence guarantees for our proposed algorithms.
By leveraging the connection between MD and natural gradient descent, we further generalize SVMD to an algorithm for unscontrained domains that can exploit the geometry of the problems.
We describe this algorithm in App.
J.
3. Experiments 3.1. Approximation quality on the simplex
We ﬁrst measure distributional approximation quality using two 20-dimensional simplexconstrained targets: the sparse Dirchlet posterior of Patterson and Teh (2013) and the
[Figure]
quadratic simplex target of Ahn and Chewi (2020).
The Dirichlet mimics the multimodal conditionals that arise in latent Dirichlet allocation (Blei et al., 2003) but induces a log concave density in η space, while the quadratic is log-concave in θ space.
To compare with standard SVGD and to prevent its particles from exiting the domain, we introduce a
Euclidean projection onto Θ following each SVGD update.
In Fig.
2, we compare MSVGD,
SVMD, and projected SVGD with 50 particles and inverse multiquadric kernel k (Gorham and Mackey, 2017) by computing the energy distance (Sz´ekely and Rizzo, 2013) to a surrogate ground truth sample.
We also compare to MSVGD with k 2 ( θ, θ (cid:48) )
=
k ( ∇ ψ ( θ ) , ∇ ψ ( θ (cid:48) )), a choice which corresponds to running SVGD in the dual space with kernel k by Thm. 4 and which ensures the convergence of MSVGD to p by Thms. 9, 11 and 12.
In the quadratic case, SVMD is favored over MSVGD as it is able to exploit the logconcavity of p ( θ ).
In contrast, for the multimodal sparse Dirichlet with p ( θ ) unbounded near the boundary, MSVGD converges slightly more rapidly than SVMD by exploiting the log concave structure in η space.
Projected SVGD fails to converge to the target in both cases and has particular diﬃculty in approximating the sparse Dirichlet target with unbounded density.
MSVGD with k and k 2 perform similarly, but k yields better approximation quality upon convergence.
Therefore, we employ k in the remaining MSVGD experiments.
3.2. Conﬁdence intervals for post-selection inference
We next apply our methods to the constrained sampling problems that arise in post-selection inference (Taylor and Tibshirani, 2015).
Speciﬁcally, we consider the task of forming valid conﬁdence intervals (CIs) for regression parameters selected by the randomized Lasso (Tian and Taylor, 2018) with data X ∈ R ˜ n × p and y ∈ R ˜ n and user-generated randomness w ∈ R p from a log-concave distribution with density g .
The Lasso returns ˆ β ∈ R p with non-zero coeﬃcients denoted by ˆ β E and their signs by s E .
It is common practice to report least squares CIs for β E by running a linear regression on the selected features E .
However, since
E is chosen based on the same data, the resulting CIs are often invalid.
Post-selection inference solves this problem by conditioning the inference on the knowledge of E and s E .
To construct valid CIs, it suﬃces to approximate the selective distribution
[Figure]
with support { ˆ β E , u − E : s E (cid:12) ˆ β E > 0 , u − E ∈ [ − 1 , 1] p −| E | } and density
[Equation]
In our experiments, we integrate out u − E analytically, following Tian and Taylor (2018), and reparameterize ˆ β E as s E (cid:12) | ˆ β E | to obtain a log-concave density of | ˆ β E | supported on the nonnegative orthant.
We choose ψ ( θ )
=
(cid:80) dj
=1 ( θ j log θ j − θ j ).
In Fig.
3 (left) we show the example of a 2D selective distribution.
We also plot the results by projected SVGD, SVMD, and MSVGD in this example.
Projected SVGD fails to approximate the target, while the samples by MSVGD and SVMD closely resemble the truth.
We then compare our methods with the standard norejection MCMC approach of the selectiveInference R package (Tibshirani et al., 2019) using the example simulation setting described in Sepehri and Markovic (2017).
To generate N total sample points we run
MCMC for N iterations after burn-in or aggregate the particles from N/n independent runs of MSVGD or SVMD with n = 50 particles.
As N ranges from 1000 to 3000 in Fig.
4( a ), the
MSVGD and SVMD CIs consistently yield higher coverage than the standard 90% CIs.
This increased coverage is of particular value for smaller sample sizes, for which the standard CIs tend to undercover.
For a much larger sample size of N = 5000 in Fig.
4( b ), the SVMD and standard CIs closely track one another across conﬁdence levels, while MSVGD overcovers.
We next apply our samplers to a post-selection inference task on the HIV-1 drug resistance dataset (Rhee et al., 2006), where we run randomized Lasso (Tian and Taylor, 2018) to ﬁnd statistically signiﬁcant mutations associated with drug resistance using susceptibility data on virus isolates.
In Fig.
3 (right) we plot the CIs of selected mutations obtained with
N = 5000 sample points.
We see that the invalid unadjusted least squares CIs can lead to premature conclusions, e.g., declaring mutation 215Y signiﬁcant when there is insuﬃcient support after conditioning on the selection event.
In contrast, mutation 184V, which has known association with drug resistance, is declared signiﬁcant by all methods even after post-selection adjustment.
The MSVGD and SVMD CIs mostly track those of the standard selectiveInference method, but their conclusions sometimes diﬀer: e.g., 62Y is ﬂagged as signiﬁcant by MSVGD and SVMD but not by selectiveInference .
Appendix A. Related Work
Our mirrored Stein operators (7) are instances of diﬀusion Stein operators in the sense of Gorham and Mackey (2017), but their speciﬁc properties have not been studied, nor have they been used to develop sampling algorithms.
There is now a large body of work on transferring algorithmic ideas from optimization to sampling (see, e.g., Dalalyan, 2017;
Welling and Teh, 2011; Durmus et al., 2018;
Ma et al., 2019; Simsekli et al., 2016).
The closest to our work in this space is the recent marriage of mirror descent and MCMC.
For example,
Hsieh et al.
(2018) propose to run Langevin Monte Carlo (LMC, an Euler discretization of the Langevin diﬀusion) in a mirror space.
Zhang et al.
(2020) analyze the convergence properties of the mirror-Langevin diﬀusion, Chewi et al.
(2020) demonstrate its advantages over the Langevin diﬀusion when using a Newton-type metric, and Ahn and Chewi (2020) study its discretization for MCMC sampling in constrained domains.
Relatedly, Patterson and Teh (2013) proposed stochastic Riemannian LMC for sampling on the simplex.
Several modiﬁcations of SVGD have been proposed to incorporate geometric information.
Riemannian SVGD (RSVGD, Liu and Zhu, 2018) generalizes SVGD to Riemannian manifolds, but, even with the same metric tensor, their updates are more complex than ours: notably they require higher-order kernel derivatives, do not operate in a mirror space, and do not reduce to natural gradient descent when a single particle is used.
They also reportedly do not perform well when with scalable stochastic estimates of ∇ log p .
Stein Variational
Newton (SVN, Detommaso et al., 2018; Chen et al., 2019) introduces second-order information into SVGD.
Their algorithm requires an often expensive Hessian computation and need not lead to descent directions, so inexact approximations are employed in practice.
Our SVNG can be seen as an instance of matrix SVGD (MatSVGD, Wang et al., 2019) with an adaptive time-dependent kernel discussed in App.
J, a choice that is not explored in Wang et al.
and which recovers natural gradient descent when n
= 1 unlike the heuristic kernel
[Figure]
constructions of Wang et al.
None of the aforementioned works provide convergence guarantees, and neither SVN nor matrix SVGD deals with constrained domains.
Appendix B. Mirror Descent, Riemannian Gradient Flow, and Natural Gradient
The equivalence between the mirror ﬂow dη t =
−∇
f
( θ t ) dt, θ t = ∇ ψ ∗ ( η t )
dt and the Riemannian gradient ﬂow in (2) is a direct result of the chain rule:
[Equation]
[Equation]
Depending on discretizing (12) or (13), there are two natural gradient descent (NGD) updates that can arise from the same gradient ﬂow:
[Equation]
With ﬁnite step sizes (cid:15) t , their updates need not be the same and can lead to diﬀerent optimization paths.
Since ∇ f ( θ t ) = ∇ 2 ψ ∗ ( η t ) − 1 ∇ η t f ( ∇ ψ ∗ ( η t )), NGD (b) is equivalent to the dual-space update by mirror descent.
This relationship was pointed out in Raskutti and
Mukherjee (2015) and has been used for developing natural gradient variational inference algorithms (Khan and Nielsen, 2018).
We emphasize, however, our SVNG algorithm developed in App.
J corresponds to the discretization in the primal space as in NGD (a).
Therefore, it does not require an explicit dual space, and allows replacing ∇ 2 ψ with more general information metric tensors.
Appendix C. Details of Example
For the entropic mirror map ψ ( θ ) = (cid:80) d +1 j =1 θ j log θ j , we have ∇ 2 ψ ( θ ) − 1 = diag( θ ) − θθ (cid:62) .
Note here θ denotes a d -dimensional vector and does not include θ d +1 = 1 − (cid:80) dj
=1 θ d .
Since
Θ is a ( d + 1)-simplex, ∂ Θ is composed of d + 1 faces with θ in the j -th face satisﬁes θ j = 0.
The outward unit normal vector n ( θ ) for the ﬁrst d faces are − e j for 1 ≤ j ≤ d , where e j denotes the j -th standard basis of R d .
The outward unit normal vector for the ( d + 1)-st face is a vector with 1 / √ d in all coordinates.
Therefore, we have
[Equation]
where in the second to last identity we used θ (cid:62) 1 = 1 − θ d +1 .
Finally, we can verify the condition in Prop. 2 as
[Equation]
Appendix D. Derivation of the Mirrored Stein Operator
We ﬁrst review the (overdamped) Langevin diﬀusion – a Markov process that underlies many recent advances in Stein’s method – along with its recent mirrored generalization.
The Langevin diﬀusion with equilibrium density p on R d is a Markov process ( θ t ) t ≥
⊂ R d satisfying the stochastic diﬀerential equation (SDE)
[Equation]
) with ( B t ) t ≥ 0 a standard Brownian motion (Bhattacharya and Waymire, 2009, Sec. 4.5).
To identify Stein operators that generate mean-zero functions under p for broad classes of targets p , Gorham and Mackey (2015) proposed to build upon the generator method of Barbour (1988):
First, identify a Markov process ( θ t ) t ≥ 0 that has p as the equilibrium density; they chose the Langevin diﬀusion of (14).
Next, build a Stein operator based on the (inﬁnitesimal) generator A of the process (Øksendal, 2003, Def. 7.3.1):
[Equation]
as the generator satisﬁes E p [( Af )( θ )] = 0 under relatively mild conditions.
We use the following theorem to derive the generator of the processes described by SDEs like (14):
Theorem 7 (Generator of Itˆo diﬀusion; Øksendal, 2003, Thm 7.3.3)
Let ( x t ) t ≥be the Itˆo diﬀusion in X ⊆ R d satisfying dx t
=
b ( x t ) dt + σ ( x t ) dB t .
For any f ∈ C 2 c ( X ) , the (inﬁnitesimal) generator A of ( x t ) t ≥ 0 is
[Equation]
Substituting ∇ log p ( · ) for b ( · ) and √ 2 I for σ ( · ) in Thm.
7, we obtain Af =
( ∇ log p ) (cid:62) ∇ f + ∇ · ∇ f .
Replacing ∇ f with a vector-valued function g gives the Langevin Stein operator in (4).
To derive a Stein operator that works well for constrained domains, we consider the
Riemannian Langevin diﬀusion (Patterson and Teh, 2013; Xifara et al., 2014; Ma et al., 2015) that extends the Langevin to non-Euclidean geometries encoded in a positive deﬁnite metric tensor G ( θ ):
[Equation]
We show in App.
E that the choice G = ∇ 2 ψ yields the recent mirror-Langevin diﬀusion (Zhang et al., 2020; Chewi et al., 2020)
[Equation]
According to Thm. 7, the generator of the mirror-Langevin diﬀusion described by (16) is
[Equation]
Now substituting g ( θ ) for ∇ f ( θ ), we obtain the associated mirrored Stein operator:
[Equation]
Appendix E. Riemannian Langevin Diﬀusions and Mirror-Langevin Diﬀusions
Zhang et al.
(2020) pointed out (15) is a particular case of the Riemannian LD.
However, they did not give an explicit derivation.
The Riemannian LD (Patterson and Teh, 2013;
Xifara et al., 2014; Ma et al., 2015) with ∇ 2 ψ ( · ) as the metric tensor is
[Equation]
To see the connection with mirror-Langevin diﬀusion, we would like to obtain the SDE that describes the evolution of η t = ∇ ψ ( θ t ) under the diﬀusion.
This requires the following theorem that provides the analog of the “chain rule” in SDEs.
Theorem 8 (Itˆo formula; Øksendal, 2003, Thm 4.2.1) Let ( x t ) t ≥ 0 be an Itˆo process in X ⊂ R d satisfying dx t
= b ( x t ) dt + σ ( x t ) dB t .
Let f ( x ) ∈ C 2 : R d → R d (cid:48) .
Then y t
= f
( x t ) is again an Itˆo process, and its i -th dimension satisﬁes
[Equation]
Substituting ∇ ψ for f in Thm. 8, we have the SDE of η t = ∇ ψ ( θ t ) as
[Equation]
p ( )    10 0.
p H ( ) 0.
1.0    p ( )    10 0.
[Figure]
where h ( θ t )
i =
Tr( ∇ 2 θ t ( ∇ θ t,i ψ ( θ t )) ∇ 2 ψ ( θ t ) − 1 ).
Moreover, we have [ ∇ 2 ψ ( θ t
[Equation]
∇
[Equation]
∇
[Equation]
= d
[Equation]
= d (cid:88)
[Equation]
Therefore, the η t diﬀusion is described by the SDE:
[Equation]
Appendix F. Mode Mismatch Under Transformations
In Fig.
5, we present two examples where the modes of density functions in θ and η space do not match.
We assume θ s follow Beta distributions on [0 , 1].
And we choose the negative entropy ψ ( θ )
= θ log θ
+ (1 − θ ) log(1 − θ ).
Then, the transformation is the logit function η = log( θ/ (1 − θ )) and its reverse is the sigmoid function θ = 1 / (1 + e − η ).
Appendix G. Background on Reproducing Kernel Hilbert Spaces
Let H be a Hilbert space of functions deﬁned on X and taking their values in R .
We say k is a reproducing kernel (or kernel) of H if ∀ x ∈ X , k ( x, · ) ∈ H and ∀ f ∈ H , (cid:104) f, k
( x, · ) (cid:105) H = f
( x ).
H is called a reproducing kernel Hilbert space (RKHS) if it has a kernel.
Kernels are positive deﬁnite (p.d.) functions, which means that matrices with the form ( k ( x i , x j )) ij are positive semideﬁnite.
For any p.d. function k , there is a unique RKHS with k as the reproducing kernel, which can be constructed by the completion of { (cid:80) ni
=1 a i
k
( x i , · ) , x i ∈ X , a i ∈ R , i ∈ N } .
Now we assume X is a metric space, k is a bounded continuous kernel with the RKHS H , and ν is a positive measure on X .
L 2 ( ν ) denote the space of all square-integrable functions w.r.t.
ν .
Then the kernel integral operator T k : L 2 ( ν )
→ L 2 ( ν ) deﬁned by
[Equation]
is compact and self-adjoint.
Therefore, according to the spectral theorem, there exists an at most countable set of positive eigenvalues { λ j } j ∈ J ⊂ R with λ 1 ≥ λ 2 ≥ . .
converging to zero and orthonormal eigenfunctions { u j } j ∈ J such that
[Equation]
and k has the representation k
( x, x (cid:48) ) = (cid:80) j ∈ J λ j u j
x ) u j ( x (cid:48) )
(Mercer’s theorem on noncompact domains), where the convergence of the sum is absolute and uniform on compact subsets of X × X (Ferreira and Menegatto, 2009).
Appendix H. Computational Details of Stein Variational Mirror Descent
The matrix eigenvalue problem is deﬁned as
[Equation]
where B t =
( k ( θ it , θ j t ))
ni,j =1 ∈ R n × n is the Gram matrix of pairwise kernel evaluations at particle locations, and the i -th element of v t,j is u t,j ( θ it ).
To compute ∇ θ K ψ,t ( θ, θ (cid:48) ), we diﬀerentiate both sides of (10) to ﬁnd that
[Equation]
This technique was used in Shi et al.
(2018) to estimate gradients of eigenfunctions w.r.t.
a continuous q .
Following their recommendations, we truncate the sum at the J -th largest eigenvalues according to a threshold ( τ ≥ (cid:80) J j =1 λ t,j / (cid:80) n j =1 λ t,j ) to ensure numerical stability.
Appendix I. Convergence Guarantees
In this section, we turn our attention to the convergence properties of our proposed methods.
For K t and (cid:15) t as in Alg. 1, let ( q ∞ t , q ∞ t,H ) represent the distributions of the mirrored Stein updates ( θ t , η t ) when θ 0 ∼ q ∞ 0 and η t +1 =
η t + (cid:15) t g ∗ q t ,K t ( θ t ) for t ≥ 0.
Our ﬁrst result, proved in App.
L.5, shows that if the Alg. 1 initialization q n 0 ,H = 1 n (cid:80) ni =
δ η i 0 converges in Wasserstein distance to a distribution q ∞ 0 ,H as n → ∞ , then, on each round t > 0, the output of Alg. 1, q nt = 1 n (cid:80)
ni
=
δ θ it , converges to q ∞ t .
Theorem 9 (Convergence of mirrored updates as n → ∞ ) Suppose Alg. 1 is initialized with q n 0 ,H = 1 n (cid:80) ni =
δ η i 0 satisfying W 1 ( q n 0 ,H , q ∞ 0 ,H ) → 0 for W 1 the L 1 Wasserstein distance.
Deﬁne the η -induced kernel K ∇ ψ ∗ ,t ( η, η (cid:48) )
(cid:44) K t ( ∇ ψ ∗ ( η ) , ∇ ψ ∗ ( η (cid:48) )) .
If, for some c 1 , c 2 > 0 ,
[Equation]
then, W 1 ( q nt,H , q ∞ t,H ) → 0 and q nt ⇒ q ∞ t for each round t .
Remark 10 The pre-conditions hold, for example, whenever ∇ log p H is Lipschitz, ψ is strongly convex, and K t = kI for k bounded with bounded derivatives.
Given a mirrored Stein operator (7), an arbitrary Stein set G , and an arbitrary matrixvalued kernel K we deﬁne the mirrored Stein discrepancy and mirrored kernel Stein discrepancy
[Equation]
and
[Equation]
The former is an example of a diﬀusion Stein discrepancy (Gorham et al., 2019) and the latter an example of a diﬀusion kernel Stein discrepancy (Barp et al., 2019).
Since the MKSD optimization problem (17) matches that in Thm.
3, we have that MKSD K ( q, p ) =
(cid:107) g ∗ q,K (cid:107) H K .
Our next result, proved in App.
L.6, shows that the inﬁnite-particle mirrored Stein updates reduce the KL divergence to p whenever the step size is suﬃciently small and drive MKSD to 0 if, for example, (cid:15) t =
Ω(MKSD K t ( q ∞ t , p ) α ) for any α > 0.
Theorem 11 (Inﬁnite-particle mirrored Stein updates decrease KL and MKSD)
Assume κ 1 (cid:44) sup θ (cid:107) K t ( θ, θ ) (cid:107) op < ∞ , and κ 2 (cid:44) (cid:80) di
=1 sup θ (cid:107)∇ 2 i,d
+ i K t ( θ, θ ) (cid:107) op < ∞ , ∇ log p H is L -Lipschitz, and ψ is α -strongly convex.
If (cid:15) t < 1 / ( (cid:107)∇ η t g ∗ q ∞ t ,K t (
[Equation]
then
[Equation]
Our last result, proved in App.
L.7, shows that q ∞ t ⇒ p if MKSD K k ( q ∞ t , p ) → 0.
Hence, by Thms. 9 and 11, n -particle MSVGD converges weakly to p if (cid:15) t decays at a suitable rate.
Theorem 12 ( MKSD K k determines weak convergence)
Assume p H is distantly dissipative
(Eberle, 2016) with ∇ log p H Lipschitz, ψ is strongly convex with continuous ∇ ψ ∗ , and k ( θ, θ (cid:48) )
=
κ ( ∇ ψ ( θ ) , ∇ ψ ( θ (cid:48) ))
for κ ( x, y )
=
( c 2 + (cid:107) x − y (cid:107) 22 ) β with β ∈ ( − 1 , 0) .
Then, q ∞ t ⇒ p if MKSD K k ( q ∞ t , p ) → 0 .
Remark 13 The pre-conditions hold, for example, for any Dirichlet target with negative entropy ψ .
Appendix J. Stein Variational Natural Gradient
The fact that SVMD recovers mirror descent as a special case is not only of relevance in constrained problems.
We next exploit the connection between MD and natural gradient descent discussed in Sec.
1 to design a new sampler – Stein Variational Natural Gradient (SVNG) – that more eﬃciently approximates unconstrained targets.
The idea is to replace the Hessian ∇ 2 ψ ( · ) in the SVMD dynamics dθ t = ∇ 2 ψ ( θ t ) − 1 g ∗ q t ,K ψ,t ( θ t ) with a general metric tensor G ( · ).
The result is the Riemannian gradient ﬂow
[Equation]
with
[Equation]
Given any initial particle approximation q 0 = 1 n (cid:80) ni =1 δ θ i 0 , we discretize these dynamics to obtain the unconstrained SVNG sampler of Alg. 2.
SVNG can be seen as an instance of MatSVGD (Wang et al., 2019) with a new adaptive time-dependent kernel G − 1 ( θ ) K G,t ( θ, θ (cid:48) )
G − 1 ( θ (cid:48) ).
However, similar to Prop. 6 and unlike the heuristic kernels of Wang et al.
(2019), SVNG reduces to natural gradient ascent for ﬁnding the mode of p ( θ ) when n = 1.
SVNG is well-suited to Bayesian inference problems where the target is a posterior distribution p ( θ ) ∝ π ( θ ) π
( y | θ ).
There, the metric tensor G ( θ ) can be set to the Fisher information matrix E π ( y | θ ) [ ∇ log π ( y | θ ) ∇ log π ( y | θ ) (cid:62) ] of the data likelihood π ( y | θ ).
Ample precedent from natural gradient variational inference (Hoﬀman et al., 2013; Khan and
Nielsen, 2018) and Riemannian MCMC (Patterson and Teh, 2013) suggests that encoding problem geometry in this manner often leads to more rapid convergence.
In App.
K.3, we demonstrate the advantages of SVNG on unconstrained large-scale posterior inference with the Fisher information metric.
Appendix K. Supplementary Experimental Details and Additional Results
In this section, we report supplementary details and additional results from the experiments of Sec. 3.
In Secs.
3.1 and 3.2, we use the inverse multiquadric input kernel k ( θ, θ (cid:48) )
= (1 + (cid:107) θ − θ (cid:48) (cid:107) 22 /(cid:96) 2 ) − 1 / 2 due to its convergence control properties (Gorham and Mackey, 2017).
In the unconstrained experiments of App.
K.3, we use the Gaussian kernel k ( θ, θ (cid:48) )
=
exp( −(cid:107) θ − θ (cid:48) (cid:107) 22 /(cid:96) 2 ) for consistency with past work.
The bandwidth (cid:96) is determined by the median heuristic (Garreau et al., 2017).
We select τ from { 0 . 98 , 0 . 99 } for all SVMD experiments.
For unconstrained targets, we report, for each method, results from the best ﬁxed step size (cid:15) ∈ { 0 . 01 , 0 . 05 , 0 . 1 , 0 .
5 , 1 } selected on a separate validation set.
For constrained targets, we select step sizes adaptively to accommodate rapid density growth near the boundary; speciﬁcally, we use RMSProp (Hinton et al., 2012), an extension of the AdaGrad algorithm (Duchi et al., 2011) used in Liu and Wang (2016), and report performance with the best learning rate.
Results were recorded on an Intel(R) Xeon(R)
CPU E5-2690 v4 @ 2.60GHz and an NVIDIA Tesla P100 PCIe 16GB.
K.1. Approximation quality on the simplex
The sparse Dirichlet posterior of Patterson and Teh (2013) extended to 20 dimensions features a sparse, symmetric Dir( α ) prior with α k = 0 .
1 for k ∈ { 1 , . .
. , 20 } and sparse
[Figure] [Figure]
target on a 20-simplex.
We plot the energy distance to a ground truth sample of size 1000 drawn by NUTS (Hoﬀman and Gelman, 2014).
count data n 1 = 90 , n 2 = n 3 = 5 , n j =
( j > 3), modeled via a multinomial likelihood.
The quadratic target satisﬁes log p ( θ )
= − 12 σ 2 θ (cid:62) Aθ
+ const, where we slightly modify the target density of Ahn and Chewi (2020) to make it less ﬂat by introducing a scale parameter σ = 0 .
A ∈ R 19 × 19 is a positive deﬁnite matrix generated by normalizing products of random matrices with i.i.d.
elements drawn from Unif[ − 1 , 1].
We initialize all methods with i.i.d samples from Dirichlet(5) to prevent any of the initial particles being too close to the boundary.
For each method and each learning rate we apply 500 particle updates.
For SVMD we set τ = 0 .
We search the base learning rates of RMSProp in { 0 . 1 , 0 . 01 , 0 . 001 } for SVMD and MSVGD.
Since projected SVGD applies updates in the θ space, the appropriate learning rate range is smaller than those of SVMD and MSVGD.
There we search the base learning rate of RMSProp in { 0 . 01 , 0 . 001 , 0 . 0001 } .
For all methods the results under each base learning rate are plotted in Fig.
K.2. Conﬁdence intervals for post-selection inference
Given a dataset X ∈ R ˜ n × p , y ∈ R ˜ n , the randomized Lasso (Tian and Taylor, 2018) solves the following problem:
[Equation]
where G is a user-speciﬁed log-concave distribution with density g .
We choose G to be zeromean independent Gaussian distributions while leaving its scale and the ridge parameter (cid:15) to be automatically determined by the randomizedLasso function of the selectiveInference package.
We initialize the particles of our SVMD and MSVGD in the following way: First, we map the solution ˆ β E to the dual space by ∇ ψ .
Next, we add i.i.d.
standard Gaussian noise to n copies of the image in the dual space.
Finally, we map the n particles back to the primal space by ∇ ψ ∗ and use them as the initial locations.
Below we discuss the remaining settings and additional results of the simulation and the HIV-1 drug resistance experiment separately.
Simulation
In our simulation we mostly follow the settings of Sepehri and Markovic (2017) except using a diﬀerent penalty level λ recommended in the selectiveInference R package.
We set ˜ n = 100 and p = 40.
The design matrix X is generated from an equi-correlated model, i.e., each datapoint x i ∈ R p is generated i.i.d.
from N (0 , Σ) with Σ ii = 1 , Σ ij = 0 .
3 ( i (cid:54) =
j ) and then normalized to have almost unit length.
The normalization is done by ﬁrst centering each dimension by subtracting the mean and dividing the standard deviation of that column of X , then additionally multiplying 1 / ˜ n 1 / 2 .
y is generated from a standard Gaussian which is independent of X , i.e., we assume the global null setting where the true value of β is zero.
We set λ to be the value returned by theoretical.lambda of the selectiveInference R package multiplied a coeﬃcient 0 .
7˜ n , where the 0 .
7 adjustment is introduced in the test examples of the R package to reduce the regularization eﬀect so that we have a reasonably large set of selected features when p = 40.
The base learning rates for SVMD and MSVGD are set to 0 .
and we run them for T = 1000 particle updates.
τ is set to 0 .
98 for SVMD.
Our 2D example in Fig.
3 (left) is grabbed from one run of the simulation where there happen to be only 2 features selected by the randomized Lasso.
The selective distribution in this case has log-density log p ( θ )
= − 8 .
39859 θ 1 +1 .
θ 2 +2 .
39751) 2 +(1 . 18099 θ 2 − 1 . 46104) 2 )
+ const , θ 1 , 2 ≥ 0.
The error bars for actual coverage levels in Fig.
4( a ) and Fig.
4( b ) are 95% Wilson intervals (Wilson, 1927), which is known to be more accurate than ± 2 standard deviation intervals for binomial proportions like the coverage.
In Fig.
8( a ) and Fig.
8( b ) we additionally plot the average length of the conﬁdence intervals w.r.t.
diﬀerent sample size N and nominal coverage levels.
For all three methods the CI widths are very close, although MSVGD consistently has wider intervals than SVMD and selectiveInference .
This indicates that
SVMD can be preferred over MSVGD when both methods produce coverage above the nominal level.
HIV-1 drug resistance
We take the vitro measurement of log-fold change under the 3TC drug as response and include mutations that had appeared 11 times in the dataset as regressors.
This results in ˜ n = 663 datapoints with p = 91 features.
We choose λ to be the value returned by theoretical.lambda of the selectiveInference R package multiplied by ˜ n .
The base learning rates for SVMD and MSVGD are set to 0 .
and we run them for
T = 2000 particle updates.
τ is set to 0 .
99 for SVMD.
[Figure]
K.3. Large-scale posterior inference with non-Euclidean geometry
Finally, we demonstrate the advantages of exploiting non-Euclidean geometry by recreating the real-data large-scale Bayesian logistic regression experiment of Liu and Wang (2016) with 581,012 datapoints and d = 54 feature dimensions.
Here, the target p is the posterior distribution over logistic regression parameters.
We adopt the Fisher information metric tensor G , compare 20-particle SVNG to SVGD and its prior geometry-aware variants RSVGD (Liu and Zhu, 2018) and MatSVGD with average and mixture kernels (Wang et al., 2019), and for all methods use stochastic minibatches of size 256 to scalably approximate each log likelihood query.
In Fig.
9, all geometry-aware methods substantially improve the log predictive probability of SVGD.
5 SVNG also strongly outperforms RSVGD and converges to its maximum test probability in half as many steps as MatSVGD (Avg) and more rapidly than MatSVGD (Mixture).
The Bayesian logistic regression model we consider is (cid:81) ˜
ni =1 p ( y i | x i , w ) p ( w ), where p ( w ) =
N ( w | 0 , I ), p ( y i | x i , w ) =
Bernoulli( σ ( w (cid:62) x i )).
The bias parameter is absorbed into into w by adding an additional feature 1 to each x i .
The gradient of the log density of the posterior distribution of w is ∇ w log p ( w |{ y i , x i } Ni =1 ) =
(cid:80) Ni
=1 x i ( y i − σ ( w (cid:62) x i ))
− w. We choose the
[Figure]
s
[Figure]
l og p r
e d i c
t
i v e p r o b a b ili t y
[Figure]
metric tensor ∇ 2 ψ ( w ) to be the Fisher information matrix (FIM) of the likelihood:
[Equation]
Following Wang et al.
(2019), for each iteration r ( r ≥ 1), we estimate the sum with a stochastic minibatch B r of size 256: ˆ F B r =
[Equation]
∈B r
[Equation]
(cid:62)
[Equation]
(cid:62) i and approximate the FIM with a moving average across iterations:
[Equation]
where ρ r = min(1 − 1 /r, 0 . 95) .
To ensure the positive deﬁniteness of the FIM, a damping term 0 .
I is added before taking the inverse.
For RSVGD and SVNG, the gradient of the inverse of FIM is estimated with ∇ w j F − 1 ≈ − ˆ F − 1 r ( ˆ
[Equation]
F ) ˆ F − 1 r , where ˆ ∇ rw
[Equation]
F
[Equation]
r .
We run each method for T = 3000 particle updates with learning rates in { 0.01, 0.05, 0.1, 0.5, 1 } and average the results for 5 random trials.
τ is set to 0 .
98 for SVNG.
For each run, we randomly keep 20% of the dataset as test data, 20% of the remaining points as the validation set, and all the rest as the training set.
The results of each method on validation sets with all choices of learning rates are plotted in Fig.
We see that
the
SVNG updates are very robust to the change in learning rates and is able to accommodate very large learning rates (up to 1) without a signiﬁcant loss in performance.
The results in
Fig.
9 are reported with the learning rate that performs best on the validation set.
Appendix L. Proofs L.1. Proof of Prop.
Proof Fix any g ∈ G ψ .
Since g and ∇ g are bounded and ∇ 2 ψ ( θ ) − 1 ∇ log p ( θ ) and ∇ · ∇ 2 ψ ( θ ) − 1 are p -integrable, the expectation E p [( M p,ψ g )
( θ )] exists.
Because Θ is convex, Θ r is bounded and convex with Lipschitz boundary.
Since p ∇ 2 ψ − 1 g ∈ C
[Equation]
we have
[Equation]
(by dominated convergence)
[Equation]
(by Cauchy-Schwarz)
[Equation]
(by assumption).
L.2. Proof of Thm. 3: Optimal mirror updates in RKHS
Proof Let e i denote the standard basis vector of R d with the i -th element being 1 and others being zeros.
Since m ∈ H K
[Equation]
we have
[Equation]
where we deﬁne the divergence of a matrix as a vector whose elements are the divergences of each row of the matrix
[Equation]
Then
[Equation]
we write (6) as
[Equation]
Therefore, the optimal direction in the H K norm ball B H K = { g : (cid:107) g (cid:107) H K ≤ 1 } that minimizes (6) is g ∗ t ∝ g ∗ q t ,
K = E q t [ M p,ψ K ( · , θ )] .
L.3. Proof of Thm. 4: Mirrored SVGD updates
Proof A p.d. kernel k composed with any map φ is still a p.d. kernel.
To prove this, let { x 1 , . . .
, x p } = { φ ( η 1 ) , . . .
, φ ( η n ) } , p ≤ n .
Then (cid:88)
[Equation]
where β (cid:96) =
(cid:80) i ∈ S (cid:96) α i , S (cid:96) =
{ i : φ ( η i ) =
x (cid:96) } .
Therefore, k ψ ( η, η (cid:48) )
=
k ( ∇ ψ ∗ ( η ) , ∇ ψ ∗ ( η (cid:48) )) is a p.d. kernel.
Plugging K = kI into Lem.
14, for any θ (cid:48) ∈ Θ and η (cid:48)
= ∇ ψ ( θ (cid:48) ), we have g
[Equation] [Equation]
[Equation]
L.4. Proof of Prop
Proof When n = 1, λ 1 = k
( θ t , θ t ) , u 1 = 1, and thus K ψ,t ( θ t , θ t )
=
k ( θ t , θ t ) ∇ 2 ψ ( θ t ).
L.5. Proof of Thm. 9: Convergence of mirrored updates as n → ∞
Proof The idea is to reinterpret our mirrored updates as one step of a matrix SVGD in η space based on Lem.
and then follow the path of Gorham et al.
(2020, Thm. 7).
Assume that q nt,H and q ∞ t,H have integrable means.
Let η n , η ∞ be an optimal Wasserstein-1 coupling of q nt,H and q ∞ t,H .
Let Φ q t ,K t denote the transform through one step of mirrored update: θ t = ∇ ψ (cid:63) ( η t ), η t +1 = η t + (cid:15) t g ∗ q t ,K t ( θ t ).
Then, with Lem.
14, we have
[Equation]
η n ) ∇ log p H ( η n ) + ∇ η n · K ∇ ψ ∗ ,t ( η (cid:48) , η n ))]
(cid:107)+ (cid:15) t (cid:107) E η n ,η ∞ [ K ∇ ψ ∗ ,t ( η (cid:48) , η n ) ∇ log p H ( η n ) + ∇ η n · K ∇ ψ ∗ ,t ( η, η n ) − ( K ∇ ψ ∗ ,t ( η (cid:48) , η ∞ ) ∇ log p H ( η ∞ ) + ∇ η ∞ · K ∇ ψ ∗ ,t ( η (cid:48) , η ∞ ))]
(cid:107)≤ (cid:107) η − η (cid:48) (cid:107) 2 + (cid:15) t c 1 (1 + E
[ (cid:107) η n (cid:107) 2 )] (cid:107) η − η (cid:48)
(cid:107) 2 + (cid:15) t c 2 (1 + (cid:13)(cid:13) η (cid:48)
(cid:13)(cid:13) 2 ) E η n ,η ∞ [ (cid:107) η n − η ∞ (cid:107) 2 ] = (cid:107) η − η (cid:48) (cid:107) 2 + (cid:15) t c 1 (1 + E q nt,H
[ (cid:107)·(cid:107) 2 ]) (cid:107) η − η (cid:48)
(cid:107) 2 + (cid:15) t c 2 (1 + (cid:13) (cid:13) η (cid:48) (cid:13) (cid:13) 2 ) W 1 ( q nt,H , q ∞ t,H ) .
Since Φ q t ,K t ( η n ) ∼ q nt +1 ,H , Φ q t ,K ( η ∞ ) ∼ q ∞ t +1 ,H , we conclude
[Equation]
The ﬁnal claim q nt ⇒ q ∞ t now follows by the continuous mapping theorem as ∇ ψ ∗ is continuous.
L
6. Proof of Thm. 11: Inﬁnite-particle mirrored Stein updates decrease KL and MKSD
Proof Let T q ∞ t ,K t denote transform of the density function through one step of mirrored update: θ t = ∇ ψ (cid:63) ( η t ), η t +1 =
η t + (cid:15) t g ∗ q ∞ t ,K t ( θ t ).
Then
[Equation]
where we have used the invariance of KL divergence under reparameterization: KL( q t (cid:107) p )
=
KL( q t,H (cid:107) p H ) .
Following Liu (2017)
[Equation]
we bound the diﬀerence of the ﬁrst two terms as
[Equation]
where η t (
[Equation]
= − (cid:90)
[Equation]
and bound the log determinant term using Lem.
[Equation]
The next thing to notice is that E η t ∼ q ∞ t,H [ ∇ log p H ( η t ) (cid:62)
g ∗ q ∞ t ,K t ( θ t )
+ Tr( ∇ η t g ∗ q ∞ t ,K t ( θ t ))] is the square of the MKSD in (17).
We can show this equivalence using the identity proved in
[Equation] [Equation]
Finally
[Equation]
we are going to bound (cid:107) g ∗ q ∞ t ,K t ( θ t ) (cid:107)
and (cid:107)∇ η t g ∗ q ∞ t ,K t ( θ t ) (cid:107) 2 F .
From the assumptions we have ψ is α -strongly convex and thus ψ ∗ is 1 α -strongly smooth (Kakade et al., 2009), therefore (cid:107)∇ 2 ψ ∗ ( · ) (cid:107) 2 ≤ 1 α .
By Lem.
[Equation]
we know
[Equation]
where ∇ 2 i,d + i K ( θ, θ ) denotes ∇ 2 θ i ,θ (cid:48)
i K ( θ, θ (cid:48) )
| θ (cid:48)
= θ .
Combining all of the above
[Equation]
we have
KL
[Equation]
q ∞
[Equation]
≤ −
[Equation]
Plugging in the deﬁnition of κ 1 and κ 2 ﬁnishes the proof.
L.7. Proof of Thm. 12: MKSD K k determines weak convergence
Proof According to Thm. 4,
[Equation]
where q H ( η ) denotes the density of η = ∇ ψ ( θ ) under the distribution θ ∼ q .
From the assumptions we have k ( θ, θ (cid:48) )
=
κ ( ∇ ψ ( θ ) , ∇ ψ ( θ (cid:48) )).
With this speciﬁc choice of k , the squared
MKSD is
[Equation]
The ﬁnal expression in (19) is the squared kernel Stein discrepancy (KSD) (Liu et al., 2016;
Chwialkowski et al., 2016; Gorham and Mackey, 2017) between q H and p H with the kernel κ : KSD κ ( q H , p H ) 2 .
Recall that it is proved in Gorham and Mackey (2017, Theorem 8) that, for κ ( x, y )
= ( c 2 + (cid:107) x − y (cid:107) 22 ) β with β ∈ ( − 1 , 0) and distantly dissipative p H with
Lipschitz score functions, q H ⇒ p
H if KSD κ ( q H , p H ) → 0.
The advertised result ( q ⇒ p if
MKSD K k ( q, p ) → 0) now follows by the continuous mapping theorem as ∇ ψ ∗ is continuous.
Appendix M. Lemmas
Lemma
Let K ∇ ψ ∗ ( η, η (cid:48) )
(cid:44) K ( ∇ ψ ∗ ( η ) , ∇ ψ ∗ ( η (cid:48) )) .
The mirrored updates g ∗ q t ,K in (8) can be equivalently expressed as
[Equation]
Proof We will use the identity proved in Lem. 15.
[Equation]
(by change-of-variable formula)
[Equation]
(by applying Lem.
15 to each row of K ( · , θ ))
[Equation]
where A : ,j denotes the j -th column of a matrix
A
[Equation]
Lemma
For a strictly convex function ψ ∈ C 2 : R d → R and any vector-valued g ∈ C 1 : R d → R d , the following relation holds:
[Equation]
Proof By the product rule of diﬀerentiation:
[Equation]
This already gives us the ﬁrst term on the right side.
Next, we have
[Equation]
Plugging the above relation into (20) proves the claimed result
[Equation]
Lemma 16 (Liu, 2017, Lemma A.1)
Let A be a square matrix, and 0 < (cid:15) < 12 (cid:107) A +
A (cid:62) (cid:107) op .
Then,
[Equation]
where (cid:107) · (cid:107) F denotes the Frobenius norm of a matrix.
Lemma 17 Let K be a matrix-valued kernel and H K be the corresponding RKHS.
Then, for any f ∈ H K ( f is vector-valued)
[Equation]
we have
[Equation]
where (cid:107) · (cid:107) op denotes the operator norm of a matrix induced by the vector 2-norm.
Proof We ﬁrst bound the (cid:107) f
( x ) (cid:107) 2 as
[Equation]
The second result follows similarly,
[Equation]
[Figure]
V a li d a t i o n l og p r e
d
i c t i v e p r o b a b ili t y
[Figure]
a )
[Figure]
og
[Figure]
V a li d a t i o n l og p r e
d
i c t i v e p r o b a b ili t y
[Figure]
( c )
[Figure]
V a li d a t i o n l og p r e
d
i c t i v e p r
[Figure]
b a b ili t y
[Figure]
V a li d a t i o n l og p r e
d
i c t
i
[Figure]