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

Abstract

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 parameterize 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 use various experiments to demonstrate its feasibility and validity in approximating solutions of low-dimensional partial differential equations with known solutions. We also explore the use of our JKO-ICNN approach in high dimensions with an experiment in controlled generation for molecular discovery.Preprint. Under review.

Neural Optimal Transport with General Cost Functionals

“…The problem is a generalization of strong OT (2), weak OT (3), regularized OT (5); we call problem (6) a general OT problem. Surprisingly, regularized OT (5) represents the same problem: it is enough to put c(x, y) ≡ 0, γ = 1 and R(π) = F(π) to obtain general OT (6) from regularized OT (5). That is, regularized OT (5) and general OT (6) can be viewed as equivalent formulations.…”

“…10] consider analogous to (10) formulations restricted to convex potentials; they use Input Convex Neural Networks [6] to approximate them. These nets are popular in OT [41,56,19,11,5] but OT algorithms based on them are outperformed [40] by the above-mentioned unrestricted formulations. In [26,66,17,20], the authors propose methods for f -divergence regularized functionals (5).…”

Input Convex Gradient Networks

“…Optimal Transport: Makkuva et al [10] explored using Input Convex Neural Networks for learning transportation maps, while Alvarez-Melis et al [25] and Mokrov et al [26] used ICNN's for the Kantorovich dual specifically in the setting of Wasserstein gradient flows [27]. Fan et al [28] also attempted solving the Wassertein Barycenter [29] problem using ICNNs.…”

Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization

“…The time and spatial discretizations of Wasserstein gradient flows are extensively studied in literature (Jordan et al, 1998;Junge et al, 2017;Carrillo et al, 2021a,b;Bonet et al, 2021;Liutkus et al, 2019;Frogner & Poggio, 2020). Recently, neural networks have been applied in solving or approximating Wasserstein gradient flows (Mokrov et al, 2021;Lin et al, 2021b,a;Alvarez-Melis et al, 2021;Bunne et al, 2021;Hwang et al, 2021;Fan et al, 2021). For sampling algorithms, di Langosco et al (2021) learns the transportation function by solving an unregularized variational problem in the family of vector-output deep neural networks.…”