Title: Pretraining Codomain Attention Neural Operators for Solving Multiphysics PDEs
Venue: NeurIPS 2024
Accepted: yes
Keywords: Neural Operator
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

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% .

1 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 on fine grids to capture the physical phenomena accurately. Consequently, these computational requirements often exceed available memory and computational budgets for realworld 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 are 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, we can obtain data by solving individual PDEs. While the solutions of the two systems can be very different, they share common features and can benefit from a combined learning framework. Can we design a systematic curriculum learning scheme for learning multiphysics systems?

CoDA-NO

Self-Supervised

Pre-training

Few Shot Supervised Finetuning

Coupled Multi-Physics System

New variables

Mask

u

x,t

+

δt

u

y,t

+

δt

p

t

+

δt

d

x,t

+

δt

d

y,t

+

δt

u

x,t

u

y,t

p

t

d

x,t

d

y,t

u

x

u

y

p

u

x

u

y

p

CoDA-NO

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 a foundation model for solving PDEs [7-9]. However, these methods only work on predetermined PDEs with a fixed number of variables, and none of them consider multi-physics PDEs, and they are mostly restricted to uniform grids, limiting their applicability. For example, standard patching-based approaches used in Vision Transformers (ViTs) [10] often struggle with discontinuities in predicted functions and changing resolutions [11]. Since they are limited to 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 as a 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 computational domain). This property is known as discretization convergence , making them agnostic to the discretization of 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, 9] use neural operators, they still cannot handle multiphysics or coupled PDEs. They also cannot adapt to new variables that are not predetermined at the beginning of training.

Our Approach: Wepropose 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 a 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 finite-dimensional vector space to an infinite-dimensional 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 a function, capturing cross-function dynamics through attention mechanisms while maintaining discretization convergence . This design empowers the architecture to handle functions discretized on grids of varying resolutions. Specifically, each token function is subjected to the following operations: (i) concatenation with a learned positional embedding, (ii) lifting to a 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 a collection of single-physics problems.

⊕

CoDA-NO Block

GNO

K

Q

V

Token

Functions

k

1

q

1

v

1

k

2

q

2

v

2

〈

q

1

, k

1

q

1

, k

2

〉

〈

q

2

, k

1

〉

〈

q

2

, k

2

〉

×

GNO

Variable 1

Variable 2

VSPE 1

VSPE 2

⊕

GNO

Output

Functions

I

per

Variable 1

Variable 2

Input Function

Target Function

Latent Function

Key

Query

Value

×	#

(a) Illustration of CoDA-NO architecture . Each of the input physical variables, 'variable 1' and 'variable 2', depicted with two different colors (yellow and blue), incorporate a learnable variable-specific positional encoder (VSPE). These variables, along with the corresponding VSPE, are passed through GNO layers to transform from non-uniform to latent uniform grids. Codomain attention tokenizes the latent functions along the codomain. Each token undergoes transformations with K , Q , and V operators yielding key, query, and value functions { k 1 , k 2 } , { q 1 , q 2 } , and { v 1 , v 2 } . The resulting function is computed using a self-attention mechanism in function space followed by an integral operator I per . Finally, the output function on the target geometry is generated by passing through stacked CoDA-NO blocks, followed by an additional GNO layer.

u

x

u

y

p

VSPE

1

VSPE

3

⊕

⊕

⊕

VSPE

1

Pre-training on Physical Quantities

of  Navier-Stokes Equations

Encoder

Reconstructor

Masking

Masking

Masking

u

x

u

y

p

VSPE

1

VSPE

3

VSPE

2

⊕

⊕

⊕

VSPE

5

VSPE

4

⊕

⊕

Fine-tuning to additional Physical

Quantities of Elastic Wave Equation

Encoder

Predictor

u

x,t

u

y,t

p

t

d

x,t

d

y,t

u

x,t

+

δt

u

y,t

+

δt

p

t

+

δt

d

x,t

+

δt

d

y,t

+

δt

(b) Self-supervised pre-training and fine-tuning with CoDA-NO. Model, pre-trained on the Navier-Stokes equations dataset (with u x , u y , and p ) in a self-supervised way, can be fine-tuned to a fluid-solid interaction dataset (new d x and d y variables) by only including two extra VSPEs and a predictor module.

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 an additional challenge of irregular meshes on a complex geometry.

Instead of directly learning to solve the full multiphysics problem, we start with a curriculum where we first learn the basic fluid dynamics without the elastic wave equation, governed by the incompressible Navier-Stokes equation, with velocity and pressure as variables. We pre-train CoDANO in a self-supervised manner on snapshots of fluid flows by masking different parts of the velocity or pressure fields. Using few-shot supervised fine-tuning, we show that our model can adapt to unseen viscosities and additional displacement fields given by the elastic wave equation. We use graph neural operator (GNO) layers [18] as encoders and decoders to handle time-varying irregular meshes of the fluid-structure interaction problems. For the few-shot learning problem, our model achieves 36 . 8% lower errors on average compared to the best-performing baseline trained from scratch on the target problem.

The second system involves Rayleigh-Bénard convection, where the Navier-Stokes and heat (energy) equations are coupled in a regular 2 D 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 a factor of two.

We also train CoDA-NO on a 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 a co-domain attention neural operator that efficiently learns solution operators to PDEs by formulating transformer operations in function space and ensuring discretization convergence.

The proposed architecture enables self-supervised 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.

2 Related Works

Transformers for solving PDEs. Recent work [20] proposes a method to weight variables/codomains of the input function based on the weights calculated from the PDE parameters. Another study [21] proposes a scalable transformer architecture by combining a projection operator to a one-dimensional domain and a learnable factorized kernel. In contrast to these works, CoDA-NO provides a complete attention operator by considering each physical variable as a token function, i.e., an infinite-dimensional vector, extending traditional transformers that only operate on finite-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 a self-supervised fashion have also gained traction in the domain of scientific computing. One recent study [8] proposes pretraining the models with autoregressive tasks on a 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-reconstruction [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-agnostic.

3 Method

Let us first define our setting and provide a brief introduction to neural operators. For further details, we refer to Sec. A in the appendix.

For an input function a : D → R d in , we will denote the d in -dimensional output space R d in 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 = [ a 1 , . . . , a d in ] with a i : D → R . The same applies to the output function u : D → R d out . We define the action of a pointwise operator H : { f : D → R d f } → { g : D → R d g } given by a function h θ : R d f → R d g with parameters θ as

Moreover, we define an integral operator T : { f : D → R d f } → { g : D → R d g } given by a kernel function k ϕ with parameters ϕ as

Problem Statement. Our objective is to construct a 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 → R d in and ˜ a : D → R ˜ d in of two different PDE with corresponding output functions u : D → R d out and ˜ u : D → R ˜ d out . In general, the functions a and ˜ a

̸

represent d in and ˜ d in physical variables over the domain D with d in = ˜ d in . We aim to design neural operator architectures G that can both be applied to a as well as ˜ a despite the different codomains of the input as well as output functions.

̸

Such property provides the possibility to evaluate or finetune the operator on PDEs with different numbers of variables than those on which it was trained. In particular, when the PDE systems have overlapping physical variables { a i } d in i =1 ∩ { ˜ a i } ˜ d in i =1 = ∅ , 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 d in functions { a 1 , a 2 , . . . , a d in } that represents different physical variables of the PDE, we seek to construct operators that act on sets of input functions with different cardinalities.

For an efficient implementation of operators on sets of functions, we mimic transformer architectures and share weights across different variables. Specifically, we can define the integral operator I per as

where a = [ a 1 , . . . , a d in ] and I is a regular integral operator as described in Eq. (2). Such construction makes the operator permutation-equivariant with respect to the order of the variables in the set. Following the same mechanism, we can also define permutation-equivariant pointwise operator H per with a shared pointwise operator H (see Eq. (1)). We will use FNO per and GNO per to denote permutation-equivariant operators using a shared GNO and FNO, respectively.

CoDA-NO Layer. To explain the CoDA-NO layer, let us assume the input function a has been processed into a latent function w : D → R d . We partition the function into a set of so-called token functions w j : D → R d ′ with w j ∈ W for j ∈ { 1 , . . . T } along the codomain, such that w = [ w 1 , . . . w T ] (and where each w j is associated with precisely one of the physical input variables). That is, w represents the codomain-wise concatenation of the token functions w j and d ′ = d T . 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).

Let us begin by introducing a single-head CoDA-NO layer. Later, we 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 j : D → R d ′ to key, query, and value functions. We define the key, query, and value operators as

Assuming d k = d q , we denote by k j = K [ w j ] , q j = Q [ w j ] , and v j = V [ w j ] the key, query, and value functions of the token functions, respectively.

Next, we calculate the output (token) functions o j : D → R d v as

where τ is the temperature hyperparameter. Here, ⟨ ., . ⟩ denotes a suitable dot product in the function space. We take the L 2 ( D , R d k ) -dot product given by ⟨ q j , k m ⟩ = ∫ D ⟨ q j ( x ) , k m ( x ) ⟩ d x, 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 ∈ { 1 , . . . H } using K h , Q h , and V h to obtain o j,h . We then concatenate these outputs o j,h along the codomain and get c j := [ o j, 1 , . . . o j,H ] . Finally, we use an operator

to get the output function o j .

We obtain the output of the attention mechanism by concatenating o j s as o = [ o 1 , o 2 , . . . o T ] . Finally, we complete the CoDA-NO layer by applying a permutation-equivariant integral operator I per on o . When CoDA-NO is acting on functions sampled on a uniform grid, the internal operators K h , Q h , V h , M , and I are implemented as FNOs.

Function Space Normalization. Normalization is a vital aspect of deep learning architectures. However, when it comes to neural operators mapping infinite-dimensional functions, this topic remains largely unexplored. We now provide a natural extension. Given a function w , let w j : D → R d ′ be a token. Then we calculate the mean µ ∈ R d ′ and standard deviation σ ∈ R d ′ for this token as

Here, ◦ r denotes the elementwise (Hadamard) r th -power. The normalization operator can be written as

Here b ∈ R d ′ and g ∈ R d ′ are learnable bias and gain vectors and ⊘ and ⊙ denote elementwise division and multiplication operation. This normalization can be seen as an extension of instance normalization [30] for function spaces. Similarly, normalization variants, such as group norm , layer norm , and batch norm , extend to operator learning with these definitions of statistics [31-33].

Variable Specific Positional Encoding (VSPE). We learn positional encoders e i : D → R d en for each physical variable i ∈ { 1 , . . . , d in } , for the given vector-valued input function a = [ a 1 , . . . , a d in ] . We concatenate each positional encoding e i with the respective variable a i : D → R along the codomain to obtain extended input functions ¯ a i = [ a i , e i ] . Next, we apply a shared pointwise lifting operator P : { ¯ a i : D → R d en +1 } → { ¯ w i : D → R D } , typically with D > d en + 1 . Finally, we concatenate ¯ w i , i ∈ { 1 , . . . d in } , to get the lifted latent function

In the previous paragraphs, we used d = D · d in and, to maintain the permutationequivariance property of the operator, d ′ must divide D .

Algorithm 1 presents the pseudocode for the CoDA-NO architecture applied to input functions [ a 1 , a 2 ] , mapping two different physical variables on a uniform grid in a 1D domain, to the solution functions [ u 1 , u 2 ] . It assumes d ′ = D while designing the CoDA-NO layer. Notably, to incorporate another function a 3 , representing a new physical variable, it is only necessary to introduce a corresponding parameter for the new VSPE, denoted as κ 3 .

To effectively handle non-uniform complex geometries, we follow the GINO architecture

Algorithm 1 Adaptation of CoDA-NO from two physical variables a 1 , a 2 to a new variable a 3 on uniform 1 D grid. Only the parts in 'Blue' are additionally required to incorporate the new variable a 3 .

[34], where a GNO is used as an encoding and decoding module. Given a set of evaluations of an input function a on a mesh, as represented by { a ( x in i ) } n i =1 , where { x in i } n i =1 ⊂ D in , our first step involves concatenation of each physical variables with respective VSPEs (see Fig. 2a).

Next, we use GNO per to transform the function a into a new function w 0 on a uniform latent grid, represented by { x grid i } n ′ i =1 . Finally, we apply l stacked CoDA-NO layers to w 0 to obtain the encoded function w l , which acts as a representation of the input function a .

The decoding module is essentially a mirrored version of the encoding module. It starts by applying another block of l stacked CoDA-NO layers to the encoded function w l to obtain w L . Subsequently, it uses another GNO per operator to transform w L on a uniform grid to an approximation u of the solution function on an arbitrary output grid { u ( x out i ) } n ′ i =1 . The architecture is summarized in Fig. 2a.

Model Training. To seamlessly adapt to multi-physics PDEs with limited data, we propose a twostage training process: Self-supervised pretraining is followed by a supervised fine-tuning stage. For a summary, we refer to Fig. 2b.

Pre-training. In the context of self-supervised pretraining, the objective is to train the model to reconstruct the original input function from its masked version. Within this phase, the model's encoding component is denoted as the Encoder , while the decoding component comprises the Reconstructor . The values of the input function at a specific percentage of mesh points are randomly masked to zero, and certain variables (channels/co-domains) 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 a randomly initialized Predictor module. The parameters of the Encoder and VSPEs are copied from pre-trained 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 a diverse set of PDEs from PDEBench [19]. The code is available at https://github.com/neuraloperator/CoDA-NO .

Modeling Fluid-Structure Interaction. We consider the following problems: (a) a fluid dynamics problem, where a Newtonian, incompressible fluid impinges on a rigid object, and (b) a fluid-structure interaction problem between a Newtonian, incompressible fluid and an elastic, compressible solid object [17]. We denote Ω f t (resp. Ω s t ) as the domain occupied by the fluid (resp. the solid) at time t . The dynamics of the fluid are governed by the Navier-Stokes equations

where u and ρ f denote the fluid velocity and density, respectively. And σ f denotes the Cauchy stress tensor, given by σ f = -p I + µ ( ∇ u + ∇ u T ) , where I 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

with F = I + ∇ d and J = det( F ) . Here d , ρ s , F , and σ s 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 are

provided in the Appendix Sec. C.1.

Modeling Rayleigh-Bénard Convection. The Rayleigh-Bénard convection system governs the flow of a fluid layer heated from below and cooled from above. The governing equations for the RayleighBénard system consist of the incompressible Navier-Stokes equations coupled with an energy equation for heat transfer. The system is modeled as follows:

Dataset Description 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 f = 10 , using a constant time-step δt = T f n , where n = 1000 . The data sets are composed of solution trajectories [ u t , p t , d t ] (resp. [ u t , p t ] ), which denote the approximate solution of the fluidstructure interaction problem (resp. the fluid dynamics problem) at times t = iδt, i ∈ { 0 , . . . , n } . These trajectories are generated on the basis 3 parameters ( µ, c 1 , c 2 ) describing combinations of fluid viscosities µ ∈ { 0 . 5 , 1 , 5 , 10 } and inlet conditions, ( c 1 , c 2 ) ∈ I .

For our setup, the fluid considered is water, with a density of 1000 kg.m -3 and a maximum inlet velocity of approximately 4 m.s -1 , 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 as a benchmark problem [12, 17] (See Sec. C.2 for a detailed explanation).

To study the Rayleigh-Bénard convention system, we degenerate two different PDE datasets. Firstly, we generate Rayleigh-Bénard convection system with Ra number 12 × 10 3 and 20 × 10 3 . We set the temperature difference between the top (cold) and bottom (hot) boundaries to 1 . We assume no-slip boundary conditions, and to 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 2 D 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: G p NS -EW and G p NS . The former is pretrained on a fluid-structure interaction dataset that combines the Navier-Stokes equation and the elastic wave equation, denoted as G p NS -EW . The latter, G p NS , is 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 state at the subsequent time step based on its current state. For the fluid-structure interaction dataset, we train an operator G NS -EW such that G NS -EW : [ u t , p t , d t ] → [ u t + δt , p t + δt , d t + δt ] , where u, p , and d are the velocity, pressure, and mesh deformation fields (see Sec. 4). For the data with only fluid motion, we train the operator G NS which maps between the current and next time step velocity and pressure field as G NS : [ u t , p t ] → [ u t + δt , p t + δt ] .

For the Rayleigh-Bénard convention system, we pretrain a CoDA-NO model, denoted as G p NS , on the incompressible Navier-Stokes equations using 40 , 000 snapshots in a self-supervised manner. The supervised task for this system is to train an operator, G NS -T : [ u t , T t ] → [ u t + δt , T t + δt ] , where u represents velocity and T represents temperature. The pretrained model G p NS is fine-tuned for the

The pretrained model for both datasets is fine-tuned for unseen viscosity µ = 5 . 0( Re = 400) with different numbers of a few shot examples. The inlet conditions of these simulations are excluded from the pretraining data. So, the target PDEs' viscosity and inlet conditions are absent in the per-taining dataset. We test the model's adaptability on a more turbulent fluid-solid interaction dataset with Re = 4000( µ = 0 . 5) by finetuning both pretrained models G p NS -EW and G p NS on each dataset.

supervised task of solving Rayleigh-Bénard convection using different numbers of a few shot training samples.

Baselines. For comparison on the supervised tasks on fluid-structure 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 a 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-Bénard convection system, we train Unet [39] and FNO [12] from scratch and compare them against our proposed model.

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 2 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. 1, we report the performance of our 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 is 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, we 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énard convention. We observe that pretrained CoDA-NO outperforms every other baseline and adapted to the new temperature variable, T , of the Rayleigh-Bénard 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.

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 to 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. 1, 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 a 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 an ablation study on fine-tuning methods. Instead of fine-tuning all the parameters, here, we freeze the parameters of 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, we 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 we 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 a 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-physics PDEs (see Appendix Sec. F.3).

Limitations. In general, CoDA-NO's performance on target PDEs is influenced by the number of training examples, and we highlight the potential for further enhancement through the integration of physics-informed approaches.

5 Conclusion

In this work, we introduce CoDA-NO, a versatile pre-trained model architecture designed for seamless adaptation to Partial Differential Equations (PDEs) featuring diverse variable compositions. Departing from conventional patch-based attention modules, CoDA-NO innovatively extends the transformer to function spaces by computing attention across co-domains. Leveraging a flexible variable encoding scheme and a 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 as a valuable tool for addressing challenges in machine learning for PDEs.

Acknowledgments

A. Anandkumar is supported in part by Bren endowed chair, ONR (MURI grant N00014-18-12624), and by the AI2050 senior fellow program at Schmidt Sciences. We thank David Pitt for his support in adding our code to the neuraloperator library, facilitating broader use and accessibility.

References

Gilbert Strang. Computational science and engineering. Optimization , 2007. 1

Gege Wen, Zongyi Li, Qirui Long, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M Benson. Real-time high-resolution co 2 geological storage prediction using nested fourier neural operators. Energy & Environmental Science , 16(4):1732-1741, 2023. 1

Burigede Liu, Nikola Kovachki, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, Andrew M Stuart, and Kaushik Bhattacharya. A learning-based multiscale method and its application to inelastic impact problems. Journal of the Mechanics and Physics of Solids , 2022. 1

Geoffrey K Vallis. Atmospheric and oceanic fluid dynamics . Cambridge University Press, 2017. 1

Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision , 2021. 2

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning . PMLR, 2021. 2, 4

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 arXiv:2306.00258 , 2023. 2, 4

Michael McCabe, Bruno Régaldo-Saint 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 arXiv:2310.02994 , 2023. 4, 22

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

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. In International Conference on Learning Representations , 2021. 2, 9

Bowen Zhang, Shuyang Gu, Bo Zhang, Jianmin Bao, Dong Chen, Fang Wen, Yong Wang, and Baining Guo. Styleswin: Transformer-based gan for high-resolution image generation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , 2022. 2, 4

Zongyi Li, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew 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, 9, 22

Kamyar Azzizadenesheli, Nikola Kovachki, Zongyi Li, Miguel Liu-Schiaffini, Jean Kossaifi, and Anima Anandkumar. Neural operators for accelerating scientific simulations and design. arXiv preprint arXiv:2309.15325 , 2023. 2

Jean Kossaifi, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, and Anima Anandkumar. Multi-grid tensorized fourier neural operator for high resolution PDEs, 2023. 2

Boris Bonev, Thorsten Kurth, Christian Hundt, Jaideep Pathak, Maximilian Baust, Karthik Kashinath, and Anima Anandkumar. Spherical fourier neural operators: Learning stable dynamics on the sphere. In Proceedings of the International Conference on Machine Learning (ICML) , 2023. 2

Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895 , 2020. 2, 4, 14

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

Zongyi Li, Nikola Borislavov Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Prakash Otta, Mohammad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Azizzadenesheli, and Anima Anandkumar. Geometry-informed neural operator for large-scale 3d pdes. arXiv preprint arXiv:2309.00583 , 2023. 3, 9

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

Makoto Takamoto, Francesco Alesiani, and Mathias Niepert. Learning neural pde solvers with parameter-guided channel attention. arXiv preprint arXiv:2304.14118 , 2023. 4

Zijie Li, Dule Shu, and Amir Barati Farimani. Scalable transformer for pde surrogate modeling. arXiv preprint arXiv:2305.17560 , 2023. 4

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

Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. 2018.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning . PMLR, 2020. 4

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research , 2023. 4

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 text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems , 2022.

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 IEEE/CVF international conference on computer vision , 2021. 4

Bonan Xu, Yuanye Zhou, and Xin Bian. Self-supervised learning based on transformer for flow reconstruction and prediction. arXiv preprint arXiv:2311.15232 , 2023. 4

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

Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing ingredient for fast stylization. arXiv preprint arXiv:1607.08022 , 2016. 6

Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning . PMLR, 2015. 6

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV) , 2018. 6

Zongyi Li, Nikola Borislavov Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Prakash Otta, Mohammad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Azizzadenesheli, et al. Geometry-informed neural operator for large-scale 3d pdes. arXiv preprint arXiv:2309.00583 , 2023. 7

Aslak W Bergersen, Andreas Slyngstad, Sebastian Gjertsen, Alban Souche, and Kristian ValenSendstad. turtlefsi: A robust and monolithic fenics-based fluid-structure interaction solver. Journal of Open Source Software , 2020. 8

Valentin Duruisseaux, Miguel Liu-Schiaffini, Julius Berner, and Anima Anandkumar. Towards enforcing hard physics constraints in operator learning frameworks. ICML 2024 AI for Science Workshop , 2024. 8

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.03193 , 2019. 9

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. 9

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention . Springer, 2015. 9

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.08481 , 2021. 14

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.03485 , 2020. 14

Dick Kachuma and Ian Sobey. Linear instability of asymmetric poiseuille flows. 2007. 15

Sukhendu Ghosh. Relative effects of asymmetry and wall slip on the stability of plane channel flow. Fluids , 2017. 15

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

Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM . 23

Appendix

The appendix is organized as follows:

In Sec. A, we provide a brief introduction of Neural Operators.

In Sec. B, we describe regular attention mechanism.

In Sec. C, we describe the detailed regarding the dataset generation.

In Sec. D, we provide a 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 are a class of deep learning architectures designed to learn maps between infinitedimensional function spaces [40]. A Neural Operator seeks to approximate an operator G that maps an input function a ∈ A to its corresponding output function u ∈ U by building a parametric map G ϕ : A → U . The typical architecture of a Neural Operator can be described as G ϕ = P ◦ T L ◦ . . . T 1 ◦ L . Here, L : a → w 0 and P : w L → u are lifting and pointwise projection operators, respectively. The action of any pointwise operator H : { f : D → R d f } → { g : D → R d g } can be defined as

where h θ : R d f → R d g is any function with parameters θ . The integral operator T l : w l -1 → w l performs a kernel integration over the input function w l -1 as

Here, D l -1 is the domain of the function w l -1 . In the case of Fourier Neural operators (FNO) [16], a convolution kernel, i.e., k l ( x, y ) = k l ( x -y ) was used. By the convolution theorem, this enables the representation of an integral operator as a pointwise multiplication of the Fourier coefficients as follows w l = F -1 ( F ( k l ) ⊙F ( w l -1 )) .

For the Graph neural operator (GNO) [41], a small neighborhood B r ( x ) ∩D l -1 around the point x is considered instead of integrating over the whole domain D l -1 , such that Eq. (2) changes to

Given a set of evaluations of the function w l -1 on points { y i } n i =1 ⊂ D l -1 , the kernel integral can be approximated by

where q i ∈ R are suitable quadrature weights [40]. The discretized kernel integral can be viewed as a message passing on graphs, where the neighborhood of each point x consists of all points within radius r .

B Attention mechanism for finite-dimensional vectors

Given three sets of vectors, so-called queries { q i } N q i =1 , keys { k i } N k i =1 , and values { v i } N v i =1 with N k = N v and matching dimensions of queries and keys, attention mechanism calculates weighted sums of the value vectors. Specifically, the set of output vectors { o i } N q i =1 can be expressed that

where a i = SoftMax [ ⟨ q i , k 1 ⟩ τ , . . . , ⟨ q i , k Nk ⟩ τ ] and τ is the temperature term. For the self-attention mechanism, the key, query, and value vectors are calculated from some input sequence { z } L i =1 using the key, query, and value matrices K , Q , and V as

C Dataset Description

u

x

at

t

u

x

at

t

+

δt

0.4

0.4

0.2

0.2

0.0

0.0

0.0

0.5

1.5

2.0

2.5

0.0

0.5

1.0

1.5

2.0

2.5

C.1 Fluid-Structure Interaction System

Here, we provide the details on generating the fluid-structure interaction dataset involving NavierStokes and Elastic wave equations.

Fluid-structure interaction model. Under the Kirchoff St-Venant model, the Cauchy stress tensor σ s verifies

where λ s and µ s are the Lame coefficients, and

Inlet Boundary Condition. Time-dependent inlet boundary conditions consist of 4 th order polynomials velocity profiles which vanish at the channel walls [42, 43]. The inlet conditions are given by

Here v is the ramp function defined as

and ( c 1 , c 2 ) ∈ I , where

are enforced at the inlet x = 0 .

Geometric setup, boundary, and initial conditions. In the considered setup (see also Figure 3), a fluid flows past a fixed cylinder of radius R = 0 . 05 centered at ( x c , y c ) = (0 . 2 , 0 . 2) in a twodimensional 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

where n denotes a unit normal vector to the fluid-solid interface. No-slip boundary conditions are imposed on the fluid velocity at the top (resp .bottom) boundaries of the channel at y = 0 (resp. y = 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 = 0 for the pressure.

The initial conditions

where the displacement d = 0 corresponds to a perfectly horizontal elastic strap and is imposed at time t = 0 .

Details regarding the data set generation. The TurtleFSI package provides a 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-Lagrangian (ALE) formulation of the problem and developed on the FEniCS computing environment [44].

The initial conditions are expressed at set X = X S ∪ X F of mesh points, corresponding to the union of the solid and fluid domains. In the ALE formulation, at each snapshot 0 ≤ t ≤ t M of the simulation, the solution is given at a set of mesh points X t = X + d t , where d t denotes the mesh displacement. In particular, the snapshots u t (resp. p t ) correspond to numerical approximations of the velocity (resp. the pressure) at the mesh points X t . Notably, while equation (10) governs the deformation field in the solid domain Ω s t , the displacements d t are obtained through an extension of the deformation field to the fluid domain Ω f t via a biharmonic extrapolation.

In all the cases considered, the values ρ f = 1 . 0 × 10 3 , ρ s = 1 . 0 × 10 3 , λ s = 4 . 0 × 10 6 and µ s = 2 . 0 × 10 6 were used. The simulations were performed using a 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 4 m.s -1 , leading to Reynolds numbers in the range 200 -2000 ( µ = 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 or only fluid motion with such a Reynolds number is quite challenging and used as a benchmark problem [17].

Modeling fluid-solid interaction with an even higher Reynolds number requires a very high computational cost. Because TurtleFSI's (used in this study) fluid solver, including its' fluid-structure interaction 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-dissipating 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, an extremely small time step ( ∆ t ) 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 a significant disparity between the pre-training and finetuning stages, particularly concerning examples with viscosities 1 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 µ ∈ { 1 , 10 } , the model faces formidable challenges in adapting due to variations in inlet conditions. The finetuning dataset with viscosity=5 has different viscosity as well as intel conditions compared to the pre-training dataset, serving as an out-of-distribution PDE setup.

C.3 Generating Rayleigh-Bénard dataset

The initial temperature field is initialized with a linear gradient between the hot bottom boundary, T bottom = 1 , and the cold top boundary, T top = 0 . To induce instability and initiate convection, temperature perturbations are introduced in localized regions of the domain. A region centered at ( L x 4 , L y 4 ) is perturbed to T = 1 , while a region near the middle of the domain, centered at ( L x 2 , L y 2 ) ,

is set to T = -1 . These perturbations break the symmetry and help to trigger the onset of convection patterns.

For the incompressible Navier-Stocks equation, we consider two-dimensional Kolmogorov flow (a form of the Navier-Stokes equations) for a viscous, incompressible fluid,

with the incompressibility constraint ∇· u = 0 on the domain [0 , 2 π ] 2 × (0 , ∞ ) . The initial condition is given as u ( · , 0) = u 0 , where u denotes the velocity, p the pressure, and Re is the Reynolds number which we set to 500 for our 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 R applied to 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 is a global operation. In CoDA-NO, however, the mixing of channels happens in the spatial domain through the attention mechanism as well as in the Fourier space because K h , Q h , V h , M , 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 as it can seamlessly incorporate additional or fewer variables during fine-tuning, avoiding retraining the whole model from scratch like FNO would have to, which can be computationally expensive.

E Visualization of Results

0.0

0.5

1.0

1.5

2.0

2.5

0.0

0.2

0.4

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

(a)

Ground Truth

0.0

0.5

1.0

1.5

2.0

2.5

0.0

0.2

0.4

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

4.5

(b)

CoDA-NO Prediction

0.0

0.5

1.0

1.5

2.0

2.5

0.0

0.2

0.4

0.25

0.20

0.15

0.10

0.05

0.00

0.05

0.10

0.15

(c)

Error

F Additional Results

F.1 Ablation of Proposed components

Table 3 shows that replacing the codomain attention with a regular attention mechanism or the removal of any of these designed components significantly impacts the model's performance. We also observe that our proposed normalization technique is crucial for effective training.

training error.

F.2 Zero-Shot Super Resolution Test

Here, we present the results of our zero-shot super-resolution models (see Tab. 4) on complex fluidsolid interaction problems. We train the models with 1317 mesh points on the domain (xy plain). However, during inference, the solution function is 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.

F.3 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. 5. It might seem that models are not compared fairly, as the CoDA-NO has a higher parameter count. However, here, we test the models on a few shot learning problems. Increasing the baselines' parameter count worsens the overfitting problem.

To demonstrate this fact, we perform experiments on a 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 during 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.

F.4 Energy Spectrum

Here, we show the energy spectrum for the NS-EW dataset for µ = 5 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. 8, we present some additional ablation studies on our 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-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training) is a large-scale pre-training approach for learning PDE representations. It utilizes an auto-regressive denoising strategy and a Fourier transformer architecture for efficient pre-training on diverse PDE datasets.

To assess the performance of these models on a diverse set of PDEs, we conduct experiments on 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

EnergySpectrum

GNN

ViT

FNO

DeepONet

CoDA-NO (NS-EW)

103

CoDA-NO (NS)

Ground Truth

Energy

102

101,

0

10

20

30

40

50

Wavenumber

achieved through this is called the 'Reconstruction error" . This allows the models to learn meaningful representations of the underlying physical systems.

Predictive task : Subsequently, we finetune the pre-trained models using a 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 t ∈ [0 , T ] to the state at time t ∈ [ T, T +5] , effectively predicting the next 5 timesteps given the history up to time T .

Both models were pre-trained and fine-tuned on this dataset for 35 epochs.

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 to the single-physics experiments, we 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 or governing equations. Both FNO and CoDA-NO were pre-trained and fine-tuned on this dataset for 35 epochs.

DPOT, with its large-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 on SWE and 0.0073 on DIFF. It is also important to 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) 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 do well on the SWE dataset as shown in 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.

It is important to highlight that DPOT was pre-trained on 12 datasets for 1000 epochs, while CoDANO was pre-trained and fine-tuned on a single dataset for 35 epochs. Despite this difference in 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 is shared physics between the pre-training and fine-tuning datasets, CoDA-NO can effectively leverage this commonality to achieve strong generalization

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 a more compact model.

In summary, these experiments on the PDEBench datasets demonstrate the effectiveness of CoDANO in learning and generalizing to different PDE systems. CoDA-NO's performance, especially considering its smaller model size and shorter pre-training, showcases its potential as a foundation model for scientific machine learning. The ability to achieve competitive results with DPOT, despite 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 a different number of parameters along with the performance of CoDA-NO.

H Implementation Details

The variable encoders are implemented using a multi-layer perceptron mapping the position x ∈ D to the embedding vector. Following transformers and NeRF [45] model, we use positional encoding

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 to apply masking on an irregular mesh, we select a point at random from the mesh. Following 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 we have masked out a predetermined portion of all mesh points.

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 are 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 on the given assumptions.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the 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 or 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are 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 no 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-specification, 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 on the scope of the claims made, e.g., if the approach was only tested on a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text 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 fairness.

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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

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 crossreferenced.

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 to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or 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 experimental results of the paper to the extent that it 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 or 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 a 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 is a specific model and empirical evaluation, it 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 is 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 a 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a 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 a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to 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., to 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

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so '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: //nips.cc/public/guides/CodeSubmissionPolicy ) 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 are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to 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 a 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:[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, confidence 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.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on 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, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 is foundation research and has limited negative societal impacts.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or 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.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper has no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 a best faith effort.

12. Licenses 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 not use 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 a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

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/model 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 as 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 crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, 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 nor 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 NeurIPS 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe authors explained the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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).\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nPlease see my concerns in **Weaknesses** and **Questions**.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- Why do we need a VSPE per physical variable? Positional encoding are usually used when there is some sort of order between the tokens?\n\nLIMITATIONS:\nOverall, 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

comments: 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.

STRENGTHS:
- The paper investigates an interesting problem of how to appropriately capture interactions across different physical variables, that allows for generalization to new codomains.
- As far as I am aware, the generalization of the Transformer architecture to function spaces is novel.
- The experimental results, especially the generalization capabilities (from fluid dynamics to fluid-solid interactions) are impressive.
- 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.

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?
- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13