Switched Flow Matching: Eliminating Singularities via Switching ODEs

TL;DR

This paper introduces Switched Flow Matching (SFM), a novel framework that eliminates singularities in neural ODEs for continuous-time generative models, leading to faster and more efficient sampling.

Contribution

The paper proposes SFM, which uses switching ODEs to address singularities in flow matching, overcoming limitations of traditional FM and improving sampling efficiency.

Findings

SFM effectively eliminates singularities in flow matching.

Theoretical proof shows FM cannot transport between simple distributions, but SFM can.

Numerical examples demonstrate improved sampling efficiency with SFM.

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.