# Beyond Propagation of Chaos: A Stochastic Algorithm for Mean Field Optimization

arXiv: https://arxiv.org/abs/2503.13115 · HTML: https://mlsanity.com/paper/2503.13115/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 51 references that are in the map.
- First posted on arXiv 17 Mar 2025.

## Builds on (most cited first, 21 of 21)

- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [Non-convex learning via Stochastic Gradient Langevin Dynamics: a nonasymptotic analysis](https://mlsanity.com/paper/1702.03849/) (2017)
- [Rapid Convergence of the Unadjusted Langevin Algorithm: Isoperimetry Suffices](https://mlsanity.com/paper/1903.08568/) (2019)
- [Analysis of Langevin Monte Carlo via convex optimization](https://mlsanity.com/paper/1802.09188/) (2018)
- [Variational inference via Wasserstein gradient flows](https://mlsanity.com/paper/2205.15902/) (2022)
- [The Randomized Midpoint Method for Log-Concave Sampling](https://mlsanity.com/paper/1909.05503/) (2019)
- [Mean-Field Langevin Dynamics: Exponential Convergence and Annealing](https://mlsanity.com/paper/2202.01009/) (2022)
- [Convex Analysis of the Mean Field Langevin Dynamics](https://mlsanity.com/paper/2201.10469/) (2022)
- [Stochastic Particle Gradient Descent for Infinite Ensembles](https://mlsanity.com/paper/1712.05438/) (2017)
- [SGD Generalizes Better Than GD (And Regularization Doesn't Help)](https://mlsanity.com/paper/2102.01117/) (2021)
- [Uniform-in-time propagation of chaos for mean field Langevin dynamics](https://mlsanity.com/paper/2212.03050/) (2022)
- [Improved Convergence Rate of Stochastic Gradient Langevin Dynamics with Variance Reduction and its Application to Optimization](https://mlsanity.com/paper/2203.16217/) (2022)
- [Improved Finite-Particle Convergence Rates for Stein Variational Gradient Descent](https://mlsanity.com/paper/2409.08469/) (2024)
- [Independent projections of diffusions: Gradient flows for variational inference and optimal mean field approximations](https://mlsanity.com/paper/2309.13332/) (2023)
- [Improved Particle Approximation Error for Mean Field Neural Networks](https://mlsanity.com/paper/2405.15767/) (2024)
- [Langevin Monte Carlo for strongly log-concave distributions: Randomized midpoint revisited](https://mlsanity.com/paper/2306.08494/) (2023)
- [Utilising the CLT Structure in Stochastic Gradient based Sampling : Improved Analysis and Faster Algorithms](https://mlsanity.com/paper/2206.03792/) (2022)
- [Provably Fast Finite Particle Variants of SVGD via Virtual Particle Stochastic Approximation](https://mlsanity.com/paper/2305.17558/) (2023)
- [Non-Asymptotic Analysis of Stochastic Approximation Algorithms for Streaming Data](https://mlsanity.com/paper/2109.07117/) (2021)
- [Faster Sampling via Stochastic Gradient Proximal Sampler](https://mlsanity.com/paper/2405.16734/) (2024)

## Built on by (most cited first, 1 of 1)

- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)

## Abstract (arXiv)

Gradient flow in the 2-Wasserstein space is widely used to optimize functionals over probability distributions and is typically implemented using an interacting particle system with n particles. Analyzing these algorithms requires showing (a) that the finite-particle system converges and/or (b) that the resultant empirical distribution of the particles closely approximates the optimal distribution (i.e., propagation of chaos). However, establishing efficient sufficient conditions can be challenging, as the finite particle system may produce heavily dependent random variables. In this work, we study the virtual particle stochastic approximation, originally introduced for Stein Variational Gradient Descent. This method can be viewed as a form of stochastic gradient descent in the Wasserstein space and can be implemented efficiently. In popular settings, we demonstrate that our algorithm's output converges to the optimal distribution under conditions similar to those for the infinite particle limit, and it produces i.i.d. samples without the need to explicitly establish propagation of chaos bounds.