Technical Summary: Sobolev-regularized Maximum Mean Discrepancy (SrMMD) Gradient Flow

Problem Statement

Sampling from unnormalized distributions and generative modeling are often treated as distinct tasks, yet both can be framed as optimization problems over the space of probability measures $\mathcal{P}_2(\mathbb{R}^d)$. Existing approaches face significant theoretical and practical limitations:

• MMD Gradient Flow: While effective for generative modeling, standard Maximum Mean Discrepancy (MMD) flows suffer from non-convexity, leading to slow convergence rates that depend heavily on the isoperimetric properties (e.g., log-Sobolev or Poincaré constants) of the target distribution. This makes them inefficient for multi-modal targets where these constants scale exponentially with mode separation.
• Sampling Methods: Methods like the Unadjusted Langevin Algorithm (ULA) and Kernel Stein Discrepancy (KSD) flows also rely on isoperimetric assumptions for convergence guarantees.
• Regularization in GANs: In Generative Adversarial Networks (GANs), gradient regularization of the witness/critic function is empirically known to improve image quality and stability. However, a theoretical justification for why gradient regularization accelerates convergence in flow-based generative models or sampling has been lacking.

The paper addresses the need for a unified framework that applies to both sampling and generative modeling, offering provable global convergence guarantees that are independent of the target distribution's isoperimetric constants.

Methodology

The authors propose Sobolev-regularized Maximum Mean Discrepancy (SrMMD) gradient flow, a novel variant of MMD/KSD flow that incorporates a gradient penalty on the witness function.

Core Definition

The SrMMD metric is defined as a supremum over a witness function $f$ in a Reproducing Kernel Hilbert Space (RKHS) $\mathcal{H}$, subject to a constraint involving both the $L^2$ norm of the function and the $L^2$ norm of its gradient:
$$\text{SrMMD}(\mu, \pi) = \sup_{f \in \mathcal{H}: \|\nabla f\|_{L^2_d(\mu)}^2 + \lambda \|f\|_{\mathcal{H}}^2 \leq 1} \int f d(\mu - \pi)$$
where $\lambda > 0$ is a regularization parameter.

The Flow Dynamics

The SrMMD flow evolves a probability measure $\mu_t$ according to the continuity equation:
$$\partial_t \mu_t - \nabla \cdot (\mu_t \nabla f_{\mu_t, \pi}) = 0$$
where the witness function $f_{\mu, \pi}$ is the maximizer of the SrMMD definition. Unlike standard MMD flow where the velocity field is $\nabla(m_\mu - m_\pi)$, the SrMMD velocity field is:
$$\nabla f_{\mu, \pi} = \nabla (S_\mu + \lambda I)^{-1} (m_\mu - m_\pi)$$
Here, $S_\mu = D_\mu^* D_\mu$ is the gradient covariance operator, mapping functions to their gradients in $L^2(\mu)$. This formulation effectively applies Tikhonov regularization to the inverse problem of finding the witness function.

Implementation

The paper derives a closed-form expression for the witness function in the finite-particle setting. For an empirical distribution $\hat{\mu} = \frac{1}{N}\sum \delta_{x_i}$, the witness function admits a tractable implementation involving matrix inversions of size $Nd \times Nd$.

• Generative Modeling: When $\pi$ is known via samples, expectations are estimated empirically.
• Sampling: When $\pi$ is known only up to a normalization constant, the method utilizes a Stein kernel $k_\pi$. In this case, the target mean embedding $m_\pi$ vanishes, simplifying the witness function to $f_{\mu, \pi} = (S_\mu + \lambda I)^{-1} m_\mu$.

Key Contributions

1. Unified Framework: The SrMMD flow is applicable to both generative modeling (using MMD) and sampling from unnormalized targets (using KSD/Stein kernels), bridging a gap between previously distinct algorithmic developments.
2. Theoretical Convergence Guarantees:
   • Continuous Time: The flow exhibits exponential decay in MMD/KSD distance, $MMD^2(\mu_T, \pi) \leq e^{-2T} MMD^2(\mu_0, \pi) + \text{residual}$. Crucially, this rate is independent of isoperimetric constants (like log-Sobolev or Poincaré constants) of the target distribution $\pi$.
   • Discrete Time: The SrMMD descent scheme achieves an iteration complexity of $\tilde{O}(\delta^{-1/\alpha})$ to reach error $\delta$, where $\alpha$ characterizes the regularity of the trajectory. This contrasts with methods like ULA or SVGD, which often have complexities scaling with isoperimetric constants or are limited to $O(\delta^{-1})$.
3. Source Condition Analysis: The convergence analysis relies on a source condition ($m_\mu - m_\pi \in \text{Ran}(S_\mu^{\alpha/2})$) regarding the regularity of the difference between kernel mean embeddings, rather than geometric properties of the target distribution.
4. Theoretical Justification for Gradient Regularization: The paper provides the first rigorous theoretical explanation for why gradient regularization of the witness function improves convergence in MMD-based flows, addressing a longstanding open question in the literature.
5. Well-Posedness: The authors establish conditions under which the flow is well-posed (unique, absolutely continuous, non-explosive), relaxing previous assumptions on kernel growth to allow for certain unbounded kernels (e.g., polynomial kernels) while excluding singular kernels.

Experimental Results

The effectiveness of SrMMD is validated across several tasks:

• Synthetic Sampling: On multi-modal Gaussian mixtures and Swiss roll distributions, SrMMD demonstrates faster convergence and better mode coverage compared to vanilla MMD flow and performs comparably to or better than DrMMD (L2-regularized) and HrMMD (hybrid).
• Student-Teacher Networks: In a high-dimensional mean-field optimization task, SrMMD achieves faster convergence and lower final error than baseline flows.
• Image Color Transfer: SrMMD successfully aligns color distributions between images, showing superior performance in MMD and Wasserstein-2 metrics compared to unregularized MMD.
• Bayesian Logistic Regression: In sampling tasks for posterior inference on real-world datasets (Breast Cancer, Ionosphere, etc.), SrMMD significantly outperforms the original KSD flow and matches the performance of Stein Variational Gradient Descent (SVGD), despite KSD flow typically underperforming SVGD.

Significance and Claims

The paper claims that SrMMD flow represents a significant advancement by providing a unified, theoretically grounded approach for both sampling and generative modeling that avoids the "curse of isoperimetry" common in other methods.

• Modest Claims on Metrics: The authors explicitly acknowledge that their convergence guarantees are stated in terms of MMD/KSD, which are weaker metrics than KL or $\chi^2$ divergence. They note that while this allows for independence from isoperimetric constants, it represents a trade-off in the strength of the metric.
• Computational Cost: The paper admits a limitation in computational complexity, as each iteration requires inverting an $Nd \times Nd$ matrix ($O(N^3d^3)$), restricting current empirical evaluations to small-scale benchmarks.
• Future Work: The authors suggest that future work should assess the method on larger-scale tasks (e.g., earthquake modeling) where SVGD has proven effective, potentially utilizing approximations like Nyström methods or random features to mitigate computational costs.

In summary, the paper posits that Sobolev regularization offers a robust mechanism to accelerate gradient flows in probability spaces, providing a theoretical basis for empirical practices in GANs and extending their utility to sampling problems without requiring strong geometric assumptions on the target distribution.