Implicit geometric regularization in flow matching via density weighted Stein operators

TL;DR

This paper introduces $oldsymbol{ ext{ extgamma}}$-Flow Matching, a density-weighted variant of flow matching that improves efficiency and smoothness by aligning regression with the data distribution using a dynamic density estimation strategy.

Contribution

It proposes a novel density-weighted flow matching method with a dynamic density estimation, theoretically grounding it in a $ ext{ extgamma}$-Stein metric and demonstrating improved empirical performance.

Findings

Enhanced vector field smoothness and sampling efficiency.

Robustness to outliers in high-dimensional datasets.

Theoretical connection to transport cost minimization on a statistical manifold.

Abstract

Flow Matching (FM) has emerged as a powerful paradigm for continuous normalizing flows, yet standard FM implicitly performs an unweighted $L^2$ regression over the entire ambient space. In high dimensions, this leads to a fundamental inefficiency: the vast majority of the integration domain consists of low-density ``void'' regions where the target velocity fields are often chaotic or ill-defined. In this paper, we propose {$\gamma$-Flow Matching ($\gamma$-FM)}, a density-weighted variant that aligns the regression geometry with the underlying probability flow. While density weighting is desirable, naive implementations would require evaluating the intractable target density. We circumvent this by introducing a Dynamic Density-Weighting strategy that estimates the \emph{target} density directly from training particles. This approach allows us to dynamically downweight the regression loss in void regions without compromising the simulation-free nature of FM. Theoretically, we establish that $\gamma$-FM minimizes the transport cost on a statistical manifold endowed with the $\gamma$-Stein metric. Spectral analysis further suggests that this geometry induces an implicit Sobolev regularization, effectively damping high-frequency oscillations in void regions. Empirically, $\gamma$-FM significantly improves vector field smoothness and sampling efficiency on high-dimensional latent datasets, while demonstrating intrinsic robustness to outliers.