FluxNet: Learning Capacity-Constrained Local Transport Operators for Conservative and Bounded PDE Surrogates
Abstract
Autoregressive learning of time-stepping operators provides an effective approach to data-driven partial differential equation (PDE) simulation, yet for conservation laws, they face a fundamental challenge: learned updates may violate global conservation over long rollouts. For the important subclass of mass-conservation-type equations, the problem is compounded by inherent physical bounds (e.g., nonnegativity or concentrations in [0,1]) whose violation further destabilizes predictions. We introduce FluxNet, which learns cumulative transport amounts representing the total conserved quantity redistributed between each cell and a configurable neighborhood over the full surrogate interval. A conservative update guarantees exact discrete conservation by construction; modular capacity-constrained transport heads (L, U, and D) enforce lower bounds, upper bounds, or near-zero dual-bound violations through architectural design. Unlike flux-rate surrogates that require temporal integration and thus inherit CFL constraints, FluxNet involves no such integration; configurable transport neighborhoods enable large-timestep prediction at full spatial resolution. Ghost cells extend the framework to non-periodic boundaries. Experiments on four benchmarks (1D convection--diffusion, 2D shallow water, 1D traffic flow, 2D Cahn--Hilliard) demonstrate exact conservation, structural bound preservation, architecture modularity, and superior stability over flux-rate surrogates at large temporal strides. The code is publicly available at: https://github.com/Lan-zs/FluxNet.
Lay Summary
Computer simulations of physical systems, including fluid flow, traffic, and material microstructure evolution, must obey fundamental laws: mass cannot appear from nowhere, and quantities like water depth cannot become negative. Neural networks can learn to run these simulations much faster than traditional methods, but they often violate these physical laws when making long-term predictions, producing nonsensical results that quickly spiral out of control. We developed FluxNet, which learns how quantities are transported between neighboring locations rather than directly predicting future states. By design, whatever leaves one location must arrive at another, guaranteeing exact conservation. Meanwhile, the amount transported cannot exceed what is physically available, preventing violations of natural limits. This approach enables physically valid, fast simulations across diverse applications, from modeling flood propagation and predicting traffic jams to simulating material microstructure evolution. In one phase-field simulation benchmark, FluxNet achieves a 17-fold speedup over conventional solvers while preserving both conservation laws and statistical accuracy. By building physical guarantees directly into the network architecture, FluxNet makes neural-network-based scientific simulation more reliable and trustworthy.