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

arXiv:2610.07602 · math.NA, cs.LG, cs.NA, math.AP, math.OC · Submitted 2026-10-06 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "A Neural JKO Scheme for Hellinger-Kantorovich Gradient Flows via Monge-Growth Pairs".

Jane: Detailed Research Summary: Mesh-Free Neural JKO Scheme for Advection–Reaction–Diffusion Equations in Hellinger–Kantorovich Geometry This research presents a novel,

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So, to summarize what we've heard about "A Neural JKO Scheme for Hellinger-Kantorovich Gradient Flows via Monge-Growth Pairs," they’re using neural networks to parameterize a spatial map and a mass-changing factor simultaneously to solve advection reaction diffusion equations.

Jane: Right. And the paper says that this framework is built on the HK geometry of unbalanced optimal transport, which is used because it couples transport with reaction in a way that standard metrics might not handle well.

Lu: They formalize the flow of a free energy F as d t rho = a div(rho grad g rho) - b rho g rho, where a and b weight spatial transport and local mass change, respectively, corresponding to the normalization HK = HK1,4 <ref:2610.07602#pg1>.

Meng: So they are essentially trying to discretize this continuous flow using a neural network approach within this specific geometric setting. What is the core problem they are solving in simpler terms?

Tom: The core problem is how to take that continuous gradient flow structure and turn it into a discrete update step, and they solve it by parameterizing the pair (T, q) using neural networks instead of some fixed, pre-defined map.

Jane: It seems like they are trying to make the discretization process itself learn the best way to redistribute mass spatially while also handling local changes in density at each step.

Lu: They show that this is achieved by defining an action term A mu(T, q) and then optimizing a combined objective function I tau(T, q mu) over all admissible Monge-growth pairs.

Meng: And they prove that this unrestricted optimization actually recovers the exact JKO problem when you have positive densities for both the source and the minimizer density. That’s a big deal for accuracy, right?

Tom: It is, because it means their neural scheme isn't just an approximation; it can recover the true JKO minimum under those conditions.

Jane: So, if we put that together, it’s a neural way to solve these complex PDEs by learning the transport and mass modification components together within one step.

Lu: And they establish existence and mass bounds for the JKO minimizer mu one under specific conditions on the potential function U and regularization parameters tau, <ref:2610.07602#pg1,establish existence and mass bounds for>.

Meng: Existence is good, but what about stability? If we are using this in production, we need to know if these solutions stay bounded.

The paper's summary: Tom: Moving on to the specific improvements they suggest for this Neural JKO Scheme for Hellinger-Kantorovich Gradient Flows via Monge-Growth Pairs, they focus on making the scheme more robust and accurate.

Jane: They introduce a few things. One is that they use cone action bounds on the squared HK distance from above as a sufficient condition for discrete energy dissipation when compared against an identity pair.

Lu: That’s important because it provides a way to get those dissipation guarantees, which is how you ensure the scheme actually works physically rather than just getting stuck in some numerical artifacts.

Meng: So that means we have a mathematical guarantee about energy loss, but what about the approximation itself? How good is this neural representation compared to something simpler?

Tom: They also show that they can achieve exact HK-JKO step recovery for any finite horizon by employing the one-step HK-JKO reparametrization, which means their cone action equals the squared HK distance at each step.

Jane: That’s a strong claim because it suggests that for any finite time horizon, you can generate a sequence of neural updates whose cone action equals the squared HK distance at each step.

Lu: And they also have this result about approximation guarantees: 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.

Meng: That convergence sounds like a good thing, but it’s conditional on some "regular-pair approximation hypothesis," so how much do we actually need to assume for that to hold?

Tom: They specify the condition as restricting the admissible pair class to Z adm, which requires the spatial map T to be a C one diffeomorphism with a specific derivative condition <ref:2610.07602#pg1>.

Jane: So, they’re telling us that if you can find a map that satisfies those smoothness conditions, you get convergence of the densities and objective values.

The paper's improvements: Tom: We’ve covered the main points of "A Neural JKO Scheme for Hellinger-Kantorovich Gradient Flows via Monge-Growth Pairs," which involves using neural networks to learn both transport and mass modification at once in this geometry.

Jane: So, the paper gives us rigorous guarantees on existence, density bounds, and error control through primal-dual gaps. The key is linking the variational structure to physical dissipation identities.

Lu: They establish metric-dissipation identities, like Corollary C.nine which states that Z grad phi squared + four phi squared d mu = HK two(mu, nu), and Corollary C <ref:2610.07602#pg2>.eleven relating the HK distance to the density rho via HK two(rho dx, nu) = tau squared I g(rho) <ref:2610.07602#pg2>.

Meng: From a practical standpoint, these identities help quantify the primal-dual gap, which controls the relative entropy integrand as R met at least zero separating a discrete Euler residual term from the primal–dual metric gap <ref:2610.07602#pg3>.

Tom: So while they don't prove convergence for stochastic optimizers or tau zero limits, they do provide strong *a posteriori* error estimates, specifically an L1 density error bound for Boltzmann entropy by a term proportional to two(nu + mu one)R <ref:2610.07602#pg3>.

Jane: That means we get a concrete way to measure how far our computed state is from the exact one based on these bounds.

Lu: And they also show that with the regular-pair approximation hypothesis, endpoint densities and objective values converge to those of an exact JKO minimizer.

Meng: It’s a solid framework for setting up a neural scheme, even if we still have some open problems regarding the required C squared control for the Euler–Lagrange equation regularity <ref:2610.07602#pg1>.

Tom: So, in summary, "A Neural JKO Scheme for Hellinger-Kantorovich Gradient Flows via Monge-Growth Pairs" gives us a powerful neural tool that links variational structure to physical dissipation identities. We’ll be looking at how this impacts the next steps in solving these types of problems.

Conclusion: Tom: So we’ve been looking at "A Neural JKO Scheme for Hellinger–Reaction–Diffusion Equations in Hellinger–Kantorovich Geometry." Basically, they used neural networks to learn a spatial map and a mass-changing factor together to solve these complex flow equations.

Jane: Right, so the core idea is that instead of using a fixed formula for how things move and change density, the AI learns both parts simultaneously in each step. It’s like giving the computer an adaptive way to redraw its own map while it’s moving mass around.

Lu: They build this on the Monge-growth formulation, which is basically a pair of a spatial transport map T and a mass modification factor q. They optimize both of those using neural networks within the Hellinger–Kantorovich geometry.

Meng: What’s the payoff here? How does this translate from a neat math trick to something that actually runs on hardware without crashing at scale? Does it stay stable?

Tom: Well, they proved existence and mass bounds for the resulting JKO minimizer mu one under certain conditions on the potential function U. They also showed that if you have positive densities for both the source and the minimizer density, you recover the exact JKO minimum.

Jane: That’s a big deal because it means their neural scheme isn't just guessing; it can find the true optimal solution under those specific conditions. They even get an L1 density error bound for Boltzmann entropy by controlling that primal-dual gap.

Lu: The identities they derive, like Corollary C.nine and C.eleven which link the HK distance to the Monge-growth pair, show a strong connection between the geometry of transport and how much energy is dissipated in each step. That’s solid theoretical grounding for what they are doing here.

Meng: But I gotta ask, if we're deploying this on a real traffic flow problem, does parameterizing T and q with two full-input dense residual networks actually keep the inference time reasonable? How much overhead is there?

Tom: The experiments showed that even with those complex network structures, they get pointwise agreement with the underlying PDE and demonstrate energy dissipation. Table eleven shows that diffusion and drift terms contribute significantly to the residual norm at specific time steps, but their sum stays relatively small.

Jane: And they did show some nice convergence guarantees for approximation under a regular-pair approximation hypothesis, meaning if you find a map that meets certain smoothness requirements, the endpoint densities converge to the exact minimizer. That’s reassuring for practical use.

Lu: The limitation they admit is that proving uniform convergence for stochastic optimizers or a tau going to zero limit across all steps is still an open problem, and they haven't established the necessary C squared control for the Euler–Lagrange equation yet. That’s where the theory still has some gaps.

Meng: So, in short, it’s a very strong mathematical foundation for this neural method, but we still need more proof that it behaves reliably when you push it to the limits of time or complexity. It’s a powerful tool for research right now.

Tom: So the final word on "A Neural JKO Scheme for Hellinger–Kantorovich Gradient Flows via Monge-Growth Pairs" is that they’ve built a rigorous variational framework using neural networks to learn both spatial map and mass modification in this geometry, giving us strong guarantees on existence and error control.

Jane: It really shows how deep the connection is between optimal transport geometry, gradient flows, and neural network representation learning.

Lu: We’re looking at how these kinds of variational methods can be parameterized by deep architectures to tackle non-standard geometries like HK space. That opens up new ways to model systems where simple Euclidean assumptions just don't cut it.

Meng: For me, the practical implication is that if we can get those convergence proofs ironed out, this could be a way to build more flexible and adaptive simulators for complex physical processes without having to manually design every single transport step.

Tom: Exactly. The next step for this research will be bridging that gap between their theoretical guarantees and the practical deployment concerns we talked about earlier. That’s where the real engineering challenge lies next, folks.

Geuntaek Seo, Cheolhyeong Kim, Hwijae Son, Hyung Ju Hwang

Pohang University of Science and Technology · Samsung Electronics · Konkuk University

math.NA, cs.LG, cs.NA, math.AP, math.OC

Submitted: 2026-10-06

Updated: 2026-10-06

Comments: 55 pages, 10 figures

Code: https://github.com/tyty4/hkflow

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 85/100

The gist: Detailed Research Summary: Mesh-Free Neural JKO Scheme for Advection–Reaction–Diffusion Equations in Hellinger–Kantorovich Geometry This research presents a novel, mesh-free neural scheme

Key concepts

Hellinger-Kantorovich (HK) Geometry: This is a specialized mathematical space used for optimal transport problems involving unbalanced masses. It provides the geometric framework necessary to analyze how mass moves and changes in a way that respects the physical constraints of the advection-reaction-diffusion system, allowing for robust energy dissipation estimates.
Monge-Growth Pair: This is the core variational structure where an update involves two parts: a spatial map (T) that dictates how positions are transported across space, and a mass-changing factor (q) that determines whether mass is created or lost locally before transport occurs. These two components are optimized together.
Primal-Dual Gap: This concept measures the difference between an upper bound derived from the optimization problem (the primal bound) and a lower bound derived from a dual formulation. Controlling this gap allows researchers to quantify how far a computed solution is from the true optimal solution, specifically controlling density errors for Boltzmann entropy.
Neural JKO Scheme: This is the proposed computational method that uses neural networks to parameterize the spatial map and mass-changing factor in the JKO (Jordan-Kinderlehrer-Otto) scheme. This allows for a flexible, mesh-free approach to solving complex partial differential equations by learning the necessary update rules directly from data.

Terminology

Gradient Flow — 1× in this paper · explained in 2 episodes

Summary

Detailed Research Summary: Mesh-Free Neural JKO Scheme for Advection–Reaction–Diffusion Equations in Hellinger–Kantorovich Geometry

This research presents a novel, mesh-free neural scheme designed to solve advection–reaction–diffusion equations characterized by a gradient-flow structure within the specialized Hellinger–Kantorovich (HK) geometry of unbalanced optimal transport. The core innovation lies in developing a variational framework that integrates spatial map evolution and mass modification into a single, unified step, all parameterized by neural networks.

Core Methodology and Framework

The scheme is built upon the Monge-growth formulation, where an update involves two components: a spatial map T that transports positions, and a mass-changing factor q that dictates local mass creation or loss before transport. This pair (T, q) defines the Monge-growth pair.

The optimization is performed within the HK geometry, leveraging its properties to provide robust estimates for energy dissipation. A key contribution is the introduction of cone action bounds on the squared HK distance from above, which serves as a sufficient condition for discrete energy dissipation when compared against an identity pair.

The scheme employs neural networks to parameterize both T and q. Specifically, the spatial map T is realized using two full-input dense residual Feedforward Neural Network (FCNN) blocks with certified contraction bounds. Scalar fields utilize trainable mixed features eta, and the density itself is represented via an exponential function of these features, incorporating a fitted mass shift. The growth factor q is similarly parameterized using a separate growth network structure involving.

Key Theoretical Results and Guarantees

The paper establishes several rigorous theoretical guarantees regarding the existence, properties, and approximation capabilities of the JKO minimizer:

Existence and Mass Bounds (Theorem 2.3 & Theorem B): The framework guarantees the existence of a JKO minimizer mu 1 = rho 1 dx under specific conditions on the potential function U and the regularization parameters tau,. Furthermore, mass bounds are established: every minimizer satisfies mu 1 at most M tau,, where M tau, is finite. If the growth potential U'(0+) = -infinity, the minimizer density rho 1 is strictly positive almost everywhere, and its support covers the entire domain.
Positivity and Regularity (Theorem 2.6): Under additional assumptions (Assumption 2.2), a discrete Euler–Lagrange equation is obtained, alongside a metric–dissipation identity. This identity implies that the self-consistent chemical potential is nonincreasing along an optimal map. Crucially, the density rho 1 of the minimizer is bounded away from zero and satisfies rho 1 at most rho, where rho depends on bounds derived from U' and other system parameters.
Exact Minimization Recovery: A significant result is that minimizing the pair objective over all admissible pairs recovers the exact JKO minimum when both the source density (mu) and the minimizer density (mu 1) possess positive densities.
Approximation Guarantees (Theorem C & Theorem D):

Parametric Approximation: 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. This is achieved by restricting the admissible pair class to Z adm, which requires the spatial map T to be a C 1 diffeomorphism with a specific derivative condition (d (x, T z(x)) = grad z(x) squared).
A Posteriori Error Bounds: Theorem D provides an a posteriori error estimate for an admissible computed measure. Under convexity and dual feasibility assumptions, the gap between the primal upper bound (Monge-growth) and the dual lower bound controls both the JKO objective suboptimality and, specifically for Boltzmann entropy, controls the L1 density error by a term proportional to 2(nu + mu 1)R.

Metric-Dissipation Identities (Corollary C.9 & C.11): The paper derives discrete identities relating the HK distance and the Monge-growth pair, such as:

Corollary C.9: Z grad phi squared + 4 phi squared d mu = HK 2(mu, nu), where phi is related to the source component.
Corollary C.11: HK 2(rho dx, nu) = tau squared I g(rho).

Analysis of the Primal-Dual Gap and Error Control

The analysis of the primal-dual gap is central to quantifying suboptimality. For Boltzmann entropy, this gap controls the relative entropy integrand as R met at least 0, separating a discrete Euler residual term from the primal–dual metric gap.

Limitations and Future Directions

Despite these strong theoretical results, the paper explicitly notes limitations regarding convergence theory for stochastic optimizers:

No Convergence Proof: The paper does not prove convergence of the stochastic optimizer or a tau 0 limit for the trained scheme.
Uniformity Concerns: While approximation by a fixed two-block architecture is shown in experiments, it is not proven to be uniform across all necessary conditions. A general tau 0 theory would require stepwise estimates uniform in the step index, which are not established here.
Regularity Gaps: Establishing the required C squared control for the Euler–Lagrange equation remains an open problem, as standard regularity theory does not yield the necessary gain for unbalanced optimal transport.

Numerical Validation

The numerical experiments confirm pointwise agreement with the underlying PDE, demonstrate energy dissipation, and show that transport, reaction, and fully implicit interactions play distinct roles in the scheme's performance. Specific validation metrics include:

Table 10: Reports held-out fit errors relative to unprojected endpoints and optimizer work in closure evaluations.
Table 11: Shows strong residuals at specific time steps (T=0.3, tau=0.0025), indicating that diffusion and drift terms contribute significantly to the residual norm, though their sum remains small (about 0.14% of either individual term).

In summary, this paper provides a mathematically rigorous foundation for a mesh-free neural JKO scheme in HK geometry, offering strong guarantees on existence, density bounds, and error control via primal-dual gaps. It successfully links the variational structure to physical dissipation identities while utilizing deep learning architectures for parameterization.

Improvements for AI systems

The AI system can perform high-fidelity, mesh-free numerical solutions for advection–reaction–diffusion equations by utilizing a neural JKO scheme that jointly learns a spatial map and a mass-changing factor in each variational step. This allows the system to treat spatial redistribution and local mass creation or loss to be treated jointly within a single variational step.
The system can achieve exact HK-JKO step recovery for any finite horizon by employing the one-step HK-JKO reparametrization (Proposition 3.1), meaning it can generate a sequence of neural updates whose cone action equals the squared HK distance at each step.
The AI system can provide rigorous, a posteriori error estimates for computed states by combining primal upper bounds and globally feasible dual lower bounds, as described in Theorem D, which controls the JKO objective gap and the convexity remainder of U and provides an L1 density error bound for Boltzmann entropy.
The system can recover exact state information from approximate neural iterates by leveraging a regular-pair approximation hypothesis (Theorem 5.2), ensuring that endpoint densities and objective values converge to those of an exact JKO minimizer under specified conditions.
The AI system can implement robust, long-horizon simulations by utilizing the finite-horizon composition of the one-step recovery (Corollary D.2), which ensures that over a horizon T, both the pairs and the represented states converge and that variational objectives converge uniformly in k to their exact counterparts.