A Neural JKO Scheme for Hellinger-Kantorovich Gradient Flows via Monge-Growth Pairs

2026-10-05

AI Summary

This study addresses the computation of non-equilibrium optimal transport gradient flows and advection-reaction-diffusion equations under the Hellinger–Kantorovich geometry. We propose a mesh-free neural JKO scheme that leverages Monge-Growth pairs to jointly handle spatial redistribution and mass variation, thereby enabling single-step variational optimization. Furthermore, by introducing a cone-action upper bound on the Hellinger–Kantorovich distance, we establish sufficient conditions for discrete energy dissipation. Our contributions include rigorous proofs of the existence, regularity, and convergence of the minimizers. Numerical experiments validate pointwise consistency with the underlying PDEs and confirm the energy dissipation properties, effectively decoupling the mechanisms of transport, reaction, and implicit interaction.

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.

Problem

• Hellinger-Kantorovich geometry
• unbalanced optimal transport
• gradient flows
• advection-reaction-diffusion equations
• JKO scheme

Innovation

• Neural JKO scheme
• Hellinger-Kantorovich geometry
• Unbalanced optimal transport
• Monge-growth pairs
• Advection-reaction-diffusion equations