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 proposes a new mirror descent like Langevin algorithm to sample from exp(-V), where V is relatively smooth/convex. The algorithm is designed as a Forward Flow algorithm. The convergence guarantees are similar to those of the vanilla Langevin algorithm (ULA), but using some newly defined Bergman divergence instead of the Wasserstein 2 distance.