Optimizing Functionals On The Space Of Probabilities With Input Convex Neural Networks.pdf

The JKO scheme can be seen as a generalization of the implicit Euler method on the probability space endowed with the Wasserstein metric. This approach has various appealing theoretical convergence properties owing to a notion of convexity of probability functionals, known as geodesic convexity. Several computational approaches to JKO have been proposed, among them an elegant method introduced in Benamou et al. (2014) that reformulates the JKO variational problem on probability measures as an optimization problem on the space of convex functions. However, computing updates involves solving an optimization over convex functions at each step, which is challenging in general.

We propose a computational approach to the JKO scheme that is scalable in high-dimensions. At the core of our approach are Input-Convex Neural Networks (ICNN), a recently proposed class of deep models that are convex with respect to their inputs. We use ICNNs to find parametric solutions to the reformulation of the JKO problem as optimization on the space of convex functions by Benamou et al. (2014). This leads to an approximation of the JKO scheme that we call JKO-ICNN.

To evaluate the soundness of our approach, we conduct experiments on well-known PDEs in low dimensions that have exact analytic solutions, allowing us to quantify the approximation quality of the gradient flows evolved with our method. We then use our approach in a high-dimensional setting, where we optimize a dataset of molecules to satisfy certain properties, such as drug-likeness (QED). The results show that our JKO-ICNN approach is successful at approximating solutions of PDEs and has the unique advantage of scalability in terms of optimizing generic probability functionals on the probability space in high dimensions.

The JKO-ICNN approach has the advantage of computational stability and amortization of computational cost, since the maps found while training JKO-ICNN on one sample from a distribution ρ0 generalize at transporting a new sample unseen during the training at no additional cost. While preparing this manuscript we became aware of concurrent work on approximating JKO with ICNNs by Mokrov et al. (2021) and Bunne et al. (2021). While the former is concerned exclusively with the Fokker-Planck equation, here we consider other classes of PDEs too. The latter tackles a different problem: learning dynamics with JKO, i.e., learning the functional whose JKO flow follows empirical observations.