Switched Flow Matching: Eliminating Singularities via Switching ODEs
Qunxi Zhu, Wei Lin
TL;DR
This work tackles singularity issues in Flow Matching (FM) by introducing Switched Flow Matching (SFM), which uses switching ODEs to transport between heterogeneous source and target distributions. The authors establish that FM can fail to transport between simple distributions due to fundamental initial-value problem (IVP) constraints, and show that switching among multiple conditional vector fields yields better regularity and tractable sampling. They develop a formal framework where sources, targets, and probability paths are mixtures conditioned on a switching signal, derive a training objective (SFM) and a gradient-equivalent SCFM variant, and propose practical switching mechanisms based on batch clustering and coupling matrices. Experiments on synthetic distributions and CIFAR-10 demonstrate that SFM halves the singularity problem, yields straighter flows, and provides competitive sampling efficiency, with potential for seamless integration with minibatch optimal transport. Overall, SFM offers a principled, scalable path to robust, simulation-free generative modeling via switching ODEs and dynamic OT concepts.
Abstract
Continuous-time generative models, such as Flow Matching (FM), construct probability paths to transport between one distribution and another through the simulation-free learning of the neural ordinary differential equations (ODEs). During inference, however, the learned model often requires multiple neural network evaluations to accurately integrate the flow, resulting in a slow sampling speed. We attribute the reason to the inherent (joint) heterogeneity of source and/or target distributions, namely the singularity problem, which poses challenges for training the neural ODEs effectively. To address this issue, we propose a more general framework, termed Switched FM (SFM), that eliminates singularities via switching ODEs, as opposed to using a uniform ODE in FM. Importantly, we theoretically show that FM cannot transport between two simple distributions due to the existence and uniqueness of initial value problems of ODEs, while these limitations can be well tackled by SFM. From an orthogonal perspective, our framework can seamlessly integrate with the existing advanced techniques, such as minibatch optimal transport, to further enhance the straightness of the flow, yielding a more efficient sampling process with reduced costs. We demonstrate the effectiveness of the newly proposed SFM through several numerical examples.