Sampling with Mirrored Stein Operators

We introduce a new family of particle evolution samplers suitable for constrained domains and non-Euclidean geometries . Stein Variational Mirror Descent and Mirrored Stein Variational Gradient Descent minimize the Kullback-Leibler ( KL ) divergence to constrained target distributions by evolving particles in a dual space defined by a mirror map . Stein Variational Natural Gradient exploits non-Euclidean geometry to more efficiently minimize the KL divergence to unconstrained targets . We derive these samplers from a new class of mirrored Stein operators and adaptive kernels developed in this work . We demonstrate that these new samplers yield accurate approximations to distributions on the simplex , deliver valid confidence intervals in post-selection inference , and converge more rapidly than prior methods in large-scale unconstrained posterior inference . Finally , we establish the convergence of our new procedures under verifiable conditions on the target distribution . 1 INTRODUCTION . Accurately approximating an unnormalized distribution with a discrete sample is a fundamental challenge in machine learning , probabilistic inference , and Bayesian inference . Particle evolution methods like Stein variational gradient descent ( SVGD , Liu & Wang , 2016 ) tackle this challenge by applying deterministic updates to particles using operators based on Stein ’ s method ( Stein , 1972 ; Gorham & Mackey , 2015 ; Oates et al. , 2017 ; Liu et al. , 2016 ; Chwialkowski et al. , 2016 ; Gorham & Mackey , 2017 ) and reproducing kernels ( Berlinet & Thomas-Agnan , 2011 ) to sequentially minimize Kullback-Leibler ( KL ) divergence . SVGD has found great success in approximating unconstrained distributions for probabilistic learning ( Feng et al. , 2017 ; Haarnoja et al. , 2017 ; Kim et al. , 2018 ) but breaks down for constrained targets , like distributions on the simplex ( Patterson & Teh , 2013 ) or the targets of post-selection inference ( Taylor & Tibshirani , 2015 ; Lee et al. , 2016 ; Tian & Taylor , 2018 ) , and fails to exploit informative non-Euclidean geometry ( Amari , 1998 ) . In this work , we derive a family of particle evolution samplers suitable for target distributions with constrained domains and non-Euclidean geometries . Our development draws inspiration from mirror descent ( MD ) ( Nemirovskij & Yudin , 1983 ) , a first-order optimization method that generalizes gradient descent with non-Euclidean geometry . To sample from a distribution with constrained support , our method first maps particles to a dual space . There , we update particle locations using a new class of mirrored Stein operators and adaptive reproducing kernels introduced in this work . Finally , the dual particles are mapped back to sample points in the original space , ensuring that all constraints are satisfied . We illustrate this procedure in Fig . 1 . In Sec . 3 , we develop two algorithms – Mirrored SVGD ( MSVGD ) and Stein Variational Mirror Descent ( SVMD ) – with different updates in the dual space ; when only a single particle is used , MSVGD reduces to gradient ascent on the log dual space density , and SVMD reduces to mirror ascent on the log target density . In addition , by exploiting the connection between MD and natural gradient descent ( Amari , 1998 ; Raskutti & Mukherjee , 2015 ) , we develop a third algorithm – Stein Variational Natural Gradient ( SVNG ) – that extends SVMD to unconstrained targets with non-Euclidean geometry . In Sec . 5 , we demonstrate the advantages of our algorithms on benchmark simplex-constrained problems from the literature , constrained sampling problems in post-selection inference ( Taylor & Tibshirani , 2015 ; Lee et al. , 2016 ; Tian & Taylor , 2018 ) , and unconstrained large-scale posterior inference with the Fisher information metric . Finally , we analyze the convergence of our mirrored algorithms in Sec . 6 and discuss our results in Sec . 7 . Related work Our mirrored Stein operators ( 6 ) are instances of diffusion Stein operators in the sense of Gorham & Mackey ( 2017 ) , but their specific properties have not been studied , nor have they been used to develop sampling algorithms . There is now a large body of work on transferring algorithmic ideas from optimization to MCMC ( e.g. , Welling & Teh , 2011 ; Simsekli et al. , 2016 ; Dalalyan , 2017 ; Durmus et al. , 2018 ; Ma et al. , 2019 ) and SVGD-like sampling methods ( e.g. , Liu et al. , 2019a ; b ; Zhu et al. , 2020 ; Zhang et al. , 2020a ) . The closest to our work in this space is the recent marriage of mirror descent and MCMC . For example , Hsieh et al . ( 2018 ) propose to run Langevin Monte Carlo ( LMC , an Euler discretization of the Langevin diffusion ) in a mirror space . Zhang et al . ( 2020b ) analyze the convergence properties of the mirror-Langevin diffusion , Chewi et al . ( 2020 ) demonstrate its advantages over the Langevin diffusion when using a Newton-type metric , and Ahn & Chewi ( 2020 ) study its discretization for MCMC sampling in constrained domains . Relatedly , Patterson & Teh ( 2013 ) proposed stochastic Riemannian LMC for sampling on the simplex . Several modifications of SVGD have been proposed to incorporate geometric information . Riemannian SVGD ( RSVGD , Liu & Zhu , 2018 ) generalizes SVGD to Riemannian manifolds , but , even with the same metric tensor , their updates are more complex than ours : notably they require higher-order kernel derivatives , do not operate in a mirror space , and do not reduce to natural gradient descent when a single particle is used . They also reportedly do not perform well when with scalable stochastic estimates of ∇ log p. Stein Variational Newton ( SVN , Detommaso et al. , 2018 ; Chen et al. , 2019 ) introduces second-order information into SVGD . Their algorithm requires an often expensive Hessian computation and need not lead to descent directions , so inexact approximations are employed in practice . Our SVNG can be seen as an instance of matrix SVGD ( MatSVGD , Wang et al. , 2019 ) with an adaptive time-dependent kernel discussed in Sec . 4.4 , a choice that is not explored in Wang et al . ( 2019 ) and which recovers natural gradient descent when n = 1 unlike the heuristic kernel constructions of Wang et al . ( 2019 ) . None of the aforementioned works provide convergence guarantees , and neither SVN nor matrix SVGD deals with constrained domains . 2 BACKGROUND : MIRROR DESCENT AND NON-EUCLIDEAN GEOMETRY . Standard gradient descent can be viewed as optimizing a local quadratic approximation to the target function f : θt+1 = argminθ∈Θ∇f ( θt ) > θ + 12 t ‖θ − θt‖ 2 2 . When Θ ⊆ Rd is constrained , it can be advantageous to replace ‖ · ‖2 with a function Ψ that reflects the geometry of a problem ( Nemirovskij & Yudin , 1983 ; Beck & Teboulle , 2003 ) : θt+1 = argminθ∈Θ∇f ( θt ) > θ + 1 tΨ ( θ , θt ) . ( 1 ) We consider the mirror descent algorithm ( Nemirovskij & Yudin , 1983 ; Beck & Teboulle , 2003 ) which chooses Ψ to be the Bregman divergence induced by a strongly convex , essentially smooth1 function ψ : Θ → R ∪ { ∞ } : Ψ ( θ , θ′ ) = ψ ( θ ) − ψ ( θ′ ) −∇ψ ( θ′ ) > ( θ − θ′ ) . When Θ is a ( d + 1 ) - simplex { θ : ∑di=1 θi ≤ 1 and θi ≥ 0 for i ∈ [ d ] } , a common choice of ψ is the negative entropy ψ ( θ ) = ∑d+1 i=1 θi log θi , for θd+1 , 1− ∑d i=1 θd . The solution of ( 1 ) is given by θt+1 = ∇ψ∗ ( ∇ψ ( θt ) − t∇f ( θt ) ) , ( 2 ) where ψ∗ ( η ) , supθ∈Θ η > θ − ψ ( θ ) is the convex conjugate of ψ and ∇ψ is a bijection from Θ to dom ( ψ∗ ) with inverse map ( ∇ψ ) −1 = ∇ψ∗ . We can view the update in ( 2 ) as first mapping θt to ηt by∇ψ , applying the update ηt+1 = ηt − t∇f ( θt ) , and mapping back through θt+1 = ∇ψ∗ ( ηt+1 ) . 1ψ is continuously differentiable on the interior of Θ with ‖∇ψ ( θt ) ‖ → ∞ whenever θt → θ ∈ ∂Θ . Mirror descent can also be viewed as a discretization of the continuous-time dynamics dηt = −∇f ( θt ) dt , θt = ∇ψ∗ ( ηt ) , which is equivalent to the Riemannian gradient flow ( see App . A ) : dθt = −∇2ψ ( θt ) −1∇f ( θt ) dt , or equivalently , dηt = −∇2ψ∗ ( ηt ) −1∇ηtf ( ∇ψ∗ ( ηt ) ) dt , ( 3 ) where ∇2ψ ( θ ) and ∇2ψ∗ ( η ) are Riemannian metric tensors . In information geometry , the discretization of ( 3 ) is known as natural gradient descent ( Amari , 1998 ) . There is considerable theoretical and practical evidence ( Martens , 2014 ) showing that natural gradient works efficiently in learning . 3 STEIN ’ S IDENTITY AND MIRRORED STEIN OPERATORS . Stein ’ s identity ( Stein , 1972 ) is a tool for characterizing a target distribution P using a so-called Stein operator . We assume P has a differentiable density p with a closed convex support Θ ⊆ Rd . A Stein operator Sp takes as input functions g from a Stein set G and outputs mean-zero functions under p : Eθ∼p [ ( Spg ) ( θ ) ] = 0 , for all g ∈ G. ( 4 ) Gorham & Mackey ( 2015 ) proposed the Langevin Stein operator given by ( Spg ) ( θ ) = g ( θ ) > ∇ log p ( θ ) +∇ · g ( θ ) , ( 5 ) where g is a vector-valued function and ∇ · g is its divergence . For an unconstrained domain with Ep [ ‖∇ log p ( θ ) ‖2 ] < ∞ , Stein ’ s identity ( 4 ) holds for this operator whenever g ∈ C1 is bounded and Lipschitz by ( Gorham et al. , 2019 , proof of Prop . 3 ) . However , on constrained domains Θ , Stein ’ s identity fails to hold for many reasonable inputs g if p is non-vanishing or explosive at the boundary . Motivated by this deficiency and by a desire to exploit non-Euclidean geometry , we propose an alternative mirrored Stein operator , ( Mp , ψg ) ( θ ) = g ( θ ) > ∇2ψ ( θ ) −1∇ log p ( θ ) +∇ · ( ∇2ψ ( θ ) −1g ( θ ) ) , ( 6 ) where ψ is a strongly convex , essentially smooth function as in Sec . 2 with ( ∇2ψ ) −1 differentiable and Lipschitz on Θ . We derive this operator from the ( infinitesimal ) generator of the mirror-Langevin diffusion ( 19 ) in App . C. The following result , proved in App . I.1 , shows that Mp , ψ generates mean-zero functions under p whenever ∇2ψ−1 suitably cancels the growth of p at the boundary . Proposition 1 . Suppose that ∇2ψ ( θ ) −1∇ log p ( θ ) and ∇ · ∇2ψ ( θ ) −1 are p-integrable . If limr→∞ ∫ ∂Θr p ( θ ) ‖∇2ψ ( θ ) −1nr ( θ ) ‖2dθ = 0 for Θr , { θ ∈ Θ : ‖θ‖∞ ≤ r } and nr ( θ ) the outward unit normal vector2 to ∂Θr at θ , then Ep [ ( Mp , ψg ) ( θ ) ] = 0 if g ∈ C1 is bounded Lipschitz . Example 1 ( Dirichlet p , Negative entropy ψ ) . When θ1 : d+1 ∼ Dir ( α ) for α ∈ Rd+1+ , even setting g ( θ ) = 1 in ( 5 ) need not cause the identity to hold when some αj ≤ 1 . However , when ψ ( θ ) = ∑d+1 j=1 θj log θj , we show in App . B that the conditions of Prop . 1 are met for any α . Remarkably , the mirror-Langevin diffusion for our choice of ψ is the Wright-Fisher diffusion ( Ethier , 1976 ) which Gan et al . ( 2017 ) recently used to bound distances to Dirichlet distributions .

Summary: The paper introduces algorithms for sampling constrained distributions. Assuming the existence of a mirror map that maps from the primal space to the dual space, the authors introduce an algorithm MSVGD that can perform gradient descent in the dual space, and map back to the primal space. They also propose SVMD, which uses a different algorithm to perform gradient descent in the primal space itself. The authors experimentally validate their algorithms on simulated data and one non-simulated dataset (Fig 4b).