Mirror Langevin Monte Carlo: the Case Under Isoperimetry
Year: 2021

Abstract:
Motivated by the connection between sampling and optimization, we study a mirror descent analogue of Langevin dynamics and analyze three different discretization schemes, giving nonasymptotic convergence rate under functional inequalities such as Log-Sobolev in the corresponding metric. Compared to the Euclidean setting, the result reveals intricate relationship between the underlying geometry and the target distribution and suggests that care might need to be taken in order for the discretized algorithm to achieve vanishing bias with diminishing stepsize for sampling from potentials under weaker smoothness/convexity regularity conditions.

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

heading: 1.1 Mirror Flow and Mirror Descent
text: 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.

heading: 1.2 Mirror Langevin Dynamics and Mirror Langevin Monte Carlo
text: 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.

heading: 2 Related Work
text: 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.

heading: 3 Notation and Assumptions
text: 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.

heading: 4 Convergence Analysis

heading: 4.1 Continuous Time Process: Mirror LSI
text: One can get using LSI (and LSI alone) the exponential convergence in KL divergence for the continuous time process, which is a manifestation of convergence under Polyak-Łojasiewicz (PL) inequality from optimization. In optimization, steepest descent gradient flow curve y⇤
t 2 Rd is
defined as min yt hrf(yt), ẏti+ 1 2 kẏtk 2 which implies ẏ⇤ t = rf(y⇤ t ) and df(y⇤ t )/dt = krf(y⇤ t )k2, from which m-gradient dominance condition on f (weaker than strong convexity)
f(y⇤ t ) min
y
f(y)  1
m krf(y⇤ t )k2
gives exponential convergence for the objective function f as d(f(y⇤ t ) miny f(y))/dt 
m(f(y⇤ t ) miny f(y)). In sampling within the space of probability measures, the objective function(al) is replaced by H⇡(⇢t), and the seminal work of [12] (see also later treatments by [2, 22]) show that the density ⇢t 2 P(Rd) followed by the (overdamped) Langevin dynamics satisfies
⇢̇t = arg min vt tangent to ⇢t
Ex⇠⇢t [hrW2H⇡(⇢t)(x), vt(x)i] + 1
2 Ex⇠⇢t [kvt(x)k22]
where rW2H⇡(⇢t)(x) = rx @H⇡(⇢t) @⇢t (x) is the Wasserstein-2 gradient of KL divergence at ⇢t. Therefore
dH⇡(⇢⇤t )
dt = Ex⇠⇢⇤t
" rx @H⇡(⇢⇤t )
@⇢t (x)
2
2
# = Ex⇠⇢⇤t " rx ✓ 1 + log ✓ ⇢ ⇤ t (x)
⇡(x)
◆◆ 2
2
# ,
where for the second step we used @H⇡(⇢t) @⇢t = 1 + log(⇢t ⇡ ) [2]. The RHS becomes the negative of the Fisher information J ⇡ (⇢) therefore LSI (6) reduces to gradient dominance condition for KL divergence (in this case with r2 = I for classical Langevin dynamics), from which one can derive exponential convergence rate to the target distribution ⇡ as carried out in the optimization land sans bias. This casts Langevin diffusion as precisely gradient flow w.r.t KL divergence in Wasserstein metric in the space of probability measures, from which LSI (weaker than SLC) suffices for convergence. We show in the lemma below that there is an extension of this result for Mirror-LSI that gives exponential convergence for the continuous process (3).
Proposition 1 (Convergence under Mirror-LSI). Along the dynamics of (3), we have that H⇡(⇢t)  e 2 t H⇡(⇢0) under Assumption 2.
Proof. Using the PDE (22), which describes the density followed by (3) and the integration by parts formula R hr (x), v(x)idx = R (x)r · v(x)dx,
d dt H⇡(⇢t) =
Z d⇢t
dt log
⇢t ⇡ dx+
Z ⇡ 1
⇡
d⇢t
dt dx
= Z r · ⇣ ⇢t[r 2 ] 1r log ⇢t
⇡
⌘ log ⇢t
⇡ dx+
d
dt
Z ⇢tdx
= Z ⇢t r log ⇢t
⇡
2 [r2 ] 1 dx+ 0
 2 ·H⇡(⇢t)
where we used Mirror LSI in the last step. This implies from Grönwall’s inequality that H⇡(⇢t)  e 2 t H⇡(⇢0).
For ⇡ that satisfies Mirror LSI with constant , thanks to the strong convexity of r2 ⌫ ↵I , it also satisfies Talagrand’s inequality [21] with parameter · ↵, which means the Wasserstein-2 distance is upper bounded by KL divergence as ↵ 2 W2(⇢,⇡) 2  H⇡(⇢) for any ⇢. Such transportation-cost inequality has the interpretation of quadratic growth in the iterate space. Therefore Mirror LSI has the additional nice property of allowing us to translate guarantee in the objective value (i.e., KL divergence) to iterate space (that involves optimal coupling between iterates x ⇠ ⇢, x0 ⇠ ⇡).
Stability One can show that Mirror-LSI, similar to its Euclidean counterpart, is stable under bounded perturbation. Therefore if the potential function f of interest is not exactly relative-smooth w.r.t to a “nice" mirror map, it suffices for it to be close to one, in the sense made precise in Appendix B. This is not something one can hope for with strong-convexity assumption, where perturbation from convex function usually breaks the assumption. This could be especially useful when the potential takes a composite form of f + g, where one part is smooth and the other part isn’t. It is also known that operations such as convolution with Gaussian (or other density that satisfy LSI) preserves LSI as well [4], which offers the option of smoothing to perform approximate sampling from a nicer proxy potential f̃ . All these reasons make target measure satisfying Mirror-LSI appealing to study compared to the two previous works [1, 24] studying mirror Langevin for r2f 0.

heading: 4.2 EM-Discretized Process: Interpolation with Weighted Dynamics
text: For analyzing the EM discretization, we build upon the idea initiated in [23] and view the discretized mirror Langevin Monte Carlo (4) as following a weighted Langevin dynamics. It is clear that yk+1 = r (xk+1) is the value at time t = hk+1 of the stochastic process
Yt = Y0 t ·rf(r ⇤(Y0)) + p 2[r2 ⇤(Y0)] 1Wt
starting from Y0 = yk, or written in differential equation form,
dYt = rf(r ⇤(Y0))dt+ p 2[r2 ⇤(Y0)] 1dWt . (9)
Through the mapping Xt = r ⇤(Yt), one can study the corresponding dynamics in X-space, which evolves following a weighted Langevin dynamics with shifted drift µ̂ as shown in Lemma 1 (that is responsible for convergence to a biased limit ⇡h 6= ⇡), i.e.,
dXt = (r ·G(Xt) G(Xt)rf(Xt) + µ̂)dt+ p 2G(Xt)dWt . (10)
Lemma 1 (Shifted Drift and Covariance). For the dynamics written in (10) following (9), we have G(Xt) = [r2 (Xt)] 1r2 (X0)[r2 (Xt)] 1 0 and
µ̂ = [r2 (Xt)] 1 rf(X0) [r 2 (Xt)] 1 Tr r 3 (Xt)[r 2 (Xt)] 1 r 2 (X0)[r 2 (Xt)] 1
r · [r2 (Xt)] 1 r 2 (X0)[r 2 (Xt)] 1 + [r2 (Xt)] 1r2 (X0)[r2 (Xt)] 1rf(Xt) .
Having established this, a combination of Mirror LSI and careful bounding of the discretization error using self-concordance and smoothness properties can be brought together to derive the per-iteration progress. Deatiled proofs for this section can be found in Appendix C. Proposition 2 (Progress in One Step of EM Discretization). In one iteration of Algorithm (4) with xk ⇠ ⇢0, xk+1 ⇠ ⇢h, under Assumption 1- 5, define D := maxu,v kr (u) r (v)k2, we have for stepsize h  min(1/2⇣L, 1/16⇣2d,D/ p ↵L,D 2 /4↵d,M/6 ),
H⇡(⇢h)  e 3 2M hH⇡(⇢0) + 24M 2 2 dh 2 + 16M⇣2d2⌘2 h h ,
where we denote M = exp(2⇣D/ p ↵) and ⌘2
h = (1 exp( 1/16⇣2h)) · (1 ⇣(hL+2
p hd)) 4 +
exp( 1/16⇣2h) ·M2. We use the convention M = 1 when ⇣ = 0 and D = 1.
Picking appropriate stepsize gives the following result – perhaps surprisingly, this particular analysis suggests that the simplest EM discretization exhibits an irreducible bias, and similar observation was also made in [24] where the authors showed convergence to a Wasserstein ball with explicit radius, even for diminishing stepsize. Theorem 1 (Convergence Guarantee for EM). Under Assumption 1-5, picking stepsize h  min(1/2⇣L, 1/16⇣2d,D/ p ↵L,D 2 /4↵d,M/6 , /44M3 2d), after k ⌦̃(M4 2d/ 2 ) iterations of Algorithm (4), we have for xk ⇠ ⇢k that H⇡(⇢k)  +Rh, where the nonvanishing bias Rh = O(M2⇣2d2/ ). In the above, D = maxu,v kr (u) r (v)k2 and M = exp(2⇣D/ p ↵).
Proof. Iterating the inequality in Proposition 2 for k iterations,
H⇡(⇢k)  e 3 2M hkH⇡(⇢0) +
24M2 2dh2
1 e 3 2M h
+ 16M⇣2d2⌘2 h h
1 e 3 2M h
 e 3 2M hkH⇡(⇢0) +
22M3 2dh
+
15M2⇣2d2⌘2 h
where we used 1 e a 3a/4 for a 2 (0, 1/4]. Now using Lemma 6 for initialization, picking the assumed stepsize, after k ⌦̃(M/ h) ⌦̃(M4 2d/ 2 ) iterations, we have H⇡(⇢k)  + Rh. As long as r2 is not constant (therefore ⇣ 6= 0, recall M 1), Rh := 15M2⇣2d2⌘2h 1 6= 0 as h ! 0 and the asymptotic bias Rh scale as O(M2⇣2d2/ ) since ⌘2h ! 1 as h ! 0.
Remark. This convergence rate for KL divergence is stronger than [24] with their guarantee in Wasserstein distance. Additionally, by Pinsker’s inequality, TV distance is upper bounded by KL divergence therefore one could also get an analogous guarantee in that metric. Variations on the argument will likely generalize to other metrics such as 2 and Rényi divergence, with other appropriate (mirror-version of) functional inequalities such as Poincaré [6, 20].
In the case when (x) = kAxk2/2, Theorem 1 gives no asymptotic bias, as we have M = 1 and ⇣ = 0. Moreover, the assumption in this case says f is smooth in k · kA>A although not necessarily convex, and our algorithm gives an update of the form xk+1 = xk hk+1(A>A) 1rf(xk) +p 2hk+1(A>A) 1zk+1, which coincides with the classical Langevin performed on the function g(x̃) = f(A 1x̃) for x = A 1x̃. One could check that g is smooth · I r2g · I and satisfies the classical LSI, which means that f is globally nice w.r.t a fixed geometry, after a change of basis. If applying the Euclidean result of [20] on g(·), Theorem 1 recover the same Õ(d/ ) complexity, which is the best-known rate without third-order smoothness assumption on f .
A closer look at the theorem points out that M plays a prominent role in the rate, where M is effectively the Hessian stability parameter (cf. Lemma 4). However, this dependence on M necessarily means that needs to be smooth on its domain (and M will be roughly the condition number of ). Although this extends beyond the Euclidean case when needs to be constant – hence allowing for slowly-changing geometry, it is still far from satisfying. An important motivation for relaxing smoothness assumption is in constrained sampling (e.g., uniform sampling from a convex body), where one typically would pick to be a self-concordant barrier function that blows up on the boundary as a proxy for the nonsmooth constraint for approximate sampling. Such functions do not have a bounded M . This, in additional to the non-vanishing bias, all suggest that EM discretization, natural as it may seem, although could improve particular parameter dependence compared to the Euclidean counterpart (as illustrated above), might not fully benefit from the use of a mirror map.

heading: 4.3 Alternative Forward Discretization Scheme
text: In practical applications, it is often the case that the cost of evaluating r is considerably cheaper than the cost of computing rf , which could involve a finite sum over a large number of data points. Taking hints from this observation, in [1], a slightly modified mirror Langevin algorithm was considered where at iteration k, with step size ⌘
xk+1/2 = argmin v
⌘rf(xk) > v +D (v, xk) = r ⇤(r (xk) ⌘rf(xk)) (11)
solve dyt = p 2[r2 ⇤(yt)] 1dWt for y0 = r (xk+1/2) (12)
xk+1 = r ⇤(y⌘) (13)
The oracle complexity (i.e., number of queries for rf ) of the above algorithm is the same as the one in (4), but aiming at a higher accuracy implementation for the diffusion part involving . The inner step (12) can be implemented approximately using e.g., Euler-Maruyama. Integrating both sides of (12), it is not hard to see that yk+1 = r (xk+1) is the value at time t = ⌘ of the continuous process
Yt = Y0 trf(X0) + p 2
Z t
0 [r2 ⇤(Ys)] 1/2 dWs (14)
given X0 = r ⇤(Y0) = xk from the previous iteration. Written in differential form,
dYt = rf(r ⇤(Y0))dt+ p 2[r2 ⇤(Yt)] 1dWt . (15)
Compared to (9), the difference is in the second term where we traded Y0 for Yt, therefore this formulation amounts to discretizing the objective but not the geometry and will turn out to be crucial for removing the asymptotic bias. This is also in line with the observation from optimization [9], from which the authors argue that ẋt = r2 (xt) 1rf(xbtc) gives a more “faithful" discretization compared to ẋt = r2 (xbtc) 1 rf(xbtc). Indeed for the process (15), one gets in the X-space another weighted Langevin dynamics (10) with G = [r2 (Xt)] 1 and µ̂ = [r2 (Xt)] 1(rf(Xt) rf(X0)), the discretization error of the gradient in the local metric.
It turns out that this algorithm based on splitting the deterministic and stochastic part of the SDE works with a weaker notion of smoothness assumption as well. In particular, this definition of relative smoothness only involves the local metric r2 at a single point, whereas the previous Assumption 4 requires Lipschitz gradient across different metrics r2 , which might be unavoidable if one is discretizing the geometry as well. Assumption 6 (Weaker -Relative Smooth). For all x, x0 2 dom( ), it holds that
krf(x) rf(x0)k[r2 (x0)] 1  · kr (x) r (x 0)k[r2 (x0)] 1 .
When (x) = kAxk22/2, Assumption 6 reduces to krf(x) rf(x0)k(A>A) 1  · kx x0kA>A; and when = f we always have = 1 (not the case for Assumption 4). We have the following result for the forward-discretized Mirror Langevin algorithm. Proofs for this section can be found in Appendix D. Proposition 3 (Convergence Guarantee for Forward Discretization). For the Algorithm in (11)- (13), under Assumption 1-3,5,6, let M = exp(2⇣D/ p ↵) and D = maxu,v kr (u) r (v)k2, picking stepsize h  min(1/2⇣L, 1/16⇣2d,D/ p ↵L,D 2 /4↵d, 1/6 , /100M 2d), after k
⌦̃(M 2d/ 2 ) iterations, we have H⇡(⇢k)  .
While it is reassuring that the algorithm has vanishing bias with diminishing stepsize for any self-concordant mirror map (at a higher cost of computation for each step), we still see the appearance of M in the rate, which as we discussed earlier, substantially limit the use case for handling weakly smooth potentials as motivation for the introduction of a mirror map. To give a concrete example, for logistic regression, one might be interested in sampling from a posterior ⇡(✓) / (✓) · exp[ P i yi✓ > xi log(1 + exp(✓>xi))], where (✓) could be a constrained prior such as uniform on `1 ball [ 1, 1]d. In such a setting, one natural choice is to pick a mirror map which is a self-concordant barrier for the constraint set to enforce the constraint on the drawn samples, e.g., (✓) = P i log((1 ✓i) 1) + log((1 + ✓i) 1) with dom( ) = ( 1, 1)d. One could check that with
the potential of interest as f(✓) = P
i yi✓
> xi + log(1 + exp(✓>xi)), it does not have a bounded
M .

heading: 4.4 Alternative Backward Discretization Scheme
text: Backward discretization is known to be more stable compared to forward discretization in optimization and it is also known to give the best rate under LSI for Langevin diffusion in the Euclidean setting with weaker assumptions [23], albeit at a higher cost of solving a proximal step xk+1 = xk ⌘rf(xk+1) = argminx f(x)+ (2⌘) 1kx xkk22 compared to xk+1 = xk ⌘rf(xk) at each step.
It is relatively straightforward to see that a backward discretization for the dynamics using the same philosophy as Section 4.3 can be implemented with step size ⌘ as (assuming ⌘f + is convex - guaranteed if ⌘ small enough)
solve dyt = p 2[r2 ⇤(yt)] 1dWt for y0 = r (xk) (16)
xk+1 = argmin v
⌘f(v) + (v) y> ⌘ v , ⌘rf(xk+1) +r (xk+1) y⌘ = 0 (17)
and yk+1 = r (xk+1) is the value at time t = ⌘ of the continuous process
Yt = Y0 trf(r ⇤(Yt)) + p 2
Z t
0 [r2 ⇤(Ys)] 1/2 dWs (18)
given xk = r ⇤(Y0) from the previous iteration. Here step (16) can again be solved iteratively using Euler-Maruyama and (17) is another convex optimization for which we can implement approximately.
We include the argument of this scheme in Appendix E, from which we see that one gets a better rate compared to the previous forward-discretized method (⌦( 1/2) vs. ⌦( 1)) while maintaining no bias. The crucial step, which was also used in the work of [23], is to relate the process (18) to an SDE for which the µ̂ and G only involve samples at Xt and no other time points (e.g., X0), contrary to what happens for EM and forward discretization. This allows us to get a tighter control on the relevant quantities for bounding the discretization error and hence a better final rate (cf. Lemma 9). The difficulty for the other two schemes lies in the fact that due to the stochastic Brownian motion term, if we were to bound the discretization error between two time points X0 and Xt, one would need a more global stability control (hence dependence on M ) to account for the small probability that they are far apart (therefore local stability implied by self-concordance doesn’t help). Our analysis for the algorithm works with the smoothness assumption stated below. Assumption 7 (Weaker -Relative Smooth). For all x, x0 2 dom( ), it holds that
r 2 (x) r2f(x) r2 (x) ,
max kr 2 f(x)[r2 (x)] 1 r2f(x0)[r2 (x0)] 1kop,
k[r2 (x)] 1r2f(x) [r2 (x0)] 1r2f(x0)kop  Kkx x 0 k .
Examples abound for such an assumption. We give an example here where f is not smooth, yet satisfy this relative smoothness condition above. For f(x) = x log(x) which does not have Lipschitz gradient, picking the strongly convex (x) = x log(x) + (1 x) log(1 x) where dom( ) = [0, 1], it’s easy to see f 00/ 00 = x 1/(x 1 + (1 x) 1) satisfy both requirements.
Below we give the result for our algorithm. As alluded to earlier, this new proposal removes the M dependence altogether. It is somewhat expected that relative-smoothness type assumption would show up in the result, the lack of which necessarily implies that the potential f is “misaligned" w.r.t the underlying changing geometry, for which sampling is unequivocally expected to be hard. But what’s surprising is that this is in fact all that’s required from for our method, and in itself doesn’t have to be smooth or self-concrodant in this case. Proposition 4 (Convergence Guarantee for Backward Discretization). For the Algorithm (16)-(17), under Assumption 2,3,5,7 and stepsize h = O(min{1/ , 1/K, 1/ , p /( 2L2 + ↵ 1d3K2)}), after k ⌦̃( p 2L2 + ↵ 1d3K2/ 1/2 3/2) iterations, we have H⇡(⇢k)  .
Let us mention in passing some possible extensions to the framework presented above. In the case when we are dealing with potential with finite sum structure, i.e., f(x) = P i fi(x), as is often the case in machine learning problems, one could execute instead of (17) the update
xk+1 = argmin v ⌘
X i2B fi(v) + (v) y > ⌘ v, (19)
where B is a random batch of data points. Such algorithm effectively assumes that we have stochastic (and therefore) noisy access to rf , where r̂f(x) = rf(x) + ⇣ for ⇣ an independent random noise vector with E[⇣] = 0 and E[k⇣k22]  d 2. Basic considerations suggest that this stochastic variant of the algorithm will converge to a noise ball with radius that scales with the variance of the noise , but is nevertheless more efficient from a computational perspective.

heading: 5 Numerical Experiments
text: In this section, we test out the induced bias and the benefit of using mirror maps in two experiments.
For the first experiment, we take the similar setup as in [6], where we pick = f (i.e., Newton) and consider uniform sampling from a 2D box [ 0.01, 0.01]⇥ [ 1, 1]. Four methods are compared. For Newton Langevin, we aim to target ⇡ / exp( · ), taking (x) = log(1 x21) log(0.012 x22) as the barrier. We test out the 3 different discretization schemes with = 10 4 so that ⇡ ⇡ ⇡. Stepsize is chosen to be h = 10 5. Projected Langevin is taken to be another option for dealing with constraints, which targets the uniform distribution ⇡ directly and simply performs ULA followed by projection onto the domain. The plot below shows the samples after 500 iterations, where r ⇤ and the proximal operator are solved with 50 steps of gradient descent steps. Diffusion term is solved with 10 inner steps of EM. From the samples, EM seems to give qualitatively different result, suggesting the possible existence of bias.
The second experiment concerns ill-conditioned Gaussian potential (no bias in this case), for which we compare the speed of convergence for Mirror-Langevin (with EM discretization) vs. ULA. We take = f = (x µ)>⌃ 1(x µ)/2, and repeat the process 200 times with d = 50, computing the error in empirical mean kµ̂ µk2 and covariance k⌃̂ ⌃kF /k⌃kF and plot them across iterations below. Stepsize h is picked to be 10 3 in both cases and initialization as N (0, I).

heading: 6 Discussion
text: Our result characterizes the interplay between and f for different discretization schemes, which can be used to guide particular choice of mirror map given the sampling problem on hand. Our newly proposed algorithm and the analysis of several previously proposed schemes in the setting of
sampling from nonsmooth and nonconvex potentials highlight several interesting distinctions that do not find parallel in the traditional Euclidean setup.
In optimization, one typically requires quite strong assumption for global stable convergence of Newton’s method (for which = f ). Are various tricks such as Trust-Region, Cubic-Regularized Newton justified here as well for either speeding up sampling or correcting for bias? As future steps, it is an unfulfilled dream of ours to formalize and confirm a lower bound. It is conceivable that the extra bias term is unavoidable and captures the price we have to pay for discretizing while asking for weaker smoothness. It’s also interesting to ask whether higher-order discretizers such as Runge-Kutta would yield better rate. In a different vein, one could explore the possibility of higher-order dynamics (à la underdampled Langevin) for which more sophisticated integrator [19, 16] could potentially be leveraged. Conventional wisdom suggests that the introduction of auxiliary variable that handles the non-smoothness of Brownian motion can often lead to design of better discrete sampling algorithms. On the probability theory front, it remains an intriguing open question to give a complete characterization for the Mirror-LSI condition.
Acknowledgments and Disclosure of Funding
We are grateful for the constructive comments and feedback from the anonymous reviewers. This work was partially supported under NSF-2032014.

Review summaries:
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.

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.    

The paper under review mainly studies different discretization schemes for mirror Langevin dynamics. The main assumption is the non-asymptotic convergence/long time convergence behavior to the target invariant measure. The main feature of this type Langevin dynamic is that they have non-gradient drift and variable diffusion matrix. The author try to covey the importance of the underlying geometry which is introduced by the variable diffusion matrix. 

This paper gives a discretization analysis of the mirror-Langevin Monte Carlo algorithm, which is the analogue of mirror descent for MCMC. The prior works studying discretization of mirror-Langevin are as follows:  [1] Zhang et al., who studied the Euler-Maruyama discretization. They used stochastic calculus coupling techniques in the vein of works by e.g. Dalalyan.  [2] Ahn and Chewi, who introduced and studied an alternate discretization which simulates the diffusion step more faithfully. Their analysis uses Wasserstein calculus in the vein of Durmus, Majewski, and Miasojedow.  A key difference between these two works is that the result of [1] has a non-vanishing bias (even with vanishing step size), whereas the discretization of [2] does not.  This paper studies both the Euler-Maruyama and Ahn-Chewi discretizations, and introduces a third one which combines the Ahn-Chewi discretization with a proximal step to obtain better guarantees. The results obtained have non-vanishing bias for the Euler-Maruyama discretization, and vanishing bias for the latter two discretizations, which is in line with previous work. The techniques are inspired by the discretization analyses of Wibisono, who is inspired by the interpretation of sampling as composite optimization in Wasserstein space. Moreover, this paper emphasizes the use of a mirror-log-Sobolev inequality as the condition ensuring rates of convergence, which allows for treating examples of non-convex potentials (which are not handled by the prior two works).  The main contributions can be summarized thus: (1) by obtaining a similar dichotomy between the bias of Euler-Maruyama vs. Ahn-Chewi, it sheds further light on the effect of discretization of the diffusion in the presence of non-isotropic diffusions; (2) it shows how to adapt a third type of discretization analysis to the mirror-Langevin setting; (3) it introduces a proximal mirror-Langevin algorithm which had previously not been considered and gives guarantees.