Pretraining Codomain Attention Neural Operators for Solving Multiphysics PDEs
Authors: Md Ashiqur Rahman
Robert Joseph George
Mogab Elleithy
Daniel Leibovici
Zongyi Li
Boris Bonev
Colin White
Julius Berner
Raymond A. Yeh
Jean Kossaifi
Kamyar Azizzadenesheli
Anima Anandkumar
Venue: NeurIPS 2024 poster (2024)
Keywords: Neural Operator
Primary area: machine_learning_for_physical_sciences

Abstract: Existing neural operator architectures face challenges when solving multiphysics problems with coupled partial differential equations (PDEs) due to complex geometries, interactions between physical variables, and the limited amounts of high-resolution training data. 
To address these issues, we propose *Codomain Attention Neural Operator* (CoDA-NO), which tokenizes functions along the codomain or channel space, enabling self-supervised learning or pretraining of multiple PDE systems. 
Specifically, we extend positional encoding, self-attention, and normalization layers to function spaces. CoDA-NO can learn representations of different PDE systems with a single model. We evaluate CoDA-NO's potential as a backbone for learning multiphysics PDEs over multiple systems by considering few-shot learning settings. On complex downstream tasks with limited data, such as fluid flow simulations, fluid-structure interactions, and Rayleigh-Bénard convection, we found CoDA-NO to outperform existing methods by over 36%.

Pretraining Codomain Attention Neural Operators for Solving Multiphysics PDEs
Md Ashiqur Rahman', Robert Joseph George', Mogab Elleithy", Daniel Leibovici', Zongyi Li", Boris Bonev", Colin White", Julius Berner", Raymond A. Yeh Jean Kossaifi", Kamyar Azizzadenesheli, Anima Anandkumar" Purdue University, CCaltech, 2NVIDIA
Abstract
Existing neural operator architectures face challenges when solving multiphysics problems with coupled partial differential equations (PDEs) due to complex ge- ometries, interactions between physical variables, and the limited amounts of high-resolution training data. To address these issues, we propose Codomain Atten- tion Neural Operator (CoDA-NO), which tokenizes functions along the codomain or channel space, enabling self-supervised learning or pretraining of multiple PDE systems. Specifically, we extend positional encoding, self-attention, and normal- ization layers to function spaces. CoDA-NO can learn representations of different PDE systems with single model. We evaluate CoDA-NO's potential backbone for learning multiphysics PDEs over multiple systems by considering few-shot learning settings. On complex downstream tasks with limited data, such fluid flow simulations, fluid-structure interactions, and Rayleigh-Bénard Bénard convection, we found CoDA-NO outperform existing methods by over 36%.
Introduction
Many science and engineering challenges involve solving partial differential equations (PDEs). A PDE can represent physical phenomena such as fluid dynamics, wave propagation, material deformation, etc., but to describe many real-world systems, multiple such PDEs must be coupled together, viz., multi-physics modeling [1]. For instance, in subsurface engineering, equations of flow, thermodynamics, and microchemistry are coupled together [2]: in materials science, physics at multiple scales are involved in modeling [3], and in weather forecasting, atmospheric processes involve interactions of wave propagation and fluid dynamics [4].
Traditionally, numerical methods have been devised to solve PDEs. However, they typically require discretization of PDEs fine grids to capture the physical phenomena accurately. Consequently, these computational requirements often exceed available memory and computational budgets for real- world applications. Beyond these obstacles present in individual PDE problems, the convergence of numerical solvers in multiphysics systems presents major difficulties arising from intricate interactions among multiple coupled PDEs.
Deep learning techniques have emerged as faster alternatives to numerical solvers for PDEs in many applications. They typically trained using supervised learning with data obtained from solvers. This becomes a challenge when only limited data is available, especially in the case of multiphysics simulations, which are expensive and challenging for numerical solvers. Instead, obtaining data from simpler simulations where only a subset of the "physics" is incorporated is more convenient and less expensive. In other words, instead of getting data from coupled PDE systems, solving individual PDEs. While the solutions of the two systems can be very different, they share common features and can benefit from combined learning framework. Can we design systematic curriculum learning scheme for learning multiphysies systems?
obtain data by
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Single Physics System
Coupled Pulti-Physics System
Ugp
Uz
Mg.3+50
y/
Wyy
4y,1+50
p
Pt
CoDA-ND
Pt+at
Mäsk
dr,
dy (-0
CoDA
dy
dy 0-
Self- Superviedd Pre-training
Few Shot Supervised Finetuning
Figure 1. CoDA-NO adapts seamlessly to new multi-physies systems. Pre-trained on fluid dynamics data (Navier-Stokes equation with Ug My and p) using the maskeddreeonstrrctionn objective, CoDA-NO easily adapts to multi-physics fluid-solid interaction systems (new and dy variables) without any architectural changes.
More generally, a foundation model trained on different kinds of PDEs can learn representations across multiple domains and then transfer them to new problems. Such foundation models have found immense success in computer vision and natural language processing [5, 6]. Foundation models are first trained in a self-supervised manner on large and often unlabeled datasets. Then they can be efficiently adapted or fine-tuned to a broad range of downstream tasks with minimal to no additional data or training.
Recent works have attempted to train foundation model for solving PDEs [7-9]. However, these methods only work on predetermined PDEs with fixed number of variables, and none of them consider multi-physics PDEs, and they are mostly restricted to uniform grids, limiting their applica- bility, For example, standard patching- baaeed approaches used Vision Transformers (ViTs) [10] often struggle with discontinuities in predicted functions and changing resolutions [11]. Since they are limited fixed uniform grids, they cannot generalize to resolutions different from the training resolutions.
To handle varying resolutions and grids, neural operators [12, 13] have been introduced deep learning framework for learning mappings between function spaces. Neural operators are guaranteed to converge to a unique operator in the limit of increasingly fine discretizations (of the computa- tional domain). This property is known as discretization convergence, making them agnostic to the discretization the input and output functions and suitable for approximating solution operators of PDEs. Neural operators can replace numerical solvers while being significantly faster in several scenarios [14, 15]. While some of the previous PDE foundation models [7, use neural operators, they still cannot handle multiphysics coupled PDEs. They also cannot adapt to new variables that are not predetermined at the beginning of training.
Our Approach:
We propose a novel transformer neural operator architecture with codomain attention (CoDA-NO) layers designed to handle varying combinations of physical phenomena modeled through coupled PDEs. We partition the input function codomain- wise into set of token functions, each corresponding to distinct physical variables of the PDE. The CoDA-NO model processes this set of functions as input, extending the transformer architecture from a dinee-dimensional vector space to an difinte-dimensionall function space. This extension is achieved by carefully redesigning positional encodings, the self-attention mechanism, and normalization techniques.
In our architecture, each token is treated as function, capturing cross- coss-function dynamics through attention mechanisms while maintaining discretization convergence. This design empowers the ar- chitecture handle functions discretized on grids of varying resolutions. Specifically, each token function is subjected to the following operations: (i) concatenation with a learned positional embed- ding, (ii) lifting to higher-dimensional co-domain, and (iii) functional attention mechanisms to compute interactions. We use Fourier neural operators (FNOs) [16] rather than traditional multi-layer perceptrons (MLPs) to create the representations for keys, values, and queries, which helps maintain the functional nature of the input data. Details can be found in Sec. 3 and Alg. 1.
CoDA-NO can be applied to varying numbers of input functions (on different geometries) and adapt to novel PDEs with fewer or additional interacting variables, as illustrated in Fig. 1. This allows us to learn multiple PDE systems in one model.
To demonstrate CoDA-NO's generalizability across diverse physical systems, we examine two settings: multiphysics problems and collection of single-physie problems.

---

CoDA NO Block' L.
GNO
K
(or ant Aa
Query
WP po 1
per
Variable
Variable
VSPE
Betthax
Variable
GNQ
Value
Input Function
Latent Function
Functions
Token
Functions
Output
Target Function
Illustration CoDA-NO
Each input physical variables, "variable 'variable depicted two ditlierent
colors (yellow and blue), incorporate leurnable variahle specific pusitional encoder (VSPE). These variables, along with the correqunndng VSPE, are passed through GNO layers to transform from nca-umiform to latent uniform grids. Codomain attention tekenizes the latent functions along the cudomain. Each token undergoes transformations with K. Q, and uperatoos yielding key. query, and value functions [A k {a .9" ), and (w' ). The resalting function is computed using a self self-attention mechanism in function space followed by integral operator Zper- Finally, the output function the target geometry generated by passing through stacked CoDA-NO blocks, followed by an additional GNO layer.
Pre-tras Pre-tralniigg Physical Quantities of Navier-Stoker Stokes Equations
Fine-tuning additional Physical Quantities of Elastic Wave Equation
U.r
VSPE
VSPE2
Ur,t+5t
Masking
U
24 y.
Uy
VSPE
U
VSPE
Pi+ōt
U
Pt
dz,t-5t
VSPE
Encoder
Reconstructor
VSPE
Encader
Predictor
d y
d
Mankin
VSPE5
(b) sepervupe pre-training fine-tuning CoDA-NO. Mudel, pre-trained and in a self supervised way, can he fine tuned to fluid salid interaction dataset (new da and dy variables) by only including two extra VSPEs prodictor module.
equations dataset (wilh uas My-
Figure (a) CoDA-NO architecture. (b) Self-supervised pre-training and fine-tuning process with CoDA-NO.
For the multiphysics scenario, we examine two distinct systems. First, we consider a fluid-structure interaction problem [17] governed by the incompressible Navier-Stokes equation and the Elastic wave equation. The fluid-structure interaction problem is representative of the multi-physics behavior of various real-world problems, e-g., climate and atmosphere modeling. It also provides additional challenge of irregular meshes on a complex geometry.
Instead of directly learning solve the full multiphysics problem, start with curriculum where we first learn the basic fluid dynamics without the clastic wave equation, governed by the incompressible Navier-Stokes equation, with velocity and pressure variables. We pre-train CoDA- NO in a self-supervised manner on snapshots of fluid flows by masking different parts the velocity pressure fields. Using few-shot supervised fine-tuning, viscosities and additional displacement fields given by the elastic wave equation. We use graph neural operator (GNO) layers [18] encoders and decoders handle time-varying irregular meshes of the fluid-structure interaction problems. For the few-shot learning problem, our model achieves 36.8% lower errors average compared to the best-performing baseline trained from scratch on the target problem.
show that model can adapt unseen
The second system involves Rayleigh-Bénard Bénard convection, where the Navier-Stoke and heat (energy) equations are coupled in a regular 2D domain. Similar to the first case, we pre-train CoDA-NO with an incompressible Navier-Stocks equation involving just the velocity term. Then, we fine-tuned the model to predict velocity and temperature using few-shot training samples. Here, too, the pre-trained CoDA-NO significantly outperforms the baseline, reducing the prediction error by factor of two.
also train CoDA-NO on diverse set of PDEs, which form a subset of PDEBench [19] and demonstrate superior performance and parameter efficiency over prior approaches in learning all of those PDE systems. CoDA-NO consistently outperforms the FNO architecture trained on the same set of PDEs, reducing test error by up to 43% while only requiring 2% of the parameters.

---

Our contributions are as follows:
We propose co-domain attention neural operator that efficiently learns solution operators PDEs by formulating transformer operations function space and ensuring discretization convergence.
The proposed architecture enables self- supervieei learning in function space for diverse physical systems by handling varying numbers of input functions and geometries.
CoDA-NO achieves state-of-the-art performance in generalizing to unknown physical systems with very limited data. That is, CoDA-NO can be viewed as the first foundation neural operator for multiphysics problems.
Related Works
Transformers for solving PDEs. Recent work [20] proposes method to weight variables/codomains of the input function based on the weights calculated from the PDE parameters. Another study [21] proposes scalable transformer architecture by combining projection operator to dimeennnsia domain and learnable factorized kernel. In contrast to these works, CoDA-NO provides a complete at- tention operator by considering each physical variable as token function, i.e., vector, extending traditional transformers that only operate on mnite-dimensional tokens.
Self-supervised learning. Self-supervised learning (SSL.) has been proposed to tackle the issue of limited labeled data [22-24]. It allows the training of large foundation models on massive amounts of unlabeled data in the field of computer vision and natural language processing. Subsequently, these models can be successfully applied to a wide range of downstream tasks with minimal to no additional task-specific data [6, 25-27].
Pre-training for PDE solving. Models that are pre-trained in self- supersseee fashion have also gained traction in the domain of scientific computing. One recent study [8] proposes pretraining the models with autoregressive tasks on diverse dataset of multiple PDEs. These models can then be fine-tuned for specific downstream PDEs. Several recent studies have investigated task-agnostic approaches through masking-and and-innnatnteeonntuuutttinn [22] and the consistency of representations under symmetry transformations [16, 28, 29]. Recent work [7] also sheds light on the transferability of these models between different systems of PDEs. While these methods achieve good performance, the target (downstream) PDE must maintain a strict resemblance to the ones used for pretraining, In addition, adapting these models for PDEs with new additional physical variables is not possible. Additionally, ViT-based patching approaches [11] disrupt the continuity and are not resolution-agnosticc.
3 Method
Let us first define our setting and provide a brief introduction to neural operators. For further details, refer to Sec. A in the appendix.
For an input function a: + Rd we will denote the d,, dimeemsiiii output space Rd as the codomain. We consider the components of the codomain as different physical variables, given by real-valued functions over the input domain D, i.e., = [a*...., hm with a' D - The same applies the output function u: Idoue. We define the action of a pointwise operator H. (f : + Rdr) + {g D + Rd, given by function ha IRds + R& with parameters 0 as
H[ʃl()) ho()(x)).
Moreover, define integral operator T : Rdd} : - RRd)) given by kemnl function KG with parameters Φ as
T(((x) - / kala, v) (v)
Problem Statement. Our objective to construct general neural operator architecture that explicitly represents the interaction between the physical variables of PDE systems. Such an architecture should be able to learn and predict various systems without being constrained to a fixed number of variables.
Let's consider two input functions a: D + IRd and à: D + Rd of two different PDE with corresponding output functions D + dod and ū: D Rode In general, the functions and

---

represent din and d3 physical variables over the domain D with din + din We aim to design neural operator architectures G that can both be applied as well as ā despite the different codomains of the input as well as output functions.
Such property provides the possibility evaluate finetune the operator PDEs with different numbers of variables than those on which it was trained. In particular, when the PDE systems have overlapping physical variables {a" } Fi=1 { * 0. this naturally allows to transfer learned knowledge from one system to the other. We will next describe the details of the CoDA-NO layers and architecture to achieve this goal.
Neural Operator on Sets. As we consider the vector-valued input function a as a set of din functions {a!, a² ...... a } that represents different physical variables of the PDE, seek to construct operators that act on sets of input functions with different cardinalities.
For an efficient implementation of operators sets of functions, we mimic transformer architectures and share weights across different variables. Specifically, we can define the integral operator Tper as
Iper [a] = ] Iaadin)
where a = [a!....,,a" _ and I is regular integral operator as described in Eq. (2). Such construction makes the operator permutation -etiitironeuivaiiiantt with respect to the order of the variables in the set. Following the same mechanism, Hper with a shared pointwise operator H (see Eq. (I)). We will use FNO per and GNO per to denote permutation-equiiariin operators using shared GNO and FNO, respectively.
,...,,
can also define crrmutation-cquivaaiann pointwise operator
CoDA-NO Layer. To explain the CoDA-NO layer, let us assume the input function a has been processed into latent function w. D + Rd. We partition the function into a set of so-called token functions w : DD Rd with w E W for j {1....T) T) along the codomain, such that = [w'....'' w²' (and where each un is associated with precisely one of the physical input variables). That is, represents the codomain d' = d' If no other value is specified, we assume that d' = 1. The CoDA-NO layer now processes the token functions using an extension of the self-attention mechanism to the function space (see Appendix Sec. B and Fig. 2a).
concatenation of the token functions w and
Let us begin by introducing a single-head CoDA-NO layer. Later, will expand the concept to multi-head codomain attention. We extend the key, query, and value matrices of the standard attention (see Appendix Sec. B for details) to operators mapping token functions w D - RRd to key, query, and value functions. We define the key, query, and value operators
K: W + (k' D Rd;,, Q : W Ea' D + ) V: W. + (v' D + Rd.
(4)
Assuming d2 = dgy we denote by A = K[w'), J. q¹ = [wy) and vs = Vlwl]] the key, query, and value functions of the token functions, respectively.
Next, we calculate the output (token) functions of + Rdd as
Softmax tmax
(5)
(gl A
where r is the lemperature hyperparameter. Here, (..-) denotes a suitable dot product in the function space. We take the LD, Rd (= product given by (q', km)) = (g) (I), (x)) dæ, where the integral can be discretized using quadrature rules, similar to the integral operator in Eq. (2).
To implement multi-head attention, we apply the (single-head) attention mechanism described above separately for multiple heads h E {1,... H} using Kch, Qh, and yd to obtain oüth We then concatenate these outputs gith along the codomain and get ed = ..........H) Finally, we use an operator
M: (c) D -RH.ddd + (0 : Rdd
to get the output function o'.
5

---

We obtain the output of the attention mechanism by concatenating o's as o [o' oo',"².... oo....oo],, Finally, we complete the CoDA-NO layer by upplying a permutation -euiiiironniiiiiiiatttta integral operator p on o. When CoDA-NO is acting on functions sampled on uniform grid, the internal operators K. O", yh, M,, and I are implemented as FNOs.
Function Space Normalization. Normalization is a vital aspect of deep learning architectures. However, when comes to neural operators mapping infinitediimnniioaal functions, this topic remains largely unexplored. We now provide natural extension. Given function w, let w D + R be token. Then we calculate the mean j E IRd" and standard deviation a E R for this token as
0
/ u (z) dx, o²
(L (1d 12)002 dx
Here, or denotes the elementwise (Hadamard) th -power. The normalization operator can be written
Algorithm 1 Adaptation of CoDA-NO from two phys- ical variables a a² new variable a² on uniform 1D grid. Only the parts in 'Blue' are additionally re- quired to incorporate the new variable a².
+b. Here b E Rd and E Rdd are learnable bias and gain vectors and @ and denote elemen- twise division and multiplication operation. This normalization can seen as an exten- sion of instance normalization [30] for func- tion spaces. Similarly, normalization variants, such as group norm, layer norm, and batch norm, extend to operator learning with these definitions of statistics [31-33].
Norm(w"](x) (g
Require: a' a²,a². a² E R
Variable Specific Positional Encoding: // Learnable Fourier coefficients
K 2² x² E Cdex
à <- concatenate(( a1 iFFT (w
a concatenate a²
iFFT
a concatenate iFFT
Lifting:
// Lifting from Bdm+ixxn to
Variable Specific Positional Encoding (VSPE). We learn positional encoders e' D + Rd for each physical variable for the given vector-valued in-
Rx* with D den +1
w + pointwiseMLP(a')
w² pointwiseMLP (a²)
w + pointviseMLP(
put function = Ia' la',.... ad.].. We con- catenate each positional encoding et with the respective variable a' 'D + I along the codomain obtain extended input func- tions o' = [a', e*). Next, we apply shared pointwise lifting operator P {a' : D > gden +1 (w² D + IR D), typically with D > den + 1. Finally, we concatenate w², i E [I, (.n). tion
Codomain Attention Block: X
k k2 k3 +- FNO,(2u (w"), FNO, (w²) FNO (w=)
q'.q² q² +- FNO, (w FNO, (w²) FNO, (w
v 02 v + FNO, (w FNO, FNO, (w²)
// Compute attention coefficients M
for i E {1, do Mi..j] < (q',"" k)
J. to get the lifted latent func-
end for
M + softmax M
(8) In the previous paragraphs, we used d = D di. and, to maintain the permutation- equivariance property of the operator, d' must divide D.
w w= " wdd]: D D.di.
// Compute attention outputs for i (1,2,3} do
o' + (), w x M[i.j] end for // Normalize and Residual
Algorithm presents the pseudocode for the CoDA-NO architecture applied to input func- tions [a a²), mapping two different physical variables a uniform grid in ID domain, to the solution functions [u 2²]. J.. It assumes d' = D while designing the CoDA-NO layer. Notably, to incorporate another function a², representing new physical variable, it is only necessary to introduce a corresponding param- eter for the new VSPE, denoted as 3
0 + norm 0 + w
0² norm 0² 2
a + norm
3 w w² w FNOx(o=), FNOr
Projection:
// Projection from DDx to Rlxn
ut + pointwiseMLP(Ne
2² + pointwiseMLP
3 + pointwisoMLPi
To effectively handle non-uniform complex geometries, we follow the GINO architecture
return w!, u² u

---

[34], where a GNO is used as an encoding and decoding module. Given set of evaluations of an input function a on mesh, as represented by {a(x;")) ))== where {a] 1a T==1 Din. our first step involves concatenation of each physical variables with respective VSPEs (see Fig. 2a).
Next, GNO per transform the function into new function up uniform latent grid, represented by {a? {a:2nd) " Ji=1- Finally, we apply / stacked CoDA-NO layers to wy to obtain the encoded function w1, which acts representation of the input function
The decoding module is essentially a mirrored version of the encoding module. It starts by applying another block of I stacked CoDA-NO layers to the encoded function to obtain wz. Subsequently, uses another GNO operator transform wi on uniform grid to an approximation u of the solution function on an arbitrary output grid {u(=²"1) J The architecture is summarized in Fig. 2a.
Model Training. To seamlessly adapt to multi-physics PDEs with limited data, we propose two- stage training process: Self-supervised pretraining is followed by a supervised fine-tuning stage. For summary, we refer to Fig. 2b.
Pre-training. In the context of sulf-supervised pretraining, the objective to train the model to reconstruct the original input function from its masked version. Within this phase, the model's s encoding component denoted as the Encoder, while the decoding component comprises the Reconstrucior. The values of the input function at specific percentage of mesh points are randomly masked to zero, and certain variables (chaneek/co-dommiin)) of the input function are entirely masked to zero. The model is then trained to reconstruct the original input from this masked version.
We emphasize that the self-supervised learning phase is agnostic of the downstream supervised task and only requires snapshots of simulations of the physical systems.
Fine-tuning, In the supervised fine-tuning phase, the Reconstructor is omitted from the decoding module and replaced by randomly initialized Predictor module. The parameters of the Encoder and VSPEs are copied from pre- traitrained weights. If the fine-tuning (target) PDE introduces variables that are not present in the pre- training PDE; we train additional variable encoders only for these newly introduced variables (see Fig. 2b). This ensures that the model adapts to the expanded set of variables needed for the fine-tuning task with minimal additional parameters.
4 Experiments
We conduct experiments on two coupled PDEs; fluid-structure interaction and Rayleigh-Bénard convection system. We also test our model on diverse of PDEs from PDEBench [19]. The code is available at / com http://uurroorrorr/////o CoDA-NO.
Modeling Fluid-Structure Interaction. We consider the following problems: (a) fluid dynamics problem, where a Newtonian, incompressible fluid impinges on rigid object. and (b) a fluid-structure interaction problem between Newtonian, incompressible fluid and elastic, compressible solid object [17]. We denote 2/ (resp. 12)) the domain occupied by the fluid (resp. the solid) time t. The dynamics of the fluid are governed by the Navier-Stokes equations
p Ou 0t
p'v (u = V v.o! of, V u 0, in of
where u and pf denote the fluid velocity and density, respectively. And of denotes the Cauchy stress tensor, given by r. -pll + v(Vu+ Vu²), where is the identity tensor, p the fluid pressure, and the fluid dynamic viscosity.
For fluid-structure interaction, the deformable solid is governed by the elastodynamics equations
Rd 1=
VV 01
(10)
with F - I+ dd and J = FF) Here d, p², F, and 0² denote the deformation field, the solid density, the deformation gradient tensor, and the Cauchy stress tensor, respectively (see Eq. (18) in the Appendix). The fluid dynamics (resp. the fluid-structure interaction) problem considers a fluid flow past a fixed, rigid cylinder with a rigid (resp. elastic) strap attached. The details regarding the geometric setup (see Fig. 3), time-dependent inlet boundary condition, and the initial conditions
7

---

provided in the Appendix Sec. C.I.
Modeling Rayleigh-Bénan of fluid layer heated from below and cooled from above. The governing equations for the Rayleigh- Bénard system consist of the incompressible Navier-Stokes equations coupled with an energy equation for heat transfer. The system is modeled as follows:
Convection. The Rayleigh-Bénard convection system governs the flow
du Vu+ VP- - -V
(11)
0T VT
+
xV
(12)
Dataset Deseription and Generation. To study the fluid-structure interaction system, two datasets, the fluid-structure interaction (NS+EW dataset) and the fluid dynamics(NS dataset), are generated using the TurtleFSI package [35].
We simulate the fluid-structure interaction and the fluid dynamics test cases described above up to time T 10, using constant time-step ōt Tr, where n 1000. The data sets are composed of solution trajectories [uk. Pt: de] (resp. [u,, Pe]). which denote the approximate solution of the fluid- structure interaction problem (resp. the fluid dynamics problem) times t iôt, i E {O..... n). These trajectories are generated on the basis 3 parameters (H4, C1, C2) describing combinations fluid viscosities ja (0.5, 1,5, 10} and inlet conditions, (cI, 02) E I.
For our setup, the fluid considered is water, with density of 1000kg.m- and maximum inlet velocity of approximately 4m.s=-, leading to Reynolds (Re) numbers in the range 200 4000 (for μ between 10 - 0.5). Modeling fluid-solid interaction or only fluid motion with such high Reynolds numbers is challenging and serves benchmark problem [12, 17] (See Sec. C.2 for detailed explanation).
To study the Rayleigh Bénnar- convention system, we degenerate two different PDE datasets. Firstly, we generate Rayleigh-Bénn Bénnar convection system with Ra number 12 x 10 and 20 10 We set the temperature difference between the top (cold) and bottom (hot) boundaries to 1. We assume no-slip boundary conditions, and start the convection process, we also add initial temperature perturbation. Additionally, we generate incompressible Navier-Stocks equations with Reynold number Re 500 with cyclic boundary condition on a uniform 2D grid [36] (for details, see Appendix Sec. C.3).
Experiment Setup. For the fluid-structure interaction system, we conduct two distinct pretraining procedures for CoDA-NO and obtain two pretrained models: - and G The former is pretrained on a fluid-structure interaction dataset that combines the Savier-Stokes equation and the elastic wave equation, denoted -i The latter, G pretrained on a fluid motion dataset governed solely by the Navier-Stokes equation. In both scenarios, the pretraining involves utilizing 8000 snapshots of flow and displacement fields with Re {200, 2000].
The supervised task involves training the model to predict the system's 's state at the subsequent time step based on its current state. For the fluid-structure interaction dataset, we train operator and d are the velocity,
GNS-EM such that Gus-ER Ex [M,, Pt, di] - [Ue+&t, Pt+65. de+5t), where pressure, and mesh deformation fields (see Sec. 4). For the data with only fluid motion, we train the operator Gus which maps between the current and next time step velocity and pressure field as Gws : : - [we+ōtr Pt+
The pretrained model for both datasets is fine-tuned for unseen viscosity μ = 5.0 = 400) with different numbers of few shot examples. The inlet conditions of these simulations excluded from the pretraining data. So, the target PDEs' viscosity and inlet conditions are absent the per-taining dataset. We test the model's adaptability on more turbulent fluid-solid interaction dataset with Re = 4000( 0.5) by finetuning both pretrained models - and Ge on each dataset.
For the Baylligh-Bénaaee convention system, we pretrain a CoDA-NO model, denoted as Ges- on the incompressible Navier-Stokes equations using 40, 000 snapshots in self-supervised manner. The supervised task for this system is to train an operator, Gus-r r : [u,, T.] -> [u++8e: Te+s where represents velocity and T' represents temperature. The pretrained model Gp is fine-tuned for the

---

Table 1. Test L2 loss for fluid dynamics (NS) and fluid-solid interaction (NS+EW) datasets with viscosity
400 and Re 1000 for different numbers of few-shot training samples.
Re
400
Re 4000
Pretrain
Model
# Few Shot Training Samples
Dataset
5
25
100
5
25
100
Evaluation Dataset
NS
NS+EW
NS
NS+EW
NS
NS+EW
NS+EW
NS+EW
NS+EW
GINO
0.200
0.122
0.047
0.053
0.022
0.043
0.717
0.292
0.136
DeepO
0.686
0.482
0.259
0.198
0.107
0.107
0.889
0.545
0,259
GNN
0.038
0.045
0.008
0.009
0.008
0.009
0.374
0.310
0.132
VIT
0.271
0.211
0.061
0.113
0.017
0.021
0.878
0.409
0.164
U-Net
13.33
3,579
0.565
0.842
0.141
0.203
3,256
0.563
0.292
0.182
0.051
0.008
0.084
0.006
0.004
0.326
0.264
0.070
Ours
NS NS+EW
0.025
0.071
0.007
0.008
0.004
0.005
0.366
0.161
0.079
0.024
0.040
0.006
0.005
0.005
0.003
0.308
0.143
0.069
supervised task of solving Rayleigh-Bénne Bénnar convection using different numbers of a few shot training samples.
interaction system, we train GINO [18]. DeepONet [37], graph neural network (GNN) [38], vision transformer (ViT) [10], and the Unet [39] model from scratch. The mesh points of the NS and NS+EW datasets are irregular and change for each sample. So, to efficiently handle irregular mesh, in the branch network of DeepONet, we use GNN layer followed by MLPs. Also, as ViT and Unet can handle irregular mesh, we follow the architecture of GINO and use a GNN layer to query the latent function on a uniform grid. We then apply Unet and VIT to the uniform grid, followed by another GNN layer, to get the output at the desired query points. For the Rayleigh-Be Bénard convection system, we train Unet [39] and FNO [12] from scratch and compare them against our proposed model.
Baselines. For comparison on the supervised tasks on
It should be noted that employing the existing models for pertaining and subsequent finetuning on the target datasets is nontrivial due to complex geometry and the changes in the number of physical variables between the pertaining and target datasets. We report the L error between the predicted and target functions, which serves as a measure of model performance. Additional implementation details are provided in the Appendix Sec. H.
Results. In Tab. we report the performance of model and the baselines for modeling the fluid-structure interaction. We observe that the pretrained CoDA-NO model performs better than the baselines. Importantly, the performance gain higher when the number of few-shot examples is very low. This demonstrates the sample efficiency and generalization capability of CoDA-NO to previously unseen physical systems.
Next, when CoDA-NO is pretrained solely on the NS dataset, it shows an impressive ability to adapt to the more challenging NS+EW dataset. Finally, when CoDA-NO is pretrained on the more intricate NS+EW dataset, it easily adapts to the simpler NS dataset through fine-tuning. This underscores the capability of the CoDA-NO to adjust between different PDEs with varying numbers of variables seamlessly.
Also, notice that pretrained CoDA-NO performs better than CoDA-NO trained from scratch, demonstrating the effectiveness of the pretraining scheme. We also provide the energy spectra of the predicted fluid flow by the different models in Sec. F.4 where we observe that the energy spectrum remains closest to the ground truth.
In Tab. 2, we present the result on modeling the Rayleigh-Bénan -énnar convention. We observe that pretrained CoDA-NO outperforms every other baseline and adapted to the new temperature variable, T, of the Rayleigh-Bént -éenadd system. Similar to the fluid-structure interaction problem, we also observe that the pretrained CoDA-NO outperforms CoDA-NO trained from scratch, which underlines the effectiveness of our pretraining and adaptation mechanism.

---

Table 2. Test Lg for Rayleigh-Bénard convection system with coupled Navier-Stokes and energy (heat)
equation with Rayleigh number Ra 12 10 and Ra x 10 for different few shot examples.
Ra 12
10
Ra 20
10
Pretrain dataset
Model
#Few Shot Training Samples
5
10
25
5
10
25
Unet
0.049
0.025
0.013
0.126
0,083
0.075
FNO
0.119
0.070
0.044
0.491
0.166
0.127
0.067
0.045
0.035
0.221
0.058
0.040
Ours
NS
0.016
0.007
0.002
0.074
0.040
0.029
Additionally, we also conduct experiments on various PDEs from the PDEBench dataset [19], where we show superior performance and parameter efficiency (see Appendix Sec. G).
Adaptation More Turbulent Fluid-Structure Interaction. We also test the adaptation capability of our pretrained model on a more turbulent fluid-solid interaction scenario with viscosity μ = 0.5 with a Reynolds number of 4000. From Tab. I, we can observe that, even though the model is pretrained on data with lower Reynold's number (200 - 2000), it can seamlessly adapt to more turbulent flow and outperform baselines with a significant margin.
Ablation Studies. To demonstrate the effect of each of the proposed components, namely, codomain attention, normalization layer, VSPE, and pertaining. we present the result of detailed ablation study in Appendix Sec. F.1. We observe that substituting the codomain attention with regular patch-based attention impacts the model's performance. In particular, removing the normalization layer prevents the model from converging.
We also provide ablation study fine-tuning methods. Instead of fine-tuning all the parameters, here, we freeze the parameters the "Encoder" and only train the parameters of the "Predictor" and VSPEs. This minimized the number of trainable parameters during fine-tuning. Also, in this case, performed significantly better than the other models (see Appendix Sec. F.5).
We also provide the results for the zero-shot super-resolution task, where directly predict the output function on a much denser mesh than the training mesh. Our findings show that CoDA-NO outperforms other baselines significantly (see Appendix Sec. F.2).
Additionally, we have conducted comparative analysis of the parameter count and computational cost for each model, which points to the overfitting problem of the baseline when learning complex multi-physies PDEs (see Appendix Sec. F.3).
Limitations. In general, CoDA-NO's performance on target PDEs influenced by the number of training examples, and we highlight the potential for further enhancement through the integration of physics-informed approaches.
Conclusion
In this work, we introduce CoDA-NO, versatile pre-trained model architecture designed for seamless adaptation to Partial Differential Equations (PDEs) featuring diverse variable compositions. Departing from conventional patch-base based attention modules, CoDA-NO innovatively extends the transformer to function spaces by computing attention across co-domains. Leveraging flexible variable encoding scheme and graph-based neural operator module, CoDA-NO exhibits adaptability to any target PDE, accommodating new and previously unseen variables with arbitrary input-output geometries during fine-tuning. Our empirical evaluations demonstrate that CoDA-NO consistently outperforms baselines across varying amounts of training data and exhibits robustness in handling missing variables. Our findings on complex multiphysics simulations underscore the efficacy and adaptability of CoDA-NO, positioning it valuable tool for addressing challenges in machine learning for PDEs.
10

---

Acknowledgments
A. Anandkumar is supported in part by Bren endowed chair, ONR (MURI grant N00014- 18-12224) and by the AI2050 senior fellow program at Schmidt Sciences. We thank David Pitt for his support in adding our code the neuraloperator library, facilitating broader use and accessibility.
References
[1] Gilbert Strang, Computational science and engineering. Optimization, 2007.
[2] Gege Wen, Zongyi Li, Qirui Long, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M Benson. Real-time high-resolution co geological storage prediction using nested fourier neural operators. Energy & Environmental Science, 16(4):1732-1744) 2023. 1
[3] Burigede Liu, Nikola Kovachki, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, Andrew M Stuart, and Kaushik Bhattacharya. learning-based multiscale method and its application to inelastic impact problems. Journal of the Mechanics and Physics of Solids, 2022.
[4] Geoffrey K Vallis, Atmospheric and oceanic fluid dynamics. Cambridge University Press, 2017.
[5] Mathilde Caron, Hugo Touvron, Ishan Mísra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. Proceedings of the IEEE/CVF international conference computer vision, 2021. 2
[6] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et Learning transferable visual models from natural language supervision. In International conference on machine learning. PMLR, 2021. 2,4
[7] Shashank Subramanian, Peter Harrington, Kurt Keutzer, Wahid Bhimji, Dmitriy Morozov, Michael Mahoney, and Amir Gholami. Towards foundation models for scientific machine learning: Characterizing scaling and transfer behavior. arXiv preprint 236.2230660025,, 2023. 2,4
[8] Michael McCabe, Bruno Régaldo-Sain Blancard, Liam Holden Parker, Ruben Ohana, Miles Cranmer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Geraud Krawezik, Francois Lanusse, et al. Multiple physics pretraining for physical surrogate models. arXiv preprint 10.0092,, 2023. 4,22
[9] Zhongkai Hao, Chang Su, Songming Liu, Julius Berner, Chengyang Ying, Hang Su, Anima Anandkumar, Jian Song, and Jun Zhu. Dpot: Auto-regressive denoising operator transformer for large-scale pde pre-training. arXiv preprint arXiv:2403.03542. 2024. 2, 19, 22
[10] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. International Conference Learning Representations, 2021. 2,9
[11) Bowen Zhang, Shuyang Gu, Bo Zhang, Jianmin Bao, Dong Chen, Fang Wen, Yong Wang, and Baining Guo. Styleswin: Transformer-base gan for high-resolution image generation. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, 2022. 2, 4
[12) Zongyi Li, Nikola Borislavov Kovachki, Kamyar Arizzadenesheli, Kaushik Bhattacharya, An- drew Stuart, Anima Anandkumar, et al. Fourier neural operator for parametric partial differential equations. In Proceedings of the International Conference on Learning Representations (ICLR),
2021. 2, 8, 22
[13] Kamyar Arzizadenesheli, Nikola Kovachki, Zongyi Li, Miguel Liu-Schiaffini, Jean Kossaifi, and Anima Anandkumar. Neural operators for accelerating scientific simulations and design. arXiv preprint arXiv:2309.15322 2023, 2

---

[14] Jean Kossaifi, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, and Anima Anandkumar. Multi-grid tensorized fourier neural operator for high resolution PDEs, 2023.
[15] Boris Bonev, Thorsten Kurth, Christian Hundt, Jaideep Pathak, Maximilian Baust, Karthik Kashinath, and Anima Anandkumar. Spherical fourier neural operators: Learning stable dy- namics on the sphere. Proceedings of the International Conference on Machine Learning (ICML), 2023.
[16] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differen- tial equations. arXiv preprint arXiv:2010.08895, 2020. 2, 4, 14
[17] Stefan Turek and Jaroslav Hron. Proposal for numerical benchmarking of fluid-structure interaction between an elastic object and laminar incompressible flow. Springer, 2006. 3, 7, 8, 16
[18] Zongyi Li, Nikola Borislavov Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Prakash Otta, Mohammad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Azizzadeneshelii and Anima Anandkumar. ceometry-informee neural operator for large-scale 3d pdes. arXiv preprint 23X.iv2339900588. 2023. 3,
[19] Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Dan MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. PDEBENCH: An extensive benchmark for scientific machine learning, 2023. 3, 7, 10, 19, 22
[20] Makoto Takamoto, Francesco Alesiani, and Mathias Niepert. Learning neural pde solvers with paramcter-guided channel attention. arXiv preprint arXiv:2304.14118, 2023.
[21] Zijie Li, Dule Shu, and Amir Barati Farimani. Scalable transformer for pde surrogate modeling. arXiv preprint 20X5:22005.15500 2023. 4
[22] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, 2022. 4
[23] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.
[24] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. International conference on machine learning. PMLR, 2020. 4
[25] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 2023. 4
[26] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic rext-to-image diffusion models with deep language understanding, Advances in Neural Information Processing Systems, 2022.
[27] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEFSCVF international conference on computer vision, 2021. 4
[28] Bonan Xu, Yuanye Zhou, and Xin Bian. Self-supervised learning based on transformer for flow reconstruction and prediction. arXiv preprint arXiv:2311.15232, 2023.
[29] Grégoire Mialon, Quentin Garrido, Hannah Lawrence, Danyal Rehman, Yann LeCun, and Bobak Kiani. Self-supervised learning with lie symmetries for partial differential equations. In ICLR 2023 Workshop on Physics for Machine Learning. 2023. 4
[30] Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky, Instance normalization: The missing ingredient for fast stylization. arXiv preprint arXiv 2016. 6
12

---

[31] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning. PMLR, 2015. 6
[32] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint (517.:6600004500 2016.
[33] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), 2018. 6
[34] Zongyi Li, Nikola Borislavov Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Prakash Otta, Mohammad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Azizzade- nesheli, et al. Geometry-informed neural operator for large-scale 3d pdes. arXiv preprint arXiv:2309.00583 2023." 7
[35] Aslak W Bergersen, Andreas Slyngstad, Sebastian Gjertsen, Alban Souche, and Kristian Valen- Sendstad. turtlefsi: robust and monolithic fenies-based fluid-structure interaction solver. Journal of Open Source Software, 2020. 8
[36] Valentin Duruisseaux, Miguel Liu-Schine Schiaffii,,,, Julius Berner, and Anima Anandkumar. Towards enforcing hard physics constraints in operator learning frameworks. ICML 2024 Al for Science Workshop, 2024. 8
[37] Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators.
arXiv preprint arXiv: 1910.0311, 2019. 9
[38] Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. Advances in neural information processing systems, 2016.
[39] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Inter- vention. Springer, 2015.
[40] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint arXiv:2108.0844 08488, 2021. 14
[41] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar, Neural operator: Graph kernel network for partial differential equations. arXiv preprint arXiv:2003.03488 2020. 14
[42] Dick Kachuma and lan Sobey. Linear instability of asymmetric poiseuille flows. 2007. 15
[43] Sukhendu Ghosh. Relative effects of asymmetry and wall slip on the stability of plane channel flow. Fluids, 2017. 15
[44] Anders Logg, Kent-Andre Mardal, and Garth Wells. Automated solution of differential equations by the finite element method: The FEniCS book. Springer Science & Business Media, 2012. 16
[45] Ben Mildenhall, Pratul Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi. and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communi- cations of the ACM. 23
13

---

Appendix
The appendix is organized follows:
In Sec. A, we provide brief introduction of Neural Operators.
Sec. B, we describe regular attention mechanism.
In Sec. C, describe the detailed regarding the dataset generation.
In Sec. D, we provide comparison between the CoDA-NO and FNO architecture.
In Sec. E, we provide the visualization of ground truth vs predicted solution function by CoDA-NO.
In Sec. F, we provide the ablation studies and additional evaluation metrics.
In Sec. G, we provide the results of our experiment on PDEs from the PDEBench dataset.
In Sec. H, we provide additional implementation details.
A Neural Operators
Neural Operators class of deep learning architectures designed learn maps between infinite- dimensional function spaces [40]. A Neural Operator seeks to approximate an operator G that maps input function
to its corresponding output function E LI by building parametric map 9s A - U. The typical architecture of a Neural Operator can be described as Gø = P - Tz o... Ti o L. Here, L: - wo and P: WL
lifting and pointwise projection operators, respectively. The action of any pointwise operator H. { D + Rdr + {g : Rdd, can be defined as
H[[(()) ho(s())),
(13)
where he Rd4 + Rd is any function with parameters 0. The integral operator Ti wi-1 performs kernel integration over the input function wi-1 as
Ti u 1 = / (ar, v) dy.
(14)
D._1
Here, D;-1 is the domain of the function w1-1- In the case of Fourier Neural operators (FNO) [16], convolution kernel, i.c., (x, y) = kz(z - y) was used. By the convolution theorem, this enables the representation of an integral operator as pointwise multiplication of the Fourier coefficients as follows wy = F (F(k)) O F(mn-1)).
For the Graph neural operator (GNO) [41], small neighborhood B, (a) n D;-1 around the point considered instead of integrating over the whole domain D; such that Eq. (2) changes
(x) = (,)nnD:i
la(2, v) dy.
(15)
Given set of evaluations of the function w1-1 points { C Di-1, the kernel integral can be approximated by
wr(xr)
Σ
la M.
(16)
(E
where & are suitable quadrature weights [40]. The discretized kernel integral can be viewed as
message passing on graphs, where the neighborhood of each point consists of all points within radius r.
B Attention mechanism for fiitt-dimmensional vectors
No Given three sets of vectors, so-called queries {q: [q keys (k, ).Y, and values {v.)/a, with N* = N₂ and matching dimensions of queries and keys, attention mechanism calculates weighted sums of the value vectors. Specifically, the set of output vectors {o;), can be expressed that
0, a (V
VN. i == 1.....
(17)
14

---

[q.kk,]] and is the temperature term. For the self-attention mechanism, the key, query, and value vectors are calculated from some input sequence {z) using the key, query, and value matrices K, Q and V as
where at Sof tMax (gckk)
QZ!, k;- = Kz
Vz.
C Dataset Description
0.2
0.2
0.0
1.0
1.5
My t
By at
Figure 3. Visualization of horizontal velocity u, at and + 56 time step.
C.1 Fluid-Structure Interaction System
Here, we provide the details on generating the fluid-structure interaction dataset involving Navier- Stokes and Elastic wave equations.
Fluid-structure interaction model. o² verifies
Under the Kirchoff St- A--enant model, the Cauchy stress tensor
1 F(X (fr(((tr())I) 2n"E)
o² J
(18)
where X and μ the Lame coefficients, and
FFT - I)).
Inlet Boundary Condition. Time dependent inlet boundary conditions consist of 4th order poly- nomials velocity profiles which vanish at the channel walls [42, 43). The inlet conditions are given by
og (v. t) = v(t) 81 H) (y - c1 # c2
(19)
H(1 - c1) -z)
Here v is the ramp function defined as
u(t) - {
70 (1 cos
if 0 <t<2
(20)
140
t22
and (c1,02) E I, where
I= (a, b) E (-6, -4, 2,,022,4666 as5
(21)
enforced the inlet 0.
Geometric setup, boundary, and initial conditions. In the considered setup (see also Figure 3), fluid flows past a fixed cylinder radius R 0.05 centered (xc, (0.2, 0.2) in two- dimensional channel of length L = 2.5 and width H = 0.41. A deformable elastic strap of length ê 0.35 and height h 0.02 is attached to the back of the cylinder. Note that, in the test cases considering fluid motion exclusively, the elastic strap is assumed to be rigid.
In the case of the fluid-structure interaction, the interaction conditions arise from the mechanical equilibrium at the boundaries of the strap, which are given by
o²
ad
2 0t
15

---

where n denotes unit normal vector to the fluid-solid interface. No-slip boundary conditions are imposed on the fluid velocity the top (resp .bottom) boundaries of the channel at y 0 (resp. = H). as well as on the boundaries of the cylinder and the elastic strap. Outflow boundary conditions are imposed at x 2.5 by enforcing the values p = for the pressure.
The initial conditions
(u, p. d) 0,0
where the displacement d 0 corresponds to a perfectly horizontal elastic strap and is imposed at time t = O.
Details regarding the data set generation. The Tur package provides monolithic solver for the fluid-structure interaction test case, that is, combining the equations describing the solid and fluid evolution into one coupled system based on an Arbitrary Eulerian-Lagrange Lagrangian (ALE) formulation of the problem and developed on the FEniCS computing environment [44].
The initial conditions are expressed at set X Xs X of mesh points, corresponding to the union of the solid and fluid domains. In the ALE formulation, at each snapshot 0 S t Σ t. of the simulation, the solution is given set of mesh points X; X + d2, where de denotes the mesh displacement. In particular, the snapshots u, (resp. Pt) correspond to numerical approximations of the velocity (resp. the pressure) the mesh points X.. Notably, while equation (10) governs the deformation field in the solid domain Qi, the displacements de are obtained through an extension of the deformation field to the fluid domain 9/ via biharmonic extrapolation.
In all the cases considered, the values of = 1.0 x 10", p² = 1.0 x 10 X 4.0 x 10t and u 2.0 x 10 were used. The simulations were performed using constant time step öt 0.01.
C.2 Justification of Experiment Design
For our setup, the fluid considered is water, with a density of 1000 kg.m-3 and a maximum inlet velocity of approximately 4m.s-1 leading to Reynolds numbers in the range 200-20000000000 2000 (===1 10 1) for our experiments. Only when the flow becomes turbulent can ample movements of the elastic strap (Fig. 4) be observed in the fluid-structure interaction case. Modeling fluid-solid interaction only fluid motion with such Reynolds number is quite challenging and used as benchmark problem [17].
Modeling fluid-solid interaction with an even higher Reynolds number requires a very high computa- tional cost. Because TurtleFSI's (used in this study) fluid solver, including its' fluid-structure interac- tion solver, uses a direct numerical simulation (DNS) of fluid dynamics and does not employ any turbulence models. This means that in order to accurately capture the small-scale energy-dissipatinn vortices that form when the flow interacts with the cylinder and strap at high Reynolds numbers, a very fine spatial domain discretization is required. Furthermore, extremely small time step (At) is necessary to ensure numerical stability. For these reasons, the contribution [17], which introduced the benchmark fluid-structure interaction problem studied here, only deals with flows that have Reynolds numbers less than or equal to 200.
It's crucial to highlight significant disparity between the pre- pre-training and finetuning stages, par- ticularly concerning examples with viscosities I and 10. This disparity arises from the utilization of distinct inlet boundary conditions during the pre-training and finetuning phases. Consequently, even though the viscosities align with the pre-training dataset during finetuning on PDEs featuring j E (1, 10), , the model faces formidable challenges in adapting due variations in inlet conditions. The finetuning dataset with v has different viscosity as well as intel conditions compared to the pre-training dataset, serving an out-of-distribution PDE setup.
C.3 Generating Bennad dataset
The initial temperature field is initialized with linear gradient between the hot bottom boundary, T bottom 1, and the cold top boundary, Top 0. To induce instability and initiate convection, temperature perturbations are introduced in localized regions of the domain. A region centered at
L.
4, 40) perturbed to T = 1, while region near the middle of the domain, centered at
[i Ly
2
16

---

is set to T = -1. These perturbations break the symmetry and help to trigger the onset of convection patterns.
For the incompressible Navier-Stocke equation, we consider two-dimensional Kolmogorov flow (a form of the Navier-Stokes equations) for viscous, incompressible fluid,
iu ät
1
Vu Vp
Au sin(ny) Re
(22)
with the incompressibility constraint V - 0 on the domain [0, 2mJ² x (0, o0). The initial condition given as u(:;0) 0) = Un, where denotes the velocity, p the pressure, and Re is the Reynolds number which we set to 500 for simulation.
D Comparison with FNO
We would like to bring out a distinction. In FNO, the mixing of channels happens in the Fourier space in the spectral layer through the linear transform applied the Fourier coefficients. This defines a weighting of some sort on the input channels and how they are mixed. The mixing is global since the Fourier transform a global operation. In CoDA-NO, however, the mixing of channels happens in the spatial domain through the attention mechanism because 1ch,, Q", yh MM, and I are all implemented as FNO's. The attention weights determine how the channels are mixed, allowing for a more flexible and input-dependent mixing. This added flexibility enables CoDA-NO to better capture complex interactions and dependencies between different physical variables, especially in multiphysics problems where the relationships between variables can be intricate and vary depending on the input conditions. This also implies that this is more efficient it can seamlessly incorporate additional or fewer variables during fine-tur mheettnnng, avoiding retraining the whole model from scratch like FNO would have to, which can be computationally expensive.
well as in the Fourier space
E
Visualization of Results
0.4
4.0
3.0
25
0.2
2.0
0.5
0.0
1.5
0.0
(a) Ground Truth
4.5
4.0
110
1.5
0
0.0
0.5
1.0
1.5 (b) CoDA-NO Prediction
2.0
2.5
0.4
0.15
0.00
0.2
0.05
-0100
0.000
-0.25
Eiror
Figure Visualization of CoDA-NO prediction. We plot horizontal velocity uy the structetructure interaction problem.
17

---

Additional Results
F.1
Ablation of Proposed components
Table shows that replacing the codomain attention with regular attention mechanism or the removal of any these designed components significantly impacts the model's performance. We also observe that our proposed normalization technique crucial for effective training.
Table Evaluating Le loss across different models using various pre-training datasets and varying numbers of few-shot training samples. "*" indicates configurations that did not converge due to excessive training error
# Few Shot Training Samples
Pretrain Datasct
CoDA-NO
VSPE
Norm
5
25
100
NS
NS+EW
NS
NS+EW
NS
NS+EW
X
X
X
X
0.271
0.211
0.061
0.113
0.017
0,020
V
X
x
0.182
0.051
0.008
0.084
0.006
0.004
NS
0.049
0.079
0.009
0.0132
0.004
0.009
V
NS EW
0.045
0,057
0,010
0.011
0.008
0.004
V
X
NS
0.023
0.008
0.006
X
V
NS
EW
0.057
0.232
0.012
0.052
0.006
0.006
4
V
NS
0,025
0,071
0.007
0.008
0.004
0,005
NS FW
0.024
0.040
0.006
0.005
0.005
0.003
F.2 Zero-Shot Super Resolution Test
Here, we present the results of our zero-shot super-resolution models (see Tab. 4) on complex fluid- solid interaction problems. We train the models with 1317 mesh points the domain (xy plain). However, during inference, the solution function queried directly on a denser and non-uniform target mesh consisting of 2193 points.
We observe that the zero-shot super-resolution performance of CoDA-No is significantly better than the other baselines.
Table 4. Zero Shot Super Resolution Performance on Fluid-Solid (NS-EW) Interaction Problem
Pretrain
Fluid Viscocities
Model
Dataset
=
1 / 10
U-Net
0.144 0.267
0.216
Vit
0.052 0.175
0.046
GINO
0.069 0.103
0.0711
DeepO
0.113 0.107
0.357
GNN
0.223 0.211
0.247
CoDA-NO NS-ES
0.041
0.063
0.048
No
NS
0.032 0.049 0.035
F3 Parameter Count and Computational Cost.
Now the present the number of parameters and training/interference time taken by the proposed model along with different baselines used in the study in Tab. might seem that models are not compared fairly, the CoDA-NO has higher parameter count. However, here, we test the models on few shot learning problems. Increasing the baselines' parameter count worsens the overfitting problem.
18

---

To demonstrate this fact, we perform experiments on fluid-solid interaction dataset with an increased parameter count. We will observe that increasing the parameter count almost always negatively impacts the performance, especially for very few hot learning scenarios (see Tab. 6).
We also note that the additional model parameters and computation are required to learn rich inter-variable dependencies during pre-training and generalize from single to multi-physics dur- ing finetuning. Furthermore, the zero-shot super-resolution capability of CoDA-NO is discussed in Sec. F.2. CoDA-NO is a justified choice due to its seamless adaptation to various PDEs, remarkable performance gap, and zero-shot super-resolution capability despite having a little more computational overhead.
Table S. Comparison of Inference Time, Training Time (in sec.) per sample, and Number of Parameters for different models.
Models
GNN GINO DeepO
VIT
Unet
CoDA-NO
Inference Time
0.012
0.012
0.006
0.071
0.024
0.440
Training Time
0.136
0.136
0.131
0.273
0.268
1.250
# Parameter x x10
0.6 60
6
27
30
43
Table 6. Overfitting of Baselines with Higher Parameters (in x x1e6) NS-EW dataset
# Parameter (Used/High)
# Train=
# #Train=25
# 100
Models
(Used / High)
(Used/ High)
(Used / High)
GINO
60/200
0.122 0.342
0.053 0.066
0.043 0.036
DeepO
6/25
0.482 / 0.495
0.198 0.303
0.107 / 0.083
GNN
0.6/7
0.009
(009//0031
0.009 / 0.061
VIT
27/100
0.211 0.266
0.113/0.125
0.020 / 0.022
U-net
30/48
3.579 /9.462
0.842/3.90
0.203 / 0.412
F.4
Energy Spectrum
Here, we show the energy spectrum for the NS-EW dataset for μ calculated from the test set (see Fig. 5). All models are trained on 100 training examples, Due to numerical error, the measured spectral energy does not decay smoothly in the high-frequency region. However, our models' energy spectrum remains closest to the ground truth.
F.5
Ablation on Finetuning Technique
Here, in Tab. 7, Tab. 9, and Tab. we present some additional ablation studies model's performance when we keeping the weight of the "Encoder" frozen during supervised fine-tuning.
F.6
Error bar
G PDEBench experiments
We finally compare CoDA-NO, FNO, and the recently proposed DPOT [9] on PDEBench [19]. DPOT (Auto- nuto-Reeressive Denoising Operator Transformer for Lurge-Scale PDE Pre- Training) is large-scale pre-training approach for learning PDE representations. It utilizes an auto-regressive denoising strategy and Fourier transformer architecture for efficient pre-training diverse PDE datasets.
To assess the performance of these models on diverse set of PDEs, conduct experiments three single-physics datasets from the PDEBench: Shallow Water Equations (SWE), Diffusion Equations (DIFF), and Navier-Stokes Equations (NS). We follow the same pretraining and finetuning procedure for both models' datasets. And we evaluate the models on the following two tasks
Reconstructive task: The models are trained on the respective datasets during pretraining to use the self-supervised learning objective to reconstruct the masked input. The error
19

---

Energy Spectrum
GNN
VIT
FNO
DeepONet
CoDA-NO (NS-EW)
103
CoDA-NO (NS)
Ground Truth
Energy
102
10
10
20
30
40
50
Wavenumber
Figure 5. Energy Spectrum of the Velocity Field of the fluid on the fluid-solid interaction dataset.
Table 7. Test errors (L2 loss) for fluid dynamics (NS) and fluid-solid interaction (NS+EW) datasets with viscosity 1.0 for different numbers of few-shot training samples. The pre-training is done with 8000 samples taken from NS and NS+EW datasets with viscosities E (1.0, 10.0).
# Few Shot Training Samples
Pretrain
Model
5
10
50
100
250
Dataset
NS
NS+EW
NS
NS+EW
NS
NS+EW
NS
NS+EW
NS
NS+EW
NS
0.0493
0.2645
0.0237
0.1955
0.0092
0.0378
0.0103
0.0604
0.0085
0.0294
Ours
NS+EW
0.0416
0.2371
0.0221
0.1786
0.0105
0.0484
0.0110
0.0380
0.0089
0.0273
CoDA-NO
0.1279
0.2435
0.0225
0.2282
0.0117
0.0745
0.0115
0.0219
0.0091
0.0148
GINO
0.3337
0.2615
0.3189
0.1817
0.0596
0.0667
0.0349
0.0636
0.0209
0.0308
GNN
0.0265
0.1800
0.0222
0.1799
0.0068
0.0867
0.0113
0.0539
0.0050
0.0193
VIT
0.2738
0.5087
0.1519
0.4146
0.0473
0.1119
0.0407
0.1106
0.0119
0.0381
U-Net
25.33
1.434
4.007
4.320
0.1495
0.6653
0.07723
0.1821
0.0934
0.1651
DeepONet
1.262
0.8186
0.6485
0.4937
0.2576
0.3198
0.1992
0.3399
0.1385 0.1916
achieved through this is called the "Reconstruction error". This allows the models to learn meaningful representations of the underlying physical systems.
20

---

Table 8. Test errors (L2 loss) for fluid dynamics (NS) and fluid-solid interaction (NS+EW) datasets with viscosity μ 5.0 for different numbers few-shot training samples. The pre-training done with 8000 samples taken from NS and NS+EW datasets with viscosities μ E (1.0, 10.0).
# Few Shot Training Samples
Pretrain
Model
5
10
50
100
250
Dataset
NS
NS+EW
NS
NS+EW
NS NS+EW NS
NS+EW
NS
NS+EW
NS
0,0190
0.1597
0.0141
0,0220
0.0043
0.0042
0.0054
0.0053
0.0033
0.0025
Ours
NS+EW
0.0201
0.1077
0.0157
0.0153
0.0053
0.0053
0.0044
0.0030
0.0037
0.0022
CoDA-NO
0.1820
0.0513
0.0107
0.0199
0.0063
0.0066
0.0062
0.0045
0.0041
0.0029
GINO
0.2004
0.1222
0.2245
0.0753 0.0359
0.0364
0.0222
0.0438
0.0163
0.0190
GNN
0.0390
0.0460
0.0280
0.0294
0.0045
0.0123
0.0086
0.0094
0.0064
0.0033
VIT
0.2719
0.2113
0.1889
0,1561
0.0271
0.0474
0.0173
0.0207
0.0077
0.0122
U-Net
13,3370
3,5790
1.1540
2.1340
0.1608
0.3178
0.1418
0.2035
0.1317
0.1180
DeepONet 0.6863
0.4821
0.6720
0.2945 0,2019
0.2024 0.1076
0.1070
0.0731
0.1085
Table 9. Test errors (Tg loss) for fluid dynamies (NS) and fluid-solid interaction (NS+EW) datasets with viscosity 10.0 for different numbers of few-shot training samples. The pre-training is done with 8000
samples taken from NS and NS+EW datasets with viscosities j E (1.0. 10.03.
Few Shot Training Samples
Pretrain Dataset
Model
10
50
100
250
NS
NS+EW
NS NS+EW
NS NS+EW
NS
NS+EW
NS
NS+
NS
0.0186
0.1203
0.0105
0.0207 0.00327 0.00444 0.00391 0.00412 0.00229 0.00215
Ours
NS+EW 0.0171
0.0925
0.0109
0.0130 0.00383 0,00360 0.00303 0.00232
0.00225 0.00133
NO
0.0859
0.0618
0.0115
0.0166
0.00494
0,00763
0.00660
0.00330
0.00374 0.00195
GINO
0,2316
0.1560
0.1679 0.0582
0.04122 0.03327 0.03074 0.0395 0.01389 0.01139
GNN
0.0715
0.0448
0.0547 0.0179
0.00789 0.00494
0.00319
0.0148 0.00547 0.00229
VIT
0.4001
0.2201
0.3388 0.1967
0.06215
0.05700 0.04299
0.01867 0.00770 0.00903
U-Net
1.2550
0.6995
0.3255
0.9148
0.3156
0.2095 0.1889
0.2085
0.1732 0.3132
DeepONet 0.7158
0.3794
0.5515 0.2533 0.1164
0.1337 0.1027
0.07814 0.08359 0.05697
Table 10. Error bar representing standard deviation over thrce runs with different number of few shot example for NS+EW datasct for Re 400
Models
Pre-training
# Few Shot Training Samples
Dataset
5
25
100
GINO
0.121 + 0.023
0.0530 + 0.0053
0.0045
DeepO
0.534 + 0.005
0.1920 = 0.0072
0.1384+0.0293
GNN
0.121 0.136
0.0304 = 0.0210
0.0200+0.0120
ViT
0.276 + 0.093
= 4.0274 0.0284
0.0208+100044
U-net
1,770 + 1.636
0.8368 = 0.3503
0.581410.5680
0.059 + 0.017
0.0096 + 0.0010
0.0038+10.0003
Ours
NS NS-EW
0.068 + 0.055
0.0078 1 0.0002
0.0036:40.0005
0.044 + 0.041
0.0057 + 0.0012
0.003410.0006
Predictive task: Subsequently, we finetune the pre-trained models using supervised learning objective, where the goal is to minimize the "Prediction error" by accurately predicting the next 5 timesteps given the input history. We try to learn the solution operator that maps the state of the system from time & E [0. T' to the state at time t E [T, + 5), effectively predicting the next 5 timesteps given the history up to time T.
Both models were pre-trained and fine-tuned this dataset for 35 epochs.
21

---

Table 11. Error bar representing standard deviation over three with different number of few shot example for NS dataset Re
400
Models
Pre-training
# Few Shot Training Samples
Dataset
5
25
100
GINO
0.143 + 0.0231
0.0371 = 0.0044
0.0330 = 0.0105
DeepO
0.621 0.2417
0.3162 = 0.1146
0.1978 = 0.0345
GNN
0.021 + 0.0124
0.0046 + 0.0012
0.0051 0.0024
ViT
0.196 + 0.0326
0.0409 + 0.0057
0.0302 + 0.0160
U-net
7.241 + 4.3200
0.6568 + 0.3635
0.2025 0.1223
0.0612 + 0.0364
0.0094 0.0006
0.0045
0.0006
Ours
NS
0.0276 + 0.0032
0.0057 + 0.0005
0.0039 + 0.0001
NS-EW
0.0273 + 0.0054
0.0056 + 0.0005
0.0040 0.0001
Table 12. Test errors(L22 error) for CoDA-NO FNO on 2D datasets from PDEBench. SWE indicates shallow
water equations, DIFF indicates diffusion equation. NS+DIFF+SWE means pretraining and fine-tuning on combined Stokes diffusion, and shallow water equations dataset.
Test Error
Model
Dataset
Prediction Error
Reconstruction Error
CoDA-NO
0.04072
0.00460
SWE
FNO
0.04631
0.03262
CoDA-NO
0.00810
0.00041
DIFF
FNO
0.01415
0.01894
CoDA-NO
0.00302
0.00006
NS+DIFF+SWE
FNO
0.00118
0.00287
Table 12 presents CoDA-NO and FNO [12] test errors on the single-physics PDEs datasets sourced from PDEBench [19]. For the single-physics experiments. CoDA-NO consistently outperforms FNO, improving generalization (predictive error) up to 43% over FNO, indicating its ability to capture complex dynamics and dependencies within these systems. We report additional details and experimental results in Sec. G comparing FNO, DPOT [9] and CoDA-NO where CoDA-NO demonstrates superior performance and parameter efficiency. It is worth noting that DPOT and MPP [8] are significantly bigger models but can also handle a larger set of PDEs.
In addition the single-physics experiments, also explore the potential of joint pretraining and finetuning across multiple PDE systems. We create a combined dataset by merging the SWE, DIFF, and NS datasets, even though these PDEs do not share any common physical variables governing equations. Both FNO and CoDA-NO were pre-trained and fine-tuned on this dataset for 35 epochs.
DPOT, with its aarge-scale pre-training approach, demonstrates strong performance on the SWE and DIFF datasets. Even with a 500 million parameter model (DPOT-L-500), DPOT achieves impressive results, reducing the test error to 0.0017 SWE and 0.0073 on DIFF. It is also important note that DPOT is larger as it was pretrained on 12 different datasets, and hence the size is justified. However, it is noteworthy that CoDA-NO, with only 11 million parameters, comes very close to achieving similar generalization performance. CoDA-NO's test error on DIFF (0.0081) I is comparable to DPOT's performance, as shown in table 15, despite having significantly fewer parameters and fewer finetuning epochs. However, on the other hand, we see that CoDA-NO doesn't 't do well on the SWE dataset as shown table 14; we assume that the case would be the fact that we would need to finetune for more epochs to achieve better results. The SWE task is also a harder dataset; increasing the model complexity and pretraining epochs would help get better results.
is important to highlight that DPOT was pre-trained on 12 datasets for 1000 epochs, while CoDA- NO was pre-trained and fine-tuned on single dataset for 35 epochs. Despite this difference in pre- pre-training data and epochs, CoDA-NO still achieves competitive results compared to DPOT's 200/500 epochs of fine-tuning.
These results suggest that when there shared physies between the pre-training and fine-tuning datasets, CoDA-NO can effectively leverage this commonality to achieve strong generalization
22

---

Table 13. Comparison of model parameter sizes for CoDA-NO, FNO, and DPOT. DPOT-FT stands for the Finetuning model used, whereas stands for tiny, stands for small, -M stands medium, and -L stands for Large. The pretrained model sizes are present in the original paper but are around the same parameter sizes as the tune-tund models.
Model
Model Parameters
CoDA-NO
11M
FNO
1.9B
DPOT-FT-T
7M
DPOT-FT-S
30M
DPOT-FT-M
100M
DPOT-FT-L
500M
Table 14. Test errors for CoDA-NO vs DPOT on 2D datasets from PDEBench. SWE indicates shallow water equations data. '12DATA' represents the 12PDE PDE datasets DPOT is pretrained on. The "-200" and "-500" suffixes denote fine-tuning each subset for 200 and 500 epochs, respectively, which is directly taken from DPOT paper. All this was fine-tuned SWE data.
Model
Pretrained Dataset
Predcition Error
CoDA-NO
SWE
0.0407
FNO
SWE
0.0463
T-200
12DATA
0,0028
S-200
12DATA
0,0022
M-200
12DATA
0.0021
L-200
12DATA
0.0019
T-500
12DATA
0.0024
S-500
12DATA
0.0023
M-500
12DATA
0,0022
L-500
12DATA
0.0017
performance. However, when there is no shared physics, as in the case of the combined dataset, CoDA-NO's performance may not be as remarkable.
Table 13 compares the model sizes of CoDA-NO. FNO, and DPOT. CoDA-NO's model size of 11 million parameters is significantly smaller than FNO's 1.9 billion parameters and DPOT's largest model size of 500 million parameters. This highlights CoDA-NO's parameter efficiency and its ability to achieve competitive performance with more compact model.
In summary, these experiments the PDEBench datasets demonstrate the effectiveness of CoDA- NO learning and generalizing to different PDE systems. CoDA-NO's performance, especially considering its smaller model size and shorter pre-training, showcases its potential foundation model for scientific machine learning. The ability to achieve competitive results with DPOT, de- spite the differences in pre-training data and epochs, further highlights CoDA-NO's efficiency and generalization capabilities.
G.1 Ablation on the Size of FNO
In Tab. 16, we present the performance of FNO with different number parameters along with the performance of CoDA-NO.
H Implementation Details
The variable encoders are implemented using multi-layer perceptron mapping the position x E D to the embedding vector. Following transformers and NeRF [45] model, positional encoding
23

---

Table 15. Test for CoDA-NO DPOT 2D datasets from PDEBench. DIFF indicates the diffuusion equation data. '12DATA' represents 12PDE datasets DPOT was trained on. The "-200" and "-500" suffixes denote fine-tuning on each subset for 200 and 500 epochs, respectively, which is directly taken from the DPOT paper. All of this was fine-tuned on DIFF data.
Model
Pretrained Dataset
Test error
CoDA-NO
DIFF
0.0081
FNO
DIFF
0.0141
T-200)
12DATA
0.0194
S-200
12DATA
0.0171
M-200
12DATA
0.0142
L-200
12DATA
0.0158
T-500
12DATA
0.0148
S-500
12DATA
0.0129
M-500
12DATA
0.0103
L-500
12DATA
0.0073
Table 16. Error in L2 norm for models in both the Shallow Water Equation and Diffusion-Reaction experiments.
The number of parameters is reported alongside the Le errors for both tasks.
Model
# Parameter
Le (SWE)
L2 (DIFF)
CoDA-NO
11M
0.0407
0.0081
FNO
1.9B
0.0463
0.0141
FNO
485M
0.0424
0,0145
ENO
120M
0.0410
0.0153
FNO
11M
0.0491
0,0268
FNO
1M
0.2355
0.2085
instead of raw coordinates. The Encoder and Reconstructor modules use three stacked CoDA-NO layers. The Predictor modules use one layer of CoDA-NO.
For every training sample, one of the following two masking choices is selected with equal probability
50% of the mesh points of 60% of variables are masked.
30% of the variables are masked out completely.
In order apply masking this, we identify the neighboring points within a fixed distance from the selected point and set their values to zero. This process is continued until have masked out predetermined portion of all mesh points.
irregular mesh, select point at random from the mesh. Following
24

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All the claims justified either in the theoretical aspect of the paper or in the experiments, We carefully make the distinctions between what is a heuristic versus what is being proposed based the given assumptions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
See line 321.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well- cell-specffication asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken low lighting. Or speech-to-texx system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess..
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't seknowleddge in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
25

---

Answer: [Yes]
Justification:
Yes, the complete and correct proof is provided in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: All the details are either in the main paper the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
26

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
All the details are either in the main paper or the appendix.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
sot.eeeiuuumiiiioooooooocoo for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: wtips.cc/pbblcc//eecccc waabbliccguidess.. CCodeSubmia ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specifyy all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: All the details are either in the main paper or the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer:[Yor Yes]
Justification: Error bars are reported for the main result.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
27

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: Provided in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that NeurlPS Code of Ethics https /ttp///eeerrpppeeeeee/ee cttppubccc/thhcccctteeeeeeeelele
didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the
Answer: [Yes]
Justification:
We follow the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: We study a fundamental framework for the downsampling of signals, which foundation research and has limited negative societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
28

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper has no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
The resources used in the work are cited.
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
29

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [Yes]
Justification: Code will be open-sourced under MIT license.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/mooel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "R1dbnIyhEM", "reviewer": "Reviewer_rAy3", "summary": "The authors propose CoDA-NO, a neural operator architecture that captures interactions across different physical variables of coupled PDE systems. The method involves a generalization of the transformer architecture, including self-attention, positional encodings, and normalization, to function spaces. On two novel datasets for fluid-structure interaction and fluid dynamics, the authors show that their method achieves state-of-the-art performance.", "questions": "- It seems that the performance of the models across the board continue to improve with increase few-shot fine-tuning samples beyond N=100. What does the scaling look like for the proposed model and where does performance saturate?\n- Similarly, the model is evaluated on the in-distribution Re=400 and the out-of-distribution Re=4000 settings, for which the performance of the model is comparable. What does the scaling look like as the task becomes further out-of-distribution (e.g. decreasing velocity)?", "limitations": "The authors address the limitations.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper investigates an interesting problem of how to appropriately capture interactions across different physical variables, that allows for generalization to new codomains.\n- As far as I am aware, the generalization of the Transformer architecture to function spaces is novel.\n- The experimental results, especially the generalization capabilities (from fluid dynamics to fluid-solid interactions) are impressive.\n- Ablation studies on the proposed architectural changes are thorough.", "weaknesses": "Overall, the experiments seem quite compelling. However, it could be illuminating to provide a graphical visualization of the data from Table 1, regarding efficiency of fine-tuning and robustness to out-of-distribution inputs: see questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate that the reviewer values our work and recognizes the main contributions to the novel generalization of the transformer architecture, including self-attention, positional encodings, and normalization to function space, along with the introduction of two new challenging datasets. We further may state that in this work, we not only based our study on the mentioned two PDEs, but we also have SWE and Diffusion Eq. Per the reviewer's request, we are running experiments on the Rayleigh-Benard convection system, with initial results stated in the reviewer qRRk’s reply. \n```\nQ1 Visualization of the Data\n```\nWe thank the reviewer for the visualization suggestion. We provide the visualization of data in the Appendix, Sec B Dataset Description.\n\n```\nQ2 Performance scaling with data (It seems that the performance of the models across the board continues to improve with an increase in few-shot fine-tuning samples beyond N=100. What does the scaling look like for the proposed model, and where does performance saturate?).\n```\nWe conduct experiments with N=250, 500, 1000 for NS+EW with Re=400. We observe that performance continues to improve when training on more data. Therefore, finding the saturation point requires the generation of a much larger dataset, which requires future development and compute allocation for the solver. The precise scaling is problem-dependent.\n\n| Model    | Pre-training    | N = 100  | N=250    | N=500    | N=1000    |\n| -------- | --------        | -------- |----------|----------|-----------|\n| CoDA-NO  | NS              | 0.005    | 0.003    | 0.003    | 0.001     |\n| CoDA-NO  | NS+EW           | 0.003    | 0.003    | 0.002    | 0.001     |\n\n```\nQ3 Similarly, the model is evaluated on the in-distribution Re=400 and the out-of-distribution Re=4000 settings, for which the performance of the model is comparable. What does the scaling look like as the task becomes further out-of-distribution (e.g. decreasing velocity)?\n```\nFor Re=4000, we observe that the performance of the pre-trained models is worse than for Re=400. We hypothesize that the reason may be intuitively twofold, \n1.     Fluid motion at Re=4000 is highly turbulent and \n2.     the data is out-of-distribution. \nInspired by other studies in machine learning, CoDA-NO will perform worse on out-of-distribution data and will require further data for finetuning. However, our evidence suggests that the performance will be better than training from scratch.\n\nAlso, finding the accurate scaling with respect to ODD, i.e., what degree of out-of-distribution the model can handle, requires extensive problem-dependent study", "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": "", "extra_scores": {}}, {"review_id": "sMZlaUfxbQ", "reviewer": "Reviewer_Jm1K", "summary": "This paper introduces Codomain Attention Neural Operator, which tokenizes function along the channel dimension. It allows to learn representations of different PDE systems within a single model. The authors shows that finetuning a pretrained CoDA-NO on different physics yields good accuracy.", "questions": "- Why do we need a VSPE per physical variable? Positional encoding are usually used when there is some sort of order between the tokens?", "limitations": "Overall, I think my main obstacle to provide a better score is the fact that the paper is not very clear due to the introduction of a lot of mathematics not needed to understand the algorithm in practice. These mathematics do not bring any theoretical insight of the algorithm. They are just expressing the architecture in function space. I think this space could have been used to be clearer to explain what is the architecture or the philosophy of the work. Usually, in papers where I see such mathematics, a study of the sample complexity is provided (I know it is almost impossible to do for neural networks).\n\nI am willing to discuss these points with the authors and to modify my score accordingly.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "- I like the problem setting and the idea of the algorithm.\n- The experimental task is quite interesting.\n- The results are convincing, and the fact that the model can generalize to higher Reynold numbers seen during training is promising.\n- I liked the used of GNO for handling non-uniform grids.\n- The code seems solid.", "weaknesses": "- To me, the main weakness of the paper is that the presentation lacks of clarity. I don't see the point of doing 3 pages of mathematics in function space if, in practice, everything is done in discrete space. I think this blurs the message of the paper and it is difficult for the reader to understand what is the relevant information for understanding the actual CoDA-NO algorithm. In my opinion, these mathematics are not essential to the algorithm and could be put in appendix. I can always express a neural network architecture in function space, but since in practice we are working on discretized space, it is never done in experimental deep learning papers. Moreover, no discussion on how to go from infinite-dimensional space to discretized space is given by the authors.\nThis space could be used to have the actual detailed architecture. I may have missed the point on the usefulness of these sections and am willing to understand the point of view of the authors regarding this. \n- I don't fully understand the CoDA-NO algorithm and I think a Figure showing the whole architecture would have clarified this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate that the reviewer values our work and recognizes that this work presented a method for learning representations of different PDE systems within a single model.\n\n```\nQ1. Formulation of the Model in the Function Space and Clarity of the Paper\n```\nWe politely disagree with the reviewer. Indeed, in practice, we typically work with discretized versions of functional data. However, this does not imply that we only deal with a \"fixed\" finite-dimensional space for which neural networks are designed. In practice, the discretizations often vary from one sample to another (i.e., the location and number of evaluation points). Many operations that work well for \"fixed\" finite-dimensional data do not behave consistently across resolutions. For example, regular CNNs (or GNNs with nearest-neighbor graph) collapse to pointwise operations when the resolution is refined (see [1]) and is thus not well suited for data that comes at different resolutions. \n\nSuch limitations of neural networks motivated the paradigm of neural operators as a generalization of neural networks to data originating from functions [2]. Designing architecture directly in the function space allows us to work across different discretizations seamlessly. It also allows us to draw parallels between various approaches to operator learning and ensure specific properties of the learning algorithms, such as the universal approximation property demonstrated in [3] for transformer architectures in function spaces. The mathematical foundation is often helpful in conveying the core ideas of the architecture design in operator learning. For example, here is reviewer ZkA5’s statement about our mathematical foundation: “The mathematical formulation of the proposed model is well-articulated in the paper. This clarity helps readers understand the underlying principles of the model's operation.”\n\nTo better explain the architecture, following the suggestion, we provide a detailed figure of the whole architecture (see Figure 2 in the rebuttal pdf).\n\n[1] Neural Operators with Localized Integral and Differential Kernels\n\n[2] Neural Operator: Learning Maps Between Function Spaces With Applications to PDEs\n\n[3] Continuum Attention for Neural Operators.\n\n```\nQ2. From infinite-dimensional space to discretized space\n```\nOur implementations are based on previous works on Neural Operators [1,2,3]. These studies discuss how operations on functional spaces are approximated on discretized data, e.g., using Riemann sum approximation of the integral operator, which scales appropriately with the resolution. We will clarify this concept in the paper.\n\n[1] Fourier Neural Operator for Parametric Partial Differential Equations (see Section 4).\n\n[2] Neural Operator: Graph Kernel Network for Partial Differential Equations (see Section 3).\n\n[3] Neural Operator: Learning Maps Between Function Spaces With Applications to PDEs.\n\n```\nQ3. Figure showing the whole architecture.\n```\nThanks for the suggestion. We added a detailed figure (Figure 2) in the rebuttal pdf.\n\n```\nQ4. Use of VSPE per physical variable\n```\nIn this work, we treat each of the input function's variables (codomains) as a token for the Codomain attention. The model needs a way to identify each variable (e.g., temperature, velocity, pressure). In LLM, it is done through position ID. In CoDANO, we propose the Variable-Specific Positional Encoding (VSPE), which is learned separately for each variable and then concatenated with the corresponding variables. It allows the model to differentiate between the various PDE variables and their interactions. The following table demonstrates the benefit of using VSPE for fine-tuning on $5$ examples for $\\mu = 5.0$ (from Appendix D1, Table 3). Please refer to Appendix D1, Table 3 for a detailed analysis.\n\n| VSPE         | Pretrain Dataset | NS    | NS+EW |\n|--------------|------------------|-------|-------|\n| x            | NS               | 0.049 | 0.079 |\n| x            | NS EW            | 0.045 | 0.057 |\n| $\\checkmark$ | NS               | 0.025 | 0.071 |\n| $\\checkmark$ | NS EW            | 0.024 | 0.040 |\n\n```\nQ5. On the clarity and organization of the paper\n```\nWe have addressed this in a previous comment, noting that developing architecture directly in the function spaces is crucial, conveying the philosophy of the work. We also agree with the reviewer that this should not overshadow the practical implementation aspect of the work. Our philosophy and goals are outlined in the introduction (lines 28-41) and in the Problem Statement (see Sec 4).\n\nWe provide additional detailed figures (see the attached PDF), and we move some background details to the Appendix, as suggested by the reviewer.\n\nTo this end, we would appreciate it if the reviewer reconsidered our explanation and problem setup of operator learning in this paper. The field of neural operators is still in its early stages, and similar to the early stages of neural networks, explaining the problem setup and detail of the architecture may still be considered crucial and insightful.", "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": "", "extra_scores": {}}, {"review_id": "cHyOspUKlZ", "reviewer": "Reviewer_qRRk", "summary": "This paper presents a new operator learning method for solving multiphysics PDEs. The attention scheme is designed on channel space to capture multiple physical variables, which is called co-domain. Moreover, positional encoding and normalization layers are considered. Such a strategy enables self-supervised pretraining of PDE systems. The experiments have shown the effectiveness of the proposed method.", "questions": "- This paper considers the generalization to different Reynolds numbers. Is it possible to generalize to different physical parameters of elastic wave equations, such as the object size or the solid density \\rho^s?\n\n- Lines 100-101, this paper claims that it considers diverse physical systems in terms of input functions, geometries, and Reynolds numbers. I would say it’s just different PDE scenarios within one PDE type (NS and NS+EW). It seems unrelated to the diversity of PDE systems, such as reaction-diffusion, convection systems, etc.", "limitations": "Please see my concerns in **Weaknesses** and **Questions**.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The proposed idea is interesting. It enables the generalization to coupled physical systems, which is of interest to the scientific machine learning (SciML) community. Also, self-supervised pretraining is one emerging tool in SciML and will gain a lot of attention in the future. \n\n- This paper provides the experiments on a Navier-Stokes equation and its coupled version with the elastic wave equation.\n\n- This paper is well-organized and well-written. The details are easy to follow.", "weaknesses": "- This paper only considers one coupled system, i.e., NS and NS+EW. It may not validate the general applicability of the proposed method. The motivation of using this case should be enhanced. Also, considering some other PDE systems might strengthen the paper, such as the Rayleigh-Benard convection system. It is also a coupled system with NS + temperature.  \n\n- The motivation for the combination of positional encoding, self-attention, and normalization layers seems to be better clarified. Although those parts are modular (claimed in Line 68), the connections between each other are also important. \n\n- In Appendix B.1, it would be good to include more details of self-supervised pretraining, such as masked ratio.  \nThe evaluation metrics might not be sufficient. This paper only considers L2 errors. There are many papers considering relative l2 error [1]. For the turbulence data, researchers also care about the infinity norm a lot. It would be better to add more evaluation metrics in this paper.\n\n**References:** \n\n[1] Hao, Zhongkai, et al. \"Dpot: Auto-regressive denoising operator transformer for large-scale pde pre-training.\" arXiv preprint arXiv:2403.03542 (2024).\n\n[2] Ren, Pu, et al. \"Superbench: A super-resolution benchmark dataset for scientific machine learning.\" arXiv preprint arXiv:2306.14070 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate that the reviewer values our work and recognizes the fact that this work is of interest to the scientific machine learning (SciML) community, presents a strategy that enables self-supervised pre-training of PDE systems, and states the importance of our experiments that have shown the effectiveness of the proposed method. \n\n```\nQ1. Additional Coupled PDE\n```\nWe discuss the difficulty of modeling the fluid-solid interaction problem and justify our problem design in Appendix B1. Given the complex geometry and turbulent flow, the problem can be considered a challenging benchmark problem.\n\nAdditionally, we provide experiments on diverse PDEs (Navier–Stokes equations, Elastic wave equations, Shallow water equations, and Diffusion reaction; please see Table 2 and Appendix Sec E, Tables 11-12) and demonstrate the architecture's performance and flexibility.\n\nAs suggested, we also provide results on the Rayleigh-Benard convection system. The CoDA-NO is pre-trained in a supervised fashion on the Navier–Stokes (NS) equations. Then it is subsequently fine-tuned on a limited few shot data ($N \\in \\{5,10,25\\}$) of the Rayleigh-Benard convection system. We show the result in the following table, where we report the $L2$ loss. We see that, in the low data regime, pre-trained CoDA-NO performs better than FNO and CoDA-NO trained from scratch. Please note that these results are primitive and yet to be finalized. Per the reviewer’s suggestion, we are running a comprehensive study on the Rayleigh-Benard convection system.\n\n| Model    | Pre-training dataset | N=5      | N=10    | N=25    |\n| -------- | --------                           | -------- |-----        |-----        |\n| FNO      | -                                | 0.439    |0.249    |0.130    |\n| CoDA-NO  | -                           | 0.315    |0.272    |0.203    |\n| CoDA-NO  | NS                      | 0.2537   |0.2223   |0.179    |\n\n\n```\nQ2.  Motivation for the components of the CoDA-NO (The motivation for the combination of positional encoding, self-attention, and normalization layers seems to be better clarified. Although those parts are modular (claimed in Line 68), the connections between each other are also important.)\n```\n\nWe outline our motivation in the problem statement (Section 4) and explain the need for Neural Operators that can handle PDEs with varying numbers of variables (lines 167-177). Given the well-established nature of the Transformer architecture and its property to handle variable length input, it has been a natural choice for our approach. We have redesigned key components of modern transformer layers, including positional encoding, self-attention, and normalization layers, to develop codomain attention in function space. We will further emphasize this in the main text.\n\n```\nQ3. Details of self-supervised pre-training, such as masked ratio\n```\nWe provide the masking ratio and some additional details in Appendix F. For every training sample, one of the following two masking choices is selected with equal probability \n1.  50\\% of the mesh points of 60\\% of variables are masked.\n2.  30\\% of the variables are masked out completely.\n\n\n```\nQ4. Additional Evaluation Metrics\n```\nL2 is the generalization of RMSE to function spaces and is considered a natural metric of evaluation. In this work, we also provide an Energy Spectrum representing the correct statistical distributions of the fluids to go beyond distance-based evaluation metrics (please see Appendix D.4). L2 norms, together with the Energy Spectrum, provide considerable evidence establishing the proposed model's superiority to the baseline. In addition to these two metrics, we also include L1 and relative L2 in our list of metrics. To this end, as the reviewer knows, any evaluation metric comes with its advantages as well as limitations. We hope that the current list plays a sufficient role in establishing evident results on the superior performance of the proposed method. \n\n### Table: L1 and relative L2 Error\n\nResults on the fluid-solid interaction dataset combining Navier-Stokes and Elastic wave equation **(NS-EW dataset)**.\n\n| Models | Pre-training Dataset | # Train = 5 (L1/Rel-L2)| # Train=25 (L1/Rel-L2)| # Train=100 (L1/Rel-L2)|\n|--------|----------------------|------------------------|-----------------------|------------------------|\n| GINO   |                      | 0.185/0.296          | 0.151/0.221          | 0.160/0.219           |\n| DeepO  |                      | 0.453/0.687           | 0.266/0.431         | 0.184/0.325          |\n| GNN    |                      | 0.083/0.130        | 0.056/0.082       | 0.059/0.082        |\n| ViT    |                      | 0.202/0.366          | 0.156/0.276         | 0.076/0.124         |\n| U-net  |                      | 0.793/1.186           | 0.284/0.463         | 0.174/0.291           |\n| Ours   |                      | 0.092/0.164         | 0.046/0.092        | 0.032/0.058          |\n| Ours   | NS                   | 0.074/0.141         | 0.032/0.072       | 0.030/0.059        |\n| Ours   | NS-EW                | 0.066/0.128         | 0.040/0.077       | 0.033/0.057         |", "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": "", "extra_scores": {}}, {"review_id": "we9wQsIB5Q", "reviewer": "Reviewer_ZkA5", "summary": "The authors introduced an innovative attention-based neural operator and evaluated it against various baselines. They employed masked pretraining and finetuning techniques, comparing the model's performance to multiple benchmarks. Their study included interesting problems such as fluid-structure interactions. The authors showed that their approach is effective for few-shot learning in their experimental evaluations.", "questions": "- What is the reasoning behind applying the attention in the in the channel space? Do the models scale better? Do they have higher expresive power?\n- Why is 1.9B FNO model used for comparisson?", "limitations": "The authors explained the limitations.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper presents a new neural operator architecture based on attention mechanisms. This architecture demonstrates superior performance compared to tested baseline models on the NS and NS+EW benchmarks, highlighting its potential advancements in solving PDE-related problems.\n- To the best of my knowledge, the authors are the first ones to use masked training for PDE learning effectively\n- The mathematical formulation of the proposed model is well-articulated in the paper. This clarity helps readers understand the underlying principles of the model's operation.\n- The study addresses a compelling and very relevant multiphysics problem involving fluid-structure interactions.\n- The authors demonstrated through empirical evidence that their approach is effective for few-shot finetuning in various scenarios.", "weaknesses": "- The study in Table 1 demonstrates the model’s performance on two specific PDEs: Navier-Stokes for fluid flow and a coupled Navier-Stokes with elastodynamics equations for fluid-solid interactions. While these cases provide some insight into the model's capabilities, they are not sufficient to generalize the model's applicability to a broader range of multiphysics problems.\n- For the NS dataset with Reynolds number Re=400, the model trained from scratch with only 25 samples matches the performance of the pretrained model. In the case of NS+EW benchmark, when the Reynolds number increases to 4000, even with just 5 samples, both the finetuned and scratch-trained models exhibit similar testing errors. This suggests that pretraining may not provide significant advantages in many cases.\n- The use of the L2 loss metric to evaluate model performance is problematic because it aggregates outputs of different physical meanings, such as pressure p, velocity u, and displacement d, into a single loss value. This can obscure individual variable contributions and lead to misleading conclusions about model accuracy.\n- The absence of prediction visualizations diminishes the interpretability of the L2 loss values. Visualizing predictions could provide more intuitive insights into model performance and clarify discrepancies in the loss metric.\n- The study in the Table 1 does not include a Fourier Neural Operator. Including such a benchmark is crucial to fairly evaluate CoDA-NO’s performance against an FNO model of similar size.\n- The FNO model used for comparison in Table 2 has 1.9 billion parameters, vastly outnumbering the CoDA-NO's 11 million parameters. This overparameterization likely affects the model's performance due to the relatively small training set sizes and makes the comparison with CoDA-NO’s performance misleading. Smaller FNO models could provide a more realistic performance benchmark. The claim that CoDA-NO’s better performance compared to a much larger FNO model demonstrates parameter efficiency is misleading. Parameter efficiency should be evaluated with models of comparable sizes, and overparameterized models may not reflect typical scenarios.\n- CoDA-NO has significantly higher inference times compared to other baseline models. \n\nThese points collectively highlight the need for more comprehensive experiments, appropriate metrics, realistic model comparisons, and practical considerations like inference time to fully evaluate the model's capabilities.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate that the reviewer values our work and recognizes its novelty and the mathematical formulation of the proposed model.\n```\nQ1. Evaluation on additional PDE datasets\n```\nIn addition to the coupled fluid-solid interaction problem, we also provide experiments on different PDE systems (please see Table 2 and Appendix Sec E, Table 11-12), where we test the proposed CoDA-NO architecture on shallow water equation, diffusion equation, and Navier Stokes equation. \nIn addition, we also include the results of our study on the Rayleigh-Benard convection system. The CoDA-NO is pre-trained in a supervised way on the Navier–Stokes (NS) equations. Then it is subsequently finetuned on a limited few shot data ($N \\in \\{5,10,25\\}$) of Rayleigh-Benard convection. We show the result in the following table, where we report the $L2$ loss. We see that, in the extremely low data regime, pre-trained CoDA-NO performs better than FNO and CoDA-NO trained from scratch. \n| Model    | Pre-training dataset | N=5      | N=10    | N=25    |\n| -------- | --------                           | -------- |-----        |-----        |\n| FNO      | -                                | 0.439    |0.249    |0.130    |\n| CoDA-NO  | -                           | 0.315    |0.272    |0.203    |\n| CoDA-NO  | NS                      | 0.2537   |0.2223   |0.179    |\n\nPlease note that these results are preliminary. As per reviewer qRRk's suggestion, we are conducting a comprehensive study on the Rayleigh-Benard convection system.\n\n```\nQ2. Performance gain from pre-training.\n```\nIn most cases, pre-trained CoDA-NO demonstrated a significant performance improvement compared to CoDA-NO trained from scratch. While we acknowledge that in some instances, the difference is not substantial, it is consistently better. Additionally, in general, pre-trained models offer benefits such as reusability and faster convergence with fewer epochs required. The lack of a significant performance gain in a few settings does not diminish the importance and potential advantages of adapting the architecture and pre-training scheme. Further, our self-supervised pretraining method is also beneficial when the cost of generating training data with numerical solvers is very high.\n\n```\nQ3. Per variable loss\n```\n\nPrevious studies have shown that aggregated loss demonstrates model accuracy across multiple variables [1, 2, 3] which helps to get a comprehensive overview of the performamnce. For instance, the results of the compressible NS equation (velocity and pressure field) are presented in an aggregated manner [1, 2, 3]. \nWe also provide the per-variable performance of each of the variables (velocity $v_x, v_y$, pressure $p$, displacement $d_x,d_y$) on the combined Navier-Stokes and Elastic wave equation for $Re=4000$.\n### Table of Horizontal velocity ($u_x$)\n| Model | Pretrain Dataset |  N=5  |  N=25 | N=100 |\n|-------|:----------------:|:-----: |:-----:|:-----:|\n| GINO  |         -        | 0.416 | 0.145 | 0.105 |\n| DeepO |         -       | 0.513 | 0.597 | 0.189 |\n| GNN   |         -        | 0.056 | 0.186 | 0.119 |\n| ViT   |         -          | 0.393 | 0.710 | 0.437 |\n| U-Net |         -        | 2.936 | 0.227 | 0.174 |\n| CoDA-NO  |         -        | 0.078 | 0.048 | 0.023 |\n|  CoDA-NO   |        NS    | 0.061 | 0.034 | 0.016 |\n|   CoDA-NO |NS+EW      | 0.047 | 0.020 | 0.010 |\n\n### Table of Verticle velocity ($u_y$)\n| Models  | Pretrain  Dataset |  N=5  |  N=25 |  N=100 |\n|---------|:---------------------:|:-----:|:-----:|:------:|\n| GINO    |           -           | 0.503 | 0.403 |  0.111 |\n| DeepO   |           -           | 1.006 | 0.865 |  0.688 |\n| GNN     |           -           | 0.462 | 0.270 |  0.169 |\n| ViT     |           -           | 1.334 | 0.726 | 0.4254 |\n| U-Net   |           -           | 0.967 | 0.417 |  0.269 |\n| CoDA-NO |           -           | 0.240 | 0.153 | 0.0714 |\n| CoDA-NO |           NS          | 0.233 | 0.132 |  0.067 |\n| CoDA-NO |         NS+EW         | 0.208 | 0.120 |  0.060 |\n\n### Table of Verticle velocity ($p$)\n| Models  | Pre-train dataset |  N =5 |  N =25 | N =100 |\n|---------|:-----------------:|:-----:|:------:|:------:|\n| GINO    |         -         | 0.478 |  0.152 |  0.110 |\n| DeepO   |         -         | 1.435 |  0.337 |  0.229 |\n| GNN     |         -         | 0.209 |  0.149 |  0.110 |\n| ViT     |         -         | 0.560 |  0.257 |  0.148 |\n| U-Net   |         -         | 0.581 |  0.337 |  0.229 |\n| CoDA-NO |         -         | 0.174 | 0.2251 |  0.077 |\n| CoDA-NO |         NS        | 0.526 | 0.1326 |  0.064 |\n| CoDA-NO |       NS+EW       | 0.398 |  0.097 | 0.0591 |\n\n### Table of Horizontal displacement($d_x$)\n| Models  | Pre-training Dataset |  N=5  |  N=25 | NS+EW |\n|---------|:--------------------:|:-----:|:-----:|:-----:|\n| GINO    |           -          | 1.326 | 0.353 | 0.117 |\n| DeepO   |           -          | 1.467 | 0.354 | 0.216 |\n| GNN     |           -          | 0.654 | 0.466 | 0.240 |\n| ViT     |           -          | 0.769 | 0.474 | 0.279 |\n| U-Net   |           -          | 1.083 | 0.730 | 0.365 |\n| CoDA-NO |           -          | 0.582 | 0.670 | 0.139 |\n| CoDA-NO |          NS          | 0.631 | 0.372 | 0.156 |\n| CoDA-NO |         NS+EW        | 0.531 | 0.350 | 0.147 |", "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": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their valuable comments. We appreciate reviewers for recognizing the presentation of a novel neural operator architecture that “captures interactions across different physical variables of coupled PDE systems and the fact that the method involves a generalization of the transformer architecture, including self-attention, positional encodings, and normalization, to function spaces. “ It has been recognized by reviewers that this study is of interest to the scientific machine learning (SciML) community.

This paper also introduces two challenging datasets based on NS and NS+EW, which are coupled PDEs that resemble problems in weather forecasting and climate modeling. 

Based on the reviewer's comments, we have the impression that reviewers may have interpreted that we only study CoDA-NO on the two mentioned datasets. We would like to state that we study CoDA-NO on four PDEs: NS, NS+EW, SWE, and Diffusion equations. Accordingly, we edited the paper's presentation, ensuring that the explanations and structure imply that we study more than two PDEs.

In addition, per the reviewer qRRk’s suggestion, we started to incorporate another PDE system, i.e., the Rayleigh-Benard convection system. Our first runs show positive results, as presented below, showing an **average improvement of 20% compared** to the Fourier neural operator (see the response to reviewer qRRk). We will add the finalized comprehensive study to the main draft.

Meta-review: The paper introduces a neural operator model for solving systems of coupled differential equations, capturing the interaction between the physical variables. It leverages several novel extensions of classical transformer components. The authors evaluate their model on representative datasets and also introduce two novel datasets developed specifically for the paper.

The reviewers agree on the importance of the problem, the novelty of the proposed model, and the significance of the experiments. During the rebuttal, the authors presented complementary experiments, including one on a new dataset of coupled equations. Three out of four reviewers increased their scores after the rebuttal discussion. The authors are encouraged to follow the reviewers' suggestions to improve the clarity of the presentation.
Decision: Accept (poster)