2-Rectifications are Enough for Straight Flows: A Theoretical Insight into Wasserstein Convergence

2024-10-19

AI Summary

This work addresses the theoretical foundation of trajectory linearization in Rectified Flows, specifically investigating the minimal number of rectification steps required to achieve optimal transport from a Gaussian prior to an arbitrary target distribution with finite first-order moments (e.g., Gaussian mixtures). Leveraging optimal transport theory, convex optimization, and manifold regularity analysis, we derive an explicit upper bound on the Wasserstein error between the sampling and target distributions. We rigorously prove— for the first time—that only two rectification steps suffice to yield a perfectly linearized, globally optimal transport path. This result breaks the conventional reliance on multi-step correction schemes and provides a solid theoretical guarantee for efficient generative modeling. Empirical evaluation demonstrates that the two-step scheme significantly improves generation quality on both synthetic and real-world datasets, while reducing the number of function evaluations (NFE) by over 80% compared to baselines such as DDPM.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Optimization
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Diffusion models have emerged as a powerful tool for image generation and denoising. Typically, generative models learn a trajectory between the starting noise distribution and the target data distribution. Recently Liu et al. (2023b) designed a novel alternative generative model Rectified Flow (RF), which aims to learn straight flow trajectories from noise to data using a sequence of convex optimization problems with close ties to optimal transport. If the trajectory is curved, one must use many Euler discretization steps or novel strategies, such as exponential integrators, to achieve a satisfactory generation quality. In contrast, RF has been shown to theoretically straighten the trajectory through successive rectifications, reducing the number of function evaluations (NFEs) while sampling. It has also been shown empirically that RF may improve the straightness in two rectifications if one can solve the underlying optimization problem within a sufficiently small error. In this paper, we make two key theoretical contributions: 1) we provide the first theoretical analysis of the Wasserstein distance between the sampling distribution of RF and the target distribution. Our error rate is characterized by the number of discretization steps and a extit{new formulation of straightness} stronger than that in the original work. 2) under a mild regularity assumption, we show that for a rectified flow from a Gaussian to any general target distribution with finite first moment (e.g. mixture of Gaussians), two rectifications are sufficient to achieve a straight flow, which is in line with the previous empirical findings. Additionally, we also present empirical results on both simulated and real datasets to validate our theoretical findings.

Problem

• Analyze Wasserstein distance in Rectified Flow.
• Prove two rectifications ensure straight flow.
• Validate theory with simulated and real data.

Innovation

• Rectified Flow straightens trajectories
• Two rectifications suffice
• Wasserstein distance theoretical analysis