Efficient constrained sampling via the mirror-Langevin algorithm
Year: 2021

Abstract:
We propose a new discretization of the mirror-Langevin diffusion and give a crisp proof of its convergence. Our analysis uses relative convexity/smoothness and self-concordance, ideas which originated in convex optimization, together with a new result in optimal transport that generalizes the displacement convexity of the entropy. Unlike prior works, our result both (1) requires much weaker assumptions on the mirror map and the target distribution, and (2) has vanishing bias as the step size tends to zero. In particular, for the task of sampling from a log-concave distribution supported on a compact set, our theoretical results are significantly better than the existing guarantees.

heading: 1 Introduction
text: 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].

heading: 2 The mirror-Langevin algorithm

heading: 2.1 Background
text: 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.

heading: 2.2 Discretization of the mirror-Langevin diffusion
text: In order to turn a continuous-time diffusion such as (MLD) into an implementable algorithm, it is necessary to first discretize the stochastic process. The discretization considered in the prior works [ZPFP20, CLGL+20] is a simple Euler-Maruyama discretization: fixing η > 0, we define a sequence of iterates (Xk)k∈N via
∇φ(Xk+1) = ∇φ(Xk)− η∇V (Xk) + √ 2η [∇2φ(Xk)] 1/2 ξk , (2.2)
where (ξk)k∈N is a sequence of i.i.d. standard Gaussians in Rd. However, many other discretizations are possible. Indeed, in many machine learning applications, the most costly step is the evaluation of∇V , which may require a sum over a large training set, whereas the mirror map φ may be chosen to have a simple form. For the purpose of obtaining a more efficient sampling algorithm, it may therefore be a favorable trade-off to use a high-precision implementation of the diffusion step at the cost of additional computation time (which nonetheless does not require additional query access to the gradients of V ). Motivated by these considerations, we propose a new discretization (see Figure 1 for an illustration):
The mirror-Langevin algorithm (MLA):
Xk+1/2 := argmin x∈Q
[〈η∇V (Xk), x〉+Dφ(x,Xk)] , (MLA:1)
Xk+1 := ∇φ?(Wη) , where
{ dWt = √ 2 [∇2φ?(Wt)] −1/2 dBt ,
W0 = ∇φ(Xk+1/2) . (MLA:2)
In MLA:2, the stochastic processes (Wt)t≥0 are assumed to be driven by independent Brownian motions at each iteration. When φ = ‖·‖2/2, MLA recovers the unadjusted Langevin algorithm.
Practicality. Although MLA:2 is defined using the exact solution of an SDE (and we analyze the exact step MLA:2 for simplicity), it should be understood as capturing the idea of discretizing the diffusion step more finely (e.g., through multiple inner iterations of an Euler-Maruyama discretization) than the gradient step. This is indeed amenable to practical implementation since, as previously discussed, the gradient step is typically much more costly than the diffusion step. Moreover, this is justified by our theoretical results in the next section, which together with the conjecture of [ZPFP20] suggest that fine discretization of MLA:2 is potentially crucial for attaining vanishing bias. Nevertheless, it is indeed the case that a single iteration of MLA is more costly than a single step of the Euler-Maruyama discretization of MLD, and this represents a limitation of our work.
Remark 1. Our proposed discretization can be understood as a more faithful discretization of the mirror-Langevin diffusion (MLD), à la [GWS21]. It can also be understood as the forward-flow discretization of MLD in the interpretation of [Wib18]. Remark 2. In order for MLA:2 to be well-defined, we require assumptions on φ such that the diffusion (Wt)t≥0 is non-explosive, i.e., it does not exit int(dom(φ
?)) in finite time. This holds under very mild assumptions on φ; see [GK96]. For situations of interest, the assumptions of [GK96] can be checked directly.

heading: 3 Convergence rates of mirror-Langevin algorithm
text: First, we state the main assumptions which are used for our main results. Assumption 1 (Self-concordance of φ). We assume that φ is Mφ-self-concordant (Definition 3). Assumption 2 (Relative Lipschitzness). We assume that V is L-relatively Lipschitz with respect to φ, in the sense that ‖∇V (x)‖[∇2φ(x)]−1 ≤ L for all x ∈ X (see Section 2.1 for the definition of the local norm used here). Assumption 3 (Relative convexity and smoothness). We assume that V is α-convex relative to φ and β-smooth relative to φ, where 0 ≤ α ≤ β ≤ ∞ (Definitions 4 and 5). Remark 3. Our analysis works under weaker assumptions than those of [ZPFP20]. In particular, our analysis does not assume the moment condition on the Hessian ((A2) therein) and the bound on the commutator between ∇2φ and ∇2V ((A5) therein). Moreover, our analysis uses weaker (and more standard) definitions of self-concordance.
Throughout this section, we assume the conditions listed above, and we present convergence results for MLA under various sets of assumptions. Our first two results pertain to the smooth case, i.e. β <∞. Define the parameter
β′ := β + 2MφL.
One might wonder how large β′ is for typical applications. First, we only need φ to be self-concordant, not a self-concordant barrier. Hence the appearance of Mφ is typically not problematic; for instance, a log barrier with m constraints is O(1)-self-concordant. The smoothness parameter could be large and dimension-dependent in general. However, such situations are actually where our approach could be potentially advantageous. In Appendix E.2, we demonstrate an example where the smoothness parameter becomes much smaller by choosing φ carefully. Theorem 1 (Weakly convex case). Suppose that Assumptions 1, 2, 3 hold with α = 0 and β′ > 0. For a target accuracy > 0, let Xk ∼ µk denote the iterates of MLA with step size η = min{ 2β′d , 1 β′ }. Then, the following convergence rate holds for the mixture distribution µN := 1 N ∑N k=1 µk:
DKL(µN ‖ π) ≤ , provided that N ≥ 4β′dDφ(π, µ0) 2 max
{ 1,
2d
} . (3.1)
Proof. See Appendix C.2.
Theorem 2 (Strongly relatively convex case). Suppose that Assumptions 1, 2, 3 hold with α, β′ > 0.
1. (Convergence in Bregman transport cost) For a target accuracy > 0, let Xk ∼ µk denote the iterates of MLA with step size η = min{ α 2β′d , 1 β′ }. Then,
Dφ(π, µN ) ≤ , provided that N ≥ 2β′d α2 ln (2Dφ(π, µ0) ) max { 1, α 2d } .
2. (Convergence in KL divergence) For a target accuracy > 0, suppose that X0 ∼ µ0 satisfies Dφ(π, µ0) ≤ /α. Let Xk ∼ µk denote the iterates of MLA with step size η = min{ α 2β′d , 1 β′ }. Then, the following convergence rate holds for the mixture distribution
µN := 1 N ∑N k=1 µk,
DKL(µN ‖ π) ≤ , provided that N ≥ 4β′d
α max
{ 1,
2d
} .
Proof. See Appendix C.3.
Note that the initialization assumption Dφ(π, µ0) ≤ /α in the second assertion of Theorem 2 can be obtained from the first guarantee of Theorem 2. Chaining together the two parts of the theorem, we therefore obtain the following guarantee: suppose that we initialize MLA at a distribution µ0. Then, with step size η = min{ α 2β′d , 1 β′ }, we obtain
DKL ( 1 N1 N0+N1∑ k=N0+1 µk ∥∥∥ π) ≤ , provided that {N0 ≥ 2β′dα ln( 2αDφ(π,µ0) )max{1, 2d} , N1 ≥ 4β ′d α max { 1, 2d } .
Observe also that for the case φ = ‖·‖ 2
2 , Theorems 1 and 2 recover the corresponding convergence guarantees for the unadjusted Langevin algorithm [Corollary 7, and Corollaries 10 and 11 respectively in DMM19].1
Next, we present our guarantee for the non-smooth case β =∞. For this result, we assume that φ is strongly convex w.r.t. a norm ‖·‖ on Rd, and that V is L̃-Lipschitz in this norm. Since the norm of the gradient should be measured in the dual norm ‖·‖?, this means precisely that
‖∇V (x)‖? ≤ L̃ , for all x ∈ X . (3.2) We also note that the next result does not require self-concordance of φ. Theorem 3 (Non-smooth case). Assume φ is 1-strongly convex w.r.t. a norm ‖·‖ on Rd, that V is L̃-Lipschitz in this norm (in the sense of (3.2)), and that α = 0 (i.e., V is convex). For a target accuracy > 0, letXk ∼ µk denote the iterates of MLA with step size η = /L̃2. Then, the following convergence rate holds for the mixture distribution µN := 1 N ∑N k=1 µk:
DKL(µN ‖ π) ≤ , provided that N ≥ 2L̃2Dφ(π, µ1/2) 2 .
Proof. See Appendix C.4.
The assumption (3.2) is stronger than relative Lipschitzness: if V satisfies (3.2) and φ is 1-strongly convex w.r.t. ‖·‖, then V is L-relatively Lipschitz with respect to φ with L ≤ L̃. When ‖·‖ is the Euclidean norm and φ = ‖·‖ 2
2 , then we recover a special case of [DMM19, Corollary 14].
We now make a number of remarks about our result. Remark 4 (Implementing µN ). One can output a sample from µN by simply outputting one of the iterates {Xk}Nk=1 chosen uniformly at random. Remark 5 (Convergence in other metrics). Using standard inequalities, our results for convergence in KL divergence imply convergence in a number of other information divergences such as the total variation distance, see [Tsy09, §2.4].
When V is α-strongly convex (w.r.t. ‖·‖ 2
2 ), then the T2 transportation-cost inequality [Vil09, Theorem 22.14] αW 22 (µ, π) ≤ DKL(µ ‖ π) implies convergence in the W2 distance as well. In general, we do not have convergence inW2, but we can always obtain convergence with respect to a different optimal transport cost, namely, the Bregman transport cost DV associated with V . This is a consequence of Corollary 1 [also see CE17, Proposition 1], which asserts that DV (µ, π) ≤ DKL(µ ‖ π). Remark 6 (Dimension dependence). Ignoring for now the dependence on β′ (which may also have a dimension dependence depending on the application), the Bregman divergence Dφ(π, µ0) term is typically of size O(d) (see Section 5 for a particular instance of this). Thus, our overall dimension dependence is O(d2) for the weakly convex case and O(d) for the strongly convex case. Overall, this is a significantly better dependence on the dimension as compared to the previous works [BDMP17, BEL18]; we perform a detailed comparison in Section 5 for a specific setting.
We also remark that mirror descent has classically been used for dimension reduction by changing the geometry of the algorithm from `2 to `1. We investigate the possibility of doing the same for sampling in Appendix E.2.
1In this case, Mφ = 0, so the Lipschitz constant L does not enter the final result. In particular, it is not contradictory to assume strong convexity (α > 0).
Remark 7 (Comparison with [HKRC18]). [HKRC18] show that for strictly log-concave targets, there exists a good mirror map φ for which the pushforward of the target distribution via∇φ enjoys the same guarantees as ordinary Langevin. However, this result is only existential and gives no guidance on how to construct the mirror map. In contrast, our theorems hold for any choice of mirror map which satisfies our assumptions, and provide guidance on how to choose the mirror map. Also, our relative smoothness condition allows for potentials which blow up at the boundary of their domain (i.e. the target distribution vanishes near the boundary of its support), whereas this is forbidden by the assumptions of [HKRC18]. Lastly, our algorithm does not require computing the third derivative of the mirror map, whereas this is required for [HKRC18]. Remark 8 (Comparison with [ZPFP20]). [ZPFP20] performs an analysis of the Euler-Maruyama discretization of MLD, which we temporarily refer to as MLA′. Our result guarantees that for any desired accuracy , it is possible to choose the step size sufficiently small so that MLA achieves the target accuracy; in contrast, the result of [ZPFP20] only guarantees that MLA′ contracts to within a ball around π of radius O( √ d) (measured w.r.t. a modified Wasserstein distance).2 Moreover, [ZPFP20] conjecture that their bias term is unavoidable.
In light of our result, we believe that it is an interesting open question to resolve their conjecture. Their conjecture, if true, suggests that replacing MLA:2 in our algorithm by a single step of the EulerMaruyama discretization has disastrous effects on the convergence of the algorithm, and therefore provides further support for considering MLA instead of MLA′.
Recently, after the first draft of our paper was published online, Li et al. [LTVW21] gave an analysis of MLA′ under a subset of the assumptions of [ZPFP20] which indeed exhibits vanishing bias, provided that the relative strong convexity parameter of the potential is sufficiently large compared to the modified self-concordance parameter of the mirror map. It remains an open question to remove this latter restriction from their work, and moreover to obtain similar results under the more usual definitions of relative convexity/smoothness and self-concordance that we adopt in this work.
In order to generalize the discretization analysis from the vanilla Langevin algorithm to the mirrorLangevin algorithm, in the next section we prove a new displacement convexity result for the entropy with respect to the Bregman transport cost which may be of independent interest.

heading: 4 Convexity of the entropy with respect to the Bregman divergence
text: It is well-known that the entropy functional H is displacement convex along W2 geodesics [AGS08, Theorem 9.4.11]. In fact, this displacement convexity is crucial in showing that DKL(· ‖ π) = E+H is displacement convex (when π is log-concave), which in turn is used to analyze the convergence of LD to the target measure. Therefore, in order to understand the convergence of MLD, it is crucial to see if such a result is true when W2 is replaced by Dφ. We prove that indeed the displacement convexity-like property holds for H under Dφ-optimal couplings. Theorem 4 (“Convexity” of the entropy with respect to the Bregman divergence). Let µ, ν be probability measures on Rd and let X ∼ µ, Y ∼ ν be coupled according to the Bregman transport cost Dφ(µ, ν). Then, it holds that
H(ν) ≥ H(µ) + E 〈 [∇W2H(µ)](X), Y −X 〉 .
As a corollary, we can use the calculus of optimal transport in order to recover the transportation-cost inequality of [CE17]. In the following, we do not carry out the approximation arguments necessary to make the proof fully correct because a rigorous proof of the statement is already given in [CE17].3 Rather, our main purpose in giving this argument is simply to point out the convexity principle which underlies the transport inequality. Corollary 1 ([CE17, Proposition 1]). For any probability measure µ on Rd,
DKL(µ ‖ π) ≥ DV (µ, π).
Proof Sketch. Let (X,Y ) be optimally coupled according to the Bregman transport cost DV (·, ·) between µ and π. We decompose DKL(µ ‖ π) = EV (X) +H(µ). On one hand, the first term is
EV (X) = EV (Y ) + E〈∇V (Y ), X − Y 〉+ EDV (X,Y ). 2Notably, the radius of this ball is comparable to the distance at initialization. 3In fact, the proof of [CE17] does not even require convexity of V .
On the other hand, the convexity result (Theorem 4) shows that
H(µ) ≥ H(π) + E〈∇W2H(π)(Y ), X − Y 〉. Putting these together, we obtain
DKL(µ ‖ π) ≥ EV (Y ) + E〈∇V (Y ), X − Y 〉+ EDV (X,Y ) +H(π) + E〈∇W2H(π)(Y ), X − Y 〉 = DKL(π ‖ π) + E 〈 [∇V +∇W2H(π)](Y ), X − Y 〉 + EDV (X,Y ) = EDV (X,Y ),
since ∇V +∇W2H(π), the W2 gradient of DKL(· ‖ π) at π, is zero. This proves the result.

heading: 5 Application to Bayesian logistic regression
text: In this section, we apply our main result to Bayesian logistic regression. For more applications of our result, such as the possibility of dimension reduction and sampling from non-smooth distributions, we refer the readers to Appendix E.
We recall the setting of Bayesian logistic regression: we observe pairs (Xi, Yi), i = 1, . . . , n, where Xi ∈ Rd and Yi ∈ {0, 1}. The data are assumed to follow the model
Yi ∼ Bernoulli ( exp 〈θ,Xi〉 1 + exp 〈θ,Xi〉 ) , independently for i = 1, . . . , n. (5.1)
Here, the parameter θ itself is assumed to be a random variable taking values in Rd. If we assume that θ has a prior density λ with respect to Lebesgue measure, then the posterior distribution is
π(θ) ∝ λ(θ) exp [ n∑ i=1 ( Yi 〈θ,Xi〉 − ln(1 + exp 〈θ,Xi〉) )] .
Since it may be computationally infeasible to explicitly compute the normalizing constant for the posterior distribution, we turn towards sampling algorithms.
When we take a prior λ which has full support on Rd, e.g. a Gaussian prior, then we may apply off-the-shelf methods such as the Langevin diffusion (LD). However, if we choose a prior which has compact support, then the unadjusted Langevin algorithm is no longer an acceptable option because it outputs samples outside the support of the posterior. In this case, we must turn to other methods, such as the projected Langevin algorithm [BEL18]. Here, we explore the use of the mirror-Langevin algorithm (MLA) for constrained sampling.
For the rest of this section, we will focus on a particular problem for concreteness and interpretability: we consider the uniform prior λ on the `∞ ball [−1, 1]d. By duality, this is an attractive model when the data (Xi) n i=1 have small `1-norm, i.e., are approximately sparse. A natural choice of mirror map for this problem is the logarithmic barrier
φ(θ) = d∑ i=1 ( ln
1
1− θ[i] + ln
1
1 + θ[i]
) ,
where we use θ[·] to denote the coordinates of θ ∈ Rd. Then, φ is 1-self-concordant ([Nes18, §5.1.3]). We remark that the separability of the mirror map in this example implies that the diffusion step MLA:2 can be simulated in O(d) steps, rather than O(d2).
For this setting, we compare the guarantees of MLA with the Projected Langevin Algorithm (PLA) [BEL18] and the Moreau-Yosida unadjusted Langevin algorithm (MYULA) [BDMP17]; see Table 1. The details of the comparison are given in Appendix E.1.
We also perform a numerical experiment to compare the practical performance of MLA with PLA. We take θ? := (0.9, . . . , 0.9) ∈ R10 as the ground truth, and we generate 1000 i.i.d. pairs (Xi, Yi) where Xi is sampled uniformly from the `1 ball and Yi is generated from Xi according to (5.1) with θ = θ?. We generate 30 samples using both MLA and PLA (both with step size η = 0.005). At each iteration, we average the samples to obtain an estimate θk for the posterior mean, and we plot the error ‖θk − θ?‖2 in Figure 2, averaged over 10 trials. We implement MLA:2 by performing 10 inner iterations of an Euler-Maruyama discretization.
Acknowledgments. Kwangjun Ahn was supported by graduate assistantship from the NSF Grant (CAREER: 1846088) and by Kwanjeong Educational Foundation. Sinho Chewi was supported by the Department of Defense (DoD) through the National Defense Science & Engineering Graduate Fellowship (NDSEG) Program. We thank Philippe Rigollet for helpful suggestions which greatly improved the presentation of this paper.

Review summaries:
This paper proposes and studies a new discretization scheme of the mirror-Langevin diffusion proposed in Zhang et al. (2020). The results build on the assumption of self-concordance of the mirror function and has vanishing bias as the step size tends to zero.

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. 

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.