Multi-Grid Tensorized Fourier Neural Operator for High-Resolution PDEs

TL;DR

This work introduces MG-TFNO, a neural operator that scales to high-resolution PDEs by jointly compressing input domain information via multi-grid domain decomposition and operator parameters via a low-rank tensor factorization in the Fourier domain. Built on an improved Fourier Neural Operator backbone, MG-TFNO achieves substantial input and weight compression (over 150x in domain and 100x in parameters) while improving accuracy, especially in low-data regimes, demonstrated on turbulent Navier–Stokes and Burgers’ equations. The method combines separable spectral convolutions, architectural refinements (normalization, channel mixing, boundary handling), and multi-grid domain decomposition to enable parallelization and resolution-invariant generalization, with ablations confirming the contribution of each component. Overall, MG-TFNO offers a scalable, memory-efficient path toward learning high-resolution PDE solution operators with strong generalization and potential impact on large-scale scientific computing tasks such as weather forecasting.

Abstract

Memory complexity and data scarcity have so far prohibited learning solution operators of partial differential equations (PDEs) at high resolutions. We address these limitations by introducing a new data efficient and highly parallelizable operator learning approach with reduced memory requirement and better generalization, called multi-grid tensorized neural operator (MG-TFNO). MG-TFNO scales to large resolutions by leveraging local and global structures of full-scale, real-world phenomena, through a decomposition of both the input domain and the operator's parameter space. Our contributions are threefold: i) we enable parallelization over input samples with a novel multi-grid-based domain decomposition, ii) we represent the parameters of the model in a high-order latent subspace of the Fourier domain, through a global tensor factorization, resulting in an extreme reduction in the number of parameters and improved generalization, and iii) we propose architectural improvements to the backbone FNO. Our approach can be used in any operator learning setting. We demonstrate superior performance on the turbulent Navier-Stokes equations where we achieve less than half the error with over 150x compression. The tensorization combined with the domain decomposition, yields over 150x reduction in the number of parameters and 7x reduction in the domain size without losses in accuracy, while slightly enabling parallelism.

Figures (10)

• Figure 1: Comparison of the performance on the relative $L^2$ and $H^1$ test errors (lower is better) on a log-scale of our approach, compared with both our improved backbone (FNO) and the original FNO, on Navier-Stokes. Our approach enables large compression for both input and parameter, while outperforming regular FNO.
• Figure 2: Overview of our approach. First (left), a multi-grid approach is used to create coarse to fine inputs that capture high-resolution details in a local region while still encoding global context. The resulting regions are fed to a tensorized Fourier operator (middle), the parameters of which are jointly represented in a single latent space via a low-rank tensor factorization (here, a Tucker form). Here $\mathcal{F}$ denotes Fourier transform. Finally, the outputs (right) are stitched back together to form the full result. Smoothness in the output is ensured via the choice of the loss function.
• Figure 3: Original FNO and Improved Backbone Architecture. The original FNO architecture FNO is composed of simply a Spectral Convolution, with a (linear) skip connection to recover high-frequency information and handle non-periodic inputs (\ref{['fig:original_arch']}). We improve the architecture as detailed in section \ref{['sec:architecture']}. In particular, we have a version with a double (sequential) skip connection (\ref{['fig:arch_double_skip']}), while our best architecture uses nested skip connections, and can be made both with and without preactivation (subfigures \ref{['fig:improved_arch_preact']} and \ref{['fig:improved_arch']}, respectively). The latter, subfigure \ref{['fig:improved_arch']}, is our best architecture.
• Figure 4: Illustration of a Tucker decomposition. For clarity , we show ${\mathbf{W}}$ as a $3^{\text{rd}}$-order tensor weight.
• Figure 5: Domain decomposition in space (\ref{['fig:patching_only']}) and our Multi-Grid based approach. (\ref{['fig:multigrid_patching']}). White squares represent the region of interest while yellow squares the larger embeddings.