Pullback Flow Matching on Data Manifolds

2024-10-06

arXiv.org

AI Summary

Existing flow matching methods rely on predefined or restrictive closed-form manifold mappings, struggling to simultaneously preserve manifold geometry and ensure generative flexibility. This paper proposes Pullback Flow Matching (PFM), the first flow matching framework incorporating pullback geometry: it jointly models the pullback metric and isometric constraints to learn continuous neural ODE flows that intrinsically respect manifold structure. PFM enables customizable latent spaces and analytically tractable manifold mappings, eliminating dependence on prior manifold parameterizations. Evaluated on synthetic data, protein dynamics modeling, and sequence generation, PFM significantly improves latent interpolation quality and generation fidelity. Notably, it successfully designs novel, functionally validated proteins—demonstrating practical utility for drug discovery.

Technology Category

• Machine Learning: Learning with Manifolds
• Search and Optimization: Sampling/Simulation-based Search
• Natural Language Processing: Generation

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Generative modeling on data manifolds with geometry preservation
• Efficient generation and precise interpolation in latent space
• Designable latent spaces using assumed metrics for improved learning

Innovation

• Leverages pullback geometry for manifold learning
• Uses Neural ODEs for scalable isometric learning
• Enables designable latent spaces with assumed metrics