Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

Metric Flow Matching for Smooth Interpolations on the Data Manifold

About

Matching objectives underpin the success of modern generative models and rely on constructing conditional paths that transform a source distribution into a target distribution. Despite being a fundamental building block, conditional paths have been designed principally under the assumption of Euclidean geometry, resulting in straight interpolations. However, this can be particularly restrictive for tasks such as trajectory inference, where straight paths might lie outside the data manifold, thus failing to capture the underlying dynamics giving rise to the observed marginals. In this paper, we propose Metric Flow Matching (MFM), a novel simulation-free framework for conditional flow matching where interpolants are approximate geodesics learned by minimizing the kinetic energy of a data-induced Riemannian metric. This way, the generative model matches vector fields on the data manifold, which corresponds to lower uncertainty and more meaningful interpolations. We prescribe general metrics to instantiate MFM, independent of the task, and test it on a suite of challenging problems including LiDAR navigation, unpaired image translation, and modeling cellular dynamics. We observe that MFM outperforms the Euclidean baselines, particularly achieving SOTA on single-cell trajectory prediction.

Kacper Kapu\'sniak, Peter Potaptchik, Teodora Reu, Leo Zhang, Alexander Tong, Michael Bronstein, Avishek Joey Bose, Francesco Di Giovanni• 2024

Related benchmarks

TaskDatasetResultRank
Trajectory InterpolationEB 50D (held-out time points)
Mean W17.753
24
Reconstruction of discrete branching dynamics5D EB t=1
W10.449
24
Reconstruction of discrete branching dynamicsCITE 50D (t=3)
W110.621
24
Trajectory InferenceEB dataset 5D (test)
W1 (t=1)0.449
23
Population Dynamics InterpolationEB scRNA 5-dim PCA representation (leave-one-out)
W1 Distance0.713
21
Continuous-Time Dynamics EstimationSynthetic Arch first snapshot as initial state
L_DTW10.86
20
Continuous-Time Dynamics EstimationSynthetic Y-shaped first snapshot as initial state
L_DTW20.86
20
Trajectory InterpolationLung Tumor
W11.984
18
Trajectory InterpolationLight V1
W1 Error2.401
18
Trajectory InterpolationDendritic Stimulus
W1 Distance/Error3.714
18
Showing 10 of 99 rows
...

Other info

Follow for update