U-NO: U-shaped Neural Operators

TL;DR

U-NO introduces a memory-efficient, U-shaped neural operator architecture that contracts and then expands function-domain representations to enable deeper operator models. By leveraging Fourier-based integral operators and skip connections, U-NO achieves notable accuracy gains over the Fourier Neural Operator on Darcy flow and Navier–Stokes benchmarks, while using substantially less training memory for deeper networks. The approach extends to 3D spatio-temporal operator learning with strong performance gains and enables zero-shot super-resolution. These results suggest that problem-specific domain-geometry exploitation can dramatically improve the scalability and effectiveness of neural operators in solving PDEs.

Abstract

Neural operators generalize classical neural networks to maps between infinite-dimensional spaces, e.g., function spaces. Prior works on neural operators proposed a series of novel methods to learn such maps and demonstrated unprecedented success in learning solution operators of partial differential equations. Due to their close proximity to fully connected architectures, these models mainly suffer from high memory usage and are generally limited to shallow deep learning models. In this paper, we propose U-shaped Neural Operator (U-NO), a U-shaped memory enhanced architecture that allows for deeper neural operators. U-NOs exploit the problem structures in function predictions and demonstrate fast training, data efficiency, and robustness with respect to hyperparameters choices. We study the performance of U-NO on PDE benchmarks, namely, Darcy's flow law and the Navier-Stokes equations. We show that U-NO results in an average of 26% and 44% prediction improvement on Darcy's flow and turbulent Navier-Stokes equations, respectively, over the state of the art. On Navier-Stokes 3D spatiotemporal operator learning task, we show U-NO provides 37% improvement over the state of art methods.

Figures (6)

• Figure 1: U-NO architecture. $a$ is an input function, $u$ is the output. Orange circles are point-wise operators, rectangles denote general operators, and smaller blue circles denote concatenations in function spaces.
• Figure 2: (A) An illustration of aggressive contraction and re-expansion of the domain (and vice versa for the co-domains) by the variant of $\textsc{U-NO}\xspace$ with a factor of $\frac{1}{2}$ on an instance of Darcy flow equation. (B) Vorticity field generated by the $\textsc{U-NO}\xspace$ variant as a solution to an instance of the two-dimensional Navier-Stokes equation with viscosity $10^{-6}.$
• Figure 3: (A) Training memory requirements (in MB) for the 3D spatio-temporal problem of Navier-Stokes equation ($\nu=1\mathrm{e}{-3}$). For different depth only the number of stacked non-linear operators is varied. For deeper models, the additive memory requirement of U-NO is negligible compared to FNO model. For addition of $7$ more integral layer memory requirement only increased by only 80MB (vs $\sim$ 400MB for FNO). (B) Relative error in percentage on 3D spatio-temporal Navier-Stokes equation ($\nu=1\mathrm{e}{-3}$) for different depth (average over three repeated experiment is reported). We can observe a gradual decrease in the relative error with the increase of depth.
• Figure 4: Result of sensitivity of $\textsc{U-NO}\xspace^\dagger$ to learning rate and number of stacked non-linear operator (depth) used. All models are trained on the dataset of the Darcy Flow equation with resolution $211 \times 211$ following the training protocol described in \ref{['sec:darcy_flow_result']}. Models for each of the configuration is trained three times and the average error rate (in %) is reported. We can notice that except for a very high ($\ge 0.01$) or very low ($\le 0.0001$) learning rate, $\textsc{U-NO}\xspace^\dagger$ achieves low error rate at every other configurations (error rate achieved by FNO is 0.85 ). We also note that at all the high-performing hyper-parameter configuration setting $\textsc{U-NO}\xspace$ has low generalization gap.
• Figure 5: The training and test set error rate (in % on log scale) of U-NO and FNO for Navier-Stokes equation with viscosity $10e^{-3}$, (A) models performing $2D$ spatio-temporal convolution (B) models performing $3D$ spatial covolution. We can notice what U-NO converges much faster than FNO. The final test set error rate for FNO after 500 epochs is achieved only after around 200 epochs by U-NO and continues to improve on the error rate.