JAWS: Enhancing Long-term Rollout of Neural PDE Solvers via Spatially-Adaptive Jacobian Regularization

TL;DR

JAWS introduces a spatially adaptive Jacobian regularization technique for neural PDE solvers, improving long-term stability and efficiency by locally adjusting regularization strength based on physical complexity.

Contribution

It proposes a novel Jacobian regularization method that adapts spatially, enabling stable long-term neural PDE simulations with reduced memory requirements.

Findings

Enhances long-term stability of neural PDE solvers.

Preserves physical properties and accuracy in complex flow simulations.

Reduces memory usage, enabling large-scale, long-horizon simulations.

Abstract

Data-driven surrogate models can significantly accelerate the simulation of continuous dynamical systems, yet the step-wise accumulation of errors during autoregressive time-stepping often leads to spectral blow-up and unphysical divergence. Existing global regularization techniques can enforce contractive dynamics but uniformly damp high-frequency features, causing over-smoothing; meanwhile, long-horizon trajectory optimization methods are severely constrained by memory bottlenecks. This paper proposes Jacobian-Adaptive Weighting for Stability (JAWS), which reformulates operator learning as a Maximum A Posteriori (MAP) estimation problem with spatially heteroscedastic uncertainty, enabling the regularization strength to adapt automatically based on local physical complexity: enforcing contraction in smooth regions to suppress noise while relaxing constraints near singular features such as shocks to preserve gradient information. Experiments demonstrate that JAWS serves as an effective spectral pre-conditioner for trajectory optimization, allowing short-horizon, memory-efficient training to match the accuracy of long-horizon baselines. Validations on the 1D viscous Burgers' equation and 2D flow past a cylinder (Re=400 out-of-distribution generalization) confirm the method's advantages in long-term stability, preservation of physical conservation properties, and computational efficiency. This significant reduction in memory usage makes the method particularly well-suited for stable and efficient long-term simulation of large-scale flow fields in practical engineering applications. Our source code and implementation are publicly available at https://github.com/jyohosyo-dot/JAWS_2D.