Total Variation Rates for Riemannian Flow Matching
Yunrui Guan, Krishnakumar Balasubramanian, Shiqian Ma
TL;DR
The paper provides the first nonasymptotic TV bounds for Riemannian flow matching samplers discretized by Euler steps on manifolds. It derives a TV-differential inequality that links the discretization error to the divergence and gradient of the vector-field mismatch, incorporating curvature and parallel transport effects, and combines this with a carefully constructed continuous-time interpolation of the Euler scheme. Under smooth population fields and learning guarantees (uniform on compact manifolds or moment-based on Hadamard manifolds), the authors obtain explicit TV bounds of the form $\mathrm{TV} \le h\,C_{\mathrm{Lip}} + \varepsilon\,C_{\varepsilon} + \text{(higher-order terms in }\varepsilon\text{)}$, separating numerical discretization from learning error. They instantiate the theory on hyperspheres and SPD$(n)$, yielding concrete iteration complexities and demonstrating practical implications for manifold-valued data sampling without relying on diffusion heat kernels.
Abstract
Riemannian flow matching (RFM) extends flow-based generative modeling to data supported on manifolds by learning a time-dependent tangent vector field whose flow-ODE transports a simple base distribution to the data law. We develop a nonasymptotic Total Variation (TV) convergence analysis for RFM samplers that use a learned vector field together with Euler discretization on manifolds. Our key technical ingredient is a differential inequality governing the evolution of TV between two manifold ODE flows, which expresses the time-derivative of TV through the divergence of the vector-field mismatch and the score of the reference flow; controlling these terms requires establishing new bounds that explicitly account for parallel transport and curvature. Under smoothness assumptions on the population flow-matching field and either uniform (compact manifolds) or mean-square (Hadamard manifolds) approximation guarantees for the learned field, we obtain explicit bounds of the form $\mathrm{TV}\le C_{\mathrm{Lip}}\,h + C_{\varepsilon}\,\varepsilon$ (with an additional higher-order $\varepsilon^2$ term on compact manifolds), cleanly separating numerical discretization and learning errors. Here, $h$ is the step-size and $\varepsilon$ is the target accuracy. Instantiations yield \emph{explicit} polynomial iteration complexities on the hypersphere $S^d$, and on the SPD$(n)$ manifolds under mild moment conditions.