Learning Velocity Prior-Guided Hamiltonian-Jacobi Flows with Unbalanced Optimal Transport

ICLR 2026 Conference Submission13960 Authors

Keywords: optimal transport, action matching, hamiltonian-jacobi equation, cell trajectory inference

Abstract: The connection between optimal transport (OT) and control theory is well established, most prominently in the Benamou–Brenier dynamic formulation. With quadratic cost, the OT problem can be reframed as a stochastic control problem in which a density $\rho_t$ evolves under a controlled velocity field $v_t$ subject to the continuity equation $\partial_t\rho_t + \nabla\cdot(\rho_tv_t)=0$. In this work, we introduce a velocity prior into the continuity equation and derive a new Hamilton–Jacobi–Bellman (HJB) formulation to learn dynamical probability flows. We further extend the approach to the unbalanced setting by adding a growth term, capturing mass variation processes common in scientific domains such as cell proliferation and differentiation. Importantly, our method requires training only a single neural network to model $v_t$, without the need for a separate model for the growth term $g_t$. Finally, by decomposing the velocity field as $v_\mathrm{total} = v_\mathrm{prior} + v_\mathrm{corr}$, our approach is able to capture complex transport patterns that prior methods struggle to learn due to the curl-free limitation.

Primary Area: generative models

Submission Number: 13960