Learning Brenier Potentials with Convex Generative Adversarial Neural Networks
Claudia Drygala, Hanno Gottschalk, Thomas Kruse, Ségolène Martin, Annika Mütze
TL;DR
This work addresses learning an optimal transport map between a simple source distribution and a Hölder-regular target by parameterizing the Brenier potential $\phi$ and enforcing the transport map $G=\nabla\phi$ to be strictly convex. It combines universal approximation results for second-order differentiable neural networks with a convexity-penalty mechanism $\mathcal{L}_{\kappa,\gamma}(\phi,D)$ to ensure a valid Brenier map while training a GAN discriminator. The authors derive a comprehensive error decomposition—generator error $\Delta_G$, discriminator error $\Delta_D$, sample error $\Delta_S(n)$, and training error $\Delta_T(n)$—and prove consistency: as the sample size grows, the learned distribution converges to the target and the potential converges to the Brenier potential, up to an additive constant. Empirically, Brenier GAN learns convex potentials and generates data across synthetic 2D, MNIST, Fashion-MNIST, and NORB, with the convexity penalty becoming inactive in training and improving diversity and stability in several setups.
Abstract
Brenier proved that under certain conditions on a source and a target probability measure there exists a strictly convex function such that its gradient is a transport map from the source to the target distribution. This function is called the Brenier potential. Furthermore, detailed information on the Hölder regularity of the Brenier potential is available. In this work we develop the statistical learning theory of generative adversarial neural networks that learn the Brenier potential. As by the transformation of densities formula, the density of the generated measure depends on the second derivative of the Brenier potential, we develop the universal approximation theory of ReCU networks with cubic activation $\mathtt{ReCU}(x)=\max\{0,x\}^3$ that combines the favorable approximation properties of Hölder functions with a Lipschitz continuous density. In order to assure the convexity of such general networks, we introduce an adversarial training procedure for a potential function represented by the ReCU networks that combines the classical discriminator cross entropy loss with a penalty term that enforces (strict) convexity. We give a detailed decomposition of learning errors and show that for a suitable high penalty parameter all networks chosen in the adversarial min-max optimization problem are strictly convex. This is further exploited to prove the consistency of the learning procedure for (slowly) expanding network capacity. We also implement the described learning algorithm and apply it to a number of standard test cases from Gaussian mixture to image data as target distributions. As predicted in theory, we observe that the convexity loss becomes inactive during the training process and the potentials represented by the neural networks have learned convexity.