A Neural JKO Scheme for Hellinger-Kantorovich Gradient Flows via Monge-Growth Pairs
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 L1 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.
Related
- Optimal tradeoffs between variance and sensitivity, via Hellinger-Kantorovich geometry
- A Structure-Preserving Scheme for Computing the Gradient of the Eikonal Equation
- Neural Sampling with Reweighted Normalizing Flows via the Wasserstein--Fisher--Rao JKO Scheme
- A coarse-to-fine modulation classification scheme utilizing Transformers and Liquid Neural Networks for Smart Grids
- Gradient Growth for Vlasov-Poisson in background charge distributed by Desingularized Dirac Delta
BibTeX
@misc{seo2026neural,
title = {{A Neural JKO Scheme for Hellinger-Kantorovich Gradient Flows via Monge-Growth Pairs}},
author = {Geuntaek Seo and Cheolhyeong Kim and Hwijae Son and Hyung Ju Hwang},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.07602},
eprint = {2610.07602},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.07602},
}
RIS
TY - UNPB
TI - A Neural JKO Scheme for Hellinger-Kantorovich Gradient Flows via Monge-Growth Pairs
AU - Geuntaek Seo
AU - Cheolhyeong Kim
AU - Hwijae Son
AU - Hyung Ju Hwang
PY - 2026
DA - 2026/10/06
PB - arXiv
DO - 10.48550/arxiv.2610.07602
UR - https://arxiv.org/abs/2610.07602
AB - 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.
ER -