StFT: Spatio-temporal Fourier Transformer for Long-term Dynamics Prediction

Da Long, Shandian Zhe, Samuel Williams, Leonid Oliker, Zhe Bai

Abstract:

Simulating the long-term dynamics of multi-scale and multi-physics systems poses a significant challenge in understanding complex phenomena across science and engineering. The complexity arises from the intricate interactions between scales and the interplay of diverse physical processes, which manifest in PDEs through coupled, nonlinear terms that govern the evolution of multiple physical fields across scales. Neural operators have shown potential in short-term prediction of such complex spatio-temporal dynamics; however,
achieving stable high-fidelity predictions and providing robust uncertainty quantification over extended time horizons remains an open and unsolved area of research. These limitations often lead to stability degradation with rapid error accumulation, particularly in long-term forecasting of systems characterized by multi-scale behaviors involving dynamics of different orders. To address these challenges, we propose an autoregressive Spatio-temporal Fourier Transformer (StFT), in which each transformer block is designed to learn the system dynamics at a distinct scale through a dual-path architecture that integrates frequency-domain and spatio-temporal representations. By leveraging a structured hierarchy of StFT blocks, the resulting model explicitly captures the underlying dynamics across both macro- and micro- spatial scales. Furthermore, a generative residual correction mechanism is introduced to learn a probabilistic refinement temporally while simultaneously quantifying prediction uncertainties, enhancing both the accuracy and reliability of long-term probabilistic forecasting. Evaluations conducted on three benchmark datasets (plasma, fluid, and atmospheric dynamics) demonstrate the advantages of our approach over state-of-the-art ML methods.

Submission Type: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission: We thank the reviewers for their thorough reading of our manuscript. In the revised manuscript, we have provided detailed analysis of the joint influence and trade-off between patch size and the number of retained Fourier modes on model performance in the Appendix.

Code: https://github.com/BerkeleyLab/StFT/

Assigned Action Editor: ~Jacek_Cyranka1

Submission Number: 5981