Wasserstein Gradient Flows for Moreau Envelopes of f-Divergences in Reproducing Kernel Hilbert Spaces
This work addresses theoretical and computational challenges in divergence measures for machine learning, offering incremental improvements in regularization and flow analysis.
The paper tackles limitations of f-divergences like Kullback-Leibler by regularizing them with a squared maximum mean discrepancy (MMD) using characteristic kernels, showing this can be expressed as a Moreau envelope in reproducing kernel Hilbert spaces and analyzing Wasserstein gradient flows for these regularized divergences, with proof-of-concept numerical examples provided for empirical measures.
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.