PaperlayerView exact revision on arXiv

Source-linked AI summary

Simulation-free Unbalanced Dynamic Optimal Transport with General Growth Penalty

Junda Ying, Yuxuan Wang, Bowen Yang, Peijie Zhou, Lei Zhang

arXiv:2609.04710v1cs.LGcs.AIq-bio.QM

TL;DR

Unpaired single-cell snapshots require methods that jointly model cellular movement and population growth, but existing UDOT solvers trade flexible penalties against computational efficiency. SUDO learns conditional paths and transport costs to solve UDOT simulation-free with general convex penalties, matching analytical WFR accuracy while improving scalability over simulation-based methods. Its conclusions identify convex penalties as the non-degenerate regime, while the method currently assumes deterministic dynamics and a pre-specified growth penalty.

  • Problem

    Existing UDOT solvers either rely on computationally expensive NeuralODE simulations or analytical conditional paths restricted to quadratic penalties, limiting efficient use of general biological growth priors.

  • Method

    SUDO learns conditional paths and transport costs, solves the induced semi-coupling, and uses unbalanced flow matching for simulation-free UDOT with general growth penalties.

  • Results

    SUDO recovers analytical solution-based accuracy on WFR problems and outperforms NeuralODE-based methods in scalability, while supporting general penalties for biological priors.

  • Takeaways & Limitations

    Convex growth penalties provide the relevant non-degenerate regime for coupled cellular dynamics, and selecting Ψ enables biological priors on proliferation and apoptosis.

  • Takeaways & Limitations

    SUDO is theoretically developed for deterministic dynamics and requires the growth penalty Ψ to be specified in advance.

Abstract

from arXiv · show

Inferring cellular dynamics from unpaired single-cell snapshots requires modeling both state transitions and population growth or death. Unbalanced dynamic optimal transport (UDOT) addresses this by penalizing growth along transport paths, making the choice of growth penalty a key way to encode biological priors on proliferation and apoptosis. However, existing UDOT solvers either rely on computationally expensive NeuralODE simulations or depend on analytical solutions of conditional paths, restricting their efficiency solely to quadratic penalties, i.e. Wasserstein-Fisher-Rao (WFR) geodesics. To enable an efficient UDOT solver for general growth penalties, we first show that concave growth penalties lead to degenerate solutions where growth and transport are separated. We then introduce \textbf{S}imulation-free \textbf{U}nbalanced \textbf{D}ynamic \textbf{O}ptimal transport (SUDO), a simulation-free framework for UDOT with general non-quadratic convex growth penalties. SUDO learns the conditional paths and transport costs, solves the induced semi-coupling problem, and subsequently leverages unbalanced flow matching to achieve a simulation-free solution. On WFR benchmarks, SUDO matches the accuracy of efficient, analytical solution-driven algorithms while outperforming simulation-based methods in computational speed. Beyond WFR, SUDO supports asymmetric penalties that encode proliferation-dominant priors and produce more plausible trajectories and growth estimates on synthetic and single-cell datasets.

1 Introduction

Single-cell snapshots require UDOT to model both state transitions and changing population size, but existing methods trade computational efficiency against flexible growth penalties. SUDO addresses this trade-off with a simulation-free framework for general growth penalties.

  • scRNA-seq provides static, unpaired snapshots, motivating distribution-based dynamic optimal transport for reconstructing continuous cellular trajectories.
  • UDOT extends dynamic optimal transport by penalizing population growth, addressing proliferation and apoptosis that strict mass conservation overlooks.
  • Existing UDOT methods either use expensive NeuralODE simulations or restrict simulation-free efficiency to quadratic growth penalties.
  • SUDO learns conditional paths and transport costs to solve UDOT simulation-free with a broad class of growth penalties.
  • Concave growth penalties produce degenerate solutions separating non-trivial growth from transport, motivating convex penalties for meaningful coupled dynamics.
  • SUDO recovers analytical solver accuracy on WFR problems while significantly outperforming NeuralODE-based baselines in scalability.

2 Related works

Related work develops dynamic, unbalanced, and simulation-free transport methods, but prior approaches leave a gap between flexible growth modeling and exact UDOT formulation.

  • Dynamic optimal transport represents continuous probability flows minimizing kinetic energy, extending static transport between distributions.
  • Flow matching learns probability paths interpolating observed marginals without simulation and has been extended to unbalanced and stochastic dynamics.
  • NeuralODE-based UDOT methods handle arbitrary growth penalties but are computationally expensive, while related flow-matching methods do not solve the exact UDOT problem.

3 Preliminaries

UDOT augments dynamic transport with a growth field and penalty, while unbalanced flow matching learns the resulting measure flow through conditional paths. Theoretical analysis identifies convex penalties as the relevant regime for coupled dynamics.

  • Optimal Transport and Extensions: Dynamic OT minimizes kinetic-energy transport under a continuity equation, whereas UDOT adds a growth rate to model changing total mass.
  • Optimal Transport and Extensions: The growth penalty Ψ controls deviations from preferred growth, and the quadratic choice defines the Wasserstein-Fisher-Rao formulation.
  • Unbalanced Flow Matching: Unbalanced flow matching parameterizes velocity and growth networks to learn a continuous measure flow interpolating endpoint measures.
  • Unbalanced Flow Matching: Conditional flow matching replaces intractable marginal quantities with conditional paths whose velocity, growth, and mass can recover the UDOT flow.
  • Growth Penalty as Biological Prior: General growth penalties encode biological preferences for proliferation and apoptosis beyond the symmetric quadratic penalty.
  • Growth Penalty as Biological Prior: For penalties concave on both sides of their preferred growth rate, UDOT separates non-trivial growth and transport; the analysis therefore focuses on convex penalties.

4 Travelling Dirac and Static Form

The travelling-Dirac analysis reduces endpoint transport and growth between two Diracs to a one-dimensional path problem, enabling a static semi-coupling formulation under suitable convexity conditions.

  • Travelling Dirac: The travelling Dirac is the optimal path between two Dirac measures, with prescribed endpoint positions and masses.
  • Travelling Dirac: Theorem 4.1 represents the travelling Dirac using one-dimensional functions depending on distance and mass ratio, with a straight spatial path.
  • Travelling Dirac: The theorem reduces the original high-dimensional travelling-Dirac optimization problem to a one-dimensional problem.
  • Static Form: The travelling-Dirac cost is homogeneous in endpoint masses and, for even penalties, symmetric.
  • Static Form: For uniformly convex penalties under a diameter condition, the travelling-Dirac cost is jointly convex and sublinear, enabling a static UDOT formulation.
  • Static Form: The semi-coupling specifies mass sent between endpoints, while travelling Dirac paths determine how that mass moves and grows through space.
  • Assumptions: The construction assumes uniform convexity and a bounded domain, with scaling or added quadratic growth used to satisfy these conditions in practice.

5 Simulation-free Training of General UDOT problem

SUDO replaces unavailable analytical travelling-Dirac solutions with learned conditional paths and costs, then solves the induced semi-coupling before unbalanced flow matching. Its training remains simulation-free through one-dimensional numerical integration and neural surrogate optimization.

  • Previous simulation-free UDOT methods require closed-form travelling Diracs, which are available only for quadratic growth penalties.
  • SUDO learns travelling-Dirac paths with neural networks parameterized by time, endpoint distance d, and mass ratio r.The parameterization learns the scalar path functions k(t) and l(t), avoiding the curse of dimensionality.
  • The learned path automatically satisfies endpoint boundary conditions through the parameterization xη(t,d,r) = x0 + (x1 − x0)(t + t(1 − t)φη(t,d,r)).
  • SUDO approximates transport costs over moderate grids of endpoint distance and mass ratio, using one-dimensional trapezoidal integration without trajectory-gradient computation.The bounded total mass justifies restricting the approximation to moderate ranges.
  • The induced semi-coupling is optimized by projected gradient descent using a neural surrogate for the true transport cost.
  • Finally, SUDO trains neural velocity and growth-rate fields with conditional unbalanced flow matching, without trajectory simulation.

6 Experiments

Experiments show that SUDO closely matches WFR-focused baselines while improving computational cost, and that general growth penalties can encode biologically motivated asymmetric proliferation and apoptosis priors.

  • 6.1 SUDO Recovers the WFR: SUDO matches WFR-FM across synthetic datasets and achieves comparable measure-matching performance while generally outperforming simulation-based UOT methods.Table 1 reports mean W1 and relative mass error on synthetic datasets, with mean and standard deviation over five runs where applicable.
  • 6.1 SUDO Recovers the WFR: SUDO delivers UDOT costs on par with WFR-FM across three synthetic datasets and the EMT scRNA-seq dataset.Both algorithms are compared against the static semi-coupling reference, and their costs are lower because the marginal measures are not precisely matched.
  • 6.2 Scalability: SUDO outperforms simulation-based solvers in computational cost on Mouse datasets ranging from 10,000 to 49,302 cells.The scalability evaluation measures training time and peak memory consumption against DeepRUOT and VarRUOT, with WFR-FM as an optimal reference.
  • 6.3 The Necessity of General Growth Penalty: An asymmetric only-growth penalty yields a trajectory with mild proliferation and approximately 700 migrating cells, matching the stated proliferation-dominant prior.
  • 6.3 The Necessity of General Growth Penalty: On mouse hematopoiesis data, WFR produces both apoptosis and proliferation, whereas only-growth produces purely positive growth rates and more realistic estimates.

7 Conclusion

SUDO solves UDOT with general growth penalties without simulation, while convex penalties avoid the degenerate coupling caused by concave penalties. Its limitations are deterministic theoretical formulation and the need to pre-specify the growth penalty.

  • SUDO is a simulation-free framework for UDOT with general growth penalties that balances penalty flexibility and computational efficiency.
  • Concave growth penalties separate non-trivial growth from transport, identifying convex penalties as the relevant regime for coupled cellular dynamics.
  • Different growth penalties encode biological priors on proliferation and apoptosis, supporting more faithful reconstructions of cellular dynamics.
  • SUDO is theoretically developed for deterministic dynamics and requires a pre-specified growth penalty.Learning or calibrating the penalty from partial trajectory supervision is identified as future work.

A.1 Proof of Theorem 3.1

The proof establishes that concave growth penalties produce degenerate UDOT solutions in which non-trivial growth and transport separate. An explicit construction shows that the optimal cost can collapse to zero through instantaneous mass decay and growth.

  • Concavity on both sides of g0 implies that the UDOT solution separates non-trivial growth and transport.
  • An explicit construction uses zero transport while mass first vanishes and later regrows to connect the endpoint measures.
  • The transport interval has zero cost because the total mass is zero, yielding an optimal UDOT cost of 0.
  • The resulting dynamics decays the initial mass to zero at an infinite rate and regrows the final mass at an infinite rate.

A.2 Proof of Theorem 4.1

Theorem 4.1 reduces the travelling Dirac optimization to one-dimensional temporal functions for position and mass. The resulting spatial trajectory is straight, with the functions determined by endpoint distance and mass ratio.

  • The functions k(t) and l(t) depend only on the endpoint distance d and mass ratio r = m1/m0.
  • The travelling Dirac has a straight spatial trajectory parameterized by a one-dimensional function k(t).
  • The mass trajectory can be written as m(t) = m0l(t), where l(t) is one-dimensional and satisfies the endpoint mass conditions.
  • The travelling-Dirac cost is homogeneous of degree 1 in the endpoint masses and is symmetric for even Ψ.

A.4 Proof of Theorem 4.3

For uniformly convex growth penalties under the stated diameter condition, the conditional cost is jointly convex and sublinear, making the semi-coupling problem convex. The implementation then uses neural path, cost, and flow models with projected-gradient optimization and numerical simulation for inference.

  • For uniformly convex Ψ with Ψ′′ ≥κ, the conditional cost Cd is jointly convex and sublinear in endpoint masses under the stated distance condition.
  • Joint convexity follows by evaluating the cost at a convex combination of minimizing mass trajectories.
  • The semi-coupling optimization is convex because its constraint is linear and the conditional cost is jointly convex.
  • SUDO represents path and cost functions with neural networks and solves large-scale semi-couplings using mini-batch OT with projected gradient descent.
  • The learned velocity and growth fields are converted into continuous measure flows with an ODE simulator, or an SDE simulator for the stochastic version.

B.10 VarRUOT training time

The section evaluates path and cost-model behavior across parameter grids and reports robustness to penalty choices, while highlighting a scalability caveat for VarRUOT.

  • Scalability comparison: VarRUOT’s fixed-sampling strategy keeps training time from increasing with dataset size but progressively worsens performance on larger datasets.Training on all 49302 cells instead would require more than 11 hours.
  • Path-model evaluation: The path model is evaluated over d and r grids, including unseen boundary conditions outside the training range.The training range is d ∈[0.1, 12], r ∈[0.01, 10], while evaluation includes d values 0.01 and 15 and r values 0.001 and 20.
  • Path-model evaluation: For d = 15, particles do not travel because d exceeds 4π, so the dynamics become pure birth-death.This explains why the largest distance is excluded from the position-variation plot.
  • Path-model evaluation: In most cases, learned paths closely align with ground truth, while poorly learned cases occur at very small r and receive low downstream weights.The low weights limit their influence on subsequent cost fitting and unbalanced flow matching.
  • Cost and penalty evaluation: The study also evaluates transport-cost approximation error under different training grids and tests penalty and parameter sensitivity on simulation-gene data.Relative L2 error compares Eξ(d, r) with the true cost Cd(1, r), while separate ablations vary δ and Ψ.

C.5 Gaussian Data

On Gaussian and single-cell data, SUDO matches target measures and learns growth behavior that can differ from WFR-FM when the growth penalty encodes asymmetric biological structure. The section also contrasts dynamic growth modeling with formulations that represent imbalance mainly through endpoint couplings.

  • Gaussian data: SUDO matches the target Gaussian measure under both tested growth-penalty settings.The Gaussian experiment uses a 1000D mixture with a proliferating upper cluster and a transporting, non-growing lower cluster.
  • Gaussian data: SUDO learns a more plausible growth rate than WFR-FM, with zero growth for the lower cluster undergoing pure transport.WFR-FM’s quadratic-time mass variation can impose growth behavior inconsistent with this construction.
  • Single-cell evaluation: On the EMT dataset, SUDO is comparable to simulation-free WFR-FM and consistently better than simulation-based baselines.The dataset contains A549 lung-cancer-cell profiles across four time points in a 10-dimensional latent space.
  • General growth penalties: For only-growth and only-death penalties, the induced travelling-Dirac mass is respectively nondecreasing and nonincreasing.These penalties therefore produce mass behavior distinct from WFR travelling Diracs.
  • Dynamic formulation: SUDO learns both velocity and growth rate, coupling transport and mass variation along the continuous-time trajectory.This differs from formulations whose flow-matching step focuses on velocity while mass variation is represented primarily through a static transport plan.
  • Method: SUDO is a simulation-free generalization of WFR-FM that learns travelling Diracs and transport costs, solves semi-coupling by projected gradient descent, and then learns the UDOT flow.The approach avoids the closed-form travelling-Dirac requirement of WFR-FM for non-quadratic penalties.
Loading 2609.04710v1…