Towards Multi-spatiotemporal-scale Generalized PDE Modeling

Jayesh K Gupta, Johannes Brandstetter

Keywords: PDE modeling, multi-spatiotemporal-scale, PDE generalization, Fourier vs U-Net, fluid mechanics

Abstract: Partial differential equations (PDEs) are central to describing complex physical system simulations. Their expensive solution techniques have led to an increased interest in deep neural network based surrogates. However, the practical utility of training such surrogates is contingent on their ability to model complex multi-scale spatio-temporal phenomena. Various neural network architectures have been proposed to model complex spatial information in PDE solutions, most notably Fourier Neural Operators (FNOs) which give a natural handle over local & global information via parameterization of different Fourier modes, and U-Nets which treat local and global information via downsampling and upsampling paths. However, generalizing across different equation parameters or different time-scales still remains a challenge. In this work, we make a comprehensive comparison between various FNO and U-Net like approaches on fluid mechanics problems in both vorticity-stream and velocity function form. For U-Nets, we transfer recent architectural improvements from computer vision, most notably from object segmentation and generative modeling. We further analyze the design considerations for using FNO layers to improve performance of U-Net architectures without major degradation of computational performance. Finally, we show promising results on generalization to different PDE parameters and time-scales with a single surrogate model.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Submission Guidelines: Yes

Please Choose The Closest Area That Your Submission Falls Into: Machine Learning for Sciences (eg biology, physics, health sciences, social sciences, climate/sustainability )

TL;DR: We present a comprehensive study of Fourier and U-Net inspired architectural choices towards generalization of multi-spatiotemporal-scale problems.

Community Implementations: [ 4 code implementations](https://www.catalyzex.com/paper/arxiv:2209.15616/code)