Source-linked AI summary
Manifold-Stable Flow Matching
Amirhossein Nazerian, Ali Pezeshki, Jianguo Zhao
TL;DR
Flow matching can produce off-manifold samples, while geometric variants commonly require known geometry and a prior supported on the manifold. MSFM combines learned tangential transport with prescribed normal contraction, using analytical projectors or local PCA proxies; it reports stronger geometric adherence and higher observed task performance across its evaluations.
Problem
Geometric flow-matching variants commonly require a manifold-supported prior and geometric knowledge, while low regression error alone does not guarantee manifold adherence.
Source: §1 · pp. 1–2Method
MSFM learns tangential transport and prescribes normal contraction, using analytical projectors for known manifolds and local PCA affine proxies for unknown geometry.
Source: §1 · p. 2Results
Takeaways & Limitations
The results support prescribed normal contraction as a complement to learned generative transport, with stronger geometric adherence and higher observed task performance.
Source: §7 · p. 9The robotic comparisons use deliberately compact networks, and Push-T's repeated-block-disturbance protocol differs from standard benchmarks.
Source: §7 · p. 9
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Flow matching is attractive for generative policies, but structured outputs often lie on lower-dimensional manifolds where off-manifold samples can be invalid or unsupported. MSFM separates learned tangential transport from prescribed normal contraction to provide convergence toward the manifold, including when starting from an ambient prior.
1 INTRODUCTION
- The framework targets structured distributions concentrated near lower-dimensional manifolds, including rotations, rigid-body poses, constrained actions, and task manifolds.Flow matching learns a time-dependent vector field through regression, with inference performed by integrating the learned system.Source: §1 · p. 1
- Off-manifold samples can be invalid outputs, require potentially problematic projection, or represent actions unsupported by expert demonstrations.In closed-loop control, adherence to expert action geometry can matter more for performance than low validation error (Pan et al., 2026).Source: §1 · pp. 1–2
- Conventional flow matching can have low regression loss yet drift from the manifold, while projection and geometric flow-matching remedies require a known manifold or manifold-supported prior.Linear conditional flow-matching paths do not necessarily stay near a curved target manifold.Source: §1 · p. 2, Figure 1 · p. 2
- MSFM combines learned transport along data geometry with prescribed convergence toward it, using analytical projectors for known manifolds and local PCA proxies for unknown geometry.Its contraction-based construction guarantees convergence from an ambient prior within a chosen time window.Source: §1 · p. 2
2 RELATED WORK
Prior work studies intrinsic manifold transport, ambient-space alternatives, contraction-based stability, and generative robotic policies. MSFM differs by combining prescribed normal contraction with learned tangential transport and allowing ambient priors and locally estimated geometry.
2 RELATED WORK
- Intrinsic flow-matching methods define paths and tangent fields on known manifolds, often starting from manifold-supported or manifold-adapted priors (Chen & Lipman, 2024; de Kruiff et al., 2025).Applications include robot policies and pose estimation (Braun et al., 2024; Ding et al., 2025).Source: §2 · pp. 2–3
- Ambient-space approaches include learned metrics, distribution-space dynamics, learned potentials, and geometry-respecting flow-matching variants (Kapu´sniak et al., 2024; Atanackovic et al., 2025; Balcerak et al., 2026; Bamberger et al., 2026).For linear-subspace data, Pi et al. (2026) use ambient Gaussian noise and analytical normal velocity; distributional bounds do not by themselves ensure exact manifold membership or a prescribed trajectory-wise contraction rate.Source: §2 · p. 3
- Contraction theory supports incremental and partial stability, and related constructions appear in structured neural dynamics, flow matching, and Riemannian policies (Ding et al., 2025).The cited work also includes developments in nonlinear control and differential growth analysis.Source: §2 · p. 3
- Generative robotic policies have explored manifold-valued poses, spatial equivariance, feasible-trajectory projection, and flow-matching action generation (Ryu et al., 2024; Bouvier et al., 2025; Chisari et al., 2025).MSFM instead permits ambient priors and uses analytical projectors or local PCA affine proxies without learning a global intrinsic parameterization.Source: §2 · p. 3
3 BACKGROUND AND PROBLEM FORMULATION
This section sets up conditional flow matching and the geometry used to measure distance from an embedded manifold. It identifies why ordinary velocity regression does not ensure stability and motivates enforcing contraction in the normal direction.
3 BACKGROUND AND PROBLEM FORMULATION
- Conditional flow matching samples source and target points and a time, constructs a conditional path and its velocity, then trains by velocity regression.At inference, a sample from the source distribution is transported by integrating the learned vector field.Source: §3 · p. 3
- The usual Euclidean interpolation is a straight path between source and target, while changing scalar interpolations affects both tangent and normal directions.Thus a generic path choice does not isolate motion relative to the data manifold.Source: §3 · p. 3
3.1 EMBEDDED-MANIFOLD GEOMETRY
- Within a tubular neighborhood of a smooth embedded manifold, each point has a unique nearest manifold point given by a smooth projection.This provides the local geometric setting for defining distance and normal displacement.Source: §3.1 · p. 3
- The offset r(x) = x − Π(x) splits a point into its projected base point and displacement, and its length measures distance to the manifold.The tangent and normal projectors at the base point separate the corresponding directions.Source: §3.1 · p. 4
- Along a generated trajectory, only the normal component of velocity changes the manifold distance; attraction requires that distance-change rate be negative off the manifold.A sufficient condition is a pointwise contraction inequality.Source: §3.1 · p. 4
- Flow-matching regression controls average velocity error on sampled training paths, so low loss does not enforce pointwise contraction or prevent inference-time drift.MSFM enforces the contraction condition by construction while learning tangential transport.Source: §3.1 · p. 4
4 MANIFOLD-STABLE FLOW MATCHING
MSFM pairs an inference-time field that learns motion along the manifold with a compatible training path. It constructs these components using known-manifold geometry or local affine proxies, and uses a contraction schedule to make normal displacement vanish near the terminal time.
4 MANIFOLD-STABLE FLOW MATCHING
- MSFM uses a sampling field that learns along-manifold motion and prescribes contraction toward the manifold, paired with a training path that agrees with it in the normal direction.The construction covers both known manifolds and unknown geometry represented by local affine proxies.Source: §4 · p. 4
4.1 KNOWN-MANIFOLD CONSTRUCTION
- For known manifolds, the learned field is tangent while prescribed normal feedback is −α(t)r(x), which vanishes on the manifold.The compatible path follows a base curve from the projected initial point to the target and evolves the offset in the changing normal space; for SO(3), an explicit offset expression is available.Source: §4.1 · p. 4
4.2 UNKNOWN MANIFOLD: LOCAL AFFINE PROXIES
- For a fixed proxy, local PCA supplies tangent bases and affine proxies; proxy selection uses a nearest anchor or minimum normal residual.The compatible path transports the tangent component linearly and contracts its normal component, with the target velocity carrying both terms.Source: §4.2 · p. 5
4.3 ALGORITHMS
- Training regresses tangential velocity and records the normal residual, while sampling integrates the learned field from the prior using current-state geometry.Unknown manifolds require local-proxy preprocessing; the training and sampling procedures apply to both geometry settings.Source: §4.3 · p. 5
4.4 CONTRACTION SCHEDULE
- The contraction schedule α(t) = α0 + β/(1 − t) yields normal displacement tending to zero at the terminal limit and supports numerical integration in log time.For a fixed base curve with α0 = 0, the normal path uniquely minimizes a time-weighted normal kinetic action; the full ambient path may still curve, and this is not an optimal-transport claim.Source: §4.4 · p. 6
- For a fixed affine proxy, α0 = 0 and β = 1 recover the usual linear conditional flow-matching path, a special case that does not extend to curved-manifold paths.The normal perturbation volume also tends to zero under the contracting dynamics.Source: §4.4 · p. 6
5 THEORETICAL GUARANTEES
The guarantees combine a parameter-independent normal contraction with learned tangential motion. They cover known manifolds and fixed or switching local proxies, with terminal distance bounds under the stated conditions.
5 THEORETICAL GUARANTEES
- The loss separates tangential velocity error from a parameter-independent normal residual, which vanishes when Qut = -α(t)r.The decomposition uses a tangent projector P and its normal complement Q.Source: §5 · p. 6
- For a known manifold, the prescribed normal feedback preserves on-manifold samples and drives off-manifold samples toward it, assuming trajectories remain in the smooth projection neighborhood.The result applies while solutions exist and remain in U.Source: §5 · p. 6
- For an unknown manifold, fixed local PCA proxies support convergence, and minimum-residual switching preserves convergence to the proxy set.The switching result applies when the selector and flow are well defined.Source: §5 · p. 6
- At numerical terminal time 1 − δ, distance to the manifold or proxy set is bounded by e^−α0(1−δ)δ^β times its initial distance.The supplied figure describes ellipse trajectories and terminal off-proxy errors for 700 samples.Source: §5 · p. 7, Figure 2 · p. 7
- Geometric stability is distinct from exact tangential-distribution matching; without normal path compatibility, the normal residual is a modeling tradeoff.The conventional flow-matching marginalization argument applies under normal path compatibility and exact population regression.Source: §5 · p. 7
6 EXPERIMENTS
The experiments assess geometric validity separately from task success across an ellipse with local proxies, disturbed Push-T control, and known-manifold Robomimic Square manipulation. The robotic tasks compare FM and MSFM under shared policy protocols.
6.1 ELLIPSE WITH LOCAL LINEAR PROXIES
- Across 700 ellipse samples, mean terminal off-proxy error is 1.10 × 10^-6 and mean analytic ellipse residual is 9.74 × 10^-5.The ellipse geometry is represented by one-dimensional local PCA proxies and is not supplied analytically to the model.Source: §6.1 · p. 7
6.2 PUSH-T
- Push-T success rises from 74% to 82%, and mean maximum reward from 0.86 to 0.92, under repeated block-position disturbances.Figure 3 also shows conventional FM has four-orders-larger off-manifold-proxy error.Source: Figure 3 · p. 8, §6.2 · p. 8
6.3 ROBOMIMIC SQUARE
- Robomimic Square success is 72% with MSFM versus 60% with FM, while mean rotation deviation before final projection falls from 10^-2 to 10^-7.The task constrains rotation blocks toward SO(3); translation and gripper components remain unconstrained.Source: §6.3 · p. 8
7 CONCLUSION
MSFM combines learned tangential transport with prescribed normal contraction, establishing manifold invariance and transverse convergence under stated assumptions. The experiments report geometric adherence and higher observed task performance, while the conclusion identifies evaluation at state-of-the-art scale as future work.
7 CONCLUSION
- Under stated geometric and solution assumptions, MSFM establishes manifold invariance and transverse convergence bounds, with suitable contraction schedules yielding vanishing normal deviation without exact tangential learning.The conclusion describes analytical projectors for known manifolds and local affine proxies for unknown geometry, plus compatible probability paths and an orthogonal loss decomposition.Source: §7 · p. 9
- The absolute Push-T and Square success rates are below state-of-the-art policy reports; whether geometric advantages persist at state-of-the-art scale remains future work.The controlled comparison uses compact networks, and Push-T's repeated-block-disturbance protocol differs from standard benchmarks, making absolute numbers not directly comparable.Source: §7 · p. 9
- Experiments show ellipse geometric adherence, higher observed success and reward under Push-T disturbances, and improved Square success with smaller rotation deviations.In Square, Figure 4 reports nearly five orders of magnitude lower pre-projection rotation deviation with MSFM and shows an MSFM completion where conventional FM failed.Source: §7 · pp. 9–10, Figure 4 · p. 9
AI USE STATEMENT
- The authors report using generative AI for language editing and final code polishing, and retain responsibility for the final content and numerical results.Source: AI USE STATEMENT · p. 10
REPRODUCIBILITY STATEMENT
- Appendices A–G document rotation specialization, path derivations, proofs, algorithms, and numerical-experiment data, optimization, checkpoint selection, and solvers.The reproducibility statement assigns these materials across the appendices.Source: REPRODUCIBILITY STATEMENT · p. 10
A ROTATION-MATRIX SPECIALIZATION
- For Square, the analytical SO(3) field applies to each rotation block, while translation and gripper components remain unconstrained.The training velocity is evaluated analytically, without an auxiliary numerical ODE solve for this geometry.Source: A ROTATION-MATRIX SPECIALIZATION · p. 14
B COMPATIBLE PROBABILITY PATHS IN TUBULAR COORDINATES
- Tubular paths account for changing normal spaces through the ˙Q_t r_t term, while their normal velocity matches the sampling field under the projection condition.The construction does not require the learned tangential velocity to equal the target; the projection condition remains a separate geometric assumption.Source: B COMPATIBLE PROBABILITY PATHS IN… · pp. 14–15
B.1 MINIMUM KINETIC ACTION OF THE NORMAL MOTION
For a fixed base path, the prescribed normal displacement uniquely minimizes a weighted normal-motion action, without optimizing tangential transport or the full generative path. The section also establishes an orthogonal tangent–normal loss decomposition.
B.1 MINIMUM KINETIC ACTION OF THE NORMAL MOTION
- The action measures normal motion only, excluding tangential velocity required by a changing normal space.Only for fixed affine geometry does the projected action also equal ordinary ambient kinetic action of the displacement.Source: §B.1 · p. 15
- The prescribed normal curve r_t is the unique minimizer of the weighted normal action for a fixed base path.The proof uses the displacement difference from any competing curve and the endpoint conditions.Source: §B.1 · pp. 15–16
- When β = 1, the minimizer recovers ordinary kinetic-action minimization.For the full schedule, the corresponding solution remains the unique minimizer; when α0 = 0, the action equals Eβ/β.Source: §B.1 · p. 16
- Normal potential decays at a bounded rate in log time, while moving normal directions and the base path can make the ambient trajectory curved.Tangential transport and normal contraction occur simultaneously, not in successive phases.Source: §B.1 · p. 16
- The variational result fixes the base path and minimizes only normal action; it neither selects tangential transport nor guarantees global minimum action across switching affine proxies.It therefore does not establish optimality of the full generative path.Source: §B.1 · p. 16
C.1 PROOF OF PROPOSITION 1
- Orthogonal tangent and normal regression terms yield a Pythagorean loss decomposition, with the normal term vanishing exactly when Q u_t = −α(t)r.The decomposition follows because the two terms lie in orthogonal spaces.Source: §C.1 · p. 16
C.2 PROOF OF THEOREM 1
The proofs establish manifold invariance and terminal convergence under the prescribed contraction. They extend distance bounds to the theorem’s proxy settings and the stated corollary.
C.2 PROOF OF THEOREM 1
- Theorem 1 establishes that the manifold M is invariant and that trajectories converge to it at the terminal time.The transverse bound implies terminal convergence because its integral diverges at t = 1.Source: §C.2 · p. 17
C.4 PROOF OF THEOREM 3
- Theorem 3 derives a distance bound using the active minimizing index when the distance is represented by a finite minimum of proxy potentials.The derivative of the minimum is bounded above by the derivative for an active index.Source: §C.4 · p. 17
C.5 PROOF OF COROLLARY 1
- The corollary’s bounds for M and the proxy sets follow by substituting the contraction scale into the earlier bounds and taking square roots.For proxy sets, the distance is expressed through the corresponding normal projector.Source: §C.5 · p. 17
D ALGORITHM AND IMPLEMENTATION DETAILS
The implementation uses local PCA proxies and separates tangential regression from any normal compatibility residual. The ellipse experiment reports small terminal off-proxy error and separately measures distance to the sampled arc and analytic ellipse.
D ALGORITHM AND IMPLEMENTATION DETAILS
- Local proxy preprocessing uses neighbor differences and PCA to form tangent projectors, with proxy rank fixed or selected by explained variance.The ellipse setup uses one-dimensional proxies and selects the active proxy by nearest anchor.Source: D ALGORITHM AND IMPLEMENTATION DETAILS · p. 17, E ADDITIONAL ELLIPSE DETAILS · p. 18
- Training draws source, target, and time samples independently of network parameters, giving identical parameter gradients for full-field regression and its tangential term.Any nonzero normal residual must still be reported.Source: D ALGORITHM AND IMPLEMENTATION DETAILS · p. 17
- Sampling transforms time using τ = −log(1 − t), and the normal-feedback coefficient lies between β and α0 + β.Stopping at τ = −log δ corresponds to t = 1 − δ.Source: D ALGORITHM AND IMPLEMENTATION DETAILS · p. 18
- For a fixed affine proxy containing x1, the linear path is compatible when α0 = 0 and β = 1; this calculation does not automatically extend to changing proxies or curved manifolds.Other schedule parameters produce a normal residual when the same linear path is retained.Source: D ALGORITHM AND IMPLEMENTATION DETAILS · p. 18
E ADDITIONAL ELLIPSE DETAILS
- Across 700 generated samples, mean terminal off-proxy error is 1.10 × 10^-6 and mean absolute analytic ellipse residual is 9.74 × 10^-5.Mean nearest-data distance is 2.75 × 10^-3, distinguishing attraction to extended affine lines from closeness to the sampled arc.Source: E ADDITIONAL ELLIPSE DETAILS · p. 18
F PUSH-T EXPERIMENTAL DETAILS
In disturbed Push-T rollouts, MSFM records higher success and mean maximum reward than FM. The experimental details describe episode-disjoint data, a locally selected PCA proxy field, and rollout evaluation under repeated perturbations.
F PUSH-T EXPERIMENTAL DETAILS
- The episode-disjoint train, validation, and test splits contain 19,232, 2,469, and 2,507 complete horizons, respectively.The 206 demonstrations are split into 164 training, 21 validation, and 21 held-out test episodes.Source: F PUSH-T EXPERIMENTAL DETAILS · p. 18
- The proxy field uses joint condition-and-horizon distance to select 64 neighbors, with PCA retaining at least 95% of squared singular-value energy.The median retained dimension is three, and the active proxy is selected without access to the target action.Source: F PUSH-T EXPERIMENTAL DETAILS · p. 18
- Because the selected proxy need not contain the target action, the linear path has a model-independent transverse residual in addition to the tangential loss.The reported MSFM selected-checkpoint validation loss is approximately 0.173 after subtracting that residual, versus 31.520 before subtraction.Source: F PUSH-T EXPERIMENTAL DETAILS · p. 19
- Rollout plans start from standard Gaussian noise, use 20 RK4 steps, and execute only the first generated target before replanning.FM integrates in t while MSFM uses s = −log(1 − t), so the integration grids differ despite matched field-evaluation counts.Source: F PUSH-T EXPERIMENTAL DETAILS · p. 19
- Under the tested disturbance protocol, success counts are 41/50 for MSFM and 37/50 for FM, with mean maximum rewards of 0.922 and 0.862, respectively.Success is defined as maximum environment reward of at least 0.95; evaluation uses 50 initial-state seeds and repeated block translations.Source: F PUSH-T EXPERIMENTAL DETAILS · p. 19
G ROBOMIMIC SQUARE EXPERIMENTAL DETAILS
In the Robomimic Square evaluation, MSFM records higher success and lower rotational deviation than FM under paired rollouts from the same 50 simulator states. The experiment uses shared data and network settings, with distinct probability paths and sampling endpoints.
Data, conditioning, and controller
- The dataset contains 200 proficient-human demonstrations, split into 180 training and 20 validation episodes, with 27,213 and 2,941 horizons, respectively.There is no separate offline test partition; targets reconstruct intended absolute controller goals from scaled incremental commands.Source: G ROBOMIMIC SQUARE EXPERIMENTAL DETAILS · p. 19
Geometry and probability paths
- Both methods learn in R64×13, but FM uses linear interpolation from Gaussian noise while MSFM uses a tubular rotational path with an analytical tangent projector and prescribed normal contraction.The saved MSFM run sets α0 = 0 and β = 1; its tubular compatibility statement applies within the smooth projection neighborhood.Source: G ROBOMIMIC SQUARE EXPERIMENTAL DETAILS · p. 19
Optimization and checkpoint selection
- Both temporal U-Nets have 859,347 parameters and train for 25,000 updates; rollout checkpoints are selected by integrated validation score over the first four actions.Training uses identical initial parameters and the same optimization schedule for both networks.Source: G ROBOMIMIC SQUARE EXPERIMENTAL DETAILS · p. 20
Sampling and paired evaluation
- Both policies use RK4 and execute the first four of 64 generated actions before replanning, with evaluation on the same bank of 50 initial states.FM integrates to t = 1, while MSFM integrates to 1 − 10^-10; both apply one final SO(3) projection before command conversion.Source: G ROBOMIMIC SQUARE EXPERIMENTAL DETAILS · p. 20
Results and scope
- MSFM succeeds on 36/50 episodes versus FM’s 30/50, while mean rotational deviation is 6.40 × 10^-7 versus 0.0451.Deviation is measured before execution projection and averaged over the generated horizon, replanning times, and episodes; MSFM has lower deviation in all 50 pairs.Source: G ROBOMIMIC SQUARE EXPERIMENTAL DETAILS · p. 20
- In paired outcomes, 25 episodes succeed with both policies, five only with FM, and eleven only with MSFM.The deviation metric measures rotational validity, not accuracy relative to the expert action.Source: G ROBOMIMIC SQUARE EXPERIMENTAL DETAILS · p. 20