Table of Contents
Fetching ...

Order-Optimal Sample Complexity of Rectified Flows

Hari Krishna Sahoo, Mudit Gaur, Vaneet Aggarwal

TL;DR

This work analyzes rectified flow models, which constrain transport trajectories to straight-line paths between base and target distributions, and proves they achieve order-optimal sample complexity $\tilde{O}(\varepsilon^{-2})$ in $W_2$ distance. The authors develop a localized Rademacher-complexity framework tailored to the rectified-flow geometry (linear paths and squared-loss objective), enabling a fast $O(1/n)$ statistical error rate. They decompose the total error into approximation, statistical, and optimization components and provide an SGD-based optimization analysis under PL and smoothness assumptions, showing the optimization error is also $O(1/n)$. A matching information-theoretic lower bound confirms the rate is optimal, and the results offer a theoretical explanation for the observed empirical efficiency of rectified flows in fast, accurate distributional modeling via only a few discretization steps.

Abstract

Recently, flow-based generative models have shown superior efficiency compared to diffusion models. In this paper, we study rectified flow models, which constrain transport trajectories to be linear from the base distribution to the data distribution. This structural restriction greatly accelerates sampling, often enabling high-quality generation with a single Euler step. Under standard assumptions on the neural network classes used to parameterize the velocity field and data distribution, we prove that rectified flows achieve sample complexity $\tilde{O}(\varepsilon^{-2})$. This improves on the best known $O(\varepsilon^{-4})$ bounds for flow matching model and matches the optimal rate for mean estimation. Our analysis exploits the particular structure of rectified flows: because the model is trained with a squared loss along linear paths, the associated hypothesis class admits a sharply controlled localized Rademacher complexity. This yields the improved, order-optimal sample complexity and provides a theoretical explanation for the strong empirical performance of rectified flow models.

Order-Optimal Sample Complexity of Rectified Flows

TL;DR

This work analyzes rectified flow models, which constrain transport trajectories to straight-line paths between base and target distributions, and proves they achieve order-optimal sample complexity in distance. The authors develop a localized Rademacher-complexity framework tailored to the rectified-flow geometry (linear paths and squared-loss objective), enabling a fast statistical error rate. They decompose the total error into approximation, statistical, and optimization components and provide an SGD-based optimization analysis under PL and smoothness assumptions, showing the optimization error is also . A matching information-theoretic lower bound confirms the rate is optimal, and the results offer a theoretical explanation for the observed empirical efficiency of rectified flows in fast, accurate distributional modeling via only a few discretization steps.

Abstract

Recently, flow-based generative models have shown superior efficiency compared to diffusion models. In this paper, we study rectified flow models, which constrain transport trajectories to be linear from the base distribution to the data distribution. This structural restriction greatly accelerates sampling, often enabling high-quality generation with a single Euler step. Under standard assumptions on the neural network classes used to parameterize the velocity field and data distribution, we prove that rectified flows achieve sample complexity . This improves on the best known bounds for flow matching model and matches the optimal rate for mean estimation. Our analysis exploits the particular structure of rectified flows: because the model is trained with a squared loss along linear paths, the associated hypothesis class admits a sharply controlled localized Rademacher complexity. This yields the improved, order-optimal sample complexity and provides a theoretical explanation for the strong empirical performance of rectified flow models.
Paper Structure (34 sections, 17 theorems, 154 equations, 1 table, 1 algorithm)

This paper contains 34 sections, 17 theorems, 154 equations, 1 table, 1 algorithm.

Key Result

Theorem 1

benton2024nearlydlinearconvergencebounds Let $\pi_0, \pi_1$ be probability distributions on $\mathbb{R}^d$. Let $v^*$ be the true velocity field and $v_\theta$ the learned field. Define $Y_t$ and $Z_t$ as the flows generated by $v_\theta$ and $v^*$ respectively, both starting from $Y_0 = Z_0 \sim \p The Wasserstein-2 distance between the generated and target distributions satisfies:

Theorems & Definitions (33)

  • Remark 1
  • Remark 2
  • Remark 3
  • Remark 4
  • Remark 5
  • Remark 6
  • Theorem 1
  • Proposition 1: Error Decomposition
  • Theorem 2: Sample Complexity of Rectified Flow
  • proof : Proof Sketch
  • ...and 23 more