Pretraining Codomain Attention Neural Operators for Solving Multiphysics PDEs

We appreciate that the reviewer values our work and recognizes its novelty and the mathematical formulation of the proposed model.

Q1. Evaluation on additional PDE datasets

In addition to the coupled fluid-solid interaction problem, we also provide experiments on different PDE systems (please see Table 2 and Appendix Sec E, Table 11-12), where we test the proposed CoDA-NO architecture on shallow water equation, diffusion equation, and Navier Stokes equation. In addition, we also include the results of our study on the Rayleigh-Benard convection system. The CoDA-NO is pre-trained in a supervised way on the Navier–Stokes (NS) equations. Then it is subsequently finetuned on a limited few shot data ($N \in \{5,10,25\}$) of Rayleigh-Benard convection. We show the result in the following table, where we report the $L2$ loss. We see that, in the extremely low data regime, pre-trained CoDA-NO performs better than FNO and CoDA-NO trained from scratch.

Please note that these results are preliminary. As per reviewer qRRk's suggestion, we are conducting a comprehensive study on the Rayleigh-Benard convection system.

Q2. Performance gain from pre-training.

In most cases, pre-trained CoDA-NO demonstrated a significant performance improvement compared to CoDA-NO trained from scratch. While we acknowledge that in some instances, the difference is not substantial, it is consistently better. Additionally, in general, pre-trained models offer benefits such as reusability and faster convergence with fewer epochs required. The lack of a significant performance gain in a few settings does not diminish the importance and potential advantages of adapting the architecture and pre-training scheme. Further, our self-supervised pretraining method is also beneficial when the cost of generating training data with numerical solvers is very high.

Q3. Per variable loss

Previous studies have shown that aggregated loss demonstrates model accuracy across multiple variables [1, 2, 3] which helps to get a comprehensive overview of the performamnce. For instance, the results of the compressible NS equation (velocity and pressure field) are presented in an aggregated manner [1, 2, 3]. We also provide the per-variable performance of each of the variables (velocity $v_x, v_y$, pressure $p$, displacement $d_x,d_y$) on the combined Navier-Stokes and Elastic wave equation for $Re=4000$.

Table of Horizontal velocity ($u_x$)

Table of Verticle velocity ($u_y$)

Table of Verticle velocity ($p$)

Table of Horizontal displacement($d_x$)

Q8. The rationale for attention in channel space

Modeling physical phenomena has two sets of variables, viz., spatial and temporal locations and physical variables such as temperature, pressure, etc. Having attention in the channel enables us to model the relationships between physical variables that interact together. Given the well-established nature of the Transformer architecture and its property to handle variable length input, it has been a natural choice for our approach. We have redesigned key components of modern transformer layers, including positional encoding, self-attention, and normalization layers, to develop codomain attention in function space.

In many applications, physics problems are modeled using coupled PDEs, and generating data is often costly. Therefore, we need to have foundation models capable of tackling issues of a similar nature but with a different number of variables or changes in the variables. Attention in the channel space helps the model to pre-train simultaneously over a wide range of PDEs with different variables without changing the architecture. For example, fluid in the subsurface is often described by 12-14 variables that obey laws of thermodynamics, fluid dynamics, Darcy, and chemical interaction. Therefore, learning about similar physics, each with different variable counts at corresponding order, representing the particular physics, would help to train for the final subsurface model. (Please See Problem Statement (Sec 4) and lines 167-177).

There are also other applications for Codano where there are different sets of physical variables with varying counts across datasets. For example, in climate modeling, there are multiple datasets developed by different countries (please refer to CMIP6 global collaboration for climate data), each using slightly different PDEs with different variable counts. Learning the relationships among such variables represented as different channels is enabled by attention.

Moreover, a recent work [1] has shown the universal approximation property for attention-based operators in function space, i.e., that CoDA-NO can approximate any continuous operator. Although it is not straightforward to compare different neural operators on the basis of expressivity in a general sense, our empirical studies suggest that it generalizes and performs well on a variety of tasks without any significant change to the hyper-parameters.

[1] Continuum Attention for Neural Operators

Q9. Reason for large parameterization of FNO

We already addressed this in detail in a previous comment. For Table 2, the models are pre-trained and subsequently fine-tuned on the large PDEbench dataset. As FNO has a relatively simple architecture, we utilize a parameter-rich architecture to benefit from the pretraining and enable it to learn rich representations for subsequent prediction tasks.

Vertical displacement ($d_y$)

[1] Dpot: Auto-regressive denoising operator transformer for large-scale PDE pre-training.

[2] Gnot: A general neural operator transformer for operator learning.

[3] Pdebench: An extensive benchmark for scientific machine learning.

Q4. Visualization of output

We provide visualization in the attached PDF (Figure 1). We will include more visualizations in the paper.

Q5. FNO as a baseline for fluid-solid interaction problem

As the domain's discretization is irregular, we don't directly apply FNO since the computationally feasible implementation of such models only supports uniform grids (through the deployment of FFT). However, for the shallow water equation, for which the discretization is uniform, we use the spherical variants of FNO as our baseline.

Q6. Parameterization of FNO

For the experiment on the PDEbench (Table 2, Table 11-12), we do a different setup compared to the fluid-solid interaction. In this case, we follow the same setup as DPOT and first pretrain the model using self-supervised learning and then fine tune with supervised loss, and in our setting, both are evenly split. Due to this relatively large dataset for both pretraining and fine-tuning, the risk of overfitting is low. Nevertheless, we have done an extensive ablation of the number of parameters of FNO, where the model is first trained in a self-supervised manner followed by supervised training. We report the $L_2$ for each of the datasets. We can see that FNO with 1.9B parameter does not overfit the data but rather performs better. We hypothesize that this is due to self-supervised pre-training done before supervised training and the size of the training set.

Table for shallow water equation

Table for diffusion-reaction

Q7. Higher inference time 

The objective of the proposed CoDA-NO is to provide suitable architecture for foundation models for solving PDEs and multi-physics problems. In this study, we demonstrated the extensive adaptability and advantages of CoDA-NO and its pre-training mechanism. Similar to other large foundation models (such as [1]), CoDA-NO has a higher inference time compared to a smaller model but is still significantly faster compared to numerical solvers. Furthermore, the computation time can be mitigated through careful implementation, parallel computation, and engineering efforts, which we plan to address in future work.

Results on fluid motion dataset governed by Navier-Stokes equation (NS Dataset).

Q5. Generalization of different physical parameters

We also consider generalization to the Reynolds number and to the inlet boundary conditions (lines 291-292). The Reynolds number is the most important factor characterizing the dynamic of the PDE, and this is reflected in different benchmarking datasets, which often contain simulations at different Reynolds numbers. Hence we followed that. It is possible to generalize different solid density $\rho^s$ and object sizes. However, to get a good result in the few-shot learning experiments, the pretraining data should contain PDEs with different $\rho^s$ and object sizes. We agree that this problem is a very important topic of study, deserving special treatment, and developing the right datasets for such studies is on the horizon for the ML community.

Q6. Application to diverse PDE system

We also provide experiments on different PDE systems (please see Table 2 and Appendix Section E Table 11-12), where we test the proposed CoDA-NO architecture on shallow water equation, diffusion equation, and Navier Stokes equation. Also, we agree that expanding to many more PDE systems is desirable, and for that, the field is in the process of making diverse PDE benchmark datasets.

We appreciate that the reviewer values our work and recognizes the fact that this work is of interest to the scientific machine learning (SciML) community, presents a strategy that enables self-supervised pre-training of PDE systems, and states the importance of our experiments that have shown the effectiveness of the proposed method.

Q1. Additional Coupled PDE

We discuss the difficulty of modeling the fluid-solid interaction problem and justify our problem design in Appendix B1. Given the complex geometry and turbulent flow, the problem can be considered a challenging benchmark problem.

Additionally, we provide experiments on diverse PDEs (Navier–Stokes equations, Elastic wave equations, Shallow water equations, and Diffusion reaction; please see Table 2 and Appendix Sec E, Tables 11-12) and demonstrate the architecture's performance and flexibility.

As suggested, we also provide results on the Rayleigh-Benard convection system. The CoDA-NO is pre-trained in a supervised fashion on the Navier–Stokes (NS) equations. Then it is subsequently fine-tuned on a limited few shot data ($N \in \{5,10,25\}$) of the Rayleigh-Benard convection system. We show the result in the following table, where we report the $L2$ loss. We see that, in the low data regime, pre-trained CoDA-NO performs better than FNO and CoDA-NO trained from scratch. Please note that these results are primitive and yet to be finalized. Per the reviewer’s suggestion, we are running a comprehensive study on the Rayleigh-Benard convection system.

Q2.  Motivation for the components of the CoDA-NO (The motivation for the combination of positional encoding, self-attention, and normalization layers seems to be better clarified. Although those parts are modular (claimed in Line 68), the connections between each other are also important.)

We outline our motivation in the problem statement (Section 4) and explain the need for Neural Operators that can handle PDEs with varying numbers of variables (lines 167-177). Given the well-established nature of the Transformer architecture and its property to handle variable length input, it has been a natural choice for our approach. We have redesigned key components of modern transformer layers, including positional encoding, self-attention, and normalization layers, to develop codomain attention in function space. We will further emphasize this in the main text.

Q3. Details of self-supervised pre-training, such as masked ratio

We provide the masking ratio and some additional details in Appendix F. For every training sample, one of the following two masking choices is selected with equal probability

1. 50% of the mesh points of 60% of variables are masked.
2. 30% of the variables are masked out completely.

Q4. Additional Evaluation Metrics

L2 is the generalization of RMSE to function spaces and is considered a natural metric of evaluation. In this work, we also provide an Energy Spectrum representing the correct statistical distributions of the fluids to go beyond distance-based evaluation metrics (please see Appendix D.4). L2 norms, together with the Energy Spectrum, provide considerable evidence establishing the proposed model's superiority to the baseline. In addition to these two metrics, we also include L1 and relative L2 in our list of metrics. To this end, as the reviewer knows, any evaluation metric comes with its advantages as well as limitations. We hope that the current list plays a sufficient role in establishing evident results on the superior performance of the proposed method.

Table: L1 and relative L2 Error

Results on the fluid-solid interaction dataset combining Navier-Stokes and Elastic wave equation (NS-EW dataset).

We appreciate that the reviewer values our work and recognizes that this work presented a method for learning representations of different PDE systems within a single model.

Q1. Formulation of the Model in the Function Space and Clarity of the Paper

We politely disagree with the reviewer. Indeed, in practice, we typically work with discretized versions of functional data. However, this does not imply that we only deal with a "fixed" finite-dimensional space for which neural networks are designed. In practice, the discretizations often vary from one sample to another (i.e., the location and number of evaluation points). Many operations that work well for "fixed" finite-dimensional data do not behave consistently across resolutions. For example, regular CNNs (or GNNs with nearest-neighbor graph) collapse to pointwise operations when the resolution is refined (see [1]) and is thus not well suited for data that comes at different resolutions.

Such limitations of neural networks motivated the paradigm of neural operators as a generalization of neural networks to data originating from functions [2]. Designing architecture directly in the function space allows us to work across different discretizations seamlessly. It also allows us to draw parallels between various approaches to operator learning and ensure specific properties of the learning algorithms, such as the universal approximation property demonstrated in [3] for transformer architectures in function spaces. The mathematical foundation is often helpful in conveying the core ideas of the architecture design in operator learning. For example, here is reviewer ZkA5’s statement about our mathematical foundation: “The mathematical formulation of the proposed model is well-articulated in the paper. This clarity helps readers understand the underlying principles of the model's operation.”

To better explain the architecture, following the suggestion, we provide a detailed figure of the whole architecture (see Figure 2 in the rebuttal pdf).

[1] Neural Operators with Localized Integral and Differential Kernels

[2] Neural Operator: Learning Maps Between Function Spaces With Applications to PDEs

[3] Continuum Attention for Neural Operators.

Q2. From infinite-dimensional space to discretized space

Our implementations are based on previous works on Neural Operators [1,2,3]. These studies discuss how operations on functional spaces are approximated on discretized data, e.g., using Riemann sum approximation of the integral operator, which scales appropriately with the resolution. We will clarify this concept in the paper.

[1] Fourier Neural Operator for Parametric Partial Differential Equations (see Section 4).

[2] Neural Operator: Graph Kernel Network for Partial Differential Equations (see Section 3).

[3] Neural Operator: Learning Maps Between Function Spaces With Applications to PDEs.

Q3. Figure showing the whole architecture.

Thanks for the suggestion. We added a detailed figure (Figure 2) in the rebuttal pdf.

Q4. Use of VSPE per physical variable

In this work, we treat each of the input function's variables (codomains) as a token for the Codomain attention. The model needs a way to identify each variable (e.g., temperature, velocity, pressure). In LLM, it is done through position ID. In CoDANO, we propose the Variable-Specific Positional Encoding (VSPE), which is learned separately for each variable and then concatenated with the corresponding variables. It allows the model to differentiate between the various PDE variables and their interactions. The following table demonstrates the benefit of using VSPE for fine-tuning on $5$ examples for $\mu = 5.0$ (from Appendix D1, Table 3). Please refer to Appendix D1, Table 3 for a detailed analysis.

Q5. On the clarity and organization of the paper

We have addressed this in a previous comment, noting that developing architecture directly in the function spaces is crucial, conveying the philosophy of the work. We also agree with the reviewer that this should not overshadow the practical implementation aspect of the work. Our philosophy and goals are outlined in the introduction (lines 28-41) and in the Problem Statement (see Sec 4).

We provide additional detailed figures (see the attached PDF), and we move some background details to the Appendix, as suggested by the reviewer.

To this end, we would appreciate it if the reviewer reconsidered our explanation and problem setup of operator learning in this paper. The field of neural operators is still in its early stages, and similar to the early stages of neural networks, explaining the problem setup and detail of the architecture may still be considered crucial and insightful.

We appreciate that the reviewer values our work and recognizes the main contributions to the novel generalization of the transformer architecture, including self-attention, positional encodings, and normalization to function space, along with the introduction of two new challenging datasets. We further may state that in this work, we not only based our study on the mentioned two PDEs, but we also have SWE and Diffusion Eq. Per the reviewer's request, we are running experiments on the Rayleigh-Benard convection system, with initial results stated in the reviewer qRRk’s reply.

Q1 Visualization of the Data

We thank the reviewer for the visualization suggestion. We provide the visualization of data in the Appendix, Sec B Dataset Description.

Q2 Performance scaling with data (It seems that the performance of the models across the board continues to improve with an increase in few-shot fine-tuning samples beyond N=100. What does the scaling look like for the proposed model, and where does performance saturate?).

We conduct experiments with N=250, 500, 1000 for NS+EW with Re=400. We observe that performance continues to improve when training on more data. Therefore, finding the saturation point requires the generation of a much larger dataset, which requires future development and compute allocation for the solver. The precise scaling is problem-dependent.

Q3 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)?

For Re=4000, we observe that the performance of the pre-trained models is worse than for Re=400. We hypothesize that the reason may be intuitively twofold,

1. Fluid motion at Re=4000 is highly turbulent and 
   

2. the data is out-of-distribution. 
   

Inspired by other studies in machine learning, CoDA-NO will perform worse on out-of-distribution data and will require further data for finetuning. However, our evidence suggests that the performance will be better than training from scratch.

Also, finding the accurate scaling with respect to ODD, i.e., what degree of out-of-distribution the model can handle, requires extensive problem-dependent study