Two-Loop Stochastic Mirror Langevin Algorithms for Constrained Sampling

Abstract

We study the problem of sampling from a target distribution $\pi(x)\propto e^{-f(x)}$ supported on a convex set $ X\subseteq\mathbb R^d$, when the potential $f$ is accessible only through a stochastic first-order oracle. Mirror Langevin algorithms provide a natural approach to constrained sampling by transporting the problem to an unconstrained dual space and discretizing the resulting Mirror Langevin diffusion. Existing implementations, however, typically use a fixed discretization step size and consequently retain a nonvanishing discretization bias at any fixed step size. Moreover, their direct extension to settings with noisy gradient information entails the challenge of controlling both discretization and stochastic-oracle error. We study a stochastic first-order version of the Mirror Langevin Algorithm (sFO-MLA) and, as our main contribution, develop a warm-started two-loop implementation in which an outer loop progressively decreases the step size while an inner loop runs sFO-MLA (with a fixed step size) for an appropriately chosen epoch length. The construction provides a principled schedule linking step sizes and epoch lengths, so that successive epochs warm-start from increasingly accurate distributions rather than repeatedly paying the cost of mixing from a cold start. We establish finite-time Wasserstein guarantees for sFO-MLA that explicitly separate mixing, Euler--Maruyama discretization, and stochastic-gradient errors. These bounds yield a fixed horizon rate of $\widetilde O(T^{-1/2})$ and show that the two-loop scheme removes the associated logarithmic penalty, attaining the canonical $O(T^{-1/2})$ rate under a geometric step-size schedule and corresponding epoch lengths. We illustrate the methodology in two statistically distinct settings.