  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.
