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

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