U-NO: U-shaped Neural Operators
Authors: Md Ashiqur Rahman
Zachary E Ross
Kamyar Azizzadenesheli
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (04/2023)
U-NO: U-shaped Neural Operators
Md Ashiqur Rahman
rahman 790purue. eduu
Deportment of Computer Science Purdue University
Zachary E. Ross
aross@caltech.co edu
Seismological Laboratory
California Institute of Technology
Kamyar Azizzadenesheli
kamyjrr Snvidia.com
NVIDIA Corporation
Reviewed OpenReview: https: opwprevieeeeveeueee net/ forum? id= j3oQFccoJd
Abstract
Neural operators generalize classical neural networks to maps between infinite-dimensionaal spaces, e.g., function spaces. Prior works on neural operators proposed a series of novel methods to learn such maps and demonstrated umprecedented success in learning solution operators of partial differential equations. Due to their close proximity to fully connected architoctures, 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), 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 hyperparameterss choices. We study the performance of U-NO on PDE benchmarks, namely, Darcy's '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 Vavier-Stokes equations, respectively, over the state of art. On Navier-Stokes 3D spatiotemporal operator learning task, we show U-NO provides 37% improvement over the state of the art methods.
Keywords:
Neural Operators, Portial Differential Equations
1 Introduction
Conventional deep learning research is mainly dominated by neural networks that allow for learning maps between finite dimensional spaces. Such developments have resulted in significant advancements in many pratticalsettings, from vision and nature language processing, to robotics and recommendation systems (Levine et al., 2016; Brown et al., 2020; Simonyan & Zisserman, 2014). Recently, in the field scientific computing, deep neural networks have been used to great success, particularly in simulating physical systems following differential equations (Mishra Molinaro, 2020; Brigham, 1988; Chandler & Kerswoll, 2013; Long et al., 2015; Chen al., 2019; Raissi & Karniadakis, 2018). Many real-world problems, however, require learning maps between infinite dimensional functions spaces (Evans, 2010). Neural operators generalizations of neural networks and allow to learn maps between infinite dimensional spaces, including function spaces (Li et al., 2020b). Neural operators are universal approximators of operators and have shown numerous applications in tackling problems in partial differential equations (Kovachki et al., 2021b).
Neural operators, an abstract form, consist of sequence of lincar integral operators, cach followed by nonlinear woint-wise operator. Various neural operator architectures have been proposed, most which focus the approximation of the Nyström type of the inner linear integral operator of neural operators (Li

---

Published in Transactions on Machine Learning Research (04/2023)
Nounal Operator Layer
(itssiiossb boo
G
Q
u
G
G
G
G
Figure 1: U-NO architecture. a is an input function, u is the output. Orange circles are point-wise operators, reclangles denote general operators, and smaller blue circles denote concatenations in function spaces.
et al., 2020a;;; Gupla et al., 2021; Tripura & Chakraborly, 2023). For instance, Li et al. (2020a) suggests lo use convolution kernel integration and proposes Fourier neural operator (FNO), a model that computes the linear integral operation the Fourier domain. The lincar integral operation can also be approximated using discrete wavelet transform (Gupta al., 2021; Tripura & Chakraborty, 2023). But FNO approximates the integration using the fast Fourier transform method (Brigham, 1988) and enjoys desirable approximation theoretic guarantee for continuous operator (Kovachki et al., 2021a). The efficient use of the built-in fast Fourier transform makes the FNO approach among the fastest neural operator architectures. Tran et al. (2021) proposes slight modification to the integral operator of FNO by factorizing the multidimensional Fourier transform along each dimension. These developments have shown successes in tackling problems in seismology and geophysics, modeling the so-called Digital Twin Earth, and modeling fluid flow in carbon capture and storage experiments (Pathak et al., 2022: Yang et al., 2021; Wen al., 2021; Shui al., 2020; Li et al., 2022). These are crucial components in dealing with climate change and natural hazards. Many of the initial neural operator architectures are inspired by fully-connected neural networks and result in models with high memory demand that, despile the successes of deep neural networks, prohibil very deep neural operators. Though there are considerable efforts in finding more effective integral operators, works investigating suitable architectures of Noural operators are scaree.
To alleviate this problem, we propose the U-shaped Neural Operator (U-NO) by devising integral operators that map between functions over different domains. We provide a rigorous mathematical formulation adapt the U-net architecture for neural networks for fnite-diimenioonal spaces to neural operators mapping between function spaces (Ronneberger et al., 2015). Following the u-shaped architecture, the U-NO- first progressively maps the input function to functions defined with smaller domains (encoding) and later reverses this operation to generate an appropriate output function (decoding) with skip connections from the encoder part (see Fig. 1). This efficient contraction and subsequent expansion of domains allow for designing the over-parametrized and memorr-efficient model, which regular architectures cannot achieve.
Over the first half (encoding) of the U-NO layers, we map each input function space to vector-valued function spaces with steadily shrunken domains over layers, all the while increasing the dimension of the co-domains, ie., the output space of each function. Throughout the second half (decoding) of U-NO layers, gradually expand the domain of cach intermediate function space while reducing the dimension of the output function co-domains and employ skip connections from the first. half to construct vector-valued functions with larger co-dimension. The gradual contraction of the domain the input function space in each layer allows for the encoding of function spaces with smaller domains. And the skip connections allow for the direct passage of information between domain sizes bypassing the bottleneck of the first half, i.e., the functions with the smallest domain (Long et al., 2015). The U-NO architecture can be adapted to the existing operator
2

---

Published in Transactions on Machine Learning Research (04/2023)
learning techniques (Gupta et al., 2021; Tripura & Chakraborty, 2023; Li et al., 2020a; Tran et al., 2021) to achieve a. more effective model. In this work, for the implementation of the inner integral operator of U-NO, we adopt the Fourier transform- based integration method developed in FNO. As U-NO contracts the domain of each input function at the encoder layers, we progressively define functions over smaller domains. At a fixed sampling rale (resolution of discretization of the domain), this requires progressively fewer data points to represent the functions. This makes the U-NO memory-efficient architecture, allowing for deeper and highly parameterized neural operators compared with prior works. It is known that over-parameterization is one of the main essential components of architecture design in deep learning (He et al., 2016), crucial for performance (Belkin et al., 2019), and optimization (Du et al., 2019). Also recent works show that the model overparameterization improves generalization performance (Neyshabur et al., 2017; Zhang et al., 2021; Dar et al., 2021). The U-NO archilecture allows efficient training of overparamettrrized model with smaller memory footprint and opens neural operator learning to deeper and highly parameterized methods.
We establish our empirical study on Darcy's flow and the Navier-Stoker equations, two PDEs which have served benchmarks for the study of neural operator models. We compare the performance of U-NO against the state of art FNO. We empirically show that the advanced structure of U-NO allows for much deeper neural operator models with smaller memory usage. We demonstrate that U-NO achieves average performance improvements of 26% on high resolution simulations of Darcy's flow equation, and 44% on Navier-Stokes equation, with the best improvement of 51%. On Navier-Stokes 3D spatio-temporal operator learning task, for which the input functions are defined on the 3D spatio-temporal domain, we show U-NO provides 37% improvement over the state of art FNO model. Also, the U-NO outperforms the bascline FNO zero-shot super resolution experiment.
It is important to note that U-NO is the first neural operator trained for mapping from function spaces with 3D domains. Prior studies on 3D domains either transform the input function space with the 3D spatio-temporal domain to another auxiliary function space with 2D domain for the training purposes, resulting in an operator that is time discretization dependent or modifics the imput function by extending the domain and co-domain (Li et al., 2020a). We further show that U-NO allows for much deeper models (3x) with far more parameters (25x), while still providing some performance improvement (Appendix A.2). Such a model can be trained with only few thousand data. points, emphasizing the data efficiency of problem-specific neural operator architecture design.
Moreover, we empirically study the sensitivity of U-NO against hyperparameters (Appendix A.2) where we remonstrate the superiority of this model against FNO. We observe that U-NO is much faster to train (Appendix A. A.d) also much easier to tune. Furthermore, for U-NO as a model that gradually contracts (expands) the domain (co-domain) of function spaces, we study its variant which applies these transformations more aggressively. A detailed empirical analysis of such variant of U-NO shows that the architecture is quite robust with respect to domain/co-domma transformation hyperparameters.
2 Neural Operator Learning
Let A and U denote input and output function spaces, such that, for any vector valued function E A, a : DA D: + Rda, with DA Rd and for any vector valued function U, u: DA Rdu, with De Rd. Given set of N data points, {(aj, u;)] we aim to train a neural operator Go A - U, parameterized by 0, to learn the underlying map from to u. Neural operators, in an abstract sense, are mainly constructed using a sequence of non-linear operators G; that are linear integral operators followed by point-wise non-linearity, i.e., for an input function vi at the ith layer, G₁ : v, - U:t1 is computed as follows,
Gevi()) == o ke(E, ()( + (ovs(x))
(1)
Here, Ky is a kernel function that acis, logether with the integral, as a global linear operator with measure μ4i W; is a matrix that acts as a point-wise operator, and a is the point-wise non-linearity. One can explicitly add a bias function b, separately. Together, these operations constitute G₁ i..., a nonlincar operator between a function space of de dimensional vector-valued functions with the domain D, C Rd and a function space of def1 dimensional vector-valued functions with the domain D C Rdd Neural operators often times are
3

---

Published in Transactions on Machine Learning Research (04/2023)
accompanied with a point-wise operator P that constitute the first block of the neural operator and mainly serves as lifting operator directly applied to input function, and point-wise operator Q, for the last block, that is mainly purposed as a projection operator to the output function space.
Previous deployments of neural operators have studied settings in which the domains and co-domains of function spaces are preserved throughout the layers, i.e,, D = D; Di+1, and d' = d4 = de+1, for all These models, e.g. PNO, impose that each layer is map between functions spaces with identical domain and co-domain spaces. Therefore, at any layer i, the imput function is v4 D + R and the output function is V++1 D Rd, defined the same domain D> and co-domain Fd This is one of the main reasons for the high memory demands of previous implementations of neural operators, which occurs primarily from the global integration step. In this work, we avoid this problem, since as we move to the innermost layers of the neural operator, we contract the domain each layer while increasing the dimension of co-domains. As a result. we sequentially encode the input functions into functions with smaller domain, learning compact. representation as well as resulting in integral operations in much smaller domains.
3 A U-shaped Neural Operator (U-NO)
In this section, we introduce the U-NO architecture, which is composed of several standard elements from prior works on Neural Operators, along with U-shaped additions tailored the structure of function spaces (and maps between them). Given imput function : DA + R we first apply point-wisc operator P to and compute vo A + Rdo, The point-wise operator P is parameterized with function Po Rdd R2 and acts as
(x) = Po(a(z)), Vo E Do
where Do = DA. The function Po can be matrix or more generally deep neural network. For the purpose of this paper, we take do >> dA, making lifting operator.
Then a soquence of L1 non-lincar integral operators is applied to to:
G {u: D; + R + (wi+1 D241 + Rd for E 20, L1),
where for each i, Dj C Rd is measurable set accompanied by measure 14g- Application of this sequence of operators results in a sequence of intermediate functions {oy D, + 18 Jiml in the encoding part of U-NO.
In this study, without loss of generality, we choose Lebesgue measure for μ4's. Over this first L1 layers, i.e., the encoding half of U-NO, we map the input function to a set of vector-value functions with increasingly contracted domain and higher dimensional co-domain, ie. for each i, we have v(D,) Z u(D;-1) and doyj1 2 du..
Then, given Uh,+1, another sequence of Ti2 non-linear integral operator layers (i.e. the decoder), is applied. In this stage, the operators gradually expand the domain size and decrease the co-domain dimension to ultimately match the domain of the lerget function, i.e., for each in these layers, have ((De+1) 2 ((D)) and de, 2 do....
To include the skip connections from the encoder part from the encoder part to the decoder in (sce Fig. 1), the operator Gr, in the decoder part, takes a vector-wise concatenation of both vr,+e and ur,-i as its input. The vector- wise concatenation of UL,+ and UL is a function o' where,
UL (z) Vo E DL,+i- 2 DL,+S + t and (()) = [UL,-((x)
of t constitutes the input to the operator GLyti, and we compute the output function of the next layer as,
UL = GitL, 2442,
Here, for simplicity, we assumed DL,+1 = D
More generally, concatenation can be defined with a map between the domains m D + Di as 1/4 (1) =
(m()))(( (((())).=,,++d(
4

---

Published in Transactions on Machine Learning Research (04/2023)
Computing 2241 for the layer L = L1 L2, conclude U-NO architecture with another point wise operator Q kernelized with the Qo IRdz+1 - dR such that. for u QUL+1 we have
u(x) Qu(v (x)), Vo D1+1
where DI+11 Dr. In this paper choose dr.+1 >> du such that Q is a projection operator. In the next section, we instantiate the U-NO architecture for two benchmark operator learning tasks.
4
Empirical Studies
In this section, we describe the problem settings, implementations, and empirical results.
4.1 Darcy Flow Equation
Darcy's law is the simplest model for fluid flow in a porous medium. The 2-d Darcy Flow can be described as a second-order lincar elliptic equation with a Dirichlet boundary condition in the following form,
V Vu(x)) f(x)
E D E 0D)
u(x) 0
where a E S L (D; represents the diffusion coefficient, u E is the velocity function, and fe H-1 (D;R) the forcing function. H represents with p 2.
HD; space
Sobolev
We take D= (0, 1)a and aim learn the solution operator G1, which maps diffusion coefficient function a to the solution i.c., (a).. Please note that due to the Dirichlet boundary condition, we have the value of the solution u along the boundary 0D
For dataset preparation, we define μ as a pushforward of the Gaussian measure by operator dे as μ (0, (-A+ - as a probability measure with zero Neumann boundary conditions on the Laplacian, i.e., the Laplacian is along the boundary. The [(0, (-A 97) describes Gaussian measure with covariance operator (-A - - The (xx) is defined as
3
i 0
(xx)
12
if r
The diffusion coefficients a(x) are generated according to a ~ and we fix f(x) = The solutions are obtained by using a second-order finite difference method (Larsson & Thomée, 2003) on a uniform 421 x 421 grid over (0, 1)2. And solutions of any other resolution are down-sampled from the high-resolution data. This setup the same as the benchmark setup in the prior works.
4.2 Dovier-Stokes Equation
The Navier-Stokes equations describe the motion of fluids taking into account the effects of viscosity and external forces. Among the different formulations, we consider the v formulation of the 2-d Navier-Stokes equations for viscous and incompressible fluid on the unit torus (T²), which can be described as,
O,w(x, t) V db Vw(r, t) = v/w(2, g(x),
x E T t (0,
- - w, w<x, 0) wo(x))
x E T2, t E (0, oo).
x T
Here u E R. T2 R2 is the velocity field. And is the out-of plane component of the vorticity field V x u (curl of u). Since we are considering 2-D flow, the velocity u can be written by splitting into components as
Jun (3n,)))) uz 22) v(21,22)
E T

---

Published in Transactions on Machine Learning Research (04/2023)
And it follows that V x u (0,0.w)
The stream function v is related to velocity by = V-b, where V L is the skew gradient operator, which is defined as
Oy. Og
A v = (-
Dy dr
The function 9 is defined as the out-of-plane component of the curl of the forcing function, J, i.e., (0, 0,
(A f) and g T + I,, and IR., is the viscosity coefficient.
We generated the initial vorticity wo from Gaussian measure as (wo ~ (0 7--5 -2.5 boundary condition with the constant mean function O and covariance 71.5 ( A + 491) We fix as
2.5) with periodic
x2) = 0.1( sin (2n(22(211 + 22)) + cos(2 + x2))), x2) E T
Following (Kovachki et al., 2021b), the equation is solved using the pseudo-specttal split-step method. The forcing terms are advanced using Heun's 's method (an improved Euler method). And the viscous terms are advanced using a Crank-Nicolson update with a time-step of 10-1 implicit method in time (Larsson Thomée, 2003). We record the solution every = 1 time unit. The data is generated on a uniform 256 x 256 grid and are downsampled to x for the low-resolution training.
update is a second-order
For this time-dependent problem, we aim to learn operator that maps the vorticity field covering time interval spanning [0, T. into the vorticity field for later time interval, (Ti
G "(OO Til:: H (T?: + C((Tin, T): Hper Hepee (T):
where Heer is the periodic Sobolev space H" with constant r 2 0. In terms of vorticity, the operator is defined by x| > (T..T.
4.3 Model Implementation
We present the specific implementation of the models used for our experiments discussed in the remainder of this section. For the internal integral operators, we adopt the approach of Li et al. (2020a), in which, evoking convolution theorem, the integral kernel (convolution) operators are computed by the multiplication of Fourier transform of kernel convolution operators and the input function in the Fourier domain. Such operation is approximately computed using fast Fourier transform, providing computational benefit compared to the other Nyström integral approximation methods developed in neural operator literature, and additional performance benefits results from computing the kernels directly in the Fourier domain, evoking the spectral convergence (Canuto et al., 2007). Therefore, for given function ua i.e. the input to the G we have,
-0 -o (F (B.))) (())) (Waa(()))))(()))
G,v(()) Gev(()==
Here F and F-1 are the Fourier and Inverse Fourier transforms, respectively. Zd - Cd.1xx, and for each Fourier mode E zd, it is the Fourier transform of the periodic convolution function in the integral convolution operator. We directly parameterize the matrix valued functions R, on zd in the Fourier domain and learn from data. For each layer i, we also truncate the Fourier series at preset number of modes,
/ = 1k E Ed ki I Kmaza for j E (1,...d)l
Thus, R2 is implicitly complex valued (Binez X do..i dr,) tensor. Assuming the discretization of the domain is regular, this non-lineat operator is implemented efficiently with the fast Fourier transform. Finally, we define the wiint-wise residual operation WTe where
Ww,(x)) Ws,, (8(x))
with x E Di+1 and 8: D,41 Dj is fixed homeomorphism between D and D..

---

Published in Transactions on Machine Learning Research (04/2023)
In the empirical study, need two classes of operators. The first one is an operator setting where the input and output functions are defined on 2D spatial domain. And in the second setting, the input and output are defined on 3D spatio temporal setting. In the following, we describe the architecture of U-NO for both of these scttings.
Operator learning functions defined on 2D spatial domains. For mapping between functions defined on 2D spatial domain, the solution operator is of form G : {a (0, 1)2 - Rda] - { (0, 1)2 + R }. The U-NO architecture consists of a series of seven non-linear integral operators, {Gi) imo? that are placod between lifting and projection operators. Each layer performs 2D integral operation in the spatial domain and contract (or expand) the spatial domain (see Fig. 1).
The first non-linear operator, Go, contracts the domain by factor of 3 uniformly along each dimension, while increasing the dimension of the co-domain by a factor of 5
Go {vo : (0, = + fur (0,3/4) + gl2d.o1..
Similarly, G1 and G2 sequentially contract the domain by factor of 3 and 1. respectively, while doubling the dimension of the co-domain. G3 is a regular Fourier operator layer and thus the domain or co-domain of its input and out function spaces stay the same. The last three Fourier operators, {G:] i-4a expand the domain and decrease the co-domain dimension, restoring the domain and co-domain to that of the input to Gu-
Ground Truth
Diffusion Coefficients a: (0.1)"
Solution Function # (0.1)" R
Initial Vorticity
Predictions
t- 10
Ve: (0.0.)) - gea
Vg (0,0.5) - g
(0.0.25) Rins
:0,00225) R!28
t- 14s
B
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 U-NO with a factor of 1 2 on an instance of Darcy flow equation. (B) Vorticity field generated by the U-NO variant as a solution to an instance the two-dimensional Navier-Stokes equation with viscosity 1
We further propose U-NOt which follows more aggressive factor of (shown Fig. 2A) while contracting (or expanding) the domain in each of the integral operators. One of the reasons for this choice is that such a selection of scaling factors is more memory efficient. than the initial U-NO architecture.
7

---

Published in Transactions on Machine Learning Research (04/2023)
The projection operator Q and lifting operator P are implemented as fully-connected neural networks with lifting dimension dg 32. And we set skip connections from the first three integral operator (Go, G1, G2)
respectively to the last three (Go: Go, G4).
For time dependent problems, where we need to map between input functions a of an initial time interval [0, m, to the solution functions u of later time interval (Tin; T), iee,, to learn an operator of form,
G : ((0, Tinl, {a: : (0, - R(4)) + C(((Tin, T, (0, - Rd4))
we use an auto regressive model for U-NO and U-NO', Here recurrently compose the neural operator time to produce solution functions for the time interval (Tim, TT.
Operator learning on functions defined on 3D spatio-temporal domains. Separately, we also use a U-NO model with non-linear operators that performs the integral operation in space and time (3D), which directly maps the input functions of interval [0, Ti] to the later lime interval (Tin; T), without any recurrent composition in time. As this model does not require recurrent composition in time, it is fast during both training and inference. We redefine both a and u on 3D spatio-temporal domain and learn the the operator
: fa: (0, x [0. Tinl + Rd} + fu: (0, 1). x (Tin, T] + Rd}
The non-linear operators constructing U-NO, (G. are defined as,
G. {or (0, x T. k + {vt+1 (0, x (Tins Tin + cTinl] - R.
Here, (0, az)² x Ti is the domain of the function v, And cf and of are respectively the expansion (or contraction) factors for spatial domain, temporal domain, and co-domain for i'th operator. Note that To [0. T (0,00) (0,)l+1) (0, 1), and TI+1 = (Tin; 7].
Here we also first contract and then expand (i.e., of < in the encoding part and of 2 1 in the decoding part) the spatial domain following the U-NO performing 2D integral operation. And we set ct 2 1 for all operators G; if the output time interval is larger than input time interval i.c. Tin < Tout Tin i.e., we gradually increase the time domain the input function to match the output. Here the projection operator Q and lifting operator are also implemented as fully connected neural network with do for the lifting operator P.
For the experiments, we use Adam optimizer (Kingma & Ba, 2014) and the initial learning rate is scaled down by a factor of 0.5 every 100 epochs. As the non-linearity, we have used the GELU (Hendrycks & Gimpel, 2016) activation function. We use the architecture and implementation of FNO provided in the original work Li et al. (2020a). All the computations are carried on single Nvidia GPU with 24GB memory. For each experiment, the data divided into train, test, and validation sets. And the weight of the best-performing model on the validation set used for evaluation. Each experiment repeated 3 times and the mean is reported.
In the following, we empirically study the performance of these models on Darcy Flow and Navier-Stokes Equations.
4.4 Results on Darcy Flow Equation
From the dataset of N 2000 simulations of the Darcy flow equation, we set aside 250 simulations for testing and use the rest for training and validation. We use U-NOt with 9 layers including the integral, lifting, and projection operators We train for 700 epochs and save the best performing model on the validation set for evaluation. The performances of U-NOt and FNO on several different grid resolutions are shown in Table. 1. The number of modes used in each integral operator slays unchanged throughout different resolutions. We demonstrate that U-NO* achieve lower relative error on every resolution with 27% improvement on high resolution (421 x 421) simulations. The U-shaped structure exploits the problem structure and gradually encodes the input function to functions with smaller domains, which enables U-NO to have efficient deep architectures, modeling complex non-linear mapping between the input and solution function spaces. We
8

---

Published in Transactions on Machine Learning Research (04/2023)
also notice that U-NO requires 23% less training memory while having 7.5 times more parameters. With vanilla FNO, designing such over paramterized deep operator to learn complex maps is impracticable due to high memory requirements.
Table 1: Benchmarks Darcy Flow. The average relative error in percentage is reported with error bars are indicated in the superscript. The average back-propagation memory requirements for single training instance is reported over different resolutions are reported. U-NO" performs better than FNO for every resolution, while requiring 23% lower training memory which allows efficient training of 7.5 times more parameters
Mem.
Resolution
Model
Req. (MB)
# Parameter
s= 421
== 211
== 141
= 85
x x106
U-NO
166
8.2
0.57
0.58
0.60
0.73 10.00.ee-2
FNO
214
1.1
0.78:4.0002
0.84
0.84
0.87#3.06
Table 2: Zero-shot super resolution result on Darcy Flow. Neural operators trained on lower spatial resolution data directly tested on higher resolution with no further training. percentage relative error rate.
We observe U-NO archives lower
Test
141
211
421
Train
8=85
4.7
6.2
8.3
8- 141
2.6
6.3
U-NO
s=211
4.5
8=85
7.5
14.1
23.9
8-141
4.3
13.1
FNO
s=211
9.6
4.5 Results Navier-Stokes Equation
For the regresegressive model, we follow the architecture for U-NO and U-NO' described in Sec. 4.3. For U-NO performing spatio-temporal (3D) integral operation, following the same spatial contracting and expansion of 2D spatial U-NO. Additionally, as the positional embedding for the domain (unit torus T²) we used the Clifford torus (Appendix A.1) which embeds the domain into Euclidean space. It is important to note that, for spatio-temporal FNO setting, the domain of the input function is extended to match the solution function as it maps between the function spaces with same domains. Capable of mapping between function spaces with different domains, no such modification is required for U-NO.
7 stacked non-lincar integral operator
For each experiment, 10% of the total number of simulations N set aside for testing, and the rest is used for training and validation. We experiment with the viscosity v E (1e-3, le le-4, le-5, 1e-6), adjusting the final time T as the flow becomes more chaotic with smaller viscosities.
Table demonstrates the results of the Navier-Stokes experiment. We observe that U-NO achieves the best performance, with nearly a 50% reduction in the relative error over FNO in some experimental settings; it even achieves nearly a 1.76% relative error on problem with Reynolds number2 of 10a (Fig. 2B). In each case, the U-NO! with aggressive contraction ( and expansion) factor, performs significantly better than FNO-while requiring less memory, Also for neural operators performing 3D integral operation, U-NO achieved 37% lower relative error while requiring 50% less memory on average with 19 times more parameters. It is important to note that designing an FNO architecture matching the parameter count of
Reynolds number is estimated Re
VO. (/23)6/23
(Chandler & Kenswell, 2013)

---

Published in Transactions on Machine Learning Research (04/2023)
Table 3: Benchmarks on Navier-Stokes (Relative Error (%)) For the models performing a 2D integral operator with recurrent structure in time, packppropagttio memory requirement per time step reported. U-NO yields superior performance compared with U-NO* and FNO, while U-NO" is the most memory efficient model. In the 3D integration setting, U-NO provides significant performance improvement with almost three times less memory requirement.
Avg.
1e-3
1e-4 T = 30s
1e-5
V/ 1e T 158
model
Mem. Req.
# Parameters
T = 50s
T 20s
Tin = 10s
Tin 10s
Tin 10s
T. 6s
(MB)
(x10%)
N 5000
N = 11000
N = 11000
= 11000
U-NO
16
15.3
0.28 1218-2
3.44 1633
2.94 4.99 -
1.76 1:467
2D
U-NO
11
6.7
0.35=2.00-2
3.56 4.3e-2
3.60#1.2e-2
2.2=1.5e-2
FNO
13
1.3
0.58-1144--
5.57 2
5.12-0.77
3.33 1.33-
U-NO
108
24.0
0.31 13.9c-2
5.59
14.6e-..6e-1
5.10
3D
FNO
216
1.3
0.68 0.68 0.
9.60
8.67#1.2e-1
7.35 6.3e 6.3e-1
U-NO is impractical due to high computational requirements (see Fig. 3). FNO architecture with similar parameter counts demands enormous computational resources with multiple GPUs and would require a long training time due to the smaller batch size. Also unlike FNO, U-NO is able to perform zero-shot super-resolution in both time and spacial domain (Appendix A.5).
Table 4: Relative Error (%) achieved by 2D U-NO and 3D U-NO trained and evaluated high resolution (256 256) Simulations of Navier-Stoker Even with highly constrained neural operator architecture and smaller training data, the U-NO architecture can achieve comparable performance.
Models U-NO
Memory Requirement
le-3
== le-4
le-5
le-6
(MB)
N = 2400
N = 6000
N = 6000
N 6000
2D U-NO
86
0.51
5.9
7.0
4.4
3D U-NO
850
0.83
8.3
11.2
8.2
Due to U-shaped encoding decoding structure which allow us to train U-NO on very high resolution (256 x 256) data on a single GPU with a 24GB memory. To accommodate training with high-resolution simulations of the Navier-Stocks equation, we have adjusted the contraction (or expansion) factors. The Operators the encoder part of U-NO follow contraction ratio of and 1 for the spatial domain. But the co-domain dimension increased only by factor of 2. And the operators in the decoder part follow the reciprocals of these ratios. We have also used smaller lifting dimension for the lifting operator P. Even with such compact architecture and lower training data, U-NO achieves lower error rate on high resolution data (Tablle 4). These gains are attributed to the deeper architecture, skip connections, and encoding/ encoding/decoding components of U-NO.
4.6 Remarks Memory Usage
Compared to FNO, the proposed U-NO architectures consume less memory during training and testing. Due to the gradual encoding of input function by contracting the domain, the U-NO allow for deeper architectures with a greater number of stacked non-linear operators and higher resolution training data (Appendix A.4). For U-NO, an increase in depth, does not significantly raise the memory requirement. The training memory requirement for 3D spatio-temporal U-NO and FNO is reported in Fig. 3A.
We can observe linear increase in memory requirement for FNO with the increase depth. This makes FNO unsuitable for designing and training very deep architectures as the memory requirement keep increasing
10

---

Published in Transactions on Machine Learning Research (04/2023)
rapidly with the increase of the depth. On the other hand, due to the repeated scaling of the domain, do not observe any significant increase in the memory requirement. for U-NO. This make the U-NO suitable architecture for designing deep neural operators. The improvement of the performance with the increase of
700
FND
Mortary Requiremmnt (MB)
600
500
400
300
200
100
11
5
8
Depth
Depth
B
Figure 3: (A) Training memory requirements (in MB) for the 3D spatio-temporal problem of Navier-Stokes equation (v = 1c-3). For different depth only the number of stackod non-lincar operators 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 400MB for FNO). (B) Relative error in percentage on 3D spatio-temporal Navier-Stokes equation (v 1e-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.
depth is shown Fig. 3B. This is expected as highly paramcterized deep architoctures allow for effective approximation of highly complex non-linear operators. U-NO style architectures are more favorable for learning complex maps between function spaces.
5 Conclusion
In this paper, we propose U-NO, new neural operator architecture with a multitude of advantages compared with prior works. Our approach is inspired by Unet, neural network architecture that is highly successful at learning maps between fmile-dimensional spaces with similar structures, e.g., image to image translation. U-NO possesses a similar set of benefits to Unet, including data efficiency, robustness to hyperparameter tuning (Appendix A.2), flexible training, memory efficiency, convergence (Appendix A.6), and non-vanishing gradients. This approach, while advancing neural operator research, significantly improves the performance and learning paradigm, still carries the fundamental and inherent memory usage limitation of integral operators. Nevertheless, U-NO allows for much deeper models incorporating the problem structures and provides highly parameterized neural operators for learning maps between function spaces. U-NOs are very easy to tune, possess faster convergence to desired accuracies, and achieve superior performance with minimal tuning.
11

---

Published in Transactions on Machine Learning Research (04/2023)
References
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine-learning practice and the classical bias-variance trade-off. Proceedings of the National Acodemy of Sciences, 116(32):
18849-15854 2019.
E Oran Brigham. The fast Fourier transform and its applications. Prentice-Hall. Inc., 1988.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877 1901, 2020.
Claudio Canuto, M Yousuff Hussaini, Alfio Quartcroni, and Thomas Zang. Spectral methods: fundamentals in single domains. Springer Science & Business Media, 2007.
Gary J Chandler and Rich R Kerswell. Invariant recurrent solutions embedded in a turbulent two-dimensional kolmogorov flow. Journal of Fluid Mechanics, 722:554-588 595, 2013.
Zhengdao Chen, Jianyu Zhang, Martin Arjovsky, and Léon Bottou. Symplectic recurrent neural networks. arXiv preprint Xr::::190.11 13334, 2019.
Yehuda Dar, Vidya Muthukumar, and Richard G Baraniuk. A farewell to the bias-variance tradeoff? an overview of the theory of overparameterized machine learning. arXiv preprint wxXiu2219.00355, 2021.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference machine learning, pp. 1675-1685. PMLR, 2019.
Lawrence C Evans. Partial differential equations, volume 19. American Mathernatical Soc., 2010.
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelett operator learning for differential equations. Advances in Neural Information Processing Systems, 34, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). ar preprint arXiv:1606. 004.11666.0,45, 2016.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. wxXiv:141.69980 2014.
ar preprint
Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds for fourier neural operators. Journal of Machine Learning Reseorch, 22:Art- Noo, 2021a.
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Avizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint orXiv:2108.08. 2021b.
Stig Larsson and Vidar Thomée. Partial differential equations with numerical methods, volume 45. Springer, 2003.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. The Journal of Machine Learning Research, 17(1):1334-133,, 2016.
Bian Li, Hanchen Wang, Xiu Yang, and Youzuo Lin. Solving seismic wave equations on variable velocity models with fourier neural operator. arXiv preprint
2022.
Zongyi Li, Nikola Kovachki, Kamyar Azizzdencsheli,, Burigode Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint ar Xiv:2000. 08895, 2020a.
12

---

Published in Transactions on Machine Learning Research (04/2023)
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differential equations. arXiv preprint (xXt::2003003483 2020b.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhattacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial differential equations. Advances in Neural Information Processing Systems, 33, 2020c.
Jonathan Long, Evan Shelhamer, and Trevor Darrell. Fully convolutional networks for semantic segmentation. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 3431 3440, 2015.
Siddhartha Mishra and Roberto Molinaro. Estimates on the generalization error of physics informed neural networks (pinns) for approximating a class of inverse problems for pdes. arXiv preprint arXiv:2007.01138 2020.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nati Srebro. Exploring generalization in deep learning. Advances in neural information processing systems, 30, 2017.
Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Avizædenesheli, et al. Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202. 11214, 2022.
Maziar Raissi and George Em Karniadakis. Hidden physics models: Machine learning of nonlinear partial differential equations. Journal of Computational Physics, 357:125-141, 2018.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference Medical image computing and computer-assisted intervention, pp. Springer, 2015.
Changjian Shui, Qi Chen, Jun Wen, Fan Zhou, Christian Gagné, and Boyu Wang. Beyond H-divergence: Domain adaptation theory with jensen-shannon divergence. ar preprint urXin:2007.1533 15567, 2020.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint crXio:1409.1556, 2014.
Alasdair Tran, Alexander Mathews, Lexing Xic, and Cheng Soon Ong. Factorized fourier neural operators. arXiv preprint axX.:2211..1802,, 2021.
Tapas Tripura and Souvik Chakraborty. Wavelet neural operator for solving parametric partial differential equations in computational mechanics problems. Compuler Methods in Applied Mechanics and Engineering, 404:115783, 2023.
Gege Wen, Zongyi Li, Kamyar Avizzadeneshelii Anima Anandkumar, and Sally M Benson. U-Ino an enhanced fourier neural operator based-deep learning model for multiphase flow. arXie preprint 1rXiv:2109. 036697, 2021.
Yan Yang, Angela F Gao, Jorge C Castellanos, Zachary E Ross, Kamyar Azizzadeneshelli and Robert W Clayton. Scismic wave propagation and inversion with neural operators. The Seismic Record, 1(3):126-33,, 2021.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107-115, 2021.
13

---

Published in Transactions on Machine Learning Research (04/2023)
A Appendix
A.1 Positional Embedding For Unit Torus
A unit torus T homcomorphic to the Cartesian product two unit circles: S S We can also consider the carlesian product the embedding of the circle. which produces Clifford torus. The Clifford torus can be defined as
S x 5 = 1 (0n(())cos(())) sin())cos(()) |0 E 2m, 0 2. Φ 1 2x 2
2
A.2 Sensitivity to Hyper-parameter Selection
Lesming Rate
Error Rate
Error
Generalization Gap
0.179
0.1002
9.399
0.4276
0.3114
0,2593
0.2484
0,8189
6,7034
0.6339
0.6104
0.402
1,6057
11
0.2518
9.3225
9.36
5,3938
0.9515
0,4896
02283
0,2249
0.1991
0.5156
33
0,8946
0.7012
0,5783
0.5944
2,3528
13
0.005
5.4159
0.3534
0.3971
1.4233
0.2851
0.2709
0.2339
0.2573
0.449
0.8125
0.7558
0.6551
0.7015
1.5204
15
0.4344
0.484)
0,4212
9.4466
14714
0.2761
0,4070
0,8244
0.4858
0,4275
Figure 4: Result of sensitivity of U-NO* to learning rate and number of stacked non-linear operator (depth) used. All models are trained the dataset of the Darcy Flow equation with resolution 211 x 211 following the training protocol described in 4.4. 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 (2 0.01) or very low (S 0.0001) learning rate, U-NO* 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 generalization gap.
configuration setting U-NO has low
Table 5: Number of Parameters in U-NO models different depth. We observe U-NO architecture can efficiently learn large number of parameter (25 times of the FNO) from very limited training data (only 1500
simulations) and can perform reasonably (See Table. 4)
Depth
7
9
11
13
15
17
Parameters (x (x10%)
5.3
8.3
9.6
18.5
22.8
27.8
A.3 PNO with Skip Connection
Table 6: Performance of vanilla 2D FNO equipped with skip connection Navier-Stokes equations described in Section 4.5 along with memory requirement during training. The result of 2D FNO is reported again for comparison. We notice that skip connection alone does not improve the performance of FNO.
Models U-NO
Memory Requirement (MB)
V = le-3
= le-4
V = le-5
1e-6
= 2400
N 6000
= 6000
N 6000
FNO w. skip con. FNO
13.14 MB
0.011
0.074
0.075
0.049
13.03 MB
0.009
0.072
0.074
0.052
A.4 Spatial Memory
The memory requirements to learn the operator (l(o,1))x( + (1.1)/1(101111 during prooapropng for the Navier-Stokes problem shown in Table 7. On average, for various tested resolutions, U-NO with
14

---

Published in Transactions on Machine Learning Research (04/2023)
contraction (and expansion) factor of 2/2 requires 40% less memory than FNO during training. This makes U-NO and its variants more suitable for problems where high-resolution data are crucial, e.g,, weather forecasting model (Pathak et al., 2022).
Table 7: Memory requirements (in MB) for a single training instance of the Navier-Stokes equations for different grid resolutions. All models have seven non-linear integral operator layers.
Spatial Resolution
FNO
U-NO
U-NO
64 x 64
13.0
16.8
11.3
128 x 128
76.1
67.3
45.4
256 x 256
304.5
269.0
181.5
512 512
1218.0
1076.0
726.0
1024 1024
4872.0
4304.0
2990.9
A.5 Zero-Shot super resolution on 3D Spatio- Temporal Data
Table Zero-shot super resolution result on 3D spatio-tempores Navier Stocks equation. The 3D FNO not resolution invariant in the time domain and due the specific construction can not process data at different. temporal resolution. But U-NO resolution invariant both in time and spatial domain.
Temporal Res.
fps
fps 1.5
fps
fps
Spatial Res.
s= 64
5.10
17.43
19.39
20.32
s=128
6.34
17.61
19.56
20.62
U-NO
8 256
8.16
17.86
19.94
20.83
A.6 Superior Convergence Rate
FND Thain
PND
101
FND Veldwiion
107
FND Val idwiioo
Realive Error Rate
10"
100
200
300
450
500
100
200
300
430
500
Eoochs
Epochs
A
B
Figure The training and test set error (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.
A.7 Domain Contraction and Expansion
In this work, the domain contraction (or expansion) is performed by setting homcomorphism or mapping between the domains of the input and output function of an integral operator. Following the notations of Sec.
15

---

Published in Transactions on Machine Learning Research (04/2023)
4.3, integral operator be defined
[.V]](a)) = V (F ((v)))( -Wox(s(2))
Where with I E Di+1 and 8 Dit1 - Dr a fixed homeomorphism between Di+1 and D,. For the discussed problems this work (Darcy flow and Navier Stokes equation), the domains the input and output functions are bounded and connected. As a result, the mapping can be trivially established by scaling operation. Lets the domain of the input and the output functions be 6)2 and (c, respectively, and the function s [c, d] - [a. b) is homcomorphism between them. The function has lincar form s(x) + (x c) with = (b- a) @ (d-c) c) where and represents element wise multiplication and division. If c = a = 0, becomes scaling operation on the domain. And the resultant output function can be efficiently computed through interpolation with appropriate factors along each dimension.
A.8
Training Memory Requirement for 2D U-NO at Different Depths
27.5
FNO
Memory Requirement (MB)
25.0
U-NO
U-NO
22.5
20.0
17.5
15.0
12.5
10.0
6
7
10
11
12
13
Depth
Figure 6: Memory requirements (in MB) for a single training instance of the Navier-Stokes equation during back-propagation input and outputs of resolution 64 x 64. Number of stacked non-lincar operators is varied and the additive memory requirement of U-NO architectures for deeper models are negligible compared to FNO model. For U-NOT the addition of 7 stacked non-linear integral operator increases the memory requirement by only 2.5 MB
A.9 Comparison of Inference Time
As across different problem settings, the number of parameters of U-NO is ~ 8 25 times more than the FNO; it has more floating point operations (FLOPs). For this reason, the inference time of U-NO is more than FNO (see Table 9). But it is important to note that even if U-NO has ~ 8 25 times more parameters, the inference time only ~ 1 - 4 times more. As U-NO in the encoding part gradually contracts the domain of the functions, it reduces the time in the forward and inverse discrete Fourier transform.
16

---

Published in Transactions on Machine Learning Research (04/2023)
Table 9: The rumning time of U-NO and FNO during inference of a single sample. For the Darcy Flow problem, we test on the resolution of 421 x 421, and for Navier Stocks, we use the problem setting with viscosity 10-3 with the resolution of 64 64. For Navier-Stocks 2D time to infer a single time step and for Navier-Stocks 3D time to infer the whole trajectory are reported.
Model U-NO
FNO (sec)
Problem
(sec)
Darcy Flow
0.13
0.12
Navier-Stcoks (2D)
0.08
0.02
Navier-Stocks (3D)
0.33
0.09
17


Reviews:
[{"review_id": "vWlz2QCC0k", "reviewer": "Reviewer_ahuP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strong points:\n\n1. This paper proposes a useful extension of FNO.\n\n2. The paper is well motivated\n\n3. The paper is well written with great-looking illustrative figures.\n\nWeak points:\n\n1. The contribution of the paper is limited. Replacing fully-connected neural networks and convolutional neural networks in FNO by a U-Net in U-NO is incremental.\n\n2. There is no theoretical justification for the proposed U-NO. For example, what is the theoretical result that supports the claim that U-NO allows deeper neural operators?\n\n3. The empirical results in the paper are not convincing.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes U-shaped Neural Operator (U-NO) that combines U-Net architecture and neural operators to allow deeper models. Other advantages of U-NO, as claimed in the paper, include fast training, data efficiency, and robustness with respect to hyperparameter values. Experimental results on Darcy’s flow and Navier-Stokes are provided to justify the advantage of the proposed method.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. The proposed U-NO requires much more parameters than the baseline FNO. Can the authors discuss this more? \n\n2. The average relative errors for Darcy Flow and Navier-Stokes given in Table 1 and 3, respectively, do not match those in the original paper [1]. Also, the average relative errors in these tables are incomprehensible. I think there is a formatting bug there.\n\n3. A comparison of the runtime of U-NO and FNO is needed.\n", "broader_impact_concerns": "The paper has no ethical concern or implication that would require adding a Broader Impact Statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "iaEQ0JH80g", "reviewer": "Reviewer_tUW5", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* In general, the paper's contributions are solid both in formulation and empirically. The paper supports its claims and evidence fairly well.\n* Notation, while dense, is precise and after some re-reading, I could follow through.\n* The paper demonstrates good applications for the scientific computing / physical simulation industry\n\n\nWeaknesses:\n* The paper can potentially relax the notation, or at least give better high-level overviews of the core idea (e.g. see my summary of contributions). I needed to read the paper multiple times until I understood the high-level strategy. \n* Furthermore, because I am not a physics expert, I briefly read, but did not fully appreciate/understand the Darcy Flow Equations / Navier-Stokes Equations nor their baseline solvers - some of their details appeared \"out of nowhere\". These could also be explained more accessibly to the common reader.\n* Please clarify: Is the main contribution that U-NO is more memory/compute-efficient (and therefore one can stack more depths to improve accuracy)? For now, it's unclear why skip-connections are needed (do they empirically help, or could they even be removed?)\n* Could you explain more on the semantic meaning of the domain compression? i.e. Because it works, does it suggest that the \"solution spaces\" of certain differential equations lie in lower dimensions?\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "**Note:** My review will be from the perspective of someone unfamiliar with this type of field (i.e. solving partial differential equations with neural networks)\n\nAt a base case, Neural Operators are the \"neural network\" equivalent over functions, rather than simple tensors, where \"input functions\" are mapped to \"output functions\" via a MLP-like transformation seen in Eq (1). Unfortunately, the regular Neural Operator forces the \"output function\" to have the same exact domain/co-domain as the \"input function\". This has the consequence of large memory usage, especially since the Neural Operator needs to perform an integration over the domain.\n\nInstead, the paper proposes a U-Net equivalent of Neural Operators, in which the operators on the \"encoder\" side can shrink the output function's domain, which leads to lower computation. Like the U-Net architecture, the outputs are also passed to the decoder side for concatenation.\n\nExperiments are conducted on well-known differential equations from physics, in which the \"input functions\" are problem parameters or initial conditions, mapped to the final solutions. Results demonstrate better accuracy (especially with more operator layers), in addition to much lower memory usage, even when multiple depths are used.\n\n\n\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please address my \"Weaknesses\" section above, as they are the largest requested changes.\n\n\nMinor Potential Typos/Questions for clarification:\n1. In Equation (1), do you mean $\\kappa_{i}(x,y)v_{i}(x) d \\mu_{i}(y)$? If not, what is the meaning of $v_{i}$ (without the input $x$)?\n2. Page 6, middle of page - an additional parenthesis is unneeded, $W_{i} v_{i}(x) = W_{i} v_{i}(s(x))$", "broader_impact_concerns": "No borader impact concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "f2vM21Ah06", "reviewer": "Reviewer_L1LD", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Disclaimer:** I am no expert in neural operators and do not know in depth the related literature. I am an expert in deep learning with a good mathematical knowledge in PDEs and functional analysis. I had read (Li et al. 2020) and (Kovachki et al. 2021) but I have never worked in the field. My review should therefore be read as a low-confidence review in which I mostly focused on methodology and presentation. I would differ opinions about prior work and deep technical details to the other reviewers.\n\n## Strenghts\n\n1. **Simple, elegant and effective technique:** The idea of using a UNet style architecture to circumvent the issues of prior FNO-style architectures is a natural decision with seemingly great results.\n2. **Clear writing**: The paper is very well written and very enjoyable to read.\n3. **Strong and thorough evaluations:** Methodologically, I really appreciate the effort made by the authors in comparing the performance of different versions of UNO to FNO with respect to many different metrics. The results in the Appendix nicely complement the main paper giving further insights about UNO (see for example ablation study using skip connections on FNO).\n4. **Reproducibility**: The authors provide easy-to-follow code implementing and reproducing their main training loops used with UNOs on the two datasets.\n\n## Weaknesses\n\n1. **Unclear how the domain contraction affects memory:** It might be due to my lack of knowledge of the field, but reading the paper I do not fully understand why contracting the domain affects the memory usage of the architecture. I clearer explanation of how this is implemented in the text would be helpful for less expert readers.\n2. **No details about FNO architecture used**: As I understand it, the key message of this paper is that UNO is a more powerful architecture than FNO. However, as far I can tell, FNO is not a particular architecture, but a familiy of architectures, and depending on the specific choices of number of layers and widths, the performance of an FNO would vary. In this sense, I am missing the details in the text (or the code for that matter) explaining which specific architecture was used as baseline FNO. In particular, I am interested in knowing if you can scale the parameter count of an FNO as high as a UNO, does it perform as well albeit with a much higher memory footprint?\n3. **Some mathematical notation is not defined**: Things like the push-forward operator $\\psi_#$, GP notation, etc are not defined in the text. I would in general make sure that every symbol that is used in the text is introduced somewhere.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a novel parameterization of Neural Operators inspired by the classical UNet architecture used to encode image-to-image functions in computer vision. This novel architecture, named UNO, compared to the prior state-of-the-art Fourier Neural Operator (FNO) is more memory efficient, trains faster, is less sensitive to hyperparameters and achieves higher generalization performances across resolutions. UNOs work by contracting the domain and expanding the co-domain of the output function in the first half of the layers, and doing the opposite in the second half of the layers. In the spirit of the classical UNets, skip connections are used to connect the representations that live in the same functional spaces at mirroring layers. Results on two standard benchmarks for Neural Operators show that UNOs exceed the performance of FNOs in all reported metrics.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I would suggest that the authors work on addressing my weaknesses, mostly focusing on 1., and 2., in their revised version of the text. I believe that will make for a much stronger submission.", "broader_impact_concerns": "I see no clear broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: I would like to cite one of the reviewers here:
"After reading the other reviews and the comments of the other reviewers, I stand by my previous assessment and I will strongly recommend acceptance of this work. I believe this paper clearly meets the bar for acceptance on the two acceptance criteria of TMLR as it is factually correct/methodologically precise, and provides a very actionable and clear contribution to readers interested in Neural Operators. The results of this paper are strong and reliable, and I believe the authors have clearly shown they followed proper evaluation practices to arrive at them. The paper is also very clear and easy to read. Specifically, the authors have made a nice effort in improving the accessibility of the paper to non-experts in Neural Operators."

I fully agree with this statement and, thus, I strongly suggest accepting the paper.
Decision: Accept (poster)