Flow Map Matching

Nicholas M. Boffi,Michael S. Albergo,Eric Vanden-Eijnden

2024-06-12

Abstract:Generative models based on dynamical transport of measure, such as diffusion models, flow matching models, and stochastic interpolants, learn an ordinary or stochastic differential equation whose trajectories push initial conditions from a known base distribution onto the target. While training is cheap, samples are generated via simulation, which is more expensive than one-step models like GANs. To close this gap, we introduce flow map matching -- an algorithm that learns the two-time flow map of an underlying ordinary differential equation. The approach leads to an efficient few-step generative model whose step count can be chosen a-posteriori to smoothly trade off accuracy for computational expense. Leveraging the stochastic interpolant framework, we introduce losses for both direct training of flow maps and distillation from pre-trained (or otherwise known) velocity fields. Theoretically, we show that our approach unifies many existing few-step generative models, including consistency models, consistency trajectory models, progressive distillation, and neural operator approaches, which can be obtained as particular cases of our formalism. With experiments on CIFAR-10 and ImageNet 32x32, we show that flow map matching leads to high-quality samples with significantly reduced sampling cost compared to diffusion or stochastic interpolant methods.

Machine Learning,Dynamical Systems

What problem does this paper attempt to address?

The paper primarily aims to address the efficiency issues in sample generation processes of generative models, particularly those based on dynamical transport of measure, such as diffusion models and flow matching models. Although these models have low training costs, they require solving ordinary differential equations or stochastic differential equations through numerical integration during sample generation, leading to high computational costs.

To solve this problem, the authors propose the "Flow Map Matching" method. The core of this method is to learn a two-time flow map, which is a mapping of an ordinary differential equation between two time points. Compared to traditional generative models, the flow map matching model has the following advantages:

1. **Efficiency**: The proposed model can generate high-quality samples in fewer steps, significantly reducing sampling costs.
2. **Flexibility**: The number of steps can be adjusted post-training to balance accuracy and computational overhead.
3. **Generality**: It is applicable to different generative model frameworks, including consistency models and consistency trajectory models, and can be combined with existing neural network techniques.

Specifically, the contributions of the paper include:
- Proposing a framework for learning the two-time flow map of any ordinary differential equation and comparing it with recent related work.
- Demonstrating that learning a two-time flow map can achieve an efficient few-step generative model, with the number of steps adjustable post-training.
- Introducing two loss functions: the Lagrangian loss function and the Eulerian loss function, for extracting flow maps from pre-trained velocity fields; and proving that these loss functions can control the Wasserstein distance between the true distribution and the learned distribution.
- Proposing a new loss function for directly training flow map models, which does not require pre-trained velocity fields and can stably and efficiently train few-step generative models.
- Introducing a map distillation objective inspired by progressive distillation, which can efficiently convert a few-step flow map model into a one-step flow map model.

The experimental section demonstrates the application and effectiveness of the flow map matching method on the CIFAR-10 and ImageNet 32x32 datasets, confirming its validity and generality.