# Efficient Gradient Flows in Sliced-Wasserstein Space

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

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 104 references that are in the map.
- First posted on arXiv 21 Oct 2021.

## Builds on (most cited first, 30 of 46)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [Masked Autoregressive Flow for Density Estimation](https://mlsanity.com/paper/1705.07057/) (2017)
- [Solving high-dimensional partial differential equations using deep learning](https://mlsanity.com/paper/1707.02568/) (2017)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Towards Principled Methods for Training Generative Adversarial Networks](https://mlsanity.com/paper/1701.04862/) (2017)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Flow++: Improving Flow-Based Generative Models with Variational Dequantization and Architecture Design](https://mlsanity.com/paper/1902.00275/) (2019)
- [Deep Generative Modelling: A Comparative Review of VAEs, GANs, Normalizing Flows, Energy-Based and Autoregressive Models](https://mlsanity.com/paper/2103.04922/) (2021)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [Stein Variational Gradient Descent as Gradient Flow](https://mlsanity.com/paper/1704.07520/) (2017)
- [Maximum Mean Discrepancy Gradient Flow](https://mlsanity.com/paper/1906.04370/) (2019)
- [Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem](https://mlsanity.com/paper/1802.08089/) (2018)
- [Max-Sliced Wasserstein Distance and its use for GANs](https://mlsanity.com/paper/1904.05877/) (2019)
- [Wasserstein-2 Generative Networks](https://mlsanity.com/paper/1909.13082/) (2019)
- [Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark](https://mlsanity.com/paper/2106.01954/) (2021)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Subspace Robust Wasserstein Distances](https://mlsanity.com/paper/1901.08949/) (2019)
- [Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization](https://mlsanity.com/paper/2012.05942/) (2020)
- [Statistical and Topological Properties of Sliced Probability Divergences](https://mlsanity.com/paper/2003.05783/) (2020)
- [On the geometry of Stein variational gradient descent](https://mlsanity.com/paper/1912.00894/) (2019)
- [Generative Modeling with Optimal Transport Maps](https://mlsanity.com/paper/2110.02999/) (2021)
- [Sliced-Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusions](https://mlsanity.com/paper/1806.08141/) (2018)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)

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

- [Energy-Based Sliced Wasserstein Distance](https://mlsanity.com/paper/2304.13586/) (2023)
- [Generative Sliced MMD Flows with Riesz Kernels](https://mlsanity.com/paper/2305.11463/) (2023)
- [Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport](https://mlsanity.com/paper/2402.05443/) (2024)
- [Quasi-Monte Carlo for 3D Sliced Wasserstein](https://mlsanity.com/paper/2309.11713/) (2023)
- [Hyperbolic Sliced-Wasserstein via Geodesic and Horospherical Projections](https://mlsanity.com/paper/2211.10066/) (2022)
- [Sliced Wasserstein with Random-Path Projecting Directions](https://mlsanity.com/paper/2401.15889/) (2024)
- [Markovian Sliced Wasserstein Distances: Beyond Independent Projections](https://mlsanity.com/paper/2301.03749/) (2023)
- [Neural Wasserstein Gradient Flows for Maximum Mean Discrepancies with Riesz Kernels](https://mlsanity.com/paper/2301.11624/) (2023)
- [Tree-Sliced Wasserstein Distance: A Geometric Perspective](https://mlsanity.com/paper/2406.13725/) (2024)
- [Measure transfer via stochastic slicing and matching](https://mlsanity.com/paper/2307.05705/) (2023)
- [Hierarchical Hybrid Sliced Wasserstein: A Scalable Metric for Heterogeneous Joint Distributions](https://mlsanity.com/paper/2404.15378/) (2024)
- [Convergence of SGD for Training Neural Networks with Sliced Wasserstein Losses](https://mlsanity.com/paper/2307.11714/) (2023)
- [Differentially Private Gradient Flow based on the Sliced Wasserstein Distance](https://mlsanity.com/paper/2312.08227/) (2023)
- [Approximation properties of slice-matching operators](https://mlsanity.com/paper/2310.10869/) (2023)
- [Understanding Learning with Sliced-Wasserstein Requires Rethinking Informative Slices](https://mlsanity.com/paper/2411.10651/) (2024)
- [Proximal Mean Field Learning in Shallow Neural Networks](https://mlsanity.com/paper/2210.13879/) (2022)
- [Computational and Statistical Asymptotic Analysis of the JKO Scheme for Iterative Algorithms to update distributions](https://mlsanity.com/paper/2501.06408/) (2025)
- [Data Generation without Function Estimation](https://mlsanity.com/paper/2507.08239/) (2025)
- [Repulsive Monte Carlo on the sphere for the sliced Wasserstein distance](https://mlsanity.com/paper/2509.10166/) (2025)
- [Subspace Langevin Monte Carlo](https://mlsanity.com/paper/2412.13928/) (2024)
- [Streaming Sliced Optimal Transport](https://mlsanity.com/paper/2505.06835/) (2025)

## Abstract (arXiv)

Minimizing functionals in the space of probability distributions can be done with Wasserstein gradient flows. To solve them numerically, a possible approach is to rely on the Jordan-Kinderlehrer-Otto (JKO) scheme which is analogous to the proximal scheme in Euclidean spaces. However, it requires solving a nested optimization problem at each iteration, and is known for its computational challenges, especially in high dimension. To alleviate it, very recent works propose to approximate the JKO scheme leveraging Brenier's theorem, and using gradients of Input Convex Neural Networks to parameterize the density (JKO-ICNN). However, this method comes with a high computational cost and stability issues. Instead, this work proposes to use gradient flows in the space of probability measures endowed with the sliced-Wasserstein (SW) distance. We argue that this method is more flexible than JKO-ICNN, since SW enjoys a closed-form differentiable approximation. Thus, the density at each step can be parameterized by any generative model which alleviates the computational burden and makes it tractable in higher dimensions.