Lightweight Fourier Neural Operator for Time-Dependent Partial Differential Equations

Dawon Ahn, Satish Chandran, Daniel Leibovici, Nikola Borislavov Kovachki, Evangelos E. Papalexakis, Jean Kossaifi

Abstract:

Fourier Neural Operators (FNOs) have shown strong performance in solving timedependent partial differential equations (PDEs). However, accurately modeling
complex spatio-temporal dynamics remains challenging and is typically addressed
in one of two ways: (i) by applying spectral convolutions over the spatial domain
with temporal dynamics handled autoregressively, or (ii) by applying spectral convolutions over the entire spatio-temporal domain. While the former is more computationally efficient, it fails to capture true spatio-temporal interactions. The
latter, though more accurate, becomes computationally prohibitive when scaling
to larger datasets. We propose LITEFNO, a novel FNO framework that achieves
both numerical accuracy and computational efficiency for time-dependent PDEs.
Specifically, we first model spatial dynamics by learning a low-rank spatial basis
of spectral convolutional weights space. We then incorporate temporal dynamics
by learning a new temporal basis through transduction. This factorized formulation enables efficient learning of full spatio-temporal dynamics with significantlyfewer parameters (99.9% reduction) and superior performance (44% improvement
in VRMSE) compared to the variants of FNO models.