Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action

Thank you for carefully reviewing our paper and for providing valuable comments. Below please find our point-by-point response:

Summary & Strengths

We appreciate your recognition of our work's strengths. Our paper derives the first-order optimality conditions for the RUOT problem and solves it by fitting a single scalar field, thereby modeling unbalanced distributions and stochastic dynamics. Compared to previous methods, our Var-RUOT finds paths with lower action, achieves stable training, and converges faster. We also discuss the biological priors associated with the RUOT penalty function. With each action guiding its corresponding dynamics, we hope our framework offers a more physical approach to modeling single-cell omics data.

[W1, Q3] On the Selection of $g^\frac{2}{15}$

Thank you for your question regarding $\psi(g)$. To clarify, we use $g^{2/15}$ as an example to show that when $\psi'(g)>0$ and $\psi''(g)<0$, $g(\boldsymbol{x},t)$ increases along $\boldsymbol{u}(\boldsymbol{x},t)$ (Theorem 4.2). We chose $\psi_2(g)=g^{2/15}$ because:

1. It is the simplest form that meets $\frac{d\psi(g)}{d|g|}>0$, $\psi(g)=\psi(-g)$, and $\psi''(g)<0$.
2. From the optimality condition we obtain

For a more stable training process, we keep $g(\lambda)$ from being too steep near $\lambda=0$. So we set $p=1$ and choose a large $q$ (i.e., $q=7$).

Note that $g(\lambda)$ is undefined at $\lambda=0$ due to the necessary discontinuity in $\psi'(g)$ at $g=0$. To address this, we linearly interpolate between $g(-\delta)$ and $g(\delta)$ (Appendix A.2). This singularity causes slowing convergence and degrading distribution reconstruction when using $\psi_2(g)$. A more reasonable way to deal with this problem is to fit the reciprocal of $\lambda$, which is our future work.

Empirically, on the 2D Mouse Blood Hematopoiesis dataset (results are shown below), $\psi(g)=g^{2/15}$ performs similarly to $\psi(g)=g^{2/21}$, while $\psi(g)=g^{2/9}$ (with a steeper $g(\lambda)$) leads to unstable training and poorer reconstruction. Our experiments demonstrate that choosing the exponent on $g$ within a certain range is feasible. We will update the paper to include the above analysis in the sensitivity analysis section of the Appendix.

As noted in our Limitations, choosing $\psi(g)$ remains open. Future work may derive it from physical principles (e.g., Branching SDE, Stochastic Thermal Dynamics) or learn it with a neural network. We will include these discussions in the Appendix.

[Q1] On Weight Initialization

Using particle methods, we approximate the true measure $\mu$ with the empirical measure $\mu^N$. We initialize by placing a particle at every $t=0$ data point with the uniform weight 1. This yields

which justifies the rationale to set $w_i(0)=1$.

[Q2] On the Ground Truth for Mass Matching

In RUOT modeling, because cell counts vary (due to growth or death), the number of data points differs over time. If we set the total mass at $t_1=0$ as the reference number 1, then the mass at time $t_i$ as $\frac{N_i}{N_1}$ (with $N_i$ as the count at $t_i$), which serves as the ground truth for mass matching loss. We will clarify this in our paper.

[Q4]On the Schrodinger Bridge Problem

We apologize for the confusion. The original Schrodinger Bridge is dynamic—minimizing the KL divergence over path space between a stochastic process and a reference one. In our paper, "static" means some method[(Tong, A.,et al) S2FM] first compute a static OT solution via entropy-regularized OT, then fit the velocity field using Conditional Flow Matching. We will clarify this in the paper.

[W4,5,6] Additional Baselines for Existing Datasets

Thanks for the great suggestion. Follow your advice, we supplemented three baselines on existing datasets; in addition, we provided results for nine baselines and Var-RUOT on two new datasets. We first supplemented the additional baselines on existing datasets with:

• Wasserstein Gradient Flow (WGF in the table) [(Neklyudov, K., et al.) Wasserstein Lagrangian Flows]
• OTCFM [(Lipman, Y., et al.) Flow Matching]
• UOTCFM [(Eyring, L., et al.)Unbalancedness in Neural Monge Maps]

Specifically, for the 2D Mouse Blood Hematopoiesis dataset, we have added three previously unevaluated baselines: MIOFLOW, S2FM, and PISDE.

Simulation Dataset

EMT Dataset

2D Mouse Dataset

Next, we show experiments on additional datasets. We use a 100-dimensional, 5-time-point ($t=0,1,2,3,4$) EB dataset [(Moon, K.R., et al) Visualizing structure and transitions...] and the NeurIPS challenge dataset [(Luecken, M.D., et al.) A sandbox for prediction and integration...]. For NeurIPS dataset, we extract gene expression data from donor 28483, cluster them via K-Means by pseudo time, and order the clusters (by centroid) into $t=0,1,2,3,4$. This new dataset is then used for further tests. The results are as follows:

EB Dataset

NeurIPS Dataset

And we compared the path action learned by TIGON, DeepRUOT and VarRUOT:

In summary, the experiments show that:

• For low-dimensional data: On the EMT dataset, Var-RUOT achieves distribution reconstruction comparable to OTCFM and UOTCFM, and on the three-gene simulation and 2D Mouse Blood Hematopoiesis datasets, it outperforms the other baselines.
• For high-dimensional data: On the EB dataset, Var-RUOT yields results similar to TIGON and DeepRUOT, and outperforms simulation-free methods (OTCFM, UOTCFM) as well as Wasserstein Gradient Flow; on the NeurIPS dataset, it outperforms existing methods. Notably, compared to DeepRUOT and TIGON, Var-RUOT finds a path with significantly lower action.

Thus, the experiments confirm that Var-RUOT remains competitive on high-dimensional data, achieving good distribution reconstruction while finding lower-action paths. We will include these discussions on new datasets & baselines in the Appendix.

[W7] On the Figures

Thank you for your questions regarding our figures. For Figures 2 and 3:

• Crosses represent real dataset points, and circles represent data generated by our algorithm.
• Trajectory colors indicate particle mass, with brighter areas showing higher mass.

We will add legends and annotations for clarity.

Regarding the EMT results, Figure 3 (a) shows DeepRUOT produces curved trajectories (even passing through sparse top regions), while Figure 3 (b) shows nearly straight lines connecting the start and end points. Since straight lines minimize action, Table 2 compares the Path Action of DeepRUOT and Var-RUOT, with Var-RUOT achieving significantly lower action.

The growth pattern in Var-RUOT emerges naturally from using the WFR Metric (Theorem 4.2), where the growth rate increases along the trajectory. Without a "minimum action" objective, many valid velocity fields and growth patterns could reconstruct the marginals. Both methods correctly reconstruct the marginals, while Var-RUOT finds a lower-action solution, which aligns with our objective.

We clarify that our goal in presenting these figures is not to show that Var-RUOT has a significantly lower reconstruction loss, but to demonstrate that both methods perform similarly in marginal reconstruction. Var-RUOT’s advantages lie in it obtains a lower-action solution, simpler scalar field fitting, faster convergence, and the elimination of the pretraining phase in DeepRUOT.

[W2,3,8,9]Other Writing Issues

Thank you for your valuable suggestions. We will revise the paper to:

• Define $\alpha$ at line 103 as the penalty coefficient for growth/death.
• Clarify in lines 248–250 and 286–288 that the $\mathcal{W}_1$, $\mathcal{W}_2$ distances and Path Action are relative indicators, with Var-RUOT performing better relative to the baseline.
• Add the unit (seconds) for Wall Clock time in Table 3.
• Add time index to time-dependent variables.

Dear authors,

Thank you again for your final remarks. Everything sounds reasonable. I confirm my final score of 4.

All the best!

Thank you for carefully reviewing our paper and raising valuable questions. Below we address your concerns in a point-to-point response:

Summary & Strengths

We appreciate your recognition of our work's strengths. Our paper derives the first-order optimality conditions for the RUOT problem and solves it by fitting a single scalar field, thereby modeling unbalanced distributions and stochastic dynamics. Compared to previous methods, our Var-RUOT finds paths with lower action, achieves stable trainin, and converges faster. We also discuss the biological priors associated with the RUOT penalty function. With each action guiding its corresponding dynamics, we hope our framework offers a more physical approach to modeling single-cell omics data.

[W1, W2, Q1] Differences from Previous Methods

Thank you for your question and sorry for the confusion. First, we sincerely apologize for not thoroughly checking our citation positions and misunderstandings caused by it. We fully agree that Action Matching is a very innovative work uses a single network to parameterize $s_{\theta}(\boldsymbol{x}, t)$ that inspires our research; we will review the manuscript to correct any improper citations and clarify the significance of previous methods.

Following your great suggestion, we also made a table to clearly highlight the contribution of Var-RUOT that will later be included in the revised Appendix:

• Unbalanced Distribution indicates that the method can handle non-normalized distributions.
• Stochastic Dynamics means the method supports stochastic dynamics.
• Unbalanced + Stochastic shows the method can handle both simultaneously (e.g., DeepRUOT).
• Least Action signifies that the method employs a first-order condition to minimize the action, rather than simply adding the action as a weighted loss.

The key difference between Var-RUOT and DeepRUOT is that Var-RUOT leverages first-order optimality conditions in its parameterization and loss design (e.g., using only a scalar field rather than three neural networks as in DeepRUOT). Compared to DeepRUOT, our approach finds paths with lower action and addresses unstable training and slow convergence, thus eliminating the need for a pre-training phase.

Var-RUOT and Action Matching differ in three key ways:

Firstly, Var-RUOT can simultaneously handle imbalanced data distributions and stochastic dynamics, both crucial for single-cell omics.

Secondly, Action Matching's training algorithm needs to sample from intermediate distribution $q_{t}(x)$, while in single-cell omics data, we only have snapshot data at a few time points as boundary sample. So acquiring intermediate samples from $q_{t}(x)$ may be challenging due to missing evolution data in-between, potentially hindering the ability to capture complex dynamics. Empirically, our experimental results (Tables 1, 2, 4) show that Var-RUOT achieves more satisfactory results over existing results including Action Matching.

Lastly, Action Matching defines the Action Gap loss as:

$L_{\text{AM}} = E_{q_{0}(x)} [s_{0}(x)] - E_{q_{1}(x)} [s_{1}(x)] + \int_{0}^{1} E_{q_{t}(x)} \left[ \dfrac{1}{2} \| \nabla s_{t}(x) \|^{2} + \dfrac{\partial s}{\partial t}(x) \right] \mathrm{d} t$

While Var-RUOT directly adopts the HJB Loss(Line 204-209) to ensure that the computed solution $\lambda(\boldsymbol{x},t)$ satisfies the HJB equation. In other words, Var-RUOT incorporates the first-order condition into the loss, enabling it to find a solution with lower action. We will add these discussions in the revised manuscript.

[Q2] On the Estimation of the Action Loss

Thank you for your question. In our experiments, we employed two methods to compute the integral losses for both the action and HJB terms. When using the Euler-Maruyama method, setting the step size too large (e.g., $\Delta t = 0.2$) results in high variance estimates for the action and HJB losses, leading to unstable training. Therefore, we use $\Delta t = 0.1$ in all experiments. The Stochastic Runge-Kutta method offers better convergence, allowing an integration step of $\Delta t = 0.2$. Since our approach is simulation-based, a longer step size significantly reduces training time. In practice, we thus choose the Stochastic Runge-Kutta method and select a larger step size as long as training remains stable. We will discuss this in our final version and cite [(Berman.) High Order Action Matching] in our paper.

[Q3] Why the Related Steps in the Original RUOT Algorithm Can Be Omitted

In the original DeepRUOT algorithm, three neural networks are employed to approximate $\boldsymbol{v}(\boldsymbol{x},t)$, $g(\boldsymbol{x},t)$, and $\log \rho(\boldsymbol{x},t)$, with all constraints imposed as soft constraints via loss functions. For instance, to ensure that $\boldsymbol{v}$, $g$, and $\rho$ satisfy the Fokker–Planck equation, DeepRUOT introduces the Fokker–Planck loss

$L_{\text{FP}} = \|\partial_{t} p_{\theta} + \nabla_{x} \cdot (p_{\theta}\boldsymbol{v})- g_{\theta} p_{\theta}\| + \lambda_{w} \| p_{\theta}(\boldsymbol{x},0) - p_{0} \|$

Training these three interdependent networks simultaneously is challenging and often causes convergence issues. To accelerate training by obtaining a good initialization, DeepRUOT incorporates a pretraining phase: (1) training $\boldsymbol{v}$ and $g$ using a reconstruction loss ${L}_{\text{Recons}}$, and (2) training $p$ with Conditional Flow Matching.

In Var-RUOT, we use a single neural network $\lambda_\theta(\boldsymbol{x},t)$ to approximate a scalar field. The velocity field $\boldsymbol{u}(\boldsymbol{x},t)$ and growth rate $g(\boldsymbol{x},t)$ are derived analytically from $\lambda(\boldsymbol{x},t)$ using first-order optimality conditions (Theorem 4.1). As a result, Var-RUOT avoids handling multiple interdependent networks and eliminates the need for the pretraining phase used in DeepRUOT. This is indeed one of the main contribution of our work; ablation experiments in DeepRUOT demonstrate that these pretraining steps are indispensable (DeepRUOT[(Zhang.), DeepRUOT], Appendix B.5), and we have removed them in Var-RUOT. We will add these discussions in the revised manuscript.

[Q4] Why we Choose $g^{2/15}$

Thanks for your question regarding $\psi(g)$. we choose $g^{2/15}$ primarily to serve as an example showing that when $\psi'(g) > 0$ and $\psi''(g) < 0$, at each time step $g(\boldsymbol{x},t)$ increases along the direction of the velocity field $\boldsymbol{u}(\boldsymbol{x},t)$ (Theorem 4.2). As for why we adopt $\psi_{2}(g) = g^{2/15}$ as the example, it is because: (1) It is the simplest choice that satisfies the conditions given in our paper (lines 224–229): $\dfrac{\mathrm{d}\psi(g)}{\mathrm{d}|g|} > 0$, $\psi(g) = \psi(-g)$, and $\psi''(g) < 0$. (2) For training considerations: when choosing $\psi(g)=g^{(2p)/(2q+1)}$, the first-order optimality condition allows us to derive

We aim for more stable training, and since $g$ is computed directly from $\lambda$, we do not want the $g(\lambda)$ to be overly steep (especially near $\lambda=0$), which requires the exponent $\dfrac{2q+1}{(2q+1)-2p}$ to be as small as possible. This means choosing $p$ as small as possible (hence $p=1$) and choosing $q$ as large as possible (we select $q=7$). Therefore, we ultimately choose the penalty function $\psi(g)=g^{2/15}$.

*Note: It is worth noting that $g(\lambda)$ is undefined at $\lambda=0$. In fact, for any function $\psi(g)$ that meets the conditions outlined in lines 224–229 along with $\psi''(g)<0$, its first derivative must be discontinuous at $g=0$, so this singularity is inevitable. To handle this issue, we select a small value $\delta$ and linearly interpolate between $g(-\delta)$ and $g(\delta)$ by redefining

to bypass the singularity (Appendix A.2). Another approach to deal with the singularity is to fit the reciprocal of $\lambda$, denoted as $\mu$, thereby moving the singularity at $\lambda=0$ to $\mu \rightarrow \infty$. In practice, $\mu \rightarrow \infty$ does not occur, just as the singularity where $|g| \rightarrow \infty$ is not reached. This method avoids artificially modifying the $g(\lambda)$ relationship and is more reasonable. This will be our future work.*

To further explain our training considerations, we experimented with two additional parameter sets on the 2D Mouse Blood Hematopoiesis dataset. The experimental results are as follows:

The experimental results indicate that choosing $\psi(g)=g^{2/15}$ produces performance comparable to $\psi(g)=g^{2/21}$, while selecting $\psi(g)=g^{2/9}$ leads to unstable training and poorer reconstruction of the distribution. Our experiments demonstrate that choosing the exponent on $g$ within a certain range is feasible. We will update the paper to include the above analysis in the sensitivity analysis section of the Appendix.

On Writing Issues

Thank you for your suggestions regarding our writing. We will review and revise our paper to:

• Move the Limitations section into the main text.
• Remove the emoji from Figure 1.

Thank you for carefully reviewing our paper and for providing valuable comments. Below please find our point-by-point response:

Summary & Strengths

We appreciate your recognition of our work's strengths. Our paper derives the first-order optimality conditions for the RUOT problem and solves it by fitting a single scalar field, thereby modeling unbalanced distributions and stochastic dynamics. Compared to previous methods, our Var-RUOT finds paths with lower action, achieves stable training, and converges faster. We also discuss the biological priors associated with the RUOT penalty function.

[W1] Regarding Var-RUOT’s Guarantee for the Minimal Action Solution

Thank you for your comment. Since Var-RUOT is a deep learning–based RUOT solver, it guarantees only a local minimum of the action. Empirically, experiments (Table 1,2,4) along with all additional results in the rebuttal show that Var-RUOT achieves a lower action than DeepRUOT, and its Path Action is comparable to that of traditional OT solvers.

We consider guaranteeing a global minimum an interesting direction for future work and will discuss it in the revised manuscript. We plan to start with analytical solutions for simple cases (e.g. Dirac-to-Dirac transport) and then extend to real data by treating complex distributions as a superposition of Dirac deltas and solving the associated velocity field via a Flow Matching–like approach.

[W2] Convergence Rate

Thanks for your insightful question. From the perspective of numerical integration, we use the particle method to estimate $L_{\text{HJB}}$ and $L_{\text{Action}}$, which is equivalent to performing integration via Monte Carlo methods. Let the number of particles be $N$. Then, for any smooth function $f(\boldsymbol{x})$ (In the Action Loss and HJB Loss, the integrand.), the central limit theorem tells us that the error of the integral $\int f(\boldsymbol{x})\,\mathrm{d}\mu^N$ is of order $\mathcal{O}\left(\dfrac{1}{\sqrt{N}}\right)$. We will revise the paper and include this discussion on the convergence rate. In practical applications of the Var-RUOT algorithm, our default setting is $N=512$, which ensures that the estimates of both loss have very small variance.

[Q1, Q4] Regarding the effect of the HJB Loss.

Thanks for the insightful question. Var-RUOT controls the penalty for $\lambda(\boldsymbol{x},t)$ violating the HJB equation by adjusting the weight $\gamma_{\text{HJB}}$: a higher weight increases the importance of the HJB loss, reducing the violation after training. Sensitivity analysis (Section C.2, Figure 7) shows that as $\gamma_{\text{HJB}}$ increases, the Path Action decreases, proving that a lower HJB loss leads to a smaller action solution. To address the concerns, we further provide sensitivity results for the EMT dataset and additional results for the 2D Mouse dataset:

• EMT Dataset

• 2D Mouse Dataset

These results confirm that incorporating the HJB loss effectively finds lower action solutions. Based on our findings, we recommend setting $\gamma_{\text{HJB}}$ between $10^{-2}$ and $10^{-1}$ to balance minimal action with data reconstruction quality. We will include these discussion in our final version.

[W3,Q2] Reversibility Issue of $\psi'(g)$

Thank you for your valuable comment and apologize for the insufficient discussions in previous paper. Here we discuss the two scenarios of $\psi$ selection in more details:

• $\psi_{1}(g)=g^{2p}$, $\psi_{1}'(g)=2p\cdot g^{2p-1}$, $p\in\mathbb{Z}^{+}$. Here, $\psi_{1}'(g)$ is a bijection from $\mathbb{R}$ to $\mathbb{R}$, so $g$ is uniquely determined from $\lambda=\alpha\psi_{1}'(g)$.

• $\psi_{2}(g)=g^{\frac{2p}{2q+1}}$, $\psi_{2}'(g)=\frac{2p}{2q+1}g^{\frac{2p-(2q+1)}{2q+1}}$, $p,q\in\mathbb{Z}^{+}$ with $2q+1>2p$. Here, $\psi_{2}(g)$ is a bijection from $\mathbb{R}\setminus\{0\}$ to $\mathbb{R}\setminus\{0\}$ (undefined at 0). Directly solving$\lambda=\alpha\psi_{2}'(g)$ yields

At $\lambda=0$, $g(\lambda)$ is undefined. To handle this inevitable singularity, we set a small $\delta$ and linearly interpolate between $g(-\delta)$ and $g(\delta)$ for $\lambda\in[-\epsilon,\epsilon)$,i.e., $g(\lambda) = \dfrac{1}{2\delta}\left(g(\delta) - g(-\delta)\right)(\lambda + \delta) + g(-\delta)$ (Appendix A.2)

Another approach is to fit the reciprocal of $\lambda$, denoted as $\mu$, thereby moving the singularity at $\lambda=0$ to $\mu \rightarrow \infty$. In practice, $\mu \rightarrow \infty$ does not occur, just as the singularity where $|g| \rightarrow \infty$ is not reached. This method avoids artificially modifying the $g(\lambda)$ relationship and is more reasonable. This will be our future work. We will update the paper for more thorough disucssions on $\psi'(g)$ choice and include these detailed analysis in the appendix.

[Q3] Comparison with Traditional OT Solvers

Thank you for your valuable comment. In additional experiment, here we built two 2D normalized and unbalanced datasets to test Var-RUOT’s ability to find minimal action paths. For normalized data, we used the pot library to obtain ground truth; for unbalanced data, we applied a Sinkhorn-based static WFR solver [(Wang Z., et al.) Wasserstein-fisher-rao document distance]. With $\sigma=0$, Var-RUOT achieves similar—or even lower—Path Action values than traditional solvers(results are shown below), demonstrating its effectiveness. We will include these analysis in the appendix.

[Q4] Extrapolation Performance

We evaluated Var-RUOT’s extrapolation capability and error propagation by extending the three-gene simulation dataset to nine time points. The first five were used for training, while the remaining four tested extrapolation. We compared TIGON, DeepRUOT, and Var-RUOT. The results are shown below:

Var-RUOT outperforms TIGON and performs similarly to DeepRUOT in long-term extrapolation. However, as time increases, error grows because the model’s non-autonomous system learns $\lambda(\boldsymbol{x}, t)$ only within the training window. Beyond this range, the $t$-dependency is less accurate, leading to larger errors. Nonetheless, this non-autonomous system improves flexibility and helps prevent underfitting (see Q6). We will include these analysis in the appendix.

[Q5] Choice of $\psi(g) = g^{2/15}$

Sorry for the confusion. We clarify that choosing $\psi(g)=g^{2/15}$ primarily to showcase an example with $\psi''(g)<0$. There are two reasons.

Firstly, the simplest function family satisfying our conditions (lines 224–229)—that is, $\frac{d\psi(g)}{d|g|}>0$, $\psi(g)=\psi(-g)$, and $\psi''(g)<0$—is given by

Secondly, for training stability, the first-order condition yields

To avoid a steep $g(\lambda)$ curve near $\lambda=0$, we choose the smallest $p$ (i.e., $p=1$) and a large $q$ (here, $q=7$), which leads to $\psi(g)=g^{2/15}$.

Here we conduct additional experiments on the 2D Mouse Blood Hematopoiesis dataset (results are shown below) to show that while $\psi(g)=g^{2/15}$ and $\psi(g)=g^{2/21}$ perform similarly, however $\psi(g)=g^{2/9}$ yields unstable training and poorer reconstruction.

Overall, our experiments confirm that, within a reasonable range, different choices for the exponent on $g$ are feasible. We will include this sensitivity analysis in the appendix.

[Q6] On Whether a Single Scalar Field Causes Underfitting

Thanks for the insightful question. Here we clarify that:

• Experiments show that Var-RUOT reconstructs the marginal distributions at all time points well—both on low-dimensional data (Sections 6.1 and 6.3) and high-dimensional data (Appendix C.4), the results are comparable to DeepRUOT (which uses three networks) without underfitting. Please also refer to the responses to Reviewer kRW9 [Q1] and Reviewer YDV3 [W4,5,6], which include quantitative results of Var-RUOT on high-dimensional datasets.
• Allowing $\lambda$ to depend on both $\boldsymbol{x}$ and $t$ enables capturing both trajectory evolution and total mass fluctuations; if $\lambda$ depended solely on $\boldsymbol{x}$, it could not fit distributions showing some specific mass changing pattern (e.g., mass that first rises then falls).
• Moreover, the DeepRUOT method employs three networks, while they are interdependent (via a Fokker–Planck loss), which makes training unstable and requires pretraining, while Var-RUOT is simpler and more stable.

I thank the authors for their comprehensive rebuttal. All of my questions have been answered satisfactorily, and I am satisfied with the responses. Therefore, I am increasing my score to 5.

Thank you for carefully reviewing our paper and for providing valuable comments. Below please find our point-by-point response:

Summary & Strengths

We appreciate your recognition of our work's strengths. Our paper derives the first-order optimality conditions for the RUOT problem and solves it by fitting a single scalar field, thereby modeling unbalanced distributions and stochastic dynamics. Compared to previous methods, our Var-RUOT finds paths with lower action, achieves stable training, and converges faster. We also discuss the biological priors associated with the RUOT penalty function. With each action guiding its corresponding dynamics, we hope our framework offers a more physical approach to modeling single-cell omics data.

[W1] Relying solely on a PINN-like HJB loss to estimate the value function can produce significant control deviations, as even small errors cause large inaccuracies due to sensitive differentiation.

Thank you for your insightful question. We fully agree that since deep learning solver for the RUOT problem converts hard constraints into soft ones via the loss function, so there's no guarantee that the solution exactly satisfies the marginal constraints nor minimizes the action at every time point. Instead, the HJB loss serves as a regularizer that promotes a lower action value without necessarily reducing to 0. On the other hand, we point out that empirically, sensitivity analysis (lines 656–670) shows that increasing the loss weight weight $\gamma_{\text{HJB}}$ would decrease the action value, indicating the effectiveness of such term to achieve the decrease of action for the comnputed path. We will add the relevant discussions in the revised manuscript.

[W2] Compute Second-order derivatives causes computational burden.

Thank you for your comments and suggestion. We fully agree that compute second-order derivatives using PyTorch autograd by default will requires retaining a larger graph and thus more memory and time than first-order derivatives. Meanwhile we clarify that

• Unlike previous RUOT methods such as DeepRUOT that train three networks for $\boldsymbol{v}$, $\rho$, and $g$, VAR-RUOT fits a single scalar field $\lambda$, reducing the search space, improving stability, and converging faster (see Section 6.2).

• Since Var-RUOT only need second-order derivatives for the $\mathcal{L}_{\text{HJB}}$ loss (to compute $\nabla_x^2 \lambda$), meaning we only require the trace of the Hessian matrix $H$. We can estimate it using the Hutchinson method:

  $$\text{trace}(H)=\mathbb{E}_{\boldsymbol{v}}[\boldsymbol{v}^{T} H \boldsymbol{v}], \quad \boldsymbol{v}\sim N(0,I)$$

  which computes the derivative along one direction, reducing cost (error $O(1/\sqrt{m})$, with $m$ samples). We can choose to use Hutchinson early in training for speed and later switch to full autograd for accuracy. This feature will be made as an option in our open-source code to speedup the computation.

[Q1] Can Var-RUOT be Generalized to high-dimensional settings?

Thank you for your question on Var-RUOT's scalability to high-dimensional datasets. In previous manuscript, we validated its scalability on several datasets: 50 and 100 dimensional Gaussian Mixture datasets (lines 709–715), a 50-dimensional Mouse Blood Hematopoiesis dataset, and a 30-dimensional β-cell Differentiation dataset (lines 716–723). For the high-dimensional Gaussian Mixture dataset, Figure 13 qualitatively shows the learned particle trajectories and growth rates on gaussian datasets; Figures 14 and 15 display the learned velocity fields and growth rates for the real 50D and 30D datasets, demonstrating Var-RUOT's applicability.

To quantitatively evaluate scalability, we further designed simulated data at three time points (t = 0, 1, 2) where each point is generated from a mixture of three Gaussians in 150 dimensions. We then compared various baselines on this dataset, with results as follows:

Var-RUOT achieves comparable reconstruction accuracy to PISDE, slightly outperforms TIGON and DeepRUOT, and significantly outperforms other methods on the 150-dimensional Gaussian Mixture dataset. It also yields a significantly lower path action than TIGON/DeepRUOT, confirming its scalability in high-dimensional settings. We will include these experiments in our revised manuscript.

[Q2]Questions on hyperparameter tuning and ablation studies.

Thank you for your question on hyperparameter tuning. Our method requires tuning three hyperparameters: the WFR metric parameter $\alpha$ (controlling the penalty for mass growth/decay) and two loss parameters, $\gamma_{\text{HJB}}$ and $\gamma_{\text{Action}}$. In Appendix C.2, we show that our method is robust to variations in $\alpha$ (using the 2D Mouse Blood Hematopoiesis dataset, lines 648–655, Table 7). Similarly, sensitivity experiments on $\gamma_{\text{HJB}}$ and $\gamma_{\text{Action}}$ (lines 656–670, Figures 7–8) demonstrate that increasing either parameter reduces the Path Action. However, very high values diminish the impact of the reconstruction loss that enforces marginal constraints, which degrades performance.

To provide further insights, here we have conducted additional sensitivity tests on $\gamma_{\text{HJB}}$ and $\gamma_{\text{Action}}$ using the Mouse Blood Hematopoiesis dataset.

• Case with fixed $\gamma_{\text{Action}} = 6.25\times {10}^{-2}$:

• Case with fixed $\gamma_{\text{HJB}}= 6.25\times {10}^{-2}$:

The results from these additional sensitivity experiments are consistent with those presented in the original manuscript: higher values of $\gamma_{\text{HJB}}$ and $\gamma_{\text{Action}}$ lead to a lower Path Action, yet if they are set too high, the reconstruction quality deteriorates significantly. In fact, all experiments in the main text use the same parameters, $\gamma_{\text{HJB}} = \gamma_{\text{Action}} = 6.25 \times 10^{-2}$ (see Table 6). In practice, we recommend setting both parameters in the range of $10^{-2}$ to $10^{-1}$ and we will make it clear in the open-source package.

[Q3]Training curve for the HJB loss.

Thank you for your question. Due to rebuttal restrictions, we cannot provide an image of the HJB loss curve. The HJB loss initially rises for 10–20 epochs (since $L_{\text{Recon}}$ dominates while parameters are near zero) and then gradually declines as $L_{\text{Recon}}$ reaches a similar scale as $L_{\text{HJB}}$ and $L_{\text{Action}}$. We will include this curve in the appendix of revised manuscript.

[Q4]Can we use L1 loss instead?

Thank you for your valuable suggestion. In additional experiment, we have modified the HJB loss per your recommendation by replacing the quadratic penalty with a linear penalty. Then we re-ran experiments on three datasets presented in the main text for the Var-RUOT method, comparing the performance when using the original $L_{2}$ loss versus the $L_{1}$ loss.

• Three Gene Simulation Dataset:

• EMT Dataset:

• 2D Mouse Blood Hematopoiesis Dataset:

The experimental results indicate that with our chosen hyperparameter values ($\gamma_{\text{HJB}} = 6.25 \times 10^{-2}$ and $\gamma_{\text{Action}} = 6.25 \times 10^{-2}$), using either the $L_{1}$ loss or the $L_{2}$ loss for the HJB loss yields similar reconstruction performance and Path Action. We will add the above analysis to the appendix during the revision of the paper.