FluxNet: Learning Capacity-Constrained Local Transport Operators for Conservative and Bounded PDE Surrogates

TL;DR

FluxNet addresses the instability of autoregressive PDE surrogates by learning capacity-constrained local transport operators that guarantee discrete conservation. Rather than predicting the next state, it outputs fluxes $F_{i\to j}$ and updates states via $u_i^{t+1} = u_i^t - \sum_{j\in\mathcal{N}(i)} F_{i\to j} + \sum_{j\in\mathcal{N}(i)} F_{j\to i}$, ensuring exact mass conservation under symmetric stencils. Bound preservation is enforced structurally through specialized heads (L-head, U-head, D-head) and, when needed, a dual-consistency loss (DCL) to harmonize outflow and inflow constraints; FluxNet-LAP tailors this to shallow water, while FluxNet-D handles dual-bounded traffic flow. Across convection–diffusion, shallow water, traffic flow, and spinodal decomposition, FluxNet achieves machine-precision conservation, reduced bound violations, and substantial speedups via large timesteps, while preserving both pointwise accuracy and statistical microstructure measures such as two-point correlations. These properties suggest a scalable, physically grounded path for fast, reliable PDE surrogates in climate, fluid dynamics, and materials science.

Abstract

Autoregressive learning of time-stepping operators offers an effective approach to data-driven PDE simulation on grids. For conservation laws, however, long-horizon rollouts are often destabilized when learned updates violate global conservation and, in many applications, additional state bounds such as nonnegative mass and densities or concentrations constrained to [0,1]. Enforcing these coupled constraints via direct next-state regression remains difficult. We introduce a framework for learning conservative transport operators on regular grids, inspired by lattice Boltzmann-style discrete-velocity transport representations. Instead of predicting the next state, the model outputs local transport operators that update cells through neighborhood exchanges, guaranteeing discrete conservation by construction. For bounded quantities, we parameterize transport within a capacity-constrained feasible set, enforcing bounds structurally rather than by post-hoc clipping. We validate FluxNet on 1D convection-diffusion, 2D shallow water equations, 1D traffic flow, and 2D spinodal decomposition. Experiments on shallow-water equations and traffic flow show improved rollout stability and physical consistency over strong baselines. On phase-field spinodal decomposition, the method enables large time-steps with long-range transport, accelerating simulation while preserving microstructure evolution in both pointwise and statistical measures.

Figures (14)

• Figure 1: Rollout MAE over time for shallow water equations. Mean Absolute Error as a function of rollout time step for FluxNet-LAP, FNO, and CNN on the shallow water equation test set.
• Figure 2: Traffic flow density profiles at selected time points. Comparison of density $\rho(x,t)$ predictions between ground truth, FluxNet-D, FNO, and CNN at three time points ($t=0$, $t=25$, $t=50$).
• Figure 3: Two-point statistics error evolution for spinodal decomposition. MAE of the radial two-point correlation function as a function of simulation time for FluxNet-D models trained with different time step sizes. Shaded regions indicate $\pm 1$ standard deviation over 100 independent rollouts. The baseline error between two independent phase-field simulations is shown for reference.
• Figure 4: Effective receptive field analysis for FluxNet-D ($10\Delta t$). Visualization of the ERF for each output channel. The localized and anisotropic patterns confirm that the learned flux operators depend only on local neighborhood information.
• Figure 5: Concentration profiles at five time points comparing FluxNet-N (unconstrained), FluxNet-P (positive flux), and FluxNet-L (lower-bounded) variants on the 1D convection-diffusion equation. All three variants maintain good agreement with ground truth in smooth regions. The L-head produces the most accurate tracking near concentration minima approaching zero.