Pullback Flow Matching on Data Manifolds

## Abstract

Abstract We propose \ac PFM, a novel framework for generative modeling on data manifolds. Unlike existing methods that assume or learn restrictive closed-form manifold mappings for training \ac RFM models, \ac 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 \ac 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.

## 1 Introduction

Figure: Figure 1: An example of isometric learning, where the goal is to create a latent space that allows for interpolation. The shortest paths (in red) on the data manifold $\mathcal{D}\subset\mathbb{R}^{3}$ correspond to the shortest paths on the latent manifold $\mathcal{M}$.
Refer to caption: x1.png

Since the rise of machine learning in the scientific domain, researchers have focused on the derivation of bigger models trained on increasingly large datasets, e.g., to solve weather forecasting (Bodnar et al., ) or protein structure and function generation (Hayes et al., ).
Relying on such scaling laws is not feasible in many areas of science where data is limited and precise modeling of physical phenomena is crucial. In research where data collection is challenging, accurate interpolation between data points is essential for making reliable predictions and generating realistic representations of complex systems. However, current methods lack the mathematical foundations to accurately interpolate in the latent space (Arvanitidis et al., ) and do not capture the geometric structure of the data (Wessels et al., ). Our goal is to develop mappings that enable precise interpolation in the latent space and facilitate efficient and accurate generation on data manifolds, ultimately advancing the ability to model complex physical phenomena with limited data.

We consider modeling the data under the manifold hypothesis, which states that high-dimensional data lies on a lower dimensional manifold. This has been successfully applied to several downstream tasks in various fields across the scientific domain (Vanderplas & Connolly, ; Dsilva et al., ; Noé & Clementi, ). Modeling the data in its intrinsic dimension allows for efficient analysis Diepeveen et al. () and generation Rombach et al. (). Furthermore, accurately capturing the geometry of the data manifold in the learning problem has shown to improve several down-stream tasks such as clustering (Ghojogh et al., ), classification (Kaya & Bilge, ; Hauberg et al., ) and generation (Arvanitidis et al., ; Sun et al., ).

One way to achieve a latent manifold that supports interpolation is to have a structured Riemannian geometry, such as the one from pullback geometry, which provides closed-form manifold mappings (Diepeveen et al., ). This requires constructing an invertible and differentiable mapping—diffeomorphism—between the data manifold and the latent manifold. Interpolation on manifolds is then performed through geodesics, shortest paths, and thus to achieve our goal we require geodesics on the data manifold to match geodesics on the latent manifold. This motivates our consideration of isometries, that is, metric-preserving diffeomorphisms $\varphi$. These mappings preserve the distances of points on the data manifold on the latent manifold, and thereby ensure proper interpolation.

Related-Work. In the literature, low-dimensional generation and generation on manifolds have typically been addressed as separate problems. Low-dimensional approaches, such as latent diffusion (Rombach et al., ) or latent flow-matching (Dao et al., ), often overlook the geometric structure of the data, leading to inaccuracies in tasks requiring a faithful representation of the underlying manifold. Conversely, manifold generation methods either assume geodesics on the data manifold for simulation-free training (Chen & Lipman, )—an approach flawed when closed-form mappings are unavailable—or attempt to learn a metric that forces the generative trajectories to have data support (Kapusniak et al., ).

The pullback framework, which involves learning a diffeomorphism for isometric learning, but also presents several challenges. For example, in works such as Cuzzolin (); Gruffaz et al. (); Lebanon (), the learning problems are tailored to specific tasks, limiting their generality and preventing the learning of isometries for broader data manifolds. In contrast, Diepeveen () offers a more general approach capable of learning the data manifold directly from a dataset, showing promising results. However, their learning problem and network architecture suffer from scalability issues, which restrict their applicability to larger high-dimensional datasets.

Our approach bridges these gaps by modeling data on a lower-dimensional latent manifold with known geometry through a diffeomorphisms parameterized and trained in a scalable way. By doing so we preserve the intrinsic properties of the data manifold and enable accurate and efficient generation through simulation-free training.

Contributions. We propose \acPFM, a novel framework for latent manifold learning and generation through isometries. The benefit of this method is that it respects the geometry of the data manifold, even without closed-form manifold mappings. Second, learning can be performed in the intrinsic dimension of the data manifold resulting in efficient and effective learning of the generative model with fewer parameters. Building on Diepeveen (), we leverage pullback geometry to define a new metric on the entire ambient space, $\mathbb{R}^{d}$, by learning an isometry $\varphi$ that preserves the geometric structure of the data manifold $\mathcal{D}$ on the latent manifold $\mathcal{M}$. We use the corresponding metric of the assumed latent manifold $\mathcal{M}$ to perform \acRFM on the latent manifold that supports interpolation. Our contributions are as follows:

- 1.
We introduce \acPFM, a novel framework that enables accurate and efficient data generation on manifolds. \AcPFM leverages the pullback geometry to preserve the underlying geometric structure of the data manifold within the latent space, facilitating precise interpolation and generation.
- 2.
We improve the parameterization of diffeomorphisms, used to learn isometries, in both expressiveness and training efficiency through \acsNODEs.
- 3.
We introduce a scalable and stable isometric learning objective. This objective relies solely on a distance measure on the data manifold, simplifying the training process compared to Diepeveen () while maintaining geometric fidelity.
- 4.
We demonstrate our methods’ effectiveness through experiments on synthetic data, high-dimensional molecular dynamics data, and experimental peptide sequences. Our framework utilizes designable latent spaces to generate novel proteins with specific properties closely matching reference samples, advancing generative modeling techniques in drug discovery and materials science.

## 2 Notation

We give a brief summary of the notation used in the paper, and give a more extensive background on Riemannian and pullback geometry in Appendix A.

A manifold $\mathcal{M}$ is a topological space that locally resembles Euclidean space. A $d$-dimensional manifold $\mathcal{M}$ around a point $\bm{p}\in\mathcal{M}$ is described by a chart $\psi:U\to\mathbb{R}^{d}$, where $U\subseteq\mathcal{M}$ is a neighborhood of $\bm{p}$. The chart provides a local coordinate system for the manifold. The tangent space at a point $\bm{p}\in\mathcal{M}$, denoted $\mathcal{T}_{\bm{p}}\mathcal{M}$, is the vector space of all tangent vectors at that point.

A smooth manifold $\mathcal{M}$ equipped with a Riemannian metric is called a Riemannian manifold and is denoted by $(\mathcal{M},(\cdot,\cdot)^{\mathcal{M}})$. The Riemannian metric $(\cdot,\cdot)^{\mathcal{M}}$ is a smoothly varying inner product defined on the tangent spaces $\mathcal{T}_{\bm{p}}\mathcal{M}$ for all points $\bm{p}\in\mathcal{M}$, and it defines lengths and angles on the manifold. A geodesic, $\gamma_{\bm{p},\bm{q}}(t)$ is the shortest path between two points $\bm{p},\bm{q}\in\mathcal{M}$, generalizing the notion of a straight line in Euclidean space.

The exponential map $\exp_{\bm{p}}:\mathcal{T}_{\bm{p}}\mathcal{M}\to\mathcal{M}$ maps a tangent vector $\Xi_{\bm{p}}$ to a point on the manifold by following the geodesic in the direction of $\Xi_{\bm{p}}$ starting from $\bm{p}$. The inverse of the exponential map is the logarithmic map, denoted by $\log_{\bm{p}}:\mathcal{M}\to\mathcal{T}_{\bm{p}}\mathcal{M}$, which returns the tangent vector corresponding to a given point on the manifold.

In this work, we consider a $d$-dimensional Riemannian manifold $\big{(}\mathcal{M},(\cdot,\cdot)^{\mathcal{M}}\big{)}$, and a smooth diffeomorphism $\varphi:\mathbb{R}^{d}\rightarrow\mathcal{M}$, such that $\varphi\big{(}\mathbb{R}^{d}\big{)}\subseteq\mathcal{M}$ is geodesically convex, meaning that any pair of points within this subset are connected by a unique geodesic. This mapping allows us to pullback the geometric structure of $\mathcal{M}$ to $\mathbb{R}^{d}$ by defining the pullback metric on $\mathbb{R}^{d}$. Specifically, for tangent vectors $\Xi,\Phi\in\mathcal{T}\mathbb{R}^{d}$, the pullback metric is defined as

$$ $(\Xi,\Phi)^{\varphi}:=\big{(}\varphi_{\ast}[\Xi],\varphi_{\ast}[\Phi]\big{)}^{ \mathcal{M}},$ (1) $$

where $\varphi_{\ast}$ is the pushforward of tangent vectors under $\varphi$. Through this construction, various geometric objects in $\mathcal{M}$, such as distances and geodesics, can be expressed in terms of their counterparts in $\mathbb{R}^{d}$ with respect to the pullback metric. The distance function $d^{\varphi}_{\mathbb{R}^{d}}:\mathbb{R}^{d}\times\mathbb{R}^{d}\rightarrow
\mathbb{R}$ on $\mathbb{R}^{d}$ with the pullback metric is given by,

$$ $d^{\varphi}_{\mathbb{R}^{d}}(\bm{x}_{i},\bm{x}_{j})=d_{\mathcal{M}}\big{(} \varphi(\bm{x}_{i}),\varphi(\bm{x}_{j})\big{)},$ (2) $$

where $d_{\mathcal{M}}$ denotes the Riemannian distance on $\mathcal{M}$. The length-minimizing geodesic connecting $\bm{x}_{i}$ and $\bm{x}_{j}$ in $\mathbb{R}^{d}$ with respect to the pullback metric $\gamma^{\varphi}_{\bm{x}_{i},\bm{x}_{j}}:[0,1]\rightarrow\mathbb{R}^{d}$ is given by,

$$ $\gamma^{\varphi}_{\bm{x}_{i},\bm{x}_{j}}(t)=\varphi^{-1}\big{(}\gamma^{ \mathcal{M}}_{\varphi(\bm{x}_{i}),\varphi(\bm{x}_{j})}(t)\big{)},$ (3) $$

here $\gamma^{\mathcal{M}}$ denotes the geodesic in $\mathcal{M}$ connecting $\varphi(\bm{x}_{i})$ and $\varphi(\bm{x}_{j})$. This enables computation of geodesics and distances in $\mathbb{R}^{d}$ using the geometry of $\mathcal{M}$, as stated in Prop. 2.1 of Diepeveen ().

Figure: Figure 2: Isometric learning for coarse-grained protein conformation data of adynalate kinase. We define a new metric $(\cdot,\cdot)^{\varphi}$ on the entire ambient space, $\mathbb{R}^{d}$ ($d=214\times 3$), by learning a diffeomorphism $\varphi:\mathbb{R}^{d}\rightarrow\mathcal{M}$ that preserves a locally Euclidean metric $(\cdot,\cdot)^{\mathcal{D}}$ on the latent manifold $\mathcal{M}=\mathcal{M}_{d^{\prime}}\times\mathbb{R}^{d-d^{\prime}}$ for $d^{\prime}=1$.
Refer to caption: x3.png

## 3 Pullback Flow Matching

We propose \acfPFM, a novel framework for generative modeling on data manifolds using pullback geometry. Our goal is to transform samples from a simple distribution $\bm{x}_{0}\sim p$ on the data manifold $\mathcal{D}$ into a complex target distribution $\bm{x}_{1}\sim q$, also on $\mathcal{D}$. Ideally, we would perform this transformation using \acfRFM, see Appendix A for a summary, on $\big{(}\mathcal{D},(\cdot,\cdot)^{\mathcal{D}}\big{)}$ by optimizing the objective from Chen & Lipman (),

$$ $\mathcal{L}_{RFM}(\bm{\eta})=\mathbb{E}_{t,q(\bm{x}_{1}),p(\bm{x}_{0})}\left\| v_{t}\Big{(}\gamma^{\mathcal{D}}_{\bm{x}_{1},\bm{x}_{0}}\big{(}\kappa(t)\big{) };\bm{\eta}\Big{)}-\dot{\gamma}^{\mathcal{D}}_{\bm{x}_{1},\bm{x}_{0}}\big{(} \kappa(t)\big{)}\right\|_{\mathcal{D}}^{2},$ (4) $$

where $\|\cdot\|_{\mathcal{D}}$ is the norm induced by the metric $(\cdot,\cdot)^{\mathcal{D}}$ on the manifold, and $\dot{\gamma}^{\mathcal{D}}_{\bm{x}_{1},\bm{x}_{0}}\big{(}\kappa(t)\big{)}$ denotes the time-derivative of the geodesic between $\bm{x}_{0}$ and $\bm{x}_{1}$ on $\mathcal{D}$. Solving this objective is intractable due to a lack of closed-form manifold mappings. Existing methods address this by learning or assuming restricted manifold mappings for simulation-free \acRFM training Chen & Lipman (); Kapusniak et al. (). We overcome this limitation by using pullback geometry to derive expressive closed-form manifold mappings, enabling flexible, simulation-free \acRFM training.

Following Diepeveen (), we define a new metric on the ambient space $\mathbb{R}^{d}$ using the pullback metric. We assume a learned isometry $\varphi_{\bm{\theta}}$ that approximates geodesics $\gamma^{\varphi_{\bm{\theta}}}$ on $\big{(}\mathbb{R}^{d},(\cdot,\cdot)^{\varphi_{\bm{\theta}}}\big{)}$ to those $\gamma^{\mathcal{D}}$ on $\big{(}\mathcal{D},(\cdot,\cdot)^{\mathcal{D}}\big{)}$. Rewriting the \acRFM objective under the pullback framework yields the objective,

$$ $\mathcal{L}_{PFM}(\bm{\eta})=\mathbb{E}_{t,q(\bm{x}_{1}),p(\bm{x}_{0})}\left\| v_{t}\Big{(}\gamma^{\varphi_{\bm{\theta}}}_{\bm{x}_{1},\bm{x}_{0}}\big{(} \kappa(t)\big{)};\bm{\eta}\Big{)}-\dot{\gamma}^{\varphi_{\bm{\theta}}}_{\bm{x} _{1},\bm{x}_{0}}\big{(}\kappa(t)\big{)}\right\|_{\varphi_{\bm{\theta}}}^{2},$ (5) $$

where $\|\cdot\|_{\varphi_{\bm{\theta}}}$ is the norm induced by the pullback metric $(\cdot,\cdot)^{\varphi_{\bm{\theta}}}$. By applying Equation 3, we can reformulate the \acPFM objective in terms of manifold mappings on $\mathcal{M}$,

$$ $\mathcal{L}_{PFM}(\bm{\eta})=\mathbb{E}_{t,q(\bm{x}_{1}),p(\bm{x}_{0})}\left\| v_{t}\Big{(}\gamma^{\mathcal{M}}_{\varphi_{\bm{\theta}}(\bm{x}_{1}),\varphi_{ \bm{\theta}}(\bm{x}_{0})}\big{(}\kappa(t)\big{)};\bm{\eta}\Big{)}-\dot{\gamma} ^{\mathcal{M}}_{\varphi_{\bm{\theta}}(\bm{x}_{1}),\varphi_{\bm{\theta}}(\bm{x} _{0})}\big{(}\kappa(t)\big{)}\right\|_{\mathcal{M}}^{2},$ (6) $$

where $\|\cdot\|_{\mathcal{M}}$ is the norm induced by the metric on $\mathcal{M}$. Assuming a latent manifold $\mathcal{M}$ with closed-form mappings allows for simulation-free training. For efficiency, we model the $d$-dimensional latent manifold as a product manifold, $\mathcal{M}=\mathcal{M}_{d^{\prime}}\times\mathbb{R}^{d-d^{\prime}}$, where $\mathcal{M}_{d^{\prime}}\subset\mathcal{M}$. By encoding samples close to the submanifold $\mathcal{M}_{d^{\prime}}$, isometric learning ensures geodesics $\mathcal{M}_{d^{\prime}}$ closely match geodesics on $\mathcal{M}$. As a result, we formulate the $d^{\prime}$-\acsPFM objective,

$$ $\mathcal{L}_{d^{\prime}-PFM}(\bm{\eta})=\mathbb{E}_{t,q(\bm{x}_{1}),p(\bm{x}_{ 0})}\left\|v_{t}\Big{(}\gamma^{\mathcal{M}_{d^{\prime}}}_{\varphi_{\bm{\theta} }(\bm{x}_{1}),\varphi_{\bm{\theta}}(\bm{x}_{0})}\big{(}\kappa(t)\big{)};\bm{ \eta}\Big{)}-\dot{\gamma}^{\mathcal{M}_{d^{\prime}}}_{\varphi_{\bm{\theta}}( \bm{x}_{1}),\varphi_{\bm{\theta}}(\bm{x}_{0})}\big{(}\kappa(t)\big{)}\right\|_ {\mathcal{M}_{d^{\prime}}}^{2},$ (7) $$

where $\|\cdot\|_{\mathcal{M}_{d^{\prime}}}$ denotes the norm on the submanifold $\mathcal{M}_{d^{\prime}}$. The $d^{\prime}$-\acsPFM objective offers two key benefits. First, defining the objective on the submanifold $\mathcal{M}_{d^{\prime}}$ results in computational speed-ups during training. Second, the known geometry on the submanifold simplifies the training dynamics of the vector field $v_{t}(\cdot;\bm{\eta})$, requiring fewer epochs and parameters $\bm{\eta}$.

Figure: Figure 3: Trajectories of \acCNF (left) trained with \acCFM, \acPFM and $1$-\acsPFM objectives on the ARCH dataset compared to the data manifold $\mathcal{D}$ (right). At $t=0$ the trajectory starts with a standard normal distribution in the data space for \acsCFM and latent submanifold for ($1$-)\acsPFM mapped back to the data space through $\varphi^{-1}_{\bm{\theta}}$.
Refer to caption: x4.png

## 4 Learning Isometries

The motivation for learning isometries $\varphi_{\bm{\theta}}$—metric-preserving diffeomorphisms—is to enable a latent (sub)manifold that supports interpolation with closed-form geometric mappings, facilitating simulation-free training of \acPFM. Building on the framework of Diepeveen (), summarized in Appendix A, we propose a more expressive parameterization of learnable diffeomorphisms $\varphi_{\bm{\theta}}$ through \acpNODE and enhance the objective for scalable isometric learning on data manifolds.

### 4.1 Parameterizing Diffeomorphisms

We parameterize diffeomorphisms, invertible and differentiable functions between two manifolds, specifically $\varphi:\mathbb{R}^{d}\rightarrow\mathcal{M}$. In practice, we construct the latent manifold as a product manifold, $\mathcal{M}=\mathcal{M}_{d^{\prime}}\times\mathbb{R}^{d-d^{\prime}}$ and the diffeomorphism $\varphi$ as,

$$ $\varphi:=[\psi^{-1},\bm{I}_{d-d^{\prime}}]\circ\phi\circ T_{\bm{\mu}},$ (8) $$

where $\psi:U\rightarrow\mathbb{R}^{d^{\prime}}$ a chart on a (geodesically convex) subset $U\subset\mathcal{M}_{d^{\prime}}$ of the $d^{\prime}$-dimensional latent submanifold $\big{(}\mathcal{M}_{d^{\prime}},(\cdot,\cdot)_{\mathcal{M}_{d^{\prime}}}\big{)}$, $\phi:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ a diffeomorphism and $T_{\bm{\mu}}(\bm{x})=\bm{x}-\bm{\mu}$. We choose this construction following Diepeveen () because the manifold hypothesis translates to assuming the data manifold is homeomorphic to $\mathcal{M}_{d^{\prime}}$. In such case, the rest of the latent manifold should be mapped close to zero, e.g. $\varphi(\bm{x}_{i})$ is close to $\mathcal{M}_{d^{\prime}}\times\bm{0}^{d-d^{\prime}}$ in terms of the metric on $\mathcal{M}$.

We generate the diffeomorphism $\phi$ by solving a \acNODE (Chen et al., ). The advantage of this approach is threefold, i) this parameterization of diffeomorphisms is more expressive and efficient to train compared to Invertible Residual Networks (Behrmann et al., ) as chosen by Diepeveen (), ii) based on some mild technical assumptions a \acNODE can be proven to generate proper diffeomorphisms, see Appendix B for the proof, and iii) numerically the accuracy and invertibility of the generated flow can be controlled through smaller step-sizes and higher-order solvers.

To define the diffeomorphism $\phi_{\bm{\theta}}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$, we start with the \acNODE governing the flow:

$$ $\frac{d\bm{z}(t)}{dt}=f(\bm{z}(t);\bm{\theta}),$ (9) $$

where $f:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ is a vector field parameterized by a \acMLP with Swish activation functions and a sine-cosine time embedding and $\bm{\theta}$ denotes the parameters of the \acMLP. Given an initial condition $\bm{z}(0)=\bm{x}$, the solution to this \acNODE is:

$$ $\phi_{\bm{\theta}}(\bm{x}):=\bm{x}+\int_{0}^{1}f(\bm{z}(t);\bm{\theta})\,dt.$ (10) $$

To obtain the inverse $\phi_{\bm{\theta}}^{-1}$ one has to integrate the differential equation backwards in time with initial condition $\bm{z}(1)$. To solve the \acNODE we implemented a Runge-Kutta solver in JAX, see Appendix E for further architectural and training related details.

### 4.2 Learning Objective

The primary objectives in learning isometries are twofold, i) to map the data manifold $\big{(}\mathcal{D},(\cdot,\cdot)^{\mathcal{D}}\big{)}$ into a low-dimensional geodesic subspace of $\big{(}\mathcal{M},(\cdot,\cdot)^{\mathcal{M}}\big{)}$, specifically $\mathcal{M}_{d^{\prime}}\subset\mathcal{M}$, and ii) to preserve local isometry around the data, as motivated by Proposition 2.1 and Theorems 3.4, 3.6, and 3.8 from Diepeveen ().

To achieve these goals, we build upon the training objective proposed by Diepeveen (), as summarized in Appendix A. We use the global isometry loss and submanifold loss to map the data manifold $\mathcal{D}$ into the lower dimensional geodesic subspace $\mathcal{M}_{d^{\prime}}$. We improve the objective by introducing graph matching loss for isometric learning to enforce global isometry between the data and latent (sub)manifolds (Zhu et al., ). Intuitively we can explain this loss as enforcing each sample to be equally isometric to all the other samples.

The original objective is limited by its reliance on the pullback metric’s Riemannian metric tensor $(\cdot,\cdot)^{\varphi}$ to enforce local isometry, which is (computationally) intractable in practice and scales poorly in higher dimensions. We use stability regularization from Finlay et al. () to impose local isometry of the resulting flow, resulting in a more scalable training objective,

$$ $\displaystyle\mathcal{L}(\bm{\theta})$ $\displaystyle=\alpha_{1}\frac{1}{n^{2}}\sum^{n}_{i=1}\sum^{n}_{j=1}\|d^{ \varphi_{\bm{\theta}}}_{\mathbb{R}^{d}}(\bm{x}_{i},\bm{x}_{j})-d_{i,j}\|^{2}$ ( global isometry loss ) $\displaystyle+\alpha_{2}\frac{1}{n}\sum^{n}_{i=1}\sum_{j\neq i}\|(\bm{d}^{ \varphi_{\bm{\theta}}}_{\mathbb{R}^{d}}(\bm{x}_{i},\bm{x}_{\cdot})-\bm{d}^{ \varphi_{\bm{\theta}}}_{\mathbb{R}^{d}}(\bm{x}_{j},\bm{x}_{\cdot}))-(\bm{d}_{i ,\cdot}-\bm{d}_{j,\cdot})\|^{2}$ ( graph matching loss ) $\displaystyle+\alpha_{3}\frac{1}{n}\sum_{i=1}^{n}\left\|\begin{bmatrix}\bm{I}_ {d-d^{\prime}}&\emptyset\\ \emptyset&\bm{0}_{d^{\prime}}\end{bmatrix}(\phi_{\bm{\theta}}\circ T_{\bm{z}}) (\bm{x}_{i})\right\|_{1}$ ( submanifold loss ) $\displaystyle+\alpha_{4}\frac{1}{n}\sum_{i=1}^{n}\int_{0}^{1}\|\bm{\varepsilon }^{T}\nabla f_{\bm{\theta}}(\bm{z}_{i}(t))\|^{2}\;dt.$ ( stability regularization ) $$

Here, $\bm{\varepsilon}\sim\mathcal{N}(\bm{0},\bm{I})$ and $\bm{d}^{\varphi_{\bm{\theta}}}_{\mathbb{R}^{d}}(\bm{x}_{i},\bm{x}_{\cdot})$ and $\bm{d}_{i,\cdot}$ denote the columns of the distance matrices induces by $(\cdot,\cdot)^{\varphi}$ and $(\cdot,\cdot)^{\mathcal{D}}$. The benefit of this formulation is that it only requires approximating geodesic distances $d_{i,j}$ on the data manifold $\mathcal{D}$, without needing to calculate or differentiate the metric tensor. In section 5, we demonstrate the effectiveness of the new regularization terms through an ablation study on synthetic and high-dimensional simulated protein dynamics trajectories.

## 5 Experiments

The goal of this paper is to learn interpolatable latent (sub)manifolds for generation on data manifolds. We achieve this through isometric learning in the framework of pullback geometry. In this section we validate our methods on synthetic, simulated and experimental datasets, for full descriptions see Appendix D. For details on the training procedure and hyperparameter settings we refer the reader to Appendix E.

We begin our experiments with an ablation study of graph matching loss and stability regularization, demonstrating the benefits of including both terms for learning isometries. Second, we compare (latent) interpolation methods with interpolation on the latent manifold $\mathcal{M}$, $(\cdot,\cdot)^{\mathcal{M}}$-interpolation, and on the latent submanifold $\mathcal{M}_{d^{\prime}}$, $(\cdot,\cdot)^{\mathcal{M}_{d^{\prime}}}$-interpolation. We demonstrate that we can accurately interpolate on the data manifold by interpolating on the latent (sub)manifold (^1^11In these experiments we do not report interpolation through the \acRAE by Diepeveen () due to the intractability of the training objective for the higher-dimensional datasets.). Third, we validate \acPFM as a generative model on data manifolds and discuss how sample generation is improved by generating on the submanifold $\mathcal{M}_{d^{\prime}}$. Finally, we inspect the designability of the latent manifold through the choice of metric $(\cdot,\cdot)^{\mathcal{D}}$ in the task of small protein design.

### 5.1 Ablation Study

The goal of the ablation study is to evaluate the effectiveness of the reformulated objective function for learning isometries. To this end, we perform an ablation study for both the graph matching loss and stability regularization on a synthetic \acsARCH dataset ($n=500$, $d=2$) in the spirit of Tong et al. () and a coarse-grained protein dynamics datasets of \acI-FABP ($n=500$, $d=131\times 3$). We report three metrics on the validation set of 20 % of the data, invertibility $\varepsilon_{inv}=\frac{1}{n}\sum_{i=1}^{n}\|\bm{x}_{i}-\varphi_{\bm{\theta}}^
{-1}\big{(}\varphi_{\bm{\theta}}(\bm{x}_{i})\big{)}\|^{2}$, low-dimensionality $\varepsilon_{ld}=\frac{1}{n}\sum_{i=1}^{n}\left\|\begin{bmatrix}\bm{I}_{d-d^{
\prime}}&\emptyset\\
\emptyset&\bm{0}_{d^{\prime}}\end{bmatrix}\phi_{\bm{\theta}}(\bm{x}_{i})\right
\|^{2}_{1}$ and isometry $\varepsilon_{iso}=\frac{1}{n^{2}}\sum^{n}_{i=1}\sum^{n}_{j=1}\|d_{i,j}-d_{
\mathcal{M}}\big{(}\varphi(\bm{x}_{i}),\varphi(\bm{x}_{j})\big{)}\|^{2}$.

Result. Table 1 demonstrates that incorporating both the graph matching loss and stability regularization improves the invertibility and isometry metrics across both datasets, with the combined approach yielding both a low $\varepsilon_{inv}$ and $\varepsilon_{iso}$ values, indicating enhanced model performance in preserving the geometry of the data in the synthetic dataset as well ass the more noisy and high dimensional simulated dataset.

**Table 1: Ablation study of isometric learning for \acsARCH dataset and \acI-FABP protein dynamics datasets for graph matching loss (GM) and stability regularization (Stability). In both cases we choose $\mathcal{M}_{d^{\prime}}=\mathbb{R}$. We report invertibility ($\downarrow$), low-dimensionality ($\downarrow$) and isometry ($\downarrow$). The distance $(\cdot,\cdot)^{\mathcal{D}}$ we preserve on the latent manifold is a locally Euclidean distance based on Isomap Tenenbaum et al. ().**
|  | \acsARCH | \acsI-FABP |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Model | Invertibility | Low-Dimensionality | Isometry | Invertibility | Low-Dimensionality | Isometry |
| None | $8.628\cdot 10^{-4}$ | $\bm{2.956\cdot 10^{-4}}$ | $2.841\cdot 10^{-3}$ | $2.513\cdot 10^{0}$ | $1.591\cdot 10^{-1}$ | $5.568\cdot 10^{-2}$ |
| GM | $5.377\cdot 10^{-5}$ | $3.020\cdot 10^{-4}$ | $\bm{1.577\cdot 10^{-3}}$ | $1.264\cdot 10^{0}$ | $\bm{1.550\cdot 10^{-1}}$ | $2.302\cdot 10^{-2}$ |
| Stability | $\bm{9.677\cdot 10^{-11}}$ | $1.409\cdot 10^{-2}$ | $2.690\cdot 10^{-3}$ | $2.357\cdot 10^{-7}$ | $1.674\cdot 10^{-1}$ | $4.342\cdot 10^{-3}$ |
| Both | $1.519\cdot 10^{-10}$ | $1.325\cdot 10^{-2}$ | $2.175\cdot 10^{-3}$ | $\bm{8.657\cdot 10^{-8}}$ | $1.653\cdot 10^{-1}$ | $\bm{2.558\cdot 10^{-3}}$ |

### 5.2 Interpolation Experiments

The goal of isometric learning is to learn an interpolatable latent (sub)manifold of the data manifold with closed-form manifold mappings. To evaluate whether interpolation on the latent (sub)manifold accurately reflects interpolation on the data manifold, we conduct an interpolation experiment using the synthetic \acsARCH dataset, as well as the molecular dynamics datasets of \acAK ($n=100$, $d=214\times 3$) and \acI-FABP. In both cases we choose $\mathcal{M}_{d^{\prime}}=\mathbb{R}$, see Appendix C for guidance on latent manifold and metric selection. We approximate the metric on the data manifold $(\cdot,\cdot)^{\mathcal{D}}$ through the length of Isomap’s geodesics Tenenbaum et al. (), see Figure 4 for an example. We compare the accuracy of $100$ randomly selected geodesics between points in the test set for \acpVAE Kingma & Welling (), $\beta-$VAEs Higgins et al. (), $(\cdot,\cdot)^{\mathcal{M}}$-interpolation and $(\cdot,\cdot)^{\mathcal{M}_{d^{\prime}}}$-interpolation.

Figure: Figure 4: Example of $(\cdot,\cdot)^{\mathcal{M}_{d^{\prime}}}$-interpolation for \acsARCH dataset in red. In blue the dataset $\{\bm{x}_{i}\}^{n}_{i=1}$, black the true submanifold $\mathcal{M}_{d^{\prime}}$, the half circle, and in orange the Isomap geodesic between orange points.

**Table 2: \AcRMSE ($\downarrow$) of $100$ random Isomap geodesics for various (latent) interpolation methods.**
| Interpolation | Latent | \acsARCH | \acsAK | \acsI-FABP |
| --- | --- | --- | --- | --- |
| Linear | ✗ | $0.089_{\pm 0.097}$ | $0.360_{\pm 0.167}$ | $0.388_{\pm 0.079}$ |
| \acsVAE | ✓ | $0.462_{\pm 0.144}$ | $0.241_{\pm 0.081}$ | $0.319_{\pm 0.045}$ |
| $\beta$-\acsVAE | ✓ | $0.486_{\pm 0.130}$ | $0.263_{\pm 0.091}$ | $0.376_{\pm 0.053}$ |
| $(\cdot,\cdot)^{\mathcal{M}}$ | ✓ | $\bm{0.056}_{\pm 0.039}$ | $0.310_{\pm 0.127}$ | $0.364_{\pm 0.056}$ |
| $(\cdot,\cdot)^{\mathcal{M}_{d^{\prime}}}$ | ✓ | $0.069_{\pm\bm{0.027}}$ | $\bm{0.206_{\pm 0.022}}$ | $\bm{0.298_{\pm 0.016}}$ |

Result. The $(\cdot,\cdot)^{\mathcal{M}_{d^{\prime}}}$- and $(\cdot,\cdot)^{\mathcal{M}}$-interpolation achieves superior interpolation accuracy with lower \acRMSE variablity compared to other models, indicating more robust and reliable interpolation. $(\cdot,\cdot)^{\mathcal{M}_{d^{\prime}}}$-interpolation specifically demonstrates improvements over other methods in the more stochastic and seemingly higher dimensional \acAK ($d=639$) and \acI-FABP ($d=642$) datasets. This improvement suggests that compressing the latent representation into a lower-dimensional space reduces the noise while accurately capturing the underlying data manifold. Our findings demonstrate that accurate interpolation of protein dynamics trajectories of \acAK and \acI-FABP can be achieved using a single-dimensional latent manifold. This method shows promise for improving protein dynamics simulations, ultimately advancing understanding of protein dynamics.

### 5.3 Generation Experiments

We evaluate the effectiveness of our proposed method \acPFM for generation on data manifolds $\mathcal{D}$. We train two \acpPFM, one using the latent manifold $\mathcal{M}$ and one using the lower dimensional latent submanifold $\mathcal{M}_{d^{\prime}}$, named \acPFM and $d^{\prime}$-\acsPFM respectively. Additionally, we train a \acCFM model on the raw data as a comparison. A visual example of the learned generative flows over time for the \acsARCH dataset can be viewed in Figure 3.

Result. Figure 3 we see that the learned isometry to the latent manifold $\mathcal{M}$ acts as a strong manifold prior, capturing the manifold structure at the start of the \acCNF trajectory ($t=0.0$). Additionally, the learned isometry to the latent submanifold $\mathcal{M}_{d^{\prime}}$ captures the noiseless manifold revealing the underlying manifold used to generate the data. Through this strong (noiseless) manifold prior, we see that both \acPFM and $1$-\acsPFM approximate the distribution on the manifold earlier in the trajectory and better.

To evaluate the performance of our generative methods on manifolds, we use the $1$-\acNN accuracy metric Lopez-Paz & Oquab (). This metric assesses how well the generated point clouds match the reference point clouds. We classify each point cloud by finding its nearest neighbor from the combined set of generated and reference point clouds of the trainingset. The $1$-\acsNN accuracy measures the proportion of correct classifications of the testset, indicating how similar the generated point clouds are to the reference set. An accuracy close to $50\%$ suggests that the generated and reference point clouds come from the same distribution, and thus indistinguishable, reflecting successful learning of the target distribution.

**Table 3: Evaluation of generative model performance across dimensionality of (latent) (sub)manifold ($\downarrow$), number of model parameters, denoted by # pars ($\downarrow$), and $1$-\acsNN accuracy ($1$-\acsNN$\rightarrow 0.5$). The $1$-\acsNN metric measures the generative quality, with values closer to $0.5$ indicating better performance.**
|  | \acsARCH | \acsAK | \acsI-FABP |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | $\bm{\dim}$ | # pars | $\bm{1}$-\acsNN | $\bm{\dim}$ | # pars | $\bm{1}$-\acsNN | $\bm{\dim}$ | # pars | $\bm{1}$-\acsNN |
| \acsCFM | $2$ | $38402$ | $0.240$ | $642$ | $4560155$ | $0.386$ | $393$ | $1715081$ | $0.340$ |
| \acsPFM | $2$ | 38402 | $0.265$ | $642$ | $4560155$ | $0.336$ | $393$ | $1715081$ | $0.390$ |
| $1$-\acsPFM | $\bm{1}$ | $\bm{2657}$ | $\bm{0.380}$ | $\bm{1}$ | $\bm{38273}$ | $\bm{0.467}$ | $\bm{1}$ | $\bm{38273}$ | $\bm{0.505}$ |

Results. Table 3 highlights the effectiveness of the $1$-\acsPFM model in generative tasks. The $1$-\acsPFM model leverages the lower-dimensional isometric latent manifold $\mathcal{M}_{d^{\prime}}$, significantly reducing the number of parameters required. Training the $1$-\acsPFM is significantly faster due to the reduction in parameters and the dimensionality of the training samples. The $1$-\acsNN accuracy for $1$-\acsPFM approaches the ideal $0.5$ across all datasets, indicating that this model better captures the underlying distribution on the data manifold compared to CFM and \acPFM.

### 5.4 Designable Latent Manifolds for Novel Protein Engineering

The goal of these experiments is to design a latent manifold that captures biologically relevant properties of protein sequences, enabling the generation of novel proteins with specific characteristics. By leveraging our method’s flexibility in defining the metric on the data manifold $(\cdot,\cdot)^{\mathcal{D}}$, we structure the latent space such that it captures protein properties, such as sequence similarity, hydrophobicity, hydrophobic moment, charge, and isoelectric point.

To achieve this, we use protein sequences of up to 25 amino acids from the \acGRAMPA dataset (see Appendix D for details). We construct the following custom metric on the data manifold,

$$ $\displaystyle d_{\mathcal{D}}(x_{i},x_{j})$ $\displaystyle=d_{\text{Levenshtein}}(x_{i},x_{j})+d_{\text{hydrophobicity}}(x_ {i},x_{j})$ (11) $\displaystyle+d_{\text{hydrophobic moment}}(x_{i},x_{j})+d_{\text{charge}}(x_{ i},x_{j})$ (12) $\displaystyle+d_{\text{isoelectric point}}(x_{i},x_{j}),$ (13) $$

where the Levenshtein distance measures the number of single-character edits (insertions, deletions, or substitutions) required to transform one sequence into another.

For the remaining four properties—hydrophobicity, hydrophobic moment, charge, and isoelectric point—distances are computed using the difference in property values between sequences. Specifically, for each property, we define the (pseudo)distance as,

$$ $d_{\text{[property]}}(x_{i},x_{j})=|f_{\text{property}}(x_{i})-f_{\text{ property}}(x_{j})|.$ (14) $$

These (pseudo)distances are standardized by dividing by the maximum observed distance in the training data. Since the Levenshtein distance is a proper metric, we ensure that the combined distance $d_{\mathcal{D}}(x_{i},x_{j})$ remains a valid distance metric.

We use the designed metric $(\cdot,\cdot)^{\mathcal{D}}$ on the space of protein sequences with at most 25 amino acids in the \acGRAMPA dataset to learn an isometry that preserves this metric on the latent manifold $\mathcal{M}$ and latent submanifold $\mathcal{M}_{d^{\prime}}$.

To generate protein sequences with specific properties, we sample from a normal distribution around the data points in the latent manifold $\bm{z}\in\mathcal{M}$ or latent submanifold $\bm{z}\in\mathcal{M}_{d^{\prime}}$. The variability of this sampling process is aligned with the latent variability of the training data $\sigma_{\bm{z}_{train}}$, scaled by a temperature factor $\tau$, resulting in the following expression,

$$ $\displaystyle\bm{z}^{(analogue)}_{i}$ $\displaystyle=\bm{z}_{i}+\tau\mathcal{N}(\bm{0},\sigma_{\bm{z}_{train}}\bm{I}) ,\text{ and}$ (15) $\displaystyle\bm{x}^{(analogue)}_{i}$ $\displaystyle=\varphi_{\bm{\theta}}^{-1}(\bm{z}_{i})\text{ for }i=1,\ldots,n_{ test}.$ (16) $$

This sampling methodology is referred to as analogue generation, as it does not involve explicitly learning the distribution over the latent manifold. Instead, it generates novel sequences by sampling around existing data points on the latent (sub)manifold of the test set.

We apply this process to both the latent manifold $\mathcal{M}$ and its submanifold $\mathcal{M}_{d^{\prime}}$. To evaluate the effectiveness of the generated sequences, we measure the number of unique sequences that were not present in the original dataset and compare the properties of the generated samples to the properties of their base points. For further specifics on hyperparameters and training procedures, refer to Appendix E.

**Table 4: Unique protein sequences generated via analogue generation on the latent manifold $\mathcal{M}$ and its submanifold $\mathcal{M}_{d^{\prime}}$ at various temperatures ($\tau$). The table presents the total sequences generated (Total), those already in the dataset (In Data), and the number of novel sequences (Novel). We perform a \acKS test at a $5\%$ significance level to compare novel sequences with their base points. Non-significant \acKS values are shown as X/Y, where X is the number of non-significant properties and Y is the total properties tested.**
|  | $\mathcal{M}$ | $\mathcal{M}_{d^{\prime}}$ |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $\bm{\tau}$ | Total | In Data | Novel | Non-Sign. \acsKS | Total | In Data | Novel | Non-Sign. \acsKS |
| 0.01 | 689 | 652 | 37 | 5/5 | 687 | 5 | 682 | 2/5 |
| 0.05 | 689 | 103 | 586 | 5/5 | 689 | 4 | 685 | 2/5 |
| 0.1 | 689 | 35 | 654 | 5/5 | 689 | 4 | 685 | 2/5 |
| 0.2 | 689 | 12 | 677 | 2/5 | 689 | 0 | 689 | 2/5 |
| 0.5 | 689 | 1 | 688 | 1/5 | 689 | 0 | 689 | 1/5 |
| 1 | 689 | 0 | 689 | 0/5 | 689 | 0 | 689 | 0/5 |

Results. The application of our designed latent manifold facilitated the generation of diverse novel protein sequences, demonstrating the effectiveness of the analogue generation methodology. As illustrated in Table 4, increasing the temperature parameter, $\tau$, directly influenced the diversity of generated sequences. At lower temperatures ($\tau\leq 0.1$), many unique sequences emerged while maintaining similarity to their base points, as indicated by non-significant KS test values. Conversely, higher temperatures ($\tau>0.1$) resulted in a significant increase in novel sequences, accompanied by significant KS values suggesting greater divergence from base sequences. This observation supports the hypothesis that novel sequences generated close to the base points are structurally similar, highlighting the effectiveness of isometric learning in structuring the latent space. Overall, our results indicate that temperature manipulation can strategically balance novelty and similarity, paving the way for innovative applications in protein engineering.

Figure: Figure 5: Latent interpolation between a protein with a high hydrophobic moment and low charge and a protein with a low hydrophobic moment and high charge.
Refer to caption: extracted/5904060/figures/5_experiments_latent_interpolation_protein_design.png

Latent interpolation experiments, illustrated in Figure 5, further demonstrate the potential of our approach. By interpolating between sequences with contrasting properties, we revealed a smooth transition of characteristics within the latent space, reinforcing our method’s capability to fine-tune specific protein attributes. This smooth transition indicates that our latent manifold can be effectively navigated to explore a continuum of properties such as hydrophobicity, hydrophobic moment, charge, and isoelectric point, which are essential for determining protein solubility, stability, and interaction behavior. This capability allows for the targeted design of protein sequences that could be optimized for specific biochemical contexts, potentially enhancing their performance in applications like enzyme catalysis or therapeutic development. In summary, the efficacy of our designed latent manifold not only expands the repertoire of available protein sequences but also ensures retention of biologically relevant properties, positioning this approach as a valuable tool for precision in protein engineering.

## 6 Conclusion

We introduce \acfPFM, a novel framework for simulation-free training of generative models on data manifolds. By leveraging pullback geometry and isometric learning, \acPFM allows for closed-form mappings on data manifolds while enabling precise interpolation and efficient generation. We demonstrated the effectiveness of \acPFM through applications in synthetic protein dynamics and small protein generation, showcasing its potential in generating novel, property-specific samples through designable latent spaces. This approach holds significant promise for advancing generative modeling in fields like drug discovery and materials science, where precise and efficient sample generation is critical.

## Appendix A Background

To achieve an interpolatable latent manifold we take a Riemannian geometric perspective. We start by introducing the notation and key concepts of differential and Riemannian geometry, for a formal description see Lee (). Second, we explain prior work on \acpRAE Diepeveen (), a framework for constructing interpolatable latent manifolds. Third, we summarize \acCFM for generative modeling Lipman et al. (), a scalable way to train generative models in a simulation-free manner. Finally, we discuss how \acRFM Chen & Lipman () generalize \acCFM to Riemannian manifolds.

### A.1 Riemmanian Geometry

A $d$-dimensional smooth manifold $\mathcal{M}$ is a topological space that locally resembles $\mathbb{R}^{d}$, such that for each point $\bm{p}\in\mathcal{M}$, there exists a neighborhood $U$ of $\bm{p}$ and a homeomorphism $\psi:U\to\mathbb{R}^{d}$, called a chart. Then the tangent space $\mathcal{T}_{\bm{p}}\mathcal{M}$ at a point $\bm{p}\in\mathcal{M}$ is a vector space consisting of the tangent vectors at $\bm{p}$ representing the space of derivations at $\bm{p}$.

A Riemannian manifold $\big{(}\mathcal{M},(\cdot,\cdot)^{\mathcal{M}}\big{)}$ is a smooth manifold $\mathcal{M}$ equipped with a Riemannian metric $(\cdot,\cdot)^{\mathcal{M}}$, which is a smoothly varying positive-definite inner product on the tangent space $\mathcal{T}_{\bm{p}}\mathcal{M}$ at each point $\bm{p}$. The Riemannian metric $(\cdot,\cdot)^{\mathcal{M}}$ defines the length of tangent vectors and the angle between them, thereby inducing a natural notion of distance on $\mathcal{M}$ based on the lengths of tangent vectors along curves between two points.

The shortest path between two points on $\mathcal{M}$ is called a geodesic, which generalizes the concept of straight lines in Euclidean space to curved manifolds. Geodesics on Riemannian manifold are found by minimizing

$$ $E(\gamma)=\frac{1}{2}\int_{0}^{1}\big{(}\dot{\gamma}(t),\dot{\gamma}(t)\big{)}dt,$ (17) $$

whereas

$$ $L(\gamma)=\int_{0}^{1}\sqrt{\big{(}\dot{\gamma}(t),\dot{\gamma}(t)\big{)}}dt$ (18) $$

defines the distance between two points on the manifold. The exponential map,

$$ $\exp_{\bm{p}}:\mathcal{T}_{\bm{p}}\mathcal{M}\to\mathcal{M},$ (19) $$

at $\bm{p}$ maps a tangent vector $\Xi_{\bm{p}}\in\mathcal{T}_{\bm{p}}\mathcal{M}$ to a point on $\mathcal{M}$ reached by traveling along the geodesic starting at $\bm{p}$ in the direction of $\Xi_{\bm{p}}$ for unit time. The logarithmic map,

$$ $\log_{\bm{p}}:\mathcal{M}\to\mathcal{T}_{\bm{p}}\mathcal{M},$ (20) $$

is the inverse of the exponential map, mapping a point $\bm{q}\in\mathcal{M}$ back to the tangent space $\mathcal{T}_{\bm{p}}\mathcal{M}$ at $\bm{p}$.

These names, ’exponential’ and ’logarithmic’ map, are geometric extensions of familiar calculus concepts. Just as the exponential function maps a number to a point on a curve, the exponential map on a manifold maps a direction and starting point to a location along a geodesic. Similarly, the logarithm in calculus reverses exponentiation, and the logarithmic map on a manifold reverses the exponential map, returning the original direction and distance needed to reach a specified point along the geodesic.

Assume $\big{(}\mathcal{M},(\cdot,\cdot)^{\mathcal{M}}\big{)}$ is a $d$-dimensional Riemannian manifold and a smooth diffeomorphism $\varphi:\mathbb{R}^{d}\rightarrow\mathcal{M}$, such that $\varphi(\mathbb{R}^{d})\subseteq\mathcal{M}$ is geodesically convex, i.e., geodesics are uniquely defined on $\varphi(\mathbb{R}^{d})$. We can then define the pushforward $\varphi_{\ast}:\mathcal{T}\mathbb{R}^{d}\to\mathcal{T}\mathcal{M}$ which in turn can be used to define the pullback metric as

$$ $(\Xi,\Phi)^{\varphi}:=\big{(}\varphi_{\ast}[\Xi],\varphi_{\ast}[\Phi]\big{)},$ (21) $$

for tangent vectors $\Xi$ and $\Phi$. These mappings allow us to define all relevant geometric mappings in $\mathbb{R}^{d}$ in terms of manifold mappings on $\mathcal{M}$, see e.g. Proposition 2.1 of Diepeveen ():

- 1.
Distances $d^{\varphi}_{\mathbb{R}^{d}}:\mathbb{R}^{d}\times\mathbb{R}^{d}\rightarrow
\mathbb{R}$ on $\big{(}\mathbb{R}^{d},(\cdot,\cdot)^{\varphi}\big{)}$ are given by,
$d^{\varphi}_{\mathbb{R}^{d}}(\bm{x_{i}},\bm{x_{j}})=d_{\mathcal{M}}\big{(}
\varphi(\bm{x}_{i}),\varphi(\bm{x}_{j})\big{)},$
(22)
- 2.
Length-minimizing geodesics $\gamma^{\varphi}_{\bm{x}_{i},\bm{x}_{j}}:[0,1]\rightarrow\mathbb{R}^{d}$ on $\big{(}\mathbb{R}^{d},(\cdot,\cdot)^{\varphi}\big{)}$ are given by,
$\gamma^{\varphi}_{\bm{x}_{i},\bm{x}_{j}}(t)=\varphi^{-1}\big{(}\gamma^{
\mathcal{M}}_{\varphi(\bm{x}_{i}),\varphi(\bm{x}_{j})}(t)\big{)}$
(23)
- 3.
Logarithmic maps $\log^{\varphi}_{\bm{x_{i}}}:\mathbb{R}^{d}\rightarrow\mathcal{T}_{\bm{x}_{i}}
\mathbb{R}^{d}$ on $\big{(}\mathbb{R}^{d},(\cdot,\cdot)^{\varphi}\big{)}$ are given by,
$\log^{\varphi}_{\bm{x}_{i}}(\bm{x}_{j})=\varphi_{\ast}^{-1}\left[\log^{
\mathcal{M}}_{\varphi(\bm{x}_{i})}\big{(}\varphi(\bm{x}_{j})\big{)}\right]$
(24)
- 4.
Exponential maps $\exp^{\varphi}_{\bm{x}_{i}}:\mathcal{G}_{\bm{x}_{i}}\rightarrow\mathbb{R}^{d}$ for $\mathcal{G}_{\bm{x}_{i}}:=\log^{\varphi}_{\bm{x}_{i}}(\mathbb{R}^{d})\subset
\mathcal{T}_{\bm{x}_{i}}\mathbb{R}^{d}$ on $\big{(}\mathbb{R}^{d},(\cdot,\cdot)^{\varphi}\big{)}$ are given by
$\exp^{\varphi}_{\bm{x}_{i}}(\Xi_{\bm{x}_{i}})=\varphi^{-1}\big{(}\exp^{
\mathcal{M}}_{\varphi(\bm{x}_{i})}\left(\varphi_{\ast}[\Xi_{\bm{x}_{i}}]\right
)\big{)}$
(25)

A visual example of pullback geometry is given in Figure 6. Pullback geometry allows us to remetrize all of space $\mathbb{R}^{d}$, including the data manifold $\mathcal{D}\subset\mathbb{R}^{d}$, through the pullback metric. We can use it to define geometric mappings on $\big{(}\mathbb{R}^{d},(\cdot,\cdot)^{\varphi}\big{)}$, including geodesics (see Equation 23), through geometric mappings on the latent manifold $\mathcal{M}$. Next, we summarize work on Riemannian Auto-Encoders, that leverage pullback geometry to create an interpolatable latent manifold.

Figure: Figure 6: Example of pullback geometry for $\varphi:\mathbb{R}^{d}\rightarrow\mathcal{M}$ with $\mathcal{M}=\mathcal{M}_{d^{\prime}}\times\mathbb{R}^{d-d^{\prime}}$ for $\mathcal{M}_{d^{\prime}}=\mathbb{R}^{d^{\prime}}$, $d=3$ and $d^{\prime}=2$. Samples $\varphi(\bm{x}_{i})$ are close to elements of $\mathcal{M}_{d^{\prime}}\times\bm{0}^{d-d^{\prime}}$.
Refer to caption: x6.png

### A.2 Riemannian Auto-Encoder

The goal of \acpRAE is to create a interpolatable latent representation of the data. This is achieved through data-driven (pullback) Riemannian geometry, encoding the data onto a latent manifold with known geometry. The benefit of this, is that interpolation on the data manifold corresponds to interpolation on the latent manifold. Resulting in a more interpretable latent space compared to traditional auto-encoders.

Similar as in Diepeveen (), we define a \acRAE as a Riemannian Encoder $\text{RE}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{r}$ and Riemannian Decoder $\text{RD}:\mathbb{R}^{r}\rightarrow\mathbb{R}^{d}$,

$$ $\displaystyle\text{RAE}(\bm{x})$ $\displaystyle:=(RD\circ RE)(\bm{x})\quad s.t.,$ (26) $\displaystyle\text{RE}(\bm{x})_{k}$ $\displaystyle:=(\log^{\varphi}_{\bm{z}}(\bm{x}),\bm{v}_{\bm{z}}^{k})^{\varphi} _{\bm{z}}\text{ for }k=1,\dots r,$ (27) $\displaystyle\text{RD}(\bm{a})$ $\displaystyle:=\exp^{\varphi}_{\bm{z}}\left(\sum_{k=1}^{r}\bm{a}_{k}\bm{v}_{ \bm{z}}^{k}\right)$ (28) $$

where $\bm{z}$ denotes a base point and $(\cdot,\cdot)^{\varphi}_{\bm{z}}$ the pullback metric at $\bm{z}$. Furthermore,

$$ $\bm{v}_{\bm{z}}^{k}:=\sum_{l=1}^{d}\bm{W}_{lk}\Phi^{l}_{\bm{z}},$ (29) $$

represents the basis vectors of the latent space in the tangent space $T_{\bm{z}}\mathbb{R}^{d}$. Let $\Phi_{\bm{z}}^{l}\in T_{\bm{z}}\mathbb{R}^{d}$ be an orthonormal basis in the tangent space at $\bm{z}$ with respect to $(\cdot,\cdot)^{\varphi}_{\bm{z}}$ and define

$$ $\bm{X}_{i,l}=\big{(}\log_{\bm{z}}^{\varphi}(\bm{x}^{i}),\Phi^{l}_{\bm{z}}\big{ )}_{\bm{z}}^{\varphi}\text{ for }i=1,\dots,n\text{ and }l=1,\dots,d.$ (30) $$

We can compute $\bm{W}$ through a Singular Value Decomposition (SVD) of $\bm{X}$.

$$ $\bm{X}=\bm{U\Sigma W}^{T},$ (31) $$

where $U\in\mathbb{R}^{N\times R}$, $\Sigma=\text{diag}(\sigma_{1},\dots,\sigma_{R})\in\mathbb{R}^{R\times R}$ with $\sigma_{1}\geq\dots\geq\sigma_{R}$, $W\in\mathbb{R}^{d\times R}$ and where $R:=\text{rank}(\bm{X})$. The first $r$ columns of $\bm{W}$, corresponding to the largest singular values, are selected to form the matrix $\bm{W}\in\mathbb{R}^{d\times r}$. This parameter $r$ allows one to set the dimensionality of the latent representation of the \acRAE, if $r=d$ then the \acRAE reduces to $\text{RAE}(\bm{x})=\exp^{\varphi}_{\bm{z}}\big{(}\log^{\varphi}_{\bm{z}}(\bm{x
})\big{)}$.

To learn a \acRAE, one needs to first construct a diffeomorphism and define an objective function. In Diepeveen () diffeomorphisms are constructed by,

$$ $\varphi:=[\psi^{-1},\bm{I}_{d-d^{\prime}}]\circ\phi\circ\bm{O}\circ T_{\bm{z}},$ (32) $$

where $\psi:U\rightarrow\mathbb{R}^{d^{\prime}}$ is a chart on a (geodesically convex) subset $U\subset\mathcal{M}^{d^{\prime}}$ of a $d^{\prime}$-dimensional Riemannian manifold $(\mathcal{M}^{d^{\prime}},(\cdot,\cdot)_{\mathcal{M}_{d}^{\prime}})$, $\phi:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ is a real-valued diffeomorphism, $\bm{O}\in\mathbb{O}(d)$ is an orthogonal matrix, and $T_{\bm{z}}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ is given by $T_{\bm{z}}(\bm{x})=\bm{x}-\bm{z}$. The learnable diffeomorphism $\varphi:=\varphi_{\bm{\theta}}$ is constructed through parameterizing $\phi:=\phi_{\bm{\theta}}$ by an invertible residual network Behrmann et al. ().

### A.3 Learning Isometries with Riemannian Auto-Encoders

After constructing the diffeomorphism and Riemannian Auto-Encoder, one can learn an isometry by find the parameters $\bm{\theta}$ of $\varphi_{\theta}$ in Diepeveen () through minimizing the objective,

$$ $\displaystyle\mathcal{L}(\bm{\theta})$ $\displaystyle=\frac{1}{N(N-1)}\sum_{i,j=1,i\neq j}^{N(N-1)}\big{(}d_{\mathbb{R }^{d}}^{\varphi_{\bm{\theta}}}(\bm{x}_{i},\bm{x}_{j})-d_{i,j}\big{)}^{2}$ ( global isometry loss ) $\displaystyle+\alpha_{\text{sub}}\frac{1}{N}\sum_{i=1}^{N}\left\|\begin{ bmatrix}\bm{I}_{d-d^{\prime}}&\emptyset\\ \emptyset&\bm{0}_{d^{\prime}}\end{bmatrix}(\phi_{\bm{\theta}}\circ\bm{O}\circ T _{\bm{z}})(\bm{x}_{i})\right\|_{1}$ ( submanifold loss ) $\displaystyle+\alpha_{\text{iso}}\frac{1}{N}\sum_{i=1}^{N}\left\|\left(\left( \bm{e}^{j},\bm{e}^{j^{\prime}}\right)^{\varphi_{\bm{\theta}}}_{\bm{x}_{i}} \right)_{j,j^{\prime}=1}^{d}-\bm{I}_{d}\right\|_{F}^{2},$ ( local isometry loss ) $$

where $\|\cdot\|_{F}$ is the Frobenius norm and $\left(\left(\bm{e}^{j},\bm{e}^{j^{\prime}}\right)^{\varphi_{\bm{\theta}}}_{\bm
{x}_{i}}\right)_{j,j^{\prime}=1}^{d}$ denotes a $d$-dimensional matrix just as $(\bm{A}_{ij})^{d}_{i,j=1}$ denotes a matrix.

First, the global isometry loss takes global geometry into account, ensuring that the learned distances under the diffeomorphism $\varphi_{\bm{\theta}}$ approximate the true pairwise distances $d_{i,j}$ between data points. Second, the submanifold loss enforces that the data manifold is mapped to $\mathcal{M}=\mathcal{M}^{d^{\prime}}\times\mathbb{R}^{d-d^{\prime}}$, preserving the submanifold structure of the data in the latent space. Finally, the local isometry loss enforces local isometry, ensuring that small-scale distances and local geometry are preserved under the transformation, which is critical for maintaining the intrinsic geometric properties of the data during dimensionality reduction. For further details on the implementation and theoretical considerations, see Diepeveen ().

### A.4 Conditional Flow Matching

To achieve the goal of accurate generative modeling on data manifolds through isometric learning, we first need to understand generative modeling on Euclidean spaces. We do this through summarizing \acCFM Lipman et al. (), a commonly used and effective framework for learning \acpCNF for generative modeling for Euclidean data Chen et al. (). \AcCFM is a method designed to map a simple base distribution to a target data distribution by learning a time-dependent vector field. The fundamental goal of \acFM is to align a target probability path $p_{t}(\bm{x})$ with a vector field $u_{t}(\bm{x})$, which generates the desired distribution. The \acFM objective is defined as follows:

$$ $\mathcal{L}_{\text{FM}}(\bm{\eta})=\mathbb{E}_{t,p_{t}(\bm{x})}\|v_{t}(\bm{x}; \bm{\eta})-u_{t}(\bm{x})\|^{2},$ (33) $$

where $\bm{\eta}$ represents the learnable parameters of the neural network that parameterizes the vector field $v_{t}(\bm{x};\bm{\eta})$, and $t\sim\mathcal{U}(0,1)$ is uniformly sampled. However, a significant challenge in \acFM is the intractability of constructing the exact path $p_{t}(\bm{x})$ and the corresponding vector field $u_{t}(\bm{x})$.

To address this Lipman et al. () introduce \acCFM, a more practical approach by constructing the probability path and vector fields in a conditional manner. The \acCFM objective is then formulated by marginalizing over the data distribution $q(\bm{x}_{1})$ and considering the conditional probability paths:

$$ $\mathcal{L}_{\text{CFM}}(\bm{\eta})=\mathbb{E}_{t,q(\bm{x}_{1}),p_{t}(\bm{x}| \bm{x}_{1})}\|v_{t}(\bm{x};\bm{\eta})-u_{t}(\bm{x}|\bm{x}_{1})\|^{2}.$ (34) $$

A key result, as established in Theorem 2 of Chen & Lipman (), is that the gradients of the \acCFM objective with respect to the parameters $\bm{\eta}$ are identical to those of the original \acFM objective, i.e.,

$$ $\nabla_{\bm{\eta}}\mathcal{L}_{\text{FM}}(\bm{\eta})=\nabla_{\bm{\eta}} \mathcal{L}_{\text{CFM}}(\bm{\eta}),$ (35) $$

ensuring that optimizing the \acCFM objective yields the same result as the original \acFM objective. This enables effective train of the neural network without needing direct access to the intractable marginal probability paths or vector fields.

Given a sample $\bm{x}_{1}$ from the data distribution $q(\bm{x}_{1})$, we define a conditional probability path $p_{t}(\bm{x}|\bm{x}_{1})$ (^2^22In this work, we use two types of indexing: $x_{t}$ to denote time indices and $x_{i}$ for different data points. It should be clear from the context which indexing is being used.). This path starts at $t=0$ from a simple distribution, typically a standard Gaussian, and approaches a distribution concentrated around $\bm{x}_{1}$ as $t\rightarrow 1$:

$$ $p_{t}(\bm{x}|\bm{x}_{1})=\mathcal{N}(\bm{x}|\mu_{t}(\bm{x}_{1}),\sigma_{t}(\bm {x}_{1})^{2}\mathbf{I}),$ (36) $$

where $\mu_{t}(\bm{x}_{1}):[0,1]\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ is the time-dependent mean, and we denote the time-dependent standard deviation as $\sigma_{t}(\bm{x}_{1}):[0,1]\times\mathbb{R}\rightarrow\mathbb{R}_{>0}$. For simplicity, we set $\mu_{0}(\bm{x}_{1})=\bm{0}$ and $\sigma_{0}(\bm{x}_{1})=1$, ensuring that all conditional paths start from the same standard Gaussian distribution. At $t=1$, the path converges to a distribution centered at $\bm{x}_{1}$ with a small standard deviation $\sigma_{\text{min}}$.

The corresponding conditional vector field $u_{t}(\bm{x}|\bm{x}_{1})$ can be defined by considering the flow:

$$ $\chi_{t}(\bm{x})=\sigma_{t}(\bm{x}_{1})\bm{x}+\mu_{t}(\bm{x}_{1}),$ (37) $$

which maps a sample from the standard Gaussian to a sample from $p_{t}(\bm{x}|\bm{x}_{1})$. The vector field $u_{t}(\bm{x}|\bm{x}_{1})$ that generates this flow, as proven by Lipman et al. () in Theorem 3, is given by:

$$ $u_{t}(\bm{x}|\bm{x}_{1})=\frac{\sigma_{t}^{\prime}(\bm{x}_{1})}{\sigma_{t}(\bm {x}_{1})}\left(\bm{x}-\mu_{t}(\bm{x}_{1})\right)+\mu_{t}^{\prime}(\bm{x}_{1}),$ (38) $$

where the primes denote derivatives with respect to time to stay consistent with the original papers notation.

In this work, we choose to use the \acOT formulation of \acCFM. Here, the mean $\mu_{t}(\bm{x}_{1})$ and standard deviation $\sigma_{t}(\bm{x}_{1})$ are designed to change linearly in time, offering a straightforward interpolation between the base distribution and the target distribution. Specifically, the mean and standard deviation are defined as:

$$ $\mu_{t}(\bm{x}_{1})=t\bm{x}_{1},\quad\sigma_{t}(\bm{x}_{1})=1-(1-\sigma_{\text {min}})t.$ (39) $$

This linear path results in a vector field $u_{t}(\bm{x}|\bm{x}_{1})$ given by:

$$ $u_{t}(\bm{x}|\bm{x}_{1})=\frac{\bm{x}_{1}-(1-\sigma_{\text{min}})\bm{x}}{1-(1- \sigma_{\text{min}})t}.$ (40) $$

The corresponding conditional flow that generates this vector field is:

$$ $\chi_{t}(\bm{x})=(1-(1-\sigma_{\text{min}})t)\bm{x}+t\bm{x}_{1}.$ (41) $$

This \acOT path is optimal in the sense that it represents the displacement map between the two Gaussian distributions $p_{0}(\bm{x}|\bm{x}_{1})$ and $p_{1}(\bm{x}|\bm{x}_{1})$ Lipman et al. ().

The final \acCFM loss under this \acOT formulation is derived by substituting the above vector field and flow into the general \acCFM objective (Equation 34) and reparameterizing $p_{t}(\bm{x}|\bm{x}_{1})$ in terms of $\bm{x}_{0}$. This yields the following objective function:

$$ $\mathcal{L}_{\text{CFM}}(\bm{\eta})=\mathbb{E}_{t,q(\bm{x}_{1}),p(\bm{x}_{0})} \left\|v_{t}(\chi_{t}(\bm{x}_{0});\bm{\eta})-\frac{\bm{x}_{1}-(1-\sigma_{\text {min}})\bm{x}_{0}}{1-(1-\sigma_{\text{min}})t}\right\|^{2}.$ (42) $$

This formulation is advantageous because the \acOT paths ensure that particles move in straight lines and with constant speed, leading to simpler and more efficient regression tasks compared to traditional diffusion-based methods. We use the \acsOT-\acsCFM objective in this work when we refer to \acCFM.

### A.5 Riemannian Flow Matching

The next step toward generation on data manifolds is understanding generation on manifolds with closed form geometric mappings. \AcRFM aims to do exactly this by generalizing \acCFM to Riemannian manifolds Chen & Lipman (). Assume a complete, connected and smooth manifold $\mathcal{M}$ endowed with a Riemannian metric $(\cdot,\cdot)^{\mathcal{M}}$. We are given a set of training samples $\bm{x}_{1}\in\mathcal{M}$ from some unknown data distribution $q(\bm{x}_{1})$ on the manifold. Then the goal is to learn a parametric map $\rho:\mathcal{M}\rightarrow\mathcal{M}$ that pushes a simple base distribution $p$ to the data distribution $q$. To achieve RFM Chen & Lipman () reparameterize the conditional flow as

$$ $\bm{x}_{t}=\chi_{t}(\bm{x}_{0}|\bm{x}_{1}),$ (43) $$

where $\chi_{t}(\bm{x}_{0}|\bm{x}_{1})$ is the solution to the \acODE defined by a time-dependent conditional vector field $u_{t}(\bm{x}|\bm{x}_{1})\in\mathcal{T}_{\bm{x}}\mathcal{M}$ that is tangent to the manifold $\mathcal{M}$. The initial condition is set as $\chi_{0}(\bm{x}_{0}|\bm{x}_{1})=\bm{x}_{0}$.

This formulation leads to the RFM objective, which ensures that the vector field $u_{t}(\bm{x}|\bm{x}_{1})$ learned by the model lies entirely within the tangent space of the manifold at each point $\bm{x}_{t}\in\mathcal{M}$:

$$ $\mathcal{L}_{\text{RFM}}(\bm{\eta})=\mathbb{E}_{t,q(\bm{x}_{1}),p(\bm{x}_{0})} \|v_{t}(\bm{x}_{t};\bm{\eta})-u_{t}(\bm{x}_{t}|\bm{x}_{1})\|_{\mathcal{M}}^{2},$ (44) $$

where $\|\cdot\|_{\mathcal{M}}$ is the norm enduced by the Riemannian metric $(\cdot,\cdot)^{\mathcal{M}}$. For manifolds with closed-form geodesic expressions, a simulation-free objective can be formulated using exponential and logarithmic maps. This approach allows models to be trained without numerically simulating particle trajectories, leveraging closed-form geodesics and mappings to directly compute vector fields and transport paths. In this case, $\bm{x}_{t}$ can be defined by the geodesic between $\bm{x}_{1}$ and $\bm{x}_{0}$ and can be explicitly expressed as

$$ $\bm{x}_{t}=\gamma^{\mathcal{M}}_{\bm{x}_{1},\bm{x}_{0}}\big{(}\kappa(t)\big{)} =\exp^{\mathcal{M}}_{\bm{x}_{1}}\big{(}\kappa(t)\log^{\mathcal{M}}_{\bm{x}_{1} }(\bm{x_{0}})\big{)},$ (45) $$

with monotonically decreasing differentiable function $\kappa(t)$ satisfying $\kappa(0)=1$ and $\kappa(1)=0$ acting as a scheduler. Furthermore, the tangent vector field $u_{t}(\bm{x}|\bm{x}_{1})$ can be evaluated through,

$$ $u_{t}(\bm{x}|\bm{x}_{1})=\dot{\gamma}^{\mathcal{M}}_{\bm{x}_{1},\bm{x}_{0}} \big{(}\kappa(t)\big{)}=\frac{d}{dt}\exp^{\mathcal{M}}_{\bm{x}_{1}}\big{(} \kappa(t)\log^{\mathcal{M}}_{\bm{x}_{1}}(\bm{x_{0}})\big{)}$ (46) $$

The objective function is then given by:

$$ $\mathcal{L}_{\text{RFM}}(\bm{\eta})=\mathbb{E}_{t,q(\bm{x}_{1}),p(\bm{x}_{0})} \left\|v_{t}(\gamma^{\mathcal{M}}_{\bm{x}_{0},\bm{x}_{1}}\big{(}\kappa(t)\big{ )};\bm{\eta})-\dot{\gamma}^{\mathcal{M}}_{\bm{x}_{0},\bm{x}_{1}}\big{(}\kappa( t)\big{)}\right\|_{\mathcal{M}}^{2}$ (47) $$

Constructing a simulation-free objective for RFM on general geometries presents significant challenges due to the absence of closed-form expressions for essential geometric operations, such as exponential and logarithmic maps, or geodesics. These operations are crucial for defining and efficiently evaluating the objective but are often computationally intensive to approximate without closed-form solutions. For a list of examples of manifolds with closed-form geometric mappings, see the appendix of Chen & Lipman ().

In the absence of such closed-form solutions, existing methods tackle these difficulties by either learning a metric that constrains the generative trajectory to align with the data support Kapusniak et al. () or by assuming a metric with easily computable geodesics on the data manifold Chen & Lipman (). However, learning a metric can be problematic as it may lead to overfitting or fail to capture the true geometry of the data, particularly when the data manifold is complex or poorly understood. On the other hand, assuming a simple metric with computable geodesics can oversimplify the problem, resulting in models that inadequately represent the underlying data structure. To overcome these challenges, we introduce Pullback Flow Matching in section 3.

## Appendix B Neural ODEs Parameterize Diffeomorphisms

We can verify that this defines a diffeomorphism by using Theorem C.15 of Younes (). According to Theorem C.15, for $\phi_{\bm{\theta}}$ to be a diffeomorphism, the vector field $f$ must satisfy $f\in L^{1}([0,1],C^{1}_{(0)}(\Omega,\mathbb{B}))$, where $\Omega$ is the domain of the vector field and $\mathbb{B}$ is a Banach space representing the target space.

In our case, $f$ is composed of smooth and continuously differentiable functions due to the \acMLP parameterization, ensuring $f$ is also smooth and continuously differentiable. Additionally, we enforce local isometry by regularizing the Jacobian of $f_{\bm{\theta}}$, which guarantees local regularity of $f$ in the data domain (see stability regularization). Thus, $f$ meets the required conditions and $\phi_{\bm{\theta}}$ defines a proper diffeomorphism.

## Appendix C Manifold and Metric Selection

Isometric learning requires three key choices to be made, first one needs to choose the Riemannian metric of the data manifold $(\cdot,\cdot)^{\mathcal{D}}$, second one needs to choose both the latent (sub)manifold and its Riemannian metric $\big{(}\mathcal{M}_{d^{\prime}},(\cdot,\cdot)^{\mathcal{M}_{d^{\prime}}}\big{)}$ and finally one needs to choose the dimensionality $d^{\prime}$. Technically one also needs to assume a metric on $\mathbb{R}^{d-d^{\prime}}$, but in this work we assume a Euclidean metric $(\cdot,\cdot)_{2}$ throughout all our experiments.

There are several options when selecting the metric on the data manifold $(\cdot,\cdot)^{\mathcal{D}}$. One can choose for example a locally euclidean approximation through Isomap Tenenbaum et al. () or a more noise-robust geodesic approximation Little et al. (). One can also design a metric to create a latent space (^3^33In this text we refer to the latent space as the concept in machine and representation learning, technically its a latent manifold endowed with a Riemannian metric, not a vector space.) structured based on properties of the data one cares about, we show how in subsection 5.4. In this work, we focus on using a proper metric and defer the exploration of learning with pseudo-metrics to future research.

When selecting a latent Riemannian (sub)manifold and metric it is crucial to select $\mathcal{M}_{d^{\prime}}$ such that $\mathcal{M}=\mathcal{M}_{d^{\prime}}\times\mathbb{R}^{d-d^{\prime}}$ it is diffeomorphic to the data manifold $\mathcal{D}$. This ensures that the latent space of the RAE can effectively capture the intrinsic structure of the data. The manifold should be chosen based on its ability to accommodate the data’s periodicity, curvature, and dimensionality. This alignment is essential for accurately representing the data manifold within the latent space. Unless otherwise stated we assume $\mathcal{M}_{d^{\prime}}=\mathbb{R}^{d^{\prime}}$. Additionally, one should select the Riemannian geometry of $\big{(}\mathcal{M}_{d^{\prime}},(\cdot,\cdot)^{\mathcal{M}_{d^{\prime}}}\big{)}$ such that geometric mappings can be explicitly defined in closed form. A list of manifolds with closed form geometric mappings can be found in the appendix of Chen & Lipman (). Unless otherwise states we select $(\cdot,\cdot)^{\mathcal{M}_{d^{\prime}}}=(\cdot,\cdot)_{2}$.

Finally, $d^{\prime}$, the dimensionality of the latent space, is a hyperparameter that could be tuned through iterative testing. Techniques such as Isomap Tenenbaum et al. () or equivalents on other manifolds such as hyperbolic space Cvetkovski & Crovella () can be employed to evaluate various dimensional and Riemannian geometric settings and determine the optimal $d^{\prime}$ that balances model complexity with the ability to accurately capture the data manifold’s structure.

## Appendix D Data Description

In this work we use several datasets, synthetic, simulated and experimental. Here we describe them in order of appearance in the experiments.

### D.1 ARCH Dataset

We create a dataset in the spirit of Tong et al. (). We sample $n=500$ data points uniform on the line $[-1,1]$ ($x_{i}\sim\mathcal{U}(-1,1)$), wrap this line around the unit half circle and add normally distributed noise with $\sigma=0.1$, i.e.

$$ $y_{i,1}=\sin(0.5\pi x_{i})+a_{i,1},\quad y_{i,2}=\cos(0.5\pi x_{i})+a_{i,2} \text{ for }a_{i,j}\sim\mathcal{N}(0,0.1^{2}).$ (48) $$

An example of the dataset can be found in Figure 4.

### D.2 \Acf AK

We consider the time-normalized open-to-close transition of \acAK. This is a dataset from coarse-grained molecular dynamics simulations consisting of $n=102$ conformations of $214$ amino-acids in $3D$, samples of the trajectory can be found in Figure 7.

Figure: (a) $t=0$
Refer to caption: extracted/5904060/figures/data/f0_101.png

### D.3 Intestinal Fatty Acid Binding Protein

The second protein dynamics dataset is that of $n=500$ conformations of \acI-FABP in water. The datasets comes from simulations in CHARMM of $500$ picoseconds (ps) with a $2$ femtoseconds (fs) timestep. The data can be found on [mdanalysis.org](https://www.mdanalysis.org/MDAnalysisData/ifabp_water.html).

### D.4 \Acs GRAMPA Dataset

The \acfGRAMPA dataset Witten & Witten () is a compilation of peptides and their antimicrobial activity against various bacteria, including E. coli and P. aeruginosa. It includes data on peptide sequences, target bacteria, bacterial strains, and \acMIC values, with additional columns providing details on sequence modifications and data sources. The dataset was created to support deep learning models aimed at predicting the antimicrobial effectiveness of peptides. The dataset is available [here](https://github.com/zswitten/Antimicrobial-Peptides). In our experiments we follow the preprocessing pipeline from Szymczak et al. () and use only the sequence data and the corresponding \acMIC scores. After preprocessing we are left with $n=3444$ sequences of maximum $25$ amino-acids with tested antimicrobial activity against E. coli.

## Appendix E Training Procedure

We explain the training procedure and hyperparameter settings for each of the experiments in section 5 in further detail for reproducibility. In all experiments the datasets where split into train and test sets. We apply early stopping and present the model with the lowest average loss on the test data.

### E.1 Ablation Study

For details of hyperparameter settings for the ablation study see Table 5.

**Table 5: Hyperparameter settings for ablation study of \acRAE on the \acARCH, \acAK and \acI-FABP datasets.**
| Hyperparameter | \acsARCH | \acsAK | \acsI-FABP |
| --- | --- | --- | --- |
| Epochs | $1000$ | $1000$ | $1000$ |
| Learning Rate | $0.0001$ | $0.0001$ | $0.0001$ |
| Optimizer | Adam | Adam | Adam |
| Train/Test Split | $0.8/0.2$ | $0.8/0.2$ | $0.8/0.2$ |
| $n_{steps}$ | $10$ | $10$ | $10$ |
| Seed | 0 0 | 0 0 | 0 0 |
| Number of Layers | $5$ | $5$ | $5$ |
| $\alpha_{1}$ | $1.0$ | $1.0$ | $1.0$ |
| $\alpha_{2}$ | $[0.0,1.0]$ | $[0.0,5.0]$ | $[0.0,5.0]$ |
| $\alpha_{3}$ | $1.0$ | $1.0$ | $1.0$ |
| $\alpha_{4}$ | $[0.0,0.01]$ | $[0.0,0.005]$ | $[0.0,0.1]$ |
| $d^{\prime}$ | $1$ | $1$ | $1$ |
| Hidden Units | $64$ | $214\cdot 3+1$ | $131\cdot 3+1$ |
| Number of Neighbors | $5$ | $2$ | $4$ |
| Batch Size | $64$ | $16$ | $64$ |
| Warmup | $50$ | $400$ | $200$ |

Specific hyperparameters worth mentioning are $n_{steps}$ which is the number of Runge-Kutta steps we use in our \acNODE, Number of Layers is the number of layers of the \acMLP with swish activation function for the vector field of the \acNODE. The Number of Neighbors is the hyperparameter used to calculate the shortest paths over the nearest neighbors graph for the Isomap geodesics in sklearn and the Warmup is the number of epochs we train with $\alpha_{1},\alpha_{2}=0$ to first learn a lower dimensional representation.

### E.2 Interpolation Experiments

For details of hyperparameter settings for the interpolation experiments of $(\cdot,\cdot)^{\mathcal{M}}$- and $(\cdot,\cdot)^{\mathcal{M}_{d^{\prime}}}$-interpolation see Table 6 and for the ($\beta$-)\acspVAE see Table 7.

**Table 6: Hyperparameter settings for interpolation experiments for $(\cdot,\cdot)^{\mathcal{M}}$- and $(\cdot,\cdot)^{\mathcal{M}_{d^{\prime}}}$-interpolation on the \acARCH, \acAK and \acI-FABP datasets.**
| Hyperparameters | \acsARCH | \acsAK | \acsI-FABP |
| --- | --- | --- | --- |
| Epochs | $1000$ | $1000$ | $1000$ |
| Learning Rate | $0.0001$ | $0.0001$ | $0.0001$ |
| Optimizer | Adam | Adam | Adam |
| Train/Test Split | $0.8/0.2$ | $0.8/0.2$ | $0.8/0.2$ |
| $n_{steps}$ | $10$ | $10$ | $10$ |
| Seed | 0 0 | 0 0 | 0 0 |
| Number of Layers | $5$ | $5$ | $5$ |
| $\alpha_{1}$ | $1.0$ | $1.0$ | $1.0$ |
| $\alpha_{2}$ | $5.0$ | $5.0$ | $5.0$ |
| $\alpha_{3}$ | $1.0$ | $1.0$ | $1.0$ |
| $\alpha_{4}$ | $0.001$ | $0.005$ | $0.1$ |
| $d^{\prime}$ | $1$ | $1$ | $1$ |
| Hidden Units | $64$ | $214\cdot 3+1$ | $131\cdot 3+1$ |
| Number of Neighbors | $5$ | $2$ | $4$ |
| Batch Size | $64$ | $16$ | $64$ |
| Warmup | $50$ | $400$ | $200$ |
| $n_{parameters}$ | $17282$ | $2486480$ | $934961$ |

**Table 7: Hyperparameter settings for interpolation experiments for ($\beta$-)\acsVAE on the \acARCH dataset. \AcpVAE have $\beta=1.0$, $\beta$-\acspVAE have $\beta=10.0$.**
| Hyperparameters | \acsARCH | \acsAK | \acsI-FABP |
| --- | --- | --- | --- |
| Epochs | $1000$ | $1000$ | $1000$ |
| Learning Rate | $0.0001$ | $0.0001$ | $0.0001$ |
| Optimizer | Adam | Adam | Adam |
| Train/Test Split | $0.8/0.2$ | $0.8/0.2$ | $0.8/0.2$ |
| Seed | 0 0 | 0 0 | 0 0 |
| Number of Encoder Layers | $5$ | $5$ | $5$ |
| Number of Decoder Layers | $5$ | $5$ | $5$ |
| Hidden Units | $64$ | $214\cdot 3+1$ | $131\cdot 3+1$ |
| Beta | $[1.0,10.0]$ | $[1.0,10.0]$ | $[1.0,10.0]$ |
| $d^{\prime}$ | $1$ | $1$ | $1$ |
| Batch Size | $64$ | $16$ | $64$ |
| $n_{parameters}$ | $34184$ | $4555655$ | $1712324$ |

### E.3 Generation Experiments

**Table 8: Hyperparameter settings for \acCFM, \acPFM and $d^{\prime}$-\acsPFM for generation experiments. The same isometry $\varphi_{\bm{\theta}}$ of the interpolation experiments is used for the \acPFM and $d^{\prime}$-\acsPFM.**
|  | \acsARCH | \acsI-FABP |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Hyperparameter | \acsCFM | \acsPFM | $d^{\prime}$-\acsPFM | \acsCFM | \acsPFM | $d^{\prime}$-\acsPFM |
| Epochs | $5000$ | $5000$ | $5000$ | $5000$ | $5000$ | $5000$ |
| Learning Rate | $0.0005$ | $0.0005$ | $0.0005$ | $0.001$ | $0.001$ | $0.0005$ |
| Scheduler | Cosine | Cosine | Cosine | Cosine | Cosine | Cosine |
| Minimum Learning Rate | $5.0\cdot 10^{-6}$ | $5.0\cdot 10^{-6}$ | $5.0\cdot 10^{-6}$ | $1.0\cdot 10^{-5}$ | $1.0\cdot 10^{-5}$ | $5.0\cdot 10^{-6}$ |
| Train/Test Split | $0.8/0.2$ | $0.8/0.2$ | $0.8/0.2$ | $0.8/0.2$ | $0.8/0.2$ | $0.8/0.2$ |
| Seed | 0 0 | 0 0 | 0 0 | 0 0 | 0 0 | 0 0 |
| Number of Layers | $10$ | $10$ | $10$ | $10$ | $10$ | $10$ |
| Hidden Units | $64$ | $64$ | $16$ | $131\cdot 3+1$ | $131\cdot 3+1$ | $64$ |
| Batch Size | $64$ | $64$ | $64$ | $64$ | $64$ | $64$ |
| $n_{\text{simulation steps}}$ | $10$ | $10$ | $10$ | $10$ | $10$ | $10$ |

### E.4 Designable Latent Manifolds for Novel Protein Engineering

In Table 9 one can find the settings for training the isometry on the \acGRAMPA dataset for the protein sequence design experiments. Specific hyperparameter worth mentioning is the embedding dimensions, we use an embedding layer from the Flax library to embed the discrete sequences into a continuous space and use a sign-cosine positional embedding, to embed the location in the sequence of the amino acids in the data.

**Table 9: Hyperparameter settings for protein design experiments of the \acpRAE on the \acGRAMPA dataset.**
| Hyperparameters | Setting |
| --- | --- |
| Epochs | $1000$ |
| Learning Rate | $0.0001$ |
| Optimizer | Adam |
| Train/Test Split | $0.8/0.2$ |
| $n_{steps}$ | $10$ |
| Seed | 0 0 |
| Number of Layers | $5$ |
| Embedding dimension | $8$ |
| $\alpha_{1}$ | $5.0$ |
| $\alpha_{2}$ | $5.0$ |
| $\alpha_{3}$ | $5.0$ |
| $\alpha_{4}$ | $0.05$ |
| $d^{\prime}$ | $128$ |
| Hidden Units | $512$ |
| Batch Size | $128$ |
| Warmup | $100$ |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: - The notations can be more clear. For example, the scheduler $κ(t)$ and vector field $v(t)$ are not clearly defined in the main body of the paper. This may lead to misunderstandings of the content of the article.
- The article will be more convincing if experiments that compare PFM with RFM are included.
- How might the properties of different manifolds, particularly those with non-trivial curvature, affect the performance of PFM?
soundness: 2
strengths: - This paper demonstrates strong mathematical foundations and theoretical rigor, which helps ensure the method’s reliability and provides theoretical guarantees for its performance.
- The neural ODE-based diffeomorphisms and design of product manifold eliminates need for expensive Riemannian metric tensor calculations and requires only geodesic distance approximations.
summary: The paper introduces Pullback Flow Matching (PFM), a novel generative modeling framework designed to operate on data manifolds, which uses pullback geometry and isometric learning to preserve the geometric structure of data manifolds while enabling efficient generation and precise interpolation in latent spaces. PFM leverages Neural Ordinary Differential Equations (Neural ODEs) to improve the expressiveness and efficiency of learning diffeomorphisms. This approach enables more accurate interpolation and generation of data. The paper demonstrates the effectiveness of PFM through experiments on synthetic datasets, molecular dynamics simulations, and protein sequence data. The results show that PFM can generate novel proteins with desired properties, indicating its potential for real-world applications in scientific domains where precise modeling of complex systems is required.
weaknesses: - While PFM is presented as an improvement over RFM, there’s no quantitative evidence showing the relative performance. The paper does not provide direct experimental comparisons between PFM and RFM.
- The paper mainly compare PFM with CFM and VAE, which misses comparasion with more state-of-the-art models, such as Diffusion, GAN, SLERP.
- While the paper claims geometric preservation through pullback geometry and isometric learning, no formal theoretical proof is provided showing that PFM actually preserves geometric structures. The experiments are insufficient to verify geometric preservation claims as well.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Questions and minor comments:
- I find the notation for the metric $(\cdot,\cdot)^M$ a bit unconventional, maybe $g$ is better? Further, the dependence with $p$ is not clear with the current notation, whereas you could use $g_p$.
- The related work section could be improved. For instance [1,2,3,4] are work that explore geometric regularization of latent space, specifically [1] regularize for geodesics. 
- The link with scaling laws in the introduction is a bit confusing. What I understand is that having prior knowledge on the geometry of the data is a great inductive bias, so learning the geometry should be useful. 
- On line 200, the authors mention that the method requires fewer epochs and parameters, could you add a reference to the section backing this claim ? Further, does this claim take into account the training of the isometries ?
- On line 287 "we demonstrate the effectiveness of the new regularization terms", what are these new regularization terms ? Is it the graph matching loss and the stability regularization ?
- In section 5.1, for the isometry computation, my understanding is that you need a ground truth geodesic distance. How do you compute it for I-FABP?
- In section 5.3 and 5.4, how many seeds did you use for the experiments ? Could the authors include standard deviation?
- Missing reference on line 1044.


[1] Huguet, Guillaume, et al. "Manifold interpolating optimal-transport flows for trajectory inference." Advances in neural information processing systems 35 (2022): 29705-29718.

[2] LEE Yonghyeon, Sangwoong Yoon, Minjun Son, and Frank C Park. Regularized autoencoders for isometric representation learning. In International Conference on Learning Representations, 2021.

[3] Philipp Nazari, Sebastian Damrich, and Fred A Hamprecht. Geometric autoencoders–what you see is what you decode. 2023.

[4] Fasina, O., Huguet, G., Tong, A., Zhang, Y., Wolf, G., Nickel, M., ... & Krishnaswamy, S. (2023, July). Neural FIM for learning Fisher information metrics from point cloud data. In International Conference on Machine Learning (pp. 9814-9826). PMLR.
soundness: 3
strengths: The introductions and background read well, especially I find the method to be well motivated. Overall, the manuscript is clear and the presentation is easy to follow. The authors validate their method on toy and real data, and they conduct an ablation study on various loss terms.
summary: The manuscript is on learning representations tailored for downstream tasks such as interpolation on a lower dimensional manifold and generation. The proposed method uses a Neural ODE to learn a diffeomorphism, which enable simulation free training of dynamics with flow matching.
weaknesses: - The contributions of the paper appear somewhat limited, as it builds upon Diepeveen (2024). I think it would be beneficial if in the main body of the text the authors explained in more details how their work is different from Diepeveen (2024).
- To improve on the current experiments, It would be great to add another toy dataset with known geodesic, e.g. a swiss roll rotated in higher dimensions..
- On line 200, the authors mention that the method requires fewer epochs, but I don't see an experiment backing this claim. Specifically, it could be interesting to compare total training time including learning isometries and dynamics. 
- My main concern with the experiments is that there are no comparison with other methods using geometric regularization. The authors mentioned a few in the related work section, but they are not used for comparisons.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Q1. If I understood correctly, the goal of the diffeomorphism (in the simplest scenario) is to map the data into a Euclidean space of the same dimensionality, where the straight-line distance between data points corresponds to the geodesic distance in the original space (as measured by the Isomap approach). I believe Fig. 1 and Fig. 2 could be improved to aid understanding, or perhaps, adding a new figure.

Q2. The objective function is rather complex and depends on many hyperparameters. A critical quantity, the geodesics between the original data are computed using a graph-based method. While this should work well for the ARCH dataset, we know that graph-based interpolations are less ideal in high-dimensional spaces. Additionally, how do the parameters $\alpha_i$ affect the learned model? I believe an ablation study would have been useful here.

Q3. The data are mapped near a lower-dimensional submanifold, $\mathcal{M}_{d'}$, where dimensionality $d'$ is an additional hyperparameter. How does the choice of $d'$ influence the results? Also, is this achieved by pushing the $d - d'$ dimensions near zero using the "submanifold loss"? I think the diagonal terms in this part of the objective function should be inverted.

Q4. Minor: I think that definitions for $\psi$ (line 229) and the concept of local isometry should be included.
soundness: 3
strengths: - The technical aspects of the paper seem to be solid.
- The approach for learning a Riemannian metric in the data space is interesting.
- The experiments demonstrate the effectiveness of the proposed model.
summary: The authors propose a Riemannian flow-matching scheme by learning a metric that pulls the associated interpolations near the support of the given data. The Riemannian metric is induced in the data space by training a diffeomorphism that isometrically maps the data to a latent manifold, preserving the distances of the original data. They propose an objective function that aims to satisfy the properties of the diffeomorphism. In the experiments, they consider a synthetic example to demonstrate the properties of the proposed model and a real-world problem on protein data to highlight its effectiveness.
weaknesses: - The paper seems to rely heavily on Diepeveen (2024), with the primary difference being an extension of the diffeomorphism, while the Riemannian flow matching part is a direct application of Chen & Lipman (2024). Although the experiment with protein data is interesting, the technical contribution feels somewhat limited.
- While the writing is technically accurate, the clarity could be improved. For example, adding figures would aid in explaining the modeling approach.
- The objective function is reasonable, but it involves many hyperparameters, which are not straightforward to select (see questions).

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. What is the dimension $d$ of the data in GRAMPA dataset in Section 5.4? Why is $d'$ chosen as 128? This should be explained.

1. In Eq. (8), how do you choose $\mu$ in $T_{\mu}$? What role does this translation play? How is $\psi$ in Eq. (8) selected, and is it simply the identity mapping when $\mathcal{M}'$ is Euclidean space $\mathbb{R}^{d'}$? Could you provide more related details in the paper?

1. In the submanifold loss (line 279), why is the notation $T_{\mu}$ changed to $T_{z}$? Is $z$ related to $z_i$ in the stability regularization (line 282)?

1. It seems that choosing the latent manifold $\mathcal{M}$ as $\mathbb{R}^d$ with Euclidean metric is  sufficient (to approximate the distance measure $(\cdot,\cdot)^{\mathcal{D}}$ on the data manifold using an isometric mapping). In fact, this is how it was selected in the experiments. Why not present equations in Section 3 in a more specific form? The presentation there in general form makes it harder to understand.

1. In all the experiments except the last one, $d'$ is chosen to be 1. Is this choice made because the datasets are essentially concentrated on one-dimensional manifolds? How to choose $d'$ in more general cases? Can the author consider examples where $d'$ is greater than 1? 

1. The sizes of the datasets in the experiments are quite small, sometimes even lower than the dimension of the data. Can the method apply to large datasets? Training the isometries (using the global isometry loss and graph matching loss) requires iterating over all point pairs, which seems to be computationally intensive when the batch size is large. Can the authors discuss the computational cost of isometric learning, especially when the size of the dataset is large?

 1.   There are several typos throughout the paper.

- Line 070, the sentence doesn't read correctly. 
- Figure 2, the direction of the arrow connecting $x_1$ and $\varphi(x_1)$ seems to be wrong. 
- Line 321:  "as well ass" should be "as well as".
- Line 284: "induces" should be "induced".
- Line 1043: Missing citation.
- Line 279: To align with Eq (8), $I_{d-d'}$ and $0_{d'}$ should be switched in the submanifold loss.
- The initial letter of the names in references, e.g. poincare, ode, riemannian, jacobian, bayes, should be capitalized.
soundness: 2
strengths: ###  Originality
This paper proposes a new generative modeling approach for datasets with manifold structure. The approach combines flow matching and isometric learning, which preserve the intrinsic geometric structure of the data manifold in the latent space.

### Quality

The research question studied in this paper, i.e. generative modeling for data with manifold structure, is definitely interesting. The proposed learning approach, i.e. isometric learning and flow matching, as well as the training objectives, are reasonable. The experiments are discussed with sufficient details.    

### Clarity

Overall, the presentation of the research question, the learning approach and the numerical experiments of this paper is  clear. 

### Significance

The proposed approach may have certain advantages over existing methods for dataset with manifold structure.
summary: This paper proposes Pullback Flow Matching, a novel framework for generative modeling on data
with manifold structures. Building on the previous study, this work leverages pullback geometry
to define a new metric on the entire ambient space $\mathbb{R}^d$, by learning an isometry that preserves the
geometric structure of the data manifold $\mathcal{D}$ on the latent manifold $\mathcal{M}$. The corresponding metric of
the assumed latent manifold $\mathcal{M}$ is then used to perform Riemannian Flow Matching (RFM) on the
latent manifold. The effectiveness of this method is demonstrated through experiments on synthetic
data, molecular dynamics data, and experimental peptide sequences.

Main contributions: 
1. A novel generative modeling framework that combines isometric learning and Flow Matching. 
2. An objective for isometric learning that only relies on distance measure on data manifold. 
3. Improved parameterization of isometries via neural ODEs.
weaknesses: 1. In the numerical experiments (Section 5), a relatively large part of the text focuses on isometry learning, and less effort is made on the study of the generation part. Results in Table 3 of Section 5.3 show that 1-PFM is better than CFM/PFM, but 1-PFM is solving a modified/simplified generative task instead of generating true data distribution. The same is true for the analogue generation in Section 5.4. Given these facts, the referee is afraid that the presented material may not be sufficient to firmly support some of the statements/claims made in the the abstract and introduction.

1. Section 2 and 3 consider a general setting, but the practical implementation only addresses a special case, which may cause confusion. Concretely, while the latent space is defined as a general manifold $\mathcal{M}$ in Section 2 and 3, it is chosen as the Euclidean space $\mathbb{R}^d$ in practical computations. In particular, if I understand correctly, Eq. (6) and (7) are actually flow matching losses in the latent space $\mathbb{R}^d$ or $\mathbb{R}^{d'}$ with the standard Euclidean metric. Writing them in a general form does not help readers understand the idea.

Meta-review: This paper presents Pullback Flow Matching (PFM), a generative modeling framework for data manifolds that uses pullback geometry and isometric learning to preserve manifold geometry. By enhancing isometric learning with Neural ODEs and a scalable training objective, PFM is argued to improve manifold learning and generative performance. The effectiveness of this method is demonstrated through experiments on synthetic data, molecular dynamics data, and experimental peptide sequences.

The paper received average scores between 3 and 5, with critics concentrating on lacks of novelty with respect to reference Diepeveen (2024) and Chen & Lipman (2024), lack of experimental validation with respect to other manifold based generative approaches and lack of theoretical guarantees about preservation of isometries. In this light, I am recommending a reject decision, and I encourage the authors to further strengthen their work on the questions raised by reviewers.

Decision: Reject