Mirror Langevin Monte Carlo: the Case Under Isoperimetry

1 Introduction . It has been widely recognized that optimization and sampling are deeply connected . One one hand , optimization can be viewed as performing sampling in the limit , and on the other , since the influential work of Jordan-Kinderlehrer-Otto [ 12 ] , Langevin dynamics takes on the interpretation as performing deterministic optimization ( gradient flow ) in the space of probability measures . This profoundly shapes the way we view and understand traditional MCMC sampling algorithms , deviating from the Markov semigroup path . While huge amount of progress has been made on the optimization front in the past few decades , its sampling counterpart , finding far-reaching applications in Bayesian statistical inference and inverse problems , hasn ’ t been fully explored to leverage the advancements offered by the optimization toolbox . In this paper , we draw inspiration from mirror descent [ 17 ] and ask the question if there ’ s an analog of it that can adapt to geometries beyond the Euclidean case for Langevin diffusion , under isoperimetric inequalities such as Log-Sobolev for the target distribution , rather than Strong-Log-Concavity , where we recall the celebrated result of Bakry and Émery [ 3 ] prescribes that the latter implies the former . 1.1 Mirror Flow and Mirror Descent . In optimization , the extension to arbitrary geometry through the choice of a mirror map can often give better smoothness/strong convexity parameter dependence , or even handle cases where strong-convexity and Lipschitz gradient of f do not hold in the Euclidean geometry . The continuous limit of mirror descent can be written as dYt = rf ( Xt ) dt , Xt = r ⇤ ( Yt ) ( 1 ) which is equivalent to dYt/dt = d [ r ( Xt ) ] /dt = r2 ( Xt ) dXt/dt = rf ( Xt ) , therefore the mirror flow can be recast in the primal variable as dXt = ( r2 ( Xt ) ) 1rf ( Xt ) dt , akin to natural gradient flow , which preconditions the update to adapt to local geometry . Equation ( 1 ) makes it clear that the mirror descent update xk+1 = r ⇤ ( r ( xk ) hk+1rf ( xk ) ) = argmin x hx , rf ( xk ) i+ h 1 k+1D ( x , xk ) ⇤Work done while at Stanford University . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . is nothing more than the forward-discretized gradient descent in the dual y-space through mirror mapping . Here D ( x , xk ) = ( x ) ( xk ) r ( xk ) > ( x xk ) 0 is the Bregman divergence and ⇤ is the Fenchel conjugate of . We assume that the mirror map : Rd ! R is of Legendre type2 and strictly convex throughout . Common choices include ( x ) = kxk22/2 , which reduces to classical gradient descent as xk+1 = xk hk+1rf ( xk ) , and ( x ) = P i xi log ( xi ) , which gives multiplicative weight update . In the special case when = f , one readily recovers Newton ’ s method . 1.2 Mirror Langevin Dynamics and Mirror Langevin Monte Carlo . For sampling , we consider the Mirror Langevin stochastic differential equation ( SDE ) where for Yt = r ( Xt ) , and target distribution ⇡ = e f , dYt = rf ( r ⇤ ( Yt ) ) dt+ p 2 [ r2 ⇤ ( Yt ) ] 1dWt ( 2 ) for Wt the standard Brownian motion in Rd . If is three-times-differentiable , it is equivalent to dXt = [ r2 ( Xt ) ] 1 ⇥Tr r 3 ( Xt ) [ r 2 ( Xt ) ] 1 rf ( Xt ) ⇤ dt+ p 2 [ r2 ( Xt ) ] 1dWt ( 3 ) and the corresponding Euler-Maruyama ( EM ) discretized version ( in dual y-space ) becomes xk+1 = r ⇤ ⇣ r ( xk ) hk+1rf ( xk ) + p 2hk+1 [ r2 ( xk ) ] · zk+1 ⌘ ( 4 ) for hk+1 the stepsize and zk+1 2 Rd an independent standard Gaussian random vector , where we used r2 ⇤ ( Yt ) = [ r2 ( Xt ) ] 1 . No particular warm start is assumed for initialization . It is worth noting that while the continuous dynamics in the primal X-space involves the 3rd-order derivative tensor , the implementation of ( 4 ) only requires access to a gradient oracle for f . For ( x ) = kxk22/2 , one recovers the classical ( overdamped ) Langevin dynamics , whose Euler-Maruyama discretization is ULA : xk+1 = xk hk+1rf ( xk ) + p 2hk+1 · zk+1 . In the case when there is no closedform expression , the inversion r ⇤ can be solved numerically ( and therefore approximately ) using r ⇤ ( z ) = argmaxx { z > x ( x ) } , which is a convex optimization problem . A derivation for the equivalence between ( 2 ) and ( 3 ) , along with the fact that ( 3 ) has ⇡ = e f as the stationary distribution are given in Appendix A . It is evident ( and reminiscent of the classical Langevin algorithm ) that the discretized algorithm will converge to a biased limit ⇡h under mild regularity conditions . A Metropolis-Hastings step can be applied on top to correct for the bias but we focus on the unadjusted case in this paper . The main question we aim to address in this paper is – what is the non-asymptotic rate of convergence for ( 2 ) using different discretization schemes , under functional inequalities such as Log-Sobolev Inequality ( LSI ) , which encompasses broader classes of distributions compared to the more restrictive and well-studied Strong-Log-Concavity ( SLC ) setting . In particular , it is known that LSI is preserved under bounded perturbation [ 10 ] and Lipschitz mapping , which capture cases when f is far from convex , e.g. , multi-modal such as Gaussian mixtures . 2 Related Work . Various discretizations of underdamped , overdamped Langevin dynamics in the Euclidean setting under LSI [ 14 , 23 , 15 ] and SLC [ 8 , 5 , 7 , 19 , 16 ] are the main focus of a series of developments , for which non-asymptotic error bounds are established for several metrics including KL , Wasserstein and TV distance . Convergence of discretized algorithm under LSI , in general , introduces considerably more challenges , as commonly used synchronous/reflection coupling techniques do not apply . For Langevin dynamics working under non-Euclidean geometry , an earlier proposal was made in [ 11 ] where an algorithm was designed to converge to ( r ) # ⇡ , being essentially a change of measure from the classical Langevin dynamics . Crucially , their dynamics is different in that the diffusion term isn ’ t scaled as the one we consider to take into account the Riemannian metric structure induced by . The study of dynamics ( 2 ) was initiated in [ 24 ] under a relaxed-SLC assumption , where the authors show convergence of ( 4 ) to a Wasserstein ball with non-vanishing bias . [ 1 ] studied under a similar relaxed-SLC assumption , but with a different discretization scheme as opposed to ( 4 ) , 2so that r is invertible and r ⇤ = ( r ) 1 is a single-valued mapping that makes sense for ( 1 ) where the bias decreases to zero with diminishing stepsize . Closer to our work is [ 6 ] , in which the authors investigated convergence of the continuous process ( 2 ) under functional inequalities resembling the ones we consider and focused on = f , leaving the analysis of discretized algorithm for future work . In this case one does not need to resort to smoothness assumption etc . for deriving contraction/handling bias issues – it is only when studying stable discretization schemes that they become necessary ( and slows down convergence from exponential to polynomial ) . Recent work of [ 13 ] leverages the mean-square framework , which heavily exploits the contraction property of the dynamics , to show that under a modified self-concordance property , algorithm ( 4 ) in fact converges without asymptotic bias . This particular property is something we do not have here , therefore it ’ s likely new ideas are needed if one were to improve the analysis of ( 4 ) in this setting . 3 Notation and Assumptions . Notation We use r · v ( x ) = P i @ vi ( x ) @ xi 2 R to denote the divergence operator for a vector field v : Rd ! Rd , and for a matrix-valued function G : Rd ! Rd⇥d , we let r · G ( x ) 2 Rd with the i-th element as P j @ G ( x ) ij @ xj . We denote the Laplacian operator as where ( x ) = Tr ( r2 ( x ) ) . Moreover , let hr2 , G ( x ) i = r · ( r · G ( x ) ) = P i , j @ 2 Gij ( x ) @ xi @ xj 2 R. The norm induced by a positive definite G is defined as kzk2 G : = z > Gz . Wasserstein-p distance is defined as Wp ( ⇢ , ⇡ ) = infx⇠⇢ , x0⇠⇡ E [ kx x0kp ] 1/p for p 1 . From Monge-Kantorvich ’ s duality W1 ( ⇢ , ⇡ ) = supkfklip1 E⇢ [ f ] E⇡ [ f ] and monotonicity property W1 ( ⇢ , ⇡ ) W2 ( ⇢ , ⇡ ) , bound on Wasserstein-2 metric captures many functions of potential interest . We define the pushforward measure ⇢̄ = r # ⇢ as ⇢̄ ( B ) = ⇢ ( r 1 ( B ) ) for every borel set B . We state the assumptions that we will make below , some of which we will relax later . Assumption 1 ( ⇣-Self-Concordance ) . There exists a constant ⇣ 0 such that the conjugate mirror map ⇤ satisfies that 8y , u , s , v , |r 3 ⇤ ( y ) [ u , s , v ] | 2⇣ · ( u > r2 ⇤ ( y ) u ) 1/2 ( s > r2 ⇤ ( y ) s ) 1/2 ( v > r2 ⇤ ( y ) v ) 1/2 . Moreover , this property is preserved under Fenchel conjugation ( with the same parameter ) , affine transformation and summation [ 18 ] . Many natural barrier and entropy functions ( e.g. , log-barrier ) satisfy such self-concordance property . This also guarantees solution of the continuous dynamics ( 2 ) , cf . Appendix A of [ 24 ] . One can show that this affine-invariant condition implies a form of Hessian stability as M 1r2 ( x ) r2 ( x0 ) Mr 2 ( x ) for some M 1 ( cf . Lemma 4 ) , entailing that the underlying geometry isn ’ t rapidly changing . This form of self-concordance also appears in interior point method in optimization and previous work on Mirror Langevin [ 1 ] and suggests that locally the function behaves like a quadratic . Assumption 2 ( -Mirror-Log-Sobolev ) . The target distribution ⇡ satisfies Mirror LSI with constant w.r.t a given mirror map , i.e. , for every locally lipschitz function g , it holds that ⇡ satisfies 2 Z krg ( x ) k2 [ r2 ( x ) ] 1d⇡ Z g ( x ) 2 log g ( x ) 2d⇡ ✓Z g ( x ) 2d⇡ ◆ log ✓Z g ( x ) 2d⇡ ◆ ( 5 ) taking g ( x ) = p d⇢ ( x ) /d⇡ ( x ) one gets that for all ⇢ , H⇡ ( ⇢ ) : = Z ⇢ ( x ) log ⇢ ( x ) ⇡ ( x ) dx 1 2 Z ⇢ ( x ) r log ⇢ ( x ) ⇡ ( x ) 2 [ r2 ( x ) ] 1 dx = : 1 2 J ⇡ ( ⇢ ) . ( 6 ) This is the gradient-domination condition for KL-divergence in the Wasserstein metric , where the ( weighted ) Fisher information on the RHS is the squared norm of the gradient for KL divergence . LSI is an isoperimetric inequality that implies concentration for the distribution ( subgaussian tails on Lipschitz functions ) and plays an important role in many results in probability theory . A discussion of its implication in our context is further expanded in Section 4.1 . We note that both [ 1 ] and [ 24 ] require relative µ-strong convexity of f w.r.t , which means r2f ⌫ µr2 0 . Therefore this assumption based on LSI allows us to move away from convex potentials . The following two recover the familiar smoothness and Lipschitz condition when ( x ) = kxk22/2 . Assumption 3 ( L-Relative Lipschitz ) . For all x , it holds that f : Rd 7 ! R is differentiable with krf ( x ) k [ r2 ( x ) ] 1 L . ( 7 ) Assumption 4 ( -Relative Smooth ) . For all x , x0 2 dom ( ) , k [ r2 ( x ) ] 1rf ( x ) [ r2 ( x0 ) ] 1rf ( x0 ) kr2 ( x0 ) · kr ( x ) r ( x 0 ) k [ r2 ( x0 ) ] 1 . ( 8 ) One could show that ( 8 ) is slightly stronger than assuming r2f̃ ( y ) r2 ⇤ ( y ) for f̃ ( y ) = f ( r ⇤ ( y ) ) or that f̃ ( y ) is smooth , i.e. , r2f̃ ( y ) c · I ( cf . Lemma 7 ) . When ( x ) = kAxk22/2 , condition ( 8 ) reduces to krf ( x ) rf ( x0 ) k ( A > A ) 1 ·kx x0kA > A , for which the dual norm on the two sides makes sense as a generalization , whereas one could check that the relative smoothness as defined in [ 24 , 1 ] do not admit such a natural interpretation . We will comment more on this assumption in later sections on discretization . Condition ( 7 ) is also common in previous works [ 1 ] . Assumption 5 ( ↵-Strongly Convex ) . The mirror map is three-times differentiable , and let ↵ : = min ( r2 ) > 0 . All best known results under LSI [ 14 , 23 ] in the Euclidean setting assume third order smoothness on f ( i.e. , Lipschitz Hessian ) and Lipschitz gradient ( i.e. , L · I r2f L · I ) , sometimes with an additional disspativity assumption [ 15 ] , whereas we only require a weak notion of smoothness on f r ⇤ . The study of discretized sampling algorithms for non-smooth and non-convex potentials , to the best of our knowledge , is a road less traversed .

Summary: In this work the authors study the convergence of mirror descent langevin dynamics -- Langevin dynamics which have been transformed appropriately to satisfy certain constraints or evolve according to a specific non-euclidean geometry.   The dynamics can be characterised in various forms, including as a riemannian manifold overdamped langevin process with a specific choice of metric tensor -- but also as an SDE in the dual (unconstrained space) which is projected onto the original state space.  The paper considers the convergence to equilibrium of the mirror descent Langevin dynamics -- both for the original continuous time process, but also for various discretisations of the scheme.  This is done under the assumption of a Logarithmic Sobolev Inequality, thus generalising common convergence analysis which typically require strong-log-concavity assumptions.  It is demonstrated that the choice of discretisation plays a key role in achieving consistency under vanishing step-size for constraints which involve non-smooth contributions (e.g. constraint to lie within a box, etc), ultimately presenting a backward discretisation scheme which is far more stable to sudden changes in the geometry underlying the constraint.  These results are illustrated through a straightforward numerical example.