Group Equivariant Fourier Neural Operators for Partial Differential Equations

About

We consider solving partial differential equations (PDEs) with Fourier neural operators (FNOs), which operate in the frequency domain. Since the laws of physics do not depend on the coordinate system used to describe them, it is desirable to encode such symmetries in the neural operator architecture for better performance and easier learning. While encoding symmetries in the physical domain using group theory has been studied extensively, how to capture symmetries in the frequency domain is under-explored. In this work, we extend group convolutions to the frequency domain and design Fourier layers that are equivariant to rotations, translations, and reflections by leveraging the equivariance property of the Fourier transform. The resulting $G$-FNO architecture generalizes well across input resolutions and performs well in settings with varying levels of symmetry. Our code is publicly available as part of the AIRS library (https://github.com/divelab/AIRS).

Jacob Helwig, Xuan Zhang, Cong Fu, Jerry Kurtin, Stephan Wojtowytsch, Shuiwang Ji · 2023

Related benchmarks

Task	Dataset	Result
2-D shallow-water equations simulation	2-D shallow-water equations (ID)	ID Error0.0307	5
Operator learning	1-D Burgers resolution 1024 to 2048 (OOD)	OOD Error1.0244	5
Spatiotemporal rollout prediction	(2+1)-D Spatiotemporal Navier-Stokes OOD, resolution 64^2 to 128^2 shift	Relative Error40.91	5
2-D shallow-water equations simulation	2-D shallow-water equations resolution 64^2 to 128^2 (OOD)	OOD Error0.7593	5
Operator learning	1-D Burgers resolution 1024 (ID)	ID Error14.06	5
Solving 2-D Burgers equation	2-D Burgers In-Distribution resolution 64^2 (test)	ID Error0.0967	5
Solving 2-D Burgers equation	2-D Burgers Out-of-Distribution resolution 64^2 to 128^2 shift (test)	OOD Error1.0131	5
Spatiotemporal rollout prediction	(2+1)-D Spatiotemporal Navier-Stokes ID	Relative Error4.59	5