Integration Flow Models
Jingjing Wang, Dan Zhang, Joshua Luo, Yin Yang, Feng Luo
TL;DR
Integration Flow addresses the limitations of traditional ODE-based generative models by bypassing ODE solvers and directly learning the integrated trajectory effect via a single function $G$. The model anchors reverse-time dynamics to the target state $\mathbf{x}_0$, enabling stable, one-step generation across diffusion, Rectified Flow, and PFGM++ with a unifying neural parameterization $g_{\boldsymbol{\theta}}(\mathbf{x}_0,\mathbf{x}_t,t)$. Theoretical results establish stability, non-intersection, and optimality within the Flow Matching/Rectified Flow framework. Empirically, Integration Flow achieves competitive or state-of-the-art one-step generation on CIFAR-10 and ImageNet 64x64 across VE diffusion, Rectified Flow, and PFGM++, demonstrating improved efficiency and scalability. The framework paves the way for solver-free, unified, and scalable ODE-based generative modeling, with potential impact on rapid sampling and broader applicability across generative paradigms.
Abstract
Ordinary differential equation (ODE) based generative models have emerged as a powerful approach for producing high-quality samples in many applications. However, the ODE-based methods either suffer the discretization error of numerical solvers of ODE, which restricts the quality of samples when only a few NFEs are used, or struggle with training instability. In this paper, we proposed Integration Flow, which directly learns the integral of ODE-based trajectory paths without solving the ODE functions. Moreover, Integration Flow explicitly incorporates the target state $\mathbf{x}_0$ as the anchor state in guiding the reverse-time dynamics. We have theoretically proven this can contribute to both stability and accuracy. To the best of our knowledge, Integration Flow is the first model with a unified structure to estimate ODE-based generative models and the first to show the exact straightness of 1-Rectified Flow without reflow. Through theoretical analysis and empirical evaluations, we show that Integration Flows achieve improved performance when it is applied to existing ODE-based models, such as diffusion models, Rectified Flows, and PFGM++. Specifically, Integration Flow achieves one-step generation on CIFAR10 with FIDs of 2.86 for the Variance Exploding (VE) diffusion model, 3.36 for rectified flow without reflow, and 2.91 for PFGM++; and on ImageNet with FIDs of 4.09 for VE diffusion model, 4.35 for rectified flow without reflow and 4.15 for PFGM++.