Preconditioned Score and Flow Matching
This addresses optimization stagnation in generative models for researchers and practitioners, though it appears incremental as it builds on existing flow matching frameworks.
The paper tackles the problem of optimization bias in flow matching and score-based diffusion models caused by ill-conditioned covariance in intermediate distributions, which leads to learning plateaus at suboptimal weights. The result is a proposed reversible, label-conditional preconditioning method that improves conditioning and yields better-trained models across datasets like MNIST and high-resolution examples.
Flow matching and score-based diffusion train vector fields under intermediate distributions $p_t$, whose geometry can strongly affect their optimization. We show that the covariance $Σ_t$ of $p_t$ governs optimization bias: when $Σ_t$ is ill-conditioned, and gradient-based training rapidly fits high-variance directions while systematically under-optimizing low-variance modes, leading to learning that plateaus at suboptimal weights. We formalize this effect in analytically tractable settings and propose reversible, label-conditional \emph{preconditioning} maps that reshape the geometry of $p_t$ by improving the conditioning of $Σ_t$ without altering the underlying generative model. Rather than accelerating early convergence, preconditioning primarily mitigates optimization stagnation by enabling continued progress along previously suppressed directions. Across MNIST latent flow matching, and additional high-resolution datasets, we empirically track conditioning diagnostics and distributional metrics and show that preconditioning consistently yields better-trained models by avoiding suboptimal plateaus.