Neural scheme models complex flows mixing transport and growth in space
A Neural JKO Scheme for Hellinger-Kantorovich Gradient Flows via Monge-Growth Pairs
Machine Learning
Summary
This paper presents a new way to simulate systems where things move around and can grow or shrink locally over time, such as in chemical reactions or biological processes. The authors combine two ideas—a map showing how stuff moves and a factor capturing how mass changes—to better track these changes together. Their method is flexible and does not rely on grids, making it easier to use in complex spaces. They prove mathematical properties about their method and show it captures the expected energy behavior in these systems. Finally, they test it numerically to show it agrees well with the theoretical equations.
What this means in practice
- •For computational physicists: Simulate physical systems with transport, reaction, and diffusion effects jointly without requiring spatial meshes.
- •For computational biologists: Model biological processes involving spatial movement and local growth or decay more accurately using a flexible computational scheme.
Authors
Geuntaek Seo, Cheolhyeong Kim, Hwijae Son, Hyung Ju Hwang
Abstract
We develop a mesh-free neural JKO scheme for advection-reaction-diffusion equations with a gradient-flow structure in the Hellinger-Kantorovich (HK) geometry of unbalanced optimal transport. Each update is parametrized by a spatial map and a mass-changing factor, allowing spatial redistribution and local mass creation or loss to be treated jointly within a single variational step. Their cone action bounds the squared HK distance from above, yielding a sufficient condition for discrete energy dissipation through comparison with the identity pair. Minimizing the pair objective over all admissible pairs recovers the exact JKO minimum when the source and a minimizer have positive densities. We establish existence and mass bounds for JKO minimizers and, under additional assumptions, obtain positivity and regularity together with a discrete Euler-Lagrange equation and a metric-dissipation identity. The self-consistent chemical potential is then nonincreasing along an optimal map. There exist parametric pairs whose endpoint densities and objective values converge to those of an exact JKO minimizer, provided a regular-pair approximation hypothesis holds. Finally, we show that a primal-dual gap controls objective suboptimality and, for Boltzmann entropy, the $L^1$ density error, assuming exact-step regularity, positive-semidefinite interactions, and global dual feasibility. Numerical experiments examine pointwise agreement with the PDE, energy dissipation, and the roles of transport, reaction, and fully implicit interactions.