Fast Statistical Rates for Flow Matching without Curvature Assumptions

원문 초록

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 exact empirical risk minimization over a fixed network class and standard assumptions on the 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 for the statistical term with the Wasserstein guarantee taking the form $W_2 \le K(\varepsilon + C\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--Lojasiewicz condition in the statistical analysis. 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 the leading statistical prefactor through the Lipschitz constant of the pointwise loss, which we characterize explicitly for standard schedules.