Kernelized Wasserstein Gradient Flows

• The kernelized Wasserstein gradient is a framework that integrates optimal transport with RKHS techniques to enable explicit, scalable particle-based approximations.
• It formulates steepest-descent flows over probability distributions using a variational approach, yielding provable convergence and stability in high-dimensional settings.
• Applications span inverse problems, variational inference, and generative modeling, where kernel regularization effectively manages computational complexity.

The kernelized Wasserstein gradient provides a unifying framework for regularized, computationally tractable, and geometrically principled steepest-descent flows on spaces of probability distributions. Central to recent advances in optimal transport, variational inference, density evolution, and inverse problems, the kernelized Wasserstein gradient combines the Riemannian geometry of Wasserstein space with functional analytic techniques from reproducing kernel Hilbert spaces (RKHS) and maximum mean discrepancy (MMD). This construction enables explicit, scalable representations of Wasserstein gradient flows, efficient particle-based approximations, and provable convergence guarantees in high-dimensional statistical learning.

1. Foundations: Wasserstein Gradient Flow and Kernelization

The 2-Wasserstein gradient flow is a canonical evolution equation on the space $\mathcal{P}_2(\mathbb{R}^d)$ of probability measures with finite second moment. Given a differentiable functional $F:\mathcal{P}_2(\mathbb{R}^d)\to\mathbb{R}$, the gradient flow is the continuity equation

$$\partial_t \rho_t(x) = \nabla\cdot\left(\rho_t(x)\nabla\frac{\delta F}{\delta \rho}(x)\right)$$

where $\frac{\delta F}{\delta \rho}$ is the first variation of $F$, and the velocity field is $v_t(x) = -\nabla\frac{\delta F}{\delta\rho}(x)$. For the Kullback-Leibler functional, this recovers the Fokker–Planck/Langevin dynamics; for general divergences or energy functionals, this framework underpins nonlinear, possibly nonlocal PDEs (Chen et al., 2024, Li et al., 2023, Duong et al., 2024).

Kernelization is introduced by restricting the admissible velocity fields to an RKHS $H^d$ induced by a scalar positive definite kernel $K(\cdot,\cdot)$. The steepest-descent direction is thus found by maximizing a variational objective regularized by the RKHS norm, giving rise to explicit, data-driven vector fields and rapid computation via the representer theorem (Stein et al., 2024, Arbel et al., 2019, Hu et al., 10 Nov 2025).

2. Kernelized Wasserstein Gradient: Variational Derivation and RKHS Form

Consider the optimization of a convex functional $F(\rho)$ or divergence $D(\rho\mid\nu)$. The kernelized Wasserstein gradient restricts $v\in H^d$ and solves the variational problem

$$u^* = \arg\max_{v \in H^d} \left\{ \mathbb{E}_{\rho} \big[ v(x)\cdot \nabla_x\frac{\delta F}{\delta\rho}(x) \big] - \lambda\,\mathbb{E}_\rho[\nabla_x\cdot v(x)] - \frac{1}{2}\|v\|_{H^d}^2 \right\}$$

where $\lambda$ is a regularization parameter (e.g. negative entropy penalty in imputation), and $F$ may be a KL, $f$-divergence, or negative-entropy-regularized objective (Chen et al., 2024, Arbel et al., 2019).

Closed-form solutions (via the representer theorem) are

$$u^*(x) = \mathbb{E}_{y \sim \rho} \big[ -\lambda\,\nabla_y K(x, y) + \nabla_y \frac{\delta F}{\delta\rho}(y)\, K(x, y) \big],$$

providing an explicit, kernel-based gradient vector that is computable from empirical samples, neural network estimators, or density models. For KL divergence, $$\nabla_y \frac{\delta F}{\delta\rho}(y) = \nabla_y \log p(y)$$, so this structure generalizes SVGD and other functional-gradient methods (Cheng et al., 2023, Chewi et al., 2020).

3. Key Frameworks and Algorithmic Instantiations

a) Diffusion Model Imputation (KnewImp)

Kernelized negative-entropy-regularized Wasserstein gradient flow is employed for numerical tabular data imputation, addressing sample diversification and training complexity. Iterative updates are performed in the RKHS-restricted velocity field; only missing entries are updated at each step via kernelized vector fields computed from the current empirical distribution (Chen et al., 2024):

• Initialization: Impute missing entries (e.g., mean imputation).
• For $T$ iterations: Train a score network to estimate $\nabla_x\log p(x)$; compute $$u_i = -\lambda \mathbb{E}_j[\nabla_y K(x_i, y_j)] + \mathbb{E}_j[\text{score}(y_j) K(x_i, y_j)]$$; update $x_i \leftarrow x_i + \eta u_i$.

This procedure yields monotonic improvement in the cost functional and empirical advantages over existing imputation strategies, with convergence guarantees derived from monotonicity in the 2-Wasserstein metric (Chen et al., 2024).

b) SVGD, Regularized SVGD, and LAWGD

Kernelized Wasserstein gradients are central to SVGD and its variants. SVGD performs gradient flows in density space for the KL or $\chi^2$ divergence, replacing the true Wasserstein gradient with its RKHS-regularized version. The regularized SVGD (R-SVGD) algorithm applies a resolvent-type preconditioner to partially "de-bias" the kernel integral operator and quantifies the trade-off between regularization bias and finite-particle error:

$$\phi_{\nu, \rho} = -\Big[(1 - \nu)\, \mathcal{T}_{k, \rho} + \nu I\Big]^{-1}\mathcal{T}_{k, \rho} [\nabla\log(\rho/\pi)]$$

with explicit rates for convergence in Fisher information and Wasserstein-1 distance, and stepwise updates detailed for continuous and discrete time (He et al., 5 Feb 2026).

LAWGD leverages kernels derived from the inverse Langevin generator to ensure scale-invariant exponential contraction in KL, achieving stronger uniform ergodicity than basic SVGD, but requiring problem-specific kernels via spectral decomposition (Chewi et al., 2020).

4. Connections to MMD, f-Divergences, and Moreau Envelopes

Several works interpret kernelized Wasserstein gradients as steepest-descent flows for regularized $f$-divergence functionals, notably by embedding distributions into the RKHS of a characteristic kernel. When the objective is an MMD-regularized $f$-divergence,

$$D_{f,\nu}^{\lambda}(\mu) = \inf_{\sigma} \left\{ D_f(\sigma||\nu) + \frac{1}{2\lambda} \| m_{\mu} - m_{\sigma} \|_{H_K}^2 \right\},$$

the gradient at $\mu$ is $\nabla_x\hat{p}_\mu(x)$, where $\hat{p}_\mu$ solves the dual Moreau-Yosida variational problem in $H_K$, and discrete particle flows are generated via ODE integration (Stein et al., 2024, Duong et al., 2024). This framework extends to generalized MMD flows, sliced Wasserstein MMDs, and non-convex functionals, with regularization via $H^1$-Sobolev terms as necessary to guarantee gradient flow existence and mass conservation (Duong et al., 2024, Bonet et al., 9 Jun 2025).

5. Applications: Inverse Problems, Generative Modeling, and Statistical Learning

Kernelized Wasserstein gradients underpin modern algorithms for:

• Inverse identification in PDEs and density flows: RKHS-based representer theorems provide closed-form estimators for unknown potentials and interaction kernels in Wasserstein flows, with convergence and stability rates as discretization is refined (Hu et al., 10 Nov 2025).
• Domain adaptation and dataset distillation: "Wasserstein over Wasserstein" (WoW) flows with sliced Wasserstein MMD kernels enable dynamics over probability distributions of distributions for structured dataset flows (Bonet et al., 9 Jun 2025).
• Variational inference and Bayesian sampling: SVGD, GWG, and Ada-GWG (adaptive selection of regularizer) utilize kernelized Wasserstein gradients for efficient, scalable approximation of posterior distributions in complex latent variable models (Cheng et al., 2023, Arbel et al., 2019).

Empirical success in high-dimensional Bayesian neural networks, generative modeling, and missing data imputation stems from the ability to tune computational cost via kernel choice, sample size, and representer-system dimension, while retaining strong control on statistical and approximation errors.

6. Theoretical Properties and Convergence Analysis

The kernelized Wasserstein gradient inherits several structural properties from the underlying geometry and regularization:

• Monotonicity: The functional evolves non-increasingly (or non-decreasingly, depending on sign conventions) along the flow.
• Convergence: Under mild convexity, smoothness of the kernel, and regularizing envelope, flows enjoy existence, uniqueness, and exponential convergence to minimizers in both functional and Wasserstein metrics (Stein et al., 2024, Chewi et al., 2020).
• Stability and Error Bounds: Quantitative estimates relate regularization bias, discretization, and finite-sample error, with non-asymptotic decay rates in statistical functionals such as Fisher information and $W_1$ (He et al., 5 Feb 2026, Hu et al., 10 Nov 2025).
• Algorithmic complexity: Kernel methods can be rendered tractable in high dimensions through low-rank or Nyström approximations, stochastic mini-batching, and efficient differential operators.

7. Extensions and Future Directions

Kernelized Wasserstein gradients are actively generalized via:

• More general cost and divergence functionals: Generalized Wasserstein metrics, adaptive regularizers, and tight RKHS variational forms (Cheng et al., 2023, Stein et al., 2024).
• Multi-level and distribution-of-distributions flows: Wasserstein over Wasserstein geometries and high-level MMDs (Bonet et al., 9 Jun 2025).
• Hamiltonian flows and symplectic structures: Extension to second-order Wasserstein and quantum dynamics (Hu et al., 10 Nov 2025).
• Efficient representation and learning: Automatic kernel learning and optimal regularizer selection for domain-specific adaptation (Arbel et al., 2019, Stein et al., 2024).

The kernelized Wasserstein gradient thus constitutes a versatile, theoretically grounded, and computationally effective framework with broad implications in applied mathematics, machine learning, and computational physics.