Towards Multi-spatiotemporal-scale Generalized PDE Modeling
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. {In recent years}, various neural network architectures have been proposed to target such phenomena, most notably Fourier Neural Operators (FNOs), which give a natural handle over local \& global spatial information via parameterization of different Fourier modes, and U-Nets which treat local and global information via downsampling and upsampling paths. However, large-scale comparisons between these convolution-based approaches are notoriously sparse. In this work, we make such comprehensive comparisons regarding performance, runtime complexity, memory requirements, and generalization capabilities. Concretely, we stress-test various FNO, (Dilated) ResNet, and U-Net like approaches to 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. Next, we use our insights on design considerations, and introduce U-FNets, i.e., modern U-Nets that are augmented with FNO downsampling layers. Those architectures further improve performance without major degradation of computational cost. Finally, we ablate and discuss various choices for parameter conditioning}, and show promising results on generalization to different PDE parameters and time-scales with a single surrogate model. Source code for our PyTorch benchmark framework is available at https://anonymous.4open.science/r/tmlr-pdemulti-6677/.
Submission Length: Regular submission (no more than 12 pages of main content)
Changes Since Last Submission: Based on reviewer feedback we have reworked our abstract, experiments and conclusion sections to provide clarity w.r.t. our line of experimentation and our findings, currently highlighted in color "red".
Assigned Action Editor: ~Vikas_Sindhwani1
License: Creative Commons Attribution 4.0 International (CC BY 4.0)
Submission Number: 738
Loading