Efficient constrained sampling via the mirror-Langevin algorithm

1 Introduction . We consider the following canonical sampling problem . Let V : Rd → R∪ { ∞ } be a convex function and let π be the density on Rd which is proportional to exp ( −V ) . The task is to output a sample which is ( approximately ) distributed according to π , given query access to the gradients of V . The sampling problem has attracted considerable attention recently within the machine learning and statistics communities . This renewed interest in sampling is spurred , on one hand , by a wide breadth of applications ranging from Bayesian inference [ RC04 , DM+19 ] and its use in inverse problems [ DS17 ] , to neural networks [ GPAM+14 , TR20 ] . On the other hand , there is a deep and fruitful connection between sampling and the field of optimization , introduced in the seminal work [ JKO98 ] , which has resulted in the rapid development of sampling algorithms inspired by optimization methods such as : proximal/splitting methods [ Ber18 , Wib18 , Wib19 , SKL20 ] , coordinate descent [ DLLW21a , DLLW21b ] , mirror descent [ HKRC18 , CLGL+20 , ZPFP20 ] , Nesterov ’ s accelerated gradient descent [ CCBJ18 , MCC+21 , DRD20 ] , and Newton methods [ MWBG12 , SBCR16 , CLGL+20 , WL20 ] . To describe this connection , we recall the Langevin diffusion , which is the solution to the following stochastic differential equation ( SDE ) : dXt = −∇V ( Xt ) dt+ √ 2 dBt . ( LD ) Under standard assumptions on the potential V , the SDE is well-posed and it converges in distribution , as t→∞ , to its unique stationary distribution π . Thus , once suitably discretized , it yields a popular algorithm for the sampling problem . The Langevin diffusion is classically studied using techniques from Markov semigroup theory [ see , e.g . BGL14 , Pav14 ] , but there is a more insightful perspective which views the diffusion ( LD ) through the lens of optimization [ JKO98 ] . Specifically , if µt denotes the law of the process ( LD ) at time t , then the curve ( µt ) t≥0 is the gradient flow of the KL divergence DKL ( · ‖ π ) in the Wasserstein space of probability measures . This perspective has not only inspired 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . new analyses of Langevin [ CB18 , Wib18 , DMM19 , VW19 ] , but has also emboldened the possibility of bringing to bear the extensive toolkit of optimization onto the problem of sampling ; see the references listed above . However , the vanilla Langevin diffusion notably fails when the support of the target distribution π is not all of Rd . This task of constrained sampling , named in analogy to constrained optimization , arises in applications such as latent Dirichlet allocation [ BNJ03 ] , ordinal data models [ JA06 ] , survival time analysis [ KM06 ] , regularized regression [ CEAM+12 ] , and Bayesian matrix factorization [ PBJ14 ] . Despite such a broad range of applications , the constrained sampling problem has proven to be challenging . In particular , most prior works have focused on domain-specific algorithms [ GSL92 , PP14 , LS16 ] , and the first general-purpose algorithms for this task are recent [ BDMP17 , BEL18 ] . In this work , we tackle the constrained sampling problem via the mirror-Langevin algorithm ( MLA ) . MLA is a discretization of the mirror-Langevin diffusion [ HKRC18 , ZPFP20 ] , which is the sampling analogue of mirror descent . Namely , if φ : Rd → R∪ { ∞ } is a mirror map , then the mirror-Langevin diffusion is the solution to the SDE Xt = ∇φ ? ( Yt ) , dYt = −∇V ( Xt ) dt+ √ 2 [ ∇2φ ( Xt ) ] 1/2 dBt . ( MLD ) Technical motivation . Recently , Zhang et al . [ ZPFP20 ] analyze an Euler-Maruyama discretization of MLD ; see ( 2.2 ) for details . The most curious aspect of their result is that their convergence guarantee has a bias term that does not vanish even when step size tends to zero and the number of iterations tends to infinity . Moreover , they conjecture that this bias term is unavoidable . This is in contrast to known results for standard Langevin , which raises the main question of this paper : Can a different discretization of MLD lead to a vanishing bias ? Our contributions . We propose a new discretization of the mirror-Langevin diffusion , given in ( MLA ) and illustrated in Figure 1 . Our proposed discretization has the same cost as the standard Euler-Maruyama discretization of MLD in terms of the number of queries to the gradient oracle for V . We remark that our scheme for the case φ = ‖·‖2/2 recovers the unadjusted Langevin algorithm . The most important aspect of our result is that the bias of our algorithm vanishes as the step size tends to zero unlike the result by Zhang et al . [ ZPFP20 ] . By adapting the analysis of Durmus et al . [ DMM19 ] , we provide a clean convergence analysis of our algorithm which theoretically validates our discretization scheme . Notably , our analysis only requires standard assumptions/definitions which are well-studied in optimization . In particular , we establish a stronger link between sampling and optimization without relying on technical assumptions of Zhang et al . [ ZPFP20 ] ( e.g . commutation conditions for Hessians ; see ( A5 ) therein ) . Moreover , our analysis combines ideas from optimization with the calculus of optimal transport . In particular , we establish a new generalization of a celebrated fact , namely that the entropy functional is displacement convex along Wasserstein geodesics , to the setting of Bregman divergences ( Theorem 4 ) . This inequality has interesting consequences in its own right ; as we discuss in Corollary 1 , our result already implies the transport inequality of Cordero-Erausquin [ CE17 ] . We provide convergence guarantees for the following classes of potentials : ( 1 ) convex and relatively smooth ( Theorem 1 ) ; ( 2 ) strongly relatively convex and relatively smooth ( Theorem 2 ) ; and ( 3 ) convex and Lipschitz ( Theorem 3 ) . Our results largely match state-of-the-art results for the discretization of the Langevin algorithm for unconstrained sampling . Our work paves the way for the practical deployment of mirror-Langevin methods for sampling applications , paralleling the successes of mirror descent in optimization [ NY83 , JN11 , Bub15 ] . In Section 5 , we demonstrate the strength of our convergence guarantees compared with the previous works [ BDMP17 , BEL18 ] in an application to Bayesian logistic regression ; further applications are given in Appendix E. We also corroborate our theoretical findings with numerical experiments . Other related works . Recently , a few works have proposed modifications of the Langevin algorithm for the task of constrained sampling . Bubeck et al . [ BEL18 ] studied the projected Langevin algorithm ( PLA ) , which simply projects each step of the Langevin algorithm onto dom ( V ) . A different approach was taken in Brosse et al . [ BDMP17 ] , which applies the Langevin algorithm to a smooth approximation of V given by the Moreau-Yosida envelope . The latter approach was later interpreted and further analyzed by Salim and Richtarik [ SR20 ] using the primal-dual optimality framework from convex optimization . A different line of work , more closely related to ours , uses a mirror map to change the geometry of the sampling problem [ HKRC18 , CLGL+20 , ZPFP20 ] . In particular , the mirror-Langevin diffusion ( MLD ) was first introduced in an earlier draft of [ HKRC18 ] , as well as in [ ZPFP20 ] . The diffusion was further studied in [ CLGL+20 ] , which provided a simple convergence analysis in continuous time using the sampling analog of Polyak-Łojasiewicz inequalities [ KNS16 ] . We also remark that the idea of changing the geometry via a mirror map also played an crucial role for the problem of sampling from the uniform distribution over a polytope [ KN12 , LV17 , CDWY18 , LV18 , GN20 , LLV20 ] . Lastly , our work follows the trend of applying ideas from optimization to the task of sampling . Specifically , our analysis adopts the framework of relative convexity and smoothness , which was advocated as a more flexible framework for optimization in [ BBT17 , LFN18 ] . 2 The mirror-Langevin algorithm . 2.1 Background . In this section , we list basic definitions and assumptions that we employ in this work . Convex functions of Legendre type . Throughout , we assume familiarity with the basic notions of convex analysis [ see e.g . Roc70 , BL06 ] . Definition 1 ( Convex functions of Legendre type [ Roc70 , §26 ] ) . A proper convex lower semicontinuous function φ : Rd → R ∪ { ∞ } is of Legendre type if ( i ) Q : = int ( dom ( φ ) ) 6= ∅ , ( ii ) φ is strictly convex and differentiable on Q , and ( iii ) limk→∞‖∇φ ( xk ) ‖ =∞ whenever { xk } k∈N is a sequence in Q converging to ∂Q . The key properties of convex functions of Legendre type are listed below : • The subdifferential ∂φ is single-valued and hence ∂φ = { ∇φ } [ Roc70 , Theorem 26.1 ] . • φ is a convex function of Legendre type if and only if its Fenchel conjugate φ ? is a convex function of Legendre type [ Roc70 , Theorem 26.5 ] . • The gradient ∇φ forms a bijection between int ( dom ( φ ) ) and int ( dom ( φ ? ) ) with ∇φ ? = ( ∇φ ) −1 [ Roc70 , Theorem 26.5 ] . We refer readers to [ Roc70 , §26 ] for more details . We henceforth assume that our mirror map φ is a convex function of Legendre type . The natural notion of “ distance ” associated with the mirror map φ is given by the Bregman divergence [ see , e.g . Bub15 , §4 ] : Definition 2 ( Bregman divergence [ Bre67 ] ) . For a convex function φ of Legendre type , the Bregman divergence Dφ ( · , · ) associated to φ is defined as Dφ ( x , y ) : = φ ( x ) − φ ( y ) − 〈∇φ ( y ) , x− y〉 , ∀x , y ∈ Q : = int ( dom ( φ ) ) . The Bregman divergence behaves like a squared distance ; indeed , as x→ y a Taylor expansion shows that Dφ ( x , y ) ∼ 12 〈x− y , ∇ 2φ ( y ) ( x− y ) 〉 . We refer to [ Bub15 , §4 ] for other basic properties of the Bregman divergence . Hereinafter , when we use a Bregman divergence , we implicitly assume that its associated function φ is a mirror map . An important case to keep in mind is the mirror map φ = ‖·‖ 2 2 , where ‖·‖ denotes the Euclidean norm , in which case the Bregman divergence simply becomes Dφ ( x , y ) = 12‖x− y‖ 2 . Self-concordance . We recall the definition of self-concordance , which has been extensively used in applications such as interior-point methods [ NN94 ] . Given a C2 strictly convex function φ , the local norm at x ∈ int ( dom ( φ ) ) with respect to φ is defined as ‖u‖∇2φ ( x ) = √ 〈∇2φ ( x ) u , u〉 for all u ∈ Rd . The dual local norm at x ∈ int ( dom ( φ ) ) with respect to φ is ‖u‖ [ ∇2φ ( x ) ] −1 = √〈 [ ∇2φ ( x ) ] −1u , u 〉 for all u ∈ Rd . Definition 3 ( Self-concordant function [ Nes18 , §5.1.3 ] ) . We say that a C3 convex function φ is self-concordant with a constant Mφ ≥ 0 if for any x ∈ int ( dom ( φ ) ) |∇3φ ( x ) [ u , u , u ] | ≤ 2Mφ ‖u‖3∇2φ ( x ) for all u ∈ R d . Relative convexity/smoothness . We recall the following definitions : Definition 4 ( Relative convexity [ BBT17 , LFN18 ] ) . V is α-convex relative to φ if V ( y ) ≥ V ( x ) + 〈∇V ( x ) , y − x〉+ αDφ ( y , x ) ∀x , y ∈ Q . Definition 5 ( Relative smoothness [ BBT17 , LFN18 ] ) . V is β-smooth relative to φ if V ( y ) ≤ V ( x ) + 〈∇V ( x ) , y − x〉+ βDφ ( y , x ) ∀x , y ∈ Q . We give some basic facts about these definitions in Appendix B . In the rest of this work , we assume that V ∈ C2 ( X ) where X : = int ( dom ( V ) ) , that X ⊆ Q , and X ∩Q 6= ∅ . Also , we assume that exp ( −V ) is integrable so that π is well-defined ; this holds if and only if V ( x ) ≥ a ‖x‖ − b for some a , b > 0 [ BGVV14 , Lemma 2.2.1 ] . Optimal transport . Given a lower semicontinuous cost function c : Rd × Rd → [ 0 , ∞ ] , we can define the optimal transport cost between two probability measures µ and ν on Rd to be inf { E c ( X , Y ) | X ∼ µ , Y ∼ ν } . ( 2.1 ) Here , the infimum is taken over pairs of random variables ( X , Y ) defined on the same probability space , with marginal laws µ and ν respectively . It is known that the infimum in ( 2.1 ) is always attained ; we refer to the standard introductory texts [ Vil03 , Vil09 , San15 ] for this and other basic facts in optimal transport . In this work , we are most concerned with the case when the cost function c is the Bregman divergence associated with a mirror map : Definition 6 ( Bregman transport cost ) . The Bregman transport cost is defined as Dφ ( µ , ν ) : = inf { EDφ ( X , Y ) | X ∼ µ , Y ∼ ν } . The Bregman transport cost was also studied in [ CE17 ] . In particular , when φ = ‖·‖ 2 2 , we obtain an important special case : Definition 7 ( 2-Wasserstein distance ) . The 2-Wasserstein distance W2 is defined as W 22 ( µ , ν ) : = inf { E [ ‖X − Y ‖2 ] | X ∼ µ , Y ∼ ν } . The W2 optimal transport cost indeed defines a metric over the space of probability measures on Rd with finite second moment [ Vil03 , Theorem 7.3 ] ; we refer to this metric space as the Wasserstein space . The W2 metric is particularly important because it arises from a formal Riemannian structure on the Wasserstein space . This perspective was introduced in [ Ott01 ] and applied to the Langevin diffusion in [ JKO98 , OV00 ] ; in particular , these latter two works justify the perspective of the Langevin diffusion as a gradient flow of the Kullback-Leibler divergence in Wasserstein space . A rigorous exposition to Wasserstein calculus can be found in [ AGS08 , Vil09 ] . Here , we give a brief and informal introduction to the calculation rules of optimal transport . For any regular curve of measures ( µt ) t≥0 , there is a corresponding family of tangent vectors ( vt ) t≥0 [ see AGS08 , Theorem 8.3.1 ] ; here , vt : Rd → Rd is a vector field on Rd . Also , if F is any wellbehaved functional defined over Wasserstein space , then at each regular measure µ one can define the Wasserstein gradient of F at µ , which we denote ∇W2F ( µ ) ; it is also a mapping Rd → Rd . Then , we have the calculation rule ∂tF ( µt ) = E〈∇W2F ( µt ) ( Xt ) , vt ( Xt ) 〉 for any regular curve of measures ( µt ) t≥0 with corresponding tangent vectors ( vt ) t≥0 , whereXt ∼ µt . We will use this calculation rule in Appendix C .

Summary: The paper under review is considering a new discretization for the MLD. The motivation comes from solving constrained optimization problem. Under proper assumption, the authors present convergence analysis for the algorithm in terms of the KL and Bregman divergence.  In order to derive the convergence analysis, the author derive displacement convexity property and treat/verify several technical conditions following previous works in the literature.