Pathway to $O(\sqrt{d})$ Complexity bound under Wasserstein metric of flow-based models

TL;DR

This work analyzes flow-based generative models under the Wasserstein metric to obtain an optimal sampling complexity of $O(\sqrt{d})$. By formulating a Lipschitz-change-of-variables framework and examining the Föllmer flow under a Gaussian-tail assumption, it shows dimension-free Lipschitz regularity of the velocity field and tight $W_2$ bounds for discretized flows. The main results establish that, for isotropic Gaussian bases, the required number of steps to achieve a prescribed Wasserstein accuracy scales as $N=\mathcal{O}(\sqrt{d}/\epsilon)$, and this extends to the $1$-rectified flow and bounded-support scenarios, with potential implications for infinite-dimensional and Bayesian inverse problems. Overall, the paper provides a rigorous pathway to achieving $O(\sqrt{d})$ sampling complexity in high-dimensional or function-space settings, highlighting the central role of Lipschitz transport maps and the Gaussian-tail structure in enabling dimension-free guarantees.

Abstract

We provide attainable analytical tools to estimate the error of flow-based generative models under the Wasserstein metric and to establish the optimal sampling iteration complexity bound with respect to dimension as $O(\sqrt{d})$. We show this error can be explicitly controlled by two parts: the Lipschitzness of the push-forward maps of the backward flow which scales independently of the dimension; and a local discretization error scales $O(\sqrt{d})$ in terms of dimension. The former one is related to the existence of Lipschitz changes of variables induced by the (heat) flow. The latter one consists of the regularity of the score function in both spatial and temporal directions. These assumptions are valid in the flow-based generative model associated with the Föllmer process and $1$-rectified flow under the Gaussian tail assumption. As a consequence, we show that the sampling iteration complexity grows linearly with the square root of the trace of the covariance operator, which is related to the invariant distribution of the forward process.

Theorems & Definitions (43)

• Definition 2.1: Föllmer flow in formal sense
• Theorem 3.5
• Theorem 3.8: Regularity of the velocity field
• Remark 3.9
• Lemma 3.10: Well-posedness
• Corollary 3.11: Lipschitzness of continuous flow
• Corollary 3.14: Lipschitzness of discrete flow
• Theorem 3.15
• Corollary 3.16
• Corollary 3.18