# Stochastic Gradient Descent over P2

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 7 of its 25 references that are in the map.
- First posted on arXiv 11 Sep 2026.

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

- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [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)
- [A mathematical perspective on Transformers](https://mlsanity.com/paper/2312.10794/) (2023)
- [Stochastic Modified Equations and Dynamics of Stochastic Gradient Algorithms I: Mathematical Foundations](https://mlsanity.com/paper/1811.01558/) (2018)
- [Overparameterization of deep ResNet: zero loss and mean-field analysis](https://mlsanity.com/paper/2105.14417/) (2021)
- [One-Step Generative Modeling via Wasserstein Gradient Flows](https://mlsanity.com/paper/2605.11755/) (2026)

## Abstract (arXiv)

Stochastic gradient descent (SGD) admits diffusion approximations that replace the complicated randomness of stochastic gradients by Gaussian noise, providing a powerful tool for understanding its dynamics and long-time behavior. We investigate whether an analogous approximation principle holds for optimization over probability measures, where the objective is a functional defined on the Wasserstein space P2. The nonlinear geometry and infinite-dimensional nature of P2 prevent a direct extension of the classical Euclidean theory. Using Lions differentiability, we lift the problem to a linear Hilbert space, where higher-order differential calculus becomes available. We then construct a Gaussian random-field approximation whose velocity field matches the mean and covariance of the original stochastic gradient. By exploiting this moment matching through higher-order Taylor expansions, we show that the Gaussian approximation captures the SGD dynamics with second-order weak accuracy. Our result provides a rigorous foundation for replacing sample-driven randomness by analytically tractable Gaussian fluctuations in stochastic optimization over probability measures.