PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiarui Cao, Zixuan Wei, Yuxin Liu

arXiv:2603.10592v1cs.LGcs.AI

TL;DR

The paper addresses limited theoretical foundations for Drifting Models and develops Gradient Flow Drifting as a KDE-based Wasserstein gradient-flow framework. It proves the Drifting Model's equivalence to a forward-KL flow, unifies related generators, and proposes mixed divergences validated by preliminary synthetic experiments. The approach is limited by minibatch KDE variance, especially in higher dimensions.

  • Problem

    Drifting Models have strong empirical results but limited theoretical analysis, including a heuristic treatment and an identifiability proof requiring additional smoothness assumptions.

  • Method

    The paper formulates KDE-level Wasserstein gradient flows for divergences, identifies Drifting Models and MMD generators as special cases, and combines divergences such as reverse KL and χ2.

  • Results

    The paper proves that the Drifting Model is equivalent to the Wasserstein gradient flow of forward KL under KDE approximation, with preliminary synthetic experiments validating the framework.

  • Takeaways & Limitations

    Mixed reverse-KL and χ2 flows combine precision-forcing and coverage-forcing behavior, addressing mode blurring and mode collapse in the reported synthetic results.

  • Takeaways & Limitations

    Minibatch approximation of KDE becomes increasingly variable with dimension, potentially affecting training stability and final convergence.

Abstract

from arXiv · show

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 $χ^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.

1 Introduction

The paper addresses underdeveloped theoretical foundations of Drifting Models by identifying their drifting field with a KDE-approximated Wasserstein gradient flow. This connection yields a unified divergence-based framework, mixed flows, and simpler identifiability.

  • Drifting Models evolve generated distributions during training rather than relying on iterative inference-time dynamics, enabling one-step generation.
  • The drifting field with a Gaussian kernel equals the difference between KDE log-density gradients, up to a bandwidth-squared scaling factor.
  • This field is the particle velocity of the Wasserstein-2 gradient flow of KL(q∥p) with KDE-approximated densities.
  • Varying the divergence produces a family of gradient-flow drifting models, including MMD-based generators as L2-distance cases and Drifting Models as KL cases.
  • Mixed gradient flows combine complementary behaviors, including global mode coverage with MMD and local sharpness with reverse KL.
  • KDE injectivity makes the equilibrium condition Vp,q = 0 imply p = q, simplifying identifiability.

2 Related Work

Prior work develops Drifting Models, Wasserstein gradient-flow generative methods, KDE and mean-shift connections, MMD generators, and f-divergence estimation. This paper positions Gradient Flow Drifting as a broader framework connecting these lines.

  • Drifting Models learn one-step pushforward maps by evolving generated distributions through kernel-based drifting fields, but earlier analyses provide limited theoretical foundations.
  • Wasserstein gradient flows provide a variational framework used for MMD sampling, divergence GANs, and scalable Wasserstein gradient descent.
  • Earlier work connects mean shift, KDE gradients, MMD-optimal quantization, and gradient flows, while this paper extends the connection to arbitrary f-divergences.
  • MMD-based generative models minimize discrepancies between generated and data distributions, and this framework treats them as members of a broader family.
  • Unlike discriminator-based f-divergence methods, the KDE-based approach avoids adversarial training while connecting f-divergence minimization to Wasserstein gradient flows.

3 Preliminaries

The preliminaries define probability-measure notation, KDE operators and kernels, kernel embeddings, characteristic kernels, Wasserstein-2 gradient flows, and the Drifting Model's field and training loss.

  • P(Rd) denotes Borel probability measures, while P2(Rd) denotes those with finite second moments.
  • The KDE operator smooths a probability measure with a kernel; for a Gaussian kernel, μkde(x) is the resulting KDE with bandwidth h.
  • A positive definite kernel induces an RKHS and a kernel mean embedding characterized by the reproducing property.
  • A characteristic kernel has an injective kernel mean embedding map on probability measures.
  • A Wasserstein-2 gradient flow is a curve of probability measures satisfying a continuity equation determined by a functional's first variation.
  • The Drifting Model defines a kernel-based field for a generated pushforward distribution and trains the generator with a stop-gradient squared-distance loss.

4 Method: Gradient Flow Drifting

The framework recasts generative models as Wasserstein gradient flows of divergences applied to KDE-smoothed distributions. It unifies Drifting Models, MMD generators, mixed divergences, identifiability, and manifold extensions within this KDE-based formulation.

  • Foundation: KDE Smoothing and Distribution Matching: KDE smoothing preserves distributional identity while providing the regularity needed for gradient-flow analysis without smoothness assumptions on p or q.Under kernel regularity, KDE densities are positive and C1, enabling well-defined log ratios and standard Wasserstein machinery.
  • Generalized Gradient Flows: The generalized framework derives Wasserstein gradient-flow velocities for f-divergences, with density-ratio weights controlling local emphasis across regions.Forward KL uses unit weight, reverse KL emphasizes high-data-density regions, and χ2 emphasizes high-generated-density regions.
  • Unified Identifiability: The KDE-level f-divergence is identifiable and non-increasing along the flow, with q = p as its unique global optimum and equilibrium.The unified identifiability theorem relies on kernel injectivity and strict convexity of the divergence.
  • The Drifting Model as Forward KL Flow: For a Gaussian kernel, the Drifting Model equals the forward-KL Wasserstein gradient-flow velocity up to a bandwidth-squared time rescaling.The equivalence identifies the drifting field with the difference between KDE log-density gradients and imports the framework’s convergence and identifiability results.
  • MMD and Mixed Gradient Flows: MMD generators arise from an L2 density-distance gradient flow, while mixed divergences produce convex-combination velocities with dissipating mixed energy.The framework therefore includes both f-divergence and MMD-based generators as sample-computable KDE flows.
  • Mixed Flows and Manifold Extension: Combining reverse KL and χ2 flows is proposed to reconcile precision-oriented and coverage-oriented behavior, and the framework extends to Riemannian manifolds.On compact boundaryless manifolds, energy dissipation holds without boundary conditions; spherical kernels also permit alternative weighting profiles.

5 Experiments

Synthetic 2D experiments compare gradient-flow drifting under different divergences and kernels, highlighting contrasting mode coverage, blurring, precision, and stability behaviors.

  • The experiments visualize particle evolution under gradient-flow velocity fields using different divergence and kernel implementations.
  • The original drifting model and L2 flow drifting cover modes but exhibit blurred samples.The comparison uses the same 2D-toy benchmark and gradient-flow visualization.
  • Reverse KL plus χ2 mixture flow drifting produces precise samples, quickly explores all modes, and avoids severe mode collapse.
  • The original Laplace-kernel drifting model can show particle jittering and distorted generation near the Swiss-roll center, unlike its RBF-kernel version.The paper attributes this to the Laplace kernel violating regularity condition K4.

6 Conclusion and Discussion

The paper concludes that Gradient Flow Drifting unifies several generative models through KDE-based Wasserstein gradient flows, while identifying practical stability and scalability boundaries.

  • Gradient Flow Drifting defines a family of generative models and identifies Drifting Models as KDE-approximated Wasserstein gradient flows of forward KL.
  • Under a suitably designed kernel, matching KDE-smoothed pushforward distributions can approximate matching the original distributions.
  • The framework extends to Riemannian manifolds, loosening kernel constraints and supporting semantic-space applications.
  • Minibatch KDE approximations become increasingly variable in higher dimensions, potentially harming training stability and final convergence.
  • Future work targets large-scale and high-dimensional datasets, conditional and multimodal generation, ablations, kernel and bandwidth studies, and acceleration.

C Regularity of KDE-Smoothed Densities

The KDE regularity results show that mild kernel conditions supply positivity, differentiability, and bounded gradients without requiring regularity of the underlying probability measure.

  • KDE smoothing inherits C1 regularity from the kernel, and differentiation commutes with integration.
  • KDE densities are strictly positive everywhere and have uniformly bounded gradients under the stated kernel conditions.
  • The regularity proof requires no smoothness or moment conditions on the underlying probability measure because a uniform gradient bound dominates the integrand.
  • For any probability distributions p and q, their KDE densities are positive and C1, making log-ratios and f-divergence analysis well-defined.

D Wasserstein Gradient Flows of f-Divergences

This section formulates Wasserstein gradient flows for KDE-smoothed f-divergences, establishes boundary and identifiability conditions, and gives computable score-based velocities.

  • The framework studies Wasserstein-2 gradient flows of f-divergences between smooth positive KDE densities.
  • An f-divergence is generated by a convex function f with f(1)=0 and applied to positive densities.
  • The WGF particle velocity is obtained from the gradient of the first variation, subject to boundary terms vanishing for the energy-dissipation argument.
  • For Euclidean KDE flows, suitable tail behavior supports the boundary condition; on compact boundaryless manifolds, it is vacuous and energy dissipation holds unconditionally.
  • If the manifold is connected and f is strictly convex, an identically vanishing WGF velocity implies equality of the two positive densities.
  • The framework lists distinct velocities for forward KL, reverse KL, and χ2 divergences, while KDE scores remain computable as expectations under the data measure.

F.1 Proof of Identifiability (Theorem 4.7)

The proof shows that vanishing f-divergence velocity forces KDE densities to match, and KDE injectivity then yields equality of the underlying distributions.

  • Step 1: Constant density ratio: The proof first shows that a zero velocity field makes f′(q_kde/p_kde) constant on the connected domain.Strict positivity and C1 regularity make the density ratio and its composition with f′ well-defined.
  • Step 1: Constant density ratio: Strict convexity makes f′ injective, so the KDE density ratio is a positive constant and q_kde = λ p_kde everywhere.The constant-ratio conclusion follows from the strict monotonicity of f′.
  • Step 2: KDE matching implies distribution matching: Integrating the constant-ratio relation gives λ = 1, hence q_kde = p_kde.The proof uses normalization of the KDE densities, including the compact-manifold case.
  • Step 2: KDE matching implies distribution matching: Characteristic-kernel KDE injectivity then converts q_kde = p_kde into q = p.Thus KDE-level identifiability transfers to the original distributions.
  • MMD identifiability: The same identifiability route applies to MMD: zero MMD velocity implies equal KDEs, which imply equal underlying distributions.Connectedness first turns a constant KDE difference into zero before KDE injectivity is applied.

G Core Equivalence: Drifting as KL Gradient Flow

This section establishes the drifting–KL equivalence and places it within a KDE-based gradient-flow framework that covers multiple divergences and extends to spherical geometry.

  • Core equivalence: For Gaussian KDEs, the drifting field equals a bandwidth-squared scaling of the Wasserstein KL velocity.The KL velocity is vKL = ∇log π − ∇log ρ with ρ = q_kde and π = p_kde.
  • Core equivalence: The equivalence follows by applying the KDE score identity to p and q and subtracting, which cancels the x terms.The resulting difference is identified with vKL for the KDE-smoothed densities.
  • MMD flow: MMD has a corresponding Wasserstein gradient-flow velocity derived from its kernel mean-embedding functional.The first variation is expressed through the KDE difference q_kde − p_kde, whose gradient determines the velocity.
  • Mixed flows: Convex combinations of divergence functionals produce mixed gradient flows whose velocity is the same combination of component velocities.This superposition principle supports combining divergence-specific flow fields.
  • Unified framework: The framework expresses gradient-flow velocities through p_kde and q_kde, making them sample-computable via the score formula.Table 2 summarizes the unified KDE-based velocity constructions.
  • Euclidean kernels: For Euclidean kernels, the Matérn family satisfies all four assumptions exactly when ν > 1.When ν ≤ 1, the gradient diverges at zero and the required C1 regularity fails.
  • Riemannian extension: The framework extends to Riemannian manifolds by replacing Euclidean gradients with Riemannian gradients and imposing manifold kernel analogues.On the sphere, adapted assumptions support kernels such as the von Mises–Fisher kernel.

J.2.2 Spherical Logarithmic Kernel

The spherical logarithmic kernel satisfies the adapted kernel assumptions under a parameter restriction and offers inverse-distance-squared weighting for broader mode coverage.

  • Kernel validity: For c > 0 and 0 < α < 1/(2+c), the spherical logarithmic kernel satisfies all assumptions K1S–K4S.The restriction keeps the logarithm’s argument strictly between zero and one, ensuring positivity and smoothness.
  • Kernel validity: The kernel’s Riemannian gradient is uniformly bounded by 1/c.The bound follows from the spherical geometry and the denominator’s lower bound c.
  • Characteristicness: Its power-series coefficients are positive, so Schoenberg’s theorem gives strict positive definiteness and characteristicness on spheres of dimension d ≥ 2.The positivity of all coefficients establishes universality as well.
  • Spherical score: The spherical KDE score uses pairwise weights W_c(x, y) = 1/(1 − x⊤y + c).The score is obtained by projecting the ambient gradient onto the tangent space.
  • Weighting profile: Unlike the vMF kernel, the logarithmic kernel has polynomial inverse-distance-squared weighting, typically producing heavier tails and better global mode coverage.This contrasts its weighting profile with the exponential-type spherical kernel.
Loading 2603.10592v1…