Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization

Flow-based models are powerful tools for designing probabilistic models with\ntractable density. This paper introduces Convex Potential Flows (CP-Flow), a\nnatural and efficient parameterization of invertible models inspired by the\noptimal transport (OT) theory. CP-Flows are the gradient map of a strongly\nconvex neural potential function. The convexity implies invertibility and\nallows us to resort to convex optimization to solve the convex conjugate for\nefficient inversion. To enable maximum likelihood training, we derive a new\ngradient estimator of the log-determinant of the Jacobian, which involves\nsolving an inverse-Hessian vector product using the conjugate gradient method.\nThe gradient estimator has constant-memory cost, and can be made effectively\nunbiased by reducing the error tolerance level of the convex optimization\nroutine. Theoretically, we prove that CP-Flows are universal density\napproximators and are optimal in the OT sense. Our empirical results show that\nCP-Flow performs competitively on standard benchmarks of density estimation and\nvariational inference.\n