# Neural Wasserstein Gradient Flows for Maximum Mean Discrepancies with Riesz Kernels

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 54 references that are in the map.
- First posted on arXiv 27 Jan 2023.

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

- [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)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Stochastic Normalizing Flows](https://mlsanity.com/paper/2002.06707/) (2020)
- [Maximum Mean Discrepancy Gradient Flow](https://mlsanity.com/paper/1906.04370/) (2019)
- [Neural Optimal Transport](https://mlsanity.com/paper/2201.12220/) (2022)
- [Flows for simultaneous manifold learning and density estimation](https://mlsanity.com/paper/2003.13913/) (2020)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Variational Wasserstein gradient flow](https://mlsanity.com/paper/2112.02424/) (2021)
- [Refining Deep Generative Models via Discriminator Gradient Flow](https://mlsanity.com/paper/2012.00780/) (2020)
- [KALE Flow: A Relaxed KL Gradient Flow for Probabilities with Disjoint Support](https://mlsanity.com/paper/2106.08929/) (2021)
- [Deep Generative Learning via Variational Gradient Flow](https://mlsanity.com/paper/1901.08469/) (2019)
- [Stochastic Normalizing Flows for Inverse Problems: a Markov Chains Viewpoint](https://mlsanity.com/paper/2109.11375/) (2021)
- [Generalized Normalizing Flows via Markov Chains](https://mlsanity.com/paper/2111.12506/) (2021)
- [Efficient Gradient Flows in Sliced-Wasserstein Space](https://mlsanity.com/paper/2110.10972/) (2021)
- [Particle-based Variational Inference with Preconditioned Functional Gradient Flow](https://mlsanity.com/paper/2211.13954/) (2022)
- [Neural Variational Gradient Descent](https://mlsanity.com/paper/2107.10731/) (2021)
- [Estimating Barycenters of Measures in High Dimensions](https://mlsanity.com/paper/2007.07105/) (2020)
- [Wasserstein Steepest Descent Flows of Discrepancies with Riesz Kernels](https://mlsanity.com/paper/2211.01804/) (2022)
- [Wasserstein Proximal of GANs](https://mlsanity.com/paper/2102.06862/) (2021)
- [The Deep Minimizing Movement Scheme](https://mlsanity.com/paper/2109.14851/) (2021)
- [Optimal transport natural gradient for statistical manifolds with continuous sample space](https://mlsanity.com/paper/1805.08380/) (2018)
- [Large-Scale Optimal Transport via Adversarial Training with Cycle-Consistency](https://mlsanity.com/paper/2003.06635/) (2020)
- [Kernelized Wasserstein Natural Gradient](https://mlsanity.com/paper/1910.09652/) (2019)
- [Wasserstein Gradient Flows of the Discrepancy with Distance Kernel on the Line](https://mlsanity.com/paper/2301.04441/) (2023)
- [Curve Based Approximation of Measures on Manifolds by Discrepancy Minimization](https://mlsanity.com/paper/1910.06124/) (2019)

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

- [Generative Sliced MMD Flows with Riesz Kernels](https://mlsanity.com/paper/2305.11463/) (2023)
- [Posterior Sampling Based on Gradient Flows of the MMD with Negative Distance Kernel](https://mlsanity.com/paper/2310.03054/) (2023)
- [Deep MMD Gradient Flow without adversarial training](https://mlsanity.com/paper/2405.06780/) (2024)
- [Weighted quantization using MMD: From mean field to mean shift via gradient flows](https://mlsanity.com/paper/2502.10600/) (2025)
- [Fast Kernel Summation in High Dimensions via Slicing and Fourier Transforms](https://mlsanity.com/paper/2401.08260/) (2024)
- [Kernel Semi-Implicit Variational Inference](https://mlsanity.com/paper/2405.18997/) (2024)
- [Importance Corrected Neural JKO Sampling](https://mlsanity.com/paper/2407.20444/) (2024)
- [Manifold learning in Wasserstein space](https://mlsanity.com/paper/2311.08549/) (2023)
- [A Kernel Approach for Semi-implicit Variational Inference](https://mlsanity.com/paper/2601.12023/) (2026)

## Abstract (arXiv)

Wasserstein gradient flows of maximum mean discrepancy (MMD) functionals with non-smooth Riesz kernels show a rich structure as singular measures can become absolutely continuous ones and conversely. In this paper we contribute to the understanding of such flows. We propose to approximate the backward scheme of Jordan, Kinderlehrer and Otto for computing such Wasserstein gradient flows as well as a forward scheme for so-called Wasserstein steepest descent flows by neural networks (NNs). Since we cannot restrict ourselves to absolutely continuous measures, we have to deal with transport plans and velocity plans instead of usual transport maps and velocity fields. Indeed, we approximate the disintegration of both plans by generative NNs which are learned with respect to appropriate loss functions. In order to evaluate the quality of both neural schemes, we benchmark them on the interaction energy. Here we provide analytic formulas for Wasserstein schemes starting at a Dirac measure and show their convergence as the time step size tends to zero. Finally, we illustrate our neural MMD flows by numerical examples.
