Pullback Flow Matching on Data Manifolds

TL;DR

Pullback Flow Matching (PFM) introduces a new framework for generative modeling on data manifolds, leveraging pullback geometry and isometric learning to improve manifold learning and sample generation.

Contribution

PFM is the first method to use pullback geometry and isometric learning for scalable, efficient generative modeling on data manifolds with designable latent spaces.

Findings

Enhanced manifold learning and interpolation in latent space.

Successful application to protein data for generating novel proteins.

Improved generative performance over existing methods.

Abstract

We propose Pullback Flow Matching (PFM), a novel framework for generative modeling on data manifolds. Unlike existing methods that assume or learn restrictive closed-form manifold mappings for training Riemannian Flow Matching (RFM) models, PFM leverages pullback geometry and isometric learning to preserve the underlying manifold's geometry while enabling efficient generation and precise interpolation in latent space. This approach not only facilitates closed-form mappings on the data manifold but also allows for designable latent spaces, using assumed metrics on both data and latent manifolds. By enhancing isometric learning through Neural ODEs and proposing a scalable training objective, we achieve a latent space more suitable for interpolation, leading to improved manifold learning and generative performance. We demonstrate PFM's effectiveness through applications in synthetic data, protein dynamics and protein sequence data, generating novel proteins with specific properties. This method shows strong potential for drug discovery and materials science, where generating novel samples with specific properties is of great interest.