Building Normalizing Flows with Stochastic Interpolants (2209.15571v3)

Abstract: A generative model based on a continuous-time normalizing flow between any pair of base and target probability densities is proposed. The velocity field of this flow is inferred from the probability current of a time-dependent density that interpolates between the base and the target in finite time. Unlike conventional normalizing flow inference methods based the maximum likelihood principle, which require costly backpropagation through ODE solvers, our interpolant approach leads to a simple quadratic loss for the velocity itself which is expressed in terms of expectations that are readily amenable to empirical estimation. The flow can be used to generate samples from either the base or target, and to estimate the likelihood at any time along the interpolant. In addition, the flow can be optimized to minimize the path length of the interpolant density, thereby paving the way for building optimal transport maps. In situations where the base is a Gaussian density, we also show that the velocity of our normalizing flow can also be used to construct a diffusion model to sample the target as well as estimate its score. However, our approach shows that we can bypass this diffusion completely and work at the level of the probability flow with greater simplicity, opening an avenue for methods based solely on ordinary differential equations as an alternative to those based on stochastic differential equations. Benchmarking on density estimation tasks illustrates that the learned flow can match and surpass conventional continuous flows at a fraction of the cost, and compares well with diffusions on image generation on CIFAR-10 and ImageNet $32\times32$. The method scales ab-initio ODE flows to previously unreachable image resolutions, demonstrated up to $128\times128$.

• The paper introduces a novel approach using stochastic interpolants instead of traditional maximum likelihood estimation to construct normalizing flows.
• It formulates continuous-time flows via ODEs to reduce computational complexity and enhance scalability for high-resolution image generation.
• Empirical results on CIFAR-10 and ImageNet demonstrate competitive performance with reduced computational costs.

Overview of "Building Normalizing Flows with Stochastic Interpolants"

This paper introduces a novel framework for constructing generative models using normalizing flows, named InterFlow, with an emphasis on computational efficiency and efficacy. The core innovation of the paper lies in utilizing stochastic interpolants instead of traditional maximum likelihood estimation (MLE) methods for generating normalizing flows, offering significant advantages in terms of reducing computational complexity and enhancing scalability. The approach is particularly promising for large-scale applications, such as high-resolution image generation.

Key Contributions

The paper details several noteworthy contributions to the field of generative modeling and optimal transport:

• Stochastic Interpolants: The introduction of stochastic interpolants serves as the foundational methodological shift. This approach constructs a continuous-time flow between any two probability densities, bypassing the need for costly backpropagation through ODE solvers that is typical in conventional normalizing flows. This leads to a quadratic loss function that is computationally efficient and readily amenable to empirical estimation.
• Continuous-Time Normalizing Flows: The proposed model constructs the flow as a solution to an ordinary differential equation (ODE) representing the transport dynamics. The velocity field is inferred from the probability current of a time-dependent density, promising a model that minimizes the path length of the interpolant density and provides a new lens on optimal transport maps.
• Practical Implementation and Numerical Experiments: The authors empirically validate the approach on various datasets, demonstrating its capability to equal or exceed the performance of standard continuous flows with reduced computational costs. Notably, the paper reports successful scaling of ODE flows to previously unachievable image resolutions, showing competitive performance on CIFAR-10 and ImageNet datasets.

Theoretical and Practical Implications

The research delineates clear implications for both theoretical exploration and practical applications:

• Theoretical Foundations and Score-Based Diffusion Models: The work provides a theoretical underpinning connecting stochastic interpolants to score-based diffusion models, offering a fresh perspective on the use of these models for generative purposes. It introduces a potential avenue for exploring dualities between diffusion models and ODE-based flows, suggesting pathways for future research.
• Generative Model Efficiency: By decoupling the optimization of the transport path from the optimization objective, the method achieves greater flexibility and computational efficiency. This separation is particularly advantageous in handling high-dimensional data, allowing more straightforward and scalable implementations.
• Transport Optimization and Scalability: The ability to optimize the transport using stochastic interpolants sets a new standard for building scalable generative models. The technique scales well to large datasets and high resolutions while retaining tractability—a critical factor for real-world applications in image and signal processing.

Future Directions

The paper opens several intriguing directions for future research and development:

• Exploration of Optimal Transport: Further exploration of optimal transport paths using the proposed interpolant method could yield substantial advancements in reducing computational overhead and improving model accuracy.
• Enhanced Models for High-Dimensional Data: Continued work on extending the framework to more complex datasets beyond image generation, such as audio and video, could demonstrate the method's full potential in diverse applications.
• Framework Extensions and Refinements: Developing alternative interpolant structures and employing adaptive learning techniques could further refine the balance between modeling complexity and computational efficiency.

In summary, "Building Normalizing Flows with Stochastic Interpolants" presents a promising pathway towards building more efficient and scalable generative models. Its innovative use of stochastic interpolants coupled with practical validation suggests a significant step forward in both the theoretical and applied domains of machine learning and generative modeling.