Order-Optimal Sample Complexity for Distribution Learning via Flow Matching

摘要

Flow-based generative models have emerged as an efficient alternative to diffusion models. We study the sample complexity of learning a target distribution in Wasserstein distance via flow matching, which learns a velocity field transporting a base distribution to the target. Under standard assumptions on the network architecture and data distribution, we show that any squared-loss flow-matching variant including rectified flow with bounded derivatives achieves $\widetilde{O}(\varepsilon^{-2})$ sample complexity with the Wasserstein guarantee taking the form $W_2 \le K(\varepsilon + C\sqrt{\varepsilon_{app}})$, where $\varepsilon_{app}$ is the approximation error of the network class. This improves upon existing $\widetilde{O}(\varepsilon^{-4})$ guarantees and does not require the Polyak \L{}ojasiewicz condition assumed in prior work. Instead, we exploit a structural property of the squared-loss objective - it induces an approximate Bernstein condition that directly ties the variance of the excess loss to the excess risk up to additive approximation-error. This condition enables a localized Rademacher complexity analysis yielding fast $O(1/n)$ rates. The flow-matching schedule enters only through the Lipschitz constant of the pointwise loss, which we characterize explicitly for standard schedules.