Low-Dimensional Adaptation of Rectified Flow: A New Perspective through the Lens of Diffusion and Stochastic Localization

TL;DR

The paper tackles accelerating sampling with Rectified Flow by exploiting intrinsic low-dimensional structure. It introduces a nonuniform, U-shaped time discretization and reveals a deep connection between RF, DDPM, and stochastic localization, culminating in a stochastic RF (Stoc-RF) that adapts to low dimensionality under weaker drift accuracy assumptions. The authors prove convergence rates whose TV error scales with the target’s intrinsic dimension $k$ and provide explicit bounds under drift-approximation errors; they also demonstrate practical gains via synthetic and T2I experiments. This work advances efficient diffusion-flows hybrids, offering theoretically grounded, dimension-aware samplers with significant potential for high-dimensional data generation and image synthesis.

Abstract

In recent years, Rectified flow (RF) has gained considerable popularity largely due to its generation efficiency and state-of-the-art performance. In this paper, we investigate the degree to which RF automatically adapts to the intrinsic low dimensionality of the support of the target distribution to accelerate sampling. We show that, using a carefully designed choice of the time-discretization scheme and with sufficiently accurate drift estimates, the RF sampler enjoys an iteration complexity of order $O(k/\varepsilon)$ (up to log factors), where $\varepsilon$ is the precision in total variation distance and $k$ is the intrinsic dimension of the target distribution. In addition, we show that the denoising diffusion probabilistic model (DDPM) procedure is equivalent to a stochastic version of RF by establishing a novel connection between these processes and stochastic localization. Building on this connection, we further design a stochastic RF sampler that also adapts to the low-dimensionality of the target distribution under milder requirements on the accuracy of the drift estimates, and also with a specific time schedule. We illustrate with simulations on the synthetic data and text-to-image data experiments the improved performance of the proposed samplers implementing the newly designed time-discretization schedules.

Figures (7)

• Figure 1: Trajectories of RF and Stoc-RF samplers for a mixture of 2-Gaussian target distribution.
• Figure 2: Histogram and kernel density estimation (KDE) plot of time-grid \ref{['eq: gemoetric time discretization']} showing U-shaped distribution.
• Figure 3: (a)-(b) TV distance between blurred low-rank Gaussian distribution and the generated samples by the RF sampler under uniform and non-uniform time-grid with varying $d$. (c)-(d) TV distance between blurred low-rank Gaussian distribution and the generated samples by the RF sampler under uniform and non-uniform time-grid with varying $d$.
• Figure 4: Generated images using Flux for specific prompts via Euler's scheme \ref{['eq: emp-ode-disc']} under non-uniform time-grid \ref{['eq: gemoetric time discretization']} (Red frame) and uniform time-grid (Blue frame).
• Figure 5: Generated images using Flux for specific prompts via Euler's scheme \ref{['eq: emp-ode-disc']} under non-uniform time-grid \ref{['eq: gemoetric time discretization']} (Columns 1), uniform time-grid (Column 2), Stoc-RF sampler (Column 3), and vanilla SDE sampler (Column 4). Blue color indicates the proposed methods in this paper. The generation quality of Stoc-Rf is generally better out of the four samplers.

Theorems & Definitions (13)

• Lemma 3.1: eldan2020taming, Equation (11)
• Lemma 3.2
• Lemma 3.3
• Proposition 3.4
• Proposition 3.5
• Theorem 4.5
• Theorem 4.7
• Remark 4.8
• Lemma 2.1: Lemma 5, liang2025low
• Proposition 2.2