Preconditioned Flow Matching
Abstract
Flow matching (FM) learns vector fields by regressing stochastic velocity targets along intermediate distributions . We identify a geometric optimization bottleneck in this regression problem: when the covariance of is ill-conditioned, gradient-based training rapidly fits high-variance directions while making slow progress along low-variance ones. In an exactly solvable Gaussian setting, we prove that the excess risk is weighted by , and that both gradient descent and stochastic gradient descent inherit condition-number-dependent convergence. We then extend the analysis to Gaussian mixtures, showing that multimodality does not average away this effect; instead, the slowest and worst-conditioned component can control optimization. Motivated by this analysis, we propose \emph{preconditioned flow matching}, a precondition-then-match framework that transforms the target distribution into a more isotropic representation, trains the main flow in the transformed space, and maps generated samples back through the inverse transformation. We show theoretically that preconditioning reshapes the intermediate FM path and improves its conditioning. Across controlled Gaussian and Gaussian-mixture experiments, latent MNIST and other high resolution image datasets up to resolution, preconditioning improves path-conditioning diagnostics, low-eigenvalue recovery, FID, MMD, precision, and recall. Compute-matched baselines and preconditioner-quality ablations further show that the gains are not explained merely by additional preconditioner parameters, but by improved geometry of the downstream flow matching problem.
Keywords
Cite
@article{arxiv.2603.02337,
title = {Preconditioned Flow Matching},
author = {Shadab Ahamed and Eshed Gal and Md Shahriar Rahim Siddiqui and Simon Ghyselincks and Moshe Eliasof and Eldad Haber},
journal= {arXiv preprint arXiv:2603.02337},
year = {2026}
}
Comments
34 pages, 16 figures, 5 tables