Moreau-Yosida $f$-divergences

TL;DR

This paper introduces a novel framework for approximating and representing $f$-divergences using Moreau-Yosida regularization, enabling flexible variational formulas and practical algorithms for GAN training.

Contribution

It defines the Moreau-Yosida $f$-divergences, generalizes variational formulas, relaxes Lipschitz constraints, and proposes an algorithm for computing the convex conjugate suitable for automatic differentiation.

Findings

Proposes the Moreau-Yosida $f$-GAN with competitive results on CIFAR-10.

Provides a practical algorithm compatible with automatic differentiation.

Offers theoretical insights into the variational representation and Lipschitz function spaces.

Abstract

Variational representations of $f$-divergences are central to many machine learning algorithms, with Lipschitz constrained variants recently gaining attention. Inspired by this, we define the Moreau-Yosida approximation of $f$-divergences with respect to the Wasserstein-$1$ metric. The corresponding variational formulas provide a generalization of a number of recent results, novel special cases of interest and a relaxation of the hard Lipschitz constraint. Additionally, we prove that the so-called tight variational representation of $f$-divergences can be to be taken over the quotient space of Lipschitz functions, and give a characterization of functions achieving the supremum in the variational representation. On the practical side, we propose an algorithm to calculate the tight convex conjugate of $f$-divergences compatible with automatic differentiation frameworks. As an application of our results, we propose the Moreau-Yosida $f$-GAN, providing an implementation of the variational formulas for the Kullback-Leibler, reverse Kullback-Leibler, $\chi^2$, reverse $\chi^2$, squared Hellinger, Jensen-Shannon, Jeffreys, triangular discrimination and total variation divergences as GANs trained on CIFAR-10, leading to competitive results and a simple solution to the problem of uniqueness of the optimal critic.