Improving and generalizing flow-based generative models with minibatch optimal transport (2302.00482v4)

Abstract: Continuous normalizing flows (CNFs) are an attractive generative modeling technique, but they have been held back by limitations in their simulation-based maximum likelihood training. We introduce the generalized conditional flow matching (CFM) technique, a family of simulation-free training objectives for CNFs. CFM features a stable regression objective like that used to train the stochastic flow in diffusion models but enjoys the efficient inference of deterministic flow models. In contrast to both diffusion models and prior CNF training algorithms, CFM does not require the source distribution to be Gaussian or require evaluation of its density. A variant of our objective is optimal transport CFM (OT-CFM), which creates simpler flows that are more stable to train and lead to faster inference, as evaluated in our experiments. Furthermore, we show that when the true OT plan is available, our OT-CFM method approximates dynamic OT. Training CNFs with CFM improves results on a variety of conditional and unconditional generation tasks, such as inferring single cell dynamics, unsupervised image translation, and Schr\"odinger bridge inference.

Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport

The paper "Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport" introduces a new set of simulation-free training objectives for Continuous Normalizing Flows (CNFs). These objectives, grouped under the generalized Conditional Flow Matching (CFM) framework, aim to address the efficiency and stability issues that have historically limited the widespread use of CNFs in generative modeling tasks. This essay will provide a detailed overview of the contributions, implications, and future prospects of this research.

Overview of Contributions

Generalized Conditional Flow Matching (CFM)

The paper proposes a unifying framework called Conditional Flow Matching (CFM), which generalizes several existing flow matching techniques. CFM uses a regression-based objective that does not require simulation of the Ordinary Differential Equation (ODE) during training. This approach is akin to the objective used in stochastic flow models like diffusion models but extends it to a more general setting.

In CFM, conditional flows are used to generate a joint probability path from a mixture of simpler conditional probability paths. This is formalized through a collection of vector fields $u_t(x | z)$ which, given a dataset, guide the learning of a marginal vector field $u_t(x)$. The training objective minimizes the difference between the learned vector field and the true conditional vector fields.

Optimal Transport Conditional Flow Matching (OT-CFM)

One of the standout contributions is the introduction of OT-CFM, which leverages the principles of optimal transport to generate more stable and efficient flows. OT-CFM conditions on pairs of data points sampled from a minibatch optimal transport plan, leading to straighter and more efficiently integrable flows. This method approximates dynamic optimal transport (DOT), providing a new avenue for solving the dynamic OT problem without the need for complex simulations.

Schrödinger Bridge Conditional Flow Matching (SB-CFM)

The research also explores a variant known as SB-CFM, which approximates the probability flow associated with Schrödinger bridges. By using entropy-regularized OT plans for the conditional flows, SB-CFM provides a simulation-free approach to approximating Schrödinger bridges. This is particularly useful for tasks requiring the interpolation between distributions derived from real-world stochastic processes.

Experimental Validation

Extensive experimental validation across various datasets demonstrates the practical viability of the proposed methods.

1. Low-Dimensional Data: By evaluating on classic benchmarking datasets like Moons, 8-Gaussians, and Scurve, the authors show that OT-CFM achieves significantly lower Normalized Path Energy (NPE) and faster training convergence compared to existing methods. This suggests that OT-CFM more accurately approximates dynamic OT paths.
2. Single-Cell Dynamics: OT-CFM is applied to the task of single-cell trajectory estimation, outperforming other methods in terms of 1-Wasserstein distance on hidden timepoints, indicating its superiority in modeling biological processes with complex dynamics.
3. High-Dimensional Image Data: The performance of OT-CFM is further corroborated in high-dimensional tasks like CIFAR-10 image generation, where it achieves superior Fréchet Inception Distance (FID) scores and requires fewer function evaluations for good quality sampling.
4. Unsupervised Image Translation: The method is also validated on the CelebA dataset for unsupervised attribute translation, successfully learning mappings between complex data distributions in a latent space as measured by Maximum Mean Discrepancy (MMD) scores.

Implications and Future Prospects

Efficiency and Stability

The introduction of OT-CFM and SB-CFM addresses the high computational cost and instability traditionally associated with CNFs. The techniques achieve lower variance in the training objective, leading to faster convergence and more stable training processes. The ability of OT-CFM to approximate dynamic OT paths with fewer computational resources is particularly noteworthy, opening up new possibilities for large-scale applications of CNFs.

Theoretical Impact

The theoretical contribution of a unifying CFM framework that generalizes existing methods lays a strong foundation for future research. This work not only captures the essence of various simulation-free training paradigms but also extends them to broader settings involving arbitrary source distributions.

Future Developments in AI

The research paves the way for more efficient and scalable generative models. Future developments could explore deeper integrations of optimal transport principles, potentially leading to new state-of-the-art models in areas such as image generation, natural language processing, and beyond. The simulation-free nature of CFM objectives makes them particularly appealing for resource-constrained environments, expanding the applicability of high-caliber AI models.

In summary, the paper "Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport" marks a significant step forward in the field of generative modeling, addressing key limitations of CNFs while broadening their applicability. The novel CFM framework, along with its OT-CFM and SB-CFM variants, showcases the power of optimal transport and conditional flow matching in achieving stable, efficient, and high-quality generative models.