Two-Loop Stochastic Mirror Langevin Algorithms for Constrained Sampling

2026-10-07

AI Summary

This study addresses the challenges of controlling discretization bias and noise error in stochastic gradient sampling under convex set constraints by proposing a two-loop stochastic mirror Langevin algorithm. Methodologically, it integrates mirror maps with Euler–Maruyama discretization and introduces a novel inner-outer double-loop warm-up mechanism, wherein the outer loop employs decaying step sizes and the inner loop performs warm-up sampling. A coupled schedule linking step sizes and periods is established to eliminate mixing costs. Theoretically, a Wasserstein distance-based analysis yields an optimal convergence rate of O(T^{-1/2}) while removing logarithmic penalty terms. Experiments across diverse statistical settings demonstrate significant improvements in convergence efficiency.

Abstract

We study the problem of sampling from a target distribution $π(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.

Problem

• constrained sampling
• stochastic first-order oracle
• Mirror Langevin algorithm
• discretization bias
• Wasserstein guarantees

Innovation

• Stochastic Mirror Langevin Algorithm
• Two-loop scheme
• Constrained sampling
• Wasserstein guarantees
• Warm-start