Towards Foundation Models for Scientific Machine Learning: Characterizing Scaling and Transfer Behavior

Summary

This work explores and demonstrates how a "pretraining+finetuning" paradigm may be leveraged for neural networks applied to PDEs. Specifically, for three different systems governed by distinct PDEs, the authors study how different factors influence performance: (1) downstream dataset scale; (2) model scale; (3) OOD-ness of the downstream data; (4) pretraining on systems governed by different physics. The authors find that pretraining systematically improves performance, delivering significant performance gains in limited downstream data settings -- which are particularly relevant to SciML, which is often more data-constrained than NLP.

Strengths

* **S1.** This work could be very significant for the community: it explores important topics for SciML, presents extensive strong results, and is easy-to-follow, well-structured, and pleasant to read. The paper is also extensively detailed (regarding data setup, architectures, etc.). * **S2.** The authors study a variety of scenarios (OOD-ness, downstream dataset size, etc.) and perform appropriate comparisons with a "from scratch" baseline. The identified strength of the approach (in the low downstream data regime) answer particularly well challenges in the field. * **S3.** The authors open-source their code, which will help with propagating their work, increasing its impact, and enabling better reproducibility.

Weaknesses

* **W1. The experiments could be more systematic across the three systems proposed.** * **W1.1.** The SYS-3 setup is not featured for downstream dataset scaling & model scaling; it would be valuable to see it in these scenarios to further validate the method, especially since it is described as the most challenging setting. * **W1.2.** The behaviour over multiple operator could be better explored: it is lacking a comparison with SYS-2, and is done only on a very "in-domain" case (while Figure 1 for instance is done with disjoints domains) thus causing the n-shot behaviour to be flat. * **W1.3.** (minor) In Figure 4d) (with extreme OOD on SYS-1), the from scratch behaviour is different, never catching-up with the TL case. This diverges from the trend and warrants some explanations. * **W2. Some of the framing around large language models is inaccurate.** * **W2.1.** l25 "An important aspect of the trained foundation model is the notion of emergence—the model is able to perform tasks seemingly different than those for which it was trained by leveraging shared features across the training tasks." It is misleading to characterise emergence with this sentence. As discussed in Wei et al., 2022 emergence is more accurately described by "abilities that are not present in smaller-scale models but are present in large-scale models". * **W2.2.** The few-shot terminology can be confusing in some cases. For large language models, zero/few-shot is synonymous with in-context learning. Mentioning few-shot alone simply means prompting the model. Here, in this work, it systematically means finetuning the model on these few shots. * **W2.3.** l227 "This is motivated by observations in NLP that after a critical model size, we can expect significantly better adaptation performance" it's unclear what the authors refer to here, as finetuning efficiency in NLP is closely related to sample efficiency and is known to increase smoothly with model size (Kaplan et al., 2020). * **W2.4.** (minor) l145 the focus on GPT models is a bit odd: finetuning applies to all classes of modern language model, and in fact the GPT (i.e., causal decoder-only) models are somewhat unique in their ability to not require finetuning. Conversely, they do not perform as well for fine-tuned tasks as other models: see What Language Model Architecture and Pretraining Objective Work Best for Zero-Shot Generalization? Wang et al., 2022. * **W2.5.** (minor) l25 the references to scaling laws are missing a citation for Hoffmann et al., 2022, which is the reference work at this time and which highlighted the need for joint scaling of data and model size instead of model size alone predominantly. * **W3.** Smaller nits: l44 "TL" to refer to transfer learning as not been introduced as an acronym.

Questions

Overall, this is an interesting paper addressing important topics for the community, with the potential to be highly valuable. It is however slightly held back by some concerns around the showcased experiments. Addressing these would further validate the wide applicability of the approach, confirming its significance and increasing the quality of the work. Accordingly, I am rating this work as a **Weak Accept (6)** but would be willing to increase my rating should my concerns/questions be addressed. **EDIT: following rebuttal, I have updated my score to an Accept (7)**. * **Q1.** (W1.1.) Could the authors provide results on SYS-3 for downstream dataset scaling and model scaling? * **Q2.** (W1.2.) Could the authors provide additional results for Figure 5./Q4, demonstrating the approach on SYS-2, and across all systems but OOD? (like in Figure 2). * **Q3.** (W1.3.) Could the authors comment on why the extreme OOD in Figure 4d) behaves differently? * **Q4.** (W2.3.) Could the authors clarify their statement on l227/provide a reference? * (W2.2.) As a suggestion, to avoid confusion for readers, I would recommend systematically saying "few-shot finetuning" and not using "few-shot learning": l73 "few-shot TL" could be unclear, l132 "few-shot learning" --> "few-shot finetuning", l221 "few-shot learning" --> "few-shot finetuning"

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The authors discuss some of the limitations of their work in the conclusion, but I would encourage a more critical look regarding when the proposed method could fail.

Summary

This paper investigates the finetunability, transfer, and generalization properties of a foundational model applied to the field of scientific machine learning, as functions of finetuning dataset size, pretrained model size, and value of physics parameters governing the equations of the system.

Strengths

The paper is well-written and well-organized, easy to follow and understand. The authors present their method clearly, and they provide a detailed evaluation of their method, with findings clearly highlighted and supported by strong empirical evidence. The experimental hypothesis are clearly stated up front and verified throughout the paper, with sound methodology and experiment design.

Weaknesses

A weakness, also acknowledged by the authors at the end of Section 3 -- but without much contextualization --, is that they only explore the questions raised in the paper in the Fourier Neural Operator (FNO) architecture, and present it as a prototypical SciML architecture. Section 3 explains that this is a good choice for problems with periodic boundary conditions like the ones explored in this paper. However, it's unclear whether any of these results would hold for other types of OOD problems or other types of architectures. Therefore, without more contextualization, it's unclear whether these findings are generalizable and would allow the community to build on top of these findings, or whether they are anecdotal, and would only hold for the architecture and type of problems explored in this work. Another weakness is that some of the findings, as currently present, merely seem to provide additional experimental evidence for notions that are already widely accepted in the community, such as the fact that, in the low downstream data regime, pretraining a model and then finetuning it on this data is preferable (and yields higher performance) than training a model on this data from scratch, or that finetuning small capacity models plateaus to a certain irreducible error even if data continues to increase. Finally, while the authors acknowledge that a concurrent paper covers the alternative approach of in-context learning (as opposed to finetuning), I would have expected the authors of this paper to compare against that alternative, as opposed to simply benchmarking against training from scratch, which is a much less tight comparison. I think this paper would benefit from the inclusion of that baseline.

Questions

Small suggestion: please provide common examples of physical systems/phenomena described by the PDEs explored in this work.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

While disclosing the limitation of this analysis being limited to FNO, the author could be more forthcoming about the consequences of this limitations and the applicability, at a broader level, of these findings. Overall, however, the authors do an excellent job at not overstating their claims and clearly explaining the experimental setup and application domain this paper applies to.

Summary

This paper analyzes the transfer performance of the pre-trained foundation model for scientific machine learning applications described by PDEs. A wide range of settings is considered: various sizes of downstream datasets, pretrained model, out-of-distribution data, as well as multi-tasks and applications. Conclusions, though not necessarily surprising, are drawn under these settings.

Strengths

- To the best of the reviewer's knowledge, this is a relatively comprehensive analysis of pre-training and transfer learning regimes over a diverse set of physical systems. - The presentation is clear and easy to follow, and the experiments are comprehensive and convincing.

Weaknesses

- For the experiments, more network architectures can be considered. - The findings provided by the paper are more or less consistent with NLP and cv tasks. This is not necessarily a drawback, but the paper would be stronger if it provides some findings that make Scientific Machine Learning different from other well-explored AI tasks.

Questions

We can observe from the figures that the performance gap is reduced when more fine-tuning examples are observed. This is not surprising, and the advantage of the pre-trained foundation model is indeed when fine-tuning examples is less. In this case, the performance is worse. I am wondering how to interpret the error, and at what magnitude, the error becomes acceptable for these problems.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

The author adequately addressed the limitations.

Summary

This paper explores the behavior and performance of neural operator models on multiple partial differential equation (PDE) systems in the transfer learning setting. The authors investigate the impact of various factors on the performance of neural operator models, including model size, downstream dataset size, underlying physics of the downstream tasks in relation to pre-training, and distributional shifts. The study employs the "pre-train and fine-tune" paradigm for SciML (Scientific Machine Learning) problems. The authors show that it is possible and beneficial to develop more general SciML models capable of solving multiple tasks with the same set of weights, even when downstream tasks involve small-to-moderate distribution shifts relative to the pre-training data. The research demonstrates the potential of neural operator models for transfer learning, paving the way towards building foundation models for SciML.

Strengths

Originality: The paper demonstrates a high level of originality in several aspects. Firstly, it explores the use of neural operator models for solving multiple partial differential equation (PDE) systems, a novel approach that combines the power of neural networks with physics-informed learning. The authors investigate the behavior of these models under various conditions, including model size and downstream dataset size, providing valuable insights into their scalability and transferability. Secondly, the study applies the "pre-train and fine-tune" paradigm to SciML problems, showing the potential of developing more general SciML foundation models. This approach is innovative and opens up possibilities for efficiently addressing diverse downstream tasks with a single pre-trained model. Furthermore, the paper explores the impact of distributional shifts in downstream tasks relative to the pre-training data, contributing to the understanding of transfer learning behavior in the context of PDEs. The inclusion of multiple downstream PDE systems and the analysis of various performance metrics demonstrate a comprehensive and original investigation. Clarity: The paper is well-written and presents its findings in a clear and organized manner. The subsection is organized logically following the order of the posted questions. Significance: The significance of the paper lies in its contribution to the field of SciML and transfer learning for PDEs. By characterizing the behavior of neural operator models under different conditions and exploring transferability across multiple downstream PDE systems, the paper provides valuable insights for researchers working in the area of computational physics and scientific modeling. The findings have practical implications, as they suggest the potential for developing more versatile and efficient foundation models for SciML tasks. The "pre-train and fine-tune" approach demonstrated in the paper can be leveraged to achieve better generalization and performance in solving PDEs across a wide range of scenarios.

Weaknesses

Motivation: While the paper extensively analyzes the performance of neural operator models under different conditions, it lacks direct motivation of understanding the transfer learning for PDEs problem in SciML, as these different conditions have been well-explored in the LLM and CV community. Potential Overfitting in Transfer Learning Settings: The paper briefly mentions that few-shot transfer learning outperforms training from scratch, but it does not thoroughly address the possibility of overfitting to the specific downstream tasks during fine-tuning. An investigation into potential overfitting issues and strategies to mitigate them would be valuable.

Questions

While exploring transferability across multiple downstream PDE systems is an intriguing and significant area of research, is there any differences in transfer learning between PDE problems and well-established domains like Language Model Learning (LLM) and Computer Vision (CV). What does the nature of PDEs and their associated physical systems introduce unique challenges and characteristic that may affect the transferability of learned representations and the generalization of neural operator models? In my opinion, understanding and addressing these differences are crucial for insights from the transfer learning perspective in PDEs.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Listed in the Weaknesses and Limitations.