Sliced-Wasserstein Gradient Flows - Archive ouverte HAL Access content directly
Preprints, Working Papers, ... Year : 2021

Sliced-Wasserstein Gradient Flows


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, this bilevel optimization problem 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. Interestingly, we also show empirically that these gradient flows are strongly related to the usual Wasserstein gradient flows, and that they can be used to minimize efficiently diverse machine learning functionals.

Dates and versions

hal-03426809 , version 1 (12-11-2021)



Clément Bonet, Nicolas Courty, François Septier, Lucas Drumetz. Sliced-Wasserstein Gradient Flows. 2021. ⟨hal-03426809⟩
103 View
0 Download



Gmail Mastodon Facebook X LinkedIn More