mlsanityData to 14 Oct 2026

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

arXiv 2106.00774 ↗First posted 1 Jun 2021stat.ML cs.LG cs.NAcited by 53

Not on the map. Why?

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.

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

28 of 87 references.

References · 28

  1. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
  2. 2013Auto-Encoding Variational Bayescited by 4,053
  3. 2013Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distancescited by 2,316
  4. 2015Variational Inference with Normalizing Flowscited by 2,170
  5. 2016Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithmcited by 507
  6. 2016Input Convex Neural Networkscited by 414
  7. 2017Learning Generative Models with Sinkhorn Divergencescited by 314
  8. 2018Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Modelscited by 275
  9. 2019Residual Flows for Invertible Generative Modelingcited by 257
  10. 2017Non-convex learning via Stochastic Gradient Langevin Dynamics: a nonasymptotic analysiscited by 243
  11. 2019Optimal transport mapping via input convex neural networkscited by 206
  12. 2017Stein Variational Gradient Descent as Gradient Flowcited by 152
16 more
  1. 2019Maximum Mean Discrepancy Gradient Flowcited by 127
  2. 2018Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problemcited by 123
  3. 2020Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimizationcited by 90
  4. 2019On the geometry of Stein variational gradient descentcited by 82
  5. 2018Sliced-Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusionscited by 77
  6. 2020Hutch++: Optimal Stochastic Trace Estimationcited by 63
  7. 2021Kernel Stein Discrepancy Descentcited by 57
  8. 2021Continuous Wasserstein-2 Barycenter Estimation without Minimax Optimizationcited by 51
  9. 2021Projected Wasserstein gradient descent for high-dimensional Bayesian inferencecited by 24
  10. 2020On the Convergence of Gradient Descent in GANs: MMD GAN As a Gradient Flowcited by 21
  11. 2021Convergence Rates of Gradient Methods for Convex Optimization in the Space of Measurescited by 20
  12. 2019Regularity as Regularization: Smooth and Strongly Convex Brenier Potentials in Optimal Transportcited by 19
  13. 2021Wasserstein Proximal of GANscited by 15
  14. 2018Approximate inference with Wasserstein gradient flowscited by 12
  15. 2020Unbalanced Sobolev Descentcited by 11
  16. 2018Proximal Recursion for Solving the Fokker-Planck Equationcited by 7

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2021: 5; 2022: 13; 2023: 13; 2024: 9; 2025: 7; 2026: 6.

53 papers.

Sort:

By area:

  1. 2021Proximal Optimal Transport Modeling of Population Dynamicscited by 93
  2. 2022Supervised Training of Conditional Monge Mapscited by 61
  3. 2021Variational Wasserstein gradient flowcited by 53
  4. 2023A Computational Framework for Solving Wasserstein Lagrangian Flowscited by 46
  5. 2021Neural Monge Map estimation and its applicationscited by 40
  6. 2022Normalizing flow neural networks by JKO schemecited by 39
  7. 2023Generative Sliced MMD Flows with Riesz Kernelscited by 37
  8. 2022Neural Optimal Transport with General Cost Functionalscited by 36
  9. 2023Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein Spacecited by 36
  10. 2023Convergence of flow-based generative models via proximal gradient descent in Wasserstein spacecited by 34
  11. 2022Kernel Neural Optimal Transportcited by 31
  12. 2022On amortizing convex conjugates for optimal transportcited by 31
38 more
  1. 2023Energy-guided Entropic Neural Optimal Transportcited by 26
  2. 2023Posterior Sampling Based on Gradient Flows of the MMD with Negative Distance Kernelcited by 25
  3. 2024Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transportcited by 23
  4. 2024Learning diffusion at lightspeedcited by 22
  5. 2021Efficient Gradient Flows in Sliced-Wasserstein Spacecited by 21
  6. 2023Unifying GANs and Score-Based Diffusion as Generative Particle Modelscited by 21
  7. 2022Particle-based Variational Inference with Preconditioned Functional Gradient Flowcited by 20
  8. 2021Input Convex Gradient Networkscited by 17
  9. 2022Neural Unbalanced Optimal Transport via Cycle-Consistent Semi-Couplingscited by 17
  10. 2022Taming Hyperparameter Tuning in Continuous Normalizing Flows Using the JKO Schemecited by 17
  11. 2023Self-Consistent Velocity Matching of Probability Flowscited by 17
  12. 2023Towards Understanding the Dynamics of Gaussian-Stein Variational Gradient Descentcited by 13
  13. 2023Particle-based Variational Inference with Generalized Wasserstein Gradient Flowcited by 10
  14. 2022Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimizationcited by 9
  15. 2023Neural Wasserstein Gradient Flows for Maximum Mean Discrepancies with Riesz Kernelscited by 9
  16. 2024Neural Sinkhorn Gradient Flowcited by 9
  17. 2025Learning of Population Dynamics: Inverse Optimization Meets JKO Schemecited by 9
  18. 2024Path-Guided Particle-based Samplingcited by 8
  19. 2024Importance Corrected Neural JKO Samplingcited by 7
  20. 2022Universal Neural Optimal Transportcited by 6
  21. 2024Non-geodesically-convex optimization in the Wasserstein spacecited by 5
  22. 2024Local Flow Matching Generative Modelscited by 5
  23. 2025Are We Really Learning the Score Function? Reinterpreting Diffusion Models Through Wasserstein Gradient Flow Matchingcited by 5
  24. 2023Real-Time Machine-Learning-Based Optimization Using Input Convex Long Short-Term Memory Networkcited by 3
  25. 2022Learning with Stochastic Orderscited by 2
  26. 2025Flow-based generative models as iterative algorithms in probability spacecited by 2
  27. 2026The physics of AI weather modelscited by 2
  28. 2022Proximal Mean Field Learning in Shallow Neural Networkscited by 1
  29. 2024Semi-Implicit Functional Gradient Flow for Efficient Samplingcited by 1
  30. 2025Computational and Statistical Asymptotic Analysis of the JKO Scheme for Iterative Algorithms to update distributionscited by 1
  31. 2025Displacement-Sparse Neural Optimal Transportcited by 1
  32. 2025Data Generation without Function Estimationcited by 1
  33. 2026Leveraging Gauge Freedom for Learning Non-Gradient Population Dynamics of Stochastic Systemscited by 1
  34. 2026Two-Parameter Flows for Learning Population Dynamics of Physical Systemscited by 1
  35. 2022A Distributed Algorithm for Measure-valued Optimization with Additive Objectivecited by 0
  36. 2023Geometry-Aware Normalizing Wasserstein Flows for Optimal Causal Inferencecited by 0
  37. 2024Functional Gradient Flows for Constrained Samplingcited by 0
  38. 2025Universal Representation of Generalized Convex Functions and their Gradientscited by 0

50 most cited of 53 All 53 (JSON)