Source-linked AI summary
U-NO: U-shaped Neural Operators
Md Ashiqur Rahman, Zachary E. Ross, Kamyar Azizzadenesheli
TL;DR
Neural operators face high memory demands that limit depth. U-NO adapts a U-shaped architecture to function spaces, contracting and expanding domains with skip connections; it improves performance across Darcy flow and Navier–Stokes benchmarks, including 3D spatiotemporal learning.
Problem
Prior neural operator architectures have high memory demands and generally prohibit very deep models.
Method
U-NO adapts the U-net architecture to neural operators by progressively contracting function domains during encoding and expanding them during decoding with skip connections.
Results
U-NO improves average performance by 26% on Darcy’s flow, 44% on Navier–Stokes, and 37% on 3D spatiotemporal Navier–Stokes operator learning over FNO or state-of-the-art methods.
Takeaways & Limitations
U-NO supports deeper, highly parameterized neural operators with lower memory usage while retaining data efficiency and robustness to hyperparameter choices.
Takeaways & Limitations
The decoder skip-connection formulation assumes matching encoder and decoder domains, specifically D_L1+i = D_L1−i−1.
Abstract
from arXiv · showhide
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.
1 Introduction
Neural operators face high memory demands that restrict depth, motivating U-NO’s U-shaped architecture for efficient, deeper function-space mappings. Across PDE benchmarks, U-NO improves performance while showing data efficiency and robustness to tuning.
- Motivation: High memory demand in prior neural operators limits the design of very deep models.These architectures are closely related to fully connected networks and often preserve function-space domains and co-domains throughout layers.
- U-NO approach: U-NO contracts function domains during encoding, expands them during decoding, and uses skip connections to construct outputs.The architecture maps between functions over different domains and uses Fourier-transform-based integration for its inner operators.
- Results: 26% average improvement is reported for high-resolution Darcy’s flow simulations, while Navier–Stokes shows 44% average improvement and a best improvement of 51%.These results are reported against the state-of-the-art FNO baseline.
- Results: 37% improvement is reported for Navier–Stokes 3D spatiotemporal operator learning over state-of-the-art methods.U-NO is described as the first neural operator trained for mappings from function spaces with 3D domains.
- Practical properties: U-NO permits models three times deeper and with 25 times more parameters while still providing performance improvement.The paper presents this comparison as evidence that the architecture supports deeper neural operators.
- Practical properties: U-NO demonstrates data efficiency, faster training, and robustness to hyperparameter choices.The paper reports training with only a few thousand data points and empirical robustness to domain/co-domain transformation hyperparameters.
2 Neural Operator Learning
Neural operator learning targets maps between function spaces using layered integral and point-wise operators. Earlier models generally preserve domains and co-domains across layers, contributing to high memory demands, whereas U-NO contracts domains to learn compact representations.
- Operator learning: Neural operators learn maps from an input function space A to an output function space U.The paper formulates this as training a parameterized operator Gθ: A → U from paired input and output functions.
- Operator learning: Each neural-operator layer combines a global linear integral operator with a point-wise matrix operator and nonlinearity.A kernel function and measure define the integral operation, while σ denotes the point-wise nonlinearity.
- Operator learning: Point-wise lifting and projection operators connect the input and output function spaces to the intermediate neural-operator layers.P primarily lifts the input function, while Q primarily projects to the output function space.
- Memory limitation: Earlier neural operators such as FNO preserve the same domain and co-domain across layers, creating high memory demands during global integration.U-NO instead contracts domains inward while increasing co-domain dimensions, producing compact representations and smaller integral-operation domains.
3 A U-shaped Neural Operator (U-NO)
U-NO adapts the U-net pattern to function spaces by encoding inputs on progressively smaller domains and decoding them back to the target domain. Skip connections link corresponding encoder and decoder representations, while lifting and projection operators handle input and output dimensions.
- Encoding: U-NO begins with a point-wise lifting operator P that maps the input function into a higher-dimensional vector-valued function.The paper chooses d0 ≫ dA so that P acts as a lifting operator.
- Encoding: The encoder applies nonlinear integral operators over layers with increasingly contracted domains and higher-dimensional co-domains.The construction satisfies µ(Di) ≥ µ(Di+1) and dvi+1 ≥ dvi across encoding layers.
- Decoding: The decoder expands domains and reduces co-domain dimensions until the target function domain is reached.Across decoder layers, µ(Di+1) ≥ µ(Di) and dvi ≥ dvi+1.
- Skip connections: Decoder operators concatenate corresponding decoder and encoder functions to form their inputs.These skip connections transfer encoder representations into the decoder, as illustrated in Figure 1.
- Skip connections: For simplicity, the construction assumes matching decoder and encoder domains when defining one concatenation.More generally, concatenation can use a map between the two domains.
- Output projection: A final point-wise projection operator Q maps the last intermediate function to the output space U.The paper chooses dL+1 ≫ dU so that Q functions as a projection operator.
4 Empirical Studies
The empirical studies evaluate U-NO and its aggressive variant on Darcy flow and Navier–Stokes operator-learning tasks. Across these experiments, U-NO combines lower memory usage with deeper architectures and improved performance, including in 3D spatio-temporal learning.
- Model Implementation: U-NO maps functions through progressively contracted domains and expanded co-domains, then reverses these transformations during decoding.Skip connections connect encoding and decoding stages, while Fourier transform-based integration is used for internal integral operators.
- Model Implementation: U-NO† applies a more aggressive contraction and expansion factor and is more memory efficient than the initial U-NO architecture.The aggressive variant uses scaling factors that reduce memory requirements during domain and co-domain transformations.
- Results on Navier-Stokes Equation: U-NO achieves lower error than FNO on Navier–Stokes benchmarks, including 37% lower relative error with 50% less memory in the 3D integration setting.The 3D comparison also uses 19 times more parameters for U-NO; some settings show nearly a 50% relative-error reduction over FNO.
- Results on Navier-Stokes Equation: U-NO supports zero-shot super-resolution in both time and space, unlike FNO, while matching or improving performance under constrained architectures and smaller training data.The high-resolution Navier–Stokes experiments report comparable performance despite a highly constrained architecture and smaller training data.
5 Conclusion
U-NO is a U-Net-inspired neural operator designed to improve data efficiency, tuning robustness, convergence, and memory efficiency while enabling deeper, highly parameterized models. It improves performance but retains the fundamental memory limitation of integral operators.
- U-NO provides data efficiency, robust hyperparameter tuning, flexible training, memory efficiency, convergence, and non-vanishing gradients.
- U-NO allows deeper models that incorporate problem structures and support highly parameterized operators for maps between function spaces.
- U-NOs are easy to tune, converge faster to desired accuracies, and achieve superior performance with minimal tuning.
- U-NO still carries the fundamental and inherent memory usage limitation of integral operators.
A.1 Positional Embedding For Unit Torus
The unit torus T2 is represented using products of unit circles, with the Clifford torus providing an embedding into Euclidean space.
- The unit torus T2 is homeomorphic to the Cartesian product S1 × S1 of two unit circles.
- The Cartesian product of embedded circles produces the Clifford torus.
A.2 Sensitivity to Hyper-parameter Selection
U-NO† maintains low error across most tested learning-rate and depth configurations, while learning large models from limited data.
- U-NO† achieves low error rates across configurations except at learning rates ≥0.01 or ≤0.0001.The reported comparison error rate for FNO is 0.85.
- At high-performing hyperparameter settings, U-NO has a low generalization gap.
- U-NO† models can learn 25 times as many parameters as FNO using only 1500 simulations.
A.3 FNO with Skip Connection
Adding skip connections to vanilla 2D FNO does not improve its performance on the reported Navier–Stokes benchmark, despite the comparison also tracking training memory requirements.
- Skip connections alone do not improve vanilla 2D FNO performance on the reported Navier–Stokes benchmark.
- The comparison reports both FNO performance and memory requirements during training.
A.4 Spatial Memory
U-NO reduces training memory relative to FNO by contracting function domains during the operator layers. With a contraction factor of 1/2, it requires 40% less memory on average.
- Spatial Memory: U-NO’s memory requirements are reported for Navier–Stokes training instances across different grid resolutions, with seven non-linear integral operator layers.
A.5 Zero-Shot super resolution on 3D Spatio-Temporal Data
U-NO supports zero-shot super resolution for 3D spatio-temporal data by contracting and expanding function domains, while reducing memory requirements and improving prediction error over FNO. Its domain construction is resolution invariant in both time and space, unlike 3D FNO in the time domain.
- A.5 Zero-Shot super resolution on 3D Spatio-Temporal Data: U-NO is resolution invariant in both temporal and spatial domains, whereas 3D FNO is not resolution invariant in time and cannot process different temporal resolutions.
- A.5 Zero-Shot super resolution on 3D Spatio-Temporal Data: U-NO implements domain contraction and expansion by mapping between input and output domains through a fixed homeomorphism.For the bounded connected domains studied, the mapping can be established by scaling.
- A.5 Zero-Shot super resolution on 3D Spatio-Temporal Data: U-NO has approximately 8–25 times more parameters than FNO, but its inference time is approximately only 1–4 times higher.The contraction reduces the time required for forward and inverse discrete Fourier transforms.