Towards Stability of Autoregressive Neural Operators

TL;DR

This paper investigates the instability issues in autoregressive neural operators used for modeling physical systems, proposing architectural improvements that enhance long-term prediction accuracy and stability.

Contribution

It introduces novel architectural and application-specific strategies to control instability in neural operators without increasing computational costs.

Findings

Significantly lower long-term forecast errors.

Extended stable prediction horizons.

Validated on fluid flow, shallow water, and weather systems.

Abstract

Neural operators have proven to be a promising approach for modeling spatiotemporal systems in the physical sciences. However, training these models for large systems can be quite challenging as they incur significant computational and memory expense -- these systems are often forced to rely on autoregressive time-stepping of the neural network to predict future temporal states. While this is effective in managing costs, it can lead to uncontrolled error growth over time and eventual instability. We analyze the sources of this autoregressive error growth using prototypical neural operator models for physical systems and explore ways to mitigate it. We introduce architectural and application-specific improvements that allow for careful control of instability-inducing operations within these models without inflating the compute/memory expense. We present results on several scientific systems that include Navier-Stokes fluid flow, rotating shallow water, and a high-resolution global weather forecasting system. We demonstrate that applying our design principles to neural operators leads to significantly lower errors for long-term forecasts as well as longer time horizons without qualitative signs of divergence compared to the original models for these systems. We open-source our \href{https://github.com/mikemccabe210/stabilizing_neural_operators}{code} for reproducibility.