Implicit geometric regularization in flow matching via density weighted Stein operators
By: Shinto Eguchi
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 {$γ$-Flow Matching ($γ$-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 $γ$-FM minimizes the transport cost on a statistical manifold endowed with the $γ$-Stein metric. Spectral analysis further suggests that this geometry induces an implicit Sobolev regularization, effectively damping high-frequency oscillations in void regions. Empirically, $γ$-FM significantly improves vector field smoothness and sampling efficiency on high-dimensional latent datasets, while demonstrating intrinsic robustness to outliers.
Similar Papers
On The Hidden Biases of Flow Matching Samplers
Machine Learning (Stat)
Makes AI learn better by fixing its "energy" mistakes.
Graph Flow Matching: Enhancing Image Generation with Neighbor-Aware Flow Fields
Machine Learning (CS)
Makes AI create more realistic pictures.
Generative Modeling with Continuous Flows: Sample Complexity of Flow Matching
Machine Learning (CS)
Makes AI create better pictures with less data.