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

2026-01-21

AI Summary

This work investigates how to enable Rectified Flow to automatically adapt to the intrinsic low-dimensional structure of the target distribution during sampling, thereby improving efficiency. By establishing an equivalence between Rectified Flow and DDPM under a stochastic localization perspective, the authors propose a novel stochastic Rectified Flow sampler that requires less stringent accuracy in drift estimation. Coupled with an adaptive time discretization strategy, this sampler effectively exploits the low-dimensional geometry of the target distribution. Theoretical analysis shows that the method achieves an iteration complexity of O(k/ε) (up to logarithmic factors), where k denotes the intrinsic dimensionality of the target distribution. Experiments on both synthetic data and text-to-image generation tasks demonstrate the superior performance of the proposed sampler under the new time scheduling scheme.

Technology Category

• Reasoning under Uncertainty: Stochastic Optimization
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• 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
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

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.

Problem

• Rectified Flow
• low-dimensional adaptation
• intrinsic dimension
• sampling acceleration
• stochastic localization

Innovation

• Rectified Flow
• Low-Dimensional Adaptation
• Stochastic Localization
• Time-Discretization Schedule
• Diffusion Models