Wasserstein Gradient Flows for Moreau Envelopes of f-Divergences in Reproducing Kernel Hilbert Spaces

TL;DR

This paper introduces a novel approach to regularize f-divergences using MMD in RKHS, enabling analysis of Wasserstein gradient flows and providing numerical insights into their behavior.

Contribution

It reformulates MMD-regularized f-divergences as Moreau envelopes in RKHS, facilitating gradient flow analysis and extending to variational formulations with numerical comparisons.

Findings

MMD-regularized f-divergences can be expressed as Moreau envelopes in RKHS.

Gradient flows of these divergences can be analyzed using Hilbert space techniques.

Numerical examples demonstrate the behavior of flows from empirical measures.

Abstract

Commonly used $f$-divergences of measures, e.g., the Kullback-Leibler divergence, are subject to limitations regarding the support of the involved measures. A remedy is regularizing the $f$-divergence by a squared maximum mean discrepancy (MMD) associated with a characteristic kernel $K$. We use the kernel mean embedding to show that this regularization can be rewritten as the Moreau envelope of some function on the associated reproducing kernel Hilbert space. Then, we exploit well-known results on Moreau envelopes in Hilbert spaces to analyze the MMD-regularized $f$-divergences, particularly their gradients. Subsequently, we use our findings to analyze Wasserstein gradient flows of MMD-regularized $f$-divergences. We provide proof-of-the-concept numerical examples for flows starting from empirical measures. Here, we cover $f$-divergences with infinite and finite recession constants. Lastly, we extend our results to the tight variational formulation of $f$-divergences and numerically compare the resulting flows.