M ESH -I NDEPENDENT
O PERATOR
L EARNING
FOR
PDE S USING S ET R EPRESENTATIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 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 dependence of 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.
I NTRODUCTION
Partial Differential equations (PDEs) are one of the most successful mathematical tools for representing the physical systems with governing equations over infinitesimal segments of the domain of interest given some problem-specific boundary conditions or forcing functions
(Mizohata, 1973).
The governing PDEs, which are globally shared on the entire domain, are interpreted as how to interact between infinitesimal segments with respect to their geometrical structures and values.
Because of the universality on the entire domain, the system can be analyzed in a continuous way with respect to system inputs and outputs.
In general, identifying appropriate governing equations for unknown systems is very challenging without domain expertise, however, there still remain numerous unknown processes for many complex systems.
Even if knowing the governing equation of the system, it requires unnecessarily lots of time and memory costs to solve with 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 the conventional numerical methods, has been getting attention, pursuing to learn a mapping between infinite-dimensional input/output function spaces in a purely data-driven way without any problem-specific knowledge of the system (Nelsen & Stuart, 2021;
Li et al., 2020a;b; 2021b; Lu et al., 2019; 2021; Gupta et al., 2021; Cao, 2021;
Kovachki et al., 2021).
Intuitively, for underlying PDE, L a u = f defined on the continuous bounded domain Ω with system parameters a ∈ A , forcing function f ∈ F , and solution of the system u ∈ U , the goal of the operator learning is to approximate the inverse operator G =
L − 1 a f :
A → U or G : F → U with parametric model G θ .
Without loss of generality, in the case of input function being a , the output function can be computed by u =
G θ ( a ) .
Since the operator G θ should be able to capture interactions between elements of system inputs a to discover the governing PDEs, G θ has been approximated by a series of integral operators with parameterized kernels which iteratively update the system input to output (Nelsen & Stuart, 2021; Li et al., 2020a;b; 2021b; Gupta et al., 2021; Cao, 2021;
Kovachki et al., 2021).
In practice, since continuous measurements of the [Table] 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 y can be expressed as u ( y ) =
[ G θ ( a )]( y ) which can be viewed as the operating G : A × Y → U with two input placeholders, a ∈ A and y ∈ Y .
Then, the following two considerations could be considered for the model with respect to input function a , and query coordinate y ; (1) output of the model should not depend on any discretization format of a , and (2) model should be able to output a solution at any query coordinate y .
The measurements of the system are often sparsely and irregularly distributed due to the geometries of the domain, environmental conditions, or occasionally inoperative equipment (Belbute-Peres et al., 2020; Lienen & G¨unnemann, 2022).
Also, 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 re-meshing schemes are commonly deployed where the regions of domain need different resolutions depending on the accuracy of the prediction (Brenner et al., 2008; Huang & Russell, 2010).
Thus, the model should aggregate global information over the measurements to process [ G θ ( a )] which can be reusable to any discretization of a , and the solution [ G θ ( a )]( y ) should be able to be queried at any coordinates y , like implicit neural representations (Sitzmann et al., 2020).
Limitations of existing architectures.
There are two representative existing frameworks for operator learning, such as operator networks, and neural operators.
Based on a framework of operator networks (Chen & Chen, 1995), (Lu et al., 2019) extended the architectures with modern deep networks, called deep operator networks (DeepONets), which is one of the 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 can be queried at any coordinate y from the trunk network, but they use the fixed discretization of system inputs a from the branch network (Lu et al., 2019; 2021).
Another promising framework is neural operators, which consists of several integral operators with parameterized kernels to map between infinitedimensional functions.
Neural operators can be adapted to different resolutions of the system inputs a in a mesh-invariant way (Li et al., 2020a;b).
However, implemented architectures of the neural operators typically assume to have the same sampling points for input and output function (Li et al., 2020a;b; 2021b; Kovachki et al., 2021; Lu et al., 2022), i.e., have not been formulated as a method for decoupling the discretization of system input and solution, leading to the solution of the neural operators not being flexibly queried at any coordinate y .
Also, widely used as successful model for operator learning due to their efficacy and accuracy (Li et al., 2021b; Pathak et al., 2022; Yang et al., 2021), Fourier neural operators (FNO) are limited to uniform grid discretization due to the usage of
FFT.
The limitations and schematics of the existing studies are presented in Table 1 and Figure 5.
Contributions.
There has been limited discussion on the generalization abilities for discretizations with extended variations such as different sampling points for inputs and outputs functions with arbitrary numbers and irregular distributions.
For these considerations, we treat the observational data as a set without intrinsic assumptions on the data and construct what we call mesh-independent operator learner (MIOL) in Figure 1.
MIOL is a fully attentional architecture consisting of encoderprocessor-decoder, where the encoder and decoder are made in a way related to set encoder-decoder frameworks (Zaheer et al., 2017; Lee et al., 2019; Maron et al., 2020) to detach the dependence of input and output meshes from the processor which processes the smaller fixed number of vectors in latent space. The attention mechanisms have been discussed as not only efficient to express the pair-wise interactions (Cao, 2021; Kovachki et al., 2021; Guibas et al., 2021; Pathak et al., 2022) but also flexible to process data in a modality-agnostic way (Vaswani et al., 2017; Jaegle et al., 2021; 2022).
Finally, we conduct several experiments on benchmark PDE datasets for operator 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 not flexibly compatible with existing ones.  R ELATED W ORKS 2. N EURAL NETWORKS FOR PDE S There has been increasing interest in utilizing neural networks for PDEs that can be divided into several lines.
In (Raissi et al., 2019; Sirignano & Spiliopoulos, 2018; Karniadakis et al., 2021), neural networks were used as approximations of the solution when given the 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 mesh-independent 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 to learn reusable governing operators for generalization to new 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 operators 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 et al., 2022; Li et al., 2021b).
Fourier neural operators (FNO), which utilize fast Fourier transform (FFT) to efficiently compute the integral operators in Fourier space, have been widely used as successful models for operator learning in a range of science and engineering for their efficacy and accuracy (Li et al., 2021b;
Yang et al., 2021; Pathak et al., 2022).
Furthermore, (Gupta et al., 2021) uses multi-wavelet transform to approximate the kernel of the operators.
Close to our works, (Cao, 2021) use Transformer-style architectures to approximate the integral operation, but they use problem-specific feature extractors and decoding regressors which make the model biased toward specific equations.
2. S ET REPRESENTATION LEARNING
In recent years, it has been witnessed that neural networks on sets gained attention starting from learning simple operations on unordered sets (e.g., sum or max) and processing the point clouds (Zaheer et al., 2017; Ravanbakhsh et al., 2016; Qi et al., 2017a;b).
The elements of the set are transformed into new representations through point-wise feed-forward neural networks which are permutation-equivariance.
Then, the representation of the set is obtained by aggregation across the entire elements of the set by permutation-invariant functions, such as sum-pooling (Zaheer et al., 2017) or max-pooling Qi et al.
(2017a;b). However, a theoretical study has been investigated to claim the limitations of the representation power of the methods (Wagstaff et al., 2019; Jurewicz & Derczynski, 2021). 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 set-valued inputs through modifying Transformers (Vaswani et al., 2017).
Set Transformer uses a cross-attention module considered as learnable parameterized pooling operations to project setvalued input into smaller latent vectors, which ensure handling arbitrary cardinality and permutation of the inputs (Lee et al., 2019).
Perceiver uses a similar strategy by emphasizing the efficiency of computational cost and the property of modality-agnostic which can be applied to a wide range of large-scale input modalities, such as images, point clouds, audio, and video (Jaegle et al., 2021).
Furthermore, beyond simple outputs like classification, Perceiver IO uses another cross-attention module to decode complex outputs modalities which can be applied to the domains of language, optical flow, audio-visual sequence, and game environments (Jaegle et al., 2022).
A similar strategy is also applied to reinforcement learning (Tang & Ha, 2021).
P ROBLEM SETTING
Let the observed data be provided as N instance of input-output pairs, { ( a i , u i ) }
Ni =1 , where a i = a i | X a and u i =
u i | Y u are finite discretization of the domain for the input function a i ∈ A at X a = { x 1 , ..., x n x } and output function u
i ∈ U at Y u
= { y 1 , ..., y n y } , with the number of discretized points n x and n y , respectively.
The training procedure of the original operator learning is to optimize the following objective to learn a model G θ : A → U ,
[Equation]
where a ∼ µ is iid on A , and the discretizations of the input and output space are usually assumed to be the same (Li et al., 2020a;b; 2021b; Kovachki et al., 2021;
Lu et al., 2022), i.e., X a = Y u .
In this paper, we consider the extended situations where the discretized points X a
and Y u are allowed to be different i.e., X a ̸ = Y u , and arbitrary on the domain along with varying the number of n x and n y for each i . In practice, we aim to build a model that is applicable to the irregularly distributed and varying number of discretization points during testing. Therefore, the goal is to learn the model G θ : A× Y → U that is expected to make the following test errors as small as possible when training with the above objective 1,
[Equation]
where X a and Y u can be arbitrary on the domain.
Note that when the discretization X a is fixed, the problem can be considered as a problem such as Lu et al.
(2019; 2021), where the model cannot handle other discretization formats of the system inputs a i .
P RELIMINARIES 4. N EURAL OPERATORS
The architectures of neural operator (Li et al., 2020a), usually consist of lifting-iterative updatesprojection, which are corresponding to encoder-processor-decoder, respectively.
The lifting v 1 ( x ) = P ( a ( x )) and projection u ( x ) = Q ( v L ( x )) are local transformations usually implemented by pointwise feed-forward neural networks for mapping input features to target dimensional features. The iterative updates G l : v l (cid:55)→ v l +1 , l ∈ [1 , L − 1] are global transformations implemented by sequence of following transformations to capture the interactions between the elements, [Equation]
where σ are nonlinear functions, W l are point-wise linear transformations, K l are kernel integral operations on v l ( x ) . While the spatial domains of input and output are usually the same the discretizations of input and output functions can be different with varying numbers and locations, X a ̸ = Y u . However, the existing implementations of the neural operators typically use the same discretizations for the input and output (Li et al., 2020a;b; 2021b; Kovachki et al., 2021; Lu et al., 2022), which leads to the predicted outputs only queried at the input meshes. Although the compatibility and performance of neural operators are mesh-invariant to inputs (at the test stage, stable performances 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 a and query coordinates y are required for decoupling the discrtizations of inputs and outputs.
The kernel integral operations are generally implemented by integration of input values weighted by kernel values κ representing the pair-wise interactions between the elements on input domain x ∈ Ω x and output domain y ∈ [Equation] y ,
[Equation]
where the parameterized kernels κ ϕ are defined on Ω y × Ω x .
The transform K can be interpreted as mapping a function v ( x ) defined on domain x ∈ Ω x to the function [ K ( v )]( y ) defined on domain y ∈ Ω y .
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., 2021; Guibas et al., 2021; Pathak et al., 2022).
Intuitively, let input vectors X ∈ R n x × d x and query vectors Y ∈ R n y × d y , then the attention can be expressed as
[Equation]
where Q = Y W q ∈ R n y ×
d q , K = XW k ∈ R n x × d q , V = XW v ∈ R n x × d v , and σ are the query, key, value matries, and softmax function, respectively.
The attention mechanism, the weighted sum of V with the attention matrix σ ( QK T ) , 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., 2021).
This attention is also known as cross-attention, where the input vectors are projected to query embedding space by the attention, Attn ( Y, X, X ) .
Note that when X =
Y , the mechanism denotes the self-attention, Attn ( X, X, X ) .
The attention blocks Attention ( Y, X, X ) used in this paper are described in the Appendix A.3.
The term Attention ( Y, X, X ) used in the following paper stands for the attention blocks.
The nonlinear functions σ , point-wise linear transformations W l , and kernel integral operations K l in the iterative update (Equation 3) are approximated by feed-forward neural networks, residual connections, and the attention modules, which are conventional modules from Transformer-like architectures (Vaswani et al., 2017).
When we consider the attention block as mapping operation
Attention ( Y, X, X ) : R n x (cid:55)→ R n y , 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
Attention ( Y, X, X ) .
A PPROACH 5. P REPROCESSING
Set representations for the discretizations.
Before starting to construct a model, we reconsider the representations of a = a | X a ∈ R n x × d a , u = u | Y u ∈ R n y × d u and query coordinates Y u ∈ R n y × d .
To compensate the positional information, we follow the common way of existing neural operators literature (Li et al., 2021b; Kovachki et al., 2021), which usually concatenate the position coordinates and the corresponding values as the input representation, a = { ( x 1 , a ( x 1 )) , ..., ( x n x , a ( x n x )) } ∈
R n x × ( d + d a ) .
Then, we treat a , u , and Y u as the set representations represented by flattened arrays, without using structured bias (e.g., 2D, 3D-structured arrays, or local neighbor connected graphs) to avoid our model bias toward specific data structure and flexibly process any of the discretization formats of input and output.
By the concatenated position coordinates for each value, the output should be permutation-invariant to the representation of the input function a , when the set is represented as feature vectors, i.e.,
[Equation]
Meanwhile, since the u is output of the model at the query coordinates Y u
= { y 1 , ..., y n y } ∈ R n y × d , only the values are used, u = { u ( y 1 ) , ..., u ( y n y ) } ∈ R n y × d u .
Since the output u ( y ) is a function of query coordinates, the permutation operation and the function operations are commutative.
Therefore, the output should be permutation-equivariant to the representation of the query coordinates Y u , when the set is represented as vectors, i.e.,
[Equation]
[Equation]
Detailed explanations of the permutation-symmetry can be found in Appendix A.2.
Positional embeddings.
Instead of using raw position coordinates for both inputs and outputs x i , y j ∈ R d , we concatenate the Fourier embeddings for the position coordinates, which is a com- mon strategy to enrich the positional information Vaswani et al.
; Mildenhall et al.
Tancik et al.
; Sitzmann et al.
The positional embeddings exploit sine and cosine functions with frequencies spanning from minimum to maximum frequencies sufficiently covering the Nyquist rates for corresponding dimensions This simple technique provides the model with the capability of representing fine-grained functions or wide-spectral components.
5. A RCHITECTURE
Following (Lee et al., 2019; Jaegle et al., 2021; 2022; Tang & Ha, 2021), we build our model with a attention-based architecture consisting of encoder-processor-decoder with modality-agnostic encoder and decoder, called mesh-independent operator learner (MIOL) as presented in Figure 1.
The significant modifications from the existing neural operators are mostly in the lifting (encoder) and projection (decoder) for detaching the dependence of input and output meshes.
The encoder encodes the input function a to fixed smaller number of latent feature vectors (discretization number: n x (cid:55)→ n z ), satisfying permutation-invariance, the processor processes the pair-wise and higher-order interactions between elements of the latent features vectors (discretization number: n z (cid:55)→ n z ), and the decoder decodes the latent features to output solutions at a set of query coordinates Y u (discretization number: n z (cid:55)→ n y ) satisfying permutation-equivariance.
The permutation-symmetric property of attention blocks are described in Appendix A.2.
Encoder.
We use a cross-attention block as the encoder to encode inputs a ∈ R n x × ( d + d a ) to a smaller fixed number n z of learnable queries Z 0 ∈ R n z × d z (typically n z < n x ), then the result of the block is
[Equation]
(8) which is permutation-invariant to the elements of a and independent of the size of the input n x .
These properties, which make the model mesh-invariant to a , are a significant difference from the lifting component in the existing neural operators, which is not permutation-invariant and outputs the same size as the inputs n x with point-wise transformation.
The encoder is also interpreted as a projection inputs domain discretized by n x elements to a latent domain consisting of the smaller number of n z 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 of self-attention blocks as the processor each of which takes z l ∈ R n z × d h as the input of the query, key, and value components.
Then the output of each self-attention block with l ∈ [1 , L − 1] is
[Equation]
which is permutation-equivariant to the elements of Z l , therefore the permutation-invariant property to the elements of a is preserved through successive modules.
Also, the results Z l +1 have fixed discretization format with fixed ordering and number of elements n z which is decoupled from discretization format of the input function n x and output functions n y .
Due to the decoupling property, the whole architecture can not only capture the global interactions by the processor but is also applicable to mesh-independent operator learning independent of discretization formats of the input and output functions.
Decoder.
We use a cross-attention block as the decoder to decode the latent vectors from the processor Z L ∈ R n z × d h at query coordinates Y u ∈ R n y × d .
Then the final output of the entire architecture is
[Equation]
[Equation]
which is permutation-equivariant to the elements of Y u , therefore every solution u j corresponds to query coordinates y j .
This property makes the model applicable to any arbitrary discretization format of Y u .
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 a .
E XPERIMENTS 6. E XPERIMENTAL SETTINGS
Datasets.
Our experiments are conducted on several PDE benchmarks datasets following (Li et al., 2021b), such as Burgers equations, Darcy flow, and Navier-Stokes equations, 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.4.
Baselines.
We took several representative architectures for operator learning problems as baselines, such as the original graph neural operator (GNO) (Li et al., 2020a), the multipole graph neural operator (MGNO)
(Li et al., 2020b), a 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).
While we follow the original training procedures of Li et al. (2021b), we take extended tasks with test sets varying discretization points for input a and solution u (or query coordinates y ) to evaluate the generalization abilities of architectures for input and output formats.
Results for experiments 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 downstream tasks.
The implementation details and additional experimental results can be found in the Appendix A.5, A.6, respectively.
6. R ESULTS
Original benchmark tasks.
Original benchmarks tasks on Burgers’ equation and Darcy flow for different resolutions are presented in Figure 2, and the corresponding qualitative results are presented in Table 5, and Table 6 (Appendix A.6), 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 every resolution in a 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 ( x 1 , x 2 ) are given at 85 × 85 grid points (where the dimension of position coordinates d = 2 and the corresponding function values d a = 1 ), MIOL flattens to treat the batch-wise inputs as flatten arrays like a ∈
R b × 7225 × 3 , while FNO treats them as 2D structured-arrays like a ∈ R b × 85 × 85 × 3 .
Tasks for evaluating mesh-independence. 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 are 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 Table 5.
- Task 1: Supervised learning. This is the original setting of the existing operator learning models which evaluate the performance under the supervised learning settings. The results are 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 2: Mesh-independence to input, but requiring the same grids for 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 a .
DeepONets are not commonly applicable due to requirements of fixed input mesh while FNO and MIOL are robustly applicable and performances
[Table]
are almost mesh-invariant.
For instance, in case of Darcy flow, the performances of MIOL are [Table] at grids of 85 × 85 , and 0.0170 at grids of 421 × 421 . - Task 3: Mesh-independence to output, but requiring the same grids for train a = test a .
This task, discussed in previous works of operator networks (Lu et al., 2019), evaluates the performance 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 requirements of the same girds for a = u , while DeepONet and MIOL are robustly applicable and performances are almost mesh-invariant. For instance, in case of Darcy flow, the performances of MIOL are 0.0172 at grids of 85 × 85 , and 0.0173 at grids of 421 × 421 . - Task 4: Randomly masked input, and queried at all points.
This task, which has not been discussed before, evaluates the performance under zero-shot learning setting where the discretizations of input and output functions are allowed to have different numbers and be in random locations.
The settings of random mask for every input function and queried at all given points are enough to 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 due to their restrictions, while MIOL are robustly applicable and the errors are consistently low when masking ratio is lower than about 50%, where the performances according to input masking ratios are found in Figureand Table 7 (Appendix A.6).
Predictions at any query point from the varying number and irregular distributed inputs.
Additionally, Figure 3 visualizes predictions of MIOL at all given query coordinates from the varying number and irregular distributed inputs on 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 which should be mapped to output space u through a series of integral transforms.
As presented, the more shots are given, the closer the predictions are to the ground truth.
Applications to time-stepping systems.
We assume that the Navier-Stokes equation is under the
Markov process that the vorticity field at time t + dt is described as w t + dt = G dt · w t , where the operator G dt during dt is not dependent on time Li et al.
(2021a).
When we set the dt = 1 , the vorticity field at time t is denoted as w t =
G dec · [ G t − 1 · · · · · G 1 ] · G enc ( w 0 ) where G enc and G dec are the encoder and decoder implemented by cross-attention blocks, respectively, and G dt = G 1 = · · · = G t − 1 are identical processors implemented by the same self-attention blocks.
The goal is to learn mapping the initial vorticity up to time T , G : w 0 | (0 , 1) 2 (cid:55)→ w | (0 , 1) 2 × (0 ,T ] .
The models are trained with minimizing the object E w 0 ∼ µ (cid:104) 1 T (cid:80) T − 1 t =
L ( G dec · G dt · G enc ( w t ) , w t +1 ) (cid:105) .
Figure 4 visualizes that discretized initial vorticity field with the reduced size is encoded and processed iteratively with the same processors in latent space.
The latent features are decoded to output the predictions of vorticity field at the corresponding time.
As shown in Table 9 (Appendix A.6), the quantitative results for Navier-Stokes equation have similar trends to the case for Darcy flows.
C ONCLUSIONS AND L IMITATIONS
Conclusions.
In this work, we raise potential issues with existing operator learning models for
PDEs in 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 the issues, we construct a fully-attentional architecture called MIOL, which treats discretized functions as set-valued data without prior data structures and structurally separates dependencies on input and output meshes.
MIOL is evaluated on the original tasks and extended downstream tasks compared with other existing representative models.
The results show that our 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 not compatible with existing representative models.
Limitations.
First, although our main contribution is constructing an operator learning model to be flexible on measurement formats, the higher the dimension in which the measurement points can be placed, the performances of our model tested on some original benchmarks get slightly worse than the problem-specific model (details in Appendix A.6), and the computational costs to obtain the flexibility grow exponentially (in Table 4).
It could be an important direction for future work to address the efficiency and scaling issues on discretization numbers of input and output functions.
Second, the experimental setting is far from real-world problems.
Applications to realistic scenarios are an interesting direction for future work.
For example, we might consider using a model trained on the dataset of Navier-Stokes to analyze the real fluid flow along complex geometries regarded as the same phenomena with different measurement formats.
A. P ERMUTATION - SYMMETRY AND ATTENTION
Permutation invariance and equivariance.
For a function whose input can be represented by a set, there can be two interesting symmetries with respect to any permutations; permutation-invariance, and permutation-equivariance.
Let F is the function with set-valued inputs
X = { x 1 , x 2 , ..., x n } , x i ∈ R d .
Since deep learning models can not directly treat the set-valued input, the input data is provided as ordered vectors representing one of n !
permutations of X .
In practice, with arbitrary permutation action π to the first dimension of the matrix-represented X ∈ R n × d , the permutationinvariance and equivariance are defined as,
[Equation]
where the output of the permutation-invariant function does not depend on ordering and the cardinality of the input set, while the ordering and the cardinality of the output of the permutation-equivariant function remain consistent with those of the input set.
A full mathematical definition and details of the properties can be found in (Zaheer et al., 2017;
Murphy et al., 2019;
Wagstaff et al., 2021; Maron et al., 2020).
Permutation-symmetric property of attention.
Here, we attempt to explain that the attention module Attn ( Y, X, X ) in Equation 5 is permutation-invariant to X and permutation-equivariant to Y .
Let the permutation matrix to the element (first) dimension, P ∈ R n × n .
First, when the permutation matrix P is multiplied to the X , the output of the attention is
[Equation]
which states that the Attn ( Y, X, X ) is permutation-invariant to X .
Also, when the permutation matrix P is multiplied to the Y , the output of the attention is
[Equation]
which states that the Attn ( Y, X, X ) is permutation-equivariant to Y .
Thus, it is easily proved that self-attention Attn ( X, X, X ) is permutation-equivariant to X ,
[Equation]
A. A TTENTION BLOCKS
Following Lee et al.
; Jaegle et 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 ∈ R n y × d y and a key-value input X ∈ R n x × d x ,
[Equation]
LayerNorm
[Equation]
[Equation]
where LayerNorm is layer normalization Ba et al.
(2016), FF consists of two point-wise feedforward neural networks with a GELU nonlinearity Hendrycks & Gimpel (2016), and the exact calculation of attention Attn is
[Equation]
(16) where Q = X q W q ∈ R n y ×
d q , K = X k W k ∈ R n x × d q , and V =
X
v W v ∈ R n x × d v for a 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 crossattention block Attention ( Y, X, X ) is also permutation-invariant to X and permutation-equivariant to Y , and the output of self-attention block
Attention ( X, X, X ) is also permutation-equivariant to X , because the other modules operate in a point-wise way, the properties of symmetry for the attention modules are preserved for the corresponding attention blocks.
A. B ENCHMARKS
Burgers’ equation. First, we consider a benchmark problem of 1D Burgers’ equation which is a non-linear parabolic PDE combining the terms of convection and diffusion. The equation with periodic boundary conditions is [Equation]
where u 0 ∼ µ is the initial state generated from µ = N (0 , 625( − ∆ + 25 I ) − 2 ) 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, G : u 0 (cid:55)→ u ( · , 1) .
Darcy flow.
Second, we consider another benchmark problem of 2D steady-state Darcy flow which is a second-order elliptic PDE describing the flow of fluid through a porous medium.
The equation of Darcy flow on the unit box is
[Equation]
(0 , 1)where u is density of the fluid, a ∼ µ is the diffusion field generated from µ = N (0 , ( − ∆ + 9 I ) − 2 ) with fixed forcing function f = 1 .
The goal of operator learning is to learn mapping the diffusion field to the solution of the density, G : a (cid:55)→ u .
Navier-Stokes equation.
Third, we consider another benchmark problem of 2D Navier-Stokes equation describing the dynamics of a viscous, incompressible fluid.
The equation in vorticity form on the unit torus is
[Equation]
where u is the velocity field, w = ∇× u is the vorticity field, w 0 ∼ µ is the initial vorticity field generated from µ = N (0 , 7 3 / 2 ( − ∆ + 49 I ) − 2 . 5 ) with periodic boundary conditions, ν is the viscosity coefficient and forcing function is kept f ( x ) = 0 .
1 (sin (2 π ( x 1 + x 2 ))
+ cos (2 π ( x 1 + x 2 )))
[Table]
A. I MPLEMENTATION DETAILS
[Table]
Burgers’ equation.
The positional embeddings at position coordinate x ∈ [0 , 1] from both input/output domain consists of [ x, sin ( f k πx ) , cos ( f k πx )] , where f k are from equally spacedfrequencies from min 1 to max resolution 64 are used, resulting to have 129 = 1 + 2 × 64 dimensions.
For the encoder, a 4-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, a 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.
Darcy flow.
The positional embeddings at position coordinate ( x 1 , x 2 ) ∈ [0 , 1] 2 of diffusion field a ( x 1 , x 2 ) consists of [ x 1 , x 2 , sin( f k πx 1 ) , sin ( f k πx 2 ) , cos ( f k πx 1 ) , cos ( f k πx 2 )] , where f k are from equally spaced 32 frequencies from min 1 to max resolution 32 are used, resulting to have 130 = 2 × (1+2 × 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 areand 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 ( x 1 , x 2 ) ∈ [0 , 1] 2 of the vorticity field w ( x 1 , x 2 ) consists of [ x 1 , x 2 , sin( f k πx 1 ) , sin ( f k πx 2 ) , cos ( f k πx 1 ) , cos ( f k πx 2 )] , where f k are from equally spaced 32 frequencies from min 1 to max resolution 32 are used, resulting to have 50 = 2 × (1 + 2 × 12) 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 128 and 64, respectively.
For the processor, two 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.
Optimization.
For all systems, the experiments are conducted on a GTX TITAN X GPU and we use Adam optimizer Kingma & Ba (2014) to 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 to 10.
Computational complexity comparisons.
In Table 4, 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 indicates 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 a 1D sequence (identical point set) before processing the inputs and query coordinates, which removes any intrinsic locality of data, MIOL is the modality-agnostic architecture 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 1D sequence instead of 85 ×sizes of 2D structured data, which leads to exponential growth computational cost compared to encoding 1D 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 a modality-agnostic way for flexible input formats and query coordinates.
Meanwhile, 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 a large amount of computational cost for gaining the flexibility to data formats, while FNO has a huge number of parameters, which already include some prior data structure, requiring less training cost but could be overfitted to the data structure.
[Table]
A. A DDITIONAL R ESULTS
[Table]
Performances according to input masking ratios.
Relative L 2 errors according to input masking ratios (%) on the extended tasks are presented in Figure 6 and Table 7.
The discretized inputs are randomly masked with corresponding ratios.
To the best of our knowledge, MIOL is the only [Table] model which is robustly applicable to this extended tasks, but not compatible with existing operator learning models
[Table]
[Table]
Comparisons with other set-encoder.
The comparison model consists of encoder-processordecoder like MIOL, where encoder is replaced by Deepsets Zaheer et al.
(2017) variant with averagepooling operations, while processor and decoder are the same as MIOL, called Deepsets-variant.
The performance of Deepsets-variants on Burgers equation is presented in Table 8, 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 representation with extra learnable parameters Z 0 and inputs a , while the outputs of the existing pooling operation are calculated by fixed coefficient weighted sum (or operation) of inputs a (in the case of avg-pooling, the weights are 1 n x ).
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]
[Table]