Mirrored Langevin Dynamics
Ya-Ping Hsieh Ali Kavis Paul Rolland Volkan Cevher



numbers, compressnatbib















 
    
       
            
       
       
       
      





toltxlabel
*appendix
equationsection





















































































ł((
)̊)
((
))

    assumptionAssumption𝐱𝐲𝐛̱𝐳𝐏𝐗𝐘𝐧𝐍𝐬𝐕𝐦Ł Lℝℂ𝐀𝐔𝐏𝐯̌𝐰𝐓𝐕Ṽ𝐈αβψΨϕΦ1𝐚𝐱^♮𝒫ℐf_minf_avgf_gen𝒜diag c_min c_maxΩ_bestGPC𝐱ℙ𝔼ℤ𝔼̂𝐱̃w_t1w_t+11a_t+1𝒮𝒳𝒵ℳℋ𝒟𝒯𝒞_𝒞^∞_𝐌𝐜̧ℬ𝐁ϵ_minϵ_avg[m]  n∇^2 h^-1/2∇^2 h^1/2  h^⋆ł(∇^2 h^⋆)̊^1/2Ł_hℒ_h ∇ h 2h∇^2 h ∇ h^⋆2hstar∇^2 h^⋆ t𝐁_t μν e^-V() e^-W() d_ℒ_h theoremTheoremlemmaLemmaremarkRemarkexampleExample h^⋆∇𝕀Δ̃þh̃h̃^⋆𝕊^+_d ℳ𝒩𝒯d∑_i=1^d ∑_i=1^d+1Mirrored Langevin Dynamics
        
      Ya-Ping Hsieh  
        
        
      
      
      Ali Kavis 
    
      Paul Rolland 
    
      Volkan Cevher 
      Laboratory for Information and Inference Systems (LIONS),
    EPFL, Lausanne, Switzerland
    
        Received: date / Accepted: date
    ===============================================================================================================================================================================
    
    
    We consider the problem of sampling from constrained distributions, which has posed significant challenges to both non-asymptotic analysis and algorithmic design. We propose a unified framework, which is inspired by the classical mirror descent, to derive novel first-order sampling schemes. We prove that, for a general target distribution with strongly convex potential, our framework implies the existence of a first-order algorithm achieving Õ(ϵ^-2d) convergence, suggesting that the state-of-the-art Õ(ϵ^-6d^5) can be vastly improved. With the important Latent Dirichlet Allocation (LDA) application in mind, we specialize our algorithm to sample from Dirichlet posteriors, and derive the first non-asymptotic Õ(ϵ^-2d^2) rate for first-order sampling. We further extend our framework to the mini-batch setting and prove convergence rates when only stochastic gradients are available. Finally, we report promising experimental results for LDA on real datasets.
    
    
    
    
    
     0
    We generalize the Langevin Monte Carlo algorithm through the mirror descent framework for first-order sampling. Our motivation is to tackle constrained sampling problems, which are ubiquitous in machine learning, but have also posed significant challenges to both theoretical analysis and practical algorithm design. In this paper, we demonstrate how the classical mirror map idea, when combined with theory of Optimal Transport, leads to simple algorithms with fast convergence rates. Furthermore, for certain important problems that were previously considered difficult, our framework allows to reduce them to the simple task of sampling from unconstrained log-concave distributions. Concretely, we present the first non-asymptotic Õ(ϵ^-2d^2 R_0) rate (where ϵ is the precision, d is the dimension, and R_0 is the initial distance) in relative entropy for Dirichlet posteriors, a distribution of central interest in topic modelling. For general constrained sampling problems, we prove the existence of a mirror map that sharpens the state-of-the-art Õ(ϵ^-6d^5) rate to Õ(ϵ^-2d), when the target distribution is strongly log-concave with compact support. We provide preliminary experimental evidence that our algorithm outperforms existing first-order methods for sampling from Dirichlet posteriors.
     0
    We generalize the Langevin Dynamics through the mirror descent framework for first-order sampling. The naïve approach of incorporating Brownian motion into the mirror descent dynamics, which we refer to as Symmetric Mirrored Langevin Dynamics (S-MLD), is shown to connected to the theory of Weighted Hessian Manifolds. The S-MLD, unfortunately, contains the hard instance of Cox–Ingersoll–Ross processes, whose discrete-time approximation exhibits slow convergence both theoretically and empirically. We then propose a new dynamics, which we refer to as the Asymmetric Mirrored Langevin Dynamics (A-MLD), that avoids the hurdles of S-MLD. In particular, we prove that discretized A-MLD implies the existence of a first-order sampling algorithm that sharpens the state-of-the-art Õ(ϵ^-6d^5) rate to Õ(ϵ^-2d), when the target distribution is strongly log-concave with compact support. For sampling on a simplex, A-MLD can transform certain non-log-concave sampling problems into log-concave ones. As a concrete example, we derive the first non-asymptotic Õ(ϵ^-2d^2 R_0) rate for first-order sampling of Dirichlet posteriors, where R_0 is the 2-Wasserstein distance between the initial distribution and the target. 
    
    
    § INTRODUCTION
    
     
    
    
    
    
    
    
    
    Many modern learning tasks involve sampling from a high-dimensional and large-scale distribution, which calls for algorithms that are scalable with respect to both the dimension and the data size. One approach <cit.> that has found wide success is to discretize the Langevin Dynamics:
    _t = -∇ V(_t) + √(2),
     
    
    
    
    
    
    where  presents a target distribution and _t is a d-dimensional Brownian motion. Such a framework has inspired numerous first-order sampling algorithms <cit.>, and the convergence rates are by now well-understood for unconstrained and log-concave distributions <cit.>.  
    
    However, applying (<ref>) to sampling from constrained distributions (i.e., when V has a bounded convex domain) remains a difficult challenge. From the theoretical perspective, there are only two existing algorithms <cit.> that possess non-asymptotic guarantees, and their rates are significantly worse than the unconstrained scenario under the same assumtions; cf., Table <ref>. Furthermore, many important constrained distributions are inherently non-log-concave. A prominent instance is the Dirichlet posterior, which, in spite of the presence of several tailor-made first-order algorithms  <cit.>, is still lacking a non-asymptotic guarantee.
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    In this paper, we aim to bridge these two gaps at the same time. For general constrained distributions with a strongly convex potential V, we prove the existence of a first-order algorithm that achieves the same convergence rates as if there is no constraint at all, suggesting the state-of-the-art Õ(ϵ^-6d^5) can be brought down to Õ(ϵ^-2d).
    
     When specialized to the important case of simplex constraint, we provide the first non-asymptotic guarantee for Dirichlet posteriors, Õ(ϵ^-2d^2R_0) for deterministic and Õł(ϵ^-2(Nd + σ^2)R_0)̊ for the stochastic version of our algorithms; cf., Example <ref> and <ref> for the involved parameters.
    
    Our framework combines ideas from the Mirror Descent<cit.> algorithm for optimization and the theory of Optimal Transport <cit.>. Concretely, for constrained sampling problems, we propose to use the mirror map to transform the target into an unconstrained distribution, whereby many existing methods apply. Optimal Transport theory then comes in handy to relate the convergence rates between the original and transformed problems. For simplex constraints, we use the entropic mirror map to design practical first-order algorithms that possess rigorous guarantees, and are amenable to mini-batch extensions.
    
    The rest of the paper is organized as follows. We briefly review the notion of push-forward measures in Section <ref>. In Section <ref>, we propose the Mirrored Langevin Dynamics and prove its convergence rates for constrained sampling problems. Mini-batch extensions are derived in Section <ref>. Finally, in Section <ref>, we provide synthetic and real-world experiments to demonstrate the empirical efficiency of our algorithms.
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
     §.§ Related Work
    First-Order Sampling Schemes with Langevin Dynamics: There exists a bulk of literature on (stochastic) first-order sampling schemes derived from Langevin Dynamics or its variants <cit.>. However, to our knowledge, this work is the first to consider mirror descent extensions of the Langevin Dynamics.
    
    The authors in <cit.> proposed a formalism that can, in principle, incorporate any variant of Langevin Dynamics for a given distribution . The Mirrored Langevin Dynamics, however, is targeting the push-forward measure  (see Section <ref>), and hence our framework is not covered in <cit.>.
    
    For Dirichlet posteriors, there is a similar variable transformation as our entropic mirror map in <cit.> (see the “reduced-natural parametrization” therein). The dynamics in <cit.> is nonetheless drastically different from ours, as there is a position-dependent matrix multiplying the Brownian motion, whereas our dynamics has no such feature; see (<ref>).
    
    
    
    
    
    Mirror Descent-Type Dynamics for Stochastic Optimization: Although there are some existing work on mirror descent-type dynamics for stochastic optimization<cit.>, we are unaware of any prior result on sampling. 
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    § PRELIMINARIES
    
    
     §.§ Notation
    
    In this paper, all Lipschitzness and strong convexity are with respect to the Euclidean norm ·. We use 𝒞^k to denote k-times differentiable functions with continuous kth derivative. The Fenchel dual <cit.> of a function h is denoted by . Given two mappings T, F of proper dimensions, we denote their composite map by T∘ F. For a probability measure μ, we write ∼μ to mean that “ is a random variable whose probability law is μ”. 
    
    
     
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
     §.§ Push-Forward and Optimal Transport
    
    Let = be a probability measure with support (V) = {∈^d | V() < +∞}, and h be a convex function on . Throughout the paper we assume:
    h is closed, proper, h ∈𝒞^2 , and ∇^2 h ≻ 0 on ⊂^d.
    
    All measures have finite second moments.
    
    All measures vanish on sets with Hausdorff dimension <cit.> at most d-1.
    
    
    The gradient map ∇ h induces a new probability measure  through ν(E) = μł(∇ h^-1(E) )̊ for every Borel set E on ^d. We say that ν is the push-forward measure of μ under ∇ h, and we denote it by ∇ h #μ = ν. If ∼μ and ∼ν, we will sometimes abuse the notation by writing # = to mean #μ = ν. 
    
    
    
    
    
    If #μ = ν, the triplet (μ, ν, h) must satisfy the Monge-Ampère equation:
    
    e^-V = e^-W ∘∇ h∇^2 h.
    
    Using (∇ h)^-1  = ∇ and ∇^2 h ∘∇ = ∇^2 ^-1, we see that (<ref>) is equivalent to
    
    e^-W = e^-V ∘∇∇^2 
    which implies ∇#ν = μ.
    
    
    
    
    
    
    The 2-Wasserstein distance between μ_1 and μ_2 is defined by[In general, (<ref>) is ill-defined; see <cit.>. The validity of (<ref>) is guaranteed by McCann's theorem <cit.> under Assumption <ref> and <ref>. ]^2_2(μ_1, μ_2) inf_T: T#μ_1 = μ_2∫ - T() ^2  _1().
    
    
    § MIRRORED LANGEVIN DYNAMICS
    
    This section demonstrates a framework for transforming constrained sampling problems into unconstrained ones. We then focus on applications to sampling from strongly log-concave distributions and simplex-constrained distributions, even though the framework is more general and future-proof.
    
    
    
     §.§ Motivation and Algorithm
    
    We begin by briefly recalling the mirror descent (MD) algorithm for optimization. In order to minimize a function over a bounded domain, say min_∈ f(), MD uses a mirror map h to transform the primal variable  into the dual space (), and then performs gradient updates in the dual: ^+ =  - β∇ f() for some step-size β. The mirror map h is chosen to adapt to the geometry of the constraint , which can often lead to faster convergence <cit.> or, more pivotal to this work, an unconstrained optimization problem <cit.>.
    
    Inspired by the MD framework, we would like to use the mirror map idea to remove the constraint for sampling problems.
    
    Toward this end, we first establish a simple fact <cit.>:
    
    
    Let h satisfy Assumption <ref>. Suppose that ∼μ and =  (). Then ∼ν∇ h #μ and () ∼μ.
    
    For any Borel set E, we have ν(E) = ł( ∈ E )̊  = ł( ∈∇ h^-1(E) )̊ =  μł(∇ h^-1 (E) )̊. Since  is one-to-one, =  () if and only if = ^-1()= ().
    
    
    In the context of sampling, Theorem <ref> suggests the following simple procedure: For any target distribution  with support , we choose a mirror map h on  satisfying Assumption <ref>, and we consider the dual distribution associated with  and h:
    #.
    Theorem <ref> dictates that if we are able to draw a sample  from , then () immediately gives a sample for the desired distribution . Furthermore, suppose for the moment that () = ^d, so that  is unconstrained. Then we can simply exploit the classical Langevin Dynamics (<ref>) to efficiently take samples from . 
    
    The above reasoning leads us to set up the Mirrored Langevin Dynamics (MLD):
    MLD≡{[ _t = -(∇ W∘∇ h)(_t)  + √(2);                 _t = ∇ (_t) ].
            .
    
    Notice that the stationary distribution of _t in MLD is , since _t is nothing but the Langevin Dynamics (<ref>) with ∇ V ←∇ W. As a result, we have _t →_∞∼.
    
    
    
    
    
    
    
    
    
    
    
    
    
    Using (<ref>), we can equivalently write the _t term in (<ref>) as _t = -  ∇^2 h(_t)^-1(∇ V(_t) + ∇log∇^2h(_t)  ) + √(2).
    
    In order to arrive at a practical algorithm, we then discretize the MLD, giving rise to the following equivalent iterations:
    
    ^t+1 - ^t = {[                               - β^t ∇ W  (^t)+ √(2β^t)ξ^t; -  β^t∇^2 h(^t)^-1(∇ V(^t) + ∇log∇^2h(^t)  ) + √(2β^t)ξ^t ]. 
    
    where in both cases ^t+1 =  (^t+1), ξ^t's are i.i.d. standard Gaussian, and β^t's are step-sizes. 
    The first formulation in (<ref>) is useful when ∇ W has a tractable form, while the second one can be computed using solely the information of V and h. 
    
    
    
    
    
    
    
    
    
    
    Next, we turn to the convergence of discretized MLD. Since _t in (<ref>) is the classical Langevin Dynamics, and since we have assumed that W is unconstrained, it is typically not difficult to prove the convergence of ^t to _∞∼. However, what we ultimately care about is the guarantee on the primal distribution . The purpose of the next theorem is to fill the gap between primal and dual convergence.
    
    We consider three most common metrics in evaluating approximate sampling schemes, namely the 2-Wasserstein distance _2, the total variation , and the relative entropy D(··). 
    
    
    For any h satisfying Assumption <ref>, we have d_(∇ h#μ_1, ∇ h#μ_2) = d_(μ_1, μ_2) and D(∇ h#μ_1 ∇ h#μ_2) = D(μ_1 μ_2). In particular, we have d_(^t, _∞) =d_(^t, _∞) and D(^t _∞) = D(^t _∞) in (<ref>).
    
    If, furthermore, h is ρ-strongly convex: ∇^2 h ≽ρ I. Then _2(^t, _∞)≤1/ρ_2(^t, _∞).
    
    
    See Appendix <ref>.
    
    
     §.§ Applications to Sampling from Constrained Distributions
    
    We now consider applications of MLD. For strongly log-concave distributions with general constraint, we prove matching rates to that of unconstrained ones; see Section <ref>. In Section <ref>, we consider the important case where the constraint is a probability simplex[More examples of mirror map can be found in Appendix <ref>.].
    
    
    
      §.§.§ Sampling from a strongly log-concave distribution with constraint
    
    As alluded to in the introduction, the existing convergence rates for constrained distributions are significantly worse than their unconstrained counterparts; see Table <ref> for a comparison. 
    
    
    
    
    
    
    The main result of this subsection is the existence of a “good” mirror map for arbitrary constraint, with which the dual distribution  becomes unconstrained:
    
    
    
    Let μ() = e^-V() be a probability measure with bounded convex support such that V ∈𝒞^2, ∇^2 V ≽ m I ≻ 0, and V is bounded away from +∞ in the interior of the support. Then there exists a mirror map h ∈𝒞^2 such that the discretized MLD (<ref>) yields
    
    Dł( ^T _∞)̊=Õł(   d/T)̊,    _2 ł( ^T, _∞)̊ = Õł(  √(d/T))̊,    d_ł(^T, _∞)̊ = Õł( √(d/T))̊.
    
    
    
    See Appendix <ref>.
    
    We remark that Theorem <ref> is only an existential result, not an actual algorithm. Practical algorithms are considered in the next subsection.
    
    
      §.§.§ Sampling Algorithms on Simplex
    
    We apply the discretized MLD (<ref>) to the task of sampling from distributions on the probability simplex Δ_d{∈^d |  ∑_i=1^dx_i≤ 1, x_i ≥ 0 }, which is instrumental in many fields of machine learning and statistics.
    
    On a simplex, the most natural choice of h is the entropic mirror map <cit.>, which is well-known to be 1-strongly convex:
    
    h() = ∑_ℓ =1^d x_i log x_ℓ + ł(1-∑_ℓ =1^dx_ℓ)̊logł(1-∑_ℓ =1^dx_ℓ)̊,  where   0 log 0  0.
    
    In this case, the associated dual distribution can be computed explicitly.
    
    Let  be the target distribution on Δ_d, h be the entropic mirror map (<ref>), and #. Then the potential W of the push-forward measure admits the expression
    
        W() 
           = V∘ ()- ∑_ℓ=1^dy_ℓ + (d+1)  ()
    where () = logł(1+ ∑_ℓ=1^d e^y_ℓ)̊ is the Fenchel dual of h, which is strictly convex and 1-Lipschitz gradient.  
    
    
    See Appendix <ref>.
    
    Crucially, we have () = ^d, so that the Langevin Dynamics for  is unconstrained.
    
    
    
    Based on Lemma <ref>, we now present the surprising case of the non-log-concave Dirichlet posteriors, a distribution of central importance in topic modeling  <cit.>, for which the dual distribution  becomes strictly log-concave. 
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    [Dirichlet Posteriors]Given parameters α_1, α_2, ..., α_d+1 >0 and observations n_1, n_2, ..., n_d+1 where n_ℓ is the number of appearance of category ℓ, the probability density function of the Dirichlet posterior is 
    
    p() = 1/C∏_ℓ=1^d+1 x_ℓ^n_ℓ+α_ℓ-1,   ∈ł(Δ_d)̊
    where C is a normalizing constant and x_d+1 1- ∑_ℓ=1^d x_ℓ. The corresponding V is
    
    V() = -log p() = log C - ∑_ℓ=1^d+1 (n_ℓ+α_ℓ-1) log x_ℓ,   ∈ł(Δ_d)̊. 
    The interesting regime of the Dirichlet posterior is when it is sparse, meaning the majority of the n_ℓ's are zero and a few n_k's are large, say of order O(d). It is also common to set α_ℓ <1 for all ℓ in practice. Evidently, V is neither convex nor concave in this case, and no existing non-asymptotic rate can be applied. However, plugging V into (<ref>) gives
    
    W() = log C - ∑_ℓ=1^d (n_ℓ + α_ℓ)y_ℓ + ł(∑_ℓ=1^d+1 (n_ℓ+α_ℓ) )̊ ()
    
    which, magically, becomes strictly convex and O(d)-Lipschitz gradient no matter what the observations and parameters are! In view of Theorem <ref> and <cit.>, one can then apply (<ref>) to obtain an Õł(ϵ^-2d^2 R_0)̊ convergence in relative entropy, where R_0 ^2_2(^0,) is the initial Wasserstein distance to the target. 
    
    § STOCHASTIC MIRRORED LANGEVIN DYNAMICS
    
    
    We have thus far only considered deterministic methods based on exact gradients. In practice, however, evaluating gradients typically involves one pass over the full data, which can be time-consuming in large-scale applications. In this section, we turn attention to the mini-batch setting, where one can use a small subset of data to form stochastic gradients. 
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    Toward this end, we assume:
    [Primal Decomposibility]
    The target distribution  admits a decomposable structure V = ∑_i=1^N V_i for some functions V_i.
    
    
    
    Consider the following common scheme in obtaining stochastic gradients. Given a batch-size b, we randomly pick a mini-batch B from { 1, 2, …, N} with |B| = b, and form an unbiased estimate of ∇ V by computing
    ∇̃V N/b∑_i∈ B∇ V_i. 
    The following lemma asserts that exactly the same procedure can be carried out in the dual.
    
    Assume that h is 1-strongly convex. For i = 1,2,...,N, let W_i be such that
    
    e^-NW_i = #e^-NV_i/∫ e^-NV_i.
    
    Define W ∑_i=1^N W_i and ∇̃W N/b∑_i∈ B∇ W_i, where B is chosen as in (<ref>). Then:
    
    
      *  Primal decomposibility implies dual decomposability: There is a constant C such that e^-(W+C) = # e^-V.
    
    
      *  For each i, the gradient ∇ W_i depends only on ∇ V_i and the mirror map h.
    
    
      *  The gradient estimate is unbiased: ∇̃W = ∇ W.
    
    
      *  The dual stochastic gradient is more accurate: ∇̃W - ∇ W ^2 ≤∇̃V - ∇ V ^2.
    
    
    
    See Appendix <ref>.
    
    Lemma <ref> furnishes a template for the mini-batch extension of MLD. The pseudocode is detailed in Algorithm <ref>, whose convergence rate is given by the next theorem.
    
    
    
    
    
    Let  be a distribution satisfying Assumption <ref>, and h a 1-strongly convex mirror map. Let σ^2 ∇̃V - ∇ V ^2 be the variance of the stochastic gradient of V in (<ref>). Suppose that the corresponding dual distribution = # satisfies LI ≽∇^2 W ≽ 0. Then, applying SMLD 
    with constant step-size β^t = β yields[Our guarantee is given on a randomly chosen iterate from {^1, ^2, ..., ^T}, instead of the final iterate ^T. In practice, we observe that the final iterate always gives the best performance, and we will ignore this minor difference in the theorem statement.]:
    
    Dł(^T )̊≤√(2^2_2ł( ^0, )̊ł( Ld+σ^2 )̊/T) = Oł(√(Ld+σ^2/T))̊,
    
    provided that β≤min{[2T ^2_2ł( ^0, )̊ł( Ld+σ^2 )̊]^-1/2, 1/L}.
    
    See Appendix <ref>.
    [SMLD for Dirichlet Posteriors]For the case of Dirichlet posteriors, we have seen in (<ref>) that the corresponding dual distribution satisfies (N+Γ)I ≽∇^2 W ≻ 0, where N ∑_ℓ=1^d+1 n_ℓ and Γ∑_ℓ=1^d+1α_ℓ. Furthermore, it is easy to see that the stochastic gradient ∇̃W can be efficiently computed (see Appendix <ref>):
    ∇̃W ()_ℓN/b∑_i∈ B∇ W_i ()_ℓ = -  ł(Nm_ℓ/b + α_ℓ)̊  +  ł(N+ Γ)̊e^y_ℓ/1+∑_k=1^d  e^y_k, 
    
    where m_ℓ is the number of observations of category ℓ in the mini-batch B. As a result, Theorem <ref> states that SMLD achieves
    
    Dł(^T )̊≤√(2^2_2ł( ^0, )̊( (N+Γ)(d+1)+σ^2 ) /T) = Oł(√((N+Γ)d + σ^2/T))̊
    with a constant step-size. □
    
    
    § EXPERIMENTS
    
    We conduct experiments with a two-fold purpose. First, we use a low-dimensional synthetic data, where we can evaluate the total variation error by comparing histograms, to verify the convergence rates in our theory. Second, We demonstrate that the SMLD, modulo a necessary modification for resolving numerical issues, outperforms state-of-the-art first-order methods on the Latent Dirichlet Allocation (LDA) application with Wikipedia corpus.
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
     §.§ Synthetic Experiment for Dirichlet Posterior
    
    We implement the deterministic MLD for sampling from an 11-dimensional Dirichlet posterior (<ref>) with n_1 = [group-separator=,]10000,n_2 = n_3 = 10, and n_4 = n_5 = ⋯ = n_11 = 0, which aims to capture the sparse nature of real observations in topic modeling. We set α_ℓ = 0.1 for all ℓ. 
    
    As a baseline comparison, we include the Stochastic Gradient Riemannian Langevin Dynamics (SGRLD) <cit.> with the expanded-mean parametrization. SGRLD is a tailor-made first-order scheme for simplex constraints, and it remains one of the state-of-the-art algorithms for LDA.
     For fair comparison, we use deterministic gradients for SGRLD.
    
    
    We perform a grid search over the constant step-size for both algorithms, and we keep the best three for MLD and SGRLD.
     For each iteration, we build an empirical distribution by running [group-separator=,]2000000 independent trials, and we compute its total variation with respect to the histogram generated by the true distribution.
    
    Figure <ref> reports the total variation error along the first dimension, where we can see that MLD outperforms SGRLD by a substantial margin.
     As dictated by our theory, all the MLD curves decay at the O(T^-1/2) rate until they saturate at the dicretization error level. In contrast, SGRLD lacks non-asymptotic guarantees, and there is no clear convergence rate we can infer from Figure <ref>.
    
    The improvement along all other dimensions (i.e., topics with less observations) are even more significant; see Appendix <ref>.
    
    
    
     §.§ Latent Dirichlet Allocation with Wikipedia Corpus
    
    An influential framework for topic modeling is the Latent Dirichlet Allocation (LDA) <cit.>, which, given a text collection, requires to infer the posterior word distributions without knowing the exact topic for each word. The full model description is standard but somewhat convoluted; we refer to the classic <cit.> for details.
    
    Each topic k in LDA determines a word distribution π_k, and suppose there are in total K topics and W+1 words. The variable of interest is therefore π (π_1, π_2, ..., π_K) ∈Δ_W×Δ_W×⋯Δ_W. Since this domain is a Cartesian product of simplices, we propose to use h̃(π) ∑_k=1^K h(π_k), where h is the entropic mirror map (<ref>), for SMLD. It is easy to see that all of our computations for Dirichlet posteriors generalize to this setting. 
    
    
    
    
    
    
    
    
    
    
      §.§.§ Experimental Setup
    
    We implement the SMLD for LDA on the Wikipedia corpus with [group-separator=,]100000 documents, and we compare the performance against the SGRLD <cit.>.
     In order to keep the comparison fair, we adopt exactly the same setting as in <cit.>, including the model parameters, the batch-size, the Gibbs sampler steps, etc. See Section 4 and 5 in <cit.> for omitted details.
    
    Another state-of-the-art first-order algorithm for LDA is the SGRHMC in <cit.>, for which we skip the implementation, due to not knowing how the B̂_t was chosen in <cit.>. Instead, we will repeat the same experimental setting as <cit.> and directly compare our results versus the ones reported in <cit.>. See Appendix <ref> for comparison against SGRHMC. 
    
    
    
    
      §.§.§ A Numerical Trick and the SMLD-approximate Algorithm
    
    A major drawback of the SMLD in practice is that the stochastic gradients (<ref>) involve exponential functions, which are unstable for large-scale problems. For instance, in python, , whereas the relevant variable regime in this experiment extends to 1600. To resolve such numerical issues, we appeal to the linear approximation[One can also use a higher-order Taylor approximation for exp(), or add a small threshold exp() ≃max{ϵ, 1+} to prevent the iterates from going to the boundary. In practice, we observe that these variants do not make a huge impact on the performance.]exp() ≃max{0, 1+}. Admittedly, our theory no longer holds under such numerical tricks, and we shall not claim that our algorithm is provably convergent for LDA. Instead, the contribution of MLD here is to identify the dual dynamics associated with (<ref>), which would have been otherwise difficult to perceive. We name the resulting algorithm “SMLD-approximate” to indicate its heuristic nature.
    
    
    
    
    
      §.§.§ Results
    
    Figure <ref> reports the perplexity on the test data up to [group-separator=,]100000 documents, with the five best step-sizes we found via grid search for SMLD-approximate. For SGRLD, we use the best step-sizes reported in <cit.>. 
    
    From the figure, we can see a clear improvement, both in terms of convergence speed and the saturation level, of the SMLD-approximate over SGRLD. One plausible explanation for such phenomenon is that our MLD, as a simple unconstrained Langevin Dynamics, is less sensitive to discretization. On the other hand, the underlying dynamics for SGRLD is a more sophisticated Riemannian diffusion, which requires finer discretization than MLD to achieve the same level of approximation to the original continuous-time dynamics, and this is true even in the presence of noisy gradients and our numerical heuristics
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
      §.§.§ Acknowledgments
    
    
    This project has received funding from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (grant agreement n^∘ 725594 - time-data).
    
    
    
    
    
    
    plainequationsection
    
    § PROOF OF THEOREM <REF>
    
    We first focus on the convergence for total variation and relative entropy, since they are in fact quite trivial. The proof for the 2-Wasserstein distance requires a bit more work.
    
    
    
     §.§ Total Variation and Relative Entropy
    
    Since h is strictly convex,  is one-to-one, and hence 
    
        d_(∇ h#μ_1, ∇ h#μ_2)
           = 1/2sup_E | ∇ h#μ_1(E) - ∇ h#μ_2(E)|  
           = 1/2sup_E| μ_1( ^-1(E) ) - μ_2 (^-1(E)) | 
           =d_(μ_1, μ_2).
    
    
    On the other hand, it is well-known that applying a one-to-one mapping to distributions leaves the relative entropy intact. Alternatively, we may also simply write (letting ν_i = #μ_i):
    
        D(ν_1 ν_2)    = ∫log_1/_2_1 
           =  ∫logł(_1/_2∘)̊_1       by (<ref>) below
           = ∫log_1/_2_1        by (<ref>)
           = D(μ_1 μ_2)
    
    
    The “in particular” part follows from noticing that ^t ∼#^t and _∞∼#_∞.
    
    
    
    
    
     §.§ 2-Wasserstein Distance
    
    Now, let h be ρ-strongly convex. The most important ingredient of the proof is Lemma <ref> below, which is conceptually clean. Unfortunately, for the sake of rigor, we must deal with certain intricate regularity issues in the Optimal Transport theory. If the reader wishes, she/he can simply assume that the quantities (<ref>) and (<ref>) below are well-defined, which is always satisfied by any practical mirror map, and skip all the technical part about the well-definedness proof.
    
    For the moment, assume h ∈𝒞^5; the general case is given at the end. Every convex h generates a Bregman divergence via B_h (, ')  h() - h(') -  (')-'. 
    
    
    
    
    
    
    
     The following key lemma allows us to relate guarantees in _2 between ^t's and ^t's. It can be seen as a generalization of the classical duality relation (<ref>) in the space of probability measures.
    
    Let μ_1, μ_2 be probability measures satisfying Assumptions <ref> and <ref>. If h is ρ-strongly convex and 𝒞^5, then the (<ref>) and (<ref>) below are well-defined:  
    _B_h(μ_1, μ_2) inf_T: T#μ_1 = μ_2∫ B_h ł(, T() )̊μ_1()
    
    and (notice the exchange of inputs on the right-hand side)
    _B_(ν_1, ν_2) inf_T: T#ν_1 = ν_2∫ B_ł(T(), )̊ν_1().
    
    Furthermore, we have
    _B_h(μ_1, μ_2) = _B_(∇ h #μ_1, ∇ h #μ_2). 
    
    
    
    Before proving the lemma, let us see that the relation in _2 is a simple corollary of Lemma <ref>. Since h is ρ-strongly convex, it is classical that, for any  and ',
    
        ρ/2 - '^2 ≤ B_h(, ' ) = B_(('), ())  ≤1/2ρ() - (')^2.
    
    Using Lemma <ref> and the fact that ^t ∼#^t and _∞∼#_∞, we conclude _2(^t, _∞) ≤1/ρ_2(^t, _∞). It hence remains to prove Lemma <ref> when h ∈𝒞^5.
    
    
    
    
      §.§.§ Proof of Lemma <ref> When h ∈𝒞^5
    
    
    
    
    We first prove that (<ref>) is well-defined by verifying the sufficient conditions in Theorem 3.6 of <cit.>. Specifically, we will verify (C0)-(C2) in p.554 of <cit.> when the transport cost is B_.
    
    
    
    Since h is ρ-strongly convex,  is injective, and hence = ()^-1 is also injective, which implies that  is strictly convex. On the other hand, the strong convexity of h implies ∇^2 ≼1/ρI, and hence B_ is globally upper bounded by a quadratic function. 
    
    We now show that the conditions (C0)-(C2) are satisfied. Since we have assumed h∈𝒞^5, we have B_∈𝒞^4. Since B_ is upper bounded by a quadratic function, the condition (C0) is trivially satisfied. On the other hand, since  is strictly convex, simple calculation reveals that, for any ', the mapping →∇_'B_(, ') is injective, which is (C1). Similarly, for any , the mapping ' →∇_B_(, ') is also injective, which is (C2). By  Theorem 3.6 in <cit.>, (<ref>) is well-defined.
    
    
    We now turn to (<ref>), which will automatically establish the well-definedness of (<ref>). We first need the following equivalent characterization of #μ = ν<cit.>:
    ∫ f  = ∫ f∘
    for all measurable f. Using (<ref>) in the definition of _B_, we get
    
        _B_(∇ h #μ_1, ∇ h #μ_2)    = inf_T∫ B_ł(T(), )̊∇ h#μ_1() 
           = inf_T∫ B_( (T∘∇ h)(), ∇ h() ) μ_1(),
    where the infimum is over all T such that T# (∇ h#μ_1) = ∇ h#μ_2. Using the classical duality B_h(, ') = B_(∇ h('), ∇ h()) and ∇ h ∘∇ ()=, we may further write
    
        _B_(∇ h #μ_1, ∇ h #μ_2)    = inf_T∫ B_h ( , (∘ T ∘ ) () ) μ_1()
    where the infimum is again over all T such that T# (∇ h#μ_1) = ∇ h#μ_2. In view of (<ref>), the proof would be complete if we can show that T# (∇ h#μ_1) = ∇ h#μ_2 if and only if (∘ T ∘ ) #μ_1 = μ_2. 
    
    For any two maps T_1 and T_2, we claim that 
    
    (T_1 ∘ T_2) #μ = T_1#ł(  T_2 #μ)̊. 
    Indeed, for any Borel set E, we have, by definition of the push-forward,
    
        (T_1 ∘ T_2) #μ (E)    = μ( (T_1 ∘ T_2)^-1 (E) ) 
           = μ( (T_2^-1∘ T_1^-1) (E) ).
    On the other hand, recursively applying the definition of push-forward to T_1#ł(  T_2 #μ)̊ gives
    
        T_1#ł(  T_2 #μ)̊(E)    =  T_2 #μ( T^-1(E) )  
           = μ(  (T_2^-1∘ T_1^-1) (E) )
    which establishes (<ref>). 
    
    Assume that T #(#μ_1) = #μ_2. Then we have
    
        (∘ T ∘ ) #μ_1    = # (T #(#μ_1))       
           = # (#μ_2)       
           = (∘)#μ_2       
           = μ_2.
    
    On the other hand, if (∘ T ∘ ) #μ_1 = μ_2, then composing both sides by  and using (<ref>) yields T# (∇ h#μ_1) = ∇ h#μ_2, which finishes the proof.
    
    
    
    
      §.§.§ When h is only 𝒞^2
    
    
    When h is only 𝒞^2, we will directly resort to (<ref>). Let T be any map such that T# (∇ h#μ_1) = ∇ h#μ_2, and consider the optimal transportation problem inf_T∫- T() ^2 #μ_1(). By (<ref>) and (<ref>), we have
    
        inf_T∫- T() ^2 #μ_1()    = inf_T∫()- (T∘)()) ^2 μ_1() 
           ≥ρ^2 inf_T∫- (∘ T∘)()) ^2 μ_1()
    where the infimum is over all T such that T# (∇ h#μ_1) = ∇ h#μ_2. But as proven in Appendix <ref>, this is equivalent to (∘ T∘) #μ_1 = μ_2. The proof is finished by noting ^t ∼#^t and _∞∼#_∞.
    
    
    
    
    
    
    
    
    
    § MORE EXAMPLES OF MIRROR MAP AND THEIR DUAL DISTRIBUTIONS
    
    
    In this section, we present more instances of mirror map other than on the simplex, and their corresponding dual distributions.
    
    
    
     §.§ Mirror map on the hypercube
    
    
    On the hypercube [-1,1]^d, a possible mirror map is 
        h() = 1/2∑_i=1^d ((1+x_i)log(1+x_i) + (1-x_i)log(1-x_i)).
    
    
    It can easily be shown that
    
        ∂ h/∂ x_i = arctanh(x_i),        ∂^2h/∂ x_i ∂ x_j = δ_ij/1 - x_i^2,         ∂ h^*/∂ y_i = tanh(y_i)
    
    which implies that h is 1-strongly convex.
    
    Since the Hessian matrix of h is diagonal, we have:
    
        log∇^2 h() = ∑_i=1^d log(1/1-x_i^2).
    
    Then, using the definition of W, we obtain
    
        W()    = V ∘∇ h^*() + ∑_i=1^d log(1/1-tanh^2(y_i)) 
           = V ∘∇ h^*() + ∑_i=1^d log(4/exp(2y_i) + 2 + exp(-2y_i)) 
           = V ∘∇ h^*() + ∑_i=1^d log(1/2(1 + cosh(2y_i))).
    
    
     §.§ Mirror map on the Euclidean ball
    
    
    On the unit ball {∈ℝ^d : ≤ 1}, where · denote the Euclidean norm, a possible mirror map is
    
        h() = -log(1-)  - .
    
    
    We can compute:
    
        ∂ h/∂ x_i = x_i/1 - ,        ∂^2h/∂ x_i ∂ x_j = δ_ij/1 -  + x_ix_j/(1-)^2,         ∂ h^*/∂ y_i = y_i/1 + .
    
    
    The Hessian matrix can thus be written as ∇^2 h() = 1/1-I + 1/(1 - )^2^T, where I is the identity matrix. Invoking the matrix determinant lemma, we get
    
        (∇^2 h()) = (1 + ^⊤/(1-))) (1/1 - I) = (1/1 - )^d+1.
    
    
    We thus obtain:
    
        W()    = V ∘∇ h^*() - (d+1)log(1 - /1 + ) 
           = V ∘∇ h^*() + (d+1) log(1 + ).
    
    
    § PROOF OF THEREOM <REF>
    
    In previous sections, we are given a target distribution e^-V and a mirror map h, and we derive the induced distribution e^-W through the Monge-Ampère equation (<ref>). The high-level idea of this proof is to reverse the direction: We start with two good distributions e^-V and e^-W, and we invoke deep results in Optimal Transport to deduce the existence of a good mirror map h.
    
    First, notice that if V has bounded domain, then the strong convexity of V implies V > -∞. Along with the assumption that V is bounded away from +∞ in the interior, we see that e^-V is bounded away from 0 and +∞ in the interior of support.
    
    Let ν () = e^-W() be any distribution such that ∇^2 W ≽ I. By Brenier's polarization theorem <cit.> and Assumption <ref>, <ref>, there exists a convex function  whose gradient solves the _2ł(ν, μ)̊ optimal transportation problem. Caffarelli's regularity theorem <cit.> then implies that the Brenier's map h^⋆ is in 𝒞^2. Finally, a slightly stronger form of Caffarelli's contraction theorem <cit.> asserts:
    ∇^2 ≼1/m I,
    
    which implies h= ()^⋆ is m-strongly convex.
     
    Let us consider the discretized MLD (<ref>) corresponding to the mirror map h. Invoking Theorem 3 of <cit.>, the convergence rate of the discretized Langevin dynamics ^T for μ is such that D(^T ν) = Õł(   d/T)̊, which in turn implies _2( ^T, ν) = Õł(  √(d/T))̊ and ( ^T,  ν) = Õł(  √(d/T))̊.Theorem <ref> then completes the proof.
    
    
    
    
    
    
    
    
    
    
    
    
    
    § PROOF OF LEMMA <REF>
    
    Straightforward calculations in convex analysis shows
    
        ∂ h/∂ x_i= logx_i/x_d+1,         ∂^2 h/∂ x_i ∂ x_j =δ_ij x_i^-1 + x_d+1^-1, 
           () = logł( 1+  ∑_i=1^d e^y_i)̊,   ∂/∂ y_i = e^y_i/1+∑_i=1^de^y_i,
    
    which proves that h is 1-strongly convex.
    
    Let μ = be the target distribution and define ν = ∇ h#μ. By (<ref>), we have 
    
        W∘   = V + log∇^2 h.
    
    Since ∇^2 h () = [x^-1_i] + x_d+1^-111^⊤ where 1 is the all 1 vector, the well-known matrix determinant lemma “(A + 𝐮𝐯^⊤) = (1+ 𝐯^⊤ A^-1𝐮) A” gives
    
        log∇^2 h()    = logł( 1+   x_d+1^-1∑_i=1^d x_i )̊·∏_i=1^d x^-1_i 
           = -∑_i=1^d+1log x_i = -∑_i=1^dlog x_i  - logł(1- ∑_i=1^d x_i )̊.
    Composing both sides of (<ref>) with  and using (<ref>), (<ref>), we then finish the proof by computing
    
        W()    = V∘ ()- ∑_i=1^dy_i + (d+1) logł(  1+ ∑_i=1^d e^y_i)̊
           = V∘ ()- ∑_i=1^dy_i + (d+1)  ().
    
    
    § PROOF OF LEMMA <REF>
    
    The proof relies on rather straightforward computations.
    
    
      *  In order to show e^-(W+C) = ∇ h # e^-V for some constant C, we will verify the Monge-Ampère equation:
    
        e^-V = e^-(W∘∇ h +C)det∇ ^2 h
    for V = ∑_i=1^NV_i and W = ∑_i=1^N W_i, where W_i is defined via (<ref>). By (<ref>), it holds that
    1/C_ie^-NV_i = e^-NW_i∘∇ h∇^2 h,    C_i 1/∫ e^-NV_i.
    
    Multiplying (<ref>) for i = 1, 2, ..., N, we get
    ∏_i=1^N 1/C_i e^-NV = e^-NW∘ł(∇^2 h )̊^N.
    
    The first claim follows by taking the Nth root of (<ref>).
    
    
      *  The second claim directly follows by (<ref>).
    
    
      *  Trivial.
    
    
      *  By (<ref>) and (<ref>) and using ∇∘∇ h() =, we get
    
        W_i    =  V_i ∘∇ +  1/Nlog∇^2 h (∇) - log C_i, 
        
        W    = V ∘ + log∇^2 h (∇) - C,
    
    which implies N∇ W_i  - ∇ W=   ∇^2 ł( N∇ V_i ∘∇  - ∇ V ∘∇)̊. Since h is 1-strongly convex,  is 1-Lipschitz gradient, and therefore the spectral norm of ∇^2 is upper bounded by 1. In the case of b=1, the final claim follows by noticing
    
        ∇̃W - ∇ W   ^2    = 1/N∑_i=1^N   N∇ W_i - ∇ W  ^2 
           = 1/N∑_i=1^N  ∇^2 ł( N∇ V_i ∘∇  - ∇ V ∘∇)̊^2 
           ≤∇^2 _^2/N∑_i=1^N  N∇ V_i ∘∇  - ∇ V ∘∇^2 
           ≤∇̃V - ∇ V   ^2.
    The proof for general batch-size b is exactly the same, albeit with more cumbersome notation.
    
    
    § PROOF OF THEOREM <REF>
    
    The proof is a simple combination of the existing result in <cit.> and our theory in Section <ref>.
    
    
    By Theorem <ref>, we only need to prove that the inequality (<ref>) holds for D(^T), where ^T is to be defined below. By assumption, W is unconstrained and satisfies LI ≽∇^2 W ≽ 0. By Lemma <ref>, the stochastic gradient ∇̃W is unbiased and satisfies 
    ∇̃W - ∇ W^2 ≤∇̃V - ∇ V^2 = σ^2.
    
    
    Pick a random index[The analysis in <cit.> provides guarantees on the probability measure ν_T 1/N∑_t=1^T ν_t where ^t ∼ν_t. The ^T defined here has law ν_T.]t∈{ 1, 2, ..., T} and set ^T ^t. Then Corollary 18 of <cit.> with D^2 = σ^2 and M_2 = 0 implies D(^T ) ≤ϵ, provided 
    β≤min{ϵ/2ł(  Ld + σ^2   )̊, 1/L},    T ≥_2^2(^0, )/βϵ.
    
    Solving for T in terms of ϵ establishes the theorem.
    
    
    
    
    
    § STOCHASTIC GRADIENTS FOR DIRICHLET POSTERIORS
    
    In order to apply SMLD, one must have, for each term V_i, the corresponding dual W_i defined via (<ref>).
    In this appendix, we derive a closed-form expression in the case of the Dirichlet posterior (<ref>). 
    
    Recall that the Dirichlet posterior (<ref>) consists of a Dirichlet prior and categorical data observations <cit.>. Let N ∑_ℓ=1^d+1 n_ℓ, where n_ℓ is the number of observations for category ℓ, and suppose that the parameters α_ℓ's are given. If the ith data is in category c_i ∈{1, 2, ..., d+1}, then we can define V_i()  - ∑_ℓ=1^d+1𝕀_{ℓ = c_i}log x_ℓ - 1/N∑_ℓ=1^d+1 (α_ℓ - 1) log x_ℓ so that Assumption <ref> holds. In view of Lemma <ref>, The corresponding dual W_i is, up to a constant, given by 
    
    W_i () = - ∑_ℓ=1^d 𝕀_{ℓ=c_i} y_ℓ -  ∑_ℓ=1^d α_ℓ/N y_ℓ +  + ł(∑_ℓ=1^d+1α_ℓ/N)̊().
    
    
    
    
    
    
    
    Similarly, if we take a mini-batch B of the data with |B| = b, then 
    N/bW̃ () N/b∑_i∈ B W_i () = - ∑_ℓ=1^d ł(Nm_ℓ/b + α_ℓ)̊ y_ℓ +  ł(N+ ∑_ℓ=1^d+1α_ℓ)̊(), 
    
    
    
    
    
    
    
    where m_ℓ is the number of observations of category ℓ in the set B. Apparently, the gradient of (<ref>) is (<ref>). 
    
    
    
    
    
    
    
    
    
    
    
    
    
    § MORE ON EXPERIMENTS
    
    
     §.§ Synthetic Data
    
    Figure <ref> reports the total variation error along the 8th dimension of the synthetic experiment in Section <ref>. Compared to Figure <ref> in the main text, it is evident that MLD achieves an even stronger performance than SGRLD, especially in the saturation error phase. 
    
    
    
    
    
     §.§ Comparison against SGRHMC for Latent Dirichlet Allocation
    
    The only difference between the experimental setting of <cit.> and the main text is the number of topics (50 vs. 100). In this appendix, we run SMLD-approximate under the setting of <cit.> and directly compare against the results reported in <cit.>. We have also included the SGRLD as a baseline.
    
    
    
    Figure <ref> reports the perplexity on the test data. According to <cit.>, the best perplexity achieved by SGRHMC up to [group-separator=,]10000 documents is approximately [group-separator=,]1400, which is worse than the 1323 by SMLD-approximate. Moreover, from Figure 3 of <cit.>, we see that the SGRHMC yields comparable performance as SGRLD for 2 out 3 independent runs, especially in the beginning phase, whereas the SMLD-approximate has sizeable lead over SGRLD at any stage of the experiment. The potential reason for this improvement is, similar to SGRLD, that the SGRHMC exploits the Riemannian Hamiltonian dynamics, which is more complicated than MLD and hence more sensitive to the discretization error.
