Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling

We thank the reviewer for the positive comments and insightful questions.

Q1: Could the authors provide correlation analyses between predicted growth rates and known proliferation markers or synthetic ground truth?

As suggested, we implement correlation analysis between the predicted growth rate and the ground truth to verify the accuracy of our growth rate modeling and learning as follows. We choose the Simulation Gene Data (Appendix B.2, Page 19) which is generated from the given dynamic systems (eq. (A-23), Page 19). The ground truth of growth rate can also be calculated analytically. Therefore, we can conduct correlation analysis on the out-of-distribution (OOD) time points, which are not included in the training data, to verify the generalization ability of VGFM. Results are reported in the following Table r4-1.

Table r4-1: correlation analysis on four OOD time points on Simulation Gene Data

The growth rate predicted by our VGFM shows a high correlation with the ground truth, and we will add this analysis in the Appendix.

Q2: Have the authors considered estimating uncertainty in predicted velocity/growth (e.g., via dropout or ensembles)?

While our current model is deterministic, dropout can be applied to our velocity/growth networks to obtain uncertainty. During inference, we calculate the averaged variance of each dimension of velocity and growth over all the points across different branches (i.e., cell types) and times. We experiment on Simulation Gene Data (Appendix B.2), set dropout rate as 0.1 and execute 5 runs on the velocity/growth network during inference for all the points. Results are shown in the following Table r4-2.

Table r4-2: uncertainty estimation of predicted velocity (2-dimensional) and growth on Simulation Gene Data. var_x1/var_x2 and var_g are the average variance of $1$-st/$2$-nd dimension of velocity and growth respectively.

We analyze the results at both spatial and temporal levels.

Spatial level: In the Simulation Gene Data, cells of branch 0 relatively stand still with little variation on state and counts. The result shows that the variance of $v_1$, $v_2$ and $g$ in this branch is small, indicating low uncertainty, while in branch 1 the variances is larger, indicating relatively higher uncertainty.

Temporal level: In the Simulation Gene Data, cells undergo significant state transitions and quantity changes in the initial phase, and the magnitude of the change decreases thereafter. Accordingly, the variance of $v_1,v_2$ and $g$ also reflects this trend, confirming our uncertainty estimation.

We will include these results as well as visulize the heatmap of uncertainty in the state space at each time points in appendix.

Q3: Can the method be extended to scenarios where timepoint alignment is unknown or uncertain, possibly by integrating OT-based time alignment?

Our current method assumes known time labels. In the absence of such labels, one possible extension is to first infer pseudotime using standard approaches (e.g., Monocle, Slingshot), then apply OT-based alignment between discretized pseudotime segments to estimate cell transitions and population growth. Alternatively, traditional pseudotime methods typically rely on low-dimensional embeddings (e.g., PCA, UMAP) followed by trajectory reconstruction. Inspired by this, we are also interested in future work that jointly learns cell evolution trajectories, velocity fields, and growth dynamics directly in the latent space, integrating pseudotime inference and dynamic modeling into a unified OT-based framework.

Q4: Effect of Sampling Bias: Since the method relies on observed cell counts for growth learning, have the authors tested robustness under non-uniform or subsampled cell distributions across time?

In the cell dynamics reconstruction, we usually assume that the observed dataset accurately reflects the evolution of true cell distribution across time. So the observed cell counts is consistent with the fact of cell proliferation and death. It's also practically meaningful to evaluate the robustness of VGFM trained on the biased observed dataset.

For the experiment settings, we use Simulation Gene Data (Appendix B.2) which has two spacial branches. To generate the non-uniform or subsampled data, we resample the original data with perturbed branch ratio. Specifically, we randomly select perturbation factor where $a_t\sim\mathcal{N}(0,\sigma^2)$ for time $t$. Then if the original branch ratio is $s_t:1-s_t$, the resampled branch ratios become $s_t-|a_t|:1-s_t+|a_t|$ . $\sigma$ indicates the bias level. We train the model on the resampled data and evaluate on the original data, to testify the robustness of VGFM.

Table r4-3: results of training on sampling bias on Simulation Gene Data with merics ($W_1$/RME) averaged over three random seeds.

We also experiment on the real-world dataset EB (Appendix B.6). Since this dataset doesn't have obvious spatial branches, we take a different dropout method to construct the biased dataset and likewise control the bias level with a factor $\eta$. Specifically, we sample several $a_t\sim \mathcal{U}([0,\eta])$ with different bias factors, where $\mathcal{U}([0,\eta])$ is the uniform distribution on $[0,\eta]$. Then, the points at time $t$ with a proportion of $a_t$ are randomly removed from the dataset. In such way, we can get a biased version of EB dataset.

Table r4-4: results of training on the biased EB datasets with merics ($W_1$/RME) averaged over three random seeds.

Tables r4-3 and r4-4 indicates that when the bias level is relatively low ($\sigma\le 0.1$ or $\eta\le 0.1$), our VGFM seems robust to the sampling bias. But if the bias level is high, the perfromance degrades. For such cases, we believe that with more biology prior beyond the cell counts integrated into growth modeling, our method could be more robust to the bias obervation. We will include these experimental results and the discussions in the appendix.

Q5: The estimation of the Monge map via barycentric projection introduces approximation error, especially when entropy regularization is non-negligible.

As stated in Appendix B.1 (Line 574-576, Page 18), we try to choose the entropy regularization coefficient $\varepsilon$ as small as possible to ensure it can approximate the exact optimal transport and monge map, while ensuring numerical stability and computational efficiency. Extensive experiments (Table 1, Table 2, Page 8) demonstrate the advantages of our VGFM despite the approximation error of exact optimal transport and Monge map. We acknowledge that theoretical research of this approximation remains an open and intereting direction, and we thank the reviewer for highlighting it.

Q6: On limitations

For the first point, we will consider adding visualizations of growth error, both spatially and per cell type, in the appendix for the simulated gene dataset. However, for real datasets, growth rates are generally not directly observable, which limits such an analysis.

For the second point, certain existing databases, such as UniProtKB, record genes that are potentially related to growth. We can leverage this information by comparing the learned $g$ with these known growth-associated genes, and incorporate such comparison as a regularization term during training to regularize or debias the growth modeling process.

We sincerely thank the reviewer for the positive comments and insightful questions.

Q1: It still seems unclear to me if the suggested model accurately reflects the potential behavior of genetic populations of cells. In particular, is it not possible for novel populations of cells to grow or appear outside the support of the current (or initial distribution) ? Or for whole populations to just disappear, and to have the remaining population be transported to the next snapshot?

Cell differentiation naturally causes transitions in the state space, which can lead to changes in both the support set and the population size at the next time point. Our proposed VGFM aims to model such dynamics.

It is a particularly challenging scenario that part of the population in a future snapshot does not originate from the first snapshot, but gradually emerges at some intermediate time from outside the support of the earlier distribution. From our current biological understanding, such a “spontaneous appearance” of an entirely novel population without any lineage connection is highly unlikely, as new cell populations generally arise through differentiation from existing ones. That said, if such a situation were to occur, we acknowledge that VGFM in its current form could not handle it, as it assumes that the population at the next snapshot evolves entirely from the previous one.

For the second scenario mentioned, where a certain population disappears while the remaining population transitions to the next snapshot, VGFM can indeed identify such disappearing populations via the semi-relaxed OT plan, as these cells would have vanishing future weights. However, since such cells are not sampled during training, their growth rates and velocity fields could not be learned. Handling this case may require additional constraints, such as enforcing the velocity field to be zero and the growth rate to be very small in those regions, or incorporating biological priors to determine the likely fate and disappearance time of such populations.

Q2: Is there a problem with non-differentiability in the pipeline?

We appreciate the reviewer’s careful observation. We now clarify how the $W_1$ loss is computed. In practice, we use the ot.emd() function from the pot library to compute the transport plan $\pi$ in $W_1$. This function solves for $\pi$ using a network flow algorithm, which is not compatible with PyTorch’s automatic differentiation. However, the gradients of the cost function $c$ can still be backpropagated. This strategy is also adopted in the implementations of Zhang et al. (ICLR 2025) and Huget et al. (NeurIPS 2022). Therefore, the gradient of $W_1$ w.r.t. network parameters is partial rather than full. Nevertheless, this does not affect the effectiveness of our method nor the feasibility of computation, as demonstrated by our results (see sec. 4). Furthermore, in the rebuttal to reviewer kHRn (Please refer to Q3 of response to reviewer kHRn), we also explored replacing $W_1$ with the Sinkhorn divergence (Feydy et al., AISTATS 2019), which is fully differentiable and enjoys many favorable properties. The Sinkhorn divergence and $W_1$ yield similar results. The corresponding results and discussion are provided in our response to Reviewer kHRn.

Q3: For Eq. 5, is there a reason you couldn't incorporate the energy from A.4 in order to get a unique specification for $g_t$?

The static semi-relaxed optimal transport in Eq. 4 is a strictly convex optimization problem: the first term is linear, the second term $\mathrm{KL}(\cdot \| q)$ is convex, and the constraints are linear. Therefore, the optimal transport plan $\pi^\ast$ is unique, which in turn makes $p_\lambda$ unique. Given the additional constraint in A.4, we can then uniquely determine $g_t$. We aim develop an dynamic understanding in Eq. (5) of the static semi-relaxed OT in Eq. (4). Directly adding the $L^2$ energy term into the objective function of Eq. (5) would break its equivalence with Eq. (4). Instead, we treat it as a constraint to select $g_t$.

Q4: Will an implementation be released?

Yes. We are currently organizing and cleaning up the code, and it will be released publicly upon acceptance of the paper.

Q5: Introducing distribution fitting loss needs simulations.

Indeed, our method is not fully simulation-free due to the additional distribution fitting loss, we stress the importance of distribution fitting loss:

Necessity of distribution fitting loss: Training the flow matching model only at the level of velocity and growth regression without an explicit global distribution-level fitting constraint may lead to inaccuracy distribution learning. Specifically, since the theoretically optimal velocity $v\_\theta(x,t)=\mathbb{E}\_{z|(x,t)}v_\theta(x,t|z)$ in flow matching, there may exist a possible velocity average due to the expectation, and the expected velocity may deviate from the real velocity. Introducing the distribution fitting loss will enforce the global distributional match constraint over the evolved population of cells.

Computational Efficiency Comparison: As the ablation study (Table 3 in Sect. 4.1, Page 9) shows, the computation cost needed for VGFM (13 mins) is much less than the totally simulation-based method DeepRUOT (90 mins), although VGFM sacrifices some computation efficiency compared with the totally simulation-free VGFM (w/o $\mathcal{L}\_{\rm OT}$). At the same time, VGFM achieves the best results on the dynamics reconstruction among the totally simulation-based and totally simulation-free methods.

We sincerely appreciate your valuable and insightful suggestions. We have conducted additional experiments and will include results and discussions in our manuscript and appendix.

Q1: The scalability of VGFM on real-world high-dimensional gene-space datasets.

To further explore the scalability of VGFM, we applied our model and compared methods explicitly modeling $g\_t(x)$ (Sha et al., 2024; Zhang et al., 2025) to the Pancreas dataset with 1000 highly variable genes.

• Stable training loss curve. We observed that VGFM is the only method showing a steadily decreasing training loss, both for $\mathcal{L}\_{\mathrm{VGFM}}$ and $\mathcal{L}\_{\mathrm{OT}}$. We compute $W\_1$ at day 15.5 between real data and generated data from VGFM and its variant VGFM (w/o) $\mathcal{L}\_{\mathrm{OT}}$, as shown in the table below:

Table r2-1: $W_1$ at day 15.5 between real data and generated data from VGFM and its variant.

This indicates that both loss terms are still effective in high-dimensional real-world data.

• Analysis on mean and variance. We calculated the means and variances of the real and generated data at day 15.5 and plotted the corresponding trend. The results show that the generated samples closely follow the mean–variance trend of the real data. To present our results, we provide a table showing binned statistics, summarizing the distribution of real and generated data across different mean and variance intervals.

Table r2-2: Distribution of real and generated data across different mean and variance intervals.

As shown in the table, the distribution of means in the generated data exhibits a high degree of consistency with that of the real data, while the variances are generally slightly smaller than those of the real data but still remain reasonably accurate.

• Analysis of the sparsity. Our current paper is applied to PCA-projected single-cell data, with the main contributions lying in the theoretical formulation and algorithmic design for the joint dynamic modeling of $g_t$ and $v_t$, from which we derive our flow-matching-based approach. In its present form, VGFM does not explicitly impose sparsity constraints on the data, generating data not as sparse as the real data. Nevertheless, recent studies have introduced encoder–decoder frameworks that operate in a latent space while explicitly modeling single-cell gene expression properties such as sparsity and the negative binomial distribution (Palma et al., 2025). By integrating VGFM into such a framework, for example, using advanced encoders to map raw count data into a latent space and appropriately designing decoders to reconstruct gene counts, one could learn the growth function from OT plans in the latent space with more accurate metrics beyond the Euclidean cost, while simultaneously preserving biologically relevant statistical properties. We believe this direct extension to raw data with explicit sparsity and negative binomial modeling is a promising direction, deserving thorough investigation in future work. We will cite these suggested papers and discuss them in the manuscript.

• Computing fate probabilities. We follow Eyring et al. (2024) to compute the averaged fate probabilities of each ground truth transition. Notably, in almost each transition, VGFM ranked among the top three performers, despite operating directly in the 1,000-dimensional gene expression space, whereas all other methods were evaluated in a 10–50-dimensional PCA space, as done in Eyring et al. (2024). This might indicate that VGFM can achieve biological interpretability even in high-dimensional settings.

Table r2-3:Averaged fate probabilities of each ground truth transition.

• Interpretable learned growth function. We highlight that our main contribution lies in modeling and training the growth function $g_t(x)$, with VGFM in the flow matching framework, to leverage snapshot mass change for learning $g_w(x,t)$. During training, the growth loss rapidly converges to a negligible value, much faster than the velocity loss. To further investigate $g_w(x,t)$, we follow Sha et al. (2024), for each time $t$, compute the element-wise absolute value $\left|\frac{\partial g(x,t)}{\partial x_i}\right|$ in the gene space for every cell, and average across cells to quantify each gene’s contribution to growth dynamics. From our analysis of $g_w(x,t)$, the three genes with the largest absolute gradients are Pnliprp1, Clps, and Ctrb1. These genes are well-known markers of Acinar cells and are highly expressed in the exocrine pancreas. Notably, at this time point, cells from the non-endocrine branch, particularly Acinar cells, are much more abundant at later stages due to high proliferation rates. This alignment between the learned genes and known biological processes indicates that VGFM captures growth dynamics in a biologically interpretable manner, successfully linking the learned growth function to meaningful cell-type-specific proliferation patterns.

Q2: The motivation behind the distribution-fitting loss.

Our ultimate goal is to generate data close to the real data distribution. Although flow matching is a powerful generative modeling framework, its learned velocity field is the expectation of the conditional velocity field, i.e.,$v\_\theta(x,t)=\mathbb{E}\_{z|(x,t)}v_\theta(x,t|z)$， where $z=(x_0,x_1)$ is the entropic semi-relaxed ot plan in our setting, which is a soft coupling. For the same $x_0$, $x_1$ may come from different branches of the trajectory. This bias could accumulate during generation, possibly leading to deviation from the target distribution. To address this bias, we introduce a distribution-fitting loss that directly measures how well the generated samples match the empirical data distribution at each time point. As done in the ablation study (Page 9, Table 3) on loss terms in our paper, this complements the flow-matching objective and ensures that even if the velocity field is biased, the model still adjusts its dynamics to match the correct data distribution.

Q3: The biological significance of explicitly modeling growth effects.

We appreciate the reviewer’s positive feedback on our methodological choice of explicitly modeling growth effects. From a biological perspective, cell development is inherently a coupled process of state transitions and mass changes. Ignoring the growth component may lead to erroneous state transition inference, for example, reconstructing cross-branch trajectories that are biologically implausible, as we demonstrated in the simulation gene dataset experiment in our paper. Moreover, modeling growth explicitly enables downstream biological interpretation. By computing element-wise absolute value $\left|\frac{\partial g(x,t)}{\partial x_i}\right|$ in the gene space, one can identify growth-related genes and proliferative cell populations, as done in Sha et al.,2024. However, their approach relies on autoencoder-based dimensionality reduction. In our response to Q1, we have analyzed our learned $g_w(x,t)$, which is trained directly on gene space, on the Pancreas dataset. We successfully identified key genes such as Pnliprp1, Clps, and Ctrb1, which are known to be associated with proliferating Acinar cells. This demonstrates that VGFM also yields biologically meaningful insights into cell population dynamics, even in high-dimensional settings where previous methods face limitations.

Q4: How VGFM could be made conditional on covariates.

A natural approach is to encode cell type as a one-hot vector $c$ and integrate it into both $v_\theta(x,t,c)$ and $g_w(x,t,c)$. The main challenge lies in determining the coupling between cells of the same type across two time points. This becomes considerably easier if prior knowledge of the transition fate for each cell type is available. While VGFM can identify proliferating cell types (as demonstrated in the Pancreas dataset), incorporating such prior knowledge could improve model accuracy. Moreover, VGFM could be extended to model perturbation responses conditioned on different drugs or gene perturbations. These embeddings could be generalized to unseen perturbations via the mode-of-action framework (Bunne et al., 2022), enabling VGFM to predict which perturbations might facilitate or inhibit cell growth.

Q5: Key dataset details.

Thanks for your suggestions, we will add this table to our appendix.

Table r2-4: Key dataset details.

We sincerely appreciate your valuable and insightful suggestions.

Q1: Choice of distribution fitting loss diverges from simulation-free training.

Necessity of distribution fitting loss: Training the flow matching model only at the level of velocity and growth regression without an explicit global distribution-level fitting constraint may lead to inaccuracy distribution learning. Specifically, since the theoretically optimal velocity $v\_\theta(x,t)=\mathbb{E}\_{z|(x,t)}v_\theta(x,t|z)$ in flow matching, there may exist a possible velocity average due to the expectation, and the expected velocity may deviate from the real velocity. Introducing the distribution fitting loss will enforce the global distributional match constraint over the evolved population of cells.

Computational Efficiency Comparison: As the ablation study (Table 3 in Sect. 4.1, Page 9) shows, the computation cost needed for VGFM (13 mins) is much less than the totally simulation-based method DeepRUOT (90 mins), although VGFM sacrifices some computation efficiency compared with the totally simulation-free VGFM (w/o $\mathcal{L}\_{\rm OT}$). At the same time, VGFM achieves the best results on the dynamics reconstruction among the totally simulation-based and totally simulation-free methods.

Q2: Influence of batch size.

We clarify that our VGFM approach did NOT use mini-batch OT either in the computation of the OT plan or the distribution fitting loss. Instead, we propose an offline sampling strategy for OT plan computation and its improved version tailored for large-scale datasets.

Offline sampling strategy: The OT plan between the adjacent snapshots for the whole population is precomputed and cached before training, as Algorithm 1 (Page 7) shows. Then, in training of the flow matching, we sample a batch of paired timestep samples using this precomputed OT plan, alleviating the sampling bias. This strategy needs the computation of OT for the whole dataset. We further propose the Big Batch strategy in Page 22, Line 668. Specifically, we uniformly partition the data at each time point into $n$ subsets as Big Batches. Then we precompute and cache the $n$ OT plans for these Big Batches and randomly select one Big Batch for constructing batches in the training process. During training the velocity field and growth function, we don't have to compute any ot plan, one just need to sample batch $(x_0,x_1)$ from the precomputed ot plan(s) $\pi$ (to train $v$) and sums up the corresponding rows of $x_0$ (to train $g$) (See Eq.(12), Page5). Since the pre-computed plan(s) is on full sample or Big Batches and is fixed during training, we have better training stability compared with other mini-batch ot approaches.

In the following Tables r1-1,r1-2,r1-3,r1-4, we report the sensitivity analysis on $n$ (Table r1-1: number of big batches in precomputing the OT over dataset), batchsize (Table r1-2: batchsize in training loss $\mathcal{L}\_{\mathrm{VGFM}}$ ), sample size (Table r1-3: sampling size in distribution fitting loss $\mathcal{L}\_{\mathrm{OT}}$ ), $n$ (Table r1-4: number of batches in computing distribution fitting loss $\mathcal{L}\_{\mathrm{OT}}$ ), on EB 50D dataset under the same experiment setting of ablation study (Table 3, Page 9).

Table r1-1: Sensitivity analysis of $n$ on EB dataset with metrics (computation time and $W_1$/RME at each time).

Table r1-2: Sensitivity analysis of batch size in $\mathcal{L}\_{\mathrm{VGFM}}$ on EB dataset with metrics ($W_1$/RME at each time t).

Table r1-3: Comparison of full sampling and mini-batch sampling in $\mathcal{L}\_{\mathrm{OT}}$ on EB dataset with metrics ($W_1$/RME at each time t).

Table r1-4: Results of using Big Batch strategy in $\mathcal{L}\_{\mathrm{OT}}$ to speed up computation on EB dataset with metrics ($W_1$/RME at each time t and the training time).

These results demonstrate the robustness of performance to the number of big batches in the offline sampling strategy (Table r1-1) and distribution fitting loss $\mathcal{L}\_{\mathrm{OT}}$ (Table r1-4), and the effect of batch size in $\mathcal{L}\_{\mathrm{VGFM}}$ and sampling size in computing distribution-fitting loss $\mathcal{L}\_{\mathrm{OT}}$ (Tables r1-2/3) on the final performance.

Q3: Clarification on if relying on the Sinkhorn algorithm and exploration of Sinkhorn divergence.

We compute the distribution fitting loss $W_1$ in Section 3.4 using EMD distance by pot library, following MIOFlow [Guillaume et al., NeurIPS22] and DeepRUOT [Zhang et al., ICLR25]. Particularly, we compute the OT using the pot libary, using the function pot.emd() and only backpropagate the gradient with respect to the cost $c$ in $W_1$ since $\pi$ is solved using a network flow algorithm, which is not compatible with PyTorch’s automatic differentiation. We will make a clear statement in Appendix.

As suggested, we replace the EMD distance with Sinkhorn divergence using CUDA/C++ based geomloss library. Table r1-5 verified the effectiveness of Sinkhorn divergence for the distribution fitting loss.

Table r1-5: comparison of EMD and sinkhorn divergence (with different entropic regularization parameters) on EB dataset with metrics ($W_1$/RME at each time step).

It can be found that the sinkhorn divergence achieves further improvement on the accuracy of dynamics reconstruction, and costs less time for computing the distribution fitting loss. We will include this result of sinkhorn divergence in the appendix and cite it in the discussion section of the manuscript.

Q4: Motivation for time-independent growth function.

The selection of time-independent growth function is because of

• First, it minimizes $L^2$ energy potential. As we have discussed in Appendix A.4, time-independent growth function is not only easy to implement but also minimizes the $L^2$ energy potential, which satisfies the principle of least action.
• Second, it can be explained by the Malthusian growth model, i.e., the exponential growth model. The Malthusian growth model is $\frac{\mathrm{d}p_t}{\mathrm{d}t}=gp_t$, where $p_t$ is the population size and $g$ is a constant growth rate. Our model on the growth $\frac{\partial p_t}{\partial t}=g_tp_t$ (refer to Eq.5, Page 4, Line 145-146) is consistent with the above Malthusian growth model. In these models, g is treated as a constant based on the assumption that the resources are abundant and the environment is stable, causing growth rates to remain relatively stable over time. In the context of scRNA-seq experiments, these conditions are typically satisfied because cells are often cultured or sampled under controlled laboratory conditions, where nutrient supply, temperature, and other environmental factors are maintained at constant levels. Thus, choosing this form is biologically reasonable, especially when there is no prior knowledge, such as the maximum carrying capacity of the environment.

Q5: On notations.

Thanks for your suggestion. Upon revisiting our paper, we realized that the discrete entropic semi-relaxed optimal transport formulation in Eq. (11) can be removed, as it is a standard result. Accordingly, we can directly use $\pi^*$ in place of $\pi^{0\to1}$. In addition, we can describe the barycentric mapping and $N_i$ in plain language (line 188) for better clarity. We will also carefully review the manuscript to refine the notations and formulations throughout.

Dear Reviewer kHRn,

We would like to thank you again for the valuable comments and suggestions. We have responded to the comments in detail and hope to address your concerns. As the discussion period is drawing to a close, we kindly remind you that if you have other concerns, please let us know. We will try our best to clarify further.

Best regards,

The Authors