Input Convex Gradient Networks

TL;DR

This paper introduces Input Convex Gradient Networks (ICGNs), a neural network approach for modeling convex gradients, with theoretical analysis and empirical comparisons to existing models, demonstrating improved fitting capabilities in simple cases.

Contribution

The paper proposes ICGNs, a novel neural network architecture for modeling convex gradients, and provides theoretical and empirical analysis comparing them to ICNNs.

Findings

Single layer ICGN fits toy examples better than ICNN.

ICGNs can be extended to deeper networks.

Connections to Riemannian geometry are explored.

Abstract

The gradients of convex functions are expressive models of non-trivial vector fields. For example, Brenier's theorem yields that the optimal transport map between any two measures on Euclidean space under the squared distance is realized as a convex gradient, which is a key insight used in recent generative flow models. In this paper, we study how to model convex gradients by integrating a Jacobian-vector product parameterized by a neural network, which we call the Input Convex Gradient Network (ICGN). We theoretically study ICGNs and compare them to taking the gradient of an Input-Convex Neural Network (ICNN), empirically demonstrating that a single layer ICGN can fit a toy example better than a single layer ICNN. Lastly, we explore extensions to deeper networks and connections to constructions from Riemannian geometry.