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

TL;DR

This work proposes Jacobian-Adaptive Weighting for Stability (JAWS), a probabilistic regularization strategy designed to mitigate limitations in data-driven surrogate models and improves long-term stability, shock fidelity, and out-of-distribution generalization while reducing training computational costs.

Abstract

Data-driven surrogate models improve the efficiency of simulating continuous dynamical systems, yet their autoregressive rollouts are often limited by instability and spectral blow-up. While global regularization techniques can enforce contractive dynamics, they uniformly damp high-frequency features, introducing a contraction-dissipation dilemma. Furthermore, long-horizon trajectory optimization methods that explicitly correct drift are bottlenecked by memory constraints. In this work, we propose Jacobian-Adaptive Weighting for Stability (JAWS), a probabilistic regularization strategy designed to mitigate these limitations. By framing operator learning as Maximum A Posteriori (MAP) estimation with spatially heteroscedastic uncertainty, JAWS dynamically modulates the regularization strength based on local physical complexity. This allows the model to enforce contraction in smooth regions to suppress noise, while relaxing constraints near singular features to preserve gradients, effectively realizing a behavior similar to numerical shock-capturing schemes. Experiments demonstrate that this spatially-adaptive prior serves as an effective spectral pre-conditioner, which reduces the base operator's burden of handling high-frequency instabilities. This reduction enables memory-efficient, short-horizon trajectory optimization to match or exceed the long-term accuracy of long-horizon baselines. Evaluated on the 1D viscous Burgers' equation, our hybrid approach improves long-term stability, shock fidelity, and out-of-distribution generalization while reducing training computational costs.

Figures (16)

• Figure 1: Evolution of relative $L^2$ error over 200 autoregressive steps (log scale). PINN exhibits severe numerical divergence beyond step 110, whereas all other models maintain bounded growth. Among them, JAWS-S achieves the lowest accumulated error.
• Figure 2: Evolution of kinetic energy $\|u\|^2$ over 200 autoregressive steps (log scale). The ground truth (dashed) decays monotonically due to viscous dissipation. JAWS-G over-dissipates slightly, while JAWS-S exhibits residual energy preservation due to localized relaxation.
• Figure 3: Jacobian spectral radii distribution. Note the "Stability Gap": JAWS-S creates a safe margin ($\rho \approx 0.35 \ll 1$) that buffers against nonlinear instabilities, whereas PINN ($\rho \approx 0.91$) and Baseline ($\rho \approx 0.93$) operate on the precarious "Edge of Stability."
• Figure 4: Wavenumber energy spectrum. The Baseline suffers from "Spectral Blocking" (energy pile-up at high $k$), a precursor to instability. JAWS-G correctly mimics viscous dissipation, while JAWS-S preserves a controlled high-frequency plateau essential for shock capturing.
• Figure 5: Relative $L^2$ error vs. Input Noise $\sigma$. JAWS variants exhibit a flat error response compared to standard baselines. The aleatoric uncertainty term $s_1$ automatically "down-weights" the noisy data, preventing overfitting to Gaussian perturbations.