Balanced conic rectified flow

TL;DR

This work analyzes distributional drift in reflow-based rectified flow, showing that training with synthetic pairs alone leads to misalignment with real data and degraded full-step generation. It introduces balanced conic rectified flow, which combines conic reflow (real-image inversions with Slerp-based perturbations) and original reflow to anchor trajectories to real data while still covering the domain with synthetic supervision. The approach yields substantial gains on CIFAR-10 (notably FID improvements with a fraction of synthetic data) and demonstrates robustness across ImageNet 64×64 and LSUN, reducing curvature and initial-velocity delta of the learned trajectories. The method is modular, scalable, and compatible with existing rectified-flow models, offering a practical path toward more data-efficient, high-fidelity ODE-based generative modeling.

Abstract

Rectified flow is a generative model that learns smooth transport mappings between two distributions through an ordinary differential equation (ODE). Unlike diffusion-based generative models, which require costly numerical integration of a generative ODE to sample images with state-of-the-art quality, rectified flow uses an iterative process called reflow to learn smooth and straight ODE paths. This allows for relatively simple and efficient generation of high-quality images. However, rectified flow still faces several challenges. 1) The reflow process requires a large number of generative pairs to preserve the target distribution, leading to significant computational costs. 2) Since the model is typically trained using only generated image pairs, its performance heavily depends on the 1-rectified flow model, causing it to become biased towards the generated data. In this work, we experimentally expose the limitations of the original rectified flow and propose a novel approach that incorporates real images into the training process. By preserving the ODE paths for real images, our method effectively reduces reliance on large amounts of generated data. Instead, we demonstrate that the reflow process can be conducted efficiently using a much smaller set of generated and real images. In CIFAR-10, we achieved significantly better FID scores, not only in one-step generation but also in full-step simulations, while using only of the generative pairs compared to the original method. Furthermore, our approach induces straighter paths and avoids saturation on generated images during reflow, leading to more robust ODE learning while preserving the distribution of real images.

Figures (39)

• Figure 1: Problem of rectified flow. (a) By randomly pairing data $X_0 \sim \pi_0$ and $X_1 \sim \pi_1$, linear interpolation trajectories are defined. (b) The 1-rectified flow connects $X_0$ and $X_1$ with a learned velocity field which is potentially inaccurate. After the 1-rectified flow, the trajectories are rewired to avoid crossing. (c) The trajectories for reflow are defined as linear interpolation trajectories between $Z_0$ and the generated $Z_1 = v(Z_0)$. Note that $Z_1$ is drifted away from $\pi_1$. (d) Consequently, the 2-rectified flow has a velocity field drifted away from $X_1$.
• Figure 1: One-step and full-simulation comparison of 2,3 Rectified Flows on CIFAR-10.
• Figure 2: (a) 2-rectified flow overfits fake samples. Following the reverse and forward 2-rectified flow, fake images inherently return at similar images with or without perturbation at $\pi_0$. In contrast, real images return at different images and it is worse with perturbation, implying overfitting. (b) Reconstruction discrepancies emerge between real and fake images due to the use of fake-only pairs.
• Figure 3: (a) As reflow steps increase, generated samples diverge from the target distribution. (b) This drift is further evidenced by the rising KL divergence from the real data distribution.
• Figure 4: Illustration of original fake pairs and our real pairs. (a) The original rectified flow supervises 2-rectified flow with fake pairs $(Z_0, v_\theta(Z_0))$. (b) Our conic reflow supervises 2-rectified flow with real pairs $(v_\theta^{-1}(X_1), X_1)$ and their conic neighbors.