Mesh-Independent Operator Learning for PDEs using Set Representations
Authors: Seungjun Lee
Venue: Submitted to ICLR 2023 (2023)
Keywords: partial differential equations
operator learning
set representations
attention-based model
implicit neural representation

Abstract: Operator learning, learning the mapping between infinite-dimensional function spaces, has been attracted as an alternative approach to traditional numerical methods to solve partial differential equations (PDEs). In practice, the functions of the physical systems are often observed by sparse or even irregularly distributed measurements, thus the functions are discretized and usually represented by finite structured arrays, which are given as data of input-output pairs. Through training with the arrays, the solution of the trained models should be independent of the discretization of the input function and can be queried at any point continuously. Therefore, the architectures for operator learning should be flexibly compatible with arbitrary sizes and locations of the measurements, otherwise, it can restrict the scalability when the observations have discrepancies between measurement formats. In this paper, we propose to treat the discretized functions as set-valued data and construct an attention-based model, called mesh-independent operator learner (MIOL), to provide proper treatments of input functions and query coordinates for the solution functions by detaching the dependencies on input and output meshes. Our models pre-trained with benchmark datasets of operator learning are evaluated by downstream tasks to demonstrate the generalization abilities to varying discretization formats of the system, which are natural characteristics of the continuous solution of the PDEs. 

func- Under review as a conference paper at ICLR 2023
MESH-INDEPENDENT OPERATOR LEARNING FOR PDEs USING SET REPRESENTATIONS
Anonymous authors Paper under double-blind review
ABSTRACT
Operator learning, which learns the mapping between tion spaces, is an attractive alternative to traditional numerical methods for solving partial differential equations (PDEs). In practice, the functions of the physical sys- tems are often observed by sparse or even irregularly distributed measurements; thus, the functions are discretized and usually represented by finite structured ar- rays, which are given as data of input-output pairs. Through training with the arrays, the solution of the trained models should be independent the discretiza- tion of the input function and can be queried any point continuously. Therefore, the architectures for operator learning should be flexibly compatible with arbitrary sizes and locations of the measurements, otherwise, scalability be restricted when the observations have discrepancies between measurement formats. In this study, we propose the proper treatment of discretized functions as set-valued data and construct an attention-based model, called meshhindeeenndnn operator learner (MIOL), provide proper treatments of input functions and query coordinates for the solution functions by detaching the dependence of the input and output meshes. Our models pre-trained with benchmark datasets of operator learning are evaluated by downstream tasks to demonstrate the generalization abilities to vary- ing discretization formats of the system, which are natural characteristics of the continuous solution of the PDEs.
INTRODUCTION
Partial Differential equations (PDEs) are among the most successful mathematical tools for repre- senting the physical systems with governing equations over infinitesimal segments of the domain of interest, given some problem-speeifie boundary conditions or forcing functions (Mizohata 1973). The governing PDEs, which are globally shared in the entire domain, are interpreted as interactions between infinitesimal segments with respect to their geometrical structures and values. Because of the universality of the entire domain, the system can be analyzed in a continuous manner with respect to the system inputs and outputs. In general, identifying appropriate governing equations for unknown systems very challenging without domain expertise, however, numerous unknown processes remain for many complex systems. Even knowing the governing equation of the system known, requires unnecessary time and memory costs be solved using conventional numerical methods, and sometimes it is intractable to compute in a complex and large-scale system.
Motivation. In recent years, operator learning, an alternative to conventional numerical methods, has been gaining attention, pursuing mapping between dimenii-ai input/output function spaces in a data-driven manner without any problem-speciff knowledge of the system (Nelsen & Stuari 2021 Li et al., 2020a b, 20216] Lu et al., 2019; 2021] [Cao] 2021] Kovachki et al.) 20221). Intuitively, for the underlying PDE, Lau = f defined on the continuous bounded domains with system parameters E A. forcing function f E F, and the solution of the system u E U, the goal of the operator learning is to approximate the inverse operator = f A + U or : F + li with parametric model Go. Without loss of generality, when the input function is the output function can be computed as u = Ga(a). Because the operator Go should be able to capture interactions between elements of system inputs to discover the governing PDEs, Go is approximated by series of integral operators with parameterized kernels that iteratively update the system input to the output Kovachki et al. 2021). In practice,
(Nelsen & Stuart 2021 Li et al. 2020a 2021b:
Cao
2021

---

Under review as a conference paper at ICLR 2023
Table 1: Comparison with other studies,
Operator networks
Neural operators
MIOL
(Lu el al
2019 2021)
(Li et al.
2020a b 2021b
(ours)
Decoupling input/ouiput mesh Mesh-independence input
Yes
No
Yes
No, require fixed input mesh
Yes
Yes
Mesh-independennce to output
Yes
No, depend on input mesh
Yes
because continuous measurements of the input/output functions are infeasible, the observed data are provided as a set of input-output pairs, which are point-wise finite discretization of the functions.
The output values at coordinate can be expressed as u(y) = Go (v)](v) which can be viewed as the operating 9 : A Y -> l with two input placeholders, a A and y E Y. Then, the following two considerations can be considered for the model with respect to the input function a, and query coordinate y/: (1) the output the model should not depend on any discretization format of a, and (2) the model should be able to output solution at any query coordinate y. The measurements of the system are often sparsely and irregularly distributed owing to the geometry of the domain, environmental conditions, or inoperative equipment (Belbute-Peres & Günneeannn] 2022). In addition, in popular numerical methods for solving the PDEs, such as finite element methods, unstructured meshes are often utilized for the discretization of the domain, and adaptive remeshing schemes are commonly deployed where the regions of the domain require different resolutions depending on the accuracy of the prediction (Brenner al.] 2008 Huang & Russell 2010). Thus, the model should aggregate global information over the measurements to process (o(a) which can be reused for any discretization of a, and the solution [Go(a) should be queried at any coordinate y, such as implicit neural representations (Sitzmann et al. 2020).
al. 2020 Lienen
Existing architectures. There are two representative frameworks for operator learning: operator networks, and neural operators. Based on the framework of operator networks (Chen & Chen 1995). (Lu et al. 2019) extended the architectures with modern deep networks, called deep oper- ator networks (DeepONets), which is successful architectures for operator learning. To process the input functions and query coordinates, DeepONets consist of two sub-networks, called branch network and trunk network, respectively. DeepONets be queried at any coordinate from the trunk network, however, they used the fixed discretization of the system inputs a from the branch network (Lu al., 2019 2021). Another promising framework is neural operators, which consists of several integral operators with parameterized kernels to map between -iffntt-dimmesiinall func- tions. Neural operators can be adapted different resolutions of system inputs in a
manner (Li et a 2020a]b) However, the implemented architectures of the neural operators are typ- ically assumed to have the same sampling points for input and output functions (Li et al. 2020a bj 2021b, Kovachki et al., 2021 Lu et al. 2022). that is, they have not been formulated as a method for decoupling the discretization of the system input and solution, leading to the solution of the neu- ral operators not being flexibly queried at any coordinate y. In addition, widely used as successful model for operator learning due to their efficacy and accuracy (Li et al. 2021b) Pathak et al. 2022), Fourier neural operators (FNO) are limited to uniform grid discretization owing to the usage of FFT. The limitations and schematics of the existing studies are presented in Table and Figure
Contributions. There has been limited discussion the generalization abilities of discretizations with extended variations, such as different sampling points for inputs and outputs functions with arbitrary numbers and irregular distributions. For these considerations, we treated the observational data as a set without intrinsic assumptions about the data and constructed what we call a mesh- independent operator learner (MIOL), as shown in Figure [ MIOL is fully attentional architecture where the encoder and decoder are made in a way
consisting of an related set encoder-decoder frameworks (Zaheer et al., 2017 Lee et al. 2019) to detach the de- pendence of input and output meshes from the processor that processes the smaller fixed number of vectors in latent space. Attention mechanisms have been discussed as not only efficient in express- ing pair-wise interactions (Cao) 2021 Kovachki et al., [021] Guibas et al. |021 Pathak et al. (022) but also flexible in processing data in a modality-agnostte way (Vaswani et al. 2017: Jaegle et al., 1 2022). Finally, we conducted several experiments on the benchmark PDE datasets for oper- ator learning (Li et al. 2021b). The results show that our model is not only competitive in original benchmark tasks but also robustly applicable when discretizations of input and output functions are allowed to be different, unstructured, and irregularly distributed, which are natural consequences of continuous solutions for physical systems but are not flexibly compatible with existing ones.
2

---

Under review as a conference paper at ICLR 2023
2 PROBLEM DEFINITION
Let the observed data be provided as N instances of input-output pairs, {(as, ui un)) a; = a:lx and u; = Wg y are finite discretizations us E ( R" where A and U are separable Banach spaces. X {.1,..., Ln., } and Y fy,...., yn, J are the finite sets of discretized points for the continuous bounded domains Q2 and Q, with the number of discretized points nu and Ny, respectively. The training procedure of the original operator learning involves optimizing the following objective to learn a model Go A U.
the input function ag E A(22g; R2h) and output function
Eorn (Ga(a), u)] / - Go(a))) &
N
14 - Ga(a:)l?
(1)
1=1
where a ~ 1 is i.i on A, and the discretizations of the input and output space are usually assumed to be the same (Li et al. 2220aa by (221bb] Kovachki et al., 2021, Lu et al.) 2022). that is, X = Y. In this study, we consider the extended situations where the discretized points and are allowed to be different, that is, X # Y. and arbitrarily
the domain, while varying the number of ng and ny for each i. In practice, we aimed to build a model that is applicable to an irregularly distributed and varying number of discretization points during testing. Therefore, the goal is learn the model Go : A Y - N that is expected to make the following test errors as small as possible when training with the above objective
Eamy [4 ([G) (v), u(y)) Eamp /
Ja, n.
|u(ly) - [Gu (alx)
(2)
where X and Y can be arbitrary discretizations of the input and output domains. Described in sec- tion 5.2 to evaluate generalization ability respect varying the size of inputs, instead of calculating the test error 2 we calculate respective empirical test errors correspond to ratio p.
N
(v - (9 Σ
(P) (()))
(3)
where (() is randomly masked discretization with masking ratio p of the given input domain X, a; = as xe) is a, evaluated on ((),, and the model is evaluated at all given discretized points of Y (using all of Y is not necessary). The empirical test error can be bounded by the sum of the approximation error and discretization errors from the discrepancy of the discretizations of the input function,
(C ()(())))( |lu, - Go -
(((a))))(()))))))
(4)
The approximation error is expected to be sufficiently small by the training procedure T and if we consider the input functions
representation, that is, (P) C ag, there exists an appropriate masking ratio p. where the discretization errors can be expected to be sufficiently small by pro- cessing them with efficient permutation-invaarint set encoder whose outputs are independent of
cardinalities and locations of input discretizations (Zaheer et al. 2017 Wagstaff et al. 2019 2021
3 PRELIMINARIES
3.1 NEURAL OPERATORS
The architectures of neural operator (Li et al. 2020a), usually consist of lifting-iterative updates- projection, which are corresponding to encoder- encoder-processor processor-decode respectively. The lifting ei (z) = P(a(z)) and projection u(x) = Q(vz(2)) are local transformations usually implemented by point- wise feed-forward neural networks for mapping input features to target dimensional features. The iterative updates G: 12 + 0141, I [1, L are global transformations implemented by sequence of following transformations to capture the interactions between the elements,
1411 = (1)))) (x) = of (Wir(()) + (x)))))))) x E Ω
(5)
3

---

Under review as a conference paper at ICLR 2023
where a nonlinear functions, Wi are point-wiss linear transformations KG are kernel integral op- erations on uz( While the spatial domains of imput and output are usually the same, the discretiza- tions of imput and output functions can be different with varying numbers and locations, X Y. However, the existing implementations of the neural operators typically use the same discretizations
for the input and output (Li et al. 2020a b. [002bb] Kovachki et al., 2021] Lu et al., 2022), which leads to the predicted outputs only queried at the input meshes. Although the compatibiliyy and performance of neural operators are mesh-invariant
inputs (stably low test errors are observed at different resolutions from training (Li et al. 2020a b) 2021b)), the outputs cannot be flexibly queried to any points of coordinate y. For these reasons, separative treatments for input functions and query coordinates are required for decoupling the discretizations of inputs and outputs.
3.2 KERNEL INTEGRAL OPERATION AND ATTENTION
The kernel integral operations are generally implemented by integration of input values weighted by kernel values representing the pair-wise interactions between the elements input domain x E S2 and output domain E Sly.
(v) - / (x(y,x))(())) 92
(x,3) )) E U2 x Sly
(6)
where the parameterized kernels Hig are defined on Sl x Sgg. The transform K. can be interpreted as mapping a function (()) defined on domain x E f2g to the function [K(e))()) defined on do- main y E 22y Recently, it has been proved that the kernel integral operation can be successfully approximated by the attention mechanism of Transformers both theoretically and empirically (Cao
2021 Kovachki et al.) [021]] Guibas et al. 2021 Pathak et al.] (002).. Intuitively, let input vectors X My Xd, and query vectors Y E 1R y, then the attention can be expressed as
( =QQ xe [ (a(Y) -(()))(())()
Ja
where Q = YW? E Ra, K = XW* Rd, = XW RR. and a are the query, key, value matries, and softmax function, respectively. Details of the derivation can be found in Ap- pendix A.3 The attention mechanism, the weighted sum of V with the attention matrix a (QKT)
can be interpreted as the kernel integral operation in which the parameterized kernel is approximated by the attention matrix (Cao 2021; Tsai et al. 2019: Xiong et al. 2021] Choromanski et al. 202T). This attention is also known cross-attention. where the input vectors are projected to query em- bedding space by the attention, (Y, Note that when X = Y, the mechanism denotes the self-attention, Al,nXXX,XX)
The attention blocks Attention(Y, X, X) used in this paper are descrihed in the Appendix A.4 The term Attention(Y, X, X) used in the following paper stands for the attention blocks. The nonlinear functions a, point-wise linear transformations W,, and kernel integral operations K; in the iterative update (Equation 5 are approximated by feed-forward neural networks, residual con- nections, and the attention modules, which are conventional modules from Transformer-lis ar- chitectures (Vaswani et al. 2017). When we consider the attention block as mapping operation Attention (Y, X.) R + RR",, it can be used to not only approximate the kernel integral operation but also decouple the cardinality and sampled locations of input vector X and output Atlention (Y, X).
4 APPROACH
4.1 PREPROCESSING
Set representations for the discretizations. Before starting to construct a model, we reconsider the representations of a = =axx [8",xd, u =u/y Xd, and query coordinates Y, g=,xd To compensate the positional information, we follow the common way of cxisting neural operators lit- erature (Li et al. 2021b Kovachki et al., 2021), which usually concatenate the position coordinates and the corresponding values as the input representation, a = {(21, a(zz)), (In,, a(xn,) })} E Rn x (d+dd,, Then, we treat a, u, and Y, as the set representations represented by flattened arrays, without using structured bias (e.g., 2D, 3D-structured arrays, local neighbor connected graphs) to
4

---

Under review as a conference paper at ICLR 2023
Input function
Query locations
Output solutions
(Ra ((x))
w(x))
((x, ((x)) hi
[So(e)](X
Latent space
AnmtttinZ,, 2,30
AstentiolZ
In
2,
2, (aa())
Figure Mesh-independent operator learner.
avoid model bias toward specific data structure and flexibly process any of the discretization for- mats of input and output. By the concatenated position coordinates for each value, the output should be permutation-invarianl to the representation of the input function a, when the is represented as feature vectors, i.e.,
[Go (Ta)] [Go(a)](Y).
Meanwhile, since the is output of the model at the query coordinates Y = {y...... Bny E R only the values are used, = {u(v1), u(yng)} E Rd=, Since the output u(y) is function of query coordinates, the permutation operation and the function operations are commutative. There- fore, the output should be cetnuutaaioo-qquvvrrttee to the representation of the query coordinates Y. when the set is represented as vectors, i.c.,
[Go(a)[(FY) = =[Go(a)](Y).
Detailed explanations of the permutation-symmeett can be found in Appendix A.3
Positional embeddings. Instead of using raw position coordinates for both inputs and outputs Fz E Rd. we concatenate the Fourier embeddings for the position coordinates, which com- mon strategy to enrich the positional information (Vaswani et al., 2017 Mildenhall et al., [2020] Tancik al., 2020; Sitzmann et al., 2020p. The positional embeddings exploit sine and cosine functions with frequencies spanning from minimum to maximum frequencies sufficiently covering Nyquist rates for corresponding dimensions This simple technique provides the model with capability of representing fine-ge ine-grained functions or wide-spectral components.
4.2 ARCHITECTURE
Following (Lee et al. 2019 Jaegle et al. 2021 2022 Tang & Ha 2021), ),, we build our model with attention- architecture consisting of encoder encoder-processor processor-decoder with modality-agnostic en- coder and decoder, called mesh-inde mosh-independent operator learner (MIOL) as presented in Figure The significant modifications from the existing neural operators are mostly the lifting (encoder) and projection (decoder) for detaching the dependence of input and output meshes. The encoder en- codes the input function to fixed smaller number of latent feature vectors (discretization number: the processor processes the pair-wise and higher-order
na n2), satisfying interactions between elements the latent features vectors (discretization number: n2 + n2),, and the decoder decodes the latent features to output solutions at a set of query coordinates Y, (dis- The
cretization number: ng + ny) satisfying cermutation-equivariann property of attention blocks are described in Appendix A.3
Encoder. We use cross-attention block the encoder to encode inputs E ((++d)
smaller fixed number n, learnable queries Zo in,, Xd, (typically n, < nz), then the result of the block is
Gene = Allention(
E R=
(10)
which commuuttiio-invariann to the elements a and independent of the size of the input ng- These properties, which make the model mesh mesh-independent a, significant difference from the lifting component in the existing neural operators, which is not Eermutation-invariii and outputs the same size the inputs n, with point-wise transformation. The encoder is also interpreted as
5

---

Under review as a conference paper at ICLR 2023
projection inputs domain discretized by n, elements to latent domain consisting of the smaller number of n, elements, which is called "inducing points" in (Lee et al. 2019) and reduces the computational complexity of the following attentional blocks (Jaegle et al. 2021 2022
Processor. We use a series self-attention blocks as the processor each of which takes 21 E R xd, as the input of the query, key, and value components. Then the output of each self-attention block with E [1, L - 1] is
Z (Z))) = Attention Zz, Z, Zt) E R"=xdd,
(11)
which is crrmuuatioo-equiiariann to the elements of Z,, therefore the permutation-invarian property to the elements of a is preserved through successive modules. Also, the results Zi+1 have fixed discretization format with fixed ordering and number of elements n, which is decoupled from dis- cretization format the input function nt and output functions Ny. Due to the decoupling property, the whole architecture can not only capture the global interactions by the processor but is also appli- cable to mesh-independent operator learning independent of discretization formats of the input and output functions.
Decoder. We use a cross-altention block as the decoder to decode the latent vectors from the proces- sor L Xdn is
xdn at query coordinates Y E Rnyxd Then the final output of the entire architecture
[a(())) = Gdeo(Zt) = Attention(Y, Zn, Zτ) E ERdydd. Xd.,
(12)
which is permutation-cquivariaat to the elements of Yu. therefore every solution uy corresponds query coordinates Y/- This property makes the model applicable to any arbitrary discretization format of Y. Since the result from the processor is independent of the discretization format of input function a, the model is also applicable to any arbitrary discretization format of
EXPERIMENTS
5.1 EXPERIMENTAL SETTINGS
Datasets. Our experiments are conducted on several PDE benchmarks datasets following (Li al., 2021b), such as Burgers equations, Darcy flow, and Navier-Stokes equations, and 3D spherical shallow water dataset following (Yin et al. 2022), to investigate the flexibility and generalization abilities of MIOL through various downstream tasks. Details of the equations and the problems are described in Appendix A..
Baselines. We took several representative architectures baselines for PDE benchmarks problems, such as the original graph neural operator (GNO) (Li et al., (020a)), the multipole graph neural operator (MGNO) (Li et al. 2020b). neural operator based on low-rank decomposition (LNO) (Li et al. 2021b), the deep operator network (DeepONet) (Lu et al., 2019). and the Fourier neural operator (FNO) Li et al. 2021b), and baselines for 3D spherical shallow water problem, such as message passing PDE solver (MP-PDE) (Brandstetter et al. 2022). and dynamic-aware implicit neural representation (DINO) (Yin et al. 2022).
While we follow the original training procedures of (Li et al. 2021b), we take extended tasks with test sets varying discretization points for input and solution (or query coordinates y) to eval- uate the generalization abilities of architectures for input and output formats. Results for experi- ments already discussed on these baselines were obtained from the related literature, and results for the extended tasks that have not been discussed before have been reproduced from their original codes. Note that the 'n/a" result will occur when the baselines cannot be applicable for some down- stream tasks. The implementation details and additional experimental results can be found in the Appendix A.6 A.7 respectively.
5.2 RESULTS
Original benchmark tasks. Original benchmarks tasks on Burgers' equation and Darcy flow for different resolutions presented in Figure and the corresponding qualitative results presented in Table 66 and Table [ (Appendix [A.]). where all of the results except ours are brought from (Li et al. 2021b). The results show that our MIOL outperforms almost all neural operators at
conrespondinggquultttivee

---

Under review as a conference paper at ICLR 2023
Darcy Flow
00'
307
066
UNO
MONO
MONO
1983
0.04
7NO
Reletive aITOr
MIC
Relathe enorr
0.08
0.02
10-
0.01
253
2048
4096
0133
86
141
2-1
Resolution
Reselution
Figure 2: Benchmarks on (a) Burgers' equation, and (b) Darcy Flow for different resolutions.
every resolution in mesh-invariant way, except FNO in case of Darcy flow. This is because FNO explicitly exploits structural bias of the 2D structured array with equally spaced grids, while MIOL does not use the structural bias. For example, in the case of 2D Darcy flow, when the input values a(x1, x2) are given at 85 x 85 grid points (where the dimension of position coordinates d = 2 and the corresponding function values dæ 1), MIOL flattens to treat the batch-wise inputs as flatten arrays like a E R 7225x3 while FNO treats them as 2D structured- arrays like a E R 85x85x3
Tasks for evaluating mesh-independeenee To distinguish the specific generalization abilities of the existing architectures and to show that our proposed MIOL has all the generalization abilities and is robustly applicable even for tasks that the existing architectures unable to solve, we compare MIOL with two main baselines, DeepONet (Lu et al.] 2019) and FNO (Li et al. 2021b), by the following selective tasks presented in T
Supervised learning. This is the original setting of the existing operator learning models which evaluate the performance under the supervised learning settings. The results included in the original benchmark tasks, which shows that the MIOL is competitive with others and slightly worse than FNO in case of Darcy flow.
Task I:
Task 2: Mesh- Mesh-independence to input, but requiring the same grids for a a== u. This task, discussed in previous works of neural operators (Li et al. 2021b), evaluates the performance under zero-shot learning settings where the solutions are provided in unseen discretization during training. However, requiring the same grids for input and output functions, the solutions u of neural operators are only able to be queried at the same grids of DeepONets are not commonly applicable due to requirements of fixed input mesh, while FNO and MIOL. are robustly applicable and performances are almost
Task 3: feesh-independence to output, but requiring the same grids for train a = test a. This task, discussed in previous works operator networks (Lu et al. 2019), evaluates the perfor- mance under zero-shot learning settings where the solutions are provided in unseen discretization during training. However, operator networks are only applicable when the discretization formats (same numbers and locations) of the input function remain the same during training and testing. FNO are not commonly applicable due to the requirements of the same grids for u, whereas DeepONet and MIOL are robustly applicable, and their performances are almost mesh-invariant.
Randomly masked input, and queried at all points. This task, which has not been discussed before, evaluates the performance under a zero-shot learning setting, wherein the dis- cretizations of input and output functions are allowed have different numbers and be in random locations. The settings of the random mask for every input function and queried at all given points are sufficient evaluate whether the solution of the model is independent of the discretization of the input function and can be queried at any point. DeepONet and FNO are not applicable because of their restrictions, whereas MIOL. is robustly applicable.
Tass 4:
Predictions at any query point from the varying number and irregular distributed inputs. Additionally, Figure 3 visualizes the predictions of MIOL at all given query coordinates from the varying number and irregularly distributed inputs in Burgers' equation (left), and Darcy Flow (right). This visualization is reminiscent of few-shot regression problems (Finn et al. 2017 Kim et al. 2019) for u(y) at query locations y, where the support set is given by relatively few samples of discretized inputs a that should be mapped to the output space u through a series of integral transforms. As presented, the more shots are given, the closer the predictions are the ground truth.
7

---

Under review as a conference paper at ICLR 2023
Table 2: Relative I errors on Burgers' equation and Darcy flow under different settings.
Train grids
Test grids
Models
Task
DeepONet
FNO
MILL (ours)
Burgers' equation
1024
1024
0. 0.1582
0.0160 0.0139
0.0104
8192
8192
n/a
0.0090
1024
1024
8192
0.1584
n/a
0.0106
512 (50% mask)
1024
n/a
n/a
0.0479
Darey flow
85 x5x85
85x85
0.0765
0.0108
0,0172
421 x421 x5
x211
0.0098
0.0170
85
x421
0.0766 n/a
n/a 0.0173
3612 (50% mask)
x5x85
n/a
0.0272
Ulecred Inpues
Peedection
Input 5.0 %
Inpur 10.0 %
Input 50.0
Me
Impur 100 %
Irour to0
Figure 3: Predictions of MIOL at all query coordinates from the varying number and irregular distributed inputs on Burgers' equation (left), and Darcy Flow (right).
Applications to time-stepping systems. For time-dependent PDEs, we assumed that the sys- tems follow Markov process in which the state time dt is described as Wi+di Ga Wi, where the operator Ga during is not dependent on time (Lict al. 2021a). When we set the dt = 1, the state of the system time is denoted wt Gáoe [9.- and Gdee are the encoder and decoder implemented by the cross-attention blocks, respectively, and Get G
G1 Gone(wo) where Gone
G--1 are identical processors implemented by the same self-attention blocks. The goal is to learn the mapping from the initial state up to time T, G wolx + w|xx(0,11- The models were trained by minimizing the object Eunrya ETD L(g Gae Geno(wr). w
Figure visualizes that a initial states were encoded and processed iteratively with the same proces- sors in latent space. The latent features were decoded to output the predictions of the state fields the corresponding time. The quantitative results for Navier-Stokes equation and spherical shallow water presented in Table X and Table respectively.
Performances according to input masking ratios. Relative L errors according to input masking ratios (%) on the extended tasks presented in Figure 5 and Table The discretized inputs randomly masked with corresponding ratios. The errors are consistently low when the masking ratio lower than approximately 50%, where the performances according input masking ratios are found in Figure and Table 3 The efficient set encoder can aggregate sufficient information from the subsampled input a to make sufficiently small discretization errors from Equation
8

---

Under review as a conference paper at ICLR 2023
Ground Irot
Figure 4: Initial states were encoded, processed, and decoded to output the predictions of states for the Navier-Stokes equation (left) and 31D spherical shallow water (right) at the corresponding time. Since MIOL can process arbitrary discretization formats of input and output, it can be applicable even when the initial input states randomly masked with masking ratio 50% (left) or can be queried even at finer grids on the 3D spherical surface (right).
Basgers
Dany
Maxian Stikes
814
812
013
1.00
6.14
004
610
Marding rato
Maxting satoo 154
Marding tatoo
Figure 5: Relative 1.2 errors according to input masking ratios (%) on (a) Burgers' equation, (b) Darcy Flow, and (c) Navier-Stokes equations.
Table 3: Relative L errors according to input masking ratios (%)
Masking ratio
Burgers
0.0104
0.0154
0.0214
0.0399
0.0340
0.0479
0.0809
0.1081
0.1516
0.1649
Darey flow
0.0122
0,0185
0.0198
0.0221
0.0243
0.0272
0.0324
0.0424
0.0569
0,1019
Navicr-Stokes
0.0349
0.0437
0.0538
0.0514
0.0538
0.0644
0.1030
0.1378
0.1504
0.1776
CONCLUSIONS AND LIMITATIONS
Conclusions. In this study, we raise potential issues with existing operator learning models for PDEs from the perspective of discretization for the continuous input/output functions of the systems when the observations are irregular and have discrepancies between training and testing measurement formats. To solve these issues, we constructed fully-attentional architecture called MIOL, which treats discretized functions as set-valued data without prior data structures and structurally separates dependencies input and output meshes. MIOL is evaluated on the original tasks and extended downstream tasks and compared with other existing representative models. The results show that model is not only competitive in original operator learning tasks but also robustly applicable in extended tasks which are natural consequences of continuous solutions for physical systems but are compatible with existing representative models.
Limitations. First, although our main contribution is to construct an operator learning model to be flexible in measurement formats, the higher the dimension in which the measurement points can be placed, the performances our model tested on some original benchmarks are slightly worse than the problem-specific model (details in Appendix A.7 and the computational costs to obtain flexibility grow exponentially (in Table 5). It could be an important direction for future work to address the efficiency and scaling issues of discretization numbers of the input and output functions. Second, the experimental setting was far from real-world problem. Applications to realistic scenarios are an interesting direction for future research. For example, we might consider using model trained on the Navier-Stokes dataset to analyze the real fluid flow along complex geometrics regarded as the same phenomena with different measurement formats.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 06450, 2016.
Filipe De Avila Belbute-Peres, Thomas Economon, and Zico Kolter. Combining differentiable pde solvers and graph neural networks for fluid flow prediction. In international conference on ma- chine learning, pp. 2402-2411, PMLR, 2020.
Johannes Brandstetter, Daniel E. Worrall, and Max Welling. Message passing neural PDE solvers. In International Conference on Learning Representations, 2022.
Susanne C Brenner, L Ridgway Scott, and L Ridgway Scott. The mathematical theory of finite element methods, volume 3. Springer, 2008.
Shuhao Cao. Choose transformer: Fourier galerkin. Advances in Neural Information Processing Systems, 34, 2021.
Tianping Chen and Hong Chen. Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application dynamical systems. IEEE Transactions on Neural Networks, 6(4):911-917, 1995.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreca Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with per- formers. In International Conference on Learning Representations, 2021.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta- meta-learnng for fast adaptation of deep networks. In International conference on machine learning. pp- 1 1126-112 1135. PMLR, 2017.
John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catan- zaro. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arXiv:2111.13589 2021.
Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet-based operator learning for differen- tial equations. Advances in Neural Information Processing Systems, 34, 2021.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint
1rXiv:1606.08415, 2016.
Weizhang Huang and Robert D Russell. Adaptive moving mesh methods, volume 174. Springer Science & Business Media, 2010.
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. In International Conference Machine Learning. pp. 4651-4664. PMLR, 2021.
Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier Henaff, Matthew Botvinick, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver IO: general archi- tecture for structured inputs & outputs. In International Conference on Learning Representations, 2022.
Mateusz Jurewicz and Leon Derezynski. Set-to-sequence methods in machine learning: a review. Journal of Artificial Intelligence Research, 71:885 924, 2021.
George Em Karniadakis, Ioannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, 3(6):422 440, 2021.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami. Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. In International Conference Learning Representations, 2019.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 14/2.6980 2014.
10

---

Under review as a conference paper at ICLR 2023
Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint wwIXi:2218..044500 2021.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Tch. Set transformer: framework for attention-based permutation-invaraa neural networks. In Interna- PMLR, 2019.
tional Conference on Machine Learning, pp.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Neural operator: Graph kernel network for partial differ- ential equations. arXiv preprint arXi::2003.03488 2020a.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, Kaushik Bhat- tacharya, and Anima Anandkumar. Multipole graph neural operator for parametric partial differ- ential equations. Advances Neural Information Processing Systems, 33:6755-6766 6666, 2020b.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, An- drew Stuart, and Anima Anandkumar. Markov neural operators for learning chaotic systems. arXiv preprint arXiv:2106.06899, 2021a.
Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadeneshelii Burigede liu, Kaushik Bhat- tacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021b.
Marten Lienen and Stephan Günnemann. Learning the dynamics of physical systems from sparse observations with finite element networks. In International Conference on Learning Representa- tions, 2022.
Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for iden- tifying differential equations based on the universal approximation theorem of operators. arXiv preprint arXiv:1910.03/ 2019.
Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):218-222) 2021.
Lu Lu, Xuhui Meng, Shengze Cai, Zhiping Mao, Somdatta Goswami, Zhongqiang Zhang, and George Em Karniadakis. comprehensive and fair comparison of two neural operators (with practical extensions) based on fair data. Computer Methods in Applied Mechanics and Engineer- ing, 393: 14778, 2022.
Haggai Maron, Or Litany, Gal Chechik, and Ethan Fetaya. On learning sets of symmetric elements. In International Conference on Machine Learning, pp. 6734-6744. PMLR, 2020.
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. In European conference on computer vision, pp. 405-421. Springer, 2020.
Sigeru Mizohata. The theory of partial differential equations. CUP Archive, 1973.
Murphy, B Srinivasan, V Rao, and B Riberio. Janossy pooling: Learning deep permutation- invariant functions for variable-size inputs. In International Conference on Learning Representa- tions, 2019.
Nicholas H Nelsen and Andrew M Stuart. The random feature model for onput- maps between banach spaces. SIAM Journal on Scientific Computing. 43(5):A3212-A3244 2021.
Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al. Four- castnet: A global data-driven high-resolution weather model using adaptive fourier neural opera- tors. arXiv preprint arXiv:2202. 1214 2022.

---

Under review as a conference paper at ICLR 2023
Charles R Qi. Hao Su, Kaichun Mo, and Leonidas J Guibas. Pointnet: Deep learning point sets for 3d classification and segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 652-660, 2017a.
Charles Ruizhongtai Qi, Li Yi, Hao Su, and Leonidas J Guibas. Pointnet++: Deep hierarchical fea- ture learning on point sets in a metric space. Advances in neural information processing systems, 30, 2017b.
Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-7077 2019.
Siamak Ravanbakhsh, Jeff Schneider, and Barnabas Poczos. Deep learning with sets and point clouds. arXiv preprint 161 1.04500, 2016.
Justin Sirignano and Konstantinos Spiliopoulos. Dgm: A deep learning algorithm for solving partial differential equations. Journal of computational physics, 375:1339-1364, 2018.
Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Im- plicit neural representations with periodic activation functions. Advances in Neural Information Processing Systems, 33, 2020.
Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Keil, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan Barron, and Ren Ng. Fourier features let networks learn high frequency functions in low dimensional domains. Advances in Neural Information Process- ing Systems, 33, 2020.
Yujin Tang and David Ha. The sensory neuron as transformer: Permutation-invariaat neural net- works for reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
Yao-Hung Hubert Tsai, Shaojic Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: unified understanding of transformer's attention via the lens of kernel. arXiv preprint arXiv:/908 1775, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz. Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
Edward Wagstaff, Fabian Fuchs, Martin Engelcke, Ingmar Posner, and Michael A Osborne. On the limitations of representing functions on sets. In International Conference on Machine Learning, pp. 6487-6494. PMLR, 2019.
Edward Wagstaff, Fabian B Fuchs, Martin Engelcke, Michael A Osborne, and Ingmar Posner. Uni- versal approximation of functions on sets. arXiv preprint arXiv:2107.0/90 2021.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nystöm-based algorithm for approximating self-attention. In Proceedings of the... AAAI Conference on Artificial Intelligence. AAAI Conference on Artificial Intelligence, volume 35, pp. 14138. NIH Public Access, 2021.
Yan Yang, Angela F Gao, Jorge C Castellanos, Zachary E Ross, Kamyar Azizzadenesheli, and Robert W Clayton. Seismic wave propagation and inversion with neural operators. The Seismic Record, 1(3):126-13 2021.
Yuan Yin, Matthieu Kirchmeyer, Jean- Franceschi, Alain Rakotomamonjy, and Patrick Galli- nari. Continuous pde dynamics forecasting with implicit neural representations. arXiv preprint arXiv:2209. 14855, 2022.
Huaiqian You, Yue Yu, Marta D' Elia, Tian Gao, and Stewart Silling. Nonlocal kernel network (nkn): A stable and resolution-indeeenddee deep neural network. Journal of Computational Physics, 469; 111536, 2022. ISSN 0021-9991.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh. Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
12

---

Under review as a conference paper at ICLR 2023
A APPENDIX
A.1 ARCHITECTURE SCHEMATICS
Neure operator
Input function
Set decoder
Query points
Figure 6: Schematics of operator learning architectures, such as operator networks (left), neural operators (middle), and our proposed MIOL (right). For each respective formulation, the gray circles are given and the red circles are the corresponding outputs.
A.2 RELATED WORKS
Neural networks for PDEs. There has been increasing interest in utilizing neural networks for PDEs that can be divided into several lines. In (Raissi al. 2019 Sirignano & Spiliopoulos 2018 Karniadakis et al. 2021), neural networks were used as approximations of the solution when given boundary conditions and collocations points constrained to known governing equations of the system, called physics-informed neural networks. It can be an alternative to traditional PDE solvers, yielding inseeiindeennnnnt and relatively high-fidelity solutions. However, these methods require full knowledge of the PDEs, and the trained model is usually not reusable for new systems, otherwise requires expensive re-training for every new boundary condition. Meanwhile, several studies have been conducted learn reusable governing operators for generalization to systems. A universal operator approximation theorem (Chen & Chen 1995) was originally implemented by single-layer neural networks, and (Lu et al., 2019] 2021) develop the theorem by extending the architectures with current deep networks, called DeepONets, leading them to be more expressive. Neural opera- tors have been originally implemented by message passing on graphs (Li et al.) 2020a|b). However, the graph-based neural operator was limited due to its stability issues and the quadratically growing costs for computing the integral operator You al. 2022 Li al., 2021b). Fourier neural opera- tors (FNO), which utilize fast Fourier transform (FFT) to eficiently compute the integral operators in Fourier space, have been widely used as successful models for operator learning
range of science and engineering for their efficacy and accuracy (Li et al. 2021b Yang al. 202 Pathak al., 2022]. Furthermore, (Gupta et al. (02T) uses multi-waveles transform to approximate the kernel of the operators. Close approximate the integral operation, but they use problem-specific feature extractors and decoding regressors which make the model biased toward specific equations.
works, (Cao) 2021) use Transformer-sty architectures to
Set representation learning. In recent years, it has been witnessed that neural networks on sets gained attention starting from learning simple operations unordered sets (e.g., sum max) and processing the point clouds (Zaheer et al., 2017 Ravanbakhsh et al., 2016 Qi et al. 2017a bj. The elements of the are transformed into new representations through point-wise feed-forward neural Then, the representation of the set is obtained by -invaiiiiinnn functions, such sum- pooling (Zaheer et al. (017)) or max- max-pooling (Qi et al., (2017a]b) b). However, a theoretical study has been investigated claim the limitations of representation power of the methods (Wagstaff al., 2019 Jurewicz & Derczynski] (021). Additionally, the methods make it difficult for the model to learn pair-wise and higher-order interactions between the elements of the sets, which are lost during those pooling operations. More recently, several groups have suggested more sophisticated models for treating sel-vel inputs through modifying Transformers (Vaswani et al. 2017). Set Transformer uses to project sal-vel input into smaller latent vectors, which ensure handling arbitrary cardinality and permutation of the inputs (Lee et al. 2019). Perceiver similar strategy by emphasizing
networks which are sermutation-equivariaan
aggregation across the entire elements of the set by
attention module considered as learnall parameterized pooling operations
13

---

Under review as a conference paper at ICLR 2023
the efficiency of computational cost and the property of modality-agnoss which can be applied to a wide range of large-scale imput modalities, such as images, point clouds, audio, and video (Jaegle) et al. 202I). Furthermore, beyond simple outputs like classification, Perceiver I0 uses another cross-attention module to decode complex outputs modalities which can be applied to the domains of language, optical flow, audio-visua sequence, and game environments (Jaegle et al. (022). A similar strategy is also applied to reinforcement learning (Tang & Ha 2021
A.3 PERMUTATION-SYMMETNT AND ATTENTION
Detail explanation of equation Here is brief explanation of the approximation with the in- tegral for the cross-attention mechanism, where the softmax for the attention matrix is ignored for simplicity.
A
(1(x))
ki (Inz)
v(x1)
Attn(Y,X,X)- =
AAttnlyn X, x)
(((nn))
kd, (aine)
(((ome))
[((1)) k((1)
a(11) k(())))
(((1))
E E (g(er) M(x)))u(as))
(()n.)) k(x1)
a(vmn)) k(In=))
u(Ina))
E (alum,) -((x)))o(())
Sa, (alyn) (())))(x)ddr
/ (()))=(=))((2)da Ja
Sa (g(w)n) k(z))o(z)dr
(13)
Here, the discretization of input is {#1, Ina (as key and value vectors), and can be changed to the discretization of output (yı,... Uny (as query vectors) with cardinality changed from n₂ to By. Using this mechanism, we can detach the dependences of discretization formats of input and output from the processor, by encoding arbitrary discretization {-i, Jnn, } to fixed size (nz) of learnable latent set vectors, and decoding the latent set vectors to output arbitrary discretization (yi,..., Vn, }. The discretization number varied as n& (arbitrary) n, (fixed) + ny (arbitrary).
Permutation invariance and equivariance. For function whose input can be represented by set, there can be two interesting symmetries with respect any permutations; permutation-invarian and ceamnntation-quivriinttn
Let F is the function with set-valued inputs X = {21, F2, .... 2n, } I; E Rdd. Since deep learning models can not directly treat the valued input, the input data provided as ordered vectors representing one of n! permutations of X. In practice, with arbitrary permutation action и to the first dimension of the matrix-represennte X [aadd the permutation- invariance and equivariance are defined as,
F(xX) F(X), F(xX) =F()).
(pernutationn-ivvaraaeee)
(14)
eermutation-equiiariiaa
where the output of the permutation-invant ity of the input set, while the ordering and the cardinality of the output of the permutation-cqqivart function remain consistent with those of the input set. A full mathematical definition and details of
function does not depend on ordering and the cardinal-
the properties can be found in (aheer et al. 2017 Murphy al. 2019 Wagstaff al.) 2021 Maron et al. 2020).
property of attention. Here, attempt to explain that the attention module Attn(Y, XY, in Equation 7 is cermutation-invarte to X and permutation-equivarii to Y. Let the permutation matrix the element (first) dimension, P Ruxn, First, when the
14

---

Under review as a conference paper at ICLR 2023
permutation matrix P is multiplied the X, the output of the attention is Attm(Y, PX, PX) =QQK") = (OYW*(PXWW' PXW"
=o (oWW pT) PXW" (Pt (xw*)(wW)) (XW+) 7) xw
(15)
a <o(Y
=o (rw (xW+) XW AY,,n(Y X),
which states that the Xt, X) is permutation-invarian to X. matrix P is multiplied the Y, the output of the attention is
Also, when the permutation
Alnn(PY, X, X) =QQK"VV
=o (PPWW((WW)) XW
(16)
Po (XWA) 7) XW =PAttn(Y,X,X))
which states that the Alln(Y, X,, X) is permutation-eqquivriaatt to Y. Thus, it is easily proved that self-attention Alln(X, X, X) permutation-equiiariaatt X,
Attn(PX, PX) = PX, PXX
= PAttn(X,X,X)
(17)
A.4 ATTENTION BLOCKS
Following (Lee al. 2019 Jaegle al. 2021 2022), mesh-independent operator learner (MIOL) consists of two types of attention blocks, cross- and self-attention blocks, which implement the respective attention mechanisms. The attention blocks have the following shared structure, which takes two input arrays, a query input Y E Rd, and key-value input X E xdz
0 Y Latn(LayroNorm/ff
Lorm Layer Attention XY X) = O + FF (aayerroomm(()))
(18)
where Lager orm is layer normalization (Ba et al. 2016), FF consists of two point-wise feed- forward neural networks with GELU nonlinearity Henddrycks & Gimpel 2016). and the exact calculation of attention Attn
QKT Vdg
Attn(X x X") softmax
V,
(19)
where Q= X E Rayxd, K X*W* Xd,, and V X"W" Ru=xd. for single headed attention. In the case of multi-headed attention, several outputs from different learnable parameters are concatenated and projected with the linear transformation. Here. the output of cross- attention block Attention(Y. ( X)) is also commuuation-invaiian to X and permutation-cqquiaaiin Y, and the output of self-attention block Attention (X, X, X) is also permutation-equi -equivaain to X, because the other modules operate in point- wise way, the properties of symmetry for the attention modules are preserved for the corresponding attention blocks.
A.5 BENCHMARKS
Burgers' equation. First, consider benchmark problem of ID Burgers' equation which is a non-linear parabolic PDE combining the terms of convection and diffusion. The equation with periodic boundary conditions is
d.u(x, t) + Oz Ou t)/2) = vdza u(x,t),
x E (0, t e (0.1) (0,1)
u(z, 0) uo (x))
15

---

Under review as a conference paper at ICLR 2023
where u4 μ is the initial state generated from (0, 625(-2 + 257) -3) and 0.1 is the viscosity coefficient. The goal of operator learning is to learn mapping the initial state to the solution at time one, uo + 1).
Darcy flow. Second, we consider another benchmark problem of 2D steady-state Darey flow which second-order elliptic PDE describing the flow of fluid through porous medium. The equation of Darey flow on the unit box is
(a(x) f(x),
I (0, 1)a
u(x) = 0, x 0 1)²
where u is density of the fluid, a ~ is the diffusion field generated from - (O (-, with fixed forcing function 1. The goal of operator learning is to learn mapping the diffusion field to the solution of the density, 9 : a-++
Navier-Stokes equation. Third, consider another benchmark problem of 2D Navier-Stokes equation describing the dynamics of viscous, incompressible fluid. The equation in vorticity form on the unit torus is
a,w(s, 1) + u(x,()) ()- Vw(x, 1) vAw(a, () f(x),
E (0,1) (0,
V ((xuxx)) O, w(z, 0) wo(x),
(0,1) E [O, (0.1)²
where u is the velocity field, V vxuin is the vorticity field, wo is the initial vorticity field gen- erated from (O, (- + 495)-2.5) with periodic boundary conditions, v is the viscosity coefficient and forcing function is kept fxx 0.1 (sin (2i (21 + x2)) + cos (2x (x1 22))).
A.6 IMPLEMENTATION DETAILS
Table 4: Implementation details for training.
System
Burgers' equation
Darey flow
Navier Navier-Stokes equation
Encoder
Input function values, a E 1 Xda
1024x1
7225x1
4096
Positional embedding
1024x 129
7225x 130 32
4096x50
Frequency bins, K
64
12
Number of heads
4
1
Inputs. g(d+dd))
1024x 130
7225x131
4096x51
Processor
Learnable queries, Zu E R Xd,
256x64
256x64
128x64
Latent channels, da. Number of heads
64
64
64
8
8
8
Number of blocks, I.
4
Decoder
Query coordinates, Y, E Rd Positional embedding Outputs, Xdy
1024x1
7225x2
4096x2
1024x 129
7225xx300
4096x50
1024x1
7225x1
4096x1
Burgers' equation. The positional embeddings at position coordinate E [0,1] 1] from both in put/output domain consists of [Ir, sin (fKTT) cos (JıTT)) where Sy. are from equally spaced 64 frequencies from min 1 to max resolution 64 are used, resulting to have 129 sions. For the encoder, a 4-headed cross-attention module is used, and the number of the elements (first) and channel (second) dimension latent space are 256 and 64, respectively. For the proces- sor, 8-headed self-attention module with the same channel dimension of latent's is used. For the decoder, a 4-headed cross-attention module with the channel dimension of 64 is used.
x 64 dimen-
Darey flow. The positional embeddings at position coordinate (I1, field (z)) 22) consists of [21, 22, sin (fLTT)) ) sin (f cos (fATT)) ), cos (fa where fx resolution 32 used, resulting have
[0, 1]2 of diffusion
from equally spaced 32 frequencies from min
16

---

Under review as a conference paper at ICLR 2023
130 = 2 x (1 + 2 x 32) dimensions. For the encoder, a single-headed cross-attention module is used, and the number of the elements (first) and channel (second) dimension of latent space are 256 and 64, respectively. For the processor, four 8-headed self-attention modules with the same channel dimension of latent's are used. For the decoder, a single-headed cross-attention module with the channel dimension of 64 is used.
Navier-Stokes equation. The positional embeddings at position coordinate (2), E [0, 1]2 of the vorticity field (2) consists of [ sin (faTT)) sin (fir)2)), cos (fATT)) ) cos (fuTTr))] where Sn are from equally spaced 32 frequencies from min to max resolution 32 are used, re- sulting to have 50 x (1+2) 12) dimensions. For the encoder, a single-headed -ross-attention module is used, and the number of the elements (first) and channel (second) dimension of latent space are 128 and 64, respectively. For the processor, two 8-headed self-attention modules with the same channel dimension of latent's 's are used. For the decoder, a single-headed cross-attention module with the channel dimension of 64 is used.
Spherical shallow water. The position coordinate (2,, U2, Ia) are transformed from spherical coor- dinates of latitude and longitude o, 0 E //2/ [-x, T.. Therefore, the positional embeddings can be [21, T2 Fy sin (fxTTz ), sin (J&TT12) sin (J cos (JaTTa)), cos (JuTTz)) cos (JuTTa))) where f& are from equally spaced 24 frequencies from min to max resolution 32 are used, re- sulting to have 147 = 3 (1+ x 24) dimensions. For the encoder, single-heade gegdeed cross-attention module is used, and the number of the elements (first) and channel (second) dimension of latent space are 256 and I 128, respectively. For the processor,
8-headed self-attention modules with the same channel dimension of latent's are used. For the decoder, a single-headed cross-attention module with the channel dimension of 128 is used.
Optimization. For all systems, the experiments are conducted on a GTX TITAN X GPU and we use Adam optimizer (Kingma & Bal 2014) train MIOL for 500 epochs with initial learning rate of 0.001 which is halved every 100 epochs. Batch sizes for Burgers' equation and Navier-Stokes equations are set to 20, and for Darcy flow is set 10.
Computational complexity comparisons. In Table 5 we compare MIOL with the baselines in perspective of the training time, memory usage, and number of parameters on the benchmarks. The training time is measured during one epoch of each benchmarks training set, the memory usage in- dicates the GPU memory usage for the active Python process recorded from nvidia-smi command. The time and memory cost of MIOL is on par with both baselines in the case of Burgers' equation but increases drastically compared with both baselines in the case of Darcy flow. By flattening the discretized inputs into ID sequence (identical point set) before processing the inputs and query architec-
coordinates, which removes any intrinsic locality of data, MIOL is the modality ture that does not exploit the structural bias of data (such as a 2D structured array). Therefore, for instance, MIOL regards the inputs of Darcy flow as 7,225 sizes of ID sequence instead of 85 x85 sizes of 2D structured data, which leads to exponential growth computational cost compared to en- coding ID data such as inputs of Burgers' equations. Exponentially more resources are required for MIOL to have the capability to handle the more highly structured data in modality-agnostte way for flexible input formats and query coordinates. Mcanwhile, while attention- based architecture consists of QKV modules where the parameters are shared across the input elements, kernel-based architecture needs a large number of parameters when the dimension of data increases due to the dimension of the kernel increasing along the dimension of data. Therefore, although the number of parameters of FNO is smaller than MIOL in the case of Burgers' equations, it becomes much larger than those of MIOL. In summary, MIOL has fewer parameters but becomes more expressive when paying large amount of computational cost for gaining the flexibility to data formats, while FNO has huge number of parameters, which already include some prior data structure, requiring less training cost but could be overfitted to the data structure.
A.7 ADDITIONAL RESULTS
Quantitative results for 1D Burgers' equation, 2D Darcy flow, Navier-Stokes equation, and 3D spherical shallow water.
Comparisons with other set-encoder. The comparison model consists of mncoder-processee decoder like MIOL, where encoder is replaced by Deepsets (Zaheer et al., 2017) variant with average-pooling operations, while processor and decoder are the same MIOL, called Deepsets-
17

---

Under review as a conference paper at ICLR 2023
Table 5: Computational complexity comparisons with baselines.
Model
Training time per epoch
Memory usage
Number of parameters
Burgers' equation
DeepONet FNO MIOL
0,645
1.34 GB
107,405
0.871
1.45 GB
57.825
0.715
1.63 GB
125,204
Darcy flow
DeepONet FNO MIOL
0.732
1.40 GB
417,550
1.881
1.96 GB
1,188.353
8.843
5.68 GB
401,514
Table 6: Benchmarks ID Burgers' equation
Models
5 256
5 512
8
1024
8 = 2048
s 4096
8 8192
FCN
0.0958
0.1407
0.1877
0.2313
0.2855
0.3238
PCANN
0.0398
0.0395
0.0391
0.0383
0.0392
0.0393
GNO
0.0555
0.0594
0.0651
0.0663
0.0666
0.0699
LNO
0.0212
0.0221
0.0217
0.0219
0.0200
0.0189
MGNO
0.0243
0.0355
0.0374
0.0360
0.0364
0.0364
FNO 0.0149
0.0158
0.0160
0.0146
0.0142
0.0139
MIOL (ours)
0.0105
0.0109
0.0104
0.0092
0.0090
0.0099
Table Benchmarks on 2D Darcy Flow
Models
85
141
211
= 421
FCN
0.0253
0.0493
0.0727
0.1097
PCANN
0.0299
0.0298
0.0298
0,0299
GNO
0.0346
0.0332
0.0342
0.0369
LNO
0.0520
0.0461
0.0445
MGNO
0.0416
0.0428
0.0428
0.0420
FNO 0.0108
0,0109
0.0109
0,0098
MIOL (ours)
0.0172
0.0173
0.0177
0.0182
Table 8: Relative I,2 errors on Navier-Stokes equation under different settings.
Train grids
Test grids
Models
a, 1
a
W
FNO
MIOL. (ours)
64x64
same with a
0.0110
0.0349
64x64
2048 (50% mask)
64x64
n/a
0.0644
Table 9: L super-resolution test errors spherical shallow water
Model
0 < T < 10
10
< 20
I-MP-PDE (Brandstetter et al. DINO ( Yı et al. 2022
2022
1.908e-3
7.240e-3
1.063e-4
6.466e-4
MIOL (ours)
4.524e-4
1.122e-3
18

---

Under review as a conference paper at ICLR 2023
variant. The performance of Deepsets-variante on Burgers equation is presented in Table 10, which shows not competitive compared to MIOL. The outputs of the attention layer can be calculated by the weighted sum of attention and inputs, where the attention values have more expressive repre- sentation with extra learnable parameters Zo and inputs a, while the outputs of the existing pooling operation calculated by fixed coefficient weighted sum (or operation) of inputs (in the case of avg-pooling, the weights are ) Therefore, from the perspective of the modeling operator, much of the information of interactions between pair-wise input elements is lost during those pooling operations.
Table 10: Comparisons with other set-encoder.
Train grids
Test grids
Models
0,
Deepsets-variant
MIOL (ours)
1024
a
0.0179
0.0104
8192
8192
0,0183
0.0090
1024
8192
0,0189
0.0106
1024
512 (50% mask),
1024
0.0748
0.0479
19


Reviews:
[{"review_id": "c6huZOXJwsr", "reviewer": "Reviewer_kSWf", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "Operator learning models for PDEs learn the parameter-to-solution mapping of families of partial differential equations. Since they map functions into functions, their approximation should be independent of the discretization of the respective function spaces. This is not always the case for the operator learning networks proposed until now: in practice they often evaluate the functions on finite discretizations and are not able to evaluate the inputs and outputs on different meshes, or to treat training and testing sets with different discretization meshes. The paper raises these issues and develops MIOL- a mesh-independent attentional-architecture that separates dependencies on input and output meshes. It compares its performance with the existing models on several tasks.", "summary_of_the_review": "I think the paper raises an important issue in the context of operator learning models. However, the paper needs some polishing of the form of the exposition, a more in depth discussion regarding discretization and error bounds, and some Figure polishing.  \n", "strength_and_weaknesses": "I think the paper has his strength in raising an important question regarding discretization-flexibility w.r.t input/output functions  for operator learning models. The paper proposes an architecture that combines properties of existing operators (such as neural operators and set transformers) to address this issue and shows in experiments its performance. However, I think there are some weaknesses in the presentation of the paper:\n\n–  Some more background on the used architectured would be appreciated\n\n– The paper lacks a discussion about the error bounds w.r.t the discretization chosen\n\n– The authors should go again through the text since it is at times redundant  and imprecise in the linking words and logical phrases\n\n– In the experiments, the figures need to be polished. In particular Figure 3 (fontsize of colorbar ticks too small, x and y axes missing....)\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "M3-s1pUxUNZ", "reviewer": "Reviewer_p5hq", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes an attention-based model called the mesh-independent operator learner (MIOL) to provide proper treatments of input functions and query coordinates for the solution functions by detaching the dependence on input and output meshes. The proposed models pre-trained with benchmark datasets of operator learning are evaluated by downstream tasks to demonstrate the system's generalization abilities to varying discretization formats. The paper shows the advantages of the method with some examples.", "summary_of_the_review": "This paper is innovative, but the presentation of results is somewhat naive. The examples shown are relatively simple. The comparison with other methods does not show the superiority of the proposed method.", "strength_and_weaknesses": "Strength\n\nThis paper proposes a method that treats discretized functions as set-valued data without prior data structures and structurally separates dependencies on input and output meshes. \n\nCompared with other representative models, the proposed method is evaluated on the original and extended downstream tasks. And the results show that this model is competitive in original operator learning tasks and robustly applicable in extended studies.\n\nWeaknesses\n\nSince this method does not require mesh, why not test some equation problems on surfaces, such as spheres?\n\nIs it possible to solve 3D problems?\n\nThe error gap in figure 2 is vast, as the comparison methods are too simple. Can you show some more reasonable comparisons? For example, show some similar mesh-free methods.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4_OLoVKjvHQ", "reviewer": "Reviewer_YQEY", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "A paper on operator learning that does not require a specific mesh structure, but allows for querying points. Presentation a bit on the heavy side and experiments could be more demonstrative.\n", "summary_of_the_review": "This is an interesting paper that considers a topical problem. The presentation is on the heavy side and the experiments only concentrate on straightforward benchmark problems.\n", "strength_and_weaknesses": "*Strengths*\n\n1. The problem is topical and has received recent attention also beyond purely theoretical work. The topic area itself should be of interest for the audience of the conference.\n\n2. The proposed approach appears sensible and described in good detail.\n\n3. The paper is rather detailed and covers the required technical details even for someone who is familiar with the problem domain but is not an expert in the area.\n\n*Weaknesses*\n\n4. Even if the paper reads rather well (with only minor typos), the presentation is on the heavy side. The first 4.5 pages are describing the problem/background/preliminaries, and the paper actually starts at the bottom of page 5. The problem itself is rather simple, and could be presented in a more lightweight/approachable way, which would probably widen the impact of the work.\n\n5. The experiments focus on rather standard benchmark problems and showing that the proposed method actually solves the problem at hand. This is valuable and provides an 'empirical' sanity check, but as a reader I would have expected this to be a first set of experiments with actual demonstrations and/or real-world problem show case data in the end. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject