U-NO: U-shaped Neural Operators

TL;DR

U-NO introduces a U-shaped neural operator architecture that enhances memory efficiency and enables deeper models, leading to significant improvements in predicting solutions to complex PDEs like Darcy's flow and Navier-Stokes equations.

Contribution

The paper proposes U-NO, a novel U-shaped neural operator architecture that improves depth, training speed, and accuracy for PDE solution learning tasks.

Findings

26% prediction improvement on Darcy's flow

44% prediction improvement on Navier-Stokes equations

37% improvement on 3D spatiotemporal Navier-Stokes tasks

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.