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

2024-02-07

arXiv.org

AI Summary

Conventional $f$-divergences suffer from ill-posedness when the supports of compared probability measures do not align. Method: This paper proposes an MMD-regularized $f$-divergence reconstruction, which—uniquely—admits a closed-form characterization as the Moreau envelope of the $f$-divergence over a reproducing kernel Hilbert space (RKHS). Contribution/Results: From this formulation, we derive an explicit Wasserstein gradient expression and rigorously establish existence and well-posedness of the associated gradient flow under both infinite and finite decay constants. We further develop a numerically stable implementation framework based on kernel mean embeddings, enabling robust optimization starting directly from empirical measures. Experiments demonstrate geometric fidelity and dynamical stability of the resulting gradient flow for canonical $f$-divergences—including KL and Hellinger divergences—substantially enhancing robustness and computational tractability in distributional discrepancy modeling.

Technology Category

• Machine Learning: Learning with Manifolds
• Search and Optimization: Non-convex Optimization
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

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.

Problem

• Reproducing Kernel Hilbert Space
• f-divergence
• Wasserstein Gradient Flow

Innovation

• MMD-adjusted f-divergence
• Wasserstein gradient flow
• Kernel mean embedding