Year: 2021

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.

text: 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.

heading: 1 INTRODUCTION
text: 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 & Thomas-Agnan, 2011) to sequentially minimize Kullback-Leibler (KL) divergence. SVGD has found great success in approximating unconstrained distributions for probabilistic learning (Feng et al., 2017; Haarnoja et al., 2017; Kim et 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 & Taylor, 2018), and fails to exploit informative non-Euclidean geometry (Amari, 1998).
In this work, we derive a family of particle evolution samplers suitable for target distributions with constrained domains and non-Euclidean geometries. Our development draws inspiration from mirror descent (MD) (Nemirovskij & Yudin, 1983), a first-order optimization method that generalizes gradient descent with non-Euclidean geometry. To sample from a distribution with constrained support, our method first maps particles to a dual space. There, we update particle locations using a 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. 1. In Sec. 3, we develop two algorithms – Mirrored SVGD (MSVGD) and Stein Variational Mirror Descent (SVMD) – with different updates in the dual space; when only a single particle is 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 on benchmark simplex-constrained problems from the literature, constrained sampling problems in post-selection inference (Taylor & Tibshirani, 2015; Lee et al., 2016; Tian & Taylor, 2018), 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.
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, 2011; Simsekli et al., 2016; Dalalyan, 2017; Durmus et al., 2018; Ma et al., 2019) and SVGD-like sampling methods (e.g., Liu et al., 2019a;b; Zhu et al., 2020; Zhang et al., 2020a). 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 diffusion) in a 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 a 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. Riemannian 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 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 Sec. 4.4, a choice that is not explored in Wang et al. (2019) and which recovers natural gradient descent when n = 1 unlike the heuristic kernel constructions of Wang et al. (2019). None of the aforementioned works provide convergence guarantees, and neither SVN nor matrix SVGD deals with constrained domains.

heading: 2 BACKGROUND: MIRROR DESCENT AND NON-EUCLIDEAN GEOMETRY
text: Standard gradient descent can be viewed as optimizing a local quadratic approximation to the target function f : θt+1 = argminθ∈Θ∇f(θt)>θ + 12 t ‖θ − θt‖ 2 2. When Θ ⊆ Rd is constrained, it can be advantageous to replace ‖ · ‖2 with a function Ψ that reflects the geometry of a problem (Nemirovskij & Yudin, 1983; Beck & Teboulle, 2003):
θt+1 = argminθ∈Θ∇f(θt)>θ + 1 tΨ(θ, θt). (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 smooth1 function ψ : Θ → R ∪ {∞}: Ψ(θ, θ′) = ψ(θ) − ψ(θ′) −∇ψ(θ′)>(θ − θ′). When Θ is a (d + 1)- simplex {θ : ∑di=1 θi ≤ 1 and θi ≥ 0 for i ∈ [d]}, a common choice of ψ is the negative entropy ψ(θ) = ∑d+1 i=1 θi log θi, for θd+1 , 1− ∑d i=1 θd. The solution of (1) is given by θt+1 = ∇ψ∗(∇ψ(θt)− t∇f(θt)), (2) where ψ∗(η) , supθ∈Θ η
>θ − ψ(θ) is the convex conjugate of ψ and ∇ψ is a bijection from Θ to dom(ψ∗) with inverse map (∇ψ)−1 = ∇ψ∗. We can view the update in (2) as first mapping θt to ηt by∇ψ, applying the update ηt+1 = ηt − t∇f(θt), and mapping back through θt+1 = ∇ψ∗(ηt+1).
1ψ is continuously differentiable on the interior of Θ with ‖∇ψ(θt)‖ → ∞ whenever θt → θ ∈ ∂Θ.
Mirror descent can also be viewed as a discretization of the continuous-time dynamics dηt = −∇f(θt)dt, θt = ∇ψ∗(ηt), which is equivalent to the Riemannian gradient flow (see App. A): dθt = −∇2ψ(θt)−1∇f(θt)dt, or equivalently, dηt = −∇2ψ∗(ηt)−1∇ηtf(∇ψ∗(ηt))dt, (3)
where ∇2ψ(θ) and ∇2ψ∗(η) are Riemannian metric tensors. In information geometry, the discretization 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.

heading: 3 STEIN’S IDENTITY AND MIRRORED STEIN OPERATORS
text: Stein’s identity (Stein, 1972) is a tool for characterizing a target distribution P using a so-called Stein operator. We assume P has a differentiable density p with a closed convex support Θ ⊆ Rd. A Stein operator Sp takes as input functions g from a Stein set G and outputs mean-zero functions under p:
Eθ∼p[(Spg)(θ)] = 0, for all g ∈ G. (4) Gorham & Mackey (2015) proposed the Langevin Stein operator given by
(Spg)(θ) = g(θ)>∇ log p(θ) +∇ · g(θ), (5) where g is a vector-valued function and ∇ · g is its divergence. For an unconstrained domain with Ep[‖∇ log p(θ)‖2] <∞, Stein’s identity (4) holds for this operator whenever g ∈ C1 is bounded and Lipschitz by (Gorham et al., 2019, proof of Prop. 3). However, on constrained domains Θ, Stein’s identity fails to hold for many reasonable inputs g if p is non-vanishing or explosive at the boundary.
Motivated by this deficiency and by a desire to exploit non-Euclidean geometry, we propose an alternative mirrored Stein operator,
(Mp,ψg)(θ) = g(θ)>∇2ψ(θ)−1∇ log p(θ) +∇ · (∇2ψ(θ)−1g(θ)), (6) where ψ is a strongly convex, essentially smooth function as in Sec. 2 with (∇2ψ)−1 differentiable and Lipschitz on Θ. 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,ψ generates mean-zero functions under p whenever ∇2ψ−1 suitably cancels the growth of p at the boundary. Proposition 1. Suppose that ∇2ψ(θ)−1∇ log p(θ) and ∇ · ∇2ψ(θ)−1 are p-integrable. If limr→∞ ∫ ∂Θr
p(θ)‖∇2ψ(θ)−1nr(θ)‖2dθ = 0 for Θr , {θ ∈ Θ : ‖θ‖∞ ≤ r} and nr(θ) the outward unit normal vector2 to ∂Θr at θ, then Ep[(Mp,ψg)(θ)] = 0 if g ∈ C1 is bounded Lipschitz. Example 1 (Dirichlet p, Negative entropy ψ). When θ1:d+1 ∼ Dir(α) for α ∈ Rd+1+ , even setting g(θ) = 1 in (5) need not cause the identity to hold when some αj ≤ 1. However, when ψ(θ) = ∑d+1 j=1 θj log θj , we show in App. B that the conditions of Prop. 1 are met for any α. Remarkably, the mirror-Langevin diffusion for our choice of ψ is the Wright-Fisher diffusion (Ethier, 1976) which Gan et al. (2017) recently used to bound distances to Dirichlet distributions.

heading: 4 SAMPLING WITH MIRRORED STEIN OPERATORS
text: 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 θt+1 = θt + tgt(θt) for a chosen mapping gt : Rd → Rd. Specifically, SVGD chooses the mapping g∗t that leads to the largest decrease in KL divergence to p in the limit as t → 0. The following theorem summarizes their main findings. Theorem 2 (Liu & Wang, 2016, Thm. 3.1). Suppose (θt)t≥0 satisfies dθt = gt(θt)dt for bounded Lipschitz gt ∈ C1 : Rd → Rd and that θt has density qt with Eqt [‖∇ log qt‖2] <∞. If KL(qt‖p) , Eqt [log(qt/p)] exists then, for the Langevin Stein operator Sp (5),
d dtKL(qt‖p) = −Eθt∼qt [(Spgt)(θt)]. (7)
2For a closed convex set whose boundary ∂Θ can be locally represented as F (θ1, . . . , θd) = 0, its unit normal vector is defined as n(θ) = ± ∇F (θ1,...,θd)‖∇F (θ1,...,θd)‖2 .
Algorithm 1 Mirrored Stein Variational Gradient Descent & Stein Variational Mirror Descent
Input: density p on Θ, kernel k, mirror function ψ, particles (θi0)ni=1 ⊂ Θ, step sizes ( t)Tt=1 Init: ηi0 ← ∇ψ(θi0) for i ∈ [n] for t = 0 : T do
if SVMD then Kt ← Kψ,t (13) else Kt ← kI (MSVGD) for i ∈ [n], ηit+1 ← ηit + t 1n ∑n j=1Mp,ψKt(θit, θ j t ) (forMp,ψKt(·, θ) defined in Thm. 3)
for i ∈ [n], θit+1 ← ∇ψ∗(ηit+1) return {θiT+1}ni=1.
To improve its current particle approximation, SVGD finds the choice of gt that most quickly decreases KL(qt‖p) at time t, i.e., it minimizes ddtKL(qt‖p) over a set of candidate directions gt. SVGD finds gt in a reproducing kernel Hilbert space (RKHS, Berlinet & Thomas-Agnan, 2011) norm ball BHd = {g : ‖g‖Hd ≤ 1}, where Hd is the product RKHS containing vector-valued functions with each component in the RKHSH of k. Then the optimal g∗t ∈ BHd that minimizes (7) is
g∗t ∝ g∗qt,k , Eθt∼qt [k(θt, ·)∇ log p(θt) +∇θtk(θt, ·)] = Eθt∼qt [SpKk(·, θt)], where we let Kk(θ, θ′) = k(θ, θ′)I , and SpKk(·, θ) denotes applying Sp to each row of Kk(·, θ). 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-Euclidean geometries.

heading: 4.1 MIRRORED DYNAMICS
text: SVGD encounters two difficulties when faced with a constrained support. First, the SVGD updates can push the random variable θt outside of its support Θ, rendering all future updates undefined. Second, Stein’s identity (4) often fails to hold for candidate directions in BHd (cf. Ex. 1). When this occurs, SVGD need not converge to p as p is not a stationary point of its dynamics (i.e., ddtKL(qt‖p)|qt=p 6= 0 when qt = p). Inspired by mirror descent (Nemirovskij & Yudin, 1983), we consider the following mirrored dynamics
θt = ∇ψ∗(ηt) for dηt = gt(θt)dt, or, equivalently, dθt = ∇2ψ(θt)−1gt(θt)dt, (8) where gt : Θ → Rd now represents the update direction in η space. The inverse mirror map ∇ψ∗ automatically guarantees that θt belongs to the constrained domain Θ. Since ψ is strongly convex and ∇2ψ−1 is bounded Lipschitz, from Thm. 2 it follows for any bounded Lipschitz gt that
d dtKL(qt‖p) = −Eθt∼qt [(Mp,ψgt)(θt)], (9)
whereMp,ψ is the mirrored Stein operator (6). In the following sections, we propose three new deterministic sampling algorithms by seeking the optimal direction gt that minimizes (9) over different function classes. Thm. 3 (proved in App. I.2) forms the basis of our analysis. Theorem 3 (Optimal mirror updates in RKHS). Suppose (θt)t≥0 follows the mirrored dynamics (8). LetHK denote the RKHS of a matrix-valued kernel K : Θ×Θ→ Sd×d (Micchelli & Pontil, 2005). Then, the optimal direction of gt that minimizes (9) in the norm ball BHK , {g : ‖g‖HK ≤ 1} is
g∗t ∝ g∗qt,K , Eθt∼qt [Mp,ψK(·, θt)], (10) whereMp,ψK(·, θ) appliesMp,ψ (6) to each row of the matrix-valued function Kθ = K(·, θ).

heading: 4.2 MIRRORED STEIN VARIATIONAL GRADIENT DESCENT
text: Following the pattern of SVGD, one can choose the K of Thm. 3 to be Kk(θ, θ′) = k(θ, θ′)I , where k is any scalar-valued kernel. In this case, the resulting update g∗qt,Kk(·) = Eθt∼qt [Mp,ψKk(·, θt)] is equivalent to running SVGD in the dual η space before mapping back to Θ. Theorem 4 (Mirrored SVGD updates). In the setting of Thm. 3, let kψ(η, η′) = k(∇ψ∗(η),∇ψ∗(η′)), pH(η) = p(∇ψ∗(η)) · |det∇2ψ∗(η)| denote the density of η = ∇ψ(θ) when θ ∼ p, and qt,H denote the distribution of ηt under the mirrored dynamics (8). If Kk = kI , g∗qt,Kk(θ ′) = Eηt∼qt,H [kψ(ηt, η′)∇ log pH(ηt) +∇ηtkψ(ηt, η′)] ∀θ′ ∈ Θ, η′ = ∇ψ(θ′). (11)
The proof is in App. I.3. By discretizing the dynamics dηt = g∗qt,Kk(θt)dt and initializing with any particle approximation q0 = 1n ∑n i=1 δθi0 , we obtain Mirrored SVGD (MSVGD), our first algorithm for sampling in constrained domains. The details are summarized in Alg. 1.
When only a single particle is used (n = 1) and the differentiable input kernel satisfies∇k(θ, θ) = 0, the MSVGD update (11) reduces to gradient descent on − log pH(η). Note however that the modes of the mirrored density pH(η) need not match those of the target density p(θ) (see the examples in App. E). Since we are primarily interested in the θ-space density, it is natural to ask whether there exists a mirrored dynamics that reduces to finding the mode of p(θ) 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.

heading: 4.3 STEIN VARIATIONAL MIRROR DESCENT
text: Our second sampling algorithm for constrained problems is called Stein Variational Mirror Descent (SVMD). We start by introducing a new matrix-valued kernel that incorporates the metric ∇2ψ and evolves with the distribution qt. Definition 1 (Kernels for SVMD). Given a continuous scalar-valued kernel k, consider the Mercer representation3 k(θ, θ′) = ∑ i≥1 λt,iut,i(θ)ut,i(θ ′) w.r.t. qt, where ut,i is an eigenfunction satisfying
Eθt∼qt [k(θ, θt)ut,i(θt)] = λt,iut,i(θ). (12)
For k1/2t (θ, θ ′) , ∑ i≥1 λ 1/2 t,i ut,i(θ)ut,i(θ ′), we define the adaptive SVMD kernel at time t,
Kψ,t(θ, θ ′) , Eθt∼qt [k 1/2 t (θ, θt)∇2ψ(θt)k1/2t (θt, θ′)]. (13)
By Thm. 3, the optimal update direction for the SVMD kernel ball is g∗qt,Kψ,t =Eqt [Mp,ψKψ,t(·, θt)]. We obtain the SVMD algorithm (summarized in Alg. 1) by discretizing dηt = g∗qt,Kψ,t(θt)dt and initializing with q0 = 1n ∑n i=1 δθi0 . Because of the discrete representation of qt, Kψ,t takes the form
Kψ,t(θ, θ ′) = ∑n i=1 ∑n j=1 λ 1/2 t,i λ 1/2 t,j ut,i(θ)ut,j(θ ′)Γt,ij ,
Γt,ij = 1 n ∑n `=1 ut,i(θ ` t)ut,j(θ ` t)∇2ψ(θ`t).
Here both λt,j and ut,j(θit) can be computed by solving a matrix eigenvalue problem involving the particle set {θit}ni=1: Btvt,j = nλt,jvt,j , where Bt = (k(θit, θjt ))ni,j=1 ∈ Rn×n is the Gram matrix of pairwise kernel evaluations at particle locations, and the i-th element of vt,j is ut,j(θit). To compute ∇θKψ,t(θ, θ′), we differentiate both sides of (12) to find that∇ut,j(θ) = 1λt,j ∑n i=1 vt,j,i∇θk(θ, θit). 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 (τ ≥∑Jj=1 λt,j/∑nj=1 λt,j) to ensure numerical stability. Notably, SVMD differs from MSVGD only in its choice of kernel, but, whenever ∇k(θ, θ) = 0, this change is sufficient to exactly recover mirror descent when n = 1. Proposition 5 (Single-particle SVMD is mirror descent). If n = 1, then one step of SVMD becomes
ηt+1 = ηt + t(k(θt, θt)∇ log p(θt) +∇k(θt, θt)), θt+1 = ∇ψ∗(ηt+1).

heading: 4.4 STEIN VARIATIONAL NATURAL GRADIENT
text: 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 a new sampler – Stein Variational Natural Gradient (SVNG) – that more efficiently approximates unconstrained targets. The idea is to replace the Hessian∇2ψ(·) in the SVMD dynamics dθt = ∇2ψ(θt)−1g∗qt,Kψ,t(θt) with a general metric tensor G(·). The result is the Riemannian gradient flow
dθt = G(θt) −1g∗qt,KG,t(θt)dt with KG,t(θ, θ ′) , Eθt∼qt [k1/2(θ, θt)G(θt)k1/2(θt, θ′)]. (14)
3See App. F for background on Mercer representations in non-compact domains.
Given any initial particle approximation q0 = 1n ∑n i=1 δθi0 , 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−1(θ)KG,t(θ, θ
′)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(θ) 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|θ)>] of the data likelihood π(y|θ). Ample precedent from natural gradient variational inference (Hoffman et al., 2013; Khan & Nielsen, 2018) and Riemannian MCMC (Patterson & Teh, 2013) suggests that encoding problem geometry in this manner often leads to more rapid convergence.

heading: 5 EXPERIMENTS
text: We next conduct a series of simulated and real-data experiments to assess (1) distributional approximation on the simplex, (2) frequentist confidence interval construction for (constrained) post-selection inference, and (3) large-scale posterior inference with non-Euclidean 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.

heading: 5.1 APPROXIMATION QUALITY ON THE SIMPLEX
text: We first measure distributional approximation quality using two 20-dimensional simplex-constrained 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 a log concave density in η space, while the quadratic is log-concave in θ space. In Fig. 2, we compare the quality of MSVGD, SVMD, and projected SVGD with 50 particles and inverse multiquadric kernel k (Gorham & Mackey, 2017) by computing the energy distance (Szé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-Turn Sampler (Hoffman & Gelman, 2014)). We also compare to MSVGD with k2(θ, θ′) = k(∇ψ(θ),∇ψ(θ′)), 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 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(θ). 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. This parallels the observation of Hsieh et 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. MSVGD with k and k2 perform very similarly, but we observe that k yields better approximation quality upon convergence. Therefore, we employ k in the remaining MSVGD experiments.

heading: 5.2 CONFIDENCE INTERVALS FOR POST-SELECTION INFERENCE
text: 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 & Taylor, 2018) with data X ∈ Rñ×p and y ∈ Rñ and user-generated randomness w ∈ Rp from a log-concave distribution with density g. The randomized Lasso returns β̂ ∈ Rp with non-zero coefficients denoted by β̂E and their signs by sE . 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 sE . To construct valid CIs, it suffices to approximate the selective distribution with support {β̂E , u−E : sE β̂E > 0, u−E ∈ [−1, 1]p−|E|} and density
ĝ(β̂E , u−E) ∝ g ( X>y − (X>EXE+ I|E| X>−EXE ) β̂E + λ ( sE u−E )) .
In our experiments, we integrate out u−E analytically, following Tian & Taylor (2018), and reparameterize β̂E as sE |β̂E | to obtain a log-concave density of |β̂E | supported on the nonnegative orthant with mirror function ψ(θ) = ∑d j=1(θj log θj − θj). In Fig. 4a we show the example of a 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 at 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 selectiveInference 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 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. 3b, the SVMD and standard CIs closely track one another across confidence levels, while MSVGD consistently yields 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 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 & Taylor, 2018) to find statistically significant mutations associated with drug resistance using susceptibility data on virus isolates. We take the vitro measurement of log-fold change under the 3TC drug as response and include mutations that had appeared at least 11 times in the dataset as regressors. In Fig. 4b 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 significant
when there is insufficient support after conditioning on 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 selectiveInference method, but their conclusions sometimes differ: e.g., 62Y is flagged as significant by MSVGD and SVMD but not by selectiveInference.

heading: 5.3 LARGE-SCALE POSTERIOR INFERENCE WITH NON-EUCLIDEAN GEOMETRY
text: 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 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 & 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. 5, all geometry-aware methods substantially improve the log predictive probability of SVGD.4 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).

heading: 6 CONVERGENCE GUARANTEES
text: We next turn our attention to the convergence properties of our proposed methods. For Kt and 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 + tg∗qt,Kt(θt) for t ≥ 0. Our first result, proved in App. I.5, shows that if the Alg. 1 initialization qn0,H = 1 n ∑n i=1 δηi0 converges in Wasserstein distance to a distribution q ∞ 0,H
as n→∞, then, on each round t > 0, the output of Alg. 1, qnt = 1n ∑n i=1 δθit , converges to q ∞ t . Theorem 6 (Convergence of mirrored updates as n → ∞). Suppose Alg. 1 is initialized with qn0,H = 1 n ∑n i=1 δηi0 satisfying W1(q n 0,H , q ∞ 0,H)→ 0 for W1 the L1 Wasserstein distance. Define the η-induced kernel K∇ψ∗,t(η, η′) , Kt(∇ψ∗(η),∇ψ∗(η′)). If, for some c1, c2 > 0, ‖∇(K∇ψ∗,t(·, η)∇ log pH(η) +∇ ·K∇ψ∗,t(·, η))‖op ≤ c1(1 + ‖η‖2), ‖∇(K∇ψ∗,t(η′, ·)∇ log pH(·) +∇ ·K∇ψ∗,t(η′, ·))‖op ≤ c2(1 + ‖η′‖2),
then, W1(qnt,H , q ∞ t,H)→ 0 and qnt ⇒ q∞t for each round t.
4Notably, on the same dataset, SVGD was shown to outperform preconditioned stochastic gradient Langevin dynamics (Li et al., 2016), a leading MCMC method imbued with geometric information (Wang et al., 2019).
Remark The pre-conditions hold, for example, whenever ∇ log pH is Lipschitz, ψ is strongly convex, and Kt = kI for k bounded with bounded derivatives.
Given a mirrored Stein operator (6), an arbitrary Stein set G, and an arbitrary matrix-valued kernel K we define the mirrored Stein discrepancy and mirrored kernel Stein discrepancy
MSD(q, p,G) , supg∈G Eq[(Mp,ψg)(θ)] and MKSDK(q, p) , MSD(q, p,BHK ). (15) The former is an example of a 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 optimization problem (15) matches that in Thm. 3, we have that MKSDK(q, p) = ‖g∗q,K‖HK . Our next result, proved in App. I.6, shows that the infinite-particle mirrored Stein updates reduce the KL divergence to p whenever the step size is sufficiently small and drive MKSD to 0 if, for example, t = Ω(MKSDKt(q ∞ t , p)
α) for any α > 0. We also provide two conditions in App. H that generalize the Stein Log-Sobolev and Stein Poincaré 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-particle mirrored Stein updates decrease KL and MKSD). Assume κ1 , supθ‖Kt(θ, θ)‖op < ∞, and κ2 , ∑d i=1 supθ ‖∇2i,d+iKt(θ, θ)‖op < ∞, ∇ log pH is L-Lipschitz, and ψ is α-strongly convex. If t < 1/(‖∇ηtg∗q∞t ,Kt(θt) +∇ηtg ∗ q∞t ,Kt (θt) >‖op), then
KL(q∞t+1‖p)− KL(q∞t ‖p) ≤ − ( t − ( Lκ1 2 + 2κ2 α2 ) 2t ) MKSDKt(q ∞ t , p) 2.
Our last result, proved in App. I.7, shows that q∞t ⇒ p if MKSDKk(q∞t , p)→ 0. Hence, by Thms. 6 and 7, n-particle MSVGD converges weakly to p if t decays at a suitable rate. Theorem 8 (MKSDKk determines weak convergence). Assume pH is distantly dissipative (Eberle, 2016) with ∇ log pH Lipschitz, ψ is strongly convex with continuous ∇ψ∗, and k(θ, θ′) = κ(∇ψ(θ),∇ψ(θ′)) for κ(x, y) = (c2 + ‖x− y‖22)β with β ∈ (−1, 0). Then, q∞t ⇒ p if MKSDKk(q ∞ t , p)→ 0.
Remark The pre-conditions hold, for example, for any Dirichlet target with negative entropy ψ.

heading: 7 DISCUSSION
text: This paper introduced the mirrored Stein operator along with three new 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 a 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(n2) time per update. Second, SVMD and SVNG are more costly than MSVGD due to the adaptive kernel construction. Low-rank kernel approximation may be needed to 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-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 et al., 2020), and de novo sampling (Chen et al., 2018; Futami et al., 2019).

heading: REPRODUCIBILITY STATEMENT
text: See App. G for experimental details and the uploaded supplementary material for Python and R code replicating all experiments.

heading: ETHICS STATEMENT
text: This is major theoretical work and we do not anticipate negative societal impact from this work.

heading: A MIRROR DESCENT, RIEMANNIAN GRADIENT FLOW, AND NATURAL GRADIENT
text: The equivalence between the mirror flow dηt = −∇f(θt)dt, θt = ∇ψ∗(ηt)dt and the Riemannian gradient flow in (3) is a direct result of the chain rule:
dθt dt = −∇ηtθt dηt dt = −(∇θtηt)−1 dηt dt = −∇2ψ(θt)−1∇f(θt), (16) dηt dt
= −∇f(θt) = −∇θtηt∇ηtf(∇ψ∗(ηt)) = −∇2ψ∗(ηt)−1∇ηtf(∇ψ∗(ηt)). (17) Depending on discretizing (16) or (17), there are two natural gradient descent (NGD) updates that can arise from the same gradient flow:
NGD (a): θt+1 = θt − t∇2ψ(θt)−1∇f(θt), NGD (b): ηt+1 = ηt − t∇2ψ∗(ηt)−1∇ηtf(∇ψ∗(ηt)).
With finite step sizes t, their updates need not be the same and can lead to different optimization paths. Since∇f(θt) = ∇2ψ∗(ηt)−1∇ηtf(∇ψ∗(ηt)), NGD (b) is equivalent to 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 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.

heading: B DETAILS OF EXAMPLE 1
text: For the entropic mirror map ψ(θ) = ∑d+1 j=1 θj log θj , we have ∇2ψ(θ)−1 = diag(θ) − θθ>. Note
here θ denotes a d-dimensional vector and does not include θd+1 = 1 − ∑d j=1 θd. Since Θ is a (d+ 1)-simplex, ∂Θ is composed of d+ 1 faces with θ in the j-th face satisfies θj = 0. The outward unit normal vector n(θ) for the first d faces are−ej for 1 ≤ j ≤ d, where ej denotes the j-th standard basis of Rd. The outward unit normal vector for the (d + 1)-st face is a vector with 1/ √ d in all coordinates. Therefore, we have∫ ∂Θ p(θ)g(θ)>∇2ψ(θ)−1n(θ)dθ = ∫ ∂Θ p(θ)g(θ)>(diag(θ)− θθ>)n(θ)dθ
= ∫ ∂Θ p(θ)(θ g(θ)− θθ>g(θ))>n(θ)dθ
= d∑ j=1 ∫ θj=0 p(θ)(θ>g(θ)− gj(θ))θjdθ−j
+ 1√ d ∫ θd+1=0 p(θ)θ>g(θ)θd+1dθ
=0,
where in the second to last identity we used θ>1 = 1− θd+1. Finally, we can verify the condition in Prop. 1 as
lim r→∞ ∫ ∂Θr p(θ)‖∇2ψ(θ)−1nr(θ)‖2dθ = sup ‖g‖∞≤1 ∫ ∂Θ p(θ)g(θ)>∇2ψ(θ)−1n(θ)dθ = 0.

heading: C DERIVATION OF THE MIRRORED STEIN OPERATOR
text: We first review the (overdamped) Langevin diffusion – a Markov process that underlies many recent advances in Stein’s method – along with its recent mirrored generalization. The Langevin diffusion with equilibrium density p on Rd is a Markov process (θt)t≥0 ⊂ Rd satisfying the stochastic differential equation (SDE)
dθt = ∇ log p(θt)dt+ √
2dBt (18) with (Bt)t≥0 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 (θt)t≥0 that has p as the equilibrium density; they chose the Langevin diffusion of (18). Next, build a Stein operator based on the (infinitesimal) generator A of the process (Øksendal, 2003, Def. 7.3.1):
(Af)(θ) = limt→0 1 t (Ef(θt)− Ef(θ0)) for f : Rd → R,
as the generator satisfies Eθ∼p[(Af)(θ)] = 0 under relatively mild conditions. We use the following theorem to derive the generator of the processes described by SDEs like (18): Theorem 9 (Generator of Itô diffusion; Øksendal, 2003, Thm 7.3.3). Let (xt)t≥0 be the Itô diffusion in X ⊆ Rd satisfying dxt = b(xt)dt+ σ(xt)dBt. For any f ∈ C2c (X ), the (infinitesimal) generator A of (xt)t≥0 is
(Af)(x) = b(x)>∇f(x) + 12 Tr(σ(x)σ(x)>∇2f(x)).
For the Langevin diffusion (18), substituting ∇ log p(·) for b(·) and √
2I for σ(·) in Thm. 9, we obtain Af = (∇ log p)>∇f + ∇ · ∇f . Replacing ∇f with a vector-valued function g gives the Langevin Stein operator in (5).
To derive a 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 non-Euclidean geometries encoded in a positive definite metric tensor G(θ):
dθt = (G(θt) −1∇ log p(θt) +∇ ·G(θt)−1)dt+ √ 2G(θt) −1/2dBt. 5
We show in App. D that the choice G = ∇2ψ yields the recent mirror-Langevin diffusion (Zhang et al., 2020b; Chewi et al., 2020)
θt = ∇ψ∗(ηt), dηt = ∇ log p(θt)dt+ √ 2∇2ψ(θt)1/2dBt. (19)
According to Thm. 9, the generator of the mirror-Langevin diffusion described by (20) is
(Ap,ψf)(θ) = (∇2ψ(θ)−1∇ log p(θ) +∇ · ∇2ψ(θ)−1)>∇f(θ) + Tr(∇2ψ(θ)−1∇2f(θ)) = ∇f(θ)>∇2ψ(θ)−1∇ log p(θ) +∇ · (∇2ψ(θ)−1∇f(θ)).
Now substituting g(θ) for ∇f(θ), we obtain the associated mirrored Stein operator: (Mp,ψg)(θ) = g(θ)>∇2ψ(θ)−1∇ log p(θ) +∇ · (∇2ψ(θ)−1g(θ)).

heading: D RIEMANNIAN LANGEVIN DIFFUSIONS AND MIRROR-LANGEVIN DIFFUSIONS
text: 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 et al., 2015) with ∇2ψ(·) as the metric tensor is
dθt = (∇2ψ(θt)−1∇ log p(θt) +∇ · ∇2ψ(θt)−1)dt+ √ 2∇2ψ(θt)−1/2dBt. (20)
To see the connection with mirror-Langevin diffusion, we would like to obtain the SDE that describes the evolution of ηt = ∇ψ(θt) under the diffusion. This requires the following theorem that provides the analog of the “chain rule” in SDEs.
5A matrix divergence∇ ·G(θ) is the vector obtained by computing the divergence of each row of G(θ).
Theorem 10 (Itô formula; Øksendal, 2003, Thm 4.2.1). Let (xt)t≥0 be an Itô process in X ⊂ Rd satisfying dxt = b(xt)dt+ σ(xt)dBt. Let f(x) ∈ C2 : Rd → Rd ′ . Then yt = f(xt) is again an Itô process, and its i-th dimension satisfies
dyt,i = (∇fi(xt)>b(xt) + 1
2 Tr(∇2fi(xt)σ(xt)σ(xt)>)dt+∇fi(xt)>σ(xt)dBt.
Substituting∇ψ for f in Thm. 10, we have the SDE of ηt = ∇ψ(θt) as
dηt = (∇ log p(θt) +∇2ψ(θt)∇ · ∇2ψ(θt)−1 + h(θt))dt+ √ 2∇2ψ(θt)1/2dBt,
where h(θt)i = Tr(∇2θt(∇θt,iψ(θt))∇2ψ(θt)−1). Moreover, we have
[∇2ψ(θt)∇ · ∇2ψ(θt)−1]i + Tr(∇2θt(∇θt,iψ(θt))∇2ψ(θt)−1)
= d∑ `=1 d∑ j=1 ∇2ψ(θt)ij∇θt,` [∇2ψ(θt)−1]j` + d∑ `=1 d∑ j=1 ∇θt,`∇2ψ(θt)ij [∇2ψ(θt)−1]j`
= d∑ `=1 ∇θt,`  d∑ j=1 ∇2ψ(θt)ij [∇2ψ(θt)−1]j`  = d∑ `=1 ∇θt,`Ii` = 0.
Therefore, the ηt diffusion is described by the SDE:
dηt = ∇ log p(θt)dt+ √ 2∇2ψ(θt)1/2dBt, θt = ∇ψ∗(ηt).

heading: E MODE MISMATCH UNDER TRANSFORMATIONS

heading: F BACKGROUND ON REPRODUCING KERNEL HILBERT SPACES
text: Let H be a Hilbert space of functions defined on X and taking their values in R. We say k is a reproducing kernel (or kernel) ofH if ∀x ∈ X , k(x, ·) ∈ H and ∀f ∈ H, 〈f, k(x, ·)〉H = f(x). H is called a reproducing kernel Hilbert space (RKHS) if it has a kernel. Kernels are positive definite (p.d.) functions, which means that matrices with the form (k(xi, xj))ij are positive semidefinite. 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 {∑ni=1 aik(xi, ·), xi ∈ X , ai ∈ R, i ∈ N}. Now we assume X is a metric space, k is a bounded continuous kernel with the RKHSH, and ν is a positive measure on X . L2(ν) denote the space of all square-integrable functions w.r.t. ν. Then the kernel integral operator Tk : L2(ν)→ L2(ν) defined by
Tkg = ∫ X g(x)k(x, ·)dν
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 {uj}j∈J such that
Tkuj = λjuj , and k has the representation k(x, x′) = ∑ j∈J λjuj(x)uj(x
′) (Mercer’s theorem on non-compact domains), where the convergence of the sum is absolute and uniform on compact subsets of X × X (Ferreira & Menegatto, 2009).

heading: G SUPPLEMENTARY EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS
text: 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(θ, θ′) = (1 + ‖θ − θ′‖22/`2)−1/2 due to its convergence control properties (Gorham & Mackey, 2017). In the unconstrained experiments of Sec. 5.3, we use the Gaussian kernel k(θ, θ′) = exp(−‖θ − θ′‖22/`2) for consistency with past work. The bandwidth ` 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 fixed step size ∈ {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; specifically, we use RMSProp (Hinton et al., 2012), an extension of the AdaGrad algorithm (Duchi et al., 2011) 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 PCIe 16GB.

heading: G.1 APPROXIMATION QUALITY ON THE SIMPLEX
text: The sparse Dirichlet posterior of Patterson & Teh (2013) extended to 20 dimensions features a sparse, symmetric Dir(α) prior with αk = 0.1 for k ∈ {1, . . . , 20} and sparse count data n1 = 90, n2 = n3 = 5, nj = 0 (j > 3), modeled via a multinomial likelihood. The quadratic target satisfies log p(θ) = − 12σ2 θ>Aθ+ const, where we slightly modify the target density of Ahn & Chewi (2020) to make it less flat by introducing a scale parameter σ = 0.01. A ∈ R19×19 is a positive definite 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.98. 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. 7.

heading: G.2 CONFIDENCE INTERVALS FOR POST-SELECTION INFERENCE
text: Given a dataset X ∈ Rñ×p, y ∈ Rñ, the randomized Lasso (Tian & Taylor, 2018) solves the following problem:
argminβ∈Rp 1 2‖y −Xβ‖ 2 2 + λ‖β‖1 − w>β + 2‖β‖ 2 2, w ∼ G.
where G is a user-specified log-concave distribution with density g. We choose G to be zeromean independent Gaussian distributions while leaving its scale and the ridge parameter 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 & Markovic (2017) except using a different 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 xi ∈ Rp is generated i.i.d. from N (0,Σ) with Σii = 1,Σij = 0.3 (i 6= j) and then normalized to have almost unit length. The normalization is 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̃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 coefficient 0.7ñ, where the 0.7 adjustment is introduced in the test examples of the R package to reduce the regularization effect 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.01 and we run them for T = 1000 particle updates. τ 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 logdensity log p(θ) = −8.07193((2.39859θ1 + 1.90816θ2 + 2.39751)2 + (1.18099θ2 − 1.46104)2) + const, θ1,2 ≥ 0.
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 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.01 and we run them for T = 2000 particle updates. τ is set to 0.99 for SVMD.

heading: G.3 LARGE-SCALE POSTERIOR INFERENCE WITH NON-EUCLIDEAN GEOMETRY
text: The Bayesian logistic regression model we consider is ∏ñ i=1 p(yi|xi, w)p(w), where p(w) = N (w|0, I), p(yi|xi, w) = Bernoulli(σ(w>xi)). The bias parameter is absorbed into into w by adding an additional feature 1 to each xi. The gradient of the log density of the posterior distribution of w is ∇w log p(w|{yi, xi}Ni=1) = ∑N i=1 xi(yi − σ(w>xi)) − w. We choose the metric tensor ∇2ψ(w) to be the Fisher information matrix (FIM) of the likelihood:
F = 1
ñ ñ∑ i=1 Ep(yi|w,xi)[∇w log p(yi|xi, w)∇w log p(yi|xi, w)>]
= 1
ñ ñ∑ i=1 σ(w>xi)(1− σ(w>xi))xix>i .
Following Wang et al. (2019), for each iteration r (r ≥ 1), we estimate the sum with a stochastic minibatch Br of size 256: F̂Br = ñ|Br| ∑ i∈Br σ(w
>xi)(1 − σ(w>xi))xix>i and approximate the FIM with a moving average across iterations:
F̂r = ρrF̂r−1 + (1− ρr)F̂Br , where ρr = min(1− 1/r, 0.95). To ensure the positive definiteness of the FIM, a damping term 0.01I is added before taking the inverse. For RSVGD and SVNG, the gradient of the inverse of FIM is estimated with ∇wjF−1 ≈ −F̂−1r (∇̂rwjF )F̂−1r , where ∇̂rwjF = ρr∇̂r−1wj F + (1− ρr)∇wj F̂Br . 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. 10. 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 significant loss in performance. The results in Fig. 5 are reported with the learning rate that performs best on the validation set.

heading: H EXPONENTIAL CONVERGENCE OF CONTINUOUS-TIME ALGORITHMS
text: We derive a time-inhomogeneous 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-time algorithms and time-inhomogeneous generalization of the Stein Poincaré inequality of Duncan et al. (2019) which guarantees exponential convergence near equilibrium (i.e., when qt is sufficiently close to p). As the results hold for a generic sequence of kernels (Kt)t≥0, the implications apply to both MSVGD and SVMD.
Definition 2 (Mirror Stein Log-Sobolev inequality). We define the Mirror Stein Log-Sobolev inequality (cf., Korba et al., 2020, Def. 2) as
KL(qt‖p) ≤ 1
2λ MKSD2Kt(qt, p) =
1
2λ Eqt [(∇2ψ−1∇ log qt p )>PKt,qt∇2ψ−1∇ log qt p ],
where MKSDKt is defined in Eq. (15); PKt,qt : L 2(qt) → L2(qt) is the kernel integral operator: (PKt,qtϕ)(·) , Eqt(θ)[Kt(·, θ)ϕ(θ)] for a general kernel Kt and vector-valued function ϕ on Θ, and the stated equality holds whenever integration-by-parts is applicable. Proposition 11. Suppose (θt)t≥0 follows the mirrored dynamics (8) with gt chosen to be g∗qt,Kt as in (10). Then, the dissipation of KL(qt‖p) is
d dt KL(qt‖p) = −MKSDKt(qt, p)2.
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
dt KL(qt‖p) ≤ −2λKL(qt‖p),
and the exponential convergence KL(qt‖p) ≤ KL(q0‖p)e−2λt follows by Gronwall’s lemma (Gronwall, 1919). Definition 3 (Mirror Stein Poincaré inequality). We say that the distribution p satisfies the Mirror Stein Poincaré inequality (cf., Duncan et al. 2019, Eq. (57)) with strongly convex ψ and constant λ if
Varp[φ] ≤ 1
λ Ep[∇φ>PKt,p∇2ψ−1∇φ]
for all φ ∈ L2(p) ∩ C∞(Θ) that is locally Lipschitz, where PKt,p is the kernel integral operator under p defined similarly as in Definition 2.
This inequality can also be viewed as a kernelized generalization of the mirror Poincaré inequality introduced in Chewi et al. (2020, Def. 1) for proving exponential convergence of mirror-Langevin diffusion. In a manner analogous to Thm. 1 of Chewi et al. (2020), the following proposition relates the Mirror Stein Poincaré inequality to chi-squared divergence. Proposition 12. Suppose (θt)t≥0 follows the mirrored dynamics (8) with gt chosen to be g∗qt,Kt as in (10). Then, the dissipation of chi-square divergence χ2(qt‖p) is
d dt χ2(qt‖p) = −2Eqt [ ∇qt p > PKt,p∇2ψ−1∇ qt p ] whenever integration-by-parts is applicable.
Proof We first note that by applying integration-by-parts, g∗qt,Kt as in (10) can be equivalently written as
g∗qt,Kt = −PKt,qt∇2ψ−1∇ log qt p .
Then using the Fokker-Planck equation of qt under the dynamics, we have
d dt χ2(qt‖p) = d dt ∫ ( qt p )2 dp = 2 ∫ qt p · dqt dt = −2Eqt [ g∗t,Kt >∇qt p ] = −2Eqt [( PKt,qt∇2ψ−1∇ log
qt p )> ∇qt p ] = −2Eqt [ ∇qt p > PKt,p∇2ψ−1∇ qt p ] .
Note that the right hand side of the equation differs from the Mirror Stein Poincaré inequality only in the base measure of the expectation. Duncan et al. (2019) proposes to replace qt with p to study the convergence near equilibrium (See their Sec. 6, where Eq. (46) is replaced with Eq. (51)). If we do the same and combine this identity with the Mirror Stein Poincaré inequality, we obtain ddtχ
2(qt‖p) ≤ −2λVarp[ qtp ] = −2λχ2(qt‖p), which implies exponential convergence in KL KL(qt‖p) ≤ χ2(qt‖p) ≤ χ2(q0‖p)e−2λt by Gronwall’s lemma (Gronwall, 1919).

heading: I PROOFS

heading: I.1 PROOF OF PROP. 1
text: 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[(Mp,ψg)(θ)] exists. Because Θ is convex, Θr is bounded and convex with Lipschitz boundary. Since p∇2ψ−1g ∈ C1, we have |Ep[(Mp,ψg)(θ)]| = |Ep[g(θ)>∇2ψ(θ)−1∇ log p(θ) +∇ · (∇2ψ(θ)−1g(θ))]|
= ∣∣∣∣∫ Θ ∇p(θ)>∇2ψ(θ)−1g(θ) + p(θ)∇ · (∇2ψ(θ)−1g(θ))dθ ∣∣∣∣
= ∣∣∣∣∫ Θ ∇ · (p(θ)∇2ψ(θ)−1g(θ))dθ ∣∣∣∣
= ∣∣∣∣ limr→∞ ∫
Θr
∇ · (p(θ)∇2ψ(θ)−1g(θ))dθ ∣∣∣∣ (by dominated convergence)
= ∣∣∣∣ limr→∞ ∫ ∂Θr (p(θ)∇2ψ(θ)−1g(θ))>nr(θ)dθ ∣∣∣∣ (by the divergence theorem)
≤ lim r→∞ ∫ ∂Θr p(θ)‖g(θ)‖2 ∥∥∇2ψ(θ)−1nr(θ)∥∥2dθ (by Cauchy-Schwarz)
≤ ‖g‖∞ limr→∞ ∫ ∂Θr p(θ) ∥∥∇2ψ(θ)−1nr(θ)∥∥2dθ = 0 (by assumption).

heading: I.2 PROOF OF THM. 3: OPTIMAL MIRROR UPDATES IN RKHS
text: Proof Let ei denote the standard basis vector of Rd with the i-th element being 1 and others being zeros. Since m ∈ HK , we have
m(θ)>∇2ψ(θ)−1∇ log p(θ) = 〈m,K(·, θ)∇2ψ(θ)−1∇ log p(θ)〉HK
∇ · (∇2ψ(θ)−1m(θ)) = d∑ i=1 ∇θi(m(θ)>∇2ψ(θ)−1ei)
= d∑ i=1 〈m,∇θi(K(·, θ)∇2ψ(θ)−1ei)〉HK
= 〈m,∇θ · (K(·, θ)∇2ψ(θ)−1)〉HK , where we define the divergence of a matrix as a vector whose elements are the divergences of each row of the matrix. Then, we write (9) as
− Eqt [m(θ)>∇2ψ(θ)−1∇ log p(θ) +∇ · (∇2ψ(θ)−1m(θ))] = −Eqt [〈m,K(·, θ)∇2ψ(θ)−1∇ log p(θ) +∇θ · (K(·, θ)∇2ψ(θ)−1)〉HK ] = −〈m,Eqt [K(·, θ)∇2ψ(θ)−1∇ log p(θ) +∇θ · (K(·, θ)∇2ψ(θ)−1)]〉HK = −〈m,Eqt [Mp,ψK(·, θ)]〉HK .
Therefore, the optimal direction in theHK norm ball BHK = {g : ‖g‖HK ≤ 1} that minimizes (9) is g∗t ∝ g∗qt,K = Eqt [Mp,ψK(·, θ)].

heading: I.3 PROOF OF THM. 4: MIRRORED SVGD UPDATES
text: Proof A p.d. kernel k composed with any map φ is still a p.d. kernel. To prove this, let {x1, . . . , xp} = {φ(η1), . . . , φ(ηn)}, p ≤ n. Then∑
i,j αiαjk(φ(ηi), φ(ηj)) = ∑ `,m β`βmk(x`, xm) ≥ 0,
where β` = ∑ i∈S` αi, S` = {i : φ(ηi) = x`}. Therefore, kψ(η, η′) = k(∇ψ∗(η),∇ψ∗(η′)) is a p.d. kernel. Plugging K = kI into Lem. 13, for any θ′ ∈ Θ and η′ = ∇ψ(θ′), we have g∗qt,Kk(θ ′) = Eηt∼qt,H [K∇ψ∗(∇ψ(θ′), ηt)∇ log pH(ηt) +∇ηt ·K∇ψ∗(∇ψ(θ′), ηt)]
= Eηt∼qt,H [k(∇ψ∗(η′),∇ψ∗(ηt))∇ log pH(ηt) + d∑ j=1 ∇ηt,jk(∇ψ∗(η′),∇ψ∗(ηt))ej ] = Eηt∼qt,H [kψ(η′, ηt)∇ log pH(ηt) +∇ηtkψ(η′, ηt)].

heading: I.4 PROOF OF PROP. 5: SINGLE-PARTICLE SVMD IS MIRROR DESCENT
text: Proof When n = 1, λ1 = k(θt, θt), u1 = 1, and thus Kψ,t(θt, θt) = k(θt, θt)∇2ψ(θt).
I.5 PROOF OF THM. 6: 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. 13 and then follow the path of Gorham et al. (2020, Thm. 7). Assume that qnt,H and q ∞ t,H have integrable means. Let ηn, η∞ be an optimal Wasserstein-1 coupling of qnt,H and q ∞ t,H . Let Φqt,Kt denote the transform through one step of mirrored update: θt = ∇ψ?(ηt), ηt+1 = ηt + tg∗qt,Kt(θt). Then, with Lem. 13, we have
‖Φqt,Kt(η)− Φqt,Kt(η′)‖2 = ‖η + tg∗qnt ,Kt(θ)− η ′ − tg∗q∞t ,Kt(θ ′)‖2
≤ ‖η − η′‖2 + t‖g∗qnt ,Kt(θ)− g ∗ q∞t ,Kt (θ′)‖2 ≤ ‖η − η′‖2 + t‖Eηn [K∇ψ∗,t(η, ηn)∇ log pH(ηn) +∇ηn ·K∇ψ∗,t(η, ηn)
− (K∇ψ∗,t(η′, ηn)∇ log pH(ηn) +∇ηn ·K∇ψ∗,t(η′, ηn))]‖2 + t‖Eηn,η∞ [K∇ψ∗,t(η′, ηn)∇ log pH(ηn) +∇ηn ·K∇ψ∗,t(η, ηn) − (K∇ψ∗,t(η′, η∞)∇ log pH(η∞) +∇η∞ ·K∇ψ∗,t(η′, η∞))]‖2 ≤ ‖η − η′‖2 + tc1(1 + E[‖ηn‖2)]‖η − η′‖2 + tc2(1 + ‖η′‖2)Eηn,η∞ [‖ηn − η∞‖2] = ‖η − η′‖2 + tc1(1 + Eqnt,H [‖·‖2])‖η − η
′‖2 + tc2(1 + ‖η′‖2)W1(qnt,H , q∞t,H).

heading: Since Φqt,Kt(η
text: n) ∼ qnt+1,H , Φqt,K(η∞) ∼ q∞t+1,H , we conclude W1(q n t+1,H , q ∞ t+1,H)
≤ E[‖Φqt,K(ηn)− Φqt,K(η∞)‖2] ≤ (1 + tc1(1 + Eqnt,H [‖·‖2]))E[‖η
n − η∞‖2] + tc2(1 + ‖η′‖2)W1(qnt,H , q∞t,H)] ≤ (1 + tc1(1 + Eqnt,H [‖·‖2]) + tc2(1 + Eq∞t,H [‖·‖2]))W1(q n t,H , q ∞ t,H).
The final claim qnt ⇒ q∞t now follows by the continuous mapping theorem as∇ψ∗ is continuous.

heading: I.6 PROOF OF THM. 7: INFINITE-PARTICLE MIRRORED STEIN UPDATES DECREASE KL AND
text: MKSD
Proof Let Tq∞t ,Kt denote transform of the density function through one step of mirrored update: θt = ∇ψ?(ηt), ηt+1 = ηt + tg∗q∞t ,Kt(θt). Then
KL(q∞t+1‖p)− KL(q∞t ‖p) = KL(q∞t ‖T−1q∞t ,Ktp)− KL(q ∞ t ‖p)
= Eηt∼q∞t,H [log pH(ηt)− log pH(ηt + tg ∗ q∞t ,Kt (θt))− log |det(I + t∇ηtg∗q∞t ,Kt(θt))|],
where we have used the invariance of KL divergence under reparameterization: KL(qt‖p) = KL(qt,H‖pH) . Following Liu (2017), we bound the difference of the first two terms as
log pH(ηt)− log pH(ηt + tg∗q∞t ,Kt(θt)) = − ∫ 1
0
∇s log pH(ηt(s)) ds, where ηt(s) , ηt + s tg∗q∞t ,Kt(θt)
= − ∫ 1
0
∇ log pH(ηt(s))>( tg∗q∞t ,Kt(θt)) ds
= − t∇ log pH(ηt)>g∗q∞t ,Kt(θt) + ∫ 1
0
(∇ log pH(ηt)−∇ log pH(ηt(s)))>( tg∗q∞t ,Kt(θt)) ds
≤ − t∇ log pH(ηt)>g∗q∞t ,Kt(θt) + t ∫ 1
0
‖∇ log pH(η)−∇ log pH(ηt(s))‖2 · ‖g∗q∞t ,Kt(θt)‖2 ds
≤ − t∇ log pH(ηt)>g∗q∞t ,Kt(θt) + L 2t 2 ‖g∗q∞t ,Kt(θt)‖ 2 2,
and bound the log determinant term using Lem. 15:
− log |det(I + t∇ηtg∗q∞t ,Kt(θt)) ≤ − t Tr(∇ηtg ∗ q∞t ,Kt (θt)) + 2 2 t‖∇ηtg∗q∞t ,Kt(θt)‖ 2 F .
The next thing to notice is that Eηt∼q∞t,H [∇ log pH(ηt)>g∗q∞t ,Kt(θt) + Tr(∇ηtg ∗ q∞t ,Kt (θt))] is the square of the MKSD in (15). We can show this equivalence using the identity proved in Lem. 14:
Eηt∼q∞t,H [g ∗ q∞t ,Kt (θt) >∇ log pH(ηt) + Tr(∇ηtg∗q∞t ,Kt(θt))]
= Eθt∼q∞t [g ∗ q∞t ,Kt (θt) >∇2ψ(θt)−1∇θt(log p(θt)− log det∇2ψ(θt))
+ Tr(∇2ψ(θt)−1∇g∗q∞t ,Kt(θt))] = Eθt∼q∞t [g ∗ q∞t ,Kt (θt) >∇2ψ(θt)−1∇ log p(θt) +∇ · (∇2ψ(θt)−1g∗q∞t ,Kt(θt))] (Lem. 14) = Eθt∼q∞t [(Mp,ψg∗q∞t ,Kt)(θt)] = MKSDKt(q ∞ t , p) 2.
Finally, we are going to bound ‖g∗q∞t ,Kt(θt)‖ 2 2 and ‖∇ηtg∗q∞t ,Kt(θt)‖ 2 F . From the assumptions we have ψ is α-strongly convex and thus ψ∗ is 1α -strongly smooth (Kakade et al., 2009), therefore ‖∇2ψ∗(·)‖2 ≤ 1α . By Lem. 16, we know
‖g∗q∞t ,Kt(θt)‖ 2 2 ≤ ‖g∗q∞t ,Kt‖ 2 HKt ‖K(θt, θt)‖op = MKSDKt(q∞t , p)2‖Kt(θt, θt)‖op,
‖∇ηtg∗q∞t ,Kt(θt)‖ 2 F = ‖∇2ψ∗(ηt)∇g∗q∞t ,Kt(θt)‖ 2 F
≤ ‖∇2ψ∗(ηt)‖22‖∇g∗q∞t ,Kt(θt)‖ 2 F
≤ 1 α2 ‖g∗q∞t ,Kt‖ 2 HKt d∑ i=1 ‖∇2i,d+iKt(θt, θt)‖op
= 1
α2 MKSDKt(q
∞ t , p) 2 d∑ i=1 ‖∇2i,d+iKt(θt, θt)‖op,
where ∇2i,d+iK(θ, θ) denotes∇2θi,θ′iK(θ, θ ′)|θ′=θ. Combining all of the above, we have
KL(q∞t+1‖p)− KL(q∞t ‖p) ≤ − ( t −
L 2t 2 sup θ ‖Kt(θ, θ)‖op − 2 2t α2 d∑ i=1 sup θ ‖∇2i,d+iKt(θ, θ)‖op ) MKSDKt(q ∞ t , p) 2.
Plugging in the definition of κ1 and κ2 finishes the proof.
I.7 PROOF OF THM. 8: MKSDKk DETERMINES WEAK CONVERGENCE
Proof According to Thm. 4,
g∗q,Kk = EqH [k(·,∇ψ∗(η))∇ log pH(η) +∇ηk(∇ψ∗(η), ·)], where qH(η) denotes the density of η = ∇ψ(θ) under the distribution θ ∼ q. From the assumptions we have k(θ, θ′) = κ(∇ψ(θ),∇ψ(θ′)). With this specific choice of k, the squared MKSD is
MKSDKk(q, p) 2 = ‖g∗q,Kk‖2HKk = Eη,η′∼qH [
1
pH(η)pH(η′) ∇η∇η′(pH(η)k(∇ψ∗(η),∇ψ∗(η′))pH(η′)) ] = Eη,η′∼qH [ 1
pH(η)pH(η′) ∇η∇η′(pH(η)κ(η, η′)pH(η′))
] . (21)
The final expression in (21) is the squared kernel Stein discrepancy (KSD) (Liu et al., 2016; Chwialkowski et al., 2016; Gorham & Mackey, 2017) between qH and pH with the kernel κ: KSDκ(qH , pH)2. Recall that it is proved in Gorham & Mackey (2017, Theorem 8) that, for κ(x, y) = (c2 + ‖x − y‖22)β with β ∈ (−1, 0) and distantly dissipative pH with Lipschitz score functions, qH ⇒ pH if KSDκ(qH , pH)→ 0. The advertised result (q ⇒ p if MKSDKk(q, p)→ 0) now follows by the continuous mapping theorem as∇ψ∗ is continuous.

heading: J LEMMAS
text: Lemma 13. Let K∇ψ∗(η, η′) , K(∇ψ∗(η),∇ψ∗(η′)). The mirrored updates g∗qt,K in (10) can be equivalently expressed as
g∗qt,K = Eqt,H [K∇ψ∗(∇ψ(·), η)∇ log pH(η) +∇η ·K∇ψ∗(∇ψ(·), η)].
Proof We will use the identity proved in Lem. 14.
g∗qt,K = Eqt [Mp,ψK(·, θ)] = Eqt [K(·, θ)∇2ψ(θ)−1∇ log p(θ) +∇θ · (K(·, θ)∇2ψ(θ)−1)] = Eqt [K(·, θ)∇2ψ(θ)−1∇θ(log pH(∇ψ(θ)) + log det∇2ψ(θ)) +∇θ · (K(·, θ)∇2ψ(θ)−1)]
(by change-of-variable formula)
= Eqt [K(·, θ)∇2ψ(θ)−1∇θ log pH(∇ψ(θ)) + d∑
i,j=1
[∇2ψ(θ)−1]ij∇θiK(·, θ):,j ]
(by applying Lem. 14 to each row of K(·, θ))
= Eqt [K(·, θ)∇2ψ(θ)−1∇θ log pH(∇ψ(θ)) + d∑ j=1 ∇ηjK(·, θ):,j ] = Eqt,H [K(·,∇ψ∗(η))∇ log pH(η) + d∑ j=1 ∇ηjK(·,∇ψ∗(η)):,j ]
= Eqt,H [K∇ψ∗(∇ψ(·), η)∇ log pH(η) +∇η ·K∇ψ∗(∇ψ(·), η)], where A:,j denotes the j-th column of a matrix A.
Lemma 14. For a strictly convex function ψ ∈ C2 : Rd → R and any vector-valued g ∈ C1 : Rd → Rd, the following relation holds:
∇ · (∇2ψ(θ)−1g(θ)) = Tr(∇2ψ(θ)−1∇g(θ))− g(θ)>∇2ψ(θ)−1∇θ log det∇2ψ(θ).
Proof By the product rule of differentiation:
∇ · (∇2ψ(θ)−1g(θ)) = Tr(∇2ψ(θ)−1∇g(θ)) + g(θ)>∇ · (∇2ψ(θ)−1). (22)
This already gives us the first term on the right side. Next, we have
[∇2ψ(θ)−1∇ log det∇2ψ(θ)]i
= d∑ j=1 [∇2ψ(θ)−1]ij Tr(∇2ψ(θ)−1∇θj∇2ψ(θ))
= d∑ j=1 [∇2ψ(θ)−1]ij d∑ `,m=1 [∇2ψ(θ)−1]`m[∇θj∇2ψ(θ)]m`
= d∑ j,`,m=1 [∇2ψ(θ)−1]ij [∇2ψ(θ)−1]`m∇θj∇2ψ(θ)m`
= d∑ j,`,m=1 [∇2ψ(θ)−1]ij∇θm∇2ψ(θ)j`[∇2ψ(θ)−1]`m = − d∑
m=1
∇θm(∇2ψ(θ)−1)im
= −[∇ · ∇2ψ(θ)−1]i.
Plugging the above relation into (22) proves the claimed result.
Lemma 15 (Liu, 2017, Lemma A.1). Let A be a square matrix, and 0 < < 12‖A+A>‖op. Then,
log |det(I + A)| ≥ Tr(A)− 2 2‖A‖2F ,
where ‖ · ‖F denotes the Frobenius norm of a matrix.
Lemma 16. Let K be a matrix-valued kernel andHK be the corresponding RKHS. Then, for any f ∈ HK (f is vector-valued), we have
‖f(x)‖2 ≤ ‖f‖HK‖K(x, x)‖ 1/2 op , ‖∇f(x)‖ 2 F ≤ ‖f‖ 2 HK d∑ i=1 ‖∇2xi,x′iK(x, x ′)|x′=x‖op,
where ‖ · ‖op denotes the operator norm of a matrix induced by the vector 2-norm.
Proof We first bound the ‖f(x)‖2 as
‖f(x)‖2 = sup ‖y‖2=1 f(x)>y = sup ‖y‖2=1 〈f,K(·, x)y〉HK ≤ ‖f‖HK sup ‖y‖2=1 ‖K(·, x)y‖HK
= ‖f‖HK sup ‖y‖2=1 (y>K(x, x)y)1/2 ≤ ‖f‖HK sup ‖y‖2=1 sup ‖u‖2=1 (u>K(x, x)y)1/2
= ‖f‖HK sup ‖y‖2=1 ‖K(x, x)y‖1/22 = ‖f‖HK‖K(x, x)‖ 1/2 op .
The second result follows similarly,
‖∇f(x)‖2F = d∑ i=1 ‖∇xif(x)‖22 = d∑ i=1 sup ‖y‖2=1 (∇xif(x)>y)2 = d∑ i=1 sup ‖y‖2=1 (∇xi〈f,K(·, x)y〉HK )2
= ∑ i=1 sup ‖y‖2=1 (〈f,∇xiK(·, x)y〉HK )2 ≤ ‖f‖2HK d∑ i=1 sup ‖y‖2=1 ‖∇xiK(·, x)y‖2HK
= ‖f‖2HK d∑ i=1 sup ‖y‖2=1 (y>∇2xi,x′iK(x, x ′)|x=x′y)
≤ ‖f‖2HK d∑ i=1 sup ‖y‖2=1 sup ‖u‖2=1 (u>∇2xi,x′iK(x, x ′)|x=x′y)
= ‖f‖2HK d∑ i=1 sup ‖y‖2=1 ‖∇2xi,x′iK(x, x ′)|x=x′y‖2
= ‖f‖2HK d∑ i=1 ‖∇2xi,x′iK(x, x ′)|x′=x‖op.

Review summaries:
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).  

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.

The paper introduces new methods to run SVGD in constrained domains and non-euclidean geometries. The authors develop the theory to combine mirror descent dynamics with the Stein's method via what they call mirrored Stein operators. They also show experimental results on two concrete problems and prove convergence guarantees.

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.