Sampling with Mirrored Stein Operators
Authors: Jiaxin Shi
Chang Liu
Lester Mackey
Venue: ICLR 2022 Spotlight (2022)
Keywords: Stein's method
Sampling
Mirror descent
Natural gradient descent
Probabilistic inference
Bayesian inference
Post-selection inference
Stein operators

Abstract: We introduce a new family of particle evolution samplers suitable for constrained domains and non-Euclidean geometries. Stein Variational Mirror Descent and Mirrored Stein Variational Gradient Descent minimize the Kullback-Leibler (KL) divergence to constrained target distributions by evolving particles in a dual space defined by a mirror map. Stein Variational Natural Gradient exploits non-Euclidean geometry to more efficiently minimize the KL divergence to unconstrained targets. We derive these samplers from a new class of mirrored Stein operators and adaptive kernels developed in this work. We demonstrate that these new samplers yield accurate approximations to distributions on the simplex, deliver valid confidence intervals in post-selection inference, and converge more rapidly than prior methods in large-scale unconstrained posterior inference. Finally, we establish the convergence of our new procedures under verifiable conditions on the target distribution.

Published as a conference paper at ICLR 2022
SAMPLING WITH MIRRORED STEIN OPERATORS
Jiaxin Shi Chang Liu
Lester Mackey!
Microsoft Research New England
2 Microsoft Research Asia
miaxinshi, chang. liu, lmackey) maackey>emicrosoft.ccm
ABSTRACT
We introduce a new family of particle evolution samplers suitable for constrained domains and non-Euclidean geometries. Stein Variational Mirror Descent and Mirrored Stein Variational Gradient Descent minimize the Kullback-Leibler (KL) divergence to constrained target distributions by evolving particles in dual space defined by a mirror map. Stein Variational Natural Gradient exploits non-Euclidean geometry to more efficiently minimize the KL divergence unconstrained targets. We derive these samplers from a new class of mirrored Stein operators and adaptive kernels developed in this work. We demonstrate that these new samplers yield accurate approximations to distributions on the simplex, deliver valid confidence in- tervals in post-selection inference, and converge more rapidly than prior methods large-scale unconstrained posterior inference. Finally, we establish the convergence of our new procedures under verifiable conditions the target distribution.
INTRODUCTION
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 (SVGD, Liu & Wang, 2016) tackle this challenge by applying deterministic updates to particles using operators based on Stein's method (Stein, 1972:
Gorham & Mackey, 2015; Oates et al., 2017: Liu et al., 2016; Chwialkowski et al., 2016; Gorham & Mackey, 2017) and reproducing kernels (Berlinet Lebber (KI.) divergence. SVGD has found great success in approximating unconstrained distributions for probabilistic learning (Feng et al., 2017: Haarnoja et al., 2017; Kim al., 2018) but breaks down for constrained targets, like distributions on the simplex (Patterson & Teh, 2013) or the targets of post-selection inference (Taylor & Tibshirani, 2015; Lee et al., 2016; Tian et al., 2016), and fails to exploit informative non-Euclidean geometry (Amari, 1998).
Thomas-Agnas 2011) 1) to sequentially minimize
In this work, we derive a family of particle evolution samplers suitable for target distributions with constrained domains and non- -on-Euulidean geometries. Our development draws inspiration from mirrrr descent (MD) (Nemirovskij & Yudin, 1983), a first-order optimization method that generalizes gradient descent with Fun-Fuclidean geometry. To sample from distribution with constrained support, our method first maps particles to a dual space. There, we update particle locations using new class of mirrored Stein operators and adaptive reproducing kernels introduced in this work. Finally, the dual particles are mapped back to sample points in the original space, ensuring that all constraints are satisfied. We illustrate this procedure in Fig. In Sec. 3, we develop two algorithms Minrored SVGD (MSVGD) and Stein Variational Mirror Descent (SVMD) with different updates in the dual space; when only a single particle used. MSVGD reduces to gradient ascent on the log dual space density, and SVMD reduces to mirror ascent on the log target density. In addition, by exploiting the connection between MD and natural gradient descent (Amari, 1998; Raskutti Mukherjee, 2015), we develop a third algorithm Stein Variational Natural Gradient (SVNG) that extends SVMD to unconstrained targets with non-Euclidean geometry.
In Sec. 5, we demonstrate the advantages of our algorithms benchmark simplex-constrained problems from the literature, constrained sampling problems in post-selection inference (Taylor & Tibshirani, 2015; Lee et al., 2016; Tian et al., 2016), and unconstrained large-scale posterior inference with the Fisher information metric, Finally, we analyze the convergence of our mirrored algorithms in Sec. 6 and discuss our results in Sec. 7.

---

Published as a conference paper at ICLR 2022
n
Og
e E.wa MMp
8
1]++1
Figure 1: Updating particle approximations in constrained domains O Standard updates like SVGD (dashed arrow) can push particles outside of the support. Our mirrored Stein updates in Alg. (solid arrows) preserve the support by updating particles in dual space and mapping back to θ.
Related work Our mirrored Stein operators (6) are instances of diffusion Stein operators in the sense of Gorham & Mackey (2017), but their specific 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 MCMC (e.g., Welling & Teh, 201 I; Simsekli et al., 6;;6; Dalalyan, 2017: Durmus al., 2018; Ma et al., 2019) and SVGD-like sampling methods (e.g., Liu et al., 2019a;;; Zhu et al., 2020; Zhang et al., 2020a). The closest to work in this space is the recent marriage of mirror descent and MCMC. For example, Hsich et al. (2018) propose to run Langevin Monte Carlo (LMC, Euler discretization of the Langevin diffusion) in mirror space. Zhang et al. (2020b) analyze the convergence properties of the mirror- Langevin diffusion, Chewi et al. (2020) demonstrate its advantages over the Langevin diffusion when using Newton-type metric, and Ahn & Chewi (2020) study its discretization for MCMC sampling in constrained domains. Relatedly, Patterson & Teh (2013) proposed stochastic Riemannian LMC for sampling on the simplex.
Several modifications of SVGD have been proposed to incorporate geometric information. Rieman- nian SVGD (RSVGD, Liu & 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
mirror space, and do not reduce to natural gradient descent when single particle is used. They also reportedly do not perform well when with scalable stochastic estimates of logp. Stein Variational Newton (SVN, Detommaso et al., 2018; Chen er 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 instance of matrix SVGD (MatVVGD, Wang et al., 2019) with an adaptive time-dependent kernel discussed in Sec. 4.4, a choice that is not explored in Wang et al. (2019) and which recovers natural gradient descent when tions of Wang et al. (2019). None of the aforementioned works provide convergence guarantees, and neither SVN nor matrix SVGD deals with constrained domains.
unlike the heuristic kernel construc-
2 BACKGROUND: MIRROR DESCENT AND NON-EUCLIDEAN GEOMETRY
Standard gradient descent can be viewed as optimizing local quadratic approximation to the target function f: 0.41 = argmingge (0,) T + 2e, | 02 11.... When C Rd is constrained, can be advantageous to replace ||- 1/2 with a function Vi that reflects the geometry of a problem (Nemirovskij & Yudin, 1983; Beck & Teboulle, 2003):
0++1 argmingge (f.)) (8,
(1)
We consider the mirror descent algorithm (Nemirovskij & Yudin, 1983; Beck & Teboulle, 2003)
which chooses to be the Bregman divergence induced by a strongly convex, essentially smooth
function & θ R U {oo]: (0 0²) v²() (00")= Vv'(0") (0 0'). When θ is (d + 1)-
simplex {0 Ed i=1 0, I and 2 for i E [d]] common choice of 1b is the negative entropy
(0) = dd+1 0; log 0,, for auta A Σí 0g. The solution of (1) is given by
O+41 V (Vi)(0.) ex (())))
where v0 (n) A supaee 17 (0)) is the convex conjugate of yb and V is bijection from θ to dom )) with inverse map (Vb)- Vv=" We can view the update in (2) as first mapping 8, to 7)
by Vvh applying the update 7/+1 п) (f,)), and mapping back through d Val (7+1))
d continuously differentiable the interior of with [())ll)) whenever 0. e de.
2

---

Published as a conference paper at ICLR 2022
Mirror descent can also be viewed as discretization of the continuous- time dynamics drte _V, 0e = Vii (nt), which is equivalent to the Riemannian gradient flow (see App. A):
d0, - -v2 f(O,)dt, or equivalenty, drp = - - (m)) Vv f(Ve' (m))) (3)
where (0) and (n) are Riemannian metric tensors. In information geometry, the discretiza- tion of (3) is known as natural gradient descent (Amari, 1998). There is considerable theoretical and practical evidence (Martens, 2014) showing that natural gradient works efficiently in learning.
3 STEIN'S IDENTITY AND MIRRORED STEIN OPERATORS
Stein's 's identity (Stein, 1972) is a tool for characterizing a target distribution P using so-called Stein operator. We assume P has differentiable density p with a closed convex support θ C Rd. A Stein operator Sp takes as input functions 9 from Stein ser G and outputs mean-zero functions under p.
Ee- (S,p)(() = 0, for all E g.
(4)
Gorham & Mackey (2015) proposed the Langevin Stein operator given by
(0) = g(0) -(0) ([0))[logr(0))
V g
(5)
where g is vector-valued function and V is its divergence. For an unconstrained domain with Il log p(0)ll2] < oo, Stein's s identity (4) holds for this operator whenever g E C² is bounded and Lipschitz by (Gorham et al., 2019, proof of Prop. 3). However, on constrained domains O Stein's s identity fails to hold for many reasonable inputs y if p is non- non-vanishing or explosive at the boundary.
Motivated by this deficiency and by desire to exploit non-Euclidean geometry, we propose an alternative mirrored Stein operator,
(Mpxg)(()- = g T (0)
+ V (V²)())))(0)),
(6)
where & is a strongly convex, essentially smooth function as in Sec. 2 with (V2))-1 differentiable and Lipschitz on O We derive this operator from the (infinitesimal) generator of the mirror-Langevin
diffusion (19) in App. C. The following result, proved in App. I.1, shows that Mp.e generates mean-zero functions under whenever 22w-1 suitably cancels the growth of p at the boundary.
Proposition Suppose that (2) logr(0) and V (20(0)-1 are p-integrable. If lim Sae, p(0)IlV²<>(0) (n, = for O & {0 E θ : 10 S r and n, the outward unit normal vector to 00, at 0, then Mp)
0 if E C³ is bounded Lipschitz. Example (Dirichlet p, Negative entropy $). When 01--11 ~ Dir(a) for a E 1 even setting g(0) = 1 in (5) need not cause the identity to hold when some aj I 1. However, when (0) = Σ 02 log we show in App. B that the conditions of Prop. are met for any a. Remarkably, the mirror-Langevin diffusion for our choice of & is the Wright-Fistee Fisher diffusion (Ethier, 1976) which Gan et al. (2017) recently used to bound distances to Dirichlet distributions.
4 SAMPLING WITH MIRRORED STEIN OPERATORS
Liu & Wang (2016) pioneered the idea of using Stein operators to approximate a target distribution with particles. Their popular SVGD algorithm updates each particle in its approximation by applying the update rule 0++1 = 0: + e.g, (0,) for a chosen mapping 9: Rd -> I&d. Specifically, SVGD
chooses the mapping of that leads to the largest decrease in KL divergence to in the limit as 6 + 0.
The following theorem summarizes their main findings.
Theorem 2 (Liu & Wang. 2016, Thm. 3.1). Suppose (0.) satisfies d6, = (O,)dt for bounded Lipschitz 9t E C Rd + Rd and that 0, has density qi with Eg, [|[V log gell2] < If KL 1p) 2 Egg [log(q]/ p)] exists then, for the Langevin Stein operator S, (5),
d KL - Ee,,ag. (Spgt)
de
FFor closed convex set whose boundary 00 can locally represented F(B....., 02) 0, unit VF(0)...9..)
normal vector is defined n0)) = + [VF Palla
3

---

Published as a conference paper at ICLR 2022
Algorithm 1 Mirrored Stein Variational Gradient Descent & Stein Variational Mirror Descent
Input: density p on θ kernel k, mirror function 1b, particles (0)) C O step sizes (et)==
Init: no < (v((%) for i E [n]
for = 0 do
if SVMD then Kt Kv, (13) else K₁ t- k1 (MSVGD)
for E [n], Th1 - ni + tem
+ Ei=1 Mp K.(0), (0t.0))) (for Me.. defined in Thm. 3)
for E [n], f: + Vib"
return (0/+1)=-1
To improve its current particle approximation, SVGD finds the choice of 9r that most quickly decreases (L- [p)) at time i, i.e., minimizes (I. [p) over a set of candidate directions gt- SVGD finds It in a reproducing kernel Hilbert space (RKHS, Berlinet & Thomas-Agnan, 201 1) norm ball By = {g: : || S 1). where Hd the product RKHS containing vector-valued functions with each component in the RKHS H of k. Then the optimal 96 E Byg that minimizes (7) is
S < Ige,k - Eg, [k((0, V log p(O,) + k(0, :] = Es.~q. (s, 0.)).
where we let (0,00) k(, 0')4, and (,,,(,,)) 0) denotes applying Sp each row of (x 0). SVGD has found great success in approximating unconstrained target distributions p but breaks down for constrained targets and fails to exploit non-Euclidean geometry. Our goal is to develop new particle evolution samplers suitable for constrained domains and non- oon-Euclidean geometries.
4.1 MIRRORED DYNAMICS
SVGD encounters two difficulties when faced with constrained support. First, the SVGD updates can push the random variable Og outside of its support O, rendering all future updates undefined. Second, Stein's identity (4) often fails hold for candidate directions in Bx (cf. Ex. 1). When this occurs, SVGD need not converge to p as p is not stationary point of its dynamics (i.e., dt (gelll)))l [gp) * O when q: = p). Inspired by mirror descent (Nemirovskij & Yudin, 1983), we consider the following mirrored dynamics
d KLL
8: V Vab* (n:) for dn, = (f, or. equivalently, d8, = V V (0,)
9.
(8)
where St : θ - Ted now represents the update direction in v) space. The inverse mirror map Vel" automatically guarantees that 0, belongs the constrained domain e. Since v is strongly convex and 22 -1 is bounded Lipschitz, from Thm. it follows for any bounded Lipschitz gr that
d (L = -Es.gg (.)......(.))).
where Mp is the mirrored Stein operator (6). In the following sections, we propose three new deterministic sampling algorithms by seeking the optimal direction şi that minimizes (9) over different function classes. Thm. 3 (proved in App. 1.2) forms the basis of our analysis.
( Theorem 3 (Optimal mirror updates in RKHS). Suppose (0,) follows the mirrored dynamics (8).
Let Hr denote the RKHS matrix-valued kernel K e θ gixd (Micchelli & Pontil, 2005).
Then, the optimal direction of 9t that minimizes (9) in the norm ball BHx & [g : 1199l1 Hx 1} is
gt. 9r K A Eg, (Mr...iMppp (M..KK(..
(10)
where Mp K(, applies Mp.y (6) to each row of the matrix-valued function Kg = K(.
4.2 MIRRORED STEIN VARIATIONAL GRADIENT DESCENT
Following the pattern of SVGD, one choose the K of Thm. 3 to be Ks 0') =(0)) 0') where
k is any scalar-valued kernel. In this case, the resulting update Ig..K, (.) (-)= Es,- MM Ka(,, 0.)] equivalent to running SVGD in the dual v space before mapping back to O.
Theorem 4 (Mirrored SVGD updates). In the setting of Thm. 3, let kg(n,,1))
ke o (n), (r))) (1) = (v))"())) det (2r0(() denote the density of n = V2)(0) when O p, and qtH denote the distribution of qt under the mirrored dynamics (8). If K2 kI,
Fq. Kv (0') Emrag. (R n/ log pm (na) V (1).)))) V0' e,, V.'(') (11)
4

---

Published as a conference paper at ICLR 2022
The proof is in App. 1.3. By discretizing the dynamics dnj = K, (8t)dt and initializing with any particle approximation do n -i=1 6og. we obtain Mirrored SVGD (MSVGD), our first algorithm for sampling in constrained domains. The details are summarized in Alg.
When only single particle is used (n 1) and the differentiable input kernel satisfies Vk(0,0) = 0. the MSVGD update (11) reduces to gradient descent on log pn (1). Note however that the modes of the mirrored density (n) need not match those of the target density p(0) (see the examples in App. E). Since we are primarily interested in the θ-space density, is natural to ask whether there exists mirrored dynamics that reduces to finding the mode of p(0) in this limiting case. In the next section, we give an answer to this question by designing an adaptive reproducing kernel that yields a mirror descent-like update.
4.3 STEIN VARIATIONAL MIRROR DESCENT
Our second sampling algorithm for constrained problems is called Stein Variational Mirror De- scent (SVMD). We start by introducing new matrix-valued kernel that incorporates the metric 2 and evolves with the distribution qt.
Definition (Kernels for SVMD). Given continuous scalar-valued kernel consider the Mercer representation' k(0, = Ei>1
w.E.f. qt, where ut,i is an eigenfunction satisfying
Ee (6, (.)),,,;(e)) = (0))
(12)
For k. (0,0' Eizi 1/2 We, ur, (0)u,,(000))) we define the adaptive SVMD kernel at time I
(0, 8') a Eg,~g. [k. /2 (0, 8, ( (0,
(13)
By Thm. 3, the optimal update direction for the SVMD kernel ball is L2. Ke.. = FEg (Mp K (.))) We obtain the SVMD algorithm (summarized in Alg. 1) by discretizing dn: Kg. Ke. (Ot)dt and initializing with go = n Ei= ōøg. Because of the discrete representation of qt. KV takes the form
1/2 1/2 K (0, Ei= Ei- Ai.3 X..j
Un (0) (f) (0')
Te,iy 1 n Σr U4 (0;)) y (0 ( (e)))
Here both Ae,j and (e;) can be computed by solving matrix eigenvalue problem involving the particle set ) B, = nt., where B, = (A:0), 8 5) E R is the Gram matrix of pairwise kernel evaluations at particle locations, and the i-th element of et. is uz (0)).. To compute Vo.w.e(0, 0'). we differentiate both sides of (12) to find that Vun, (0) = A.. Σ Ut.j.a (o (0, 0;). This technique was used in Shi et al. (2018) to estimate gradients of cigenfunctions w.r.t. a continuous Following their recommendations we truncate the sum at the J-th largest eigenvalues according to a threshold (T 2 E/= / Ej=1 Xt,z) to ensure numerical stability.
Notably, SVMD differs from MSVGD only in its choice of kernel, but, whenever Vk((0,) 0) 0, this change sufficient to exactly recover mirror descent when n Proposition 5 (Single-particle SVMD is mirror descent). Ifn = 1, then one step of SVMD becomes
M++1 = 1/- e (k(0,, 0,) V log p(0.) + Vk(0., 0,)), 0.41 7pb"((e+))
4.4 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. 2 to design new sampler- - Stein Variational Natural Gradient (SVNG) - that more efficiently approximates unconstrained targets. The idea is to replace the Hessian V-+((-) in the SVMD dynamics d6, (2)) Ky. (0z) (f:) with general metric tensor G(.). The result is the Riemannian gradient flow
d0,
G(0.)-+ Ko, (0.) with Ko., 0') A Eo,ng (,72(0,00))0)))))2300000
(14)
'See App. F for background Mercer representations in non-compact domains.
5

---

Published as a conference paper at ICLR 2022
Spazse Ditichlet
Quadratic
1.5
Projectdd SVGD
SVMD
Fawrgy distamne
MSVGD, &
1.0
0.2
MSVGD, &
Prujected SVGD
SVMD
0.1
MSVGD, &
MSVGD, 4
0.0
0.0
200
400
200
400
Number partiele updates,
Number particle updates, T
Figure 2: Quality of 50-particle approximations 20-dimensional distributions the simplex after particle updates. (Left) Sparse Dirichlet posterior of Patterson & Teh (2013). (Right) Quadratic simplex target of Ahn & Chewi (2020). Details of the target distributions are in App. G.1.
Given any initial particle approximation qo 1. -i= dog. we discretize these dynamics to obtain the unconstrained SVNG sampler of Alg. 2 in the appendix. SVNG can be seen as an instance of MatSVGD (Wang et al., 2019) with a new adaptive time-dependent kernel G- (=, (108.80'G--1 " However, similar to Prop. 5 and unlike the heuristic kernels of Wang et al. (2019), SVNG reduces to natural gradient ascent for finding the mode of p(0) when n = 1. SVNG is well-suited Bayesian inference problems where the target is a posterior distribution p(0) & (v()) There, the metric tensor G(Ø) can be set to the Fisher information matrix Evloo) V log (y)0) V log (y 0) " of the data likelihood (y]0). Ample precedent from natu- ral gradient variational inference (Hoffman et al., 2013; 3; Khan & Nielsen, 2018) and Riemannian MCMC (Patterson & Teh, 2013) suggests that encoding problem geometry in this manner often leads to more rapid convergence.
EXPERIMENTS
We next conduct series of simulated and real-data experiments to assess (1) distributional approxi- mation on the simplex, (2) frequentist confidence interval construction for (constrained) post-selection inference, and (3) large-scale posterior inference with -un-Buclidean geometry. To compare with standard SVGD on constrained domains and to prevent its particles from exiting the domain θ we introduce a Euclidean projection onto θ following each SVGD update. For SVMD, we need to solve an eigenvalue problem, which costs O(n²) time. In practice the number of particles used for particle evolution algorithms is relatively small, even for SVGD, due to the ((n²) cost of updates. We have produced a practical SVMD implementation that is computationally competitive with MSVGD and SVGD for standard particle counts like n = 50 (used in all experiments).
5.1 APPROXIMATION QUALITY ON THE SIMPLEX
We first measure distributional approximation quality using two 20-dimensional targets: the sparse Dirchlet posterior of Patterson & Teh (2013) extended to 20 dimensions and the quadratic simplex target of Ahn & Chewi (2020). The Dirichlet target mimics the multimodal sparse conditionals that arise in latent Dirichlet allocation (Blei et al., 2003) but induces log concave density in space, while the quadratic is log-concave in 0 space. In Fig. 2, we compare the quality of MSVGD, SVMD, and projected SVGD with 50 particles and inverse multiquadric kernel & (Gorham Mackey, 2017) by computing the energy distance (Székely & Rizzo, 2013) to a surrogate ground truth sample of size 1000 (drawn i.i.d. or, in the quadratic case, from the No-U-Tur Sampler (Hoffman & Gelman, 2014)). We also compare to MSVGD with (0, 0¹) = k ((V:)(0), V (O)),, a choice which corresponds running SVGD in the dual space with kernel l by Thm. 4 and which ensures the convergence of MSVGD to p by the upcoming Thms. 6 to 8.
In the quadratic case, SVMD is favored over MSVGD as it is able to exploit the log-concavity of p(0). In contrast, for the multimodal sparse Dirichlet with p(0) unbounded near the boundary, MSVGD converges slightly more rapidly than SVMD by exploiting the log concave structure in 17 space. This parallels the observation of Hsich al. (2018) that LMC in the mirror space outperforms Riemannian LMC for sparse Dirichlet distributions. Projected SVGD fails to converge to the target in both cases and has particular difficulty in approximating the sparse Dirichlet target with unbounded density.

---

Published as a conference paper at ICLR 2022
0..0
Artual cnverig
coverien
0.00
0.90
0.65
Standezd MSVGD
Standann SVMD
0.90
SVMD
1500
Number of sumple 25:00
3000
0.50
0.14
1.00
Naminal
Nominal coverage: 0.9
(b) N 5000 sample points
Figure 3: Coverage of post-selection CIs across (a) 500 (b) 200 replications of simulation of Sepehri & Markovic (2017).
MSVGD with & and k2 perform very similarly, but we observe that yields better approximation quality upon convergence. Therefore, we employ k in the remaining MSVGD experiments.
5.2 CONFIDENCE INTERVALS FOR POST-SELECTION INFERENCE
We next apply our algorithms to the constrained sampling problems that arise in post-selection inference (Taylor Tibshirani, 2015; Lee et al., 2016). Specifically, we consider the task of forming valid confidence intervals (CIs) for regression parameters selected using the randomized Lasso (Tian al., 2016) with data X axp and y E TR and user-generated randomness we E R from Log-concave distribution with density 9. The randomized Lasso returns 3 E IR' with non-zero coefficients denoted by åg and their signs by 8g. is common practice to report least squares Cls for 3g by running linear regression 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 8g. To construct valid Cls, suffices to approximate the selective distribution with support {3g, u-E sg 2 BE > O, u_E E [-1, 1|m-1E1 and density
0(3n. u- x g(XT, y
x Xxxxx+tIi ) X x
In our experiments, we integrate out E analytically, following Tian et al. (2016), and reparameterize 3R ④ 18xz to obtain a log- concave density of |3²z supported on the nonnegative orthant with mirror function (0) Σ (0, log 0,). In Fig. 4a we show the example of 2D selective distribution using samples drawn by NUTS (Hoffman & Gelman, 2014). We also plot the results by projected SVGD, SVMD, and MSVGD in this example. Projected SVGD fails to approximate the target with many samples gathering the truncation boundary, while the samples by MSVGD and SVMD closely resemble the truth.
We then compare our methods with the standard norejection MCMC approach of the R package (Tibshirani et al., 2019) using the example simulation setting described in Sepehri & Markovic (2017) and a penalty factor 0.7. 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. 3a. the MSVGD and SVMD Cls consistently yield higher coverage than the standard 90% Cls. This increased coverage of particular value for smaller sample sizes, for which the standard CIs tend to undercover. For much larger sample size of N 5000 in Fig. 3b, the SVMD and standard CIs closely track one another across confidence levels, while MSVGD consistently yiclds longer CIs with high coverage. The higher coverage of MSVGD is only of value for larger confidence levels at which the other methods begin to undercover.
We next apply samplers to post-selection inference task the HIV-1 drug resistance dataset (Rhee et al., 2006), where we run randomized Lasso (Tian et al., 2016) to find statisti- cally significant mutations associated with drug resistance using susceptibility data virus isolates.
7

---

Published as a conference paper at ICLR 2022
We take the vitro measurement of log-fold change under the 3TC drug response and include mutations that had appeared at least 11 times in the dataset as regressors. In Fig. 4b we plot the Cls of selected mutations obtained with N 5000 sample points. We that the invalid unad- justed least squares Cls can lead to premature conclusions, e.g., declaring mutation 215Y significant when there is insufficient support after conditioning the selection event. In contrast, mutation 184V, which has known association with drug resistance, is declared significant by all methods even after post-selection adjustment. The MSVGD and SVMD CIs mostly track those of the standard selecti i ve veInferr Infernnce method, but their conclusions sometimes differ: e.g., 62Y is flagged as significant by MSVGD and SVMD but by selectiveInference
Trith
Projected SVGn
Unadjusted
2.4
Stardaad
SVMD
MSVGD
SVMD
MSVGD
0.4
0.2
0.0
-02
P4lL
P62V
Posht
Po7N
PER
ProR
Paak
PISIM
PISIC
PISAV
Pzinw
P2isY
(b)
Figure 4: (a) Sampling from 2D selective density; (b) Unadjusted and post-selection CIs for the mutations selected by the randomized Lasso as candidates for HIV-I drug resistance (see Sec. 5.2).
5.3 LARGE-SCALE POSTERIOR INFERENCE WITH NON- VON-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 & Wang (2016) with 581,012 datapoints and d = 54 feature dimensions. Here, the target is the posterior distribution over logistic regression parameters. We adopt the Fisher information metric tensor G. compare 20- 20-particle SVNG to SVGD and its prior awaae variants RSVGD (Liu & Zhu, 2018) and MatSVGD with average and mixture kernels (Wang et al., 2019), and for all methods use stochastic minibatches of methods
size 256 to scalably approximate each log likelihood query. In Fig. 5, all geometry-ave substantially improve the log predictive probability of SVGD. SVNG also strongly outperforms RSVGD and converges to its maximum test probability in half as many steps MatSVGD (Avg) and more rapidly than MatSVGD (Mixture).
6 CONVERGENCE GUARANTEES
We next turn our attention to the convergence properties of our proposed methods. For K₃ and ca as in Alg. let (q." 94. represent the distributions of the mirrored Stein updates (0,, ne) when Oo ~ の0 and 1 1/- + exD" q.. (0.) for 2 0. Our first result, proved in App. I.5, shows that E 5n converges in Wasserstein distance to distribution प as n oc, then, on each round > the output of Alg. q.¹ - dø; converges qf
the Alg. initialization 66,1
Theorem 6 (Convergence of mirrored updates as n + oc). Suppose Alg. I is initialized with G. Ei= 6ng satisfying W, (90. 90.H 0 for W1 the L Wasserstein distance. Define the q-induced kernel Kv (nn. 17") & (V (16"((m (n), V (TT)). If. for some
> 0.
(V(K ll, 11) V PM (1) + V K .ls m))loop < ci(1 (1 + [1ml12) [V ( V (. + v Kvo- Mlpp I ca(1 + || 1n'1)2),
then, Wa (97 (.. 0 and oi = of for each round t.
8

---

Published as a conference paper at ICLR 2022
-0.33
153
Maan SVCD (Mlatadd
-4.33
-0.30
23
503D
-137
-0.0
Kender pa pecadde spilao.
Stader pealee rplet,
Figure 5: Value of non-Euclidean geometry in large-scale Bayesian logistic regression.
The pre-conditions hold, for example, whenever V log is Lipschitz, & is strongly convex, and K, kl for bounded with bounded derivatives.
Remark
Given mirrored Stein operator (6), an arbitrary Stein set 9, and an arbitrary matrix-valued kernel K we define the mirrored Stein discrepancy and mirrored kernel Stein discrepancy
and MKSD (g.p) P) - MSD(1, BHx) (15) MSD(4 9) _ supgeg [.g[(Mp...))
The former is an example diffusion Stein discrepancy (Gorham et al., 2019) and the latter an example of a diffusion kernel Stein discrepancy (Barp et al., 2019). Since the MKSD optimiza- tion problem (15) matches that in Thm. 3, we have that MKSD (q.P) p) = 190, Our next result, proved in App. 1.6, shows that the infinit--particle mirrored Stein updates reduce the KL divergence to p whenever the step size is sufficiently small and drive MKSD to if, for example, 4 = 9(MKSD (nt", p))) for any a O. We also provide two conditions in App. H that generalize the Stein Log-Sobolev and Stein Poincaré inequalities in Duncan et al. (2019); Korba et al. (2020) and which imply exponential convergence rates of our algorithms in continuous time.
Theorem 7 (Infinite-partick mirrored Stein updates decrease KL and MKSD). Assume K1 (up,)lK((O,O))ll. < oo, R2 E supa | K, O)lp))lopppp)p)) < OC, V Vogg is L-Lipschitz and zb is a-strongly convex. c, < (2 |
,K, (0) + Vig. K. (0) (2-((0)-11lo))
KL(Gft
(g Ip) S (c LKi
2x2
MKSD (9:0 .p)?.
Our last result, proved in App. 1.7, shows that oieh p if MKSD Kú (4F pp) 0. Hence, by Thms. and 7, n-particle MSVGD converges weakly to if e decays at a suitable rate.
Theorem 8 (MKSDK. determines weak convergence). Assume pa is distantly dissipative (Eberle, 2016) with V log pH Lipschitz, vp is strongly convex with continuous Vii" and k(0, 0¹) R (()())) v+()))) for k(x,y) (c² 2 y1l2)² with 3 E (-1,0). Then. of p if MKSDK, (f ,P) 0.
Remark
The pre-conditions hold, for example, for any Dirichlet target with negative entropy ab.
DISCUSSION
This paper introduced the mirrored Stein operator along with three particle evolution algorithms for sampling with constrained domains and non-Euclidean geometries. The first algorithm MSVGD performs SVGD updates in a mirrored space before mapping to the target domain. The other two algorithms are different discretizations of the same continuous dynamics for exploiting non-Euclidean geometry. SVMD is multi-particle generalization of mirror descent for constrained domains, while SVNG is designed for unconstrained problems with informative metric tensors.
We highlight three limitations. First, like SVGD, our MSVGD require O(n²) time per update. Second, SVMD and SVNG are more costly than MSVGD due to the adaptive kernel construction. Low-rank kernel approximation may needed reduce their complexity. Third, we leave open the question of convergence when stochastic gradient estimates are employed, but we suspect the results of Gorham et al. (2020, Thm. 7) can be extended to our setting. In the future, we hope to deploy our mirrored Stein operators for other inferential tasks on constrained domains including sample quality measurement (Gorham & Mackey. 2015; Huggins & Mackey, 2018), goodness-of-fe fit testing (Chwialkowski et al., 2016; Liu et al., 2016; Jitkrittum et al., 2017), graphical model inference (Zhuo et al., 2018; Wang et al., 2018), parameter estimation (Barp et al., 2019), thinning (Riabiz al., 2020), and de novo sampling (Chen et al., 2018: Futami et al., 2019).

---

Published as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
See App. G for experimental details and
https: http:p////ttthhooooo/////////tta
for Python and R code replicating all experiments.
REFERENCES
Kwangjun Ahn and Sinho Chewi. Efficient constrained sampling via the mirror-Langevin algorithm. arXiv preprint arXiv:2010. 16212, 2020.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 10(2):251-276 1998.
Andrew Barbour. Stein's 's method and Poisson process convergence. Journal of Applied Probability, pp. 175-184, 1988.
Alessandro Barp, Francois-Xaves Briol, Andrew Duncan, Mark Girolami, and Lester Mackey. Minimum Stein discrepancy estimators. In Advances in Neural Information Processing Systems, pp. 12964-12976, 2019.
Amir Beck and Mare Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 316))1677-75,, 2003.
Alain Berlinet and Christine Thomas-Agnan, Reproducing kernel Hilbert spaces in probability and statistics. Springer Science & Business Media, 2011.
Rabi Bhattacharya and Edward Waymire. Stochastic processes with applications. SIAM, 2009.
David M Blei, Andrew Y Ng, and Michael Jordan. Latent Dirichlet allocation. Journal of Machine Learning Research, 3:923-1022, 2003.
Peng Chen, Keyi Wu, Joshua Chen, Tom PLeary-Roseberry, and Omar Ghattas. Projected Stein variational Newton: A fast and scalable Bayesian inference method in high dimensions. In Advances in Neural Information Processing Systems, volume 32, 2019.
Wilson Ye Chen, Lester Mackey, Jackson Gorham, François-Xavve Xavier Briol, and Chris Oates. Stein points. In International Conference on Machine Leurning, pp. 844-853, 2018.
Sinho Chewi, Thibaut Le Gouic, Chen Lu, Tyler Maunu, Philippe Rigollet, and Austin Stromme. Exponential ergodicity of mirror-Langevin diffusions. arXiv preprint arXiv:2005.0666 2020.
Kacper Chwialkowski, Heiko Strathmann, and Arthur Gretton. A kernel test of goodness of fit. In International Conference Machine Learning, pp. 2606-2615, 2016.
Arnak Dalalyan. Further and stronger analogy between sampling and optimization: Langevin Monte Carlo and gradient descent. In Conference on Learning Theory, pp. 777-669,,7 2017.
Gianluca Detommaso, Tiangang Cui, Youssef Marzouk, Alessio Spantini, and Robert Scheichl. A Stein variational Newton method. In Advances in Neural Information Processing Systems, volume 31, pp. 9187-9197, 2018.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal Machine Learning Research, 12(7), 2011.
Andrew Duncan, Nikolas Nüsken, and Lukasz Szpruch. On the geometry of Stein variational gradient descent. arXiv preprint 1912:0991200094 2019.
Alain Durmus, Eric Moulines, and Marcelo Pereyra. Efficient Bayesian computation by proximal Markov chain Monte Carlo: when Langevin meets Moreau. SIAM Journal on Imaging Sciences, 11(1):473-500 2018
10

---

Published as a conference paper at ICLR 2022
Andreas Eberle. Reflection couplings and contraction rates for diffusions. Probability Theory and Related Fields, 166(3):851-88,, 2016.
Stewart N Ethier. A class of degenerate diffusion processes occurring in population genetics. Communications on Pure and Applied Mathematics, 29(5))483 493, 1976.
Yihao Feng, Dilin Wang, and Qiang Liu. Learning to draw samples with amortized Stein variational gradient descent. Uncertainty in Artificial Intelligence, 2017.
JC Ferreira and VA Menegatto. Eigenvalues of integral operators defined by smooth positive definite kernels. Integral Equations and Operator Theory, 64(1):61-81, 2009.
Futoshi Futami, Zhenghang Cui, Issei Sato, and Masashi Sugiyama. Bayesian posterior approximation via greedy particle optimization. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 3606-3613, 2019.
Han I. Gan, Adrian Röllin, and Nathan Ross. Dirichlet approximation of equilibrium distributions in Cannings models with mutation. Advances in Applied Probability, 49(3):927-959, 2017.
Damien Garreau, Wittawat Jitkrittum, and Motonobu Kanagawa, Large sample analysis of the median heuristic. arXiv preprint arXiv: 1707,0760, 2017.
Jackson Gorham and Lester Mackey. Measuring sample quality with Stein's method. In Advances in Neural Information Processing Systems, pp. 226-234, 2015.
Jackson Gorham and Lester Mackey. Measuring sample quality with kemnlss. International Conference Machine Learning, pp. 1292-1301, 2017.
Jackson Gorham, Andrew B Duncan, Sebastian J Vollmer, and Lester Mackey. Measuring sample quality with diffusions. The Annals of Applied Probability, 29(5))2884-292) 2019.
Jackson Gorham, Anant Raj, and Lester Mackey. Stochastic stein discrepancies. arXiv preprint arXiv:2007.028577 2020.
Tuomas Haarnoja, Haoran Tang. Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International Conference on Machine Learning, pp. 1352-1361, 2017.
Geoffrey Hinton, Nitish Srivastava, and Kevin Swersky. Neural networks for machine learning lecture 6a: overview of mini-batch gradient descent. 2012.
Matthew D Hoffman and Andrew Gelman. The No-U-Tum sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15(1):1593-1633, 2014.
Matthew Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 14(5), 2013.
Ya-Ping Hsich, Ali Kavis, Paul Rolland, and Volkan Cevher. Mirrored Langevin dynamics. In Advances in Neural Information Processing Systems, pp. 2878-2887, 2018.
Jonathan Huggins and Lester Mackey. Random feature Stein discrepancies. In S. Bengio. H. Wal- lach, Larochelle, Grauman, N. Cesa-Biaa Bianchi, and R. Garnett (eds.), Advances in Neural 2018.
Information Processing Systems, pp.
Wittawat Jitkrittum, Wenkai Xu, Zoltán Szabó, K. Fukumizu, and A. Gretton. A Linear- Time Kernel Goodness-of- Test. In Advances in Neural Information Processing Systems, 2017.
Sham Kakade, Shai -awaar-S Ambuj Tewari, et al. On the duality of strong convexity and strong smoothness: Learning applications and matrix regularization. Unpublished Manuscript, (1), 2009.
Mohammad Emtiyaz Khan and Didrik Nielsen. Fast yet simple natural- -aauurllggrdientttt descent for variational inference in complex models. In 2018 International Symposium on Information Theory and Its Applications (ISITA), pp- 31-35. IEFF, 2018
11

---

Published as a conference paper at ICLR 2022
Taesup Kim, Jaesik Yoon, Ousmane Dia, Sungwoong Kim, Yoshua Bengio, and Sungiin Ahn. Bayesian model-agnostic meta-learning. Advances in Neural Information Processing Systems, 2018.
Anna Korba, Adil Salim, Michael Arbel, Giulia Luise, and Arthur Gretton. A non-asymptotic analysis for Stein variational gradient descent. Advances in Neural Information Processing Systems, 33, 2020.
Jason Lee, Dennis Sun, Yuekai Sun, and Jonathan E Taylor. Exact post-selection inference, with application to the lasso. Annals of Statistics, 44(3):907-99 2016.
Chang Liu and Jun Zhu. Riemannian Stein variational gradient descent for Bayesian inference. In Proceedings of the AAAI Conference on Artificial Intelligence, pp. 3627-3634, 2018.
Chang Liu, Jingwei Zhuo, Pengyu Cheng. Ruiyi Zhang, and Jun Zhu. Understanding and accelerating particle-based variational inference. In International Conference on Machine Learning, pp. 4082- 4092, 2019a.
Chang Liu, Jingwei Zhuo, and Jun Zhu. Understanding MCMC dynamics as flows on the Wasserstein space. In Proceedings of the 36th International Conference on Machine Learning, pp. 4093 4103, 2019b.
Qiang Liu. Stein variational gradient descent as gradient flow. In Advances in Neural Information Processing Systems, pp. 3115-3123, 2017.
Qiang Liu and Dilin Wang. Stein variational gradient descent: general purpose Bayesian inference algorithm. Advances in Neural Information Processing Systems, 19:2378-2338 2016.
Qiang Liu, Jason Lee, and Michael Jordan. A kemelized Stein discrepancy for goodness-of fit tests. In International Conference on Machine Learning, pp. 276-284, 2016.
Yi-An Ma, Tianqi Chen, and Emily Fox. A complete recipe for stochastic gradient MCMC. In Advances in Neural Information Processing Systems, pp. 2917-2925, 2015.
Yi-An Ma, Niladri Chatterji, Xiang Cheng, Nicolas Flammarion, Peter Bartlett, and Michael Jordan. Is there analog of Nesterov acceleration for MCMC? arXiv preprint arXiv:1902.00 1902.00996, 2019.
James Martens. New insights and perspectives on the natural gradient method. arXiv preprint 1rXiv:1412.1193, 2014.
Charles A Micchelli and Massimiliano Pontil. On learning vector-valued functions. Neural Compula- tion, 17(1):177-204, 2005.
Arkadij Semenovic Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. 1983.
Chris Oates, Mark Girolami, and Nicolas Chopin. Control functionals for Monte Carlo integration. Journal of the Royal Statistical Society: Series B (Methodological), 79(3):695-718, 2017.
Bernt Øksendal. Stochastic Differential Equations: An Introduction with Applications. Springer Science & Business Media, 2003.
Sam Patterson and Yee Whye Teh. Stochastic gradient Riemannian Langevin dynamics the probability simplex. Advances in Neural Information Processing Systems, pp. 5102-3110, 2013.
Garvesh Raskutti and Sayan Mukherjee. The information geometry of mirror descent. IEEE Transactions Information Theory, 61(3):1451-1457, 2015.
Soo-Yon Rhee, Jonathan Taylor, Gauhar Wadhera, Asa Ben-Hur, Douglas Brutlag, and Robert W Shafer. Genotypic predictors of human immunodeficiency virus type drug resistance. Proceedings of the National Academy of Sciences, 103(46): 2006.
Marina Riabiz, Wilson Chen, Jon Cockayne, Pawel Swietach, Steven Niederer, Lester Mackey, Chris Oates, al. Optimal thinning of MCMC output. arXiv preprint 2rXiv22005.03952 2020.
12

---

Published as a conference paper at ICLR 2022
Amir Sepehri and Jelena Markovic. Non-reversible, tuning-and rejection- Markov chain Monte Carlo via iterated random functions. arXiv preprint 17Xi:1771007177,2000
Jiaxin Shi, Shengyang Sun, and Jun Zhu. A spectral approach to gradient estimation for implicit distributions. In International Conference on Machine Learning, pp. 4644-4653, 2018.
Umut Simsekli, Roland Badeau, Taylan Cemgil, and Gaël Richard. Stochastic quasi-Newton Langevin Monte Carlo. In International Conference on Machine Learning, pp. 642 65166,1,220 2016.
Charles Stein. A bound for the enror in the normal approximation to the distribution of sum of dependent random variables. In Proceedings of the Sixth Berkeley Symposium on Mathemati- cal Statisties and Probability, Volume 2: Probability Theory. The Regents of the University of California, 1972.
Gábor Székely and Maria L Rizzo. Energy statistics; A class of statistics based on distances. Journal of Statistical Planning and Inference, 143(6):1249-1272, 2013.
Jonathan Taylor and Robert Tibshirani. Statistical learning and selective inference. Proceedings of the National Academy of Sciences, 1 112(25):722-77332 2015.
Xiaoying Tian, Nan Bi, and Jonathan Taylor. MAGIC: general, powerful and tractable method for selective inference. arXiv preprint 1rXiv:160.0066113 2016.
Ryan Tibshirani, Rob Tibshirani, Jonatha Taylor, Joshua Loftus Stephen Reid, and Jelena Markovic. selectivelnference: Tools for Post-Selection Inference, 2019. URL https: //CANN. R
R package version 1.2.5.
Dilin Wang. Zhe Zeng, and Qiang Liu. Stein variational message passing for continuous graphical models. In International Conference on Machine Learning, pp. 5219-5227, 20199,222,2111 2018.
Dilin Wang, Ziyang Tang, Chandrajit Bajaj, and Qiang Liu. Stein variational gradient descent with matrix-valued kernels. In Advances in Neural Information Processing Systems, pp. 7836-7846, 2019.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient Langevin dynamics. In International Conference on Machine Learning, pp. 681-688. 2011.
Edwin Wilson. Probable inference, the law of succession, and statistical inference. Journal of the American Statistical Association, 22(158):209-212 1927.
Tatiana Xifara, Chris Sherlock, Samuel Livingstone, Simon Bymne, and Mark Girolami. Langevin diffusions and the Metropolis-adjustel Langevin algorithm. Statistics Probability Letters, 91: 14-19,20 2014.
Jianyi Zhang, Yang Zhao, and Changyou Chen. Variance reduction in stochastic particle-optmmizatio sampling. In Proceedings of the 37th International Conference on Machine Learning, pp. 11307- 11316, 2020a.
Kelvin Shuangjian Zhang, Gabriel Peyré, Jalal Fadili, and Marcelo Pereyra. Wasserstein control of mirror Langevin Monte Carlo. arXiv preprint 2003:20002.04363. 2020b.
Michael Zhu, Chang Liu, and Jun Zhu. Variance reduction and quasi-Newton for particle-based variational inference. In Proceedings of the 37th International Conference on Machine Learning, pp. 11576-11587 2020.
Jingwei Zhuo, Chang Liu, Jiaxin Shi. Jun Zhu, Ning Chen, and Bo Zhang, Message passing Stein variational gradient descent. In International Conference on Machine Learning, pp. 6018-6027, 2018.
13

---

Published as a conference paper at ICLR 2022
Algorithm 2 Stein Variational Natural Gradient (SVNG)
Input: density p(0) on IR.d, kernel k, metric tensor G(0), particles (0h)==1> step sizes (et)²=1
for tor 0: do for E [n]. f/+11 f- of + (c((0)) (0..0)))''...00)) where
(0)
(Kc.(0 (see (14))
return (0++1)==1-
A MIRROR DESCENT, RIEMANNIAN GRADIENT FLOW, AND NATURAL GRADIENT
The equivalence between the mirror flow dn, = - (/)edd 0, = Vib" (nt)dt and the Riemannian gradient flow in (3) is direct result of the chain rule: d0, V q. D: drp
dipt
= oV Te) dt
= 2 (0,)
(16)
dt
dt
dnc v (0,) = Va, Th V n
S(V (T:) = - -v²"" (VVb
(17)
dt
Depending on discretizing (16) or (17), there are two natural gradient descent (NGD) updates that arise from the same gradient flow:
NGD (a): 0c+1 0 (V00)))V
NGD (b):
7++1 7ht V (m.
With finite step sizes ft. their updates need not be the same and can lead to different optimization paths. Since (0.)) V
ff (77)). NGD is equivalent the dual-space update by mirror descent. This relationship was pointed out in Raskutti & Mukherjee (2015) and has been used for developing natural gradient variational inference algorithms (Khan & Nielsen, 2018). We emphasize, however, our SVNG algorithm developed in Sec. 4.4 corresponds to the discretization the primal space as NGD (a). Therefore, does not require an explicit dual space, and allows replacing 2i with more general information metric tensors.
B DETAILS OF EXAMPLE 1
For the entropic mirror map (0) - - 0, log fj. we have (e) = diag( 00T Note
here 0 denotes d-dimensional vector and does not include 02+1 = 1 - ΣØ= O2. Since is (d + -isimmplex de composed of + faces with unit normal vector n0)) for the first faces are -cy for 1 I j S d, where ej denotes the j-th standard
the j-th face satisfies 0, = 0. The outward
basis of Rdd The outward unit normal vector for the (d + 1)-st face vector with 1/vd in all coordinates. Therefore, we have
/ 00 ((0))(0)TV<+00)) (00)))))(0)aa)))
(0)) (diag(O) 00 000 (0)
Jae
p()) 3 g(0) - 00 g(0) (0)
p()) 9 93 (0))0,d0
0041 p000 g0)
Vd
=0,
where in the second to last identity we used 0 = Ou41. Finally, we can verify the condition in
Prop. 1 as
lim +000
(00) (t))(() = sup p())g(0) (( (1) (n = 0.
89.
Illll $1 ae
14

---

Published as a conference paper at ICLR 2022
C DERIVATION OF THE MIRRORED STEIN OPERATOR
We first review the (overdamped) Langevin diffusson - a Markov process that underlies many recent advances in Stein's method along with its recent mirrored generalization. The Langevin diffusion with equilibrium density on I2d is a Markov process (8t)e20 C R satisfying the stochastic differential equation (SDE)
d0.- (og + VdBB
(18)
with (Br) a standard Brownian motion (Bhattacharya & Waymire, 2009, Sec. 4.5).
To identify Stein operators that satisfy (4) for broad classes of targets p. Gorham & Mackey (2015)
proposed to build upon the generator method of Barbour (1988): First, identify a Markov process (0)) that has p as the equilibrium density; they chose the Langevin diffusion of (18). Next, build Stein operator based on the (infinitesimal) generator A of the process (Øksendal, 2003, Def. 7.3.1):
(A)) = lim, / (Ef(0.) - E.((+)) for f : R.
as the generator satisfies Eø (()) 0 under relatively mild conditions. We use the following theorem to derive the generator of the processes described by SDEs like (18):
Theorem (Generator of Itô diffusion; Øksendal, 2003, Thm 7.3.3). Let (zz)620 be the Itô diffusion X' C I& satisfying dr, = b(x,)) cll + a (xt)d For any J E (X), the (infiitesimml)) generator A of (ax 120 is
(Af) = b(x) fx(x) + + [T((())
For the Langevin diffusion (18), substituting V logg(.) for b()) and v2T for a(-) Thm. 9, we obtain Aʃ (V logp) p) 'V V. V.. Replacing Vf with vector-valued function g gives the Langevin Stein operator (5).
To derive Stein operator that works well for constrained domains, we consider the Riemannian Langevin diffusion (Patterson & Teh, 2013; Xifara et al., 2014; Ma et al., 2015) that extends the Langevin diffusion to Mon-Euclidean geometries encoded in a positive definite metric tensor G(0);
d0,- (G(O,) log p(O:) + V G(O,) + dB²
We show in App. D that the choice G = 2 yields the recent mirror-Langevin diffusion (Zhang al., 2020b; Chewi al., 2020)
0. (n)) dnj v llog + v2V
(19)
According to Thm. 9, the generator of the mirror-Langevin diffusion described by (20) is
(0))
+ v V=
- -1
V (V-r(0) -
Now substituting g for (0), we obtain the associated mirrored Stein operator:
(Mp.eg) (() V (0) VVlo) ((g V (
D RIEMANNIAN LANGEVIN DIFFUSIONS AND MIRROR-LANGEVI DIFFUSIONS
Zhang et al. (2020b) pointed out (19) is a particular case of the Riemannian LD. However, they did not give an explicit derivation. The Riemannian LD (Patterson & Teh, 2013; Xifara et al., 2014; Ma al., 2015) with (-)((-) as the metric tensor is
d0, 2 (0,) log ((.)) + V V (1.)) dt + (0.)
(20)
To see the connection with mirror-Langeve diffusion, we would like to obtain the SDE that describes the evolution of (e)):) under the diffusionn. This requires the following theorem that provides the analog of the "chain rule" in SDEs.
A matrix divergence V G(D) is the vector obtained by computing the divergence of each of G(O).
15

---

Published as a conference paper at ICLR 2022
Theorem 10 (Itō formula; Øksendal, 2003, Thm 4.2.1). Let (z)) be Itō process in X R satisfying dæ b(x²)dt + ( Let f(x) C² _2 R Then yt f(x)) is again an ltô process, and its i-th dimension satisfies
1 Tr(V ((L))())))))))))))))))))))))))))
dye, ([(((,,)(())) +
Substituting V for f in Thm. 10, we have the SDE of T/2 - Vv)(0,) as
dn. V log (V, V2 + h(t,)) + (2V2
where h(t:); = Tr
((0))) Moreover, we have
(V (() TV (V.,, (0)) V') (0:))))))
2 (0)) Vo. (.. The
0r. [
{=1j=1
E=1Ø=1
Σ Dr.F
V (b:)) (v2
Vo, 4e +==
Therefore, the Ti diffusion is described by the SDE:
dnc = loggp(f)) (()) 8, (n))
E MODE MISMATCH UNDER TRANSFORMATIONS
0.12
(60)
0.08
0.04
0.00
0.0
0.2
0.4
0.6
0.8
-10
10
0.3
p100
0.2
2
0.1
0.0
0.0
0.2
0.4
0.6
0.8
1.0
-10
-5
0
10
11
Figure 6: The density functions of the same distribution in 0 (left) and (right) space under the transformation 7) V.(#). Each 0 follows a Beta distributions on [0, 1]. We choose the negative entropy vi(() = 0 log f + (1 8) log(1 [1 0). Then, the transformation is the logit function 1 log(0/(1-0) f)) and its reverse the sigmoid function 0 1/(11 Top: Beta(0.5, 0.5). Dashed lines mark the mode of the transformed density prr (1)) and the corresponding 0, which gives lowest value p(0); Bottom: 0 Beta(1.1, 10). Dashed lines mark the mode of the target density p(0) and the corresponding 7, which clearly does not match the mode of pa (n))
16

---

Published as a conference paper at ICLR 2022
F BACKGROUND ON REPRODUCING KERNEL HILBERT SPACES
Let H be Hilbert space of functions defined on X' and taking their values in R. We say l: is reproducing kernel (or kernel) of H. if V' E X. k(z,)) -) E H. and vff E H. (f. k(æ, ") fz).. H. is called reproducing kernel Hilbert space (RKHS) if has a kemel. Kernels are positive definite (p.d.) functions, which means that matrices with the form (x);,4;)); are positive semidefinite. For any p.d. function k, there a unique RKHS with & as the reproducing kernel, which can be constructed by the completion of (Ei=1 0,k( ) X, a, RR, i N).
metric space, is bounded continuous kernel with the RKHS H. and is positive measure X. (v( denote the space of all square-integrable functions w.r.t Then the kernel integral operator Ta (v))) - (v) defined by
Now we assume
Tag [ g(x)k(2,)))u
is compact and self-adjoint. Therefore, according to the spectral theorem, there exists an at most countable set of positive eigenvalues {Ai}jes C R with d1 2 X2 2 ... converging zero and orthonormal cigenfunctions {ujbee such that
Thu = Xylli.
and k has the representation k(z,z') Σjer Azu (z))((()))) (Mercer's theorem non-compact domains), where the convergence of the sum is absolute and uniform on compact subsets of X x X (Ferreira Menegatto, 2009).
G SUPPLEMENTARY EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS
In this section, we report supplementary details and additional results from the experiments of Sec. 5. In Secs. 5.1 and 5.2, we use the inverse multiquadric input kernel k(0,, 8*) (1 + |0 8)133/2)-14 due to its convergence control properties (Gorham Mackey, 2017). In the unconstrained experiments of Sec. 5.3, we use the Gaussian kernel k(0,, 0*) 0 (1) for consistency with past work. The bandwidth is determined by the median heuristic (Garreau al., 2017). We select from (0.98, 0.99) for all SVMD experiments. For unconstrained targets, report, for each method, results from the best fixed step size e E {0.01, 0.05, 0.1, 0.5, 1} selected on a separate validation set. For constrained targets, select step sizes adaptively to accommodate rapid density growth near the boundary; specifically, we use RMSProp (Hinton et al., 2012), an extension of the AdaGrad algorithm (Duchi al., 2011) 1) used in Liu & 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 PCIc 16GB.
G.1 APPROXIMATION QUALITY ON THE SIMPLEX
The sparse Dirichlet posterior of Patterson & Teh (2013) extended to 20 dimensions features sparse, symmetric Dir(a) prior with
0.1 for & {I...., 20} and sparse count data 90, n2 = n3 5, n2 (j > 3). modeled via multinomial likelihood. The quadratic target satisfies 1() 202 T AO + const, where we slightly modify the target density of Ahn & Chewi (2020)
to make it less flat by introducing scale parameter 0.01. A E RR 19x is a positive definite matrix generated by normalizing products of random matrices with i.i.d. elements drawn from 1ni
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 SVMD and MSVGD. Since projected SVGD applies updates in the 0 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. 7.
0.98. We search the base learning rates of RMSProp in {0.1, 0.01, 0.001 for
17

---

Published as a conference paper at ICLR 2022
Projectel SVCD
SVMID
MIVGID
MSVGI
1.5
1.5
Eaeery destam
LR-0.1
LR-6.1
LR-0.01
LRR=0.01
LR-0.00
LR-0,00
0.5
1R-0.01
0.5
LIt-0.1
0.5
0.5
1RR=0.00 (00)
LR=001
1R-0,0001
-000L
0.0
200
9.00
700
400
200
100
2000
100
Naraber of particle upčates, T
Number of particle upcates, T
Number of partiile updates, T
Number of paraile updatess T
Figure 7: Sampling from a Dirichlet target on a 20-simplex. We plot the energy distance to a ground truth sample of size 1000.
Projected SVCD
SVMD
MBVGD &
MSVCD by
LR-0.01
0.3
LRt-0.1
0.3
0.3
1R=00001
[RR-001
Eterey detaaree
LR-00.0011
LR-0000
0.2
0.1
[R-000
[RR-001
0.0
0.0
0.0
0.0
200
600
200
200
400
Namber particle updates, T
Numaber of particle upGater,
Number of paritle updatee, T
Number of parsicle updates,
Figure 8: Sampling from quadratic target 20-simplex. We plot the energy distance to ground truth sample of size 1000 drawn by NUTS (Hoffman & Gelman, 2014).
G.2 CONFIDENCE INTERVALS FOR POST-SELECTION INFERENCE
Given a dataset X E Ruxp, y E R² the randomized Lasso (Tian et al., 2016) solves the following problem:
argmingeR. I X x ||ll||, - w 3 1113112 w G.
where G user-specified log-concave distribution with density g. We choose to zero- mean independent Gaussian distributions while leaving its scale and the ridge parameter to be automatically determined by the randomi zedLasso function of the selectivee seleetiveInference package. We initialize the particles of our SVMD and MSVGD in the following way: First, we map the solution ßE the dual space by Vib. 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 Vit" and use them as the initial locations. Below discuss the remaining settings and additional results of the simulation and the HIV-1 drug resistance experiment separately.
Simulation In our simulation mostly follow the settings of Sepehri & Markovic (2017) except using different penalty level À recommended in the select iecctivve iiveivienn fencee R package. We ń 100 and p 40. The design matrix X is generated from an equi-correlated model, i.c.. each datapoint क E TRP is generated i.i.d. from (0, Σ) with Σ, = 1, Eij 0.3 (i * j) and then normalized to have almost unit length. The normalization done by first centering each dimension by subtracting the mean and dividing the standard deviation of that column of X, then additionally multiplying 1/n /2 is generated from a standard Gaussian which is independent of X, i.e., we assume the global null setting where the true value of zero. We set be the value returned by theoret i ca 1 1 ambda of the sel ect i ve Tnnferennce R package multiplied coefficient 0.7ù, where the 0.7 adjustment
introduced in the test examples of the package reduce the regularization effect so that we have reasonably large set of selected features when p 40. The base learning rates for SVMD and MSVGD are set to 0.01 and we run them for T 1000 particle updates. 5 is set to 0.98 for SVMD.
Our 2D example in Fig. 4a 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 ((0) = 8.0..0991122.388511 + 1.908166, + 2.39751)² 2 + (1.180998, 1.46104)2) const, 01,2 > 0.
18

---

Published as a conference paper at ICLR 2022
The error bars for actual coverage levels in Fig. 3a and Fig. 3b 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. 9a and Fig. 9b we additionally plot the average length of the confidence intervals w.r.t. different 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 se select i LvveInn Inffreence. 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 the dataset as regressors. This results in ñ 663 datapoints with = 91 features. We choose À to be the value returned by theoret ca 1. lambda of the select i ve Iffernnce R package multiplied by n. The base learning rates for SVMD and MSVGD are set to 0.01 and we run them for T 2000 particle updates. t is set to 0.99 for SVMD.
6
Wi2tb
2
Standard MSVCD
Standard MISVGD
SVMD
SVMD
1000
1500
2000
2500
30X0
0.80
0.85
0.500
0.95
1.00
Number of sumple
Numinal
Nominal coverage: 0.9
N 5000 sample points
Figure 9: Width of post-selection CIs across (a) 500 (b) 200 replications of simulation of Sepehri & Markovic (2017).
G.3 LARGE-SCALL POSTERIOR INFERENCE WITH NON-EUCLIDEAN GEOMETRY
The Bayesian logistic regression model we consider is I[ P(y i, w)p(w) where p(w) = (w I), p(y F; w) = Bernoulli(a(u) x;)). The bias parameter is absorbed into into w by adding an additional feature each zt. The gradient of the log density of the posterior distribution of is v. logr(w| (yi, 2; ) = Σi= a,(ys a (ww" T xx)) w. We choose the metric tensor 22 (ww be the Fisher information matrix (FIM) the likelihood:
F
B (w,,) V pog xE, w) V. log log w)
i=1
(w x1 aw xx))) xz)))aa..
Following Wang et al. (2019), for each iteration r (r 2 1), we estimate the sum with a stochastic minibatch B, size 256: FB. B. EieB, ow T xx)(1 - oww" =;)),;;] and approximate the FIM with a moving average across iterations:
F p. (1 (1- where Pr. = min 1/r, 0.95).
To ensure the positive definiteness the FIM, a damping term 0.01/ is added before taking the inverse. For RSVGD and SVNG, the gradient of the inverse of FIM is estimated with V F R3
- F where マi F = F V fs..
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, randomly
19

---

Published as a conference paper at ICLR 2022
keep 20% of the dataset 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 learning rates are very robust to the change in learning rates and
plotted in Fig. 10. We that the SVNG updates able to accommodate very large learning rates (up to 1) without a significant loss in performance. The results in Fig. 5 are reported with the learning rate that performs best on the validation set.
Validation log predictive probability
SVGD
Valieiation log, predite protodality
SVNG
-0.52
-0.52
-0.31
-0.51
-0.56
-0.56
-0.58
0.58
0,35
0.06
-0.60
0.1
-0.00
0.1
0.5
0.5
-0.62
1.0
1000
2300
3000
1000
Number of particle updates T
Number of particle updates T
RSVGD
Matrix SVGD (Avg)
Validation log predictive probabiity
Valldation log, prodiciive probudallity
-0.52
-0.51
-0.51
-0.56
-0.56
-0.58
-0.88
0.01
0.05
-0.60
0.31
-0.00
0.1
0.95
0.5
-0.62
-0.02
1000
2000
3000
1000
2000
3000
Number of partiele updates, T
Number of purticle updatess
Valisation log ieve probuiility
Matrix SVGD (Mixture)
-0.52
-0.30
-0.36
-0.38
0.01
0.05
-0.00
0.1
0.5
-0.62
1.0
1000
2000
3001
Number particle updates,
Figure 10: Logistic regression results validation sets with learning rates in (0.01, 0.05, 0.1, 0.5, 1]. Running RSVGD with learning rates 0.5 and 1 produces numerical errors. Therefore, we did not include them in the plot.
H EXPONENTIAL CONVERGENCE OF CONTINUOUS-TIME ALGORIHTMM
We derive time-inhiee tammiiihhmogeeeuus generalization of the Stein Log- Sobolev inequality of Duncan et al. (2019) and Korba et al. (2020) which ensures the exponential convergence of our continuous-ti time algorithms and tim--inhomogeneuus generalization of the Stein Poincaré inequality of Duncan et al. (2019) which guarantees exponential convergence near equilibrium (i.e., when ye is sufficiently close to p). As the results hold for a generic sequence of kernels (Kt)zzo. the implications apply to both MSVGD and SVMD.
20

---

Published as a conference paper at ICLR 2022
Definition 2 (Mirror Stein Log- Nog-Sobolev inequality). We define the Mirror Stein Log-Sobolev inequal- ity (af., Korba et al., 2020, Def. 2) as
KL(gellp)= S ..MSS 22
1
K. (gc, p) 2A E. ([ v log
T Pk... V2
log -
where MKSDK, K. is defined in Eq. (15); PK..." (g)) - (g)) is the kernel integral operator: (PK.. (k.))))()) Ey. [K,( ()((O))) for a general kernel Ke and vector-valued fanction yp on O and the stated equality holds whenever integration-by--oe
is applicable.
Proposition 11. Suppose (0,) )e20 follows the mirrored dynamics (8) with chosen to be Sg. as in (10). Then, the dissipation of (L (lp) is
d (L(gell)) (Ip) -MKSD K, (gt, p)
dl
Proof The proof directly follows from Thm. 3 since the optimization problem there matches the definition of MKSD in (15).
Therefore, when the Mirror Stein Log-Sobolev inequality holds, we have
d (L(gellp) - (2IIL(
dt
and the exponential convergence (L(2: < (L(gollp)e -2x follows by Gronwall's lemma (Gron- wall, 1919).
Definition 3 (Mirror Stein Poincaré inequality). We say that the distribution p satisfies the Mirror
Stein Poincaré inequality (cf., Duncan et al. 2019, Eq. (57)) with strongly corvex ip and constant À if
1 Varplø] S ス R PK.P E V-ab
for all @ E (p) (+) that is locally Lipschitz, where Pku. is the kernel integral operator under defined similarly as in Definition 2.
This inequality can also be viewed introduced Chewi et al. (2020, Def. 1) for proving exponential convergence mirror-Langevin diffusion. In manner analogous to Thm. of Chewi et al. (2020), the following proposition relates the Mirror Stein Poincaré inequality to chi chi-squared divergence.
kernelized generalization of the mirror Poincaré inequality
Proposition 12. Suppose (0, 10)) follows the mirrored dynamics (8) with chosen to be 9 K, as in (10). Then, the dissipation of chi-square divergence x² dt
(qull)) is
V 9t T
x2 = -2Esg
PK.P 2
whenever integration-by-par is applicable.
Proof We first note that by applying integration-be by-rrrtss,----sssss Sg, as in (10) can be equivalently written as
gg.K. = -KKL.a V2 V log St
Then using the Pokker-Plancl equation of a: under the dynamics, have
d
(aell)) dt
d
qt dp 2 qt dq /(
2Tg
gi.k.
V qt
di
p
p
dt
P
-2
(KKE 2
V
yt
V
4:
V at
PK.PP
PK.
V2
Zt
p
P
P
Note that the right hand side of the equation differs from the Mirror Stein Poincaré inequality only in the base measure of the expectation. Duncan et al. (2019) proposes to replace with to study the convergence near equilibrium (See their Sec. 6, where Eq. (46) replaced with Eq. (51)). If we do the same and combine this identity with the Mirror Stein Poincaré inequality, d.
obtain de (xx((qlllp) < 2 = -2xx2 which implics exponential convergence in KL
(T (p) x² (Ip) x² (gollp)e -22 22 by Gronwall's lemma (Gronwall, 1919).
21

---

Published as a conference paper at ICLR 2022
1 PROOFS
I.1 PROOF OF PROP.
Proof
Fix any g E Gv. Since 9 and Vg bounded and (2) p(0) and V (
are p-integrable, the expectation Eg~p ((M, (MM...)(()) exists. Because θ is convex, O is bounded and convex with Lipschitz boundary. Since p g E C¹ we have
[Bp((Mp.e)))0) Eplg(()) (2) log V (V (0) s(0))|l
+ (V) 90)
(p(0) (
lim
(p(0) (0))
(by dominated convergence)
2-000 e.
lim
(0(0)V²~(() (0) T n1 (0)d0
(by the divergence theorem)
1-000 000
lim
p())l((0) V2 (0) ~,, n, (0) d0 (by Cauchy-Schwarz)
ae.
||g||
lim
x (0) d8 (by assumption).
--000
50.
1.2 PROOF OF THM. 3: OPTIMAL MIRROR UPDATES IN RKHS
Proof Let e, denote the standard basis vector of Rdd with the i-th element being and others being zeros. Since E HK. we have
m(0) V (0) p(f) = (m, (( log.()))
V (V²
Vo, (m))
Σ (m, Ve, [K (K (2 -c)))
(m, Vo (K(, 0) 0) (0)
where we define the divergence of matrix vector whose elements are the divergences of each row of the matrix. Then, we write (9) as
-E,, (m)) V2 (0)
+ (V
ET = --E, [m,, K (log Vo (K(: 2)) Hxc -}])
-mm, Em. [K 8) () ()) log (0) Ve
-mm, E,, MMp.wKK K H)
Therefore, the optimal direction in the Hx norm ball BHg = {g |1 Hx $ 1} that minimizes (9) is of x 92.K = IM,.,M, (.,))).
I.3 PROOF OF THM. 4: MIRRORED SVGD UPDATES
Proof A p.d. kernel & composed with any map & is still a p.d. kernel. To prove this, let (21....,Hp)) ( ,(T))) p n. Then
Σ (om)))) = Σ BeBmklag, xm) 2 0,
E,m
22

---

Published as a conference paper at ICLR 2022
where Be se dg St {i ((p)) xz). Therefore, (y 1/) k(V (1). (r)
p.d. kernel. Plugging K kl into Lem. 13, for any 0' @ and r/ = Veb(6") we have
Ky..Ke Ke (0') Bngra [K (Vv) VoggPgp))) - Vn. K (Vw-(Vi((0') m))
Eng M (V (1)) Vi (ln))) + Σ
V n, kVV (1)) V (m))e)
5=1
En (1 (1', V log (m)) + V kn mp).
1.4
PROOF OF PROP. 5: SINGLE-PARTICLE SVMD MIRROR DESCENT
Proof
When = 1, A1 k(0,, 01). 1, and thus (0,, 0r) = k(0, ((0r,8))32(())
1.5 PROOF OF THM. 6: CONVERGENCE MIRRORED UPDATES nI
Proof The idea is to reinterpret our mirrored updates one step of matrix SVGD in ry space based on Lem. 13 and then follow the path of Gorham et al. (2020, Thm. 7). Assume that gi and 20 have integrable means. Let v7" 11 be an optimal Wasserstein-1 1 coupling of H. and q. Let qg denote the transform through one step of mirrored update: O. Val" (Tt), 7h+1 The eggg. (0t). Then, with Lem. 13, we have
I K. (n) + Pu. (1/)
|| + (199 K, (0) n/ Ager K, (f2)
S ||æ - r/ll2 + cell (0) go.= Kr (8))
< ||1 1'l12
+ ez (KV (( 1f V log Vn (vor,().n")
-KKvw-,(1) n log vr ( ,11" ,7"))]]
@. @. ( (og V,, Kvp-.e(1 17")
(KVve, x Kv , (1/ 0//,1010 ))]l2
Vogg
||7 - 1/'12 G ||||n" (
+ (0 + 1/ 11 Ba p Illn" [12]
1| n''l c Eng (1l-111))1~~/1 c (1 I|n ll2) Wi (i.
Since o (11") ~ 1 L x (1) ~ v. we conclude
Wi ( 9
< [[ 1 q. K (1)" 9 x (1) [||2]
$ c (1 Egg |ll. |||l" 1 |22] + cc2(1 + 111l)) (g [.r)
< (1 + cc01 (1 + Egg , |l|. (1z)) (102(1 Eng,, (1l:ll)))W((( (g.. H.
The final claim of T0e0 now follows by the continuous mapping theorem as V is continuous. 11
1.6 PROOF THM. 7: CHANITE-PARTICLL MIRRORED STEIN UPDATES DECREASE KL AND MKSD
Proof
Let To K. denote transform of the density function through one step of mirrored update: (nt) 71t1 = Tt + c.ggp .K. (0c). Then
KL(gff llp) - kL(g;
-KL(g qi (L (1p)
= -Egenggs [log (1((1)) - log log (m1 + C199p Ks, (0:) - log (det(T + 6, V Vo K, (0.)).,
23

---

Published as a conference paper at ICLR 2022
where have used the invariance of KL divergence under reparameterization:: KL(9)lp) (L (PH) Following Liu (2017), we bound the difference of the first two terms as PM (m))
log e.gar. ., (0:))
v. logPH((((s))) ds, where ne (s(s) - ny + 88. Ke (0p)
v (o)))( (frggg K. (0:)) ds
v (1)) Fop K. (0t) ! (VlogpH(1)) (og))(n(ss)) (erg
(() ... (0t)) ds - v logpu (1)) Fur K. (0c) + Ft / VVlogpn(n) V Vlogpn(((())lll> 11 x. (0e)lla ds
Le?
S - V (H)))( Fer K. (0t) +
ll ,k. (0:)1
2
and bound the log determinant term using Lem. 15:
log det(I + c Io (0d) & -t Tr(V If K. (0:)) + 2c VV 9. K, (0.)
The next thing to notice is that Bn-w,, Po (1:) D6 K. (0.) + Tr(V Fyp K. (0c)) is the square of the MKSD in (15), We can show this equivalence using the identity proved in Lem. 14:
Em- W K., (0.) T log (n):) + Tr(V m. K. (0:))
-Ei-oo L0 K. (0.) (0.) (log (logp(0.) log det (O.)
+ TT(V2)((0)) Vig. K. (0+))
-Egingg 10 (0.) V (0x) -Vog p(.))) + (V2 ' K. (0:)) (Lem. 14)
= -Biing [MMp pf K. (0:)
MKSDK, (9f p)?
Finally, we are going to bound 10 K. (02) and V. .K, (0.)) From the assumptions we have va is a-strongly convex and thus di is -troonggly smooth (Kakade al., 2009), therefore || 20 (1)())))2 S 2 By Lem. 16, we know
10 K. (0a)112 & 1197 K. 11 [K((,, 0t) MKSDK, ( ,)) 0z)pppp V Vo (1. V/ (m:) V .K. (0.)
< 2 & 11211 K. (6:)
d
1
S a 11 ifr K. 1ites
1 Ke(e,, 0:)l
1=1
d MKSDK, (p. .p)
Vi,4,,K (0c, le)p
where K0,,4 () denotes 0 K(0, (((0,0))ee l0r=0- Combining all the above, we have
KL 1pp) - (L(gg;91p))
Led
2c²
4 sup (0, 0)
Σ sup | K.
MKSD (g,", ,p)
Plugging in the definition K₁ and K2 finishes the proof.
24

---

Published as a conference paper at ICLR 2022
1.7
PROOF OF THM. 8: MKSDK, DETERMINES WEAK CONVERGENCE
Proof
According to Thm. 4,
Sy.K, Bgn [k(, Vet (ng)) V,,k(V"(n))
where qu (17) denotes the density of Vib(() under the distribution 0 ~ q. From the assumptions we have k(0, 8*) k(V)(O), Vg)(6"). With this specific choice of k. the squared MKSD is
MKSD KR (q. p)² Ke ll
Eng -qg
PH (T)PP (1)) V (P (n), vo'
V Vr (nw))
(21)
PH (T)
The final expression in (21) is the squared kernel Stein discrepancy (KSD) (Liu et al.. 2016;
Chwialkowski et al., 2016; Gorham & Mackey, 2017) between Qи and pи with the kernel k: KSD,, PH)2 Recall that it proved
Gorham & Mackey (2017, Theorem 8) that, for k(x, v) = (c + I| y13))² with B E (-1,0) 0) and distantly dissipative Ри with Lipschitz score functions, QH PHr if KSDx Pı)) 0. The advertised result (q = p if MKSD K2 (g.P) 0)
now follows by the continuous mapping theorem Vib" is continuous.
J LEMMAS
Lemma 13. Let Kv (n, n/) & (VV))(()) (n))) The mirrored updates I2. in (10) can be (n.
equivalently expressed as
Iq..K - ((vo-(vv() n) (og + V,, (Vv)(Ve()
Proof
We will use the identity proved in Lem. 14.
Iy..K = Bgg [M (M,
(E,, log Vo (K(,))())))))))))))) (logPH
= [ - + det V2 (0)) + V, (K(: (K(,))V²<(0)>> (by change-of-variable formula)
= -E,](((,0)V'<0)) (V
V 2 10) Vo, ((,0)
(by applying Lem. 14 to each row of ((: 0))
=E,(K(;,)V''00) V Log (V)()))
V K(, 0).
= E [. Val" log +
Vn, K(, (()())))
1=1
=Bn. v (V)()), 1) V lo (1) V, K - (V (1), 7))
where A., denotes the j-th column of a matrix A.
Lemma 14. For strictly convex function & E C2, Rd + R and any vector-valued g E C* Rd Dd the following relation holds:
V
-1 det (2
Proof
By the product rule of differentiation:
V (V22111))~'0(0)) (()) T + T (V2+(0)-1)
(22)
25

---

Published as a conference paper at ICLR 2022
This already gives us the first term on the right side. Next, we have (o) det (0));
V ()) 11s Tr (() Vog (2)(
(P)(0)-1)1
V (0)) - Am Vo, v+) me
2,10=
V2 ( " Jem Vo, V2 v0) ml
j.m==
(1y Ve ~000)/1V>>100 1em)
j.hm=l
d
Ven. (V2 Jm
-
(1)
Plugging the above relation into (22) proves the claimed result.
Lemma 15 (Liu, 2017, Lemma A.1). Lel A be square matrix, and 0 < E 24 Top" Then, log det(T cA)) - 2
where || ||e denotes the Frobenius norm of a matrix.
Lemma 16. Let K be a matrix-valued kernel and HK be the corresponding RKHS. Then, for any T HK (f is vector-valued), we have
llf(z))l2 S 1lffll (x, x) ???
|V S I1
EIVe K(x N'
where
||
llop denotes the operator norm of a matrix induced by the vector 2-norm.
Proof We first bound the [[((x)]l2
[f(()))2=
sup f(x) T sup (f,K(,z)))H& 111fl
sup
|K [1(,,))lll
|=|l==1
1=1l==1
Ilvll==1
|lfllH sup (y) K k(z,z)y)
x)y) 1/2 < 1lfll 1lvb=1 ||u||==1
sup (u K ((x,
sup
11v1l==1
= =Ilfll> sup | x)ylb 1/2 =1151 (K(x,))) 1/2. l1vl|==1
The second result follows similarly,
|| Σ f())l
Σ sup (V (v.,f(x)Ty) Σ
sup ((f, (z,, xll==1
llvil==1
< 1l
sup V R. K (,, x)y1 11 = 11111
lvll== بq K(x, x0) w
sup (g 2
livll, =1
$ 11 xx
sup
sup
(u V2
24
K
K(x, )z=z))
livll =1 ||ا|ا=|
im
-11/11
2 2 K (I, x Ma'= llop-
26


Reviews:
[{"review_id": "raz1UT750VO", "reviewer": "Reviewer_Zj9c", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper tackles the important problem of non-Euclidean or constrained space with particle methods. Its significance is therefore very high.\nThe paper is very complete and introduces all the different notions needed to understand the theory of the proposed methods.\nThe only comment I would have is that with the appendix being itself ~14 pages, this submission would better fit as a journal paper as it is.\n\nQuestions/Comments:\n- How does the algorithm compare with vanilla SVGD combined with a bijection to solve the constraint problem (for example x->exp(x) for positive real values)?\n- Could you explain why \"the modes of the mirrored density $p_H(η)$ need not match those of the target density $p(θ)$\"? \n- Could you comment on the cost of computing eigenfunctions and eigenvalues for SVMD?", "summary_of_the_paper": "The paper introduces new methods to run SVGD in constrained domains and non-euclidean geometries.\nThe authors develop the theory to combine mirror descent dynamics with the Stein's method via what they call mirrored Stein operators.\nThey also show experimental results on two concrete problems and prove convergence guarantees.", "summary_of_the_review": "This paper addresses an interesting problem and its development looks very exciting.\nThe work is very complete and addresses both theoretical and experimental aspects of the proposed methods.\nThe theory proposed looks sound with a strong background and the experiments look very convincing.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "fdPZc9lslVk", "reviewer": "Reviewer_dVdw", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1. The premise of the paper is clear and interesting : in traditional optimization, mirror descent allows you to optimize over constrained feasible sets, and this paper attempts a similar approach to sampling.\n\n1. In order to achieve the sampling analogue of mirror descent, the authors define an infinitesimal generator for a Markov process that has an equilibrium density $ p $. Since $ p $ has constraints (for example, it should be a sparse Dirichlet distribution), the authors construct the generators via the mirror maps associated with the constraints. I think the ideas used to construct these generators are quite non-trivial.\n\n1. The authors show that their algorithm mixes quicker in simulated experiments.\n\nWeaknesses:\n1. The writing can be improved. The authors try to give a lot of information about gradient flows that are often repetitive and not particularly friendly to readers. For example, the paragraph between Eqn (4) and (5) will not help an unfamiliar reader. Additionally, I did not understand the details of the experiments in section 5.2 \n\n1. Some of the notation is quite confusing -- the authors use $ \\theta $ and $\\theta_t$ interchangeably at times. This is not a problem in some cases, like Eqn (10), but Eqn (11) is very confusing to me. In theorem 4, $ \\eta $ is a function of $ \\theta$, and $ \\theta$ is a random variable distributed according to $ p $. But then $ q_{t, H} (\\eta ) $ is a density of $ \\eta_t$, and then Eqn (11) contains an average over $ \\eta_t$, and no term except $ k_\\psi$ has a $ \\eta$ term. I'm still not sure I understand Eqn 11 completely. This issue gets even worse in Eqns (12) and (13), where a lot of terms are functions of $ t $ although the variables do not include $ t $.\n\n1. The limitations of this work should be addressed a little bit -- from what I understand, you need to be able to compute kernel functions quickly, and require closed form solutions for $ \\nabla \\psi $ and $ \\nabla \\psi^* $. \n\nMinor comments:\n1. Above eqn (2), the equality should be $ \\theta_{d+1} = \\sum_{i=1}^d \\theta_i$.\n\n1. The authors keep saying they want to generalize to non-Euclidean geometries. If SVGD considers the geometry induced by the KL-divergence, then isn't this already non-Euclidean?", "summary_of_the_paper": "The paper introduces algorithms for sampling constrained distributions. Assuming the existence of a mirror map that maps from the primal space to the dual space, the authors introduce an algorithm MSVGD that can perform gradient descent in the dual space, and map back to the primal space. They also propose SVMD, which uses a different algorithm to perform gradient descent in the primal space itself. The authors experimentally validate their algorithms on simulated data and one non-simulated dataset (Fig 4b).\n\n", "summary_of_the_review": "I think the result is intuitive and requires advanced techniques to construct the infinitesimal generator used for constrained sampling.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "TxQytmFb2Cg", "reviewer": "Reviewer_5TzE", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- The paper is dense and offers a large number of both algorithmic and theoretical contributions. The use of mirror transformations to build specialized Stein's operators show a remarkable technical mastery.\n- The addressed problems are relevant. While several prior works have already addressed the Riemannian case, the convergence properties of this method seems to be superior. The application to constrained problems is novel and in my opinion particularly useful.\n- The writing is clear. Both prior work and contributions are clearly explained. \n- Good balance between theoretical and algorithmic contributions.\n\nWeaknesses: \n- The abundance of content makes the overall narrative difficult to follow. the paper introduces several related methods in order to solve several loosely related problems. I would have appreciated a more focused presentation.\n- The experiments are unfocused, with too little space devoted to each application. I would have preferred to see less experiments with more detailed discussion and analysis. ", "summary_of_the_paper": "the paper introduces a new family of deterministic particle samplers which apply mirror transformations to the Stein's variational gradient descent method. This approach can be used to generalize SVGD to both constrained and non-euclidean settings. The authors introduce two algorithms of this type that differ by the choice of the kernel function. One of these, SVMD is shown to reduce to regular mirror descent when a single particle is used. ", "summary_of_the_review": "The contribution is novel and solid with a wide range of potential applications. The paper is slightly too dense and somewhat unfocused. However, the main contributions are presented in a clear and understandable way. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "LHSiOs0tS6s", "reviewer": "Reviewer_BkpK", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "[Strengths]\nThe paper is overall well-written and makes reasonable attempts to provide sufficient background to understand the proposed algorithms. The literature review is thorough.\n\n[Weakness]\n- In the discussion in this work, constrained domains seem to be restricted to simplex. I wonder if MSVGD and SVMD are applicable to constrained domains in general, for example, a conjunction of non-overlapping simplexes and whether the convergence analysis still holds in this case.\n\n- At the end of Section 4.2, I find the statement that MSVGD does not necessarily match the modes of the target density lacks theoretical or empirical evidence. In the multimodal sparse Dirichlet experiments, MSVGD and SVMD have pretty similar performances. Since this statement is the motivation for SVMD, it would be more convincing to show why or when SVMD is able to capture the modes while MSVGD is not either theoretically or empirically. Also, a more detailed comparison of these two algorithms would be appreciated.\n\n\n[Minor Comments]\n- Some references are not properly presented. There are some cases where \\citep should have been used but \\citet is used; the opposite cases also happen. A careful pass on citations might be beneficial.\n- In Figure 2, the curves for MSVGD and SVMD are partially invisible when they are close to 0.0.", "summary_of_the_paper": "The aim of this work is to extend the existing particle evolution method Stein variational gradient descent or SVGD to constrained domains and non-Euclidean geometry. Three algorithms are proposed for this purpose based on mirrored Stein operators. The first is Stein variational mirror descent that runs SVGD in dual space s.t. the updated particles stay in constrained domains. The second one is Stein variational mirror descent defined with some adaptive kernels and it is also applicable to constrained domain problems. The third one is Stein variational natural gradient which is intended for unconstrained problems with informative metric tensors. Empirical and convergence analyses are further provided for all three proposed algorithms.", "summary_of_the_review": "The theoretical contribution of this work is solid. Still, quality and clarity need to be improved in several points as mentioned in main review.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes to extend mirror descent to sampling with stein operator when the density is defined on a constrained domain and non euclidean geometry. All reviewers agreed on the novelty and the merits of the paper. Accept
Decision: Accept (Spotlight)