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

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{ot}$, our approach is able to capture complex transport patterns—including looping and divergence —that prior methods struggle to learn due to the curl-free limitation.