Manifold-Stable Flow Matching

Flow matching (FM) learns generative dynamics through velocity regression. A common assumption in geometric FM variants is that the prior distribution for the data is, in fact, a distribution on the data manifold. However, this requires knowledge of the data manifold, which in many applications is not available. In the absence of such prior knowledge, a low regression error alone does not guarantee that the FM conforms to the true data manifold. Manifold adherence keeps generated samples within valid configurations and is empirically associated with better task performance. In this paper, we introduce manifold-stable flow matching (MSFM), which can start from an arbitrary ambient prior (i.e., a prior distribution that is not necessarily on the manifold). The method leverages tools from nonlinear dynamics, namely contraction theory, to learn tangential transport with prescribed normal contraction to the manifold. Contraction theory provides a theoretical guarantee of convergence to the data manifold within a desired time window (e.g., one second). The construction uses analytical projectors when the data manifold is known, and local affine proxies estimated by principal component analysis of data samples when the data manifold is not known. We derive a family of probability paths compatible with our contraction guarantees and separate the training loss into a learnable manifold tangential term and a manifold normal residual. We establish manifold invariance and transverse convergence under explicit projection and proxy-selection assumptions. An ellipse experiment attains a mean terminal off-manifold error of order $10^{-6}$ (controlled by our set numerical tolerance). In our Push-T robotic experiments, MSFM raises the success rate from $74\%$ to $82\%$. In the Robomimic Square task, the success rate increases from $60\%$ to $72\%$, while the rotation error order decreases from $10^{-2}$ to $10^{-7}$ (controlled by our set numerical tolerance). These results show stronger geometric adherence and higher observed task performance, produced by learning the data manifold in generative AI tasks.