Preconditioned Score and Flow Matching

2026-03-02

AI Summary

This work addresses the issue of suboptimal stagnation in generative model training caused by ill-conditioned covariance matrices in intermediate distributions, which leads to the neglect of low-variance directions during optimization. To mitigate this directional bias without altering the generative model itself, the authors propose an invertible label-conditional preprocessing map that reshapes the geometry of intermediate distributions and improves the condition number of their covariance matrices. This approach departs from conventional acceleration paradigms and is compatible with both flow matching and score-based diffusion frameworks. Experiments on the MNIST latent space and high-resolution datasets demonstrate a significant reduction in suboptimal plateaus, yielding improved training dynamics and enhanced distributional fidelity.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Generative Adversarial Networks (GANs) for Vision
• Natural Language Processing: Generation

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

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.

Innovation

• preconditioning
• flow matching
• score-based diffusion
• covariance conditioning
• optimization geometry