Sobolev Regularized MMD Gradient Flow

2026-05-12

AI Summary

This work addresses the convergence challenges of Maximum Mean Discrepancy (MMD) gradient flows arising from non-convexity by proposing a Sobolev-regularized MMD gradient flow method with gradient penalty on the witness function. The approach eliminates the need for isoperimetric assumptions on the target distribution and, for the first time, guarantees global convergence simultaneously in both continuous and discrete time. It provides a unified framework applicable to sampling and generative modeling tasks involving unnormalized densities. By integrating Sobolev regularization, kernel mean embeddings, and Stein kernel techniques, the proposed method demonstrates superior convergence properties and generalization performance over existing approaches across a range of experiments.

Abstract

We propose Sobolev-regularized Maximum Mean Discrepancy (SrMMD) gradient flow, a regularized variant of maximum mean discrepancy (MMD) gradient flow based on a gradient penalty on the witness function. The proposed regularization mitigates the non-convexity of the MMD objective and yields provable \emph{global} convergence guarantees in MMD in both continuous and discrete time. A more surprising appeal is that our convergence analysis does not rely on isoperimetric assumptions on the target distribution. Instead, it is based on a regularity condition on the difference between kernel mean embeddings. A key highlight of the proposed flow is that it is applicable in both sampling (from an unnormalized target distribution) -- using Stein kernels -- and generative modeling settings, unlike previous works, where a gradient flow is suitable for only generative modeling or sampling but not both. The effectiveness of the proposed flow is empirically verified on a broad range of tasks in both generative modelling and sampling.

Problem

• Maximum Mean Discrepancy
• gradient flow
• non-convexity
• global convergence
• sampling and generative modeling

Innovation

• Sobolev regularization
• Maximum Mean Discrepancy
• gradient flow
• global convergence
• Stein kernel