Towards Stability of Autoregressive Neural Operators

Published: 26 Oct 2023, Last Modified: 26 Oct 2023Accepted by TMLREveryoneRevisionsBibTeX
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 code for reproducibility.
Submission Length: Regular submission (no more than 12 pages of main content)
Changes Since Last Submission: We made the requested changes to the submission, namely: - Updated the figure reference in section 4.2 - Added labels to figure 3 and updated caption to describe each section. - Updated notation in appendix to match updates during discussion phase - Clarified extent of claims for stability improvement.
Assigned Action Editor: ~Andriy_Mnih1
License: Creative Commons Attribution 4.0 International (CC BY 4.0)
Submission Number: 1241