Gradient Flow Drifting: Generative Modeling via Wasserstein Gradient Flows of KDE-Approximated Divergences

TL;DR

This paper introduces Gradient Flow Drifting, a new framework for generative models that unifies and extends existing approaches using Wasserstein gradient flows and KDE approximations, with theoretical guarantees and preliminary experiments.

Contribution

It establishes a mathematical equivalence between Drifting Models and Wasserstein gradient flows of KL divergence under KDE, and extends the framework to include MMD-based generators and Riemannian manifolds.

Findings

Proves the equivalence between Drifting Model and Wasserstein gradient flow of KL divergence.

Includes MMD-based generators as special cases within the framework.

Demonstrates preliminary validation on synthetic benchmarks.

Abstract

We reveal a precise mathematical framework about a new family of generative models which we call Gradient Flow Drifting. With this framework, we prove an equivalence between the recently proposed Drifting Model and the Wasserstein gradient flow of the forward KL divergence under kernel density estimation (KDE) approximation. Specifically, we prove that the drifting field of drifting model (arXiv:2602.04770) equals, up to a bandwidth-squared scaling factor, the difference of KDE log-density gradients $\nabla \log p_{\mathrm{kde}} - \nabla \log q_{\mathrm{kde}}$, which is exactly the particle velocity field of the Wasserstein-2 gradient flow of $KL(q\|p)$ with KDE-approximated densities. Besides that, this broad family of generative models can also include MMD-based generators, which arises as special cases of Wasserstein gradient flows of different divergences under KDE approximation. We provide a concise identifiability proof, and a theoretically grounded mixed-divergence strategy. We combine reverse KL and $\chi^2$ divergence gradient flows to simultaneously avoid mode collapse and mode blurring, and extend this method onto Riemannian manifold which loosens the constraints on the kernel function, and makes this method more suitable for the semantic space. Preliminary experiments on synthetic benchmarks validate the framework.