# Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks

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

## Key facts

- Cited by 53 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 87 references that are in the map.
- First posted on arXiv 1 Jun 2021.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm](https://mlsanity.com/paper/1608.04471/) (2016)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Learning Generative Models with Sinkhorn Divergences](https://mlsanity.com/paper/1706.00292/) (2017)
- [Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Models](https://mlsanity.com/paper/1811.12823/) (2018)
- [Residual Flows for Invertible Generative Modeling](https://mlsanity.com/paper/1906.02735/) (2019)
- [Non-convex learning via Stochastic Gradient Langevin Dynamics: a nonasymptotic analysis](https://mlsanity.com/paper/1702.03849/) (2017)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (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)
- [Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization](https://mlsanity.com/paper/2012.05942/) (2020)
- [On the geometry of Stein variational gradient descent](https://mlsanity.com/paper/1912.00894/) (2019)
- [Sliced-Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusions](https://mlsanity.com/paper/1806.08141/) (2018)
- [Hutch++: Optimal Stochastic Trace Estimation](https://mlsanity.com/paper/2010.09649/) (2020)
- [Kernel Stein Discrepancy Descent](https://mlsanity.com/paper/2105.09994/) (2021)
- [Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimization](https://mlsanity.com/paper/2102.01752/) (2021)
- [Projected Wasserstein gradient descent for high-dimensional Bayesian inference](https://mlsanity.com/paper/2102.06350/) (2021)
- [On the Convergence of Gradient Descent in GANs: MMD GAN As a Gradient Flow](https://mlsanity.com/paper/2011.02402/) (2020)
- [Convergence Rates of Gradient Methods for Convex Optimization in the Space of Measures](https://mlsanity.com/paper/2105.08368/) (2021)
- [Regularity as Regularization: Smooth and Strongly Convex Brenier Potentials in Optimal Transport](https://mlsanity.com/paper/1905.10812/) (2019)
- [Wasserstein Proximal of GANs](https://mlsanity.com/paper/2102.06862/) (2021)
- [Approximate inference with Wasserstein gradient flows](https://mlsanity.com/paper/1806.04542/) (2018)
- [Unbalanced Sobolev Descent](https://mlsanity.com/paper/2009.14148/) (2020)
- [Proximal Recursion for Solving the Fokker-Planck Equation](https://mlsanity.com/paper/1809.10844/) (2018)

## Built on by (most cited first, 30 of 53)

- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Variational Wasserstein gradient flow](https://mlsanity.com/paper/2112.02424/) (2021)
- [A Computational Framework for Solving Wasserstein Lagrangian Flows](https://mlsanity.com/paper/2310.10649/) (2023)
- [Neural Monge Map estimation and its applications](https://mlsanity.com/paper/2106.03812/) (2021)
- [Normalizing flow neural networks by JKO scheme](https://mlsanity.com/paper/2212.14424/) (2022)
- [Generative Sliced MMD Flows with Riesz Kernels](https://mlsanity.com/paper/2305.11463/) (2023)
- [Neural Optimal Transport with General Cost Functionals](https://mlsanity.com/paper/2205.15403/) (2022)
- [Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein Space](https://mlsanity.com/paper/2304.05398/) (2023)
- [Convergence of flow-based generative models via proximal gradient descent in Wasserstein space](https://mlsanity.com/paper/2310.17582/) (2023)
- [Kernel Neural Optimal Transport](https://mlsanity.com/paper/2205.15269/) (2022)
- [On amortizing convex conjugates for optimal transport](https://mlsanity.com/paper/2210.12153/) (2022)
- [Energy-guided Entropic Neural Optimal Transport](https://mlsanity.com/paper/2304.06094/) (2023)
- [Posterior Sampling Based on Gradient Flows of the MMD with Negative Distance Kernel](https://mlsanity.com/paper/2310.03054/) (2023)
- [Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport](https://mlsanity.com/paper/2402.05443/) (2024)
- [Learning diffusion at lightspeed](https://mlsanity.com/paper/2406.12616/) (2024)
- [Efficient Gradient Flows in Sliced-Wasserstein Space](https://mlsanity.com/paper/2110.10972/) (2021)
- [Unifying GANs and Score-Based Diffusion as Generative Particle Models](https://mlsanity.com/paper/2305.16150/) (2023)
- [Particle-based Variational Inference with Preconditioned Functional Gradient Flow](https://mlsanity.com/paper/2211.13954/) (2022)
- [Input Convex Gradient Networks](https://mlsanity.com/paper/2111.12187/) (2021)
- [Neural Unbalanced Optimal Transport via Cycle-Consistent Semi-Couplings](https://mlsanity.com/paper/2209.15621/) (2022)
- [Taming Hyperparameter Tuning in Continuous Normalizing Flows Using the JKO Scheme](https://mlsanity.com/paper/2211.16757/) (2022)
- [Self-Consistent Velocity Matching of Probability Flows](https://mlsanity.com/paper/2301.13737/) (2023)
- [Towards Understanding the Dynamics of Gaussian-Stein Variational Gradient Descent](https://mlsanity.com/paper/2305.14076/) (2023)
- [Particle-based Variational Inference with Generalized Wasserstein Gradient Flow](https://mlsanity.com/paper/2310.16516/) (2023)
- [Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization](https://mlsanity.com/paper/2205.13098/) (2022)
- [Neural Wasserstein Gradient Flows for Maximum Mean Discrepancies with Riesz Kernels](https://mlsanity.com/paper/2301.11624/) (2023)
- [Neural Sinkhorn Gradient Flow](https://mlsanity.com/paper/2401.14069/) (2024)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [Path-Guided Particle-based Sampling](https://mlsanity.com/paper/2412.03312/) (2024)

## Abstract (arXiv)

Gradient flows are a powerful tool for optimizing functionals in general metric spaces, including the space of probabilities endowed with the Wasserstein metric. A typical approach to solving this optimization problem relies on its connection to the dynamic formulation of optimal transport and the celebrated Jordan-Kinderlehrer-Otto (JKO) scheme. However, this formulation involves optimization over convex functions, which is challenging, especially in high dimensions. In this work, we propose an approach that relies on the recently introduced input-convex neural networks (ICNN) to parametrize the space of convex functions in order to approximate the JKO scheme, as well as in designing functionals over measures that enjoy convergence guarantees. We derive a computationally efficient implementation of this JKO-ICNN framework and experimentally demonstrate its feasibility and validity in approximating solutions of low-dimensional partial differential equations with known solutions. We also demonstrate its viability in high-dimensional applications through an experiment in controlled generation for molecular discovery.