Data-Efficient Operator Learning via Unsupervised Pretraining and In-Context Learning

Authors: wuyang  chen
jialin  song
pu  ren
shashank  subramanian
dmitriy  morozov
michael w mahoney
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-04 14:34:50

Abstract: Recent years have witnessed the promise of coupling machine learning methods and physical domain-specific insight for solving scientific problems based on partial differential equations (PDEs). However, being data-intensive, these methods still require a large amount of PDE data. This reintroduces the need for expensive numerical PDE solutions, partially undermining the original goal of avoiding these expensive simulations. In this work, seeking data efficiency, we design unsupervised pretraining and in-context learning methods for PDE operator learning. To reduce the need for training data with simulated solutions, we pretrain neural operators on unlabeled PDE data using reconstruction-based proxy tasks. To improve out-of-distribution (OOD) generalization, we further assist neural operators in flexibly leveraging in-context OOD examples, without incurring extra training costs or designs. Extensive empirical evaluations on a diverse set of PDE equations demonstrate that our method is highly data-efficient, more generalizable, and even outperforms conventional vision-pretrained models.

Main paper: 
 introduction: 
 Recent advancements in machine learning methodology have shown promise in solving partial differential equations (PDEs) [24, 67, 38, 45, 44, 50, 39, 25] . A significant development in this area is the concept of operator learning for PDEs. This approach differs from traditional neural network methods, which are restricted to fixed-dimension input and output, since neural operators focus on learning mappings between function spaces [38, 45, 44] . Like other neural network methods, neural operators are recognized to be universal approximators for any continuous operator [46, 38] , enabling them to approximate any physical operator, including solution operators for various parametric PDE families. A solution operator is defined as a function that maps physical inputs to output solutions. Previous work has shown that in simple settings, neural operators can effectively capture complex, multi-scale dynamic processes [45, 46, 74, 73] .However, neural operators tend to suffer from a problem common to other deep networks, namely the need for enormous quantities of data. Limited availability of data is common in science and 38th Conference on Neural Information Processing Systems (NeurIPS 2024).engineering. High-fidelity numerical simulations are computationally costly or even infeasible for many applications [71] . For example, an extreme-scale simulation of magnitude 7.0 earthquake at frequencies up to 10 Hz in San Francisco requires 3600 Summit GPU nodes and 42.7 hour [53] .Motivated by this data-efficiency challenge, recent works, particularly in natural language processing (NLP) and computer vision (CV), have focused on unsupervised (or self-supervised) pretraining 2 to reduce the cost of collecting or generating labeled data. Such pretrained models have been shown to be highly data-efficient in downstream fine-tuning [7, 29, 8] , and they can even become few-shot learners without any downstream data [3] . In CV, researchers collect large amounts of natural images without any manual labels, and then pretrain visual encoders with proxy tasks, such as Noise-Contrastive Estimation (NCE) [59] , masked reconstruction [28] , rotation and jigsaw prediction [58, 20] . In NLP, people typically pretrain models via next-word prediction or masked tokens [3, 12] .However, unsupervised pretraining is still largely underexplored in Scientific Machine Learning (SciML). Therefore, our core question is: How can we design unsupervised pretraining for operator learning to reduce the data simulation costs?Fine-tuning reduced simulation costs In this work, we resort to unsupervised pretraining for neural operator learning to achieve data efficiency in SciML. We overview our framework in Figure 1 . First, we define unlabeled data for PDEs, which avoids heavy computation costs for simulating PDE solutions. We propose two physics-inspired reconstruction-based proxy tasks, and we pretrain neural operators on unlabeled PDE data. We demonstrate that with unsupervised pretraining, our neural operators not only improve counterparts trained with more simulated data, but also they outperform off-the-shelf pretrained checkpoints from other popular domains (such as CV) that are ready for fine-tuning. Then, to further improve the data efficiency during out-ofdistribution (OOD) inference, we design a similarity-based method that learns in-context examples [3, 49, 80, 81, 47] . This approach introduces zero overhead during training: one just maintains the standard training pipeline, and it can be seamlessly plugged in for OOD inference, without further fine-tuning. In more detail, we summarize our main contributions:Inputs1. We introduce unlabeled PDE data and unsupervised pretraining for data-efficient neural operator learning. We show that our method can achieve better performance than models trained with more simulated PDE solutions, or fine-tuned from public checkpoints pretrained on other benchmarks, demonstrating the importance of unsupervised pretraining on domain-specific PDE data. 2. We propose a similarity-based method to improve the OOD generalization of neural operators, which is flexible and can scale up to a large number of unseen in-context examples ("demos"). 3. We provide detailed empirical evaluations on both diverse PDE benchmarks and also several realworld scenarios, demonstrating that we can achieve both strong forward modeling performance and significant savings in PDE simulations.2 Related Works 
machine learning for scientific modeling: 
 There has been a long history of using learning-based methods to model physical scientific phenomena [41, 42, 6, 5] . A representative line of work is so-called physics-informed neural networks (PINNs) [67, 87, 19, 18, 69] , which try to incorporate physics in neural networks by including the differential form of the PDE as an additional physics loss regularization term. However, this paradigm is confined to specific PDE scenarios (e.g., fixed PDE coefficients), instead of being more physics-agnostic. Moreover, recent work has highlighted several fundamental "issues" with PINNbased methods [40, 15] . On the other hand, operator learning methods, including Fourier Neural Operators [45, 44, 38] and Deep Operator Network [50] , have achieved progress in approximating the solution operators of PDEs. Although these data-driven approaches show promise in learning PDE solutions, they (like many neural network-based methods) rely on vast quantities of high-fidelity labeled data. For operator learning, such data are usually computationally expensive to simulate [67, 2, 84] .More recently, people have tried to generate synthetic PDE solutions to train SciML models [27] . In contrast, our method works on unlabeled PDE data. This approach is distinct from the generation of synthetic PDE data, and it could be further combined as a semi-supervised learning strategy.Unsupervised (or self-supervised) pretraining is a key method in CV and NLP to achieve meaningful representations [7] , data-efficient fine-tuning [31] , and foundation models [1] . In CV, contrastive learning learns meaningful features by distinguishing between similar (positive) and different (negative) samples [59, 76, 7, 56, 29] . Masked Autoencoder (MAE) [28] uses a reconstructive approach where parts of the input are masked and the model learns to predict masked parts. In NLP, among the most prominent works are large language models (LLMs) such as GPT [3, 64, 65] and BERT [12, 66] , which leverage token predictions for pretraining. Similar directions also show progress in SciML. For example, [2] and [54] propose to create augmented views in the solution space via Lie Symmetries; [73] study the scaling behavior of supervised pretraining and OOD generalization, charting directions for foundational models for SciML; [43] target learning astronomical foundation models with crossmodal contrastive learning; and [52] build large task-agnostic models with a broad understanding of common physical behavior to serve as foundation models for SciML.In-context learning (ICL) is a promising paradigm that helps deep networks generalize to unseen domains with a few in-context examples. Early works in CV seek to learn feature-level correspondence between the target and a few "shots," such that models can generalize to open-set unseen objects [17, 83] . In NLP, people find LLMs are naturally few-shot learners [3] , and thus tuning or optimizing prompts becomes extremely important to improve the in-context learning performance of LLMs [85, 72] . More recently, within SciML, a different operator learning strategy, termed "in-context operator learning," has been proposed [80, 81, 47] . During both training and inference, the neural operator is asked to make predictions by explicitly leveraging a predefined number of so-called "demo" examples (pairs of physical parameters and simulated solutions). This approach provides a balance between model generalization and addressing data scarcity in the scenario of OOD testing. 
methods: 
 In this section, we introduce our framework (outlined in Figure 1 ). We propose first to pretrain the model with unsupervised pretraining (Sec. 3.1), which will contribute to the data efficiency and reduced PDE simulation costs during standard training of neural operators. When we move to OOD scenarios during inference, we test our models with in-context examples (Sec. 3.2) to avoid further fine-tuning costs.The core idea in unsupervised (or self-supervised) pretraining is to train a neural network with properly designed proxy tasks. These proxy tasks do not require labeled data, but they are designed to be highly related to the supervised learning objectives of interest. While popular in CV and NLP, unsupervised pretraining on unlabeled data has never been explored in PDE operator learning in SciML. This is due to two unresolved questions: 1) What kinds of unlabeled data can we use to train neural operators? 2) How to design proxy tasks for PDE data? We address these questions below.How to Define Unlabeled PDE Data? In general, when a neural operator is trained on PDE datasets [45, 74, 73] , it learns to map inputs (physical parameters, coordinates, forcing functions, initial conditions, etc.) to PDE solutions. Therefore, given a set of PDE data (collected via simulations or observations), its unlabeled version is defined as the one without PDE solutions. Our unlabeled PDE data is a broader concept of related inputs in modeling PDE systems. Let us consider the second-order linear differential equation as a general example. It is formulated asn i,j=1 a ij (x)u xixj + n i=1 b i (x)u xi + c(x)u = f (x),where x ∈ R n represents physical space that varies in different systems (e.g. n = 3 for 2D timedependent PDEs); the coefficients (or physical parameters) a ij , b i , c are known from the physical process; u is the target solution; and f denotes an external forcing function [57] . We can consider two situations where solutions are unavailable:• Time-independent equations: following [45, 73] , our unlabeled PDE data include physical parameters (a ij , b i , c), forcing functions (f ), and coordinates (grids of the discrete physical space). • Time-dependent equations (e.g., forecasting problems [74, 52] ): without simulating the temporal dynamics, our unlabeled PDE data only include initial snapshot u 0 (x) that defines PDE systems. Note that collecting snapshots with temporal dynamics in large-scale scenes is more complex than capturing individual snapshots. For example, weather forecasting [30] and smoke dispersion [14] require continuous monitoring and multiple sensors, whereas single measurements are simpler and less resource-intensive. Long-term data collection often involves extensive networks and processing, unlike one-time measurements.For concrete examples of different PDEs and unlabeled data that we will study, see Appendix A.There are two main reasons for pretraining only on unlabeled PDE data, as discussed below.Cheap Generation of Unlabeled PDE Data. One critical reason that leads to expensive computational costs when collecting PDE data is the time marching scheme [23, 22] in numerical simulation. However, only generating unlabeled PDE data and snapshots without temporal dynamics will be much cheaper than simulating solutions (Table 4 ), making our unsupervised pretraining highly feasible in practice. Our pretraining strategy will be very data-efficient, and can avoid the heavy computational cost of simulating complex time-dependent equations for massive high-fidelity labeled solutions. We defer our results and experimental details in Figure 4 in Sec. 4.1.To illustrate our general approach of constructing proxy tasks, we choose two variants of reconstruction as our core proxy tasks. In particular, we will input unlabeled PDE data to our neural operators, and after a decoder network, we will force the output to be close to the input. We consider two perturbation variants (or augmented views). They are inspired by real-world settings when people collect scientific data, and they are important invariances we need to introduce to SciML models.Masked Autoencoder. Masked autoencoders (MAEs) have been shown to be scalable selfsupervised learners [28] . The method is conceptually simple: remove a portion of the input data, and learn to predict the removed content. These methods enable training in NLP and CV of generalizable models containing over one hundred billion parameters [12, 3] . Here, we investigate the potential of MAE for scientific modeling. 70% ) of spatial locations are masked, followed by a Gaussian blur. After the encoder and decoder, the full set of input is required to be reconstructed.Motivation: PDE dynamics are invariant to sparse sensing of the full field scientific data. It is very common [4, 16, 32] that scientific data need to be collected from sparse sensors, and that people need to reconstruct or generate data for domains without sensors. We enforce our model to learn sensor invariance via random masking, and we extract the invariant features over distorted views of the same unlabeled PDE data. The invariance to sparse sensing of scientific data will facilitate the robustness of the representations from MAE.Therefore, we consider MAE as a proxy task. Specifically, our MAE is a straightforward autoencoding approach that reconstructs the original signal, given its partial observation. Like all autoencoders, our approach has an encoder that maps the observed signal to a latent representation, and a decoder that reconstructs the original signal from the latent representation. We randomly sample masks according to a certain masking ratio (i.e., the ratio of removed areas), and the values of masked areas of unlabeled PDE data are set to zero. Our loss function computes the mean squared error (MSE) between the reconstructed and original input. We compute the loss only on masked areas; if no mask is applied (i.e., a vanilla autoencoder), the MSE loss will be applied to all spatial areas. Super-resolution. Super-resolution (SR) techniques have emerged as powerful tools for enhancing data resolution, improving the overall quality and fidelity of data representation, and retrieving fine-scale structures. SR is a task that involves recovering fine-scale data from corresponding coarsegrained data. It is also a popular task on PDE learning [9, 68] , which is to train SciML models to preserve the inherent physical properties of scientific data.Motivation: Numerical solutions of PDEs are expected to exhibit invariance to filtering blur or different resolutions of inputs [37, 79] . For instance, in turbulence simulations, the traditional numerical methods always fail to model the expected physical phenomenon with low-resolution meshes due to substantial numerical errors. SR has emerged as a powerful tool for subgrid modeling of PDE dynamics, especially helping to capture the critical patterns of turbulence [51, 37] . Given a specific input distribution, after fitting the SR objective, neural operators are expected to preserve the inherent physical properties and exhibit invariance to filtering blur. Therefore, we introduce SR as another proxy task. Our objective shares the same motivation as recent SciML works for SR [68] . Specifically, we enforce the model to learn invariant features of unlabeled PDE data that are immune to resolution and blur. Blurry snapshots often occur when the resolution is too low to accurately represent the details in the original content. To do so, we apply a Gaussian filter to blur the unlabeled PDE data, and the autoencoder will reconstruct the high-resolution input with fine-grained details. Instead of applying a fixed blurring, we randomly sample the variance of the Gaussian filter from a certain range as augmentations.After pretraining on unlabeled PDE data, we fine-tune neural operators on simulated solutions of PDEs. We study two time-independent PDEs (Poisson, Helmholtz) and two time-dependent PDEs (Reaction-Diffusion, Navier-Stokes). We include details of these PDEs in Appendix A.We consider two popular architectures for fair comparisons with previous works. These are encoderdecoder architectures designed to reconstruct the original input given partial observations, where the encoder maps observed unlabeled PDE data to a latent space, and the decoder reconstructs the original conditions. We include visualizations of these architectures in Appendix D.Fourier Neural Operator. Fourier Neural Operator (FNO) targets learning PDE data in the Fourier space. The original model backbone (encoder) employs Fourier transform and learns lower Fourier modes with linear transforms. The FNO backbone outputs features back to the spatial domain (i.e., the embeddings are on the pixel level). We refer readers to the original paper for details [45, 46] .• Pretraining: We build the decoder to be identical to the encoder (except for the input/output dimension). Unlabeled PDE data are randomly masked at the pixel level.• Fine-tuning: After the pretraining, we discard the decoder, and we follow the original design to append two fully-connected layers (with ReLU activations) to predict final spatial-wise solutions.Transformer. Transformers, which mainly employ self-attention and linear transform blocks, have shown promise in both NLP and CV [78, 13] . Different from FNO, which directly operates on grids, transformers tokenize and group grids into patches, i.e., each tokenized patch embeds a local neighborhood of subgrids. We follow the 3D transformer architecture of Video-MAE [77, 52] . Our encoder embeds patches by a linear projection with added positional embeddings (just as in a standard ViT), and it then processes the resulting set via a series of Transformer blocks. For the transformer, the unlabeled PDE data are randomly masked at the patch level.• Pretraining: For the transformer encoder, we only apply it on the subset of tokens that are visible (i.e., unmasked patches), and masked patches are removed. This allows us to train encoders efficiently. The input to the MAE decoder is the full set of tokens consisting of (i) encoded visible patches, and (ii) the mask token. The mask token is a shared and learned vector that indicates the presence of a missing patch to be predicted. We add positional embeddings to all tokens in this full set. Without this, mask tokens would have no information about their location in the input. Following [28, 77] , we adopt an asymmetric design where the decoder is more lightweight (shallower and narrower) than the encoder.• Fine-tuning: After the pretraining, the decoder is preserved during fine-tuning, since we need to reconstruct the tokenized patches back to the input.Out-of-distribution (OOD) generalization is a critical technical challenge, not only in SciML but also across multiple domains in AI for science [84] . To improve the OOD generalizability of neural operators and to reduce the extra effort of downstream fine-tuning, the following inference paradigm has been proposed: given a query input, the model is also provided with a few supporting examples (dubbed "demos"), together with their ground-truth solutions, to make the final prediction. This approach enables the "open-set" generalization of the model to make predictions on unseen samples.Originally, in the literature on few-shot learning [82, 83, 55, 70, 35, 48, 61] Similarity by Prediction. We find spatial-wise and temporal-wise similar demos by calculating their distance in the output space. That means, for two input locations over the spatial and temporal domains, if we find their outputs of the trained neural operator similar, then we treat them as similar samples. Following [80, 81] , we assume demos share the same distribution of physical parameters with the query.Aggregation. For each spatial-temporal location of the query, after finding its similar samples in demos, we aggregate and average their solutions as the prediction. Data resolution: x-axis (W ), y-axis (H), output temporal steps (T ), output channel dimensions for the solution (Cout). Input: Query input (x). Paired unlabeled PDE data (X) and solutions (Y ∈ R J×H×W ×T ×C out ) as J demos. Trained Neural Operator Model M. TopK (k) demo solutions to aggregate. ▷ Spatial-wise and temporal-wise aggregation of solutions from similar demos. Shape:ŷ = M(x) ▷ Shape: H × W × T × Cout Ŷ = M(X) ▷ Shape: J × H × W × T × Cout δ = ŷ.reshape(-1, 1, Cout) -Ŷ .reshape(1, -1, Cout) ▷ Query-Demo Distance. Shape: H × W × T × (J • H • W • T ) × Cout δ = absolute( δ).H × W × T × Cout × k. Return: ŷicl .mean(-1) 
empirical results: 
 To illustrate the benefits of our approach, we perform empirical evaluations on both PDE benchmarks and real-world observations. Most importantly, our unsupervised pretraining (on unlabeled PDE data, followed by fine-tuning) outperforms neural operators trained from scratch, while requiring fewer PDE simulations (Sec. 4.1 and Sec. 4.2). Moreover, our in-context examples can help the model generalize better to OOD cases (Sec. 4.3). In our experiments, we trained models three times with different random seeds for statistical significance. For more experimental details, see Appendix B. We also include ablation studies about pretraining hyperparameters in Appendix J. For visualizations of our pretraining, see Appendix K. Data Efficiency. We first demonstrate that, by leveraging unsupervised pretraining, neural operators can achieve improved errors with less simulated data. In Figure 3 , on each PDE, compared with directly training from scratch ("random init."), pretraining on unlabeled PDE data can help neural operators achieve better performance, which can further help reduce the amount of simulated data. Specifically, in our experiments, when we target achieving the best test error of the baseline ("random init."), our method can save 5 × 10 2 ∼ 8 × 10 5 simulated solutions across diverse PDE systems.Among all PDEs, we find that Helmholtz (Figure 3 (b)) is the most challenging. Both two curves failed to improve the error until we increased the number of simulated data points to over 1024. Meanwhile, the generalization gaps remain high (Figure 12 (b)), indicating low training errors. We suspect that learning on the Helmholtz equation may be exhibiting the grokking issue [63, 86] , where the network quickly memorizes the training data, but the improvement in generalizability is delayed.Unsupervised Pretraining Outperforms Off-the-Shelf Pretrained Checkpoints. Pretraining on unlabeled PDE data is not the only way to save simulation costs. As pretraining is widely adopted in CV, pretrained checkpoints on vision data become publicly available and are ready for fine-tuning. We choose to compare with Video-MAE [77] for state-of-the-art video understanding pretrained on SSV2 [21] . As shown in Figure 3 (e), vision-pretrained Video-MAE can only outperform the random initialization with high volumes of simulated data, while its performance suffers when fine-tuned with limited simulations. In contrast, our unsupervised pretraining on unlabeled PDE data can save a significant amount of simulated data.As errors during testing may quickly accumulate with further timestamps, we also report results with more unrolled steps. We use checkpoints trained with the largest amount of simulated data from above for this study. As shown in Figure 3 (f), at each rollout step, our unsupervised pretraining achieves much better performance. Similarly, vision-pretrained Video-MAE is eventually outperformed by the random initialization at the long rollout step.Benefits of Pretraining on Unlabeled PDE Data. Pretraining on unlabeled PDE data is beneficial beyond achieving better performance with reduced simulations. First, when training on extremely low volumes of data, neural operators tend to overfit, resulting in poor generalization. In Figure 4 -left, pretraining on unlabeled PDE data can reduce the generalization gap (testing errortraining error). We further show that better generalization gaps persist across all PDEs we studied (see Figure 12 ). Second, pretraining on unlabeled PDE data can lead to faster convergence during fine-tuning. Unlike standard random initializations from Gaussian distributions, pretraining on unlabeled PDE data will provide neural operators with domain-adapted initializations and facilitate a much faster convergence rate, as shown in Figure 4 -middle 3 . Third, unsupervised pretraining can also help models extract useful representations for subsequent operator learning. From Figure 4 -right, we find that even with pre-extracted features (i.e., fixed encoder and only fine-tuned decoder), our neural operators can still outperform the baselines where both the encoder and decoder are updated during fine-tuning. We now move to a broader range of benchmarks. We will study real-world and noisy data instead of toy datasets, providing even more comprehensive experiments. These benchmarks are widely studied in previous works [60, 62, 52] .Datasets. We brief the background, with more details in Appendix L and visualizations in Fig. 15 .• ECMWF Reanalysis v5 (ERA5) [30] is a public extensive dataset, which delivers hourly data on multiple atmospheric variables spanning from 1979 to today. ERA5 represents a type of atmospheric reanalysis dataset [34] , integrating observations from a variety of measurement sources with numerical models through data assimilation [33] . Essentially, it reconstructs the most accurate estimation of the Earth's atmospheric conditions over time. This dataset has been extensively utilized in prior SciML studies [60, 68] . We focus on forecasting the important and challenging temperature atmospheric variable.• ScalarFlow [14] is a reconstruction of real-world smoke plumes. It assembles the first large-scale dataset of realistic turbulent flows. The availability of a large, volumetric data set opens up a wide range of applications, including re-simulations, novel visualization, and metric evaluations.The dataset contains 104 real-world smoke flow density reconstructions. Each reconstruction is captured from five different viewpoints for 150 temporal frames spanning 2.5 seconds.• Airfoil [75] is a large-scale dataset that contains the pressure and velocity simulations of the flow around real airfoils. This dataset has 53880 samples for training and 90 samples for testing. Each sample contains 6 channels. The 3 input channels include the binary spatial mask of the airfoil and the initial velocity of the freestream condition in the x and y directions. The output channels include the pressure and the x and y velocity components from the simulation at the steady state. 3 column "N.S. (PDEBench)". For the time-independent steady-state Airfoil dataset (Sec. L.3), we adopt the 2D-FNO architecture. We train 2D-FNO with Adam, with other hyperparameters the same as in Table 3 column "Poisson".Results. As shown in Figure 5 , compared with directly training operators from scratch ("random init."), pretraining on unlabeled data (2D snapshots of ERA5/ScalarFlow without temporal dynamics, or freestream velocities of Airfoil) can help neural operators achieve better performance on both temperature/flow forecasting and predictions of the steady-state pressure and velocity around airfoils.Figure 5 : For real-world scientific problems, pretraining neural operators on unlabeled PDE data improves its performance and data efficiency. We study VideoMAE [77] pretrained with unlabeled snapshots (no temporal dynamics), and then fine-tune across different numbers of temporal snapshots on ERA5 (left) and ScalarFlow (middle). We also pretrain 2D-FNO [45] on freestream velocities and fine-tune on time-independent steady-state airflow pressure and velocities (right). "random init.": models are trained from scratch with random initialization.We now move to OOD settings, where models will be tested on PDE data simulated with physical parameters unseen during fine-tuning/training. We include how to simulate OOD samples in Appendix B.2. Neural operators suffer from poor OOD generalization [73, 80] . Traditionally, improvements heavily depend on further fine-tuning on simulated data, which requires extra simulation and training costs. We study the benefits of leveraging test-time in-context examples. As shown in Figure 6 , when we flexibly scale up the number of demos, we can keep improving FNO's OOD generalization on diverse PDEs. We follow [80, 81] that demos are randomly sampled from the same distribution used to generate the OOD test set. When the number of demos is 0, we have the baseline in the OOD setting. Notably, we introduce zero training overhead: we keep the standard training pipeline, and our mining of in-context examples can be seamlessly plugged in during OOD inference.We further provide a baseline, which uses features extracted by the backbone of the neural operator (high-dimensional features before the final output layer) to find similar samples. As we can see, this baseline is worse than our method (both performance and confidence), indicating that the final output of the neural operator can more accurately indicate true similar samples.See Appendix M for further discussions about the benefits of leveraging in-context examples, and Appendix N for visualizations.  
conclusion: 
 In this work, we focus on improving the data efficiency of solving partial differential equations (PDEs) using deep learning, with a particular emphasis on unsupervised pretraining and in-context learning (ICL) methods. Our key contributions include introducing unsupervised pretraining for operator learning and a flexible ICL approach that enhances out-of-distribution (OOD) generalization without increasing training costs. Through extensive evaluations, we demonstrate that our method is not only more data-efficient, but it also achieves greater generalizability compared to existing approaches. By improving the data efficiency of neural operators for solving PDEs, our approach can significantly reduce the computational costs and energy demands of high-fidelity numerical PDE simulations. Additionally, by making advanced PDE solutions more accessible through efficient pretraining, our method has the potential to accelerate scientific and engineering progress across various fields, ultimately benefiting society. We hope our work will inspire the scientific machine learning (SciML) community to further address the high simulation costs and limited OOD generalization of neural operators, contributing to advancements that support both scientific innovation and environmental sustainability. 
limitations: 
 Current limitations of our work: 1) We could design more physics-inspired proxy tasks and data augmentation methods for scientific data; 2) We could study more PDE systems in our unsupervised pretraining and in-context learning; 3) We could consider more different neural operator architectures. We expect that addressing these limitations will lead to broader impacts in future works.ν signifies the viscosity of a fluid, and f is the external forcing function. Dirichlet boundary conditions are employed in this system. Unlabeled PDE Data: We follow the training settings from previous works. 1) For FNO [45] , we focus on mapping the vorticity from the initial condition to the full solution ( w 0 → w| [0,T ) , with T = 33). The input includes four channels: the vorticity (which is the initial condition, duplicated for T times), and three (xy-spatial and temporal) linear mesh position embedding channels. Thus, similar to Reaction-Diffusion above, during pretraining, the model will take individual snapshots of vorticity and position embeddings as unlabeled data. 2) For PDEBench [74] , we forecast both xy velocities and pressure (i.e. three input channels), with T = 15. The model will take individual snapshots of vorticity and xy velocities as unlabeled data.We summarize detailed inputs and outputs in Table 1 . 4 For Helmholtz OOD, we choose a narrow range of coefficients ( [15, 20] ) mainly because its solution is very sensitive to the wavenumber, and FNO's performance significantly drops when we move to more extreme OOD settings.  
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] .Justification: This work is about foundational research on data augmentations. It does not lead to any societal impact because: 1) our motivations, goals, and results are too distant from social applications to make any social impact; 2) we mainly use public data and benchmarks; and 3) we use public model architectures and follow previous works which have open access.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). 
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 poses no such risks because all datasets studied in our work are from public benchmarks, and we have clearly cited original papers.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: All datasets studied in our work are from public benchmarks, and we have clearly cited original papers.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. • 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. 
 
 Appendix: Unlabeled PDE Data. We generate data for Poisson and Helmholtz [73] , Reaction-Diffusion on PDE-Bench [74] and 2D incompressible Navier-Stokes on PINO Dataset [46] following the procedure mentioned in the paper. For unlabeled data generation, we bypass the computation of solvers, which expedites the generation speed, as shown in Table 4 .OOD Samples. The OOD data generation procedure is similar to the unlabeled data, except for the changes in the physical parameters coefficients. For Poisson and Helmholtz, we consider changing the range of diffusion eigenvalue and waver number respectively. For Navier-Stokes equation, we change the Reynolds number. We list the coefficients in Table 2 .We summarize our hyperparameters used during pretraining and fine-tuning/training in Table 3 . These hyperparameters strictly follow previous works [45, 73, 74, 46, 52] . We conducted our experiments on four A100 GPUs, each with 40GB of memory. In Table 4 , we demonstrate the cheap simulation of only unlabeled PDE data, versus simulating both unlabeled PDE data and solutions, on 2D incompressible Navier-Stoke on PINO Dataset [46] and Reaction-Diffusion on PDE-Bench [74] . We can see that unlabeled PDE data are extremely cheap to simulate. Therefore, our pretraining method can boost the performance and meanwhile save the heavy cost of data simulations. Contrastive learning is an important self-supervised pretraining technique studied in computer vision. For a fair comparison, we directly compare with MoCo v2 [8] , a highly-cited self-supervised learning method also originally implemented in PyTorch, whose core method is closely related to SimCLR [7] (originally implemented in TensorFlow).We compare MoCo v2 with our method on a broader real-world benchmark ERA5 [30] . As shown in Figure 8 , our unsupervised learning method can largely outperform MoCo v2. This extra comprehen- sive result demonstrates that our method can be widely adopted in real-world problems, outperforming previous unsupervised learning methods. Beyond Figure 3 (e) for the transformer on time-dependent PDE (Navier-Stokes), we further verify that pretraining on unsupervised PDE data makes FNO outperform its off-the-shelf vision-pretrained checkpoints. Specifically, we first pretrain the 2D FNO model on ImageNet [11] , then fine-tune on the downstream PDE simulation data. As shown in Figure 9 , vision-pretrained FNO performs worse during downstream fine-tuning on time-independent PDEs (including both Poisson and Helmholtz equations), confirming that domain-specific pretraining, even on unsupervised PDE data, is more beneficial than conventional checkpoints pretrained on unrelated domains like computer vision.We also study if joint pretraining on unlabeled multiple PDE data can bring extra benefits. We combine all 46,080 unlabeled Poisson samples and all 46,080 unlabeled Helmholtz samples (see Table 3 ). We choose this setting because recent works on SciML foundation models [52, 26] also use all samples from each PDE for pretraining. We do zero-paddings for mismatched channels. From Figure 10 , we can see that joint pretraining can further improve the performance of fine-tuning on different PDEs. We also try to fine-tune neural operators on unseen PDEs (i.e. PDEs different from pretraining). Mismatched channels are padded with zeros. We find this will lead to worse performance compared with models pretrained on the same PDE, as shown in Figure 11 . Beyond the generation gap we have shown in Figure 4 (left) for the Navier Stokes equation from PDEBench, we further collect generation gaps of our models learned on other PDEs. As shown in Figure 12 , on diverse PDE systems, our method can contribute to universally reduced overfitting. We study the optimal magnitude of perturbations during pretraining: "Mask Ratio" (1 indicates no visible grids, and 0 means no masks); and "Blur Sigma" for the variance of Gaussian kernel for blur (larger the more degradation). We show our results in Table 5 , 6, 7, 8. For example, on 2D incompressible Navier-Stokes with FNO, as shown in Table 8 , we can find that when training with a low volume of data, we should use much stronger perturbations (high masking ratios and strong blur), whereas a high volume of data only requires mild perturbations. As shown in Table 9 , the more unlabeled PDE data we use for pretraining, the better quality the pretrained model will be. To demonstrate the efficacy of our MAE-based pretraining, we show the unlabeled PDE data and its reconstructed version in Figure 13 (MAE pretraining on different PDEs) and Figure 14 (MAE pretraining on 2D incompressible Navier-Stokes with varying mask ratios). We can see that all inputs are accurately reconstructed with low errors and similar patterns. We utilize data from 2006 to 2015, with the snapshots taken every 6 hours and a spatial resolution of 360 × 360. The total number of snapshots is 14600. We apply the mean-standard deviation normalization to the data and downsample the snapshots to a spatial resolution of 180 × 180. We split 75% of the data for pretrain and 25% for finetune. For each split, we further separate 80% of the data for training, 10% for validation, and 10% for testing.The original spatial resolution is 1062 × 600. We crop the snapshots to 900 × 600 to remove the padding and background. We remove the first 15 timeframes for each simulation to avoid the initial transient phase at the beginning of the smoke flow generation. We then downscale the snapshots to 180 × 120 and apply mean-standard deviation normalization to the data. With a total of 70200 snapshots, we split 80% of the data for pretrain and 20% for finetune. For each split, we further separate 80% of the data for training, 10% for validation, and 10% for testing.We split 80% of the training data for pretrain and 20% for finetune. The original spatial resolution of each sample is 128 × 128. We crop the snapshots to 96 × 45 to remove the background. We then apply min-max normalization channel-wise to the sample.We show visualizations of real-world scientific data in Figure 15 . We try to understand the benefit of leveraging in-context examples by decomposing the relative MSE error into "scale" and "shape." "Scale" is the slope of a linear regression between targets and predictions (closer to 1 the better), indicating the alignment of the range of model outputs with targets. "Shape" is the normalized relative MSE (i.e., model outputs or targets are normalized by their own largest magnitude before MSE), indicating the alignment of scale-invariant spatial/temporal structures. We show results in Figure 16 . We find that the benefit of in-context examples lies in that the scale of the model's output keeps being better calibrated ("scale" being closer to 1) when adding more demos.In numerical simulations or predictions of PDEs, there are settings where the scale or magnitude of solutions is more important than the exact shape/pattern:1. Heat Transfer: In large-scale systems, the focus might be on overall temperature and extreme values. For instance, in evaluating a cooling system, the key concern might be the peak temperature rather than the detailed temperature distribution. 2. Fluid Dynamics: For applications like aerodynamics, the overall drag or lift force on an object is often more critical than capturing every detail of the flow pattern, such as in airfoil design. 3. Environmental Modeling: The concentration of pollutants at specific locations or total pollutant transport is often more crucial than the exact distribution, such as in groundwater flow studies. Justification: Our work does not involve any potential harm, societal impact, or violation.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).In this paper, we use the terms "unsupervised pretraining" and "self-supervised pretraining" interchangeably.Training curves when the number of training samples of 14760.For Navier Stokes from PDEBench, we use the original data. We could not generate our own pretraining/finetuning data with different ranges of physics parameters, due to a possible mismatch of the provided configuration files and the version of Phiflow used in PDEBench (see GitHub issue at https://github.com/pdebench/PDEBench/issues/36). 


Decision: Dear Authors,

Thank you for submitting the draft.

Both reviewers agree that the presented work is interested. However, both reviewers do also raise some major points of concern, especially regarding the clarity of the presentation. It is expected that authors will be addressing comments by the reviewers in the final draft. 

regards

AC

Reviews: review: title: title: Insufficient Depth, as too Many Topics are Combined into one Paper
main_review: main_review: ##  Summary of Paper
The paper aims to demonstrate two things: First the authors aim to reduce the simulation-cost in producing vasts amounts of examples to train neural operators. To this end, they propose to create data examples that do not require to solve the differential equation at hand. On this data they train the neural operator to perform common pre-training tasks such as gap-filling, de-noising or de-blurring. 
Second, the paper proposes in-context learning for neural operators. The motivation is taken from the emerging abilities of LLM that can solve new tasks when prompted with very few examples. Specifically, they propose to find examples with similar **outputs** and replace the prediction for the sample at hand by the average of the sample and the similar examples.
The authors present empirical evaluations for both of these methods.
Since both the contributions are mostly independent, I will in the following refer to the first as Pre-Training (PT) and the second as In-Context-Learning (ICL).


## Fit for this Workshop
I think that both contributions are interesting for this workshop. While I think that the ICL contribution might be more uncertain overall, I think that the PT contribution will be of great interest to a lot of participants at this workshop.

## Evaluation of the General Structure and Idea
### Pre-Training
I think the pretraining topic is interesting. The authors claim that "Unsupervised pretraining on unlabeled data has never been explored in PDE operator learning in SciML." Independent of, whether this claim is true, I think pretraining for PDE operator learning SciML is not that different to PDE operator learning in other domains and the solutions proposed in this paper are not different to pretraining of neural networks in general. To this end, the article needs a discussion of related  work. I do not agree with the authors that this discussion can be put into the appendix, especially if it is not referenced in the article once. 
Further, the explanations in Section 2.1.1 need more explanations. Especially the first bullet point is still unclear. On the other hand, The statement of the second-order linear differential equation does not add much to understanding the paper. 
Finally, empirical evaluation is way to short. No details are provided. In the current form, The empirical results are difficult to interpret and almost impossible to reproduce as no information on the architectures or training algorithms is given. It is unclear which pretraining and how mch pretraining is done in each example. This is especially important, as the random initialization results fall short of the values of other publications. For example [1] reports an error of 0.005 at 10 000 examples for the Poisson equation beating the numbers reported for the random initialization by an order of magnitude and being on par with the pretraining.
The description of the Proxy tasks (2.1.2) is nice and I appreciate the fact that the authors do not only introduce the tasks but share their intuition why these tasks should improve the performance.

### In-Context-Learning
I am not sure about this idea of ICL. I understand the motivation, but I see a big discrepancy between the motivation and the actual implementation. In the current form, the method seems to be similar to imposing a prior on the distributions where predictions by the operator are corrected towards areas of high density in the prediction space. If the solution of the PDE is continous, this is further similar to test time augmentation (eg. [2]). I think that both of these links need to be discussed and evaluated.
The description in Section 2.2 is way to short and general. The empirical results have the same shortcomings as the empirical evaluations of the PT approach. They need a lot more explanation and introduction than provided by the authors. In the current form, the results are not convincing.

## Structure and Length of the Paper
I know that it is challenging to explain a complex topic in four pages. This makes it even more difficult to put two paper into four pages. I think the paper is not long enough to discuss both topics well enough to allow the reader to follow and unfortunately, I do not think that any of the two proposed advances are well enough discussed to be helpful to the reader. This is underpined by the fact that the authors provide multiple Appendices (A1, A2, A3, B, C, D1, D2, D3, E, F, G, H, I, J1, J2, K1; 9pages) only three of which are ever referenced in the paper.
I would recommend to focus on one of the two topics instead. 

## Conclusion
In conclusion I think the paper tries to do to much on four pages to allow for the needed scientific rigor. I find especially the PT topic interesting but I can not follow the authors arguments and conclusions in this paper. I think that the discussion of the pre training tasks in 2.1.2 is a highlight but especially the empirical evaluation is to poorly described to understand whether there is any effect. Therefore, while I think the underlying research is interesting and might be well done, the paper in the current state is not good enough. I recommend to reject this paper.

##References
[1] Hasani E, Ward RA. Generating synthetic data for neural operators. arXiv preprint arXiv:2401.02398. 2024 Jan 4.

[2] Perez F, Vasconcelos C, Avila S, Valle E. Data augmentation for skin lesion analysis. InOR 2.0 Context-Aware Operating Theaters, Computer Assisted Robotic Endoscopy, Clinical Image-Based Procedures, and Skin Image Analysis: First International Workshop, OR 2.0 2018, 5th International Workshop, CARE 2018, 7th International Workshop, CLIP 2018, Third International Workshop, ISIC 2018, Held in Conjunction with MICCAI 2018, Granada, Spain, September 16 and 20, 2018, Proceedings 5 2018 (pp. 303-311). Springer International Publishing.
title: title: A promising contribution for reaching data efficiency, low computational costs and to avoid overfitting.
main_review: main_review: This work is well written and it is a significant state of the art contribution. It is well supported with up-to-date literature on the topic and the results look very impressive in terms of the magnitude of the errors. 
I feel that more details should have been included in the explanation of the methods to be a more appealing work.

Pros: 
* The authors clearly highlight the novelty of their contribution and the accuracy of the methods proposed. 
* The methods achieve data efficiency and low computational costs.
* It is a nice example of how a model can be trained when having low amounts of data avoiding overfitting.
* The magnitude of the errors achieved after the training is impressive.

Cons:
* I felt the lack of details in the methods. It would have been nice to add more explanations/reasoning about the selection of the input data and the split of the data into training/validation/testing samples.
* It seems that they did not conduct a hyperparameter tuning. I wonder if the results would have changed significantly with different hyperparameters.
* More context about the physical meaning of the parameters in the methods to inform the machine learning-based model is also missing. For instance, the ranges shown in Table 1 Section D2 are not properly explained. I feel more discussion about this topic is needed because they mention in the abstract the promising venues of coupling machine learning methods and physical domain-specific insights.