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: This paper studies discretized mirror Langevin diffusions under the assumption of mirror-log Sobolev inequalities as well as relative smoothness assumptions. In particular, they show that a new backward method obtains perhaps the best dependence in terms of finite sample guarantees in this setting. On the other hand, Euler-Maruyama discretization has an irreducible bias and a different forward discretization has an exponential dependence on the diameter of the mirror set.