Flow Matching for Generative Modeling (2210.02747v2)

Abstract: We introduce a new paradigm for generative modeling built on Continuous Normalizing Flows (CNFs), allowing us to train CNFs at unprecedented scale. Specifically, we present the notion of Flow Matching (FM), a simulation-free approach for training CNFs based on regressing vector fields of fixed conditional probability paths. Flow Matching is compatible with a general family of Gaussian probability paths for transforming between noise and data samples -- which subsumes existing diffusion paths as specific instances. Interestingly, we find that employing FM with diffusion paths results in a more robust and stable alternative for training diffusion models. Furthermore, Flow Matching opens the door to training CNFs with other, non-diffusion probability paths. An instance of particular interest is using Optimal Transport (OT) displacement interpolation to define the conditional probability paths. These paths are more efficient than diffusion paths, provide faster training and sampling, and result in better generalization. Training CNFs using Flow Matching on ImageNet leads to consistently better performance than alternative diffusion-based methods in terms of both likelihood and sample quality, and allows fast and reliable sample generation using off-the-shelf numerical ODE solvers.

Flow Matching for Generative Modeling

The paper "Flow Matching for Generative Modeling" introduces a novel training paradigm for Continuous Normalizing Flows (CNFs) that enhances scalability and efficiency in generative modeling tasks. This new method, termed Flow Matching (FM), addresses the computational challenges associated with previous CNF training methods by offering a simulation-free approach that leverages vector fields of fixed conditional probability paths. Specifically, FM enables CNFs to operate efficiently with general Gaussian probability paths, encompassing existing diffusion paths as special cases and extending to other probability paths defined by Optimal Transport (OT) principles.

Key Contributions and Methods

1. Flow Matching Paradigm:
   • FM is predicated on regressing CNF vector fields to target vector fields that generate specific probability paths.
   • The core objective is minimizing the FM loss, which quantifies the difference between the model's vector field and the target vector field. This loss is computed as $$\mathcal{L}_{FM}(\theta) = \mathbb{E}_{t,p_t(x)}\| v_t(x) - u_t(x) \|^2$$.
2. Conditional Flow Matching (CFM):
   • To mitigate the intractability of direct FM, the paper introduces the CFM objective. This objective operates on per-example probability paths, thus it does not require explicit knowledge of the marginal target vector field.
   • The equivalence in gradient between FM and CFM allows training to proceed using simplified conditional paths, yet ensuring the CNF models the correct marginal distribution at convergence.
3. Gaussian Probability Paths:
   • The authors define a flexible family of Gaussian conditional probability paths that include both traditional diffusion paths and OT paths.
   • OT paths, characterized by linear mean and variance changes over time, demonstrate empirical benefits in terms of faster training and more efficient sampling compared to diffusion paths.
4. Implementation and Empirical Validation:
   • The paper includes extensive experiments on image datasets such as CIFAR-10 and ImageNet (at various resolutions).
   • Empirical results show that models trained with FM and OT paths consistently outperform those trained with standard diffusion-based methods with respect to both likelihood estimation and sample quality metrics like FID.
   • Notably, FM-OT results in significant efficiency gains during sampling, evident from reduced function evaluations required to achieve high-quality generative samples.

Results and Implications

The primary findings include:

• Performance: On ImageNet, FM-OT achieves better performance in terms of likelihood and FID scores in comparison to models trained using traditional score matching and diffusion-based methods.
• Efficiency: By enabling direct control over probability paths, FM-OT models achieve faster training convergence and reduced sampling time, potentially reining in the computational expense typically associated with generative models.
• Stability: The deterministic nature of CNFs and the stability afforded by FM make the training process more robust, leading to fewer epochs required to achieve significant performance improvements.

Practical and Theoretical Implications

Practical Applications:

• Scalability: The reduction in training and sampling costs makes FM-OT particularly appealing for large-scale applications, such as high-resolution image synthesis and other high-dimensional data generation tasks.
• Flexibility: The ability to craft and utilize non-diffusion probability paths opens new avenues for custom-tailored generative models suited to specific application needs, including domains with stringent efficiency or performance constraints.

Theoretical Insights:

• Generalization Beyond Diffusion: FM provides a framework to generalize beyond the specific probabilistic constraints of diffusion models, suggesting that direct probability path reasoning can offer both theoretical clarity and practical benefits.
• Optimal Transport: Leveraging OT paths not only leads to more efficient computations but also highlights an elegant theoretical advancement in the use of optimal transport theory in the context of generative modeling.

Future Directions

• Diverse Probability Paths: Future research can explore the efficacy of various non-isotropic Gaussian or other more complex probability paths within the FM framework.
• Broader Applications: Extending the principles of FM and OT to other generative tasks, such as text generation or time-series modeling, could prove highly beneficial.
• Optimization Strategies: Investigating alternative optimization techniques to further expedite the FM training process or enhance generalization capabilities of the resulting models.

In summary, the "Flow Matching for Generative Modeling" paper delivers a significant advancement in the training of CNF models, outlining a method that outperforms traditional diffusion-based generative modeling approaches both in performance metrics and computational efficiency. The integration of OT principles showcases the potential for further innovations in the field, broadening the scope and applicability of generative modeling techniques.