Decentralized Proximal Stochastic Gradient Langevin Dynamics
| Authors | Mohammad Rafiqul Islam & Lingjiong Zhu |
| Year | 2026 |
| Field | Statistics / ML |
| arXiv | 2605.00723 |
| Download | |
| Categories | stat.ML, cs.LG |
Abstract
We propose Decentralized Proximal Stochastic Gradient Langevin Dynamics (DE-PSGLD), a decentralized Markov chain Monte Carlo (MCMC) algorithm for sampling from a log-concave probability distribution constrained to a convex domain. Constraints are enforced through a shared proximal regularization based on the Moreau-Yosida envelope, enabling unconstrained updates while preserving consistency with the target constrained posterior. We establish non-asymptotic convergence guarantees in the 2-Wasserstein distance for both individual agent iterates and their network averages. Our analysis shows that DE-PSGLD converges to a regularized Gibbs distribution and quantifies the bias introduced by the proximal approximation. We evaluate DE-PSGLD for different sampling problems on synthetic and real datasets. As the first decentralized approach for constrained domains, our algorithm exhibits fast posterior concentration and high predictive accuracy.
Engineering Breakdown
Plain English
This paper proposes DE-PSGLD, a decentralized algorithm for sampling from constrained probability distributions across multiple agents or nodes in a network. The key innovation is using a proximal regularization technique (Moreau-Yosida envelope) to enforce convex constraints without explicitly projecting iterates, allowing each agent to perform unconstrained updates while staying consistent with the target constrained posterior. The authors prove non-asymptotic convergence guarantees in 2-Wasserstein distance for both individual agents and network averages, quantifying the bias introduced by the proximal approximation. This addresses a gap in distributed MCMC methods where handling constraints across decentralized networks was previously difficult or required expensive projection steps.
Core Technical Contribution
The core novelty is combining three previously separate ideas: decentralized MCMC (agents sampling locally with network communication), stochastic gradient Langevin dynamics (adding noise for proper posterior sampling), and proximal methods for constrained optimization. Rather than explicitly projecting onto the constraint set—which is expensive and breaks the Markov chain structure—the authors embed constraints into the update rule via the Moreau-Yosida envelope, a smooth approximation. This allows unconstrained updates that implicitly respect constraints while maintaining theoretical guarantees. The non-asymptotic convergence analysis in 2-Wasserstein distance with explicit bias quantification is new, showing how the proximal approximation trades off computational cost against posterior approximation quality.
How It Works
Each agent in the network maintains a local parameter estimate and performs a stochastic gradient Langevin update at each iteration, with a noise term scaled by temperature for proper Bayesian sampling. Rather than projecting the updated parameters onto the constraint set (expensive, breaks MCMC guarantees), the algorithm incorporates constraints through a proximal term based on the Moreau-Yosida envelope—a smooth approximation of the constraint indicator function that penalizes constraint violations. Agents communicate with their neighbors (network topology), averaging local iterates to achieve consensus while the algorithm converges. The algorithm's update rule includes three components: a stochastic gradient step (using mini-batch data), a Langevin noise term (for exploration and sampling correctness), and the proximal regularization term (constraint enforcement). The theoretical analysis tracks both individual agent convergence and network-averaged iterate convergence, measuring distance to the regularized Gibbs distribution using 2-Wasserstein metrics.
Production Impact
For distributed Bayesian inference systems—think federated learning with privacy constraints, multi-agent optimization under resource limits, or decentralized sensor networks computing posteriors—this removes the expensive constraint-projection bottleneck that typically requires synchronization and centralized computation. In production pipelines, you could deploy inference workloads across edge devices without gathering data centrally, and constraints (e.g., parameter bounds, domain restrictions) would be enforced automatically through the proximal term rather than explicit projections that slow convergence and create synchronization overhead. The non-asymptotic convergence guarantees let you set iteration budgets with confidence: you can compute a priori how many passes over data are needed to reach a target error tolerance, essential for resource-constrained environments. Trade-offs include added computational cost per iteration (computing the Moreau-Yosida regularization) and the accuracy-cost trade-off via the proximal approximation parameter, which engineers must tune. Network communication cost scales with network size, so this approach suits moderate-sized clusters (10s-100s of agents) more than massive-scale systems.
Limitations and When Not to Use This
The analysis assumes log-concave target distributions and convex constraint sets, which excludes many real problems with non-convex constraints or multimodal posteriors common in modern deep learning. The algorithm requires full knowledge or good estimates of the log-likelihood gradient, making it impractical for black-box or simulation-based inference where gradients are unavailable. Convergence rates are likely slower than centralized MCMC due to the decentralization overhead and proximal approximation bias, and the paper doesn't empirically compare wall-clock time against baselines—only iteration counts or synthetic experiments. The method assumes a fixed network topology and reliable communication, failing gracefully under Byzantine agents or adversarial node failures. Scalability to high-dimensional problems (e.g., Bayesian neural networks with millions of parameters) is unclear from the abstract; proximal operator computation itself can become expensive in high dimensions.
Research Context
This work bridges decentralized optimization (consensus-based distributed algorithms) with Bayesian MCMC sampling, building on recent progress in proximal stochastic gradient methods and networked Langevin dynamics. It extends prior work in decentralized MCMC by removing the need for explicit projections onto constraint sets—a known bottleneck in algorithms like decentralized projected gradient methods. The proximal envelope technique is borrowed from constrained optimization literature but applied here to the MCMC setting for the first time with formal non-asymptotic guarantees. This opens a research direction toward decentralized sampling for constrained Bayesian problems in federated learning, distributed robotics, and privacy-preserving inference where constraints and decentralization are both essential.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::