Low-Dimensional Adaptation of Rectified Flow: A Diffusion and Stochastic Localization Perspective

TL;DR

This paper demonstrates that rectified flow methods can adapt to the low intrinsic dimensionality of target distributions, leading to more efficient sampling, and introduces a stochastic RF variant leveraging stochastic localization for improved performance.

Contribution

It reveals the low-dimensional adaptation property of rectified flow and proposes a stochastic RF method based on stochastic localization, with new time-discretization schemes and theoretical iteration complexity bounds.

Findings

RF achieves $O(k/)$ iteration complexity for intrinsic dimension $k$

Stochastic RF adapts to low-dimensionality with milder drift accuracy requirements

Simulations show improved performance of the proposed methods

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.