Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Energy Guided Geometric Flow Matching

Aaron Zweig, Mingxuan Zhang, Elham Azizi, David Knowles

TL;DR

This work addresses the challenge of learning trajectories that stay on the data manifold in temporal data. It introduces Energy Guided Geometric Flow Matching (EGGFM), which combines denoising score matching with annealed energy distillation to learn a data-driven metric tensor, then derives geodesics and a manifold-aware flow to realize geometry-respecting interpolations. The approach uses iterative density refinement, stratified sampling, and a metric-based geodesic loss to produce geometry-aware trajectories, demonstrating improvements over baselines on synthetic manifolds and single-cell RNA datasets (EB and CITE). The results suggest that explicitly incorporating manifold geometry into flow matching can improve interpolation fidelity and could enhance downstream analyses of developmental and disease trajectories.

Abstract

A useful inductive bias for temporal data is that trajectories should stay close to the data manifold. Traditional flow matching relies on straight conditional paths, and flow matching methods which learn geodesics rely on RBF kernels or nearest neighbor graphs that suffer from the curse of dimensionality. We propose to use score matching and annealed energy distillation to learn a metric tensor that faithfully captures the underlying data geometry and informs more accurate flows. We demonstrate the efficacy of this strategy on synthetic manifolds with analytic geodesics, and interpolation of cell

Energy Guided Geometric Flow Matching

TL;DR

This work addresses the challenge of learning trajectories that stay on the data manifold in temporal data. It introduces Energy Guided Geometric Flow Matching (EGGFM), which combines denoising score matching with annealed energy distillation to learn a data-driven metric tensor, then derives geodesics and a manifold-aware flow to realize geometry-respecting interpolations. The approach uses iterative density refinement, stratified sampling, and a metric-based geodesic loss to produce geometry-aware trajectories, demonstrating improvements over baselines on synthetic manifolds and single-cell RNA datasets (EB and CITE). The results suggest that explicitly incorporating manifold geometry into flow matching can improve interpolation fidelity and could enhance downstream analyses of developmental and disease trajectories.

Abstract

A useful inductive bias for temporal data is that trajectories should stay close to the data manifold. Traditional flow matching relies on straight conditional paths, and flow matching methods which learn geodesics rely on RBF kernels or nearest neighbor graphs that suffer from the curse of dimensionality. We propose to use score matching and annealed energy distillation to learn a metric tensor that faithfully captures the underlying data geometry and informs more accurate flows. We demonstrate the efficacy of this strategy on synthetic manifolds with analytic geodesics, and interpolation of cell

Paper Structure

This paper contains 29 sections, 1 theorem, 31 equations, 2 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Setup
  3. Denoising Score Matching
  4. Conditional Flow Matching
  5. Stochastic Interpolants along Geodesics
  6. Related Work
  7. Flow Matching
  8. Energy distillation from Score Matching
  9. Learning Geodesics along Data-driven Manifolds
  10. Methods
  11. Score and Energy Matching
  12. Iterative Density Refinement
  13. Self normalized importance sampling
  14. Density annealing
  15. Stratified Sampling
  16. ...and 14 more sections

Key Result

Theorem 4.1

Assume for each cluster $C_j$, $r_i > 0$ for all $i \in C_j$ and the importance weights have finite second moment under per-cluster measure $q_j$, as number of samples $m_j$ in $C_j$ goes to infinity we have Additionally, if $\mathcal{L}(x; \theta)$ is the denoising score matching loss, the objective converges to $\mathbb{E}_{p^\star}[\mathcal{L}(x;\theta)]$

Figures (2)

  • Figure 1: A visualization of our method. We fit an initial score and energy (left), then use annealing and self-normalized importance sampling to fit an updated energy that is less biased by unequal density in the data (middle), and we clip the energy to calculate a balanced metric tensor that captures the manifold (right).
  • Figure 2: The negative energy of EGGFM (left) and negative metric tensor of MFM (right) on the 10 dimensional sphere, projected onto the first two dimensions.

Theorems & Definitions (2)

  • Theorem 4.1
  • proof