Original: S AMPLING WITH M IRRORED S TEIN O PERATORS
Revised: S AMPLING WITH M IRRORED S TEIN O PERATORS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We introduce a new family of particle evolution samplers suitable for constrained domains and non-Euclidean geometries.
Revised: We introduce a new family of particle evolution samplers suitable for constrained domains and non-Euclidean geometries.

Original: Stein Variational Mirror Descent and
Revised: Stein Variational Mirror Descent and

Original: Mirrored Stein Variational Gradient Descent minimize the Kullback-Leibler (KL) divergence to constrained target distributions by evolving particles in a dual space deﬁned by a mirror map.
Revised: Mirrored Stein Variational Gradient Descent minimize the Kullback-Leibler (KL) divergence to constrained target distributions by evolving particles in a dual space deﬁned by a mirror map.

Original: Stein Variational Natural Gradient exploits non-Euclidean geometry to more efﬁciently minimize the KL divergence to unconstrained targets.
Revised: Stein Variational Natural Gradient exploits non-Euclidean geometry to more efﬁciently minimize the KL divergence to unconstrained targets.

Original: We derive these samplers from a new class of mirrored Stein operators and adaptive kernels developed in this work.
Revised: We derive these samplers from a new class of mirrored Stein operators and adaptive kernels developed in this work.

Original: We demonstrate that these new samplers yield accurate approximations to distributions on the simplex, deliver valid conﬁdence intervals in post-selection inference, and converge more rapidly than prior methods in large-scale unconstrained posterior inference.
Revised: We demonstrate that these new samplers yield accurate approximations to distributions on the simplex, deliver valid conﬁdence intervals in post-selection inference, and converge more rapidly than prior methods in large-scale unconstrained posterior inference.

Original: Finally, we establish the convergence of our new procedures under veriﬁable conditions on the target distribution.
Revised: Finally, we establish the convergence of our new procedures under veriﬁable conditions on the target distribution.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Accurately approximating an unnormalized distribution with a discrete sample is a fundamental challenge in machine learning, probabilistic inference, and Bayesian inference.
Revised: Accurately approximating an unnormalized distribution with a discrete sample is a fundamental challenge in machine learning, probabilistic inference, and Bayesian inference.

Original: 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;
Revised: 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;

Original: Gorham & Mackey, 2015; Oates et al., 2017; Liu et al., 2016; Chwialkowski et al., 2016; Gorham &
Revised: Gorham & Mackey, 2015; Oates et al., 2017; Liu et al., 2016; Chwialkowski et al., 2016; Gorham &

Original: Mackey, 2017) and reproducing kernels (Berlinet & Thomas-Agnan, 2011) to sequentially minimize
Revised: Mackey, 2017) and reproducing kernels (Berlinet & Thomas-Agnan, 2011) to sequentially minimize

Original: Kullback-Leibler (KL) divergence.
Revised: Kullback-Leibler (KL) divergence.

Original: 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).
Revised: 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).

Original: In this work, we derive a family of particle evolution samplers suitable for target distributions with constrained domains and non-Euclidean geometries.
Revised: In this work, we derive a family of particle evolution samplers suitable for target distributions with constrained domains and non-Euclidean geometries.

Original: Our development draws inspiration from mirror descent (MD) (Nemirovskij & Yudin, 1983), a ﬁrst-order optimization method that generalizes gradient descent with non-Euclidean geometry.
Revised: Our development draws inspiration from mirror descent (MD) (Nemirovskij & Yudin, 1983), a ﬁrst-order optimization method that generalizes gradient descent with non-Euclidean geometry.

Original: To sample from a distribution with constrained support, our method ﬁrst maps particles to a dual space.
Revised: To sample from a distribution with constrained support, our method ﬁrst maps particles to a dual space.

Original: There, we update particle locations using a new class of mirrored Stein operators and adaptive reproducing kernels introduced in this work.
Revised: There, we update particle locations using a new class of mirrored Stein operators and adaptive reproducing kernels introduced in this work.

Original: Finally, the dual particles are mapped back to sample points in the original space, ensuring that all constraints are satisﬁed.
Revised: Finally, the dual particles are mapped back to sample points in the original space, ensuring that all constraints are satisﬁed.

Original: We illustrate this procedure in Fig.
Revised: We illustrate this procedure in Fig.

Original: In Sec. 3, we develop two algorithms –
Revised: In Sec. 3, we develop two algorithms –

Original: 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.
Revised: 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.

Original: In addition, by exploiting the connection between MD and natural gradient descent (Amari, 1998; Raskutti &
Revised: In addition, by exploiting the connection between MD and natural gradient descent (Amari, 1998; Raskutti &

Original: Mukherjee, 2015), we develop a third algorithm – Stein Variational Natural Gradient (SVNG) – that extends SVMD to unconstrained targets with non-Euclidean geometry.
Revised: Mukherjee, 2015), we develop a third algorithm – Stein Variational Natural Gradient (SVNG) – that extends SVMD to unconstrained targets with non-Euclidean geometry.

Original: 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 &
Revised: 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 &

Original: Tibshirani, 2015; Lee et al., 2016; Tian & Taylor, 2018), and unconstrained large-scale posterior inference with the Fisher information metric.
Revised: Tibshirani, 2015; Lee et al., 2016; Tian & Taylor, 2018), and unconstrained large-scale posterior inference with the Fisher information metric.

Original: Finally, we analyze the convergence of our mirrored algorithms in Sec. 6 and discuss our results in Sec. 7.
Revised: Finally, we analyze the convergence of our mirrored algorithms in Sec. 6 and discuss our results in Sec. 7.

Original: [Equation]
Revised: [Equation]

Original: ∇ ψ
Revised: ∇ ψ

Original: [Equation]
Revised: [Equation]

Original: θ t θ t + η t η t
Revised: θ t θ t + η t η t

Original: + (cid:15) t
Revised: + (cid:15) t

Original: E θ ∼ q t [ M p,ψ K ( θ t , θ )]
Revised: E θ ∼ q t [ M p,ψ K ( θ t , θ )]

Original: Related work Our mirrored Stein operators (6) are instances of diffusion Stein operators in the sense of Gorham & Mackey (2017), but their speciﬁc properties have not been studied, nor have they been used to develop sampling algorithms.
Revised: Related work Our mirrored Stein operators (6) are instances of diffusion Stein operators in the sense of Gorham & Mackey (2017), but their speciﬁc properties have not been studied, nor have they been used to develop sampling algorithms.

Original: 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;
Revised: 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;

Original: 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).
Revised: 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).

Original: The closest to our work in this space is the recent marriage of mirror descent and MCMC.
Revised: The closest to our work in this space is the recent marriage of mirror descent and MCMC.

Original: For example, Hsieh et al.
Revised: For example, Hsieh et al.

Original: (2018) propose to run
Revised: (2018) propose to run

Original: Langevin Monte Carlo (LMC, an Euler discretization of the Langevin diffusion) in a mirror space.
Revised: Langevin Monte Carlo (LMC, an Euler discretization of the Langevin diffusion) in a mirror space.

Original: Zhang et al.
Revised: Zhang et al.

Original: (2020b) analyze the convergence properties of the mirror-Langevin diffusion, Chewi et al.
Revised: (2020b) analyze the convergence properties of the mirror-Langevin diffusion, Chewi et al.

Original: (2020) demonstrate its advantages over the Langevin diffusion when using a Newton-type metric, and
Revised: (2020) demonstrate its advantages over the Langevin diffusion when using a Newton-type metric, and

Original: Ahn & Chewi (2020) study its discretization for MCMC sampling in constrained domains.
Revised: Ahn & Chewi (2020) study its discretization for MCMC sampling in constrained domains.

Original: Relatedly,
Revised: Relatedly,

Original: Patterson & Teh (2013) proposed stochastic Riemannian LMC for sampling on the simplex.
Revised: Patterson & Teh (2013) proposed stochastic Riemannian LMC for sampling on the simplex.

Original: Several modiﬁcations of SVGD have been proposed to incorporate geometric information.
Revised: Several modiﬁcations of SVGD have been proposed to incorporate geometric information.

Original: 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.
Revised: 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.

Original: They also reportedly do not perform well when with scalable stochastic estimates of ∇ log p .
Revised: They also reportedly do not perform well when with scalable stochastic estimates of ∇ log p .

Original: Stein Variational Newton (SVN, Detommaso et al., 2018; Chen et al., 2019) introduces second-order information into SVGD.
Revised: Stein Variational Newton (SVN, Detommaso et al., 2018; Chen et al., 2019) introduces second-order information into SVGD.

Original: Their algorithm requires an often expensive Hessian computation and need not lead to descent directions, so inexact approximations are employed in practice.
Revised: Their algorithm requires an often expensive Hessian computation and need not lead to descent directions, so inexact approximations are employed in practice.

Original: 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.
Revised: 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.

Original: 4.4, a choice that is not explored in Wang et al.
Revised: 4.4, a choice that is not explored in Wang et al.

Original: and which recovers natural gradient descent when n = 1 unlike the heuristic kernel constructions of Wang et al.
Revised: and which recovers natural gradient descent when n = 1 unlike the heuristic kernel constructions of Wang et al.

Original: None of the aforementioned works provide convergence guarantees, and neither SVN nor matrix SVGD deals with constrained domains.
Revised: None of the aforementioned works provide convergence guarantees, and neither SVN nor matrix SVGD deals with constrained domains.

Original:  B ACKGROUND : M IRROR D ESCENT AND N ON -E UCLIDEAN G EOMETRY
Revised:  B ACKGROUND : M IRROR D ESCENT AND N ON -E UCLIDEAN G EOMETRY

Original: Standard gradient descent can be viewed as optimizing a local quadratic approximation to the target function f : θ t +1 = argmin θ ∈ Θ ∇ f ( θ [Equation] ) (cid:62) θ + 12 (cid:15) t (cid:107) θ − θ t (cid:107) 22 .
Revised: Standard gradient descent can be viewed as optimizing a local quadratic approximation to the target function f : θ t +1 = argmin θ ∈ Θ ∇ f ( θ t ) (cid:62) θ + 12 (cid:15) t (cid:107) θ − θ t (cid:107) 22 .

Original: When Θ ⊆ R d is constrained, it can be advantageous to replace (cid:107) · (cid:107) 2 with a function Ψ that reﬂects the geometry of a problem (Nemirovskij & Yudin, 1983; Beck & Teboulle, 2003):
Revised: When Θ ⊆ R d is constrained, it can be advantageous to replace (cid:107) · (cid:107) 2 with a function Ψ that reﬂects the geometry of a problem (Nemirovskij & Yudin, 1983; Beck & Teboulle, 2003):

Original: [Equation]
Revised: [Equation]

Original: 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 smoothfunction ψ : Θ → R ∪ {∞} : Ψ( θ, θ (cid:48) )
Revised: 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 smoothfunction ψ : Θ → R ∪ {∞} : Ψ( θ, θ (cid:48) )

Original: = ψ ( θ ) − ψ ( θ (cid:48) )
Revised: = ψ ( θ ) − ψ ( θ (cid:48) )

Original: − ∇ ψ ( θ (cid:48) )
Revised: − ∇ ψ ( θ (cid:48) )

Original: (cid:62) ( θ − θ (cid:48) ) .
Revised: (cid:62) ( θ − θ (cid:48) ) .

Original: When Θ is a ( d + 1) - simplex { θ : (cid:80) di
Revised: When Θ is a ( d + 1) - simplex { θ : (cid:80) di

Original: =1 θ i ≤ 1 and θ i ≥ 0 for i ∈
Revised: =1 θ i ≤ 1 and θ i ≥ 0 for i ∈

Original: [ d ] } , a common choice of ψ is the negative entropy ψ ( θ )
Revised: [ d ] } , a common choice of ψ is the negative entropy ψ ( θ )

Original: = (cid:80) d
Revised: = (cid:80) d

Original: +1 i
Revised: +1 i

Original: =
Revised: =

Original: θ
Revised: θ

Original: i log θ i , for θ d +1 (cid:44) 1 − (cid:80) di
Revised: i log θ i , for θ d +1 (cid:44) 1 − (cid:80) di

Original: =1 θ d .
Revised: =1 θ d .

Original: The solution of (1) is given by
Revised: The solution of (1) is given by

Original: [Equation]
Revised: [Equation]

Original:  where ψ ∗ ( η ) (cid:44) sup θ ∈ 
Revised: ) where ψ ∗ ( η ) (cid:44) sup θ ∈ Θ

Original: Θ η (cid:62) θ − ψ ( θ ) is the convex conjugate of ψ and ∇ ψ is a bijection from Θ to dom( ψ ∗ ) with inverse map ( ∇ ψ ) − 1 = ∇ ψ ∗ .
Revised: η (cid:62) θ − ψ ( θ ) is the convex conjugate of ψ and ∇ ψ is a bijection from Θ to dom( ψ ∗ ) with inverse map ( ∇ ψ ) − 1 = ∇ ψ ∗ .

Original: We can view the update in (2) as ﬁrst mapping θ t to η t by ∇ ψ , applying the update η t +1 =
Revised: We can view the update in (2) as ﬁrst mapping θ t to η t by ∇ ψ , applying the update η t +1 =

Original: η t − (cid:15) t ∇ f ( θ t ) , and mapping back through θ t +1 = ∇ ψ ∗ ( η t +1 ) .
Revised: η t − (cid:15) t ∇ f ( θ t ) , and mapping back through θ t +1 = ∇ ψ ∗ ( η t +1 ) .

Original: 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 ﬂow (see App.
Revised: 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 ﬂow (see App.

Original: A):
Revised: A):

Original: [Equation]
Revised: [Equation]

Original: or 
Revised: or equivalently

Original: [Equation]
Revised: [Equation]

Original: where ∇ 2 ψ ( θ ) and ∇ 2 ψ ∗ ( η ) are Riemannian metric tensors.
Revised: where ∇ 2 ψ ( θ ) and ∇ 2 ψ ∗ ( η ) are Riemannian metric tensors.

Original: In information geometry, the discretization of (3) is known as natural gradient descent (Amari, 1998).
Revised: In information geometry, the discretization of (3) is known as natural gradient descent (Amari, 1998).

Original: There is considerable theoretical and practical evidence (Martens, 2014) showing that natural gradient works efﬁciently in learning.
Revised: There is considerable theoretical and practical evidence (Martens, 2014) showing that natural gradient works efﬁciently in learning.

Original:  S TEIN ’ S I DENTITY AND M IRRORED S TEIN O PERATORS
Revised:  S TEIN ’ S I DENTITY AND M IRRORED S TEIN O PERATORS

Original: Stein’s identity (Stein, 1972) is a tool for characterizing a target distribution P using a so-called Stein operator .
Revised: Stein’s identity (Stein, 1972) is a tool for characterizing a target distribution P using a so-called Stein operator .

Original: We assume P has a differentiable density p with a closed convex support Θ ⊆ R d .
Revised: We assume P has a differentiable density p with a closed convex support Θ ⊆ R d .

Original: A Stein operator S p takes as input functions g from a Stein set G and outputs mean-zero functions under p :
Revised: A Stein operator S p takes as input functions g from a Stein set G and outputs mean-zero functions under p :

Original: [Equation]
Revised: [Equation]

Original: for all g ∈ G 
Revised: for all g ∈ G .

Original: [Equation]

Original: Gorham & Mackey (2015) proposed the Langevin Stein operator given by
Revised: Gorham & Mackey (2015) proposed the Langevin Stein operator given by

Original: [Equation]
Revised: [Equation]

Original: (5) where g is a vector-valued function and ∇ · g is its divergence.
Revised: (5) where g is a vector-valued function and ∇ · g is its divergence.

Original: For an unconstrained domain with E p [ (cid:107)∇ log p ( θ ) (cid:107) 2 ] < ∞ , Stein’s identity (4) holds for this operator whenever g ∈ C 1 is bounded and Lipschitz by (Gorham et al., 2019, proof of Prop. 3).
Revised: For an unconstrained domain with E p [ (cid:107)∇ log p ( θ ) (cid:107) 2 ] < ∞ , Stein’s identity (4) holds for this operator whenever g ∈ C 1 is bounded and Lipschitz by (Gorham et al., 2019, proof of Prop. 3).

Original: 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.
Revised: 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.

Original: Motivated by this deﬁciency and by a desire to exploit non-Euclidean geometry, we propose an alternative mirrored Stein operator ,
Revised: Motivated by this deﬁciency and by a desire to exploit non-Euclidean geometry, we propose an alternative mirrored Stein operator ,

Original: [Equation]
Revised: [Equation]

Original: where ψ is a strongly convex, essentially smooth function as in Sec.
Revised: where ψ is a strongly convex, essentially smooth function as in Sec.

Original: 2 with ( ∇ 2 ψ ) − 1 differentiable and Lipschitz on Θ .
Revised: 2 with ( ∇ 2 ψ ) − 1 differentiable and Lipschitz on Θ .

Original: We derive this operator from the (inﬁnitesimal) generator of the mirror-Langevin diffusion (19) in App.
Revised: We derive this operator from the (inﬁnitesimal) generator of the mirror-Langevin diffusion (19) in App.

Original: C.
Revised: C.

Original: The following result, proved in App.
Revised: The following result, proved in App.

Original: I.1, shows that M p,ψ generates mean-zero functions under p whenever ∇ 2 ψ − 1 suitably cancels the growth of p at the boundary.
Revised: I.1, shows that M p,ψ generates mean-zero functions under p whenever ∇ 2 ψ − 1 suitably cancels the growth of p at the boundary.

Original: Proposition 1.
Revised: Proposition 1.

Original: Suppose that ∇ 2 ψ ( θ ) − 1 ∇ log p ( θ ) and ∇ · ∇ 2 ψ ( θ ) − 1 are p -integrable.
Revised: Suppose that ∇ 2 ψ ( θ ) − 1 ∇ log p ( θ ) and ∇ · ∇ 2 ψ ( θ ) − 1 are p -integrable.

Original: Iflim
Revised: If

Original: [Equation]
Revised: [Equation]

Original: ψ
Revised: n

Original: [Equation]
Revised: [Equation]

Original: θ )
Revised: θ )

Original: [Equation]
Revised: [Equation]

Original: 1 n r ( θ ) (cid:107) 2 dθ = 0 for Θ r (cid:44) { θ ∈ Θ : (cid:107) θ (cid:107) ∞ ≤ r } and n r ( θ ) the outward unit normal vector 2 to ∂ Θ r at θ , then E p [( M p,ψ g )
Revised: 2 dθ = 0 for Θ r (cid:44) { θ ∈ Θ : (cid:107) θ (cid:107) ∞ ≤ r } and n r ( θ ) the outward unit normal vector 2 to ∂ Θ r at θ , then E p [( M p,ψ g )

Original: ( θ )]
Revised: ( θ )]

Original: = 0 if g ∈ C 1 is bounded Lipschitz.
Revised: = 0 if g ∈ C 1 is bounded Lipschitz.

Original: Example 1 ( Dirichlet p , Negative entropy ψ ) .
Revised: Example 1 ( Dirichlet p , Negative entropy ψ ) .

Original: When θ 1: d +1 ∼ Dir( α ) for α ∈ R d +1+ , even setting g ( θ ) = 1 in (5) need not cause the identity to hold when some α j ≤ 1 .
Revised: When θ 1: d +1 ∼ Dir( α ) for α ∈ R d +1+ , even setting g ( θ ) = 1 in (5) need not cause the identity to hold when some α j ≤ 1 .

Original: However, when ψ ( θ ) = (cid:80) d +1 j =1 θ j log θ j , we show in App.
Revised: However, when ψ ( θ ) = (cid:80) d +1 j =1 θ j log θ j , we show in App.

Original: B that the conditions of Prop. 1 are met for any α .
Revised: B that the conditions of Prop. 1 are met for any α .

Original: Remarkably, the mirror-Langevin diffusion for our choice of ψ is the Wright-Fisher diffusion (Ethier, 1976) which Gan et al.
Revised: Remarkably, the mirror-Langevin diffusion for our choice of ψ is the Wright-Fisher diffusion (Ethier, 1976) which Gan et al.

Original: (2017) recently used to bound distances to Dirichlet distributions.
Revised: (2017) recently used to bound distances to Dirichlet distributions.

Original:  S AMPLING WITH M IRRORED S TEIN O PERATORS
Revised:  S AMPLING WITH M IRRORED S TEIN O PERATORS

Original: Liu & Wang (2016) pioneered the idea of using Stein operators to approximate a target distribution with particles.
Revised: Liu & Wang (2016) pioneered the idea of using Stein operators to approximate a target distribution with particles.

Original: Their popular SVGD algorithm updates each particle in its approximation by applying the update rule θ t +

Revised: Their popular SVGD algorithm updates each particle in its approximation by applying the update rule θ t +

Original: =
Revised: =

Original: θ t + (cid:15)
Revised: θ t + (cid:15)

Original: t g t ( θ t ) for a chosen mapping g t : R d → R d .
Revised: t g t ( θ t ) for a chosen mapping g t : R d → R d .

Original: Speciﬁcally, SVGD chooses the mapping g ∗ t that leads to the largest decrease in KL divergence to p in the limit as (cid:15) t → 0 .
Revised: Speciﬁcally, SVGD chooses the mapping g ∗ t that leads to the largest decrease in KL divergence to p in the limit as (cid:15) t → 0 .

Original: The following theorem summarizes their main ﬁndings.
Revised: The following theorem summarizes their main ﬁndings.

Original: Theorem 2 (Liu & Wang, 2016, Thm. 3.1) .
Revised: Theorem 2 (Liu & Wang, 2016, Thm. 3.1) .

Original: Suppose ( θ t ) t ≥ 0 satisﬁes dθ t
Revised: Suppose ( θ t ) t ≥ 0 satisﬁes dθ t

Original: =
Revised: =

Original: g t ( θ t ) dt for bounded
Revised: g t ( θ t ) dt for bounded

Original: Lipschitz g t ∈ C 1 : R d → R d and that θ t has density q t with E q t
Revised: Lipschitz g t ∈ C 1 : R d → R d and that θ t has density q t with E q t

Original: (cid:107)∇ log q t (cid:107) 2 ] < ∞ .
Revised: (cid:107)∇ log q t (cid:107) 2 ] < ∞ .

Original: If KL ( q t (cid:107) p ) (cid:44) E q t [log( q t /p )] exists then, for the Langevin Stein operator S p (5) ,
Revised: If KL ( q t (cid:107) p ) (cid:44) E q t [log( q t /p )] exists then, for the Langevin Stein operator S p (5) ,

Original: [Equation]
Revised: [Equation]

Original: Input: density p on Θ , kernel k , mirror function ψ , particles ( θ i 0 ) ni
Revised: Input: density p on Θ , kernel k , mirror function ψ , particles ( θ i 0 ) ni

Original: =1 ⊂ Θ , step sizes ( (cid:15) t )
Revised: =1 ⊂ Θ , step sizes ( (cid:15) t )

Original: Tt =
Revised: Tt =

Original: Init: η i 0 ← ∇ ψ ( θ i 0 ) for i ∈ [ n ] for t = 

Revised: Init: η i 0 ← ∇ ψ ( θ i 0 ) for i ∈ [ n ] for t =

Original: T do if SVMD then K t ← K ψ,t (13) else K t ← kI (MSVGD) for i ∈ [ n ] , η it +1 ← η it + (cid:15) t 1 n (cid:80) nj =1 M p,ψ K t ( θ it , θ jt ) (for M p,ψ K t ( · , θ ) deﬁned in Thm. 3) for i ∈ [ n ] , θ it +1 ← ∇ ψ ∗ ( η it +1 ) return { θ iT +1 } ni =1 .
Revised: T do if SVMD then K t ← K ψ,t (13) else K t ← kI (MSVGD) for i ∈ [ n ] , η it +1 ← η it + (cid:15) t 1 n (cid:80) nj =1 M p,ψ K t ( θ it , θ jt ) (for M p,ψ K t ( · , θ ) deﬁned in Thm. 3) for i ∈ [ n ] , θ it +1 ← ∇ ψ ∗ ( η it +1 ) return { θ iT +1 } ni =1 .

Original: To improve its current particle approximation, SVGD ﬁnds the choice of g t that most quickly decreases KL ( q t (cid:107) p ) at time t , i.e., it minimizes ddt KL ( q t (cid:107) p ) over a set of candidate directions g t .
Revised: To improve its current particle approximation, SVGD ﬁnds the choice of g t that most quickly decreases KL ( q t (cid:107) p ) at time t , i.e., it minimizes ddt KL ( q t (cid:107) p ) over a set of candidate directions g t .

Original: SVGD ﬁnds g t in a reproducing kernel Hilbert space (RKHS, Berlinet & Thomas-Agnan, 2011) norm ball B H d
Revised: SVGD ﬁnds g t in a reproducing kernel Hilbert space (RKHS, Berlinet & Thomas-Agnan, 2011) norm ball B H d

Original: = { g : (cid:107) g (cid:107) H d ≤ 1 } , where H d is the product RKHS containing vector-valued functions with each component in the RKHS
Revised: = { g : (cid:107) g (cid:107) H d ≤ 1 } , where H d is the product RKHS containing vector-valued functions with each component in the RKHS

Original: H of k .
Revised: H of k .

Original: Then the optimal g ∗ t ∈ B H d that minimizes (7) is
Revised: Then the optimal g ∗ t ∈ B H d that minimizes (7) is

Original: [Equation]
Revised: [Equation]

Original: where we let K k ( θ, θ (cid:48) )
Revised: where we let K k ( θ, θ (cid:48) )

Original: =
Revised: =

Original: k ( θ, θ (cid:48) )
Revised: k ( θ, θ (cid:48) )

Original: I , and S p K k ( · , θ ) denotes applying S p to each row of K k ( · , θ ) .
Revised: I , and S p K k ( · , θ ) denotes applying S p to each row of K k ( · , θ ) .

Original: SVGD has found great success in approximating unconstrained target distributions p but breaks down for constrained targets and fails to exploit non-Euclidean geometry.
Revised: SVGD has found great success in approximating unconstrained target distributions p but breaks down for constrained targets and fails to exploit non-Euclidean geometry.

Original: Our goal is to develop new particle evolution samplers suitable for constrained domains and non-Euclidean geometries.
Revised: Our goal is to develop new particle evolution samplers suitable for constrained domains and non-Euclidean geometries.

Original: 4. M IRRORED DYNAMICS
Revised: 4. M IRRORED DYNAMICS

Original: SVGD encounters two difﬁculties when faced with a constrained support.
Revised: SVGD encounters two difﬁculties when faced with a constrained support.

Original: First, the SVGD updates can push the random variable θ t outside of its support Θ , rendering all future updates undeﬁned.
Revised: First, the SVGD updates can push the random variable θ t outside of its support Θ , rendering all future updates undeﬁned.

Original: Second,
Revised: Second,

Original: Stein’s identity (4) often fails to hold for candidate directions in B H d (cf.
Revised: Stein’s identity (4) often fails to hold for candidate directions in B H d (cf.

Original: Ex. 1).
Revised: Ex. 1).

Original: When this occurs,
Revised: When this occurs,

Original: SVGD need not converge to p as p is not a stationary point of its dynamics (i.e., ddt KL ( q t (cid:107) p )
Revised: SVGD need not converge to p as p is not a stationary point of its dynamics (i.e., ddt KL ( q t (cid:107) p )

Original: | q t = p (cid:54) =when q t = p ).
Revised: | q t = p (cid:54) =when q t = p ).

Original: Inspired by mirror descent (Nemirovskij & Yudin, 1983), we consider the following mirrored dynamics
Revised: Inspired by mirror descent (Nemirovskij & Yudin, 1983), we consider the following mirrored dynamics

Original: [Equation]
Revised: [Equation]

Original: where g t : Θ → R d now represents the update direction in η space.
Revised: where g t : Θ → R d now represents the update direction in η space.

Original: The inverse mirror map ∇ ψ ∗ automatically guarantees that θ t belongs to the constrained domain Θ .
Revised: The inverse mirror map ∇ ψ ∗ automatically guarantees that θ t belongs to the constrained domain Θ .

Original: Since ψ is strongly convex and ∇ 2 ψ − 1 is bounded Lipschitz, from Thm.
Revised: Since ψ is strongly convex and ∇ 2 ψ − 1 is bounded Lipschitz, from Thm.

Original: 2 it follows for any bounded Lipschitz g t that
Revised: 2 it follows for any bounded Lipschitz g t that

Original: [Equation]
Revised: [Equation]

Original: where M p,ψ is the mirrored Stein operator (6).
Revised: where M p,ψ is the mirrored Stein operator (6).

Original: In the following sections, we propose three new deterministic sampling algorithms by seeking the optimal direction g t that minimizes (9) over different function classes.
Revised: In the following sections, we propose three new deterministic sampling algorithms by seeking the optimal direction g t that minimizes (9) over different function classes.

Original: Thm. 3 (proved in App. I.2) forms the basis of our analysis.
Revised: Thm. 3 (proved in App. I.2) forms the basis of our analysis.

Original: Theorem 3 (Optimal mirror updates in RKHS) .
Revised: Theorem 3 (Optimal mirror updates in RKHS) .

Original: Suppose ( θ t ) t ≥ 0 follows the mirrored dynamics (8) .
Revised: Suppose ( θ t ) t ≥ 0 follows the mirrored dynamics (8) .

Original: Let H K denote the RKHS of a matrix-valued kernel K : Θ × Θ → S d × d (Micchelli & Pontil, 2005).
Revised: Let H K denote the RKHS of a matrix-valued kernel K : Θ × Θ → S d × d (Micchelli & Pontil, 2005).

Original: Then, the optimal direction of g t that minimizes (9) in the norm ball B H K (cid:44) { g : (cid:107) g (cid:107) H K ≤ 1 } is
Revised: Then, the optimal direction of g t that minimizes (9) in the norm ball B H K (cid:44) { g : (cid:107) g (cid:107) H K ≤ 1 } is

Original: [Equation]
Revised: [Equation]

Original: where M p,ψ K ( · , θ ) applies M p,ψ (6) to each row of the matrix-valued function K θ =
Revised: where M p,ψ K ( · , θ ) applies M p,ψ (6) to each row of the matrix-valued function K θ =

Original: K ( · , θ ) .
Revised: K ( · , θ ) .

Original: 4. M IRRORED S TEIN V ARIATIONAL G RADIENT D ESCENT
Revised: 4. M IRRORED S TEIN V ARIATIONAL G RADIENT D ESCENT

Original: Following the pattern of SVGD, one can choose the K of Thm.
Revised: Following the pattern of SVGD, one can choose the K of Thm.

Original: 3 to be K k
Revised: 3 to be K k

Original: ( θ, θ (cid:48) )
Revised: ( θ, θ (cid:48) )

Original: =
Revised: =

Original: k ( θ, θ (cid:48) )
Revised: k ( θ, θ (cid:48) )

Original: I , where k is any scalar-valued kernel.
Revised: I , where k is any scalar-valued kernel.

Original: In this case, the resulting update g ∗ q t ,K k ( · ) =
Revised: In this case, the resulting update g ∗ q t ,K k ( · ) =

Original: E θ t ∼ q t [ M p,ψ K k ( · , θ t )] is equivalent to running SVGD in the dual η space before mapping back to Θ .
Revised: E θ t ∼ q t [ M p,ψ K k ( · , θ t )] is equivalent to running SVGD in the dual η space before mapping back to Θ .

Original: Theorem 4 (Mirrored SVGD updates) .
Revised: Theorem 4 (Mirrored SVGD updates) .

Original: In the setting of Thm. 3, let k ψ ( η, η (cid:48) )
Revised: In the setting of Thm. 3, let k ψ ( η, η (cid:48) )

Original: =
Revised: =

Original: k ( ∇ ψ ∗ ( η ) , ∇ ψ ∗ ( η (cid:48) )) , p H ( η ) = p ( ∇ ψ ∗ ( η )) · | det ∇ 2 ψ ∗ ( η ) | denote the density of η = ∇ ψ ( θ ) when θ ∼ p , and q t,H denote the distribution of η t under the mirrored dynamics (8) .
Revised: k ( ∇ ψ ∗ ( η ) , ∇ ψ ∗ ( η (cid:48) )) , p H ( η ) = p ( ∇ ψ ∗ ( η )) · | det ∇ 2 ψ ∗ ( η ) | denote the density of η = ∇ ψ ( θ ) when θ ∼ p , and q t,H denote the distribution of η t under the mirrored dynamics (8) .

Original: If K k = kI , g ∗ q t ,K k ( θ (cid:48) ) =E η t ∼ q
Revised: If K k = kI , g ∗ q t ,K k ( θ (cid:48) ) =

Original: [Equation] H [ k ψ ( η [Equation] , η (cid:48) ) ∇ log 
Revised: E η t ∼ q t,H [ k ψ ( η t , η (cid:48) ) ∇ log p H ( η t ) + ∇ η t k ψ ( η

Original: [Equation]
Revised: [Equation]

Original: η 
Revised: η (cid:48) )

Original: [Equation] + ∇ [Equation] k
Revised: ] ∀ θ (cid:48) ∈ Θ

Original: [Equation]
Revised: [Equation]

Original: η [Equation] )] ∀ [Equation]
Revised: η (cid:48)

Original: = ∇ ψ ( [Equation] [Equation]
Revised: = ∇ ψ ( θ (cid:48) ) .

Original: The proof is in App.
Revised: The proof is in App.

Original: I.3.
Revised: I.3.

Original: By discretizing the dynamics dη t =
Revised: By discretizing the dynamics dη t =

Original: g ∗ q t ,K k ( θ t ) dt and initializing with any particle approximation q 0 = 1 n (cid:80) ni
Revised: g ∗ q t ,K k ( θ t ) dt and initializing with any particle approximation q 0 = 1 n (cid:80) ni

Original: =

Revised: =

Original: δ θ i 0 , we obtain Mirrored SVGD (MSVGD) , our ﬁrst algorithm for sampling in constrained domains.
Revised: δ θ i 0 , we obtain Mirrored SVGD (MSVGD) , our ﬁrst algorithm for sampling in constrained domains.

Original: The details are summarized in Alg. 1.
Revised: The details are summarized in Alg. 1.

Original: When only a single particle is used ( n = 1 ) and the differentiable input kernel satisﬁes ∇ k
Revised: When only a single particle is used ( n = 1 ) and the differentiable input kernel satisﬁes ∇ k

Original: ( θ, θ ) = 0 , the MSVGD update (11) reduces to gradient descent on − log p H ( η ) .
Revised: ( θ, θ ) = 0 , the MSVGD update (11) reduces to gradient descent on − log p H ( η ) .

Original: Note however that the modes of the mirrored density p H ( η ) need not match those of the target density p ( θ ) (see the examples in
Revised: Note however that the modes of the mirrored density p H ( η ) need not match those of the target density p ( θ ) (see the examples in

Original: App. E).
Revised: App. E).

Original: Since we are primarily interested in the θ -space density, it is natural to ask whether there exists a mirrored dynamics that reduces to ﬁnding the mode of p ( θ ) in this limiting case.
Revised: Since we are primarily interested in the θ -space density, it is natural to ask whether there exists a mirrored dynamics that reduces to ﬁnding the mode of p ( θ ) in this limiting case.

Original: In the next section, we give an answer to this question by designing an adaptive reproducing kernel that yields a mirror descent-like update.
Revised: In the next section, we give an answer to this question by designing an adaptive reproducing kernel that yields a mirror descent-like update.

Original: 4. S TEIN V ARIATIONAL M IRROR D ESCENT
Revised: 4. S TEIN V ARIATIONAL M IRROR D ESCENT

Original: Our second sampling algorithm for constrained problems is called Stein Variational Mirror Descent (SVMD) .
Revised: Our second sampling algorithm for constrained problems is called Stein Variational Mirror Descent (SVMD) .

Original: We start by introducing a new matrix-valued kernel that incorporates the metric ∇ 2 ψ and evolves with the distribution q t .
Revised: We start by introducing a new matrix-valued kernel that incorporates the metric ∇ 2 ψ and evolves with the distribution q t .

Original: Deﬁnition 1 (Kernels for SVMD) .
Revised: Deﬁnition 1 (Kernels for SVMD) .

Original: Given a continuous scalar-valued kernel k , consider the Mercer representation 3 k ( θ, θ (cid:48) )
Revised: Given a continuous scalar-valued kernel k , consider the Mercer representation 3 k ( θ, θ (cid:48) )

Original: = (cid:80) i ≥ 1 λ t,i u t,i ( θ ) u t,i ( θ (cid:48) ) w.r.t.
Revised: = (cid:80) i ≥ 1 λ t,i u t,i ( θ ) u t,i ( θ (cid:48) ) w.r.t.

Original: q t , where u t,i is an eigenfunction satisfying
Revised: q t , where u t,i is an eigenfunction satisfying

Original: [Equation]
Revised: [Equation]

Original: For k 1 / 2 t ( θ, θ (cid:48) ) (cid:44) (cid:80) i ≥ 1 λ 1 / 2 t [Equation] i u t,i ( θ ) u t,i ( θ
Revised: For k 1 / 2 t ( θ, θ (cid:48) ) (cid:44) (cid:80) i ≥ 1 λ 1 / 2 t,i u t,i ( θ ) u t,i ( θ

Original: [Equation]
Revised: [Equation]

Original: ) , we deﬁne the adaptive SVMD kernel at time t ,
Revised: ) , we deﬁne the adaptive SVMD kernel at time t ,

Original: [Equation]
Revised: [Equation]

Original: By Thm. 3, the optimal update direction for the SVMD kernel ball is g ∗ q t ,K ψ,t =
Revised: By Thm. 3, the optimal update direction for the SVMD kernel ball is g ∗ q t ,K ψ,t =

Original: E q t [ M p,ψ K ψ,t ( · , θ t )] .
Revised: E q t [ M p,ψ K ψ,t ( · , θ t )] .

Original: We obtain the SVMD algorithm (summarized in Alg. 1) by discretizing dη t =
Revised: We obtain the SVMD algorithm (summarized in Alg. 1) by discretizing dη t =

Original: g ∗ q t ,K ψ,t ( θ t ) dt and initializing with q 0 = 1 n (cid:80)
Revised: g ∗ q t ,K ψ,t ( θ t ) dt and initializing with q 0 = 1 n (cid:80)

Original: ni
Revised: ni

Original: =

Revised: =

Original: δ θ i 0 .
Revised: δ θ i 0 .

Original: Because of the discrete representation of q t , K ψ,t takes the form
Revised: Because of the discrete representation of q t , K ψ,t takes the form

Original: [Equation]
Revised: [Equation]

Original: Here both λ t,j and u t,j ( θ it ) can be computed by solving a matrix eigenvalue problem involving the particle set { θ it } ni =1 : B t v t,j = nλ t,j v t,j , where B t =
Revised: Here both λ t,j and u t,j ( θ it ) can be computed by solving a matrix eigenvalue problem involving the particle set { θ it } ni =1 : B t v t,j = nλ t,j v t,j , where B t =

Original: ( k ( θ it , θ jt ))
Revised: ( k ( θ it , θ jt ))

Original: ni,j =1 ∈ R n × n is the Gram matrix of pairwise kernel evaluations at particle locations, and the i -th element of v t,j is u t,j ( θ it ) .
Revised: ni,j =1 ∈ R n × n is the Gram matrix of pairwise kernel evaluations at particle locations, and the i -th element of v t,j is u t,j ( θ it ) .

Original: To compute ∇ θ K ψ,t ( θ, θ (cid:48) ) , we differentiate both sides of (12) to ﬁnd that ∇ u 
Revised: To compute ∇ θ K ψ,t ( θ, θ (cid:48) ) , we differentiate both sides of (12) to ﬁnd that ∇ u t

Original: [Equation]
Revised: [Equation]

Original: This technique was used in Shi et al.
Revised: This technique was used in Shi et al.

Original: (2018) to estimate gradients of eigenfunctions w.r.t.
Revised: (2018) to estimate gradients of eigenfunctions w.r.t.

Original: a continuous q .
Revised: a continuous q .

Original: Following their recommendations, we truncate the sum at the J -th largest eigenvalues according to a threshold ( τ ≥ (cid:80)
Revised: Following their recommendations, we truncate the sum at the J -th largest eigenvalues according to a threshold ( τ ≥ (cid:80)

Original: Jj =1 λ t,j / (cid:80) nj =1 λ t,j ) to ensure numerical stability.
Revised: Jj =1 λ t,j / (cid:80) nj =1 λ t,j ) to ensure numerical stability.

Original: Notably, SVMD differs from MSVGD only in its choice of kernel, but, whenever ∇ k ( θ, θ ) = 0 , this change is sufﬁcient to exactly recover mirror descent when n = 1 .
Revised: Notably, SVMD differs from MSVGD only in its choice of kernel, but, whenever ∇ k ( θ, θ ) = 0 , this change is sufﬁcient to exactly recover mirror descent when n = 1 .

Original: Proposition 5 (Single-particle SVMD is mirror descent) .
Revised: Proposition 5 (Single-particle SVMD is mirror descent) .

Original: If n = 1 , then one step of SVMD becomes
Revised: If n = 1 , then one step of SVMD becomes

Original: [Equation]
Revised: [Equation]

Original: 4. S TEIN V ARIATIONAL N ATURAL G RADIENT
Revised: 4. S TEIN V ARIATIONAL N ATURAL G RADIENT

Original: The fact that SVMD recovers mirror descent as a special case is not only of relevance in constrained problems.
Revised: The fact that SVMD recovers mirror descent as a special case is not only of relevance in constrained problems.

Original: We next exploit the connection between MD and natural gradient descent discussed in
Revised: We next exploit the connection between MD and natural gradient descent discussed in

Original: Sec. 2 to design a new sampler – Stein Variational Natural Gradient (SVNG) – that more efﬁciently approximates unconstrained targets.
Revised: Sec. 2 to design a new sampler – Stein Variational Natural Gradient (SVNG) – that more efﬁciently approximates unconstrained targets.

Original: The idea is to replace the Hessian ∇ 2 ψ ( · ) in the SVMD dynamics dθ t = ∇ 2 ψ ( θ t ) − 1 g ∗ q t ,K ψ,t ( θ t ) with a general metric tensor G ( · ) .
Revised: The idea is to replace the Hessian ∇ 2 ψ ( · ) in the SVMD dynamics dθ t = ∇ 2 ψ ( θ t ) − 1 g ∗ q t ,K ψ,t ( θ t ) with a general metric tensor G ( · ) .

Original: The result is the Riemannian gradient ﬂowdθ
Revised: The result is the Riemannian gradient ﬂow

Original: [Equation]
Revised: [Equation]

Original: t ( θ t ) dt with K G,t ( θ, θ (cid:48) )

Original: (cid:44) E

Original: [Equation]

Original: ∼

Original: [Equation] [ k [Equation] θ
Revised: dt with K

Original: [Equation]
Revised: [Equation]

Original: θ
Revised: θ,

Original: [Equation]
Revised: [Equation]

Original: θ

Original: [Equation]
Revised: [Figure]

Original: θ [Equation] [Figure] Given any initial particle approximation q 0 = 1 n (cid:80) ni =1 δ θ i 0 , we discretize these dynamics to obtain the unconstrained SVNG sampler of Alg. 2 in the appendix.
Revised: Given any initial particle approximation q 0 = 1 n (cid:80) ni =1 δ θ i 0 , we discretize these dynamics to obtain the unconstrained SVNG sampler of Alg. 2 in the appendix.

Original: SVNG can be seen as an instance of MatSVGD (Wang et al., 2019) with a new adaptive time-dependent kernel
Revised: SVNG can be seen as an instance of MatSVGD (Wang et al., 2019) with a new adaptive time-dependent kernel

Original: G − 1 ( θ ) K G,t ( θ, θ (cid:48) )
Revised: G − 1 ( θ ) K G,t ( θ, θ (cid:48) )

Original: G − 1 ( θ (cid:48) ) .
Revised: G − 1 ( θ (cid:48) ) .

Original: However, similar to Prop. 5 and unlike the heuristic kernels of Wang et al.
Revised: However, similar to Prop. 5 and unlike the heuristic kernels of Wang et al.

Original: (2019), SVNG reduces to natural gradient ascent for ﬁnding the mode of p ( θ ) when n = 1 .
Revised: (2019), SVNG reduces to natural gradient ascent for ﬁnding the mode of p ( θ ) when n = 1 .

Original: SVNG is well-suited to Bayesian inference problems where the target is a posterior distribution p ( θ ) ∝ π ( θ ) π
Revised: SVNG is well-suited to Bayesian inference problems where the target is a posterior distribution p ( θ ) ∝ π ( θ ) π

Original: ( y | θ ) .
Revised: ( y | θ ) .

Original: There, the metric tensor G ( θ ) can be set to the Fisher information matrix E π ( y | θ ) [ ∇ log π ( y | θ ) ∇ log π ( y | θ ) (cid:62) ] of the data likelihood π ( y | θ ) .
Revised: There, the metric tensor G ( θ ) can be set to the Fisher information matrix E π ( y | θ ) [ ∇ log π ( y | θ ) ∇ log π ( y | θ ) (cid:62) ] of the data likelihood π ( y | θ ) .

Original: Ample precedent from natu- ral gradient variational inference (Hoffman et al., 2013; Khan & Nielsen, 2018) and Riemannian
Revised: Ample precedent from natu- ral gradient variational inference (Hoffman et al., 2013; Khan & Nielsen, 2018) and Riemannian

Original: MCMC (Patterson & Teh, 2013) suggests that encoding problem geometry in this manner often leads to more rapid convergence.
Revised: MCMC (Patterson & Teh, 2013) suggests that encoding problem geometry in this manner often leads to more rapid convergence.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We next conduct a series of simulated and real-data experiments to assess (1) distributional approximation on the simplex, (2) frequentist conﬁdence interval construction for (constrained) post-selection inference, and (3) large-scale posterior inference with non-Euclidean geometry.
Revised: We next conduct a series of simulated and real-data experiments to assess (1) distributional approximation on the simplex, (2) frequentist conﬁdence interval construction for (constrained) post-selection inference, and (3) large-scale posterior inference with non-Euclidean geometry.

Original: 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.
Revised: 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.

Revised: For SVMD, we need to solve an eigenvalue problem, which costs O ( n 3 ) time.

Revised: In practice the number of particles used for particle evolution algorithms is relatively small, even for SVGD, due to the O ( n 2 ) cost of updates.

Revised: We have produced a practical SVMD implementation that is computationally competitive with MSVGD and

Revised: SVGD for standard particle counts like n = 50 .

Revised: We sticked to n = 50 in all experiments.

Original: 5. A PPROXIMATION QUALITY ON THE SIMPLEX
Revised: 5. A PPROXIMATION QUALITY ON THE SIMPLEX

Original: We ﬁrst 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).
Revised: We ﬁrst 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).

Original: 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.
Revised: 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.

Original: In Fig.
Revised: In Fig.

Original: 2, we compare the quality of
Revised: 2, we compare the quality of

Original: MSVGD, SVMD, and projected SVGD with 50 particles and inverse multiquadric kernel k (Gorham &
Revised: MSVGD, SVMD, and projected SVGD with 50 particles and inverse multiquadric kernel k (Gorham &

Original: Mackey, 2017) by computing the energy distance (Sz´ekely & Rizzo, 2013) to a surrogate ground truth sample of size 1000 (drawn i.i.d.
Revised: Mackey, 2017) by computing the energy distance (Sz´ekely & Rizzo, 2013) to a surrogate ground truth sample of size 1000 (drawn i.i.d.

Original: or, in the quadratic case, from the No-U-Turn Sampler (Hoffman &
Revised: or, in the quadratic case, from the No-U-Turn Sampler (Hoffman &

Original: Gelman, 2014)).
Revised: Gelman, 2014)).

Original: We also compare to MSVGD with k 2 ( θ, θ (cid:48) )
Revised: We also compare to MSVGD with k 2 ( θ, θ (cid:48) )

Original: =
Revised: =

Original: k ( ∇ ψ ( θ ) , ∇ ψ ( θ (cid:48) )) , a choice which corresponds to running SVGD in the dual space with kernel k by Thm. 4 and which ensures the convergence of MSVGD to p by the upcoming Thms.
Revised: k ( ∇ ψ ( θ ) , ∇ ψ ( θ (cid:48) )) , a choice which corresponds to running SVGD in the dual space with kernel k by Thm. 4 and which ensures the convergence of MSVGD to p by the upcoming Thms.

Original: 6 to 8.
Revised: 6 to 8.

Original: In the quadratic case, SVMD is favored over MSVGD as it is able to exploit the log-concavity of p ( θ ) .
Revised: In the quadratic case, SVMD is favored over MSVGD as it is able to exploit the log-concavity of p ( θ ) .

Original: 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.
Revised: 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.

Original: This parallels the observation of Hsieh et al.
Revised: This parallels the observation of Hsieh et al.

Original: (2018) that LMC in the mirror space outperforms Riemannian
Revised: (2018) that LMC in the mirror space outperforms Riemannian

Original: LMC for sparse Dirichlet distributions.
Revised: LMC for sparse Dirichlet distributions.

Original: Projected SVGD fails to converge to the target in both cases and has particular difﬁculty in approximating the sparse Dirichlet target with unbounded density.
Revised: Projected SVGD fails to converge to the target in both cases and has particular difﬁculty in approximating the sparse Dirichlet target with unbounded density.

Revised: [Figure]

Original: MSVGD with k and k 2 perform very similarly, but we observe that k yields better approximation quality upon convergence.
Revised: MSVGD with k and k 2 perform very similarly, but we observe that k yields better approximation quality upon convergence.

Original: Therefore, we employ k in the remaining MSVGD experiments.[Figure]
Revised: Therefore, we employ k in the remaining MSVGD experiments.

Original: 5. C ONFIDENCE INTERVALS FOR POST - SELECTION INFERENCE
Revised: 5. C ONFIDENCE INTERVALS FOR POST - SELECTION INFERENCE

Original: We next apply our algorithms to the constrained sampling problems that arise in post-selection inference (Taylor & Tibshirani, 2015; Lee et al., 2016).
Revised: We next apply our algorithms to the constrained sampling problems that arise in post-selection inference (Taylor & Tibshirani, 2015; Lee et al., 2016).

Original: Speciﬁcally, we consider the task of forming valid conﬁdence intervals (CIs) for regression parameters selected using the randomized Lasso (Tian & Taylor, 2018) with data X ∈ R ˜ n × p and y ∈ R ˜ n and user-generated randomness w ∈ R p from a log-concave distribution with density g .
Revised: Speciﬁcally, we consider the task of forming valid conﬁdence intervals (CIs) for regression parameters selected using the randomized Lasso (Tian & Taylor, 2018) with data X ∈ R ˜ n × p and y ∈ R ˜ n and user-generated randomness w ∈ R p from a log-concave distribution with density g .

Original: The randomized Lasso returns ˆ β ∈ R p with non-zero coefﬁcients denoted by ˆ β E and their signs by s E .
Revised: The randomized Lasso returns ˆ β ∈ R p with non-zero coefﬁcients denoted by ˆ β E and their signs by s E .

Original: It is common practice to report least squares CIs for β E by running a linear regression on the selected features E .
Revised: It is common practice to report least squares CIs for β E by running a linear regression on the selected features E .

Original: However, since E is chosen based on the same data, the resulting CIs are often invalid.
Revised: However, since E is chosen based on the same data, the resulting CIs are often invalid.

Original: Post-selection inference solves this problem by conditioning the inference on the knowledge of E and s E .
Revised: Post-selection inference solves this problem by conditioning the inference on the knowledge of E and s E .

Original: To construct valid CIs, it sufﬁces to approximate the selective distribution with support { ˆ β E , u − E : s E (cid:12) ˆ β E > 0 , u − E ∈ [ − 1 , 1] p −| E | } and densityˆ g (ˆ β E , u − E ) ∝ g (cid:0) X (cid:62) y − (cid:0)
Revised: To construct valid CIs, it sufﬁces to approximate the selective distribution with support { ˆ β E , u − E : s E (cid:12) ˆ β E > 0 , u − E ∈ [ − 1 , 1] p −| E | } and density

Original: [Equation]
Revised: [Equation]

Original: In our experiments, we integrate out u − E analytically, following Tian & Taylor (2018), and reparameterize ˆ β E as s E (cid:12) | ˆ β E | to obtain a log-concave density of | ˆ β E | supported on the nonnegative orthant with mirror function ψ ( θ )
Revised: In our experiments, we integrate out u − E analytically, following Tian & Taylor (2018), and reparameterize ˆ β E as s E (cid:12) | ˆ β E | to obtain a log-concave density of | ˆ β E | supported on the nonnegative orthant with mirror function ψ ( θ )

Original: = (cid:80) dj
Revised: = (cid:80) dj

Original: =1 ( θ j log θ j − θ j ) .
Revised: =1 ( θ j log θ j − θ j ) .

Original: In Fig.
Revised: In Fig.

Original: 4a we show the example of a 2D selective distribution using samples drawn by NUTS (Hoffman & Gelman, 2014).
Revised: 4a we show the example of a 2D selective distribution using samples drawn by NUTS (Hoffman & Gelman, 2014).

Original: We also plot the results by projected SVGD, SVMD, and MSVGD in this example.
Revised: We also plot the results by projected SVGD, SVMD, and MSVGD in this example.

Original: Projected SVGD fails to approximate the target with many samples gathering at the truncation boundary, while the samples by MSVGD and
Revised: Projected SVGD fails to approximate the target with many samples gathering at the truncation boundary, while the samples by MSVGD and

Original: SVMD closely resemble the truth.
Revised: SVMD closely resemble the truth.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: As N ranges from 1000 to 3000 in Fig.
Revised: As N ranges from 1000 to 3000 in Fig.

Original: 3a, the MSVGD and SVMD CIs consistently yield higher coverage than the standard 90% CIs.
Revised: 3a, the MSVGD and SVMD CIs consistently yield higher coverage than the standard 90% CIs.

Original: This increased coverage is of particular value for smaller sample sizes, for which the standard CIs tend to undercover.
Revised: This increased coverage is of particular value for smaller sample sizes, for which the standard CIs tend to undercover.

Original: For a much larger sample size of N = 5000 in Fig.
Revised: For a much larger sample size of N = 5000 in Fig.

Original: 3b, the SVMD and standard CIs closely track one another across conﬁdence levels, while MSVGD consistently yields longer CIs with high coverage.
Revised: 3b, the SVMD and standard CIs closely track one another across conﬁdence levels, while MSVGD consistently yields longer CIs with high coverage.

Original: The higher coverage of MSVGD is only of value for larger conﬁdence levels at which the other methods begin to undercover.
Revised: The higher coverage of MSVGD is only of value for larger conﬁdence levels at which the other methods begin to undercover.

Original: 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 ﬁnd statistically signiﬁcant mutations associated with drug resistance using susceptibility data on virus isolates.
Revised: 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 ﬁnd statistically signiﬁcant mutations associated with drug resistance using susceptibility data on virus isolates.

Original: 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.
Revised: 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.

Original: In Fig.
Revised: In Fig.

Original: 4b we plot the CIs of selected mutations obtained with N = 5000 sample points.
Revised: 4b we plot the CIs of selected mutations obtained with N = 5000 sample points.

Original: We see that the invalid unadjusted least squares CIs can lead to premature conclusions, e.g., declaring mutation 215Y signiﬁcant when there is insufﬁcient support after conditioning on the selection event.
Revised: We see that the invalid unadjusted least squares CIs can lead to premature conclusions, e.g., declaring mutation 215Y signiﬁcant when there is insufﬁcient support after conditioning on the selection event.

Original: In contrast, mutation 184V, which has known association with drug resistance, is declared signiﬁcant by all methods even after post-selection adjustment.
Revised: In contrast, mutation 184V, which has known association with drug resistance, is declared signiﬁcant by all methods even after post-selection adjustment.

Original: The MSVGD and SVMD CIs mostly track those of the standard selectiveInference method, but their conclusions sometimes differ: e.g., 62Y is ﬂagged as signiﬁcant by MSVGD and SVMD but not by selectiveInference .
Revised: The MSVGD and SVMD CIs mostly track those of the standard selectiveInference method, but their conclusions sometimes differ: e.g., 62Y is ﬂagged as signiﬁcant by MSVGD and SVMD but not by selectiveInference .

Original: [Figure]
Revised: [Figure]

Original: 5. L ARGE - SCALE POSTERIOR INFERENCE WITH NON -E UCLIDEAN GEOMETRY
Revised: 5. L ARGE - SCALE POSTERIOR INFERENCE WITH NON -E UCLIDEAN GEOMETRY

Original: 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, datapoints and d = 54 feature dimensions.
Revised: 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, datapoints and d = 54 feature dimensions.

Original: Here, the target p is the posterior distribution over logistic regression parameters.
Revised: Here, the target p is the posterior distribution over logistic regression parameters.

Original: 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.
Revised: 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.

Original: In Fig.
Revised: In Fig.

Original: 5, all geometry-aware methods substantially improve the log predictive probability of SVGD.
Revised: 5, all geometry-aware methods substantially improve the log predictive probability of SVGD.

Original: 4 SVNG also strongly outperforms
Revised: SVNG also strongly outperforms

Original: RSVGD and converges to its maximum test probability in half as many steps as MatSVGD (Avg) and more rapidly than MatSVGD (Mixture).
Revised: RSVGD and converges to its maximum test probability in half as many steps as MatSVGD (Avg) and more rapidly than MatSVGD (Mixture).

Original:  C ONVERGENCE G UARANTEES
Revised:  C ONVERGENCE G UARANTEES

Original: We next turn our attention to the convergence properties of our proposed methods.
Revised: We next turn our attention to the convergence properties of our proposed methods.

Original: For K t and (cid:15) t as in Alg. 1, let ( q ∞ t , q ∞ t,H ) represent the distributions of the mirrored Stein updates ( θ t , η t ) when θ 0 ∼ q ∞ 0 and η t +1 =
Revised: For K t and (cid:15) t as in Alg. 1, let ( q ∞ t , q ∞ t,H ) represent the distributions of the mirrored Stein updates ( θ t , η t ) when θ 0 ∼ q ∞ 0 and η t +1 =

Original: η t + (cid:15) t g ∗ q t ,K t ( θ t ) for t ≥ 0 .
Revised: η t + (cid:15) t g ∗ q t ,K t ( θ t ) for t ≥ 0 .

Original: Our ﬁrst result, proved in App.
Revised: Our ﬁrst result, proved in App.

Original: I.5, shows that if the Alg. 1 initialization q n 0 ,H = 1 n (cid:80) ni =

Revised: I.5, shows that if the Alg. 1 initialization q n 0 ,H = 1 n (cid:80) ni =

Original: δ η i 0 converges in Wasserstein distance to a distribution q ∞ 0 ,H as n → ∞ , then, on each round t > 0 , the output of Alg. 1, q nt = 1 n (cid:80)
Revised: δ η i 0 converges in Wasserstein distance to a distribution q ∞ 0 ,H as n → ∞ , then, on each round t > 0 , the output of Alg. 1, q nt = 1 n (cid:80)

Original: ni
Revised: ni

Original: =

Revised: =

Original: δ θ it , converges to q ∞ t .
Revised: δ θ it , converges to q ∞ t .

Original: Theorem 6 (Convergence of mirrored updates as n → ∞ ) .
Revised: Theorem 6 (Convergence of mirrored updates as n → ∞ ) .

Original: Suppose Alg. 1 is initialized with q n 0 ,H = 1 n (cid:80) ni =

Revised: Suppose Alg. 1 is initialized with q n 0 ,H = 1 n (cid:80) n i =

Original: δ η i 0 satisfying W 1 ( q n 0 ,H , q ∞ 0 ,H ) → 0 for W 1 the L 1 Wasserstein distance.
Revised: δ η i 0 satisfying W 1 ( q n 0 ,H , q ∞ 0 ,H ) → 0 for W 1 the L 1 Wasserstein distance.

Original: Deﬁne the η -induced kernel K ∇ ψ ∗ ,t ( η, η (cid:48) )
Revised: Deﬁne the η -induced kernel K ∇ ψ ∗ ,t ( η, η (cid:48) )

Original: (cid:44) K t ( ∇ ψ ∗ ( η ) , ∇ ψ ∗ ( η (cid:48) )) .
Revised: (cid:44) K t ( ∇ ψ ∗ ( η ) , ∇ ψ ∗ ( η (cid:48) )) .

Original: If, for some c 1 , c 2 > 0 , (cid:107)∇ ( K ∇ ψ ∗ ,t ( · , η ) ∇ log p H ( η ) + ∇ · K ∇ ψ∗ ,t ( · , η )) (cid:107) op ≤ c 1 (1 + (cid:107) η (cid:107) 2 ) , (cid:107)∇ ( K ∇ ψ ∗ ,t ( η (cid:48) , · ) ∇ log p H ( · ) + ∇ · K ∇ ψ ∗ ,t ( η (cid:48) , · )) (cid:107) op ≤ c
Revised: If, for some c 1 , c 2 > 0 , (cid:107)∇ ( K ∇ ψ ∗ ,t ( · , η ) ∇ log p H ( η ) + ∇ · K ∇ ψ

Original: [Equation]
Revised: [Equation]

Original: then, W 1 ( q nt,H , q ∞ t,H ) → 0 and q nt ⇒ q ∞ t for each round t . 
Revised: then, W 1 ( q n t,H , q ∞ t,H ) → 0 and q nt ⇒ q ∞ t for each round t . [Figure] s t

Original: [Figure]
Revised: [Figure]

Original: p r e d i c t i v
Revised: p r e d i c t i v

Original: [Figure]
Revised: [Figure]

Original: p r o b a b ili t
Revised: p r o b a b ili t

Original: [Figure]
Revised: [Figure]

Original: Remark
Revised: Remark

Original: The pre-conditions hold, for example, whenever ∇ log p H is Lipschitz, ψ is strongly convex, and K t = kI for k bounded with bounded derivatives.
Revised: The pre-conditions hold, for example, whenever ∇ log p H is Lipschitz, ψ is strongly convex, and K t = kI for k bounded with bounded derivatives.

Original: Given a mirrored Stein operator (6), an arbitrary Stein set G , and an arbitrary matrix-valued kernel K we deﬁne the mirrored Stein discrepancy and mirrored kernel Stein discrepancyMSD [Equation] G
Revised: Given a mirrored Stein operator (6), an arbitrary Stein set G , and an arbitrary matrix-valued kernel K we deﬁne the mirrored Stein discrepancy and mirrored kernel Stein discrepancy

Original: [Equation] (cid:44) sup g ∈G E q [ [Equation] M p,ψ g )( θ )] and MKSD K ( q, p ) (cid:44) MSD ( q
Revised: MSD ( q, p, G ) (cid:44) sup g ∈G E q [( M p,ψ g )( θ )] and MKSD K ( q, p ) (cid:44) MSD ( q, p, B H K ) .

Original: [Equation]

Original: p [Equation] K ) . 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).
Revised: 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).

Original: Since the MKSD optimization problem (15) matches that in Thm.
Revised: Since the MKSD optimization problem (15) matches that in Thm.

Original: 3, we have that MKSD K ( q, p ) =
Revised: 3, we have that MKSD K ( q, p ) =

Original: (cid:107) g ∗ q,K (cid:107) H K .
Revised: (cid:107) g ∗ q,K (cid:107) H K .

Original: Our next result, proved in App. I.6, shows that the inﬁnite-particle mirrored Stein updates reduce the KL divergence to p whenever the step size is sufﬁciently small and drive MKSD to 0 if, for example, (cid:15) t
Revised: Our next result, proved in App. I.6, shows that the inﬁnite-particle mirrored Stein updates reduce the KL divergence to p whenever the step size is sufﬁciently small and drive MKSD to 0 if, for example, (cid:15) t

Original: = Ω( MKSD K t ( q ∞ t , p ) α ) for any α > 0 .
Revised: = Ω( MKSD K t ( q ∞ t , p ) α ) for any α > 0 .

Original: We also provide two conditions in App.
Revised: We also provide two conditions in App.

Original: H that generalize
Revised: H that generalize

Original: the Stein Log-Sobolev and Stein Poincar´e inequalities in Duncan et al.
Revised: the Stein Log-Sobolev and Stein Poincar´e inequalities in Duncan et al.

Original: ; Korba et al.
Revised: ; Korba et al.

Original: and which imply exponential convergence rates of our algorithms in continuous time.
Revised: and which imply exponential convergence rates of our algorithms in continuous time.

Original: Theorem 7 (Inﬁnite-particle mirrored Stein updates decrease KL and MKSD) .
Revised: Theorem 7 (Inﬁnite-particle mirrored Stein updates decrease KL and MKSD) .

Original: Assume κ 1 (cid:44) sup θ (cid:107) K t ( θ, θ ) (cid:107) op < ∞ , and κ 2 (cid:44) (cid:80) di
Revised: Assume κ 1 (cid:44) sup θ (cid:107) K t ( θ, θ ) (cid:107) op < ∞ , and κ 2 (cid:44) (cid:80) di

Original: =1 sup θ (cid:107)∇ 2 i,d
Revised: =1 sup θ (cid:107)∇ 2 i,d

Original: + i K t ( θ, θ ) (cid:107) op < ∞ , ∇ log p H is L -Lipschitz, and ψ is α -strongly convex.
Revised: + i K t ( θ, θ ) (cid:107) op < ∞ , ∇ log p H is L -Lipschitz, and ψ is α -strongly convex.

Original: If (cid:15) t < 1 / ( (cid:107)∇ η t g ∗ q ∞ t ,K t ( θ t ) + ∇ η [Equation] g ∗ q ∞ t ,K t ( θ t ) (cid:62) (cid:107) op ) , then
Revised: If (cid:15) t < 1 / (2 (cid:107)∇ η t g ∗ q ∞ t ,K t ( θ t ) + ∇ η t g ∗ q ∞ t ,K t ( θ t ) (cid:62) (cid:107) op ) , then

Original: [Equation]
Revised: [Equation]

Original: Our last result, proved in App.
Revised: Our last result, proved in App.

Original: I.7, shows that q ∞ t ⇒ p if MKSD K k ( q ∞ t , p ) → 0 .
Revised: I.7, shows that q ∞ t ⇒ p if MKSD K k ( q ∞ t , p ) → 0 .

Original: Hence, by Thms.and 7, n -particle MSVGD converges weakly to p if (cid:15) t decays at a suitable rate.
Revised: Hence, by Thms.and 7, n -particle MSVGD converges weakly to p if (cid:15) t decays at a suitable rate.

Original: Theorem 8 (MKSD K k determines weak convergence) .
Revised: Theorem 8 (MKSD K k determines weak convergence) .

Original: Assume p H is distantly dissipative (Eberle, 2016) with ∇ log p H Lipschitz, ψ is strongly convex with continuous ∇ ψ ∗ , and k ( θ, θ (cid:48) )
Revised: Assume p H is distantly dissipative (Eberle, 2016) with ∇ log p H Lipschitz, ψ is strongly convex with continuous ∇ ψ ∗ , and k ( θ, θ (cid:48) )

Original: =
Revised: =

Original: κ ( ∇ ψ ( θ ) , ∇ ψ ( θ (cid:48) ))
Revised: κ ( ∇ ψ ( θ ) , ∇ ψ ( θ (cid:48) ))

Original: for κ ( x, y )
Revised: for κ ( x, y )

Original: =
Revised: =

Original: ( c 2 + (cid:107) x − y (cid:107) 22 ) β with β ∈ ( − 1 , 0) .
Revised: ( c 2 + (cid:107) x − y (cid:107) 22 ) β with β ∈ ( − 1 , 0) .

Original: Then, q ∞ t ⇒ p if MKSD K k ( q ∞ t , p ) → 0 .
Revised: Then, q ∞ t ⇒ p if MKSD K k ( q ∞ t , p ) → 0 .

Original: Remark The pre-conditions hold, for example, for any Dirichlet target with negative entropy ψ .
Revised: Remark The pre-conditions hold, for example, for any Dirichlet target with negative entropy ψ .

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: This paper introduced the mirrored Stein operator along with three new particle evolution algorithms for sampling with constrained domains and non-Euclidean geometries.
Revised: This paper introduced the mirrored Stein operator along with three new particle evolution algorithms for sampling with constrained domains and non-Euclidean geometries.

Original: The ﬁrst algorithm MSVGD performs SVGD updates in a mirrored space before mapping to the target domain.
Revised: The ﬁrst algorithm MSVGD performs SVGD updates in a mirrored space before mapping to the target domain.

Original: The other two algorithms are different discretizations of the same continuous dynamics for exploiting non-Euclidean geometry.
Revised: The other two algorithms are different discretizations of the same continuous dynamics for exploiting non-Euclidean geometry.

Original: SVMD is a multi-particle generalization of mirror descent for constrained domains, while
Revised: SVMD is a multi-particle generalization of mirror descent for constrained domains, while

Original: SVNG is designed for unconstrained problems with informative metric tensors.
Revised: SVNG is designed for unconstrained problems with informative metric tensors.

Original: We highlight three limitations.
Revised: We highlight three limitations.

Original: First, like SVGD, our MSVGD require O ( n 2 ) time per update.
Revised: First, like SVGD, our MSVGD require O ( n 2 ) time per update.

Original: Second, SVMD and SVNG are more costly than MSVGD due to the adaptive kernel construction.
Revised: Second, SVMD and SVNG are more costly than MSVGD due to the adaptive kernel construction.

Original: Low-rank kernel approximation may be needed to reduce their complexity.
Revised: Low-rank kernel approximation may be needed to reduce their complexity.

Original: Third, we leave open the question of convergence when stochastic gradient estimates are employed, but we suspect the results of Gorham et al.
Revised: Third, we leave open the question of convergence when stochastic gradient estimates are employed, but we suspect the results of Gorham et al.

Original: (2020, Thm. 7) can be extended to our setting.
Revised: (2020, Thm. 7) can be extended to our setting.

Original: 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-ﬁt testing (Chwialkowski et al., 2016;
Revised: 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-ﬁt testing (Chwialkowski et al., 2016;

Original: Liu et al., 2016;
Revised: Liu et al., 2016;

Original: 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).
Revised: 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).

Original: R EPRODUCIBILITY S TATEMENT
Revised: R EPRODUCIBILITY S TATEMENT

Original: See App.
Revised: See App.

Original: G for experimental details and the uploaded supplementary material for Python and R code replicating all experiments.
Revised: G for experimental details and https://github.com/thjashin/ mirror-stein-samplers for Python and R code replicating all experiments.

Original: E THICS S TATEMENT
Revised: E THICS S TATEMENT

Original: This is major theoretical work and we do not anticipate negative societal impact from this work.
Revised: This is mainly theoretical work and we do not anticipate negative societal impact from this work

Original: A M IRROR D ESCENT , R IEMANNIAN G RADIENT F LOW , AND N ATURAL G RADIENT
Revised: A M IRROR D ESCENT , R IEMANNIAN G RADIENT F LOW , AND N ATURAL G RADIENT

Original: The equivalence between the mirror ﬂow dη t =
Revised: The equivalence between the mirror ﬂow dη t =

Original: −∇
Revised: −∇

Original: f
Revised: f

Original: ( θ t ) dt, θ t = ∇ ψ ∗ ( η t )
Revised: ( θ t ) dt, θ t = ∇ ψ ∗ ( η t )

Original: dt and the Riemannian gradient ﬂow in (3) is a direct result of the chain rule:
Revised: dt and the Riemannian gradient ﬂow in (3) is a direct result of the chain rule:

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Depending on discretizing (16) or (17), there are two natural gradient descent (NGD) updates that can arise from the same gradient ﬂow:
Revised: Depending on discretizing (16) or (17), there are two natural gradient descent (NGD) updates that can arise from the same gradient ﬂow:

Original: [Equation]
Revised: [Equation]

Original: With ﬁnite step sizes (cid:15) t , their updates need not be the same and can lead to different optimization paths.
Revised: With ﬁnite step sizes (cid:15) t , their updates need not be the same and can lead to different optimization paths.

Original: Since ∇ f ( θ t ) = ∇ 2 ψ ∗ ( η t ) − 1 ∇ η t f ( ∇ ψ ∗ ( η t )) , NGD (b) is equivalent to the dual-space update by mirror descent.
Revised: Since ∇ f ( θ t ) = ∇ 2 ψ ∗ ( η t ) − 1 ∇ η t f ( ∇ ψ ∗ ( η t )) , NGD (b) is equivalent to the dual-space update by mirror descent.

Original: This relationship was pointed out in Raskutti & Mukherjee (2015) and has been used for developing natural gradient variational inference algorithms (Khan & Nielsen, 2018).
Revised: This relationship was pointed out in Raskutti & Mukherjee (2015) and has been used for developing natural gradient variational inference algorithms (Khan & Nielsen, 2018).

Original: We emphasize, however, our SVNG algorithm developed in Sec.
Revised: We emphasize, however, our SVNG algorithm developed in Sec.

Original: 4.4 corresponds to the discretization in the primal space as in NGD (a).
Revised: 4.4 corresponds to the discretization in the primal space as in NGD (a).

Original: Therefore, it does not require an explicit dual space, and allows replacing ∇ 2 ψ with more general information metric tensors.
Revised: Therefore, it does not require an explicit dual space, and allows replacing ∇ 2 ψ with more general information metric tensors.

Original: B D ETAILS OF E XAMPLE
Revised: B D ETAILS OF E XAMPLE

Original: For the entropic mirror map ψ ( θ ) = (cid:80) d +1 j =1 θ j log θ j , we have ∇ 2 ψ ( θ ) − 1 = diag( θ ) − θθ (cid:62) .
Revised: For the entropic mirror map ψ ( θ ) = (cid:80) d +1 j =1 θ j log θ j , we have ∇ 2 ψ ( θ ) − 1 = diag( θ ) − θθ (cid:62) .

Original: Note here θ denotes a d -dimensional vector and does not include θ d +1 = 1 − (cid:80) dj
Revised: Note here θ denotes a d -dimensional vector and does not include θ d +1 = 1 − (cid:80) dj

Original: =1 θ d .
Revised: =1 θ d .

Original: Since Θ is a ( d + 1) -simplex, ∂ Θ is composed of d + 1 faces with θ in the j -th face satisﬁes θ j = 0 .
Revised: Since Θ is a ( d + 1) -simplex, ∂ Θ is composed of d + 1 faces with θ in the j -th face satisﬁes θ j = 0 .

Original: The outward unit normal vector n ( θ ) for the ﬁrst d faces are − e j for 1 ≤ j ≤ d , where e j denotes the j -th standard basis of R d .
Revised: The outward unit normal vector n ( θ ) for the ﬁrst d faces are − e j for 1 ≤ j ≤ d , where e j denotes the j -th standard basis of R d .

Original: The outward unit normal vector for the ( d + 1) -st face is a vector with 1 / √ d in all coordinates.
Revised: The outward unit normal vector for the ( d + 1) -st face is a vector with 1 / √ d in all coordinates.

Original: Therefore, we have(cid:90)
Revised: Therefore, we have

Original: [Equation]
Revised: [Equation]

Original: where in the second to last identity we used θ (cid:62) 1 = 1 − θ d +1 .
Revised: where in the second to last identity we used θ (cid:62) 1 = 1 − θ d +1 .

Original: Finally
Revised: Finally

Original: [Equation]
Revised: [Equation]

Original: we can verify the condition in
Revised: we can verify the condition in

Original: Prop. 1 as
Revised: Prop. 1 as

Original: [Equation]
Revised: [Equation]

Original: C D ERIVATION OF THE M IRRORED S TEIN O PERATOR
Revised: C D ERIVATION OF THE M IRRORED S TEIN O PERATOR

Original: We ﬁrst review the (overdamped) Langevin diffusion – a Markov process that underlies many recent advances in Stein’s method – along with its recent mirrored generalization.
Revised: We ﬁrst review the (overdamped) Langevin diffusion – a Markov process that underlies many recent advances in Stein’s method – along with its recent mirrored generalization.

Original: The Langevin diffusion with equilibrium density p on R d is a Markov process ( θ t ) t ≥ 

Revised: The Langevin diffusion with equilibrium density p on R d is a Markov process ( θ t ) t ≥

Original: ⊂ R d satisfying the stochastic differential equation (SDE)
Revised: ⊂ R d satisfying the stochastic differential equation (SDE)

Original: [Equation]
Revised: [Equation]

Original: (18) with ( B t ) t ≥ 0 a standard Brownian motion (Bhattacharya & Waymire, 2009, Sec. 4.5).
Revised: (18) with ( B t ) t ≥ 0 a standard Brownian motion (Bhattacharya & Waymire, 2009, Sec. 4.5).

Original: 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):
Revised: 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):

Original: First, identify a Markov process ( θ t ) t ≥ 0 that has p as the equilibrium density; they chose the Langevin diffusion of (18).
Revised: First, identify a Markov process ( θ t ) t ≥ 0 that has p as the equilibrium density; they chose the Langevin diffusion of (18).

Original: Next, build a
Revised: Next, build a

Original: Stein operator based on the (inﬁnitesimal) generator A of the process (Øksendal, 2003, Def. 7.3.1):
Revised: Stein operator based on the (inﬁnitesimal) generator A of the process (Øksendal, 2003, Def. 7.3.1):

Original: [Equation]
Revised: [Equation]

Original: as the generator satisﬁes E θ ∼ p [( Af )( θ )]
Revised: as the generator satisﬁes E θ ∼ p [( Af )( θ )]

Original: = 0 under relatively mild conditions.
Revised: = 0 under relatively mild conditions.

Original: We use the following theorem to derive the generator of the processes described by SDEs like (18):
Revised: We use the following theorem to derive the generator of the processes described by SDEs like (18):

Original: Theorem 9 (Generator of Itˆo diffusion; Øksendal, 2003, Thm 7.3.3) .
Revised: Theorem 9 (Generator of Itˆo diffusion; Øksendal, 2003, Thm 7.3.3) .

Original: Let ( x t ) t ≥ 0 be the It ˆ o diffusion in X ⊆ R d satisfying dx t
Revised: Let ( x t ) t ≥ 0 be the It ˆ o diffusion in X ⊆ R d satisfying dx t

Original: = b ( x t ) dt + σ ( x t ) dB t .
Revised: = b ( x t ) dt + σ ( x t ) dB t .

Original: For any f ∈ C 2 c ( X ) , the (inﬁnitesimal) generator A of ( x t ) t ≥ 0 is
Revised: For any f ∈ C 2 c ( X ) , the (inﬁnitesimal) generator A of ( x t ) t ≥ 0 is

Original: [Equation]
Revised: [Equation]

Original: For the Langevin diffusion (18), substituting ∇ log p ( · ) for b ( · ) and √ 2 I for σ ( · ) in Thm.
Revised: For the Langevin diffusion (18), substituting ∇ log p ( · ) for b ( · ) and √ 2 I for σ ( · ) in Thm.

Original: 9, we obtain Af =
Revised: 9, we obtain Af =

Original: ( ∇ log p ) (cid:62) ∇ f + ∇ · ∇ f .
Revised: ( ∇ log p ) (cid:62) ∇ f + ∇ · ∇ f .

Original: Replacing ∇ f with a vector-valued function g gives the Langevin Stein operator in (5).
Revised: Replacing ∇ f with a vector-valued function g gives the Langevin Stein operator in (5).

Original: To derive a Stein operator that works well for constrained domains, we consider the Riemannian
Revised: To derive a Stein operator that works well for constrained domains, we consider the Riemannian

Original: Langevin diffusion (Patterson & Teh, 2013; Xifara et al., 2014; Ma et al., 2015) that extends the
Revised: Langevin diffusion (Patterson & Teh, 2013; Xifara et al., 2014; Ma et al., 2015) that extends the

Original: Langevin diffusion to non-Euclidean geometries encoded in a positive deﬁnite metric tensor G ( θ ) :
Revised: Langevin diffusion to non-Euclidean geometries encoded in a positive deﬁnite metric tensor G ( θ ) :

Original: [Equation]
Revised: [Equation]

Original: We show in App.
Revised: We show in App.

Original: D that the choice G = ∇ 2 ψ yields the recent mirror-Langevin diffusion (Zhang et al., 2020b; Chewi et al., 2020)
Revised: D that the choice G = ∇ 2 ψ yields the recent mirror-Langevin diffusion (Zhang et al., 2020b; Chewi et al., 2020)

Original: [Equation]
Revised: [Equation]

Original: According to Thm. 9, the generator of the mirror-Langevin diffusion described by (20) is
Revised: According to Thm. 9, the generator of the mirror-Langevin diffusion described by (20) is

Original: [Equation]
Revised: [Equation]

Original: Now substituting g ( θ ) for ∇ f ( θ ) , we obtain the associated mirrored Stein operator:
Revised: Now substituting g ( θ ) for ∇ f ( θ ) , we obtain the associated mirrored Stein operator:

Original: [Equation]
Revised: [Equation]

Original: D R IEMANNIAN L ANGEVIN D IFFUSIONS AND M IRROR -L ANGEVIN D IFFUSIONS
Revised: D R IEMANNIAN L ANGEVIN D IFFUSIONS AND M IRROR -L ANGEVIN D IFFUSIONS

Original: Zhang et al.
Revised: Zhang et al.

Original: (2020b) pointed out (19) is a particular case of the Riemannian LD.
Revised: (2020b) pointed out (19) is a particular case of the Riemannian LD.

Original: However, they did not give an explicit derivation.
Revised: However, they did not give an explicit derivation.

Original: The Riemannian LD (Patterson & Teh, 2013; Xifara et al., 2014; Ma et al., 2015) with ∇ 2 ψ ( · ) as the metric tensor is[Equation]
Revised: The Riemannian LD (Patterson & Teh, 2013; Xifara et al., 2014; Ma et al., 2015) with ∇ 2 ψ ( · ) as the metric tensor is

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: This requires the following theorem that provides the analog of the “chain rule” in SDEs.
Revised: This requires the following theorem that provides the analog of the “chain rule” in SDEs.

Original: Theorem 10 (Itˆo formula; Øksendal, 2003, Thm 4.2.1) .
Revised: Theorem 10 (Itˆo formula; Øksendal, 2003, Thm 4.2.1) .

Original: Let ( x t ) t ≥ 0 be an Itˆo process in X ⊂ R d satisfying dx t
Revised: Let ( x t ) t ≥ 0 be an Itˆo process in X ⊂ R d satisfying dx t

Original: = b ( x t ) dt + σ ( x t ) dB t .
Revised: = b ( x t ) dt + σ ( x t ) dB t .

Original: Let f ( x ) ∈ C 2 : R d → R d (cid:48) .
Revised: Let f ( x ) ∈ C 2 : R d → R d (cid:48) .

Original: Then y t
Revised: Then y t

Original: = f
Revised: = f

Original: ( x t ) is again an Itˆo process, and its i -th dimension satisﬁes
Revised: ( x t ) is again an Itˆo process, and its i -th dimension satisﬁes

Original: [Equation]
Revised: [Equation]

Original: Substituting ∇ ψ for f in Thm.
Revised: Substituting ∇ ψ for f in Thm.

Original: 10, we have the SDE of η t = ∇ ψ ( θ t ) as
Revised: 10, we have the SDE of η t = ∇ ψ ( θ t ) as

Original: [Equation]
Revised: [Equation]

Original: where h ( θ t )
Revised: where h ( θ t )

Original: i =
Revised: i =

Original: Tr( ∇ 2 θ t ( ∇ θ t,i ψ ( θ t )) ∇ 2 ψ ( θ t ) − 1 ) .
Revised: Tr( ∇ 2 θ t ( ∇ θ t,i ψ ( θ t )) ∇ 2 ψ ( θ t ) − 1 ) .

Original: Moreover, we have
Revised: Moreover, we have

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]

Original: θ t,(cid:96)

Original: [Equation]
Revised: [Equation]

Original: d

Original: [Equation]
Revised: [Equation]

Original: =

Original: [Equation]

Original: Therefore
Revised: Therefore

Original: [Figure]
Revised: [Figure]

Original: the η t diffusion is described by the SDE:
Revised: the η t diffusion is described by the SDE:

Original: [Equation]
Revised: [Equation]

Original: E M ODE M ISMATCH U NDER T RANSFORMATIONS
Revised: E M ODE M ISMATCH U NDER T RANSFORMATIONS

Original: [Figure]
Revised: [Figure]

Original: F B ACKGROUND ON R EPRODUCING K ERNEL H ILBERT S PACES
Revised: F B ACKGROUND ON R EPRODUCING K ERNEL H ILBERT S PACES

Original: Let H be a Hilbert space of functions deﬁned on X and taking their values in R .
Revised: Let H be a Hilbert space of functions deﬁned on X and taking their values in R .

Original: We say k is a reproducing kernel (or kernel) of H if ∀ x ∈ X , k ( x, · ) ∈ H and ∀ f ∈ H , (cid:104) f, k
Revised: We say k is a reproducing kernel (or kernel) of H if ∀ x ∈ X , k ( x, · ) ∈ H and ∀ f ∈ H , (cid:104) f, k

Original: ( x, · ) (cid:105) H = f ( x ) .
Revised: ( x, · ) (cid:105) H = f ( x ) .

Original: H is called a reproducing kernel Hilbert space (RKHS) if it has a kernel.
Revised: H is called a reproducing kernel Hilbert space (RKHS) if it has a kernel.

Original: Kernels are positive deﬁnite (p.d.) functions, which means that matrices with the form ( k ( x i , x j )) ij are positive semideﬁnite.
Revised: Kernels are positive deﬁnite (p.d.) functions, which means that matrices with the form ( k ( x i , x j )) ij are positive semideﬁnite.

Original: For any p.d. function k , there is a unique RKHS with k as the reproducing kernel, which can be constructed by the completion of { (cid:80) ni
Revised: For any p.d. function k , there is a unique RKHS with k as the reproducing kernel, which can be constructed by the completion of { (cid:80) ni

Original: =1 a i
Revised: =1 a i

Original: k
Revised: k

Original: ( x i , · ) , x i ∈ X , a i ∈ R , i ∈ N } .
Revised: ( x i , · ) , x i ∈ X , a i ∈ R , i ∈ N } .

Original: Now we assume X is a metric space, k is a bounded continuous kernel with the RKHS H , and ν is a positive measure on X .
Revised: Now we assume X is a metric space, k is a bounded continuous kernel with the RKHS H , and ν is a positive measure on X .

Original: L 2 ( ν ) denote the space of all square-integrable functions w.r.t.
Revised: L 2 ( ν ) denote the space of all square-integrable functions w.r.t.

Original: ν .
Revised: ν .

Original: Then the kernel integral operator T k : L 2 ( ν )
Revised: Then the kernel integral operator T k : L 2 ( ν )

Original: → L 2 ( ν ) deﬁned by
Revised: → L 2 ( ν ) deﬁned by

Original: [Equation]
Revised: [Equation]

Original: is compact and self-adjoint.
Revised: is compact and self-adjoint.

Original: Therefore, according to the spectral theorem, there exists an at most countable set of positive eigenvalues { λ j } j ∈ J ⊂ R with λ 1 ≥ λ 2 ≥ . .
Revised: Therefore, according to the spectral theorem, there exists an at most countable set of positive eigenvalues { λ j } j ∈ J ⊂ R with λ 1 ≥ λ 2 ≥ . .

Original: converging to zero and orthonormal eigenfunctions { u j } j ∈ J such that
Revised: converging to zero and orthonormal eigenfunctions { u j } j ∈ J such that

Original: [Equation]
Revised: [Equation]

Original: and k has the representation k
Revised: and k has the representation k

Original: ( x, x (cid:48) ) = (cid:80) j ∈ J λ j u j
Revised: ( x, x (cid:48) ) = (cid:80) j ∈ J λ j u j

Original: x ) u j ( x (cid:48) )
Revised: x ) u j ( x (cid:48) )

Original: (Mercer’s theorem on non-compact domains), where the convergence of the sum is absolute and uniform on compact subsets of X ×
Revised: (Mercer’s theorem on non-compact domains), where the convergence of the sum is absolute and uniform on compact subsets of X ×

Original: X (Ferreira & Menegatto, 2009).
Revised: X (Ferreira & Menegatto, 2009).

Original: G S UPPLEMENTARY E XPERIMENTAL D ETAILS AND A DDITIONAL R ESULTS
Revised: G S UPPLEMENTARY E XPERIMENTAL D ETAILS AND A DDITIONAL R ESULTS

Original: In this section, we report supplementary details and additional results from the experiments of Sec. 5.
Revised: In this section, we report supplementary details and additional results from the experiments of Sec. 5.

Original: In Secs.
Revised: In Secs.

Original: 5.1 and 5.2, we use the inverse multiquadric input kernel k ( θ, θ (cid:48) )
Revised: 5.1 and 5.2, we use the inverse multiquadric input kernel k ( θ, θ (cid:48) )

Original: = (1 + (cid:107) θ − θ (cid:48) (cid:107) 22 /(cid:96) 2 ) − 1 /due to its convergence control properties (Gorham & Mackey, 2017).
Revised: = (1 + (cid:107) θ − θ (cid:48) (cid:107) 22 /(cid:96) 2 ) − 1 /due to its convergence control properties (Gorham & Mackey, 2017).

Original: In the unconstrained experiments of Sec. 5.3, we use the Gaussian kernel k ( θ, θ (cid:48) )
Revised: In the unconstrained experiments of Sec. 5.3, we use the Gaussian kernel k ( θ, θ (cid:48) )

Original: =
Revised: =

Original: exp( −(cid:107) θ − θ (cid:48) (cid:107) 22 /(cid:96) 2 ) for consistency with past work.
Revised: exp( −(cid:107) θ − θ (cid:48) (cid:107) 22 /(cid:96) 2 ) for consistency with past work.

Original: The bandwidth (cid:96) is determined by the median heuristic (Garreau et al., 2017).
Revised: The bandwidth (cid:96) is determined by the median heuristic (Garreau et al., 2017).

Original: We select τ from { 0 . 98 , 0 . 99 } for all SVMD experiments.
Revised: We select τ from { 0 . 98 , 0 . 99 } for all SVMD experiments.

Original: For unconstrained targets, we report, for each method, results from the best ﬁxed step size (cid:15) ∈ { 0 . 01 , 0 . 05 , 0 . 1 , 0 .
Revised: For unconstrained targets, we report, for each method, results from the best ﬁxed step size (cid:15) ∈ { 0 . 01 , 0 . 05 , 0 . 1 , 0 .

Original: 5 , 1 } selected on a separate validation set.
Revised: 5 , 1 } selected on a separate validation set.

Original: For constrained targets, we select step sizes adaptively to accommodate rapid density growth near the boundary; speciﬁcally, we use RMSProp (Hinton et al., 2012), an extension of the AdaGrad algorithm (Duchi et al., 2011) used in Liu & Wang (2016), and report performance with the best learning rate.
Revised: For constrained targets, we select step sizes adaptively to accommodate rapid density growth near the boundary; speciﬁcally, we use RMSProp (Hinton et al., 2012), an extension of the AdaGrad algorithm (Duchi et al., 2011) used in Liu & Wang (2016), and report performance with the best learning rate.

Original: Results were recorded on an Intel(R) Xeon(R) CPU E5-2690 v4 @ 2.60GHz and an
Revised: Results were recorded on an Intel(R) Xeon(R) CPU E5-2690 v4 @ 2.60GHz and an

Original: NVIDIA Tesla P100 PCIe 16GB.
Revised: NVIDIA Tesla P100 PCIe 16GB.

Original: G. A PPROXIMATION QUALITY ON THE SIMPLEX
Revised: G. A PPROXIMATION QUALITY ON THE SIMPLEX

Original: The sparse Dirichlet posterior of Patterson & Teh (2013) extended to 20 dimensions features a sparse, symmetric Dir( α ) prior with α k = 0 .
Revised: The sparse Dirichlet posterior of Patterson & Teh (2013) extended to 20 dimensions features a sparse, symmetric Dir( α ) prior with α k = 0 .

Original: 1 for k ∈ { 1 , . .
Revised: 1 for k ∈ { 1 , . .

Original: , 20 } and sparse count data n 1 = 90 , n 2 = n 3 = 5 , n j = 

Revised: , 20 } and sparse count data n 1 = 90 , n 2 = n 3 = 5 , n j =

Original: ( j > 3) , modeled via a multinomial likelihood.
Revised: ( j > 3) , modeled via a multinomial likelihood.

Original: The quadratic target satisﬁes log p ( θ )
Revised: The quadratic target satisﬁes log p ( θ )

Original: = − 12 σ 2 θ (cid:62) Aθ
Revised: = − 12 σ 2 θ (cid:62) Aθ

Original: + const , where we slightly modify the target density of Ahn & Chewi (2020) to make it less ﬂat by introducing a scale parameter σ = 0 .
Revised: + const , where we slightly modify the target density of Ahn & Chewi (2020) to make it less ﬂat by introducing a scale parameter σ = 0 .

Original: A ∈ R 19 × 19 is a positive deﬁnite matrix generated by normalizing products of random matrices with i.i.d.
Revised: A ∈ R 19 × 19 is a positive deﬁnite matrix generated by normalizing products of random matrices with i.i.d.

Original: elements drawn from
Revised: elements drawn from

Original: Unif[ − 1 , 1] .
Revised: Unif[ − 1 , 1] .

Original: 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.
Revised: 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.

Original: For each method and each learning rate we apply 500 particle updates.
Revised: For each method and each learning rate we apply 500 particle updates.

Original: For SVMD we set τ = 0 .
Revised: For SVMD we set τ = 0 .

Original: We search the base learning rates of RMSProp in { 0 . 1 , 0 . 01 , 0 . 001 } for
Revised: We search the base learning rates of RMSProp in { 0 . 1 , 0 . 01 , 0 . 001 } for

Original: SVMD and MSVGD.
Revised: SVMD and MSVGD.

Original: Since projected SVGD applies updates in the θ space, the appropriate learning rate range is smaller than those of SVMD and MSVGD.
Revised: Since projected SVGD applies updates in the θ space, the appropriate learning rate range is smaller than those of SVMD and MSVGD.

Original: There we search the base learning rate of
Revised: There we search the base learning rate of

Original: RMSProp in { 0 . 01 , 0 . 001 , 0 . 0001 } .
Revised: RMSProp in { 0 . 01 , 0 . 001 , 0 . 0001 } .

Original: For all methods the results under each base learning rate are plotted in Fig.
Revised: For all methods the results under each base learning rate are plotted in Fig.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: 8: Sampling from a quadratic target on a 20 -simplex.
Revised: : Sampling from a quadratic target on a 20 -simplex.

Original: We plot the energy distance to a ground truth sample of size 1000 drawn by NUTS (Hoffman & Gelman, 2014).
Revised: We plot the energy distance to a ground truth sample of size 1000 drawn by NUTS (Hoffman & Gelman, 2014).

Original: G. C ONFIDENCE INTERVALS FOR POST - SELECTION INFERENCE
Revised: G. C ONFIDENCE INTERVALS FOR POST - SELECTION INFERENCE

Original: Given a dataset X ∈ R ˜ n × p , y ∈ R ˜ n , the randomized Lasso (Tian & Taylor, 2018) solves the following problem:
Revised: Given a dataset X ∈ R ˜ n × p , y ∈ R ˜ n , the randomized Lasso (Tian & Taylor, 2018) solves the following problem:

Original: [Equation]
Revised: [Equation]

Original: where G is a user-speciﬁed log-concave distribution with density g .
Revised: where G is a user-speciﬁed log-concave distribution with density g .

Original: We choose G to be zeromean independent Gaussian distributions while leaving its scale and the ridge parameter (cid:15) to be automatically determined by the randomizedLasso function of the selectiveInference package.
Revised: We choose G to be zeromean independent Gaussian distributions while leaving its scale and the ridge parameter (cid:15) to be automatically determined by the randomizedLasso function of the selectiveInference package.

Original: We initialize the particles of our SVMD and MSVGD in the following way: First, we map the solution ˆ β E to the dual space by ∇ ψ .
Revised: We initialize the particles of our SVMD and MSVGD in the following way: First, we map the solution ˆ β E to the dual space by ∇ ψ .

Original: Next, we add i.i.d.
Revised: Next, we add i.i.d.

Original: standard Gaussian noise to n copies of the image in the dual space.
Revised: standard Gaussian noise to n copies of the image in the dual space.

Original: Finally, we map the n particles back to the primal space by ∇ ψ ∗ and use them as the initial locations.
Revised: Finally, we map the n particles back to the primal space by ∇ ψ ∗ and use them as the initial locations.

Original: Below we discuss the remaining settings and additional results of the simulation and the HIV-1 drug resistance experiment separately.
Revised: Below we discuss the remaining settings and additional results of the simulation and the HIV-1 drug resistance experiment separately.

Original: 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.
Revised: 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.

Original: We set ˜ n = 100 and p = 40 .
Revised: We set ˜ n = 100 and p = 40 .

Original: The design matrix X is generated from an equi-correlated model, i.e., each datapoint x i ∈ R p is generated i.i.d.
Revised: The design matrix X is generated from an equi-correlated model, i.e., each datapoint x i ∈ R p is generated i.i.d.

Original: from N (0 , Σ) with Σ ii = 1 , Σ ij = 0 .
Revised: from N (0 , Σ) with Σ ii = 1 , Σ ij = 0 .

Original: 3 ( i (cid:54) =
Revised: 3 ( i (cid:54) =

Original: j ) and then normalized to have almost unit length.
Revised: j ) and then normalized to have almost unit length.

Original: The normalization is done by ﬁrst centering each dimension by subtracting the mean and dividing the standard deviation of that column of X , then additionally multiplying 1 / ˜ n 1 / 2 .
Revised: The normalization is done by ﬁrst centering each dimension by subtracting the mean and dividing the standard deviation of that column of X , then additionally multiplying 1 / ˜ n 1 / 2 .

Original: 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.
Revised: 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.

Original: We set λ to be the value returned by theoretical.lambda of the selectiveInference R package multiplied a coefﬁcient 0 .
Revised: We set λ to be the value returned by theoretical.lambda of the selectiveInference R package multiplied a coefﬁcient 0 .

Original: 7˜ n , where the 0 .
Revised: 7˜ n , where the 0 .

Original: 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 .
Revised: 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 .

Original: The base learning rates for SVMD and MSVGD are set to 0 .
Revised: The base learning rates for SVMD and MSVGD are set to 0 .

Original: and we run them for T = 1000 particle updates.
Revised: and we run them for T = 1000 particle updates.

Original: τ is set to 0 .
Revised: τ is set to 0 .

Original: 98 for SVMD.
Revised: 98 for SVMD.

Original: Our 2D example in Fig.
Revised: Our 2D example in Fig.

Original: 4a is grabbed from one run of the simulation where there happen to be only 2 features selected by the randomized Lasso.
Revised: 4a is grabbed from one run of the simulation where there happen to be only 2 features selected by the randomized Lasso.

Original: The selective distribution in this case has logdensity log p ( θ )
Revised: The selective distribution in this case has logdensity log p ( θ )

Original: = − 8 .
Revised: = − 8 .

Original: 39859 θ 1 + 1 .
Revised: 39859 θ 1 + 1 .

Original: θ 2 + 2 . 39751) 2 + (1 . 18099 θ 2 − 1 . 46104) 2 )
Revised: θ 2 + 2 . 39751) 2 + (1 . 18099 θ 2 − 1 . 46104) 2 )

Original: + const , θ 1 , 2 ≥ 0 .
Revised: + const , θ 1 , 2 ≥ 0 .

Original: The error bars for actual coverage levels in Fig. 3a and Fig.
Revised: The error bars for actual coverage levels in Fig. 3a and Fig.

Original: 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.
Revised: 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.

Original: In Fig.
Revised: In Fig.

Original: 9a and Fig.
Revised: 9a and Fig.

Original: 9b we additionally plot the average length of the conﬁdence intervals w.r.t.
Revised: 9b we additionally plot the average length of the conﬁdence intervals w.r.t.

Original: different sample size N and nominal coverage levels.
Revised: different sample size N and nominal coverage levels.

Original: For all three methods the CI widths are very close, although MSVGD consistently has wider intervals than SVMD and selectiveInference .
Revised: For all three methods the CI widths are very close, although MSVGD consistently has wider intervals than SVMD and selectiveInference .

Original: This indicates that SVMD can be preferred over MSVGD when both methods produce coverage above the nominal level.
Revised: This indicates that SVMD can be preferred over MSVGD when both methods produce coverage above the nominal level.

Original: 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.
Revised: 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.

Original: This results in ˜ n = 663 datapoints with p = 91 features.
Revised: This results in ˜ n = 663 datapoints with p = 91 features.

Original: We choose λ to be the value returned by theoretical.lambda of the selectiveInference R package multiplied by ˜ n .
Revised: We choose λ to be the value returned by theoretical.lambda of the selectiveInference R package multiplied by ˜ n .

Original: The base learning rates for SVMD and MSVGD are set to 0 .
Revised: The base learning rates for SVMD and MSVGD are set to 0 .

Original: and we run them for T = 2000 particle updates.
Revised: and we run them for T = 2000 particle updates.

Original: τ is set to 0 .
Revised: τ is set to 0 .

Original: 99 for SVMD.
Revised: 99 for SVMD.

Original: [Figure]
Revised: [Figure]

Original: G. L ARGE - SCALE POSTERIOR INFERENCE WITH NON -E UCLIDEAN GEOMETRY
Revised: G. L ARGE - SCALE POSTERIOR INFERENCE WITH NON -E UCLIDEAN GEOMETRY

Original: The Bayesian logistic regression model we consider is (cid:81) ˜
Revised: The Bayesian logistic regression model we consider is (cid:81) ˜

Original: ni =1 p ( y i | x i , w ) p ( w ) , where p ( w ) =
Revised: ni =1 p ( y i | x i , w ) p ( w ) , where p ( w ) =

Original: N ( w | 0 , I ) , p ( y i | x i , w ) =
Revised: N ( w | 0 , I ) , p ( y i | x i , w ) =

Original: Bernoulli( σ ( w (cid:62) x i )) .
Revised: Bernoulli( σ ( w (cid:62) x i )) .

Original: The bias parameter is absorbed into into w by adding an additional feature 1 to each x i .
Revised: The bias parameter is absorbed into into w by adding an additional feature 1 to each x i .

Original: The gradient of the log density of the posterior distribution of w is ∇ w log p ( w |{ y i , x i } Ni =1 ) =
Revised: The gradient of the log density of the posterior distribution of w is ∇ w log p ( w |{ y i , x i } Ni =1 ) =

Original: (cid:80) Ni
Revised: (cid:80) Ni

Original: =1 x i ( y i − σ ( w (cid:62) x i ))
Revised: =1 x i ( y i − σ ( w (cid:62) x i ))

Original: − w. We choose the metric tensor ∇ 2 ψ ( w ) to be the Fisher information matrix (FIM) of the likelihood:
Revised: − w. We choose the metric tensor ∇ 2 ψ ( w ) to be the Fisher information matrix (FIM) of the likelihood:

Original: [Equation]
Revised: [Equation]

Original: Following Wang et al.
Revised: Following Wang et al.

Original: (2019), for each iteration r ( r ≥ 1 ), we estimate the sum with a stochastic minibatch B r of size 256 : ˆ 
Revised: (2019), for each iteration r ( r ≥ 1 ), we estimate the sum with a stochastic minibatch B r of size 256 : ˆ F B

Original: [Equation]
Revised: [Equation]

Revised: =

Original: [Equation]
Revised: [Equation]

Revised: n

Revised: [Equation]

Original:  i and approximate the
Revised: [Equation] − σ [Equation] )) x i x (cid:62) i and approximate the

Original: FIM with a moving average across iterations:
Revised: FIM with a moving average across iterations:

Original: [Equation]
Revised: [Equation]

Original: where ρ r = min(1 − 1 /r, 0 . 95) .
Revised: where ρ r = min(1 − 1 /r, 0 . 95) .

Original: To ensure the positive deﬁniteness of the FIM, a damping term 0 .
Revised: To ensure the positive deﬁniteness of the FIM, a damping term 0 .

Original: I is added before taking the inverse.
Revised: I is added before taking the inverse.

Original: For RSVGD and SVNG, the gradient of the inverse of FIM is estimated with ∇ w j F − 1 ≈ − ˆ F − 1 r ( ˆ ∇ rw j F ) ˆ F − 1 r , where ˆ ∇ rw j F = ρ r ˆ ∇ r − 1 w j F + (1 − ρ r ) 
Revised: For RSVGD and SVNG, the gradient of the inverse of FIM is estimated with ∇ w j F − 1 ≈ − ˆ F − 1 r ( ˆ ∇ rw j F ) ˆ F − 1 r , where ˆ ∇ rw j F = ρ r ˆ ∇ r − 1 w j F + (1 − ρ r ) ∇

Original: [Equation]
Revised: [Equation]

Original: 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.
Revised: 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.

Original: τ is set to 0 .
Revised: τ is set to 0 .

Original: 98 for SVNG.
Revised: 98 for SVNG.

Original: 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.
Revised: 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.

Original: The results of each method on validation sets with all choices of learning rates are plotted in Fig.
Revised: The results of each method on validation sets with all choices of learning rates are plotted in Fig.

Original: We see that the SVNG updates are very robust to the change in learning rates and is able to accommodate very large learning rates (up to 1) without a signiﬁcant loss in performance.
Revised: We see that the SVNG updates are very robust to the change in learning rates and is able to accommodate very large learning rates (up to 1) without a signiﬁcant loss in performance.

Original: The results in Fig.
Revised: The results in Fig.

Original: 5 are reported with the learning rate that performs best on the validation set.
Revised: 5 are reported with the learning rate that performs best on the validation set.

Original: [Figure]
Revised: [Figure]

Original: V a li d a t i o n l og
Revised: V a li d a t i o n l og

Original: [Table]
Revised: [Table]

Original: r 
Revised: r e

Original: [Table]
Revised: [Table]

Original: i c t i v e p r
Revised: i c t i v e p r

Original: [Table]
Revised: [Table]

Original:  a b ili t y
Revised: b a b ili t y

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table] 
Revised: [Table] [Figure] [Table]

Original: [Figure]
Revised: [Figure]

Original: V a li d a t i o n l og p r e
Revised: V a li d a t i o n l og p r e

Original: d
Revised: d

Original: i c t i v e p r o b a b ili t y
Revised: i c t i v e p r o b a b ili t y

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: og
Revised: [Figure] [Table] [Figure] [Table]

Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: H E XPONENTIAL C ONVERGENCE OF C ONTINUOUS - TIME A LGORITHMS
Revised: H E XPONENTIAL C ONVERGENCE OF C ONTINUOUS - TIME A LGORITHMS

Original: We derive a time-inhomogeneous generalization of the Stein Log-Sobolev inequality of Duncan et al.
Revised: We derive a time-inhomogeneous generalization of the Stein Log-Sobolev inequality of Duncan et al.

Original: (2019) and Korba et al.
Revised: (2019) and Korba et al.

Original: (2020) which ensures the exponential convergence of our continuous-time algorithms and time-inhomogeneous generalization of the Stein Poincar´e inequality of Duncan et al.
Revised: (2020) which ensures the exponential convergence of our continuous-time algorithms and time-inhomogeneous generalization of the Stein Poincar´e inequality of Duncan et al.

Original: (2019) which guarantees exponential convergence near equilibrium (i.e., when q t is sufﬁciently close to p ).
Revised: (2019) which guarantees exponential convergence near equilibrium (i.e., when q t is sufﬁciently close to p ).

Original: As the results hold for a generic sequence of kernels ( K t ) t ≥ 0 , the implications apply to both
Revised: As the results hold for a generic sequence of kernels ( K t ) t ≥ 0 , the implications apply to both

Original: MSVGD and SVMD.
Revised: MSVGD and SVMD.

Original: Deﬁnition 2 (Mirror Stein Log-Sobolev inequality) .
Revised: Deﬁnition 2 (Mirror Stein Log-Sobolev inequality) .

Original: We deﬁne the Mirror Stein Log-Sobolev inequality (cf., Korba et al., 2020, Def. 2) as
Revised: We deﬁne the Mirror Stein Log-Sobolev inequality (cf., Korba et al., 2020, Def. 2) as

Original: [Equation]
Revised: [Equation]

Original: where MKSD K t is deﬁned in Eq.
Revised: where MKSD K t is deﬁned in Eq.

Original: (15) ; P K t ,q t : L 2 ( q t ) → L 2 ( q t ) is the kernel integral operator: ( P K t ,q t ϕ )( · ) (cid:44) E q t ( θ )
Revised: (15) ; P K t ,q t : L 2 ( q t ) → L 2 ( q t ) is the kernel integral operator: ( P K t ,q t ϕ )( · ) (cid:44) E q t ( θ )

Original: [ K t ( · , θ ) ϕ ( θ )] for a general kernel K t and vector-valued function ϕ on Θ , and the stated equality holds whenever integration-by-parts is applicable.
Revised: [ K t ( · , θ ) ϕ ( θ )] for a general kernel K t and vector-valued function ϕ on Θ , and the stated equality holds whenever integration-by-parts is applicable.

Original: Proposition 11.
Revised: Proposition 11.

Original: Suppose ( θ t ) t ≥ 0 follows the mirrored dynamics (8) with g t chosen to be g ∗ q t ,K t as in (10) .
Revised: Suppose ( θ t ) t ≥ 0 follows the mirrored dynamics (8) with g t chosen to be g ∗ q t ,K t as in (10) .

Original: Then, the dissipation of KL ( q t (cid:107) p ) is
Revised: Then, the dissipation of KL ( q t (cid:107) p ) is

Original: [Equation]
Revised: [Equation]

Original: Proof
Revised: Proof

Original: The proof directly follows from Thm. 

Revised: The proof directly follows from Thm.

Original: since the optimization problem there matches the deﬁnition of MKSD in (15).
Revised: since the optimization problem there matches the deﬁnition of MKSD in (15).

Original: Therefore, when the Mirror Stein Log-Sobolev inequality holds, we have
Revised: Therefore, when the Mirror Stein Log-Sobolev inequality holds, we have

Original: [Equation]
Revised: [Equation]

Original: , and the exponential convergence KL( q t (cid:107) p ) ≤ KL( q 0 (cid:107) p ) e − 2 λt follows by Gronwall’s lemma (Gronwall, 1919).
Revised: , and the exponential convergence KL( q t (cid:107) p ) ≤ KL( q 0 (cid:107) p ) e − 2 λt follows by Gronwall’s lemma (Gronwall, 1919).

Original: Deﬁnition 3 (Mirror Stein Poincar ´ e inequality) .
Revised: Deﬁnition 3 (Mirror Stein Poincar ´ e inequality) .

Original: We say that the distribution p satisﬁes the Mirror
Revised: We say that the distribution p satisﬁes the Mirror

Original: Stein Poincar´e inequality (cf., Duncan et al.
Revised: Stein Poincar´e inequality (cf , Duncan al

Original: 2019, Eq.

Original: (57)) with strongly convex ψ andconstant λ if
Revised: 57)) with strongly convex ψ and λ if

Original: [Equation]
Revised: [Equation]

Original: for all φ ∈ L 2 ( p ) ∩ C ∞ (Θ) that is locally Lipschitz, where P K t ,p is the kernel integral operator under p deﬁned similarly as in Deﬁnition 2.
Revised: for all φ ∈ L 2 ( p ) ∩ C ∞ (Θ) that is locally Lipschitz, where P K t ,p is the kernel integral operator under p deﬁned similarly as in Deﬁnition 2.

Original: This inequality can also be viewed as a kernelized generalization of the mirror
Revised: This inequality can also be viewed as a kernelized generalization of the mirror

Original: Poincar ´ e inequality introduced in Chewi et al.
Revised: Poincar ´ e inequality introduced in Chewi et al.

Original: (2020, Def. 1) for proving exponential convergence of mirror-Langevin diffusion.
Revised: (2020, Def. 1) for proving exponential convergence of mirror-Langevin diffusion.

Original: In a manner analogous to Thm.
Revised: In a manner analogous to Thm.

Original: 1 of Chewi et al.
Revised: 1 of Chewi et al.

Original: , the following proposition relates the Mirror Stein Poincar ´ e inequality to chi-squared divergence.
Revised: , the following proposition relates the Mirror Stein Poincar ´ e inequality to chi-squared divergence.

Original: Proposition 12.
Revised: Proposition 12.

Original: Suppose ( θ t ) t ≥ 0 follows the mirrored dynamics (8) with g t chosen to be g ∗ q t ,K t as in (10) .
Revised: Suppose ( θ t ) t ≥ 0 follows the mirrored dynamics (8) with g t chosen to be g ∗ q t ,K t as in (10) .

Original: Then, the dissipation of chi-square divergence χ 2 ( q t (cid:107) p ) is
Revised: Then, the dissipation of chi-square divergence χ 2 ( q t (cid:107) p ) is

Original: [Equation]
Revised: [Equation]

Original: whenever integration-by-parts is applicable.
Revised: whenever integration-by-parts is applicable.

Original: Proof
Revised: Proof

Original: We ﬁrst note that by applying integration-by-parts, g ∗ q t ,K t as in (10) can be equivalently written as
Revised: We ﬁrst note that by applying integration-by-parts, g ∗ q t ,K t as in (10) can be equivalently written as

Original: [Equation]
Revised: [Equation]

Original: Then using the Fokker-Planck equation of q t under the dynamics, we have
Revised: Then using the Fokker-Planck equation of q t under the dynamics, we have

Original: [Equation]
Revised: [Equation]

Original: Note that the right hand side of the equation differs from the Mirror Stein Poincar ´ e inequality only in the base measure of the expectation.
Revised: Note that the right hand side of the equation differs from the Mirror Stein Poincar ´ e inequality only in the base measure of the expectation.

Original: Duncan et al.
Revised: Duncan et al.

Original: (2019) proposes to replace q t with p to study the convergence near equilibrium (See their Sec. 6, where Eq.
Revised: (2019) proposes to replace q t with p to study the convergence near equilibrium (See their Sec. 6, where Eq.

Original: (46) is replaced with Eq.
Revised: (46) is replaced with Eq.

Original: If we do the same and combine this identity with the Mirror Stein Poincar´e inequality, we obtain dd t χ 2 ( q t (cid:107) p ) ≤ − 2 λ
Revised: If we do the same and combine this identity with the Mirror Stein Poincar´e inequality, we obtain dd t χ 2 ( q t (cid:107) p ) ≤ − 2 λ

Original: Var p [ q t p ] = − 2 λχ 2 ( q t (cid:107) p ) , which implies exponential convergence in KL
Revised: Var p [ q t p ] = − 2 λχ 2 ( q t (cid:107) p ) , which implies exponential convergence in KL

Original: KL( q t (cid:107) p ) ≤ χ 2 ( q t (cid:107) p ) ≤ χ 2 ( q 0 (cid:107) p ) e − 2 λt by Gronwall’s lemma (Gronwall, 1919).
Revised: KL( q t (cid:107) p ) ≤ χ 2 ( q t (cid:107) p ) ≤ χ 2 ( q 0 (cid:107) p ) e − 2 λt by Gronwall’s lemma (Gronwall, 1919).

Original: I P ROOFS I. P ROOF OF P ROP .
Revised: I P ROOFS I. P ROOF OF P ROP .

Original: Proof
Revised: Proof

Original: Fix any g ∈ G ψ .
Revised: Fix any g ∈ G ψ .

Original: Since g and ∇ g are bounded and ∇ 2 ψ ( θ ) − 1 ∇ log p ( θ ) and ∇ · ∇ 2 ψ ( θ ) − 1 are p -integrable, the expectation E θ ∼ p
Revised: Since g and ∇ g are bounded and ∇ 2 ψ ( θ ) − 1 ∇ log p ( θ ) and ∇ · ∇ 2 ψ ( θ ) − 1 are p -integrable, the expectation E θ ∼ p

Original: ( M p,ψ g )
Revised: ( M p,ψ g )

Original: ( θ )] exists.
Revised: ( θ )] exists.

Original: Because Θ is convex, Θ r is bounded and convex with Lipschitz boundary.
Revised: Because Θ is convex, Θ r is bounded and convex with Lipschitz boundary.

Original: Since p ∇ 2 ψ − 1 g ∈ C 1 , we have
Revised: Since p ∇ 2 ψ − 1 g ∈ C 1 , we have

Original: [Equation]
Revised: [Equation]

Revised: [Table]

Original: by the divergence theorem)
Revised: divergence theorem)

Original: [Equation]
Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

Revised: [Equation]

Original: I. P ROOF OF T HM . 3: O PTIMAL MIRROR UPDATES IN RKHSProof
Revised: I. P ROOF OF T HM . 3: O PTIMAL MIRROR UPDATES IN RKHS

Revised: Proof

Original: Let e i denote the standard basis vector of R d with the i -th element being 1 and others being zeros.
Revised: Let e i denote the standard basis vector of R d with the i -th element being 1 and others being zeros.

Original: Since m ∈ H K
Revised: Since m ∈ H K

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: where we deﬁne the divergence of a matrix as a vector whose elements are the divergences of each row of the matrix.
Revised: where we deﬁne the divergence of a matrix as a vector whose elements are the divergences of each row of the matrix.

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: we write (9) as
Revised: we write (9) as

Original: [Equation]
Revised: [Equation]

Original: Therefore, the optimal direction in the H K norm ball B H K = { g : (cid:107) g(cid:107) H K ≤ 1 } that minimizes (9) is g ∗ t ∝ g ∗ q t ,
Revised: Therefore, the optimal direction in the H K norm ball B H K = { g : [Equation] g

Original:  K = E q t [ M p,ψ K ( · , θ )] .
Revised: [Equation] H K ≤ 1 } that minimizes (9) is g ∗ t ∝ [Equation]

Original: I. P ROOF OF T HM . 4: M IRRORED SVGD UPDATES
Revised: I. P ROOF OF T HM . 4: M IRRORED SVGD UPDATES

Original: Proof
Revised: Proof

Original: A p.d. kernel k composed with any map φ is still a p.d. kernel.
Revised: A p.d. kernel k composed with any map φ is still a p.d. kernel.

Original: To prove this, let { x 1 , . . .
Revised: To prove this, let { x 1 , . . .

Original: , x p } = { φ ( η 1 ) , . . .
Revised: , x p } = { φ ( η 1 ) , . . .

Original: , φ ( η n ) } , p ≤ n .
Revised: , φ ( η n ) } , p ≤ n .

Original: Then (cid:88)
Revised: Then (cid:88)

Original: [Equation]
Revised: [Equation]

Original: where β (cid:96) =
Revised: where β (cid:96) =

Original: (cid:80) i ∈ S (cid:96) α i , S (cid:96) =
Revised: (cid:80) i ∈ S (cid:96) α i , S (cid:96) =

Original: { i : φ ( η i ) =
Revised: { i : φ ( η i ) =

Original: x (cid:96) } .
Revised: x (cid:96) } .

Original: Therefore, k ψ ( η, η (cid:48) )
Revised: Therefore, k ψ ( η, η (cid:48) )

Original: =
Revised: =

Original: k ( ∇ ψ ∗ ( η ) , ∇ ψ ∗ ( η (cid:48) )) is a p.d. kernel.
Revised: k ( ∇ ψ ∗ ( η ) , ∇ ψ ∗ ( η (cid:48) )) is a p.d. kernel.

Original: Plugging K = kI into Lem.
Revised: Plugging K = kI into Lem.

Original: 13, for any θ (cid:48)
Revised: 13, for any θ (cid:48)

Original: ∈ Θ and η (cid:48)
Revised: ∈ Θ and η (cid:48)

Original: = ∇ ψ ( θ (cid:48) ) , we have g 
Revised: = ∇ ψ ( θ (cid:48) ) , we have g [Equation] t , [Equation]

Revised: t

Original: [Equation]
Revised: [Equation]

Original: I. P ROOF OF P ROP . 5: S INGLE - PARTICLE SVMD IS MIRROR DESCENTProof
Revised: I. P ROOF OF P ROP . 5: S INGLE - PARTICLE SVMD IS MIRROR DESCENT

Revised: Proof

Original: When n = 1 , λ 1 = k ( θ t , θ t ) , u 1 = 1 , and thus K ψ,t ( θ t , θ t )
Revised: When n = 1 , λ 1 = k ( θ t , θ t ) , u 1 = 1 , and thus K ψ,t ( θ t , θ t )

Original: =
Revised: =

Original: k ( θ t , θ t ) ∇ 2 ψ ( θ t ) .
Revised: k ( θ t , θ t ) ∇ 2 ψ ( θ t ) .

Original: I. P ROOF OF T HM . 6: C ONVERGENCE OF MIRRORED UPDATES AS n → ∞
Revised: I. P ROOF OF T HM . 6: C ONVERGENCE OF MIRRORED UPDATES AS n → ∞

Original: Proof
Revised: Proof

Original: The idea is to reinterpret our mirrored updates as one step of a matrix SVGD in η space based on Lem.
Revised: The idea is to reinterpret our mirrored updates as one step of a matrix SVGD in η space based on Lem.

Original: and then follow the path of Gorham et al.
Revised: and then follow the path of Gorham et al.

Original: (2020, Thm. 7).
Revised: (2020, Thm. 7).

Original: Assume that q nt,H and q ∞ t,H have integrable means.
Revised: Assume that q nt,H and q ∞ t,H have integrable means.

Original: Let η n , η ∞ be an optimal Wasserstein-1 coupling of q nt,H and q ∞ t,H .
Revised: Let η n , η ∞ be an optimal Wasserstein-1 coupling of q nt,H and q ∞ t,H .

Original: Let Φ q t ,K t denote the transform through one step of mirrored update: θ t = ∇ ψ (cid:63) ( η t ) , η t +1 = η t + (cid:15) t g ∗ q t ,K t ( θ t ) .
Revised: Let Φ q t ,K t denote the transform through one step of mirrored update: θ t = ∇ ψ (cid:63) ( η t ) , η t +1 = η t + (cid:15) t g ∗ q t ,K t ( θ t ) .

Original: Then, with Lem. 13, we have
Revised: Then, with Lem. 13, we have

Original: [Equation]
Revised: [Equation]

Original: Since Φ q t ,K t ( η n ) ∼ q nt +1 ,H , Φ q t ,K ( η ∞ ) ∼ q ∞ t +1 ,H , we conclude[Equation] [Equation]
Revised: Since Φ q t ,K t ( η n ) ∼ q nt +1 ,H , Φ q t ,K ( η ∞ ) ∼ q ∞ t +1 ,H , we conclude

Original: [Equation]
Revised: [Equation]

Original: The ﬁnal claim q nt ⇒ q ∞ t now follows by the continuous mapping theorem as ∇ ψ ∗ is continuous.
Revised: The ﬁnal claim q nt ⇒ q ∞ t now follows by the continuous mapping theorem as ∇ ψ ∗ is continuous.

Original: I. P ROOF OF T HM . 7: I NFINITE - PARTICLE MIRRORED S TEIN UPDATES DECREASE KL AND MKSD
Revised: I. P ROOF OF T HM . 7: I NFINITE - PARTICLE MIRRORED S TEIN UPDATES DECREASE KL AND MKSD

Original: Proof
Revised: Proof

Original: Let T q ∞ t
Revised: Let T q ∞ t

Original: [Equation]
Revised: [Equation]

Original: K t denote transform of the density function through one step of mirrored update: θ t
Revised: K t denote transform of the density function through one step of mirrored update: θ t

Original: [Equation]
Revised: [Equation]

Original: t
Revised: η

Original: [Equation]
Revised: [Equation]

Original: , η

Original: [Equation]

Original: [Equation]

Original: t

Original: [Equation]
Revised: [Equation]

Original: t g ∗ q ∞ t ,K t ( θ t ) .
Revised: t g ∗ q ∞ t ,K t ( θ t ) .

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: where we have used the invariance of KL divergence under reparameterization: KL ( q t (cid:107) p )
Revised: where we have used the invariance of KL divergence under reparameterization: KL ( q t (cid:107) p )

Original: =
Revised: =

Original: KL ( q t,H (cid:107) p H ) .
Revised: KL ( q t,H (cid:107) p H ) .

Original: Following Liu (2017)
Revised: Following Liu (2017)

Original: [Equation]
Revised: [Equation]

Original: we bound the difference of the ﬁrst two terms as
Revised: we bound the difference of the ﬁrst two terms as

Original: [Equation]
Revised: [Equation]

Original: and bound the log determinant term using Lem. 15:
Revised: and bound the log determinant term using Lem. 15:

Original: [Equation]
Revised: [Equation]

Original: The next thing to notice is that E η t ∼ q ∞ t,H [ ∇ log p H ( η t ) (cid:62)
Revised: The next thing to notice is that E η t ∼ q ∞ t,H [ ∇ log p H ( η t ) (cid:62)

Original: g ∗ q ∞ t ,K t ( θ t )
Revised: g ∗ q ∞ t ,K t ( θ t )

Original: + Tr( ∇ η t g ∗ q ∞ t ,K t ( θ t ))] is the square of the MKSD in (15).
Revised: + Tr( ∇ η t g ∗ q ∞ t ,K t ( θ t ))] is the square of the MKSD in (15).

Original: We can show this equivalence using the identity proved in Lem.
Revised: We can show this equivalence using the identity proved in Lem.

Original: [Equation]
Revised: [Equation]

Original: Finally, we are going to bound (cid:107) g ∗ q ∞ t ,K t ( θ t ) (cid:107) 

Revised: Finally, we are going to bound (cid:107) g ∗ q ∞ t ,K t ( θ t ) (cid:107)

Original: and (cid:107)∇ η t g ∗ q ∞ t ,K t ( θ t ) (cid:107) 2 F .
Revised: and (cid:107)∇ η t g ∗ q ∞ t ,K t ( θ t ) (cid:107) 2 F .

Original: From the assumptions we have ψ is α -strongly convex and thus ψ ∗ is 1 α -strongly smooth (Kakade et al., 2009), therefore (cid:107)∇ 2 ψ ∗ ( · ) (cid:107) 2 ≤ 1 α .
Revised: From the assumptions we have ψ is α -strongly convex and thus ψ ∗ is 1 α -strongly smooth (Kakade et al., 2009), therefore (cid:107)∇ 2 ψ ∗ ( · ) [Equation] 2 ≤ 1 α .

Original: By Lem.
Revised: By Lem.

Original: 16, we know
Revised: 16, we know

Original: [Equation]
Revised: [Equation]

Original: where ∇ 2 i,d + i K ( θ, θ ) denotes ∇ 2 θ i ,θ (cid:48)
Revised: where ∇ 2 i,d + i K ( θ, θ ) denotes ∇ 2 θ i ,θ (cid:48)

Original: i K ( θ, θ (cid:48) )
Revised: i K ( θ, θ (cid:48) )

Original: | θ (cid:48)
Revised: | θ (cid:48)

Original: = θ .
Revised: = θ .

Original: Combining all of the above, we have
Revised: Combining all of the above, we have

Original: KL ( q ∞ t +1 (cid:107) p ) − KL ( q ∞ t (cid:107) p ) ≤ − (cid:32) (cid:15) t −
Revised: KL ( q ∞ t +1 (cid:107) p ) − KL ( q ∞ t (cid:107) p ) ≤ − (cid:32) (cid:15) t −

Original: L(cid:15) 2 t 2 sup θ (cid:107)
Revised: L(cid:15) 2 t 2 sup θ (cid:107)

Original: [Equation]
Revised: [Equation]

Original: Plugging in the deﬁnition of κ 1 and κ 2 ﬁnishes the proof.
Revised: Plugging in the deﬁnition of κ 1 and κ 2 ﬁnishes the proof

Revised: [Equation]

Original: I. P ROOF OF T HM . 8: MKSD K k DETERMINES WEAK CONVERGENCE
Revised: I. P ROOF OF T HM . 8: MKSD K k DETERMINES WEAK CONVERGENCE

Original: Proof
Revised: Proof

Original: According to Thm. 4,
Revised: According to Thm. 4,

Original: [Equation]
Revised: [Equation]

Original: where q H ( η ) denotes the density of η = ∇ ψ ( θ ) under the distribution θ ∼ q .
Revised: where q H ( η ) denotes the density of η = ∇ ψ ( θ ) under the distribution θ ∼ q .

Original: From the assumptions we have k ( θ, θ (cid:48) )
Revised: From the assumptions we have k ( θ, θ (cid:48) )

Original: =
Revised: =

Original: κ ( ∇ ψ ( θ ) , ∇ ψ ( θ (cid:48) )) .
Revised: κ ( ∇ ψ ( θ ) , ∇ ψ ( θ (cid:48) )) .

Original: With this speciﬁc choice of k , the squared MKSD is
Revised: With this speciﬁc choice of k [Equation] the squared MKSD is

Original: [Equation]
Revised: [Equation]

Original: The ﬁnal expression in (21) is the squared kernel Stein discrepancy (KSD) (Liu et al., 2016;
Revised: The ﬁnal expression in (21) is the squared kernel Stein discrepancy (KSD) (Liu et al., 2016;

Original: Chwialkowski et al., 2016; Gorham & Mackey, 2017) between q H and p H with the kernel κ :
Revised: Chwialkowski et al., 2016; Gorham & Mackey, 2017) between q H and p H with the kernel κ :

Original: KSD κ ( q H , p H ) 2 .
Revised: KSD κ ( q H , p H ) 2 .

Original: Recall that it is proved in Gorham & Mackey (2017, Theorem 8) that, for κ ( x, y )
Revised: Recall that it is proved in Gorham & Mackey (2017, Theorem 8) that, for κ ( x, y )

Original: = ( c 2 + (cid:107) x − y (cid:107) 22 ) β with β ∈ ( − 1 , 0) and distantly dissipative p H with Lipschitz score functions, q
Revised: = ( c 2 + (cid:107) x − y (cid:107) 22 ) β with β ∈ ( − 1 , 0) and distantly dissipative p H with Lipschitz score functions, q

Original: H ⇒ p
Revised: H ⇒ p

Original: H if KSD κ ( q H , p H ) → 0 .
Revised: H if KSD κ ( q H , p H ) → 0 .

Original: The advertised result ( q ⇒ p if MKSD K k ( q, p ) → 0 ) now follows by the continuous mapping theorem as ∇ ψ ∗ is continuous.
Revised: The advertised result ( q ⇒ p if MKSD K k ( q, p ) → 0 ) now follows by the continuous mapping theorem as ∇ ψ ∗ is continuous.

Original: J L EMMAS
Revised: J L EMMAS

Original: Lemma 13.
Revised: Lemma 13.

Original: Let K ∇ ψ ∗ ( η, η (cid:48) )
Revised: Let K ∇ ψ ∗ ( η, η (cid:48) )

Original: (cid:44) K ( ∇ ψ ∗ ( η ) , ∇ ψ ∗ ( η (cid:48) )) .
Revised: (cid:44) K ( ∇ ψ ∗ ( η ) , ∇ ψ ∗ ( η (cid:48) )) .

Original: The mirrored updates g ∗ q t ,K in (10) can be equivalently expressed as
Revised: The mirrored updates g ∗ q t ,K in (10) can be equivalently expressed as

Original: [Equation]
Revised: [Equation]

Original: Proof
Revised: Proof

Original: We will use the identity proved in Lem.
Revised: We will use the identity proved in Lem.

Original: [Equation]
Revised: [Equation]

Original: by change-of-variable formula)
Revised: (by change-of-variable formula)

Original: [Equation]
Revised: [Equation]

Original: (by applying Lem.
Revised: (by applying Lem.

Original: 14 to each row of K ( · , θ ) )
Revised: 14 to each row of K ( · , θ ) )

Original: [Equation]
Revised: [Equation]

Original: where A : ,j denotes the j -th column of a matrix A .
Revised: where A : ,j denotes the j -th column of a matrix A .

Original: Lemma 14.
Revised: Lemma 14.

Original: For a strictly convex function ψ ∈ C 2 : R d → R and any vector-valued g ∈ C 1 : R d → R d , the following relation holds:
Revised: For a strictly convex function ψ ∈ C 2 : R d → R and any vector-valued g ∈ C 1 : R d → R d , the following relation holds:

Original: [Equation]
Revised: [Equation]

Original: Proof
Revised: Proof

Original: By the product rule of differentiation:
Revised: By the product rule of differentiation:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: This already gives us the ﬁrst term on the right side.
Revised: This already gives us the ﬁrst term on the right side.

Original: Next, we have
Revised: Next, we have

Original: [Equation]
Revised: [Equation]

Original: Plugging the above relation into (22) proves the claimed result
Revised: Plugging the above relation into (22) proves the claimed result

Original: [Equation]
Revised: [Equation]

Original: Lemma 15 (Liu, 2017, Lemma A.1) .
Revised: Lemma 15 (Liu, 2017, Lemma A.1) .

Original: Let A be a square matrix
Revised: Let A be a square matrix, and 0 < (cid:15) <

Original: [Equation]
Revised: [Equation]

Original: and 0 < (cid:15) < 12 (cid:107)

Original: A + A (cid:62) (cid:107) op .
Revised: (cid:107) op .

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

Original: where (cid:107) · (cid:107) F denotes the Frobenius norm of a matrix.
Revised: where (cid:107) · (cid:107) F denotes the Frobenius norm of a matrix.

Original: Lemma 16.
Revised: Lemma 16.

Original: Let K be a matrix-valued kernel and H K be the corresponding RKHS.
Revised: Let K be a matrix-valued kernel and H K be the corresponding RKHS.

Original: Then, for any f ∈ H K ( f is vector-valued), we have
Revised: Then, for any f ∈ H K ( f is vector-valued), we have

Original: [Equation]
Revised: [Equation]

Original: where (cid:107) · (cid:107) op denotes the operator norm of a matrix induced by the vector 2-norm.
Revised: where (cid:107) · (cid:107) op denotes the operator norm of a matrix induced by the vector 2-norm.

Original: Proof
Revised: Proof

Original: We ﬁrst bound the (cid:107) f
Revised: We ﬁrst bound the (cid:107) f

Original: ( x ) (cid:107) 2 as
Revised: ( x ) (cid:107) 2 as

Original: [Equation]
Revised: [Equation]

Original: The second result follows similarly,
Revised: The second result follows similarly,

Original: [Equation]
Revised: [Equation]

Revised: = x (cid:107) op .