Manifold-stable flow matching improves geometric accuracy and task success
Manifold-Stable Flow Matching
Machine LearningRobotics
Summary
Generating realistic data often means keeping the generated samples on a specific curved space called a manifold, which is tricky without knowing the manifold's structure. The authors propose a new method that ensures the generated data stays on or close to this manifold by combining learned directions of movement with a way to pull points back to the manifold. This approach works both when the manifold is known and when it is estimated from data. Their method reduces errors drastically and improves success rates in robot tasks, showing better adherence to the correct geometry helps performance.
What this means in practice
- •For robotics engineers: Improve robot control success by generating motions that strictly follow a robot's valid movement manifold.
- •For machine learning practitioners: Train generative models that respect data geometry even when the underlying manifold is unknown or complex.
Authors
Amirhossein Nazerian, Ali Pezeshki, Jianguo Zhao
Abstract
Flow matching (FM) learns generative dynamics through velocity regression. Geometric FM variants commonly assume a prior supported on the data manifold, requiring geometric knowledge that is often unavailable. Without such knowledge, low regression error alone does not guarantee manifold adherence. Adherence keeps generated samples within valid configurations and is empirically associated with better task performance. We introduce manifold-stable flow matching (MSFM), which can start from an arbitrary ambient prior, not necessarily supported on the manifold. Using tools from nonlinear dynamics, namely contraction theory, MSFM combines learned tangential transport with prescribed normal contraction. The construction uses analytical projectors for known manifolds and local affine proxies estimated by principal component analysis for unknown data geometry. By implementing contraction theory in both cases of known and unknown manifolds, we guarantee manifold invariance and transverse convergence to the manifold within a desired time window (e.g., one second). We derive a family of compatible probability paths and decompose the training loss into a learnable tangential term and a normal residual. An ellipse experiment attains a mean terminal off-manifold error of order $10^{-6}$. In Push-T robotic experiments, MSFM raises success from $74\%$ to $82\%$. In the Robomimic Square task, success increases from $60\%$ to $72\%$, while rotation-manifold deviation decreases from order $10^{-2}$ to $10^{-7}$. The MSFM terminal geometric errors are controlled by the chosen numerical tolerance. These results demonstrate stronger geometric adherence and higher observed task performance, supporting prescribed normal contraction as a complement to learned generative transport.