Poseidon: Efficient Foundation Models for PDEs
Authors: Maximilian Herde
Bogdan Raonic
Tobias Rohner
Roger Käppeli
Roberto Molinaro
Emmanuel de Bezenac
Siddhartha Mishra
Venue: NeurIPS 2024 poster (2024)
Keywords: PDEs
operators
foundation models
transformers
sample efficiency
Primary area: machine_learning_for_physical_sciences

Abstract: We introduce Poseidon, a foundation model for learning the solution operators of PDEs. It is based on a multiscale operator transformer, with time-conditioned layer norms that enable continuous-in-time evaluations. A novel training strategy leveraging the semi-group property of time-dependent PDEs to allow for significant scaling-up of the training data is also proposed. Poseidon is pretrained on a diverse, large scale dataset for the governing equations of fluid dynamics. It is then evaluated on a suite of 15 challenging downstream tasks that include a wide variety of PDE types and operators. We show that Poseidon exhibits excellent performance across the board by outperforming baselines significantly, both in terms of sample efficiency and accuracy. Poseidon also generalizes very well to new physics that is not seen during pretraining. Moreover, Poseidon scales with respect to model and data size, both for pretraining and for downstream tasks. Taken together, our results showcase the surprising ability of Poseidon to learn effective representations from a very small set of PDEs during pretraining in order to generalize well to unseen and unrelated PDEs downstream, demonstrating its potential as an effective, general purpose PDE foundation model. Finally, the Poseidon model as well as underlying pretraining and downstream datasets are open sourced, with code being available at https://github.com/camlab-ethz/poseidon and pretrained models and datasets at https://huggingface.co/camlab-ethz.

POSEIDON: Efficient Foundation Models for PDEs
Maximilian Herde
Bogdan Raonicl22
Tobias Rohner' 1
Roger Käppeli
Roberto Molinaro
Emmanuel de Bézenac*
Siddhartha Mishra
Seminar for Applied Mathematics, ETH Zurich, Switzerland
'ETH AI Center, Zurich, Switzerland Correspondence to herdemeethz.ch
Abstract
We introduce POSEIDON, a foundation model for learning the solution operators of PDEs. It is based on a multiscale operator transformer, with time-conditioned layer norms that enable continuous-in-time evaluations. A novel training strategy leveraging the semi- semi-group property of time-dependent PDEs to allow for significant scaling-up of the training data also proposed. POSEIDON pretrained on a diverse, large scale dataset for the governing equations of fluid dynamics. It is then evaluated on suite of 15 challenging downstream tasks that include a wide variety of PDE types and operators. We show that POSEIDON exhibits excellent performance across the board by outperforming baselines significantly, both in terms of sample efficiency and accuracy. POSEIDON also generalizes very well to new physics that not seen during pretraining. Moreover, POSEIDON scales with respect model and data size, both for pretraining and for downstream tasks. Taken together, our results showcase the surprising ability POSEIDON to learn effective representations from very small set of PDEs during pretraining in order to generalize well to unseen and unrelated PDEs downstream, demonstrating its potential as an effective, general purpose PDE foundation model. Finally, the POSEIDON model as well as underlying pretraining and downstream datasets are open sourced, with code being available at http:/huww.ooo//maaabbeeeeeeeeae and pretrained models and datasets at https://tugginffccc.oo/mmmm http://aww....cccabbcebe.aaa
Introduction
are referred as the language of physics as they mathemat- ically model very wide variety of physical phenomena across vast range of spatio-temporal scales. Numerical methods such as finite difference, finite element, spectral methods etc. [59] are commonly used to approximate simulate PDEs. However, their (prohibitive) computational cost, particularly for the so-called many-query problems [58), has prompted the design of various data-driven machine and references therein. Among them, operator
Partial Differential Equations (PDEs) 15
learning (ML) methods for simulating PDEs, 24 51
learning algorithms have gained increasing traction in recent years.
These methods aim to learn the underlying PDE solution operator, which maps function space inputs (initial and boundary conditions, coefficients, sources) to the PDE solution. They include algorithms which approximate discretization, on fixed grid, of the underlying solution operator. These can be based on convolutions [75] [8), graph neural networks [56] [65 or transformers Other operator learning algorithms are neural operators which can directly
12] 55
26
20 33)
"Equal contribution
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

process function space inputs and outputs, possibly sampled on multiple grid resolutions 227
3
These include DeepONets [3 42], Fourier Neural Operator [3]. SFNO 7 Geo-FNO 32 Low- rank NO 3
and Convolutional Neural Operator [60] among many others.
However, existing operator learning methods not sample efficient as they require very large number of training examples to learn the target solution operator with desired accuracy (see Figure or Figure of [60). This impedes their widespread use as task-specific training data is very expensive to generate either with numerical simulations or measurements of the underlying physical system.
Task-specific
POSEIDON: Foundation Model for PDEs
Operator Learning
pretrain
NS model
+
ScOT
compressible incomp. flow
compressible incomp.
flow
finetune on out-of-distriuutio downstream tasks
Comp. Euler model
Posckion U
Belasive
Wave model
10
10
Number of trsjectonis
Poisson model
64
64
0096
64
0096
NS
Comp. Euler
Wave
Poisson
Figure 1: As opposed to PDE-specific operator learning, our pretrained model POSEIDON is up to multiple orders magnitude more sample efficient than task-specific neural operator while also being able to transfer to unseen physics during finetuning.
How can the number of training samples for PDE learning be significantly reduced? In this context, can we learn from language modeling and computer vision where similar question often arises and the current paradigm is to build foundation models [6] These generalist models are pretrained, at-scale, large datasets drawn from diverse set of data distributions. They leverage the intrinsic ability of neural networks to learn effective representations from pretraining and are then successfully deployed on a variety of downstream tasks by finetuning them on a few task-specific samples. Examples of such models include highly successful large language models
[TO] [72,, large multi-
modal models [7 52 and foundation models for robotics [9] chemistry [4, biology 663 medicine 664
and climate
(44)
66]. the challenge of designing such foundation models for PDEs is formidable, given the sheer variety of PDEs (linear and nonlinear, steady and evolutionary, elliptic, parabolic, hyperbolic and mixed etc.), the immense diversity of data distributions, wide range of underlying spatio-tempore scales and the paucity of publicly available high-quality datasets. In particular, the very feasibility of designing PDE foundation models rests on the fundamental and unanswered science question of why pretraining model on a (very) small set of PDEs and underlying data-dit dist-iiitiibutionss can allow it to learn effective representations and generalize to unseen and unrelated PDEs and data-distributions via finetuning?
Despite very recent preliminary attempts 67
68
The investigation of this open question motivates here present the POSEIDON family of PDE foundation models. POSEIDON, see Figures and is based on i) scalable Operator Transformer or scOT, multiscale vision transformer with (shifted) windowed for operator learning, ii) novel all2all training strategy for efficiently leveraging trajectories of solutions of time-dependent PDEs to scale up the volume of training data and iii) an open source large-scale pretraining dataset, containing a set of novel solution operators of the compressible Euler and incompressible Navier-Stoker Stokes equations of fluid dynamics. We evaluate POSEIDON on challenging suite of 15 downstream tasks, comprising of well rell-established benchmarks in computational physies that encompass linear and nonlinear, time-dependent and independent and elliptic, parabolic, hyperbolic and mixed type PDEs. All of these tasks are out-of-distribution with respect to the pretraining data. Moreover, nine out of the 15 tasks even involve PDEs (and underlying physical processes) which are not encountered during pretraining.
Swin attention [38 37 adapted

---

Through extensive experiments, we find that i) POSEIDON shows impressive performance across the board and outperforms baselines on the downstream tasks, with significant gains in accuracy and order of magnitude gains in sample efficiency. For instance, on an average (median) over the downstream tasks, POSEIDON requires mere 20 samples to attain the same error level as the widely-used FNO does with 1024 samples. ii) These gains in accuracy and sample efficiency are also displayed on tasks which involve PDEs not encountered during pretraining, allowing us to conclude that POSEIDON can generalize to unseen and priori unrelated physical processes and phenomena with a few task-specific training examples and iii) POSEIDON scales with model and dataset size, both for the pretraining well as for downstream tasks and iv) through case studies, elucidate possible mechanisms via which POSEIDON is able to learn effective representations during pretraining, which then leveraged generalize unrelated PDEs downstream. Taken together, these results provide the first positive answers to the afore-mentioned fundamental question of the very feasibility of PDE foundation models and pave the way for the further development and deployment of POSEIDON an efficient general purpose PDE foundation model. Finally, we also open source the POSEIDON model and the entire pretraining and downstream task datasets within the PDEGYM database.
2 Approach
Problem Formulation. We denote a generic time-dependent PDE as,
0,u(2, t) + C (u, V V.....) = 0, Vx E D C g E (O,T), ((u)= = 0, v(z, t) E aD x (0, m(0,x) = a(x), x x
(1)
Here, with function space I c R") for some 1 1 oc, C(O0, is the solution of (T) a E X the initial datum and L, B are the underlying differential and boundary operators, respectively. Note that (T) accommodates both PDEs with high-order time-derivatives as well as PDEs with (time-inde (ime-independent) coefficients and sources by including the underlying functions within the solution vector and augmenting L accordingly (see SM B.2 for examples).
Even limeeindeeendet PDEs can be recovered from D by taking the long-time limit, i.e., lim; T7 which will be the solution of the (generic) time-independent PDE,
L (T(z), V U. V
0, Vx E D. B 0, Vr OD.
(2)
Solutions of the PDE ① are given in terms of the underlying solution operator S [0, x +++ X such that u(t) = 8(t, a) is the solution of (1) at any time & E (), Given a data distribution 1 E Prob(X), the underlying operator learning task (OLT) is,
OLT: Given any initial datum a ~ p. find an approximation S* =s to the solution operator 8 of ①). in order to generate the entire solution trajectory {8 (1, a)} for all 1. E [0, T).
essential to emphasize here that the learned operator 8* has to generate the entire solution trajectory for (I. given only the initial datum (and boundary conditions), as this is what the underlying solution operator 8 (and any numerical approximation to it) does.
Model Architecture. The backbone for the POSEIDON foundation model is provided by scOT or scalable Operator Transformer, see Figure (a-c) for an illustrated summary. scOT is hierarchical multiscale vision transformer with teae-time conditioning that processes lead time t and function space valued initial data input to approximate the solution operator S(t, a) of the PDE (T)
(a-c)
For simplicity of exposition, we set d = and D = [0, 1j2 as the underlying domain. As in a vision transformer [14), any underlying input is first partitioned into patches and (linearly) embedded into latent space. At the level of function inputs a E C(D; TR.),, this amounts to the action of the patch
partitioning and embedding operator v = E(a). with E defined in SM (12) This operator transforms the input function into piecewise constant function, which is constant within patches (subdivisions of the domain D) by taking weighted averages and then transforming these piecewise constant values into a C-dimensional latent space resulting in output v E C(D; CR.. In practice, a discrete version of this operator is used and is described in SMA.2 A.2
3

---

Atlention
Input
Output
Ag
Inc MLP HC+
scaled
Sanslarity
Value
Swin Block
Query Key!
Embedding
Recovery Mixup-
Softmax
ConvNeXt Blocks
dt - time-conditioned Layer Norm
Swin V2
Block i+
MLP
S2 Stage
Patch Expansion
time-conditioned Layer Norm
Patch Merging
Figure 2: (a) scOT, the model underlying POSEIDON; (b) V2 Transformer block; (c) Shifting Window over patch-based tokens with window (patch) boundaries with black (white); (d) all2all Training for time-dependent PDEs.
As shown in Figure (a), this patch embedded output then processed through sequence of SwinV2 transformer blocks 38 37, each of which has the structure of SW; C(D; + C(D;R(),
SWi(ve-) = v/ LN LN (MLP(v))),
V/2 Vz-1 + LN LN of (W -MSA(v--)).
for layer index 6 = I,..,, L. The main building block of SwinV2 transformer block (3) (sce Figure 2 (b)) is the windowed multi-head self attention operator defined in SM (14) (see SM A.2 for its discrete version). In particular, the attention operator acts only inside each window, which is defined by another (coarser) sub-division of D (see Figure (c)), making it more computationally efficient than standard vision transformer [14]. Moreover, the windows are shifted across layers, depicted
in Figure (c), so that all the points in the domain can be attended to, by iteratively shifting windows across multiple layers, SM A.2 for detailed description of the Swin V2 block.
The MLP in (3) is defined by SM (5). We follow [55 to propose a time-conditioning strategy by introducing lead-time conditioned layer norm
LN (v) a(t) o v(x) (v(
+ B(),
(v(
1
pv(x) = C Σ j=1
vy(c), (x) à
Here, att) at + ã and B(t) 3t + 3. with learnable a, a. 3, although more general (small) MLPs can also be considered. This choice of time embedding enables continuous-ing -nn-ss evaluations.
Finally, as depicted in Figure 2 (a), the S2 transformer blocks (3) are arranged in hierarchical, multiscale manner, within U-Net style encoder-decoder architecture [1I. by employing patch merging (downscaling) and patch expansion (upscaling) operations, (see SM A.2 for a detailed description). Moreover, layers the same scale, but within the encoder and decoder stages of scOT, respectively, are connected through ConvNeXi convolutional layers [39], specified in SMA.2
Training and Inference. We denote scOT by 8 [0, 7) x X < X, with trainable parameters E C IRP For scOT to approximate the solution operator of (T). the parameters 0 need to be determined by minimizing the mismatch between the predictions of scOT and ground truth training data, given the form of trajectories (S(te, a.)). for < < K and 1 i M, with a, ~ μ and 0 to < <...te ... < fx T, being the time points which the data is sampled. We assume that the data is sampled at the same timepoints for each sample a; for simplicity. For training, it is natural to consider the loss function,
K
I(0) = M(K
18 -(te, a))
(5)

---

with the (spatial) integral in ( being replaced by a quadrature at some underlying sampling points and p in our paper. Thus, we use K + samples per trajectory in order to train our model.
Given the fact that scaling up available training data is necessary for successful foundation models propose novel training strategy that further leverages the structure of the time-dependent to increase the amount of training data. To this end, we consider the modified loss function,
23
PDE
1 MK
Σ(0)
Σ
|S(te th: (tx))) St tkk u, u,
(6)
kkk
with (tx) = S(tR, a;) (approximately) solving and K [K+1)(+2) In other words, leverage the fact that the solution operator of ① possesses a semi- group property and one can realize,
u(tt) = S(t", a) = S(t" - t, ut)) = S(t" - t, a,)(t,))), v
(7)
and any initial condition a. We term this use of all possible data pairs (u(lz), u(tz)) with S k see Figure 2 (d) for visual representation, within trajectory as all2all training and observe that it allows us to utilize quadratic O ((²2) samples per trajectory, when compared to the linear K samples used for training corresponding to the vanilla loss function (5). In practice, consider relative form of Equation 6 balance out different scales of different operator outputs, see SM for details.
Once scOT has been trained with (stochastic) gradient descent to find (local) minimum 0" of the all2all loss function (6), the trained model, denoted as 8j. can be deployed for inference for any initial condition E X and for any E R,, by directly applying 8g- (t. a) to provide continuous-in-tin evaluation of the entire trajectory. However, it might be advantageous to infer using autoregressive rollouts (6). To this end, we consider a sequence 0 = 1/ < li
< 1x = L. Then, the rollout, a(t,a) R Sg- (t% 10--x-11 1 (. (......... Sg- (t2 11, 8/- ((Ga)))
(8)
of я successive applications of the trained scOT approximates the solution operator at any time
Pretraining. The key point in the development of any foundation model is the pretraining step. in which the model is trained on a diverse set of data distributions, rather than just data drawn from one specific operator. To formulate pretraining and subsequent steps precisely, we introduce index sets A, E and let À A and & E E correspond to indexing the PDE type and the data-distribution, respectively. To see this, we fix any X E A, & E E and tag the differential and boundary operators f, B in the PDE ① by LA and Bd Similarly the initial distribution p in ① is tagged by μ and the resulting solution operator for PDE (T) with L*, Ba and initial datum a ~ jo is denoted by s2.5, In other words, A, indexes the entire set of PDEs and data distributions that consider.
Next, fix index sets, À C A and _ =gt and consider set of PDEs [ indexed by À Å and with data distributions 115, indexed by E the pretraining dataset, which consists of the corresponding trajectories, {8*iE(t, }}, for all t and (x, E SA.
Let _ be the maximum number of components of the solution vectors for all the operators in the pretraining dataset. By including additional (constant O over space and time) components, augment the relevant solution operators (for which the number components is below n ) such that for each À E Å, E E E, all the input functions have the same number of n components (channels). These inputs are fed into a scOT model I [o. T x LP(D; Rn ) C((0, 7 LD;; RT )), with T being the supremum over all the final times in the pretraining dataset. The trainable parameters of this pretrained model are then determined by minimizing the mismatch between model predictions and ground truth over all PDES and data distributions in the pretraining dataset resulting
T = 3. with 0. = =argmingge
A|E
with A obtained by replacing S and 8 in (6) with 8A. and ITAS respectively.
Finetuning. To finetune the pretrained foundation model n. for any downstream task, correspond- ing any specific solution operator S for any À A. E, we decompose the vector of learnable
5

---

C R" as 6 = 00, 8 with & R2, H E R2, and @V E Rex and p+ PN with p. PN < p. A gradient descent step for finetuning is then written as,
parameters O E ④
Vr 1, 0,+1, 1/41, r+1 (0, B,, 0-0 -Te Ne. The I (o
(10)
0o 0., 0 0NN, 00~P, P, Prob(RP).
Hence, during finetuning, a subset of parameters A of the foundation model are trained from scratch with random initializations, whereas the complementary. much larger subset of 9 and 0N is initialized by transferring the corresponding parameters from the pretrained model. When À & A, 0 consists of the embedding/recoveed parameters. On the other hand, if À E A, then all trainable parameters, including the patch enbedddins/recovery. are initialized with the corresponding parameters of the pretrained model. However, the corresponding learning rate 77 > 7. in (T)) is much higher. Similarly, the time embeddings @N, i... the trainable parameters in the layer-norm operators 4 always initialized from the corresponding time embeddings in the pretrained model but finetuned with a higher learning rate v.a.
3 Experiments
Pretraining Dataset. We pretrain POSEIDON on dataset containing 6 operators, defined on the space-time domain [0, 1 x [0, 1]. 4 of these operators (CE-RP, CE-KH, CE-CRP, CE-Gauss) pertain to the compressible Euler equations (SM (37) of gas dynamics and 2 (NS-Sines NS-Gauss) to the incompressible Navier-Stokes equations (SM (II) of fluid dynamics, see SM Table 3 for abbreviations and SM B. for a detailed description of these datasets. These datasets have been selected to highlight different aspects of the PDEs governing fluid flows (shocks and shear layers, global and local turbulent features, and mixing layers etc.). The pretraining dataset contains 9640 and 19640 trajectories for the Euler and Navier-Stokes operators, respectively, leading to a total of 77840 trajectories. Each trajectory is uniformly sampled at 11 time snapshots. Within the all2all training procedure (Section[2) 2). this implies a total of 66 input-output pairs per trajectory, leading to approx 5. training examples in the pretraining dataset.
Downstream Tasks. To evaluate POSEIDON (and the baselines), we select a suite of 15 challenging downstream tasks, see SM Table 4 for abbreviations and SM B.2 for detailed description. Each of these tasks is a (variant of) computational physics literature and corresponds to distinct PDE solution operator. They have also been selected for their diversity terms of the PDE types as they contain linear (4) and nonlinear (11), time-dependent (12) and time-independent (3), elliptic (2), parabolic (1), hyperbolic (4) and mixed- type (8). The tasks also cover a wide gamut of physical processes across a range of spatio-temporal scales. Morcover, emphasize that each of the downstream tasks is auttefddistributionn with respect equations seen
benchmarks for PDES in the numerical analysis and
to the pretraining data. While 6 of them do pertain to the Euler and Navier-Stolie in the pretraining dataset but with very different data distributions, the remaining 9 involve PDEs not seen during pretraining. These include 3 ENS-Tracer-PwC. FNS-KF, GCE-RT) which add new physical processes (tracer transport, forcing, gravity) to the Havier-Stoker and Euler equations. more Gause, Wave-Layer, ACE) involve completely new time-dependent PDEs (Wave Equ., Allen-Cahn Eqn.) and the final 3 (SE-AF, Poisson-Gauss, Helmholtz) even consider time-independent PDEs, which is in stark contrast to the pretraining dataset where only 2 time-dependent PDEs are covered. For these steady state PDES, we finetune them by using the interpretation of the PDE (2) as a long-time limit of the time-dependent PDE ① with a normalized lead time of 1. Finally, the tasks have also been selected to probe the ability of the foundation model to handle different task or operator types. To this end, we point out that all the operators in the pretraining dataset simply map the initial conditions the solution at later times in time-dependent fluid flows on the mro-dimensional unit square with periodic boundary conditions. While some of the downstream tasks (8 out of 15) do pertain this type of operators, the remaining (7 out of 15) tasks involve different types of operators which include operators mapping the coefficients or PDE parameters to the PDE solution (5 out of 15), forcing term to the PDE solution (2) and domain shape to the PDE solution. Moreover, many of the downstream tasks are with non-periodic boundary conditions while one of them is even on a mon-Cartesian domain. Thus, these downstream tasks deviate from the setup of the pretraining operators and provide a hierarchy of challenges for any foundation model.

---

Models and Baselines. We consider three different POSEIDON models: i) POSEIDON-T with 21M parameters, ii) POSEIDON-B with = 158M parameters, and iii) POSEIDON-L with 629M parameters. The detailed specifications of each of these models is provided in SM C.T As baselines, in addition to the standalone ScOT, we use trained from scratch neural operators in the form of the widely used FNO instance normalizations, Foundation model baselines are provided by MPP-aVIT (MPP) [49] and we also pretrain CNO (60] model (see details in SM C on pretraining dataset, resulting in an additional foundation model baseline termed CNO- M, see SM C for details on baselines.
[3]] and recently proposed CNO [60], each augmented with cinecconditioned
Evaluation Metrics. All the models and baselines are evaluated on each task in terms of the relative L' error at the underlying final time. This choice is motivated by the fact that errors tend to grow over time, making final time prediction harder than any time-averaged quantities, see SM D.6.3 This also corresponds well to the interpretation of time-independent PDEs as long-time limits of (T) Following [23] that advocates this approach for LLMs, we evaluate all models in terms of scaling curves which plot the test error for each task vs. the number of task task-specific training examples, see SMD.1 To extract further information from scaling plots, we introduce two evaluation metrics,
(moodl == (FNO) y (model)'
EGg(model) == where (modell = E5NOO),, (11)
with (modell being the relative error (at final time) for the model with S trajectories. Thus, Accuracy Gain 5 measures how accurate the model is w.I.t FNO for given number (S) of samples while Efficiency Gain EGş measures how much fewer (greater) number of samples the model needs to attain the same error level as FNO trained on S samples. AG is the relevant metric for the limited compute regime whereas EG is relevant for the limited data regime.
POSEIDON performs very well on all downstream tasks. From the scaling plots SM Figures 7 to [ we observe that POSEIDON readily outperforms FNO on all the 15 downstream tasks. This point is further buttressed by Table 1 where the EG and AG (1I) metrics are presented (see also & for these metries for the POSEIDON-B and -T models). We observe from this table that POSEIDON requires far fewer task specific samples to attain the same error level as FNO does with = 1024 samples for time-dependent PDEs (S = 4096 for iimeinndeenneett PDEs). In fact, there are 4 tasks for which mere 3 task-specifir samples suffice for POSEIDON to attain the same error as FNO with 1024 samples. From SM Table g we observe that, on an average (median). only 20 samples needed for POSEIDON-L reach the errors of FNO with 1024 samples and in 13 (of the 15) tasks, POSEIDON-L. needs an order of magnitude fewer samples than FNO. Similarly from Table
SM T
T and SM Table 9 we see that for the same number (S- = 128 for time-dependent, and S = 512 for -ndependenn PDEs) of samples, POSEIDON-L has significantly lower error than FNO, with gains ranging from anywhere between 10% to factor of 25, with the mean gain of accuracy being an entire order of magnitude.
Among the trained- frrom-rrratet neural operator baselines, CNO and scOT are comparable in perfor- mance to each other, while both outperform FNO significantly on almost all tasks (see Table I and SM Table 9) However, POSEIDON is much superior to both of them, in terms of gains in sample efficiency (median gain of an order of magnitude) as well as accuracy (average gain of a factor of 4).
POSEIDON generalizes well to unseen physics. This impressive performance of POSEIDON is particularly noteworthy as all the downstream tasks are distri-iiitiibution with respect to the pre- training dataset. This performance is also consistent across the 9 tasks which involve PDEs not seen during pretraining. POSEIDON is the best performing model on 8 of these tasks, including all the time-dependent PDEs. It is only for 1 of the time-indepedent PDEs, which constitute the hardest generalization challenge, that POSEIDON is outperformed by CNO, but only marginally. These results underscore the ability of POSEIDON to learn completely new physical processes and contexts from few downstream task-specific samples.
Architecture of the foundation model matters. We observe from SM D.I and Table I (see also SM Table 9 that POSEIDON outperforms CNO-FM clearly on 14 out of 15 downstream tasks. On average (median over all tasks). CNO-FM requires approximately 100 task-specific examples to attain the error levels of FNO with 1024 samples, whereas POSEIDON only requires approximately 20. As CNO-FM and POSEIDON have been pretrained on exactly the same dataset, this difference in performance can be largely attributed to architectural differences as CNO-FM based on multiscale CNNs. in contrast to the multiscale vision transformer which is the backbone of POSEIDON.
7

---

Table 1: Efficiency gain EG (II)) with 1024 for time-dependent and S 4096 for time-
independent PDEs) and Accuracy Gain (AG) (II)) with S = 128 for time-dependent and S = 512 for indeeenndennden PDEs) for all models and downstream tasks.
Pretrained Models
Models trained from Scratch
POSEIDON-L
CNO-FM
MPP-B
CNO
ScOT
FNO
EG
AG
EG
AG
EG
AG
EG
AG
EG
AG
EG
AG
NS-PwC
890.6
24.7
16.6
3.3
7.4
2.3
3.7
1.5
2.0
I
NS-SVS
502.9
7.3
59.6
3.1
34.8
2.2
73.2
3.4
10.2
1.2
1
NS-BB
552.5
29.3
10.6
3.9
4.6
2.6
2.7
1.7
3.4
2.7
1
NS-SL
21.9
5.5
0.4
0.8
0.3
0.8
0.8
1.2
0.3
0.8
1
NS-Tracer-Pwe PWC
49.8
8.7
17.8
3.6
8.5
2.7
4.6
1.9
4.6
1.9
1
FNS-KF
62.5
7.4
13.2
2.7
2.0
1.6
3.1
1.5
3.3
0.9
1
CE-RPUI
352.2
6.5
33.2
2.3
0.0
1.2
12.5
1.8
15.6
2.7
1
CE-RM
4.6
1.2
0.6
1.0
0.0
0.2
1.7
1.1
0.4
1.0
1
SE-AF
3.4
1.2
4.8
1.3
2.2
1.1
5.5
1.5
1.2
1.0
1
GCE-RT
5.3
2.0
1.2
1.0
0.0
0.3
1.2
1.4
1.7
Wave-Layer
46.5
5.6
2.2
0.0
11.4
3.0
13.0
2.9
1
Wave-Gauss
62.1
5.6
6.0
1.8
0.0
0.8
14.0
2.6
9.2
2.7
/
ACE
17.0
11.6
1.7
2.0
0.0
0.3
4.5
4.6
6.5
5.2
1
Poisson-Gause
42.5
20.5
25.0
9.2
17.0
7.3
21.1
7.0
9.8
5.3
/
Helmholtz
78.3
6.7
54.0
5.1
22.4
3.0
68.9
7.3
60.4
9.0
/
The second baseline foundation model, MPP-B of [49]. is based on a transformer with axial attention and is pretrained on the PDEBench dataset [71] However, it has been trained predict the next time step, given context window of 7 previous time steps, with 7 = 16 as the default. We emphasize that this next step prediction, given context window, does not solve the underlying operator learning task OLT directly as OLT requires that the entire trajectory needs to be generated, given the initial data. Hence, we had to finetune the pretrained MPP model with varying context windows (starting with window size of 1), see SM 6M6 for details. We see from Table Tan SM Table that the finetuned MPP modestly outperformed FNO on some (8 out of 15) of the downstream tasks but it failed on the rest of them, where MPP simply could not attain the error levels of FNO, it did not converge even blew up with increasing number of downstream samples (see scaling plots in SM D.I.
In this context, it can be argued that the POSEIDON-L model is larger in size than both CNO-FM and MPP-B and perhaps, it is this size difference which explains the differential in performance. However, this is far from the case. As shown in all the scaling plots of S.I. and SM Tables [anles(aand]] and both CNO-FM and MPP-B are significantly inferior to the POSEIDON-E model, which is comparable in size. In fact, we can see from these tables that even the POSEIDON-T model, which is order of magnitude smaller in size, outperforms CNO-FM and MPP-B B handily. also readily outperforms all the
neural operators (CNO, FNO and scOT) which are of comparable size to it, leading us to conclude that it is the combination of the pretraining dataset as well as the underlying architecture, rather than just model size, that underpins the superior performance of POSEIDON.
POSEIDON scales with model size. Nevertheless, the model size of POSEIDON does matter. As seen from SM Figure 22 both the training as well as evaluation (validation) errors on the pretraining dataset clearly decrease with increasing model size of POSEIDON. However, does this scaling with model size lead to any impact on the performance of these models, when finetuned on downstream tasks? We see from the scaling plots in SM D that POSEIDON-L consistently outperforms the smaller POSEIDON-B on most downstream tasks. This trend is reinforced by SM Tables & and g where we find that, on an average, increasing model size correlates with a consistent decrease in test error as well as increase in sample efficiency of the pretrained model on downstream tasks.
8

---

POSEIDON scales with dataset size. In SM Figure 23 we show how by increasing the size of the pretraining dataset, in terms of the number of trajectories, the training and validation losses for the pretrained POSEIDON-B model decrease. Moreover, from SM Figures| 24 to 38 where we plot the test error versus number of downstream task-specific samples for 2 different models, POSEIDON-B trained on the full pretraining dataset and on one-eighth of the pretraining dataset, we find that for most (9 of the 15) of the downstream tasks, increasing the number of samples in the pretraining dataset, by an order of magnitude, does lead to significantly greater accuracy even at the downstream task level. For the remaining tasks, the models trained with less data are either on par or marginally inferior to the model trained with the full dataset.
The quality/diversity of the pretraining dataset matters. To demonstrate this point, we consider two different datasets: one in which half the trajectories of the pretraining dataset for POSEIDON- B are randomly dropped (from every operator), and the other where less diversity of the pretraining dataset is imposed by dropping all the trajectories corresponding to 3 out of 6 operators, namely CE-CRP, CE-Gauss and NS-Sines. Thus, the total size of both datasets is the same but one clearly less diverse than the other. The respective POSEIDON-B -B models are then evaluated on all the downstream tasks. As shown SM Figures 24 to 38 the model trained on less diverse data performs worse than its counterpart on 10 out of the 15 tasks and is par on 4 of them. Thus, we demonstrate that in large majority of downstream tasks, the quality//iversity of the pretraining dataset matters.
How does POSEIDON generalize to unseen physics? In order to understand the surprising ability of POSEIDON to generalize so well to unseen and a priori unrelated PDEs and physical processes downstream, we present three case studies in SMD44 to uncover some of the inner workings of this foundation model. In particular, we first consider the CE-RPUI downstream task. This task pertains to the compressible Euler equations, which are included in the pretraining dataset. However, the underlying initial data distribution is not seen during pretraining, making the task out-of-distribution. We show in SM 1 how POSEIDON leverages different features of different operators from the pretraining dataset to learn this task accurately with very few samples (see SM Figure (9). In particular, the diversity of the pretraining dataset is more instrumental in ensuring better generalization to this unseen initial condition than the size of the dataset.
In D.4.3. we study the Poisson- Gauss task to understand arguably the most surprising finding about the POSEIDON foundation models, i.e., their ability to generalize well to PDEs that are completely unrelated to the Euler and Navier-Stokes equations of fluid dynamics. This task pertains to the Poisson equation 668) with a forcing term, which is a superposition of Gaussians. The task is very different from those seen during pretraining in multiple ways, namely the underlying PDE is not only indeeenndennntt (in contrast to the time-dependent PDEs of pretraining) but also elliptic (whereas the PDEs during pretraining are either hyperbolic or convection-dominater the boundary conditions are Dirichlet (instead of Periodic) leading to very different physics, that of diffusion and smoothing, being manifested for this task, when contrasted with the physics seen during pretraining which is dominated by transport, shock wave propagation and fluid mixing. Given this context, one would not expect POSEIDON to perform well this task. Yet. from SM Figures 20 and 74 we know that POSEIDON performs exceptionally well, learning the solution operator accurately with few samples. As we elaborate in SM D.4.3] POSEIDON does not use the first few training examples to forget the physics that it has learned during pretraining and learn the new physics for this task after that. Rather surprisingly, as illustrated in SM Figure [3] already with one task specific training example, POSEIDON outputs an (approximation of the) input, rather than the expected dynamic evolution of fluids with Gaussian inputs (see SM Figures (56 and 66)) seen during pretraining. Then, with very few (16) examples, it is able to learn the rudiments of diffusion and smoothing of features (SM Figure 43). which are characteristics of clliptic equations. To further test how the foundation model leverages physics representations learned during pretraining, we froze the latent space by only finetuning the embeddings and freezing the latent space parameters by setting 0, = 0. for all r, in (10) for finetuning. As shown in (SM Figure 44), even this fruzen latent version of POSEIDON is very effective at learning the underlying solution operator, demonstrating that very rich physical representations were learned during pretraining.
and
Further results on the robustness of POSEIDON for different factors and ablations as well as compar- isons with other foundation models is provided in SM and details of computational resources are described in SM E

---

4 Discussion
Summary. In this paper, we have presented POSEIDON, family of foundation models for learning PDEs. The backhone of POSEIDON is ScOT, a multiscale vision transformer with shifted-windowad (SwinV2) attention that maps input functions (initial data, coefficients, sources) etc. to the solution (trajectory) of a PDE. Lead-time conditioning through a time-modulated layer norm allows for continuous-in-time evaluation and a novel all2all training strategy enables the scaling up of training data by leveraging the semi-group structure of solutions of time-dependent PDEs. POSEIDON is pretrained on a diverse large-scale dataset of operators for the compressible Euler and incompressible Navier-Stokes PDEs. Its performance is evaluated on a challenging suite of 15 out oft-oo--istte dfitiiiibittttii downstream tasks covering wide variety of PDEs and data distributions. POSEIDON displays excellent downstream performance and is the best performing model on 14 of the 15 tasks. In particular, requires orders of magnitude (median of 50) fewer task asee-spcc samples to attain the same error as the widely used FNO. This large gain in sample efficiency as well as order of magnitude gains in accuracy also holds for PDEs that are not seen during pretraining, making us conclude that POSEIDON generalizes well to new physics. POSEIDON also scales with model and dataset size, with respect to pretraining and even downstream task performance. To the best of our knowledge, this is the first time that it has been clearly demonstrated that by pretraining on a very small set of PDEs, a foundation model can generalize to a wider variety of unseen and unrelated PDEs and data distributions downstream. Thus, we provide affirmative answer to the very fundamental question of whether foundation models for PDEs are even feasible. Moreover, we investigate possible mechanisms via which POSEIDON can effectively leverage representations, learnt during pretraining, to accurately learn downstream tasks by finetuning on a few spec-fie examples. Our case studies suggest hitherto undiscovered relationships between different PDEs that enable this transfer to materialize. Finally, all the models are made publicly available, as well as the pretraining and downstream datasets are open sourced in the PDEGYM collection
Related Work. Foundation models for PDEs are of very recent vintage. The foundation model of (67 is limited to very specific elliptic Poisson and Helmholtz PDEs with a FNO backbone whereas ICON 774 considers a very small 1-D dataset. Neither of these models are comparable in scope to POSEIDON. Universal physics transformers [I] employs transformers but its focus is incompressible fluid flows and the ability to generalize across Eulerian and Lagrangian data. Thus, direct comparison with POSEIDON is not possible. On the other hand, MPP [49] and DPOT [19] are designed to be general purpose foundation models for PDEs that can be compared to POSEIDON. We have already extensively compared MPP with POSEIDON in Section to demonstrate the very large superiority of POSEIDON across various metrics. Although DPOT has a different architecture (Adaptive FNO) and was trained on more datasets than MPP, it follows a similar training and evaluation strategy of next time-step prediction, given a context window of previous time-steps. As argued before. this does not solve the operator learning task of generating the entire trajectory, given initial data. At the time of writing this paper, DPOT was not publicly available but it was released by the time this paper has been revised, enabling us modify the fine-tuning procedure of DPOT and to perform comparisons between it and POSEIDON. While directing the interested reader to SM5. for details, we summarize our findings by observing that POSEIDON models are significantly better performing than DPOT foundation models, both in terms of accuracy and sample efficiency.
Limitations. The range of PDEs and underlying data distributions is huge and POSEIDON was only trained and evaluated on a few of them. Although the results here clearly demonstrate its ability to learn unseen physics from few task-specific training examples, we anticipate that given that it is scaling with respect to both data quantity and quality, POSEIDON'S performance as a general purpose PDE foundation model will significantly improve when it is pretrained with even more diverse PDE datasets in the future. In particular, pretraining with lime-independent PDEs (particularly elliptic PDEs) as well larger range of time-scales in time-dependent PDEs will greatly enhance POSHIDON. The focus here was on Cartesian geometries although POSEIDON displayed the ability to generalize to non- con-Cartesian geometries, via masking, on the SE-AF task. We plan to add several non-Cartesian examples in the pretraining dataset to augment POSEIDON's 's performance on general geometries/boundary conditions. Moreover, given the fact that POSEIDON serves as a fast and accurate neural PDE surrogate, its extension to qualitatively different downstream tasks such as uncertainty quantification 45], inverse problems [53] and PDE-constrained optimization [46] is fairly straightforward and will be considered in future work.
10

---

Acknowledgments and Disclosure of Funding
This work supported by grant from the Swiss National Supercomputing Centre (CSCS) under project ID 1217.
References
[1] B. Alkin, A. Fürst, S. Schmid, L. Gruber, M. Holzleitner, and J. Brandstetter. Universal physies transformers: A framework for efficiently scaling neural operators, 2024.
[2] J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization, 2016.
[3] Bartolucci, E. de Bézenac, B. Raonic, R. Molinaro, S. Mishra, and R. Alaifari. Representation equivalent neural operators: framework for alias-frec operator learning. arXiv:2305.199113 2023.
[4] Batatia, P. Benner, Chiang, A. M. Elena, D. P. Kovács, J. Riebesell, R. Advincula, M. Asta, M. Avaylon, W.J. Baldwin, F. Berger, N. Bernstein, A. Bhowmik, S. M. Blau, V. Cărare, J. P. Darby, S. De, F. D. Pia, V. Deringer, R. Elijošius, Z. El-Machachi, F. Falcioni, E. Fako, J. George, R. E. A. Goodall, C. P. Grey, P. Grigorev, S. Han, W. Handley, H. H. Hcenen, K. Hermansson, C. Holm, Jaafar, Hofmann, K. Jakob, H. Jung. V. Kapil, A. D. Kaplan, N. Karimitari, J. R. Kermode, N. Kroupa, J. Kullgren, M. C. Kuner, D. Kuryla, G. Liepuoniute, J. T. Margraf, I.-B. Magdäu, A. Michaelides, J. H. Moore, A. A. Naik, S. P. Niblett, S. W. Norwood, N. O'Neill, C. Ortner, K. A. Persson, K. Reuter, A. S. Rosen, L. L. Schaaf, C. Schran, B. X. Shi, E. Sivonxay, T. K. Stenczel, V. Svahn, C. Sutton, T. D. Swinburne, J. Tilly, C. van der Oord, E. Varga-Umbrich, T. Vegge, M. Vondrák, Y. Wang, W. C. Witt, F. Zills, and G. Csányi. A foundation model for atomistic materials chemistry, 2024.
A. C. Ferrari, A. Genreith-Schriewe
[5] J. Bell, Collela, and H. M. Glaz. A second-order projection method for the incompressible Navier-Stokes equations. J. Comput. Phys., 85:257-283, 1989.
[6] R. Bommasani, D. A. Hudson, E. Adeli, R. Altman, S. Arora, S. von Arx, M. S. Bernstein,
J. Bohg, A. Bosselut, F. Brunskill, F. Brynjolfsson, S. Buch, D. Card, Castellon, N. Chatterji, A. Chen, K. Creel, J. Q. Davis, D. Demszky, C. Donahue, M. Doumbouya, E. Durmus, S. Ermon, J. Etchemendy, K. Ethayarajh, L. Fei-Fei, C. Finn, Gale, L. Gillespie, K. Goel, N. Goodman, S. Grossman, N. Guha, T. Hashimoto, P. Henderson, J. Hewitt, D. E. Ho, J. Hong, K. Hsu, Khani, O. Khattab, P. W. Koh, M. Krass, R. Krishna, R. Kuditipudi, A. Kumar, F. Ladhak, M. Lee, T. Lee, J. Leskovee, I. Levent, X. L. Li, X. Li, T. Ma, A. Malik, C. D. Manning, S. Mirchandani, E. Mitchell, Z. Munyikwa, S. Nair. A. Narayan, D. Narayanan, B. Newman, A. Nie. J. C. Niebles, H. Nilforoshan, J. Nyarko, G. Ogut, L. Orr, Papadimitriou. J. Park, C. Piech, E. Portelance, C. Potts, A. Raghunathan, R. Reich, H. Ren, F. Rong. Y. Roohani, C. Ruiz, J. Ryan, C. Ré, D. Sadigh, S. Sagawa, K. Santhanam, A. Shih, K. Srinivasan, A. Tamkin, R. Taori, A. W. Thomas, F. Tramèr, R. E. Wang. W. Wang, B. Wu, J. Wu, Y. Wu, S. M. Xie, M. Yasunaga, J. You, M. Zaharia, M. Zhang, T. Zhang, X. Zhang, Y. Zhang, L. Zheng, K. Zhou, and P. Liang. On the opportunities and risks of foundation models, 2022.
J. Huang, T. Icard, Jain, D. Jurafsky, P. Kalluri, S. Karamcheti, G. Keeling,
[7] B. Bonev, T. Kurth, C. Hundt, J. Pathak, M. Baust, K. Kashinath, and A. Anandkumar. Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere, June 2023. arXiv:2306.0388 038838 [physics).
[8] J. Brandstetter, D. E. Worrall, and M. Welling, Message passing neural PDE solvers. In International Conference on Learning Representations, 2022.
[9] A. Brohan, N. Brown, J. Carbajal, Chebotar, X. Chen, K. Choromanski, Ding, D. Driess, A. Dubey, C. Finn, P. Florence, C. Fu, M. G. Arenas, K. Gopalakrishnan, K. Han, K. Hausman, A. Herzog, Hsu, B. Ichter, A. Irpan, N. Joshi, R. Julian, D. Kalashnikov, Y. Kuang, Leal, L. Lee, T.-W. E. Lee, S. Levine, Y. Lu, H. Michalewski, I. Mordatch, K. Pertsch, K. Rao, K. Reymann, M. Ryoo, G. Salazar, Sanketi, Sermanet, J. Singh, A. Singh, R. Soricut, H. Tran, V. Vanhoucke, Q. Vuong, A. Wahid, S. Welker, P. Wohlhart, J. Wu, F. Xia, T. Xiao, Xu, S. Xu, T. Yu, and B. Zitkovich. Rt-2: Vision-language-action models transfer web knowledge robotic control, 2023.

---

[10] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, Shyam, G. Sastry, A. Askell, Agarwal, A. Herbert- Voss, G. Krueger. Henighan, R. Child, Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, Clark, C. Berner, S. McCandlish, Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners, 2020.
A.
[11] H. Cao, Y. Wang. J. Chen, D. Jiang. X. Zhang, Q. Tian, and M. Wang. Swin-Unet: Unet- Like Pure Transformer for Medical Image Segmentation. In L. Karlinsky, T. Michaeli, and K. Nishino, editors, Computer Vision - ECCV 2022 Workshops, Lecture Notes in Computer Science, pages 205-218, Cham, 2023. Springer Nature Switzerland.
[12] Cao. Choose a transformer: Fourier or galerkin. In 35th conference on neural information processing systems, 2021.
[13] T. Chen and H. Chen. Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems. IEEE Transactions on Neural Networks, 6(4):911-917) 1995.
[14] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby. An Image is Worth 16x16 Transformers for Image Recognition at Scale, June 2021. arXiv:2010.11929 [cs].
Words:
[15] L. C. Evans. Partial differential equations, volume 19. American Mathematical Soc., 2010.
[16] U. S. Fjordholm, R. Käppeli, S. Mishra, and E. Tadmor. Construction of approximate entropy measure valued solutions for hyperbolic systems of conservation laws. Found. Comput. Math., 17(3):763-827, 2017.
[17] Google Gemini Team. Gemini: A family of highly capable multimodal models, 2024.
[18] J. K. Gupta and J. Brandstetter. Towards multi-spatiotemporalscaall generalized pde modeling, 2022.
[19] Z. Hao, C. Su, S. Liu, Berner, C. Ying, H. Su, A. Anandkumar, Song, and Zhu. DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training Mar. 2024. arXiv:2403.03544 [cs, math].
[20] Z. Hao, Z. Wang, H. Su, C. Ying, Y. Dong, S. Liu, Z. Cheng, Song, and J. Zhu. Gnot: A general neural operator transformer for operator learning, 2023.
[21] D. Hendrycks and K. Gimpel. Gaussian error linear units (gelus), 2023.
[22] J. S. Hesthaven. Numerical methods for conservation laws: From analysis to algorithms, SIAM, 2018.
[23] Kaplan, McCandlish, T. Henighan, B. Brown, B. Chess, R. Child, Gray, A. Radford,
J. Wu, and D. Amodei. Scaling Laws for Neural Language Models, Jan. 2020. arXiv:2001.08 .08361 [cs, stat].
[24] G. E. Karniadakis, G. Kevrekidis, Lu, P. Perdikaris, S. Wang, and L. Yang. Physics informed machine learning. Nature Reviews Physics, pages 1-19, may 2021.
[25] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:14122 6980, 2014.
[26] G. Kissas, H. Seidman, L. Guilhoto, V. M. Preciado, G. J. Pappas, and Perdikaris. Learning operators with coupled attention. Journal of Machine Learning Research, 23(215):1- 63, 2022.
[27] N. Kovachki, Z. Li, B. Liu, K. Azizzadensheli, K. Bhattacharya, A. Stuart, and Anandkumar. Neural operator: Learning maps between function spaces. arXiv preprint 0rX1v:2108.0848113, 2021.
[28] R. Käppeli and S. Mishra. Well-balanced schemes for the euler equations with gravitation. Journal of Computational Physics, 199-199- 2014.
12

---

[29] L. D. Landau and E. M. Lipschitz. Fluid Mechanics, 2nd edition. Butterworth Heinemann, 1987.
[30] S. Lanthaler and S. Mishra. On the convergence of the spectral viscosity method for the wwo-dimensional incompressible euler equations with rough initial data. Foundations of Compu- tational Mathematics, 20(5): 1000::109--33-119 10 2020.
[31] S. Lanthaler, S. Mishra, and C. Parés-Pulido Statistical solutions of the incompressible euler equations. Mathematical Models and Methods in Applied Sciences, 31(0))2223-22) 292, Feb 2021.
[32] Z. Li, D. Z. Huang, B. Liu, and A. Anandkumar. Fourier neural operator with learned deforma- tions for pdes on general geometries, 2022.
[33] Z. Li, N. B. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, and A. Anand- kumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021.
[34] Z. L.,iN. N. B. Kovachki, Azizzadenesheli, B. Liu, A. M. Stuart, K. Bhattacharya, and A. Anand- kumar. Multipole graph neural operator for parametric partial differential equations. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances Neural Information Processing Systems (NeurIPS), volume 33, pages 6755-6766. Curran Associates, Inc., 2020.
[35] Z. Li, K. Meidani, and Farimani. Transformer for partial differential equations' operator learning, 2023.
[36] P. Lippe and B. S. Veeling. PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers.
[37] Z. Liu, H. Hu, Lin, Yao, Xie, Y. Wei, Ning, Cao, Zhang, L. Dong, F. Wei, and B. Guo. Swin Transformer V2: Scaling Up Capacity and Resolution, Apr. 2022. arXiv:2111 11101110888 [cs].
[38] Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo. Swin Transformer: Hierarchical Vision Transformer using Shifted Windows, Aug. 2021. arXiv:2103.40030 [cs].
[39] Z. Liu, H. Mao, C.-Y. Wu, C. Feichtenhofer, T. Darrell, and S. Xie. A ConvNet for the 2020s, Mar. 2022. arXiv:2201.0335 [cs].
[40] A. Logg, K.-A. Mardal, and G. Wells. Automated solution differentia equations by the finite element method: The FERiCS book. Springer, 2012.
[41] I. Loshchilov and F. Hutter. Decoupled weight decay regularization, 2019.
[42] L. Lu, P. Jin, G. Pang, Z. Zhang, and G. E. Karniadakis. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):118 229, 2021.
[43] F. Luporini, M. Louboutin, M. Lange, N. Kukreja, P. Witte, J. Hückelheim, C. Yount, P. H. J. Kelly, F. J. Herrmann, and G. Gorman. Architecture and performance of devito, system for automated stencil computation. ACM Trans. Math. Softw., 46(1), apr 2020.
[44] K. O. Lye. Computation of statistical solutions of hyperbolic systems of conservation laws. PhD thesis, 2020.
[45] K. O. Lye, S. Mishra, and D. Ray. Deep learning observables in computational fluid dynamics. Journal of Computational Physics, page 109339, 2020.
[46] K. O. Lye, S. Mishra, D. Ray, and P. Chandrasekhar. Iterative Surrogate Model Optimization (ISMO): An active learning algorithm for PDE constrained optimization with deep neural networks. Computer Methods in Applied Mechanics and Engineering, 174:113575, Feb. 2021. arXiv:2008.0571 05730 [cs, math].
[47] A. J. Majda and L. Bertozzi. Vorticity and Incompressible Flow. Cambridge Texts in Applied Mathematics. Cambridge University Press, 2001.
13

---

[48] D. A. Masters, N. J. Taylor, T. Rendall, C. B. Allen, and J. Poole. Geometric comparison of aerofoil shape parameterization methods, AIAA Journal, pages 1575-1589, 2017.
[49] M. McCabe, B. R.-S. Blancard, L. H. Parker, R. Ohana, M. Cranmer, A. Bietti, M. Eickenberg, S. Golkar, G. Krawezik, F. Lanusse, M. Pettee, T. Tesileanu, K. Cho, and S. Ho. Multiple Physics Pretraining for Physical Surrogate Models, Oct. 2023. arXiv:2310.02994 [cs, stat].
[50] S. Mishra, C. Schwab, and J. Sukys. Multi-level Monte Carlo finite volume methods for nonlinear systems of conservation laws mulli-ddimension. J. Comput. Phys., 231(8)33365- 3388, 2012.
[51] S. Mishra and A. E. Townsend. Numerical Analysis meets Machine Learning. Handbook of Numerical Analysis. Springer, 2024.
[52] D. Mizrahi, R. Bachmann, O. F. Kar, T. Yeo, M. Gao, A. Dehghan, and A. Zamir. 4M: Massively Multimodal Masked Modeling, Dec, 2023, arXiv:2312.06644 [cs].
[53] R. Molinaro, Y. Yang, B. Engquist, and Mishra. Neural inverse operators for solving pde inverse problems, 2023.
[54] Nguyen, J. Brandstetter, A. Kapoor, K. Gupta, and A. Grover. Climax: A foundation model for weather and climate, 2023.
[55] E. Perez, F. Strub, H. de Vries, V. Dumoulin. and A. C. Courville, Film: Visual reasoning with a general conditioning layer. CoRR, abs/1709.07871, 2017.
[56] T. Pfaff, M. Fortunato. A. Sanchez-Gonzaler and P. W. Battaglia. Learning Mesh-Based Simulation with Graph Networks, June 2021. arXiv:2010.03409 [cs].
[57] M. Prasthofer, T. De Ryck, and S. Mishra. Variable input deep operator networks. arXiv preprint wwXi::2205.11404 2022.
[58] A. Quarteroni, A. Manzoni, and F. Negri. Reduced basis methods for partial differential equations: an introduction, volume 92. Springer, 2015.
[59] A. Quarteroni and A. Valli. Numerical approximation of Partial differential equations, vol- ume 23. Springer, 1994.
[60] B. Raonić, R. Molinaro, T. De Ryck, T. Rohner, F. Bartolucci, R. Alaifari, S. Mishra, and E. de Bézenac. Convolutional Neural Operators for robust and accurate learning of PDEs, Dec. 2023. arXiv:2302.01178 [cs].
[61] R. Krasny. A study of singularity formation in a vortex sheet with a point vortex approximation.
J. Fluid Mech., 167:65-93, 1986.
[62] T. Rohner and S. Mishra. Efficient computation of large-scale statistical solutions to incompress- ible fluid flows. In Proceedings of the Platform for Advanced Scientific Computing Conference, PASC '24, ACM, June 2024.
[63] Y. Rosen, Y. Roohani, A. Agarwal, L. Samotorcan, S. R. Quake, and J. Leskovec. Universal cell embeddings: A foundation model for cell biology. bioRxiv, 2023.
[64] K. Saab, T. Tu, W.-H. Weng. Tanno, D. Stutz, E. Wulczyn, Zhang, T. Strother, C. Park, E. Vedadi, J. Z. Chaves, S.-Y. Hu, M. Schaekermann, A. Kamath, Y. Cheng, D. G. T. Barrett, Liu, J. baptiste Alayrac, N. Houlsby, N. Tomasev, J. Freyberg, C. Lau, Kemp, J. Lai, S. Azizi, K. Kanada, S. Man, K. Kulkarni, R. Sun, S. Shakeri, L. He, B. Caine, A. Webson, N. Latysheva, M. Johnson, P. Mansfield, J. Lu, E. Rivlin, J. Anderson, B. Green, R. Wong, J. Krause, J. Shlens, E. Domi- nowska, S. M. A. Eslami, K. Chou, C. Cui, O. Vinyals, K. Kavukcuoglu, J. Manyika, J. Dean, D. Hassabis, Y. Matias, D. Webster, J. Barral, G. Corrado, C. Semturs, S. S. Mahdavi, J. Got- tweis, A. Karthikesalingam, and V. Natarajan. Capabilities of gemini models medicine, 2024.
C. Cheung, B. Mustafa, A. Palepu, D. McDuff, Hou, T. Golany,
14

---

[65] A. Sanchez- Sanchez-Gonzalez J. Godwin, T. Pfaff, R. Ying, J. Leskovec, and P. Battaglia. Learning Simulate Complex Physies with Graph Networks. In Proceedings of the 37th International Conference Machine Learning, pages 8459 8968. PMLR, Nov. 2020. ISSN: 2640-3498.
[66] J. Shen, T. Marwah, and A. Talwalkar. Ups: Efficiently building foundation models for pde solving via cross- moosl adaptation, 2024.
[67] S. Subramanian, P. Harrington, K. Keutzer, W. Bhimji, D. Morozov, M. Mahoney, and A. Gho- lami. Towards Foundation Models for Scientific Machine Learning: Characterizing Scaling and Transfer Behavior, May 2023. arXiv:2306.00 00258 [cs, math].
[68] J. Sun, Y. Liu, Z. Zhang, and H. Schaeffer. Towards foundation model for partial differential equations: Multi-operaaor learning and extrapolation. arXiv preprint (ttt:24404.2235554 2024.
[69] Tadmor. Convergence of spectral methods for nonlinear conservation laws. SIAM Journal on Numerical Analysis, 26(1):30-444 44, 1989.
[70] E. Tadmor. Burgers' Equation with Vanishing Hyper- Viscosity. Communications in Mathemati- cal Sciences, 2(2):317 324, 2004.
[71] M. Takamoto, T. Praditia, R. Leiteritz, D. MacKinlay, F. Alesiani, D. Pflüger, and M. Niepert. PDEBENCH: An Extensive Benchmark for Scientific Machine Learning, Mar. 2023. arXiv:2210.07182 [physics].
[72] H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, Joulin, E. Grave, and G. Lample. Llama: Open and efficient foundation language models, 2023.
[73] T. Wolf, L.. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, C. Ma, Y. Jernite, J. Plu, C. Xu, T. Le Scao, S. Gugger, M. Drame, Q. Lhoest, and A. M. Rush. Transformers: State-of-the-Art Natural Language Processing, pages 38-45, Association for Computational Linguistics, Oct. 2020.
[74] L. Yang, S. Liu, T. Meng, and S. J. Osher. In-context operator learning with data prompts for differential equation problems. Proceedings of the National Academy of Sciences, 120(399::230114212,, Sept. 2023. Publisher: Proceedings of the National Academy of Sciences.
[75] Y. Zhu and N. Zabaras. Bayesian deep convolutional encoder-decoder networks for surrogate modeling and uncertainty quantification. Journal of Computational Physics, 2018.
15

---

Supplementary Material for:
POSEIDON: Efficient Foundation Models for PDEs.
Table of Contents
Architecture of the scalable Operator Transformer (scOT)
A.I Operator Learning with scOT
17
A.2 Computational Realization of scO1
17
B Datasets
21
Pretraining Datasets
21
Downstream Tasks
24
Models and Baselines
34
C.1
POSEIDON Models
34
scOT
37
CNO
37
C.4 FNO
39
C.5
CNO-FM
40
MPP
41
D
Results
43
D.
Performance on Downstream Tasks
43
D.2
Scaling with respect to Model Size
51
D.3
Scaling with respect to Pretraining Datasct Size and Quality
52
D.4
Case Studies
58
Results with DPOT
63
Further Ablations Results
65
E
Computational Resources
72
F Pretrained Models, Datascts, and Source Code
72
Visualizations
73
16

---

A Architecture of the scalable Operator Transformer (scOT)
A.1 Operator Learning with scOT
First, we describe how scOT (Section 2 of Main Text and Figure transforms function space inputs into function outputs below.
For simplicity of exposition, we set d 2 and specify D = (0, 1]2 as the underlying domain. On this domain, an uniform computational grid, with grid spacing A, of J equally spaced points Thend = (i j,A), with J = 1/A, set. Let < p < such that Jmodp = 0 and set = J/p. We divide the domain D Uر into a set of p2 non-overlapping and equal (in measure) patches. Any underlying input function C(D; I" is then partitioned into a function, that is piecewise constant on patches and embedded into a C-dimensional latent representation by applying the operator,
1² v(x) = (xa)(x) = ΣF (Wir)a(addd (x),
(12)
p=1
with F E Rxx is learnable matrix and the weight function W is defined in terms of the underlying computational grid as, W (3))
Σ
Wajdayas with 8 denoting the Dirac measure, and the
SJ
shared (across all patches) learnable weights given by,
Wjade
Jasjy S
W
(13)
Wi m.
otherwise.
The (patched and embedded) output function v of (12) is then processed through a sequence of Swin transformer blocks [38 37), each of which has the structure of SW, C(D:RC) + C(D; R), for layer index { = 1, ... L. formulated in Main Text (3)
The main building block of SwinV2 transformer block is the windowed multi-heat self attention operator,
H
W - (V)((v)()))
w
Viv(y)dy,
(14)
((ccm(( (B. D2
S
h=1
DE
for any E C(D; 9.. Here, h denotes the h-th attention head, W E IR CRm be the output matrix and Q K, Vh Rmxx be the query, key and value matrices, respectively. For any two vectors a,33 3, the cosine similarity is defined as (a, 3) |a||a| cos(a, 3) and B x R is general form for positional encodings. To be more specific, we use relative log position encodings by setting the inputs to B to be the logarithm of the relative positions (k: k) within the window and the function B itself to be small MLP. Finally, the domain of integration D. is simply the window where the point of interest x lies, i.e., 1 < 9z M such that x E D Underlying
(14. is the partition of the domain into windows such that D = م D with 1 S E S L indexing the underlying layer within a Swin transformer block and with 2 denoting the number of windows. Moreover, the windows are shifted across layers, as depicted in Figure (c), so that all the points can be attended to, by iteratively shifting windows across multiple layers/blocks.
The MLP in Main Text 3 of the form, MLP C(D:R)) - C(D:RO) with
MLP(v)(z) = Wo (Wv(x)+B).
(15)
for learnable weight matrices W E CXC W CCC bias vector Ó _º and nonlinear activation function a IR + R. The Layer Norm LN in Main Text (3) is given by Main Text ( The remaining operations in scOT (see Main Text Figure 2 are described in their discrete form below.
A.2 Computational Realization of scOT
The scalable Operator Transformer (scOT), forming the underlying model architecture for POSEIDON, is constructed as encoder-decoder architecture. Starting from patching and embedding, embedded
17

---

tokens are inputted into multiple stages of Swin transformer blocks, each followed by patch merging. The encoder is connected every level to the decoder through ConvNeXt [3] blocks. whereas the bottleneck is convolution-free Finally, through patch recovery and mixup, the output is assembled. We refer to Figure (a) in the Main Text for an illustration of the overall architecture and concrete computational realizations by presenting discrete versions of the continuous operators described in the subsection above as well as claborating on other operators used in scOT.
Patch Partitioning. The encoder consists of the patch partitioning operation, creating visual tokens from n. discretized (on the uniform computational grid described Section A.T] input functions a. E [RJJJ i (1....,n), n). Each a; is divided into non-overlapping patches of size p (with
p < J) such that p2 =
patches arise. For an illustration, we refer to Figure 2 (c) of the
Main Text where P = Patches are combined for every a; such that a sequence of a Rexpxn
j E (1,..., ]2] visual tokens can be fed to the embedding operation.
Embedding. Each of these patches is linearly embedded using shared learnable weight We Rxxxxxx (E Q) and bias bg RC ()),
(v3): = (bb); + +>> Σ (We) (W)))...((()amm
(16)
k=1 m1,mm
where (.); denotes the i-th component (for all siss C) and C is the embedding dimension. It straightforward to observe that (16) is a discretization of the operator (12). with an additional bias term. The resulting embedding is then passed through (time- time-conditioned) layer norm (see Main Text Equation 23
V2 Stage. At each level i fo...., I. 1} of the U-Net-style architecture, Swin stage S; is employed consisting of t, chained SwinV2 transformer blocks Tr,,
S Tr, o T .... TI.
(17)
This is done in both encoder and decoder, and the same number t, of Swin transformer blocks is used on each level.
V2 Transformer Block.
A V2 transformer block T is built as follows
v'(v = (N o A)(() +v T(v) = (N M))(v(v))
(18)
(19)
where e /4'xx200 is the sequence of embedded tokens. A the shifted-window multi-head self-attention operation, A the (time-conditioned) Layer Norm, .M a MLP. The attention mechanism A acts only on windows of size M M patches/tokens that shift from block to block TT+1 by doing cyclic displacement of M/2 M/2 tokens (when the sequence is interpreted in 2D; see
Figure 2 (c) of Main Text). So, with an input window v E R
RMM M²x'2''
A(v) Concat (o(cat(A((v),...... Ar, (v) + I
(20)
where Ar, 1 S h; is attention in head / with the maximum number of heads depending the stage with Wo E 2C2''x.2'' bo C-2' being learnable parameters (E Ô). At then given by
cos( cos((vW) + +as= -b) (vWk))) (vw' + 1MM- -(/)
A((v) Softmax Bi(v) +
(21)
with W v, W O W E R.C-2'xC-2/// and byg by E CRC//// E R (all learnable E Ô), cos(:.)
the cosine similarity, 1 E RM* a vector of ones, Bi(v) E IR MMMxxM* the relative position bias matrix generated from the (logarithmic) relative positions of each patch [Ar, Ay] within window, parametrized through shared MLP P for all heads:
P(Ar, Ay) ReLU sign Ar) log( 1+ Axl), sign( Av) (1+ TWllllWBBttee bb. b). B.2 (22)
18

---

WB.1 E R2x112 1 E R512
WB2 E 5122A are all learnable (E +). Note that (2) is by replacing the spatial integrals therein with uniform quadrature.
discretization of the operator 14
The tokens, coming from the attention module are then fed to layer norm 2 if the PDE to be learned is time-indepenneent if it is time-dependent (also in the case of POSEIDON), it goes through
time-conditioned layer norm [55 N
C-2'
m(v)
Σ (v),
(23)
C. 2
0-2¹
1
d2(v) C 2'
((V) - vvv)))
(24)
1=1
N(v,t) = a(t) o --((v) (²(v)
16.2 + 3(t)
(25)
with v E RC22 be a token resulting from the attention module, t E R20. 1 E RC a vector of ones, and (t) = Wat ba, B(t) W by being learnable gain and bias (Wa, We. p, ba by RC-2' all E Q')
The last building block of the Swin V2 transformer block is hingle-hideen 21 as pointwise activation function and four times the width of the latent dimension C 2
MLP with GeLU
M(v) GeLU(vWW + b,) + b2
(26)
with W, C222440222 bj E R402' W2 C+C2'x02' b2 RO-2 all learnable parameters (E Q).
Patch Merging. At each (resolution) level of the architecture, after each Swin V2 stage in the encoder, linear downsampling operation D; is performed the output of the stage added to its input (additional residual connection) such that the resolution halves. This amounts to a linear transformation on four non-overlapping, stacked patches/tokens time E R40-22
Di = NWW v,
(27)
with Icarnable WD, e RC----xx40--- such that the latent dimension doubles. Here, an additional (time-conditioned) layer norm is applied.
ConvNeXt Block. Outputs from each encoder stage S, 5 - 2 are fed to n. chained (time- conditioned) ConvNe blocks [39 Qj: for that, the token sequence is reshaped to two-dimensional grid of tokens v E P/2' x P/2 xC-2 and transformed by
Q t) = ((el.U((((DWConnvV),.)
be, b We.2 boz) @ We +v
(28)
We,1 R-22'x4 b E R4-022 We.2 R+022x02 b E R W e. R all
learnable parameters (E θ) and DwConv is a depthwise convolution with kernel size 7 (and a padding of 3) and bias.
Patch Expansion. Similar to patch merging, after a Swin stage in the decoder, each output token V E C-2 is linearly upsampled through U4 to double the resolution and half the latent dimension, ui(v,t) = -N(Reshape (Wu. v), Wu..
(29)
where Wi,, E 2 x0 Wu,, E CC-22x0-22' are both learnable (E θ), and Reshape an operation that reshapes vector of size C- 2i+2 into matrix of size 4 x 2
Patch Recovery and Mixup. Having passed through the last stage of the decoder, every patch/visus token V E RCC is linearly transformed back from the latent space to form patches of the discretized output function u² R.xxxx..
(ug): (bR))I+ E
WR)... (V,)s
(30)
k=1
19

---

where (.); denotes the i-th component (for all 1 < S Cu) and cu is the number of components of the discretized output function. WR Rxxxxxxxp and bæ. R are shared across tokens and learnable )). These outputs are assembled on grid to form i E RJxJxe. which is transformed to the final output with convolution with kernel size 5 (and padding 2 to keep the dimensionality), without bias, with all parameters being in θ
Summary of Hyperparameters. In Table 2 give an overview of the hyperparameters to instantiate a scOT. To reduce the number of hyperparameters, we fix 4, M 16, L = 4. [h1, h2 h3, h4] [3, 6, 12, 24), and n. 2 in this work.
Table 2: Hyperparameters of scOT.
Hyperparameter
Description
p
patch size
M
window size
C
embedding/latent dimension
L
number of levels (L downssampiin/uussapling operations)
t;
number of Swin transformer blocks in level
h4
number attention heads in level
Te
number of ConvNeXt blocks at each level
20

---

B Datasets
We describe the various datasets used for pretraining and for the downstream tasks below. All these datasets are publicly available with the PDEGYM collection
B.1 Prefraining Datasets
Table 3: Abbreviations/Summarr for all the pretraining datasets. IC refers to initial conditions.
Abbreviation
PDE
Defining Feature
Visualization
NS-Sines
Navier-Stoker
Sine IC
Fig. 55
NS-Gaussians
Navier-Stokes
Gaussians (in Vorticity) IC
Fig. 56
CE-RP
Euler (37)
4 4-Qadraan Riemann Problem IC
Fig. 57
CE-CRP
Euler (37)
Multiple Curved Riemann Problems
Fig.
CE-KH
Fuler (77
Shear IC
Fig. 59
CE-Gauss
Euler 37)
Gaussians (in Vorticity) IC
Fig. 60
We pretrain POSEIDON models and CNO-FM dataset containing 6 operators, defined the space-time domain [0, 1J2 x [0, 1]. We include 2 operators governed by the Navier-Stokes equations (NS-Sines, NS NS-Gauss) and 4 operators governed by the Compressible Euler equations. The pretraining datasets encompass problems that exhibit a wide range of scales and complex, nonlinear dynamics.
The Incompressible Stoker equations of fluid dynamics are given by
Ue (u V)u + Vp vAu,, div == 0, in the domain D 10, 1 with suitable boundary conditions. Here, 10,T] D + R is the velocity field and p : [o, 7 x D + R4 is the pressure. In this work, a small viscosity y/ = 4 x 10 only applied high-enough Fourier modes to approximate the inviscid limit.
(31)
To generate the pretraining data, all the benchmarks for the Navier-Stokes equations are simulated until the time T = 1. Furthermore, we store 2/ snapshots of the numerically simulated velocity field u, uniformly spaced time. Each snapshot has a spatial resolution of 128 x 128. The initial conditions are drawn from various distributions, which we will describe later. The distribution of these initial conditions is crucial for determining the complexity of the samples and the overall dynamics.
All the Navier-Stokes experiments, both for the pretraining dataset and the downstream tasks, are simulated with the following spectral method. Fix mesh width A N for some N E N. We consider the following discretization of the Navier-Stokes equations in the Fourier domain
Uru + (u VuA) + vp = = nA
V
0
(32)
t=
PN
where PN is the spatial Fourier projection operator mapping a function f(z f) to its first N Fourier modes: PN EN fi(etee The artificial viscosity term we use for the stabilization of the solver consists of resolution-deeenderr viscosity e and a Fourier multiplier QN controlling the strength at which different Fourier modes are dampened. This allows us to not dampen the low frequency modes, while applying some diffusion to the problematic higher frequencies. The Fourier multiplier QN is of the form
(x)
(33)
kEZd,IkiSa
In order for the solver to converge, the Fourier coefficients of QN need to fulfill
66
70
31
Q = for |k| S mN, 1
MN
<
Qx S 1
(34)
k
21

---

where we have introduced an additional parameter 0 > 0. The quantities mn and eN are required to scale as
1
1
MN
No EEN
0 <<00
(35)
N
2
For the experiment described here, we choose the viscosity 4- 10-4 mentioned above. The Fourier multipliers QN are chosen according to and are equal to
EN
= 0.05 N' and N = 128. This gives rise
30
k
Q Smooth = -1- exp
(36)
ko
The Navier-Stokes simulations were performed with the AZEBAN spectral hyperviscosity code [6]
The Compressible Euler equations of gas dynamics are given by
Wt + div F(u)
(e, pe, F) F (pu, pu 2 (F +)]p]+
(37)
in the domain [0, 1] with suitable boundary conditions, with density p. velocity [Uz, Uy], pressure p and total Energy E related by the ideal gas equation of state:
1 E 2/p
P
(38)
1-1'
where = 1.4. All the trajectories are simulated until time T 1. The simulations for the compressible Euler equations were performed with the ALSVINN 44 code. which is based on high-resolution finite volume scheme with piecewise quadratic WENO reconstructions and HLLC Riemann solvers.
During pretraining, our goal is to predict four variables: [p, U2. py. p), where represents density, 2/2 is the horizontal velocity, By is the vertical velocity, and p is the pressure. As in the Savier-Stokes benchmarks, all the trajectories for compressible Euler are simulated until time T 1. Furthermore, we store 2/ snapshots of the numerically simulated solution, uniformly spaced in time. Each snapshot has a spatial resolution of 128 x 128, though being generated on 512 x 512 and downsampled.
Next we describe each constituent of the pretraining dataset (summarized in Table
B.1.1 NS-Sines
This dataset considers the incompressible Navier-Stolle Stokks equations 31 with the following initial conditions,
P
w (x,)) Σ
ai,j sin + Bi,)) sin(2 + n..))
2n (i +)
ij=1
(39)
(2,
cos (2rir + Br.s) cos(2 + 7h.i)
2i jj
iJ=1
where the random variables chosen Cij U[11J Big 2x and Yijj Mj.2m) The number of modes is chosen to be p = 10. Thus, the initial conditions amount to a linear combination of sines and cosines.
The underlying solution operator 8(1, is given by 8(t, u uw,(()) with Mgy My solving the Navier-Stokes equations 31 with periodic boundary conditions.
We generated 20000 NS-Sines trajectories of which the first 19640 belong to the training set, the next 120 to the validation set, and the last 240 to the test set. Note that we included 11 time steps in the pretraining dataset, with every other time step selected, starting from step 0 up to step 21. A visualization of random sample and the predictions made by POSEIDON-B B (trained on 128 training trajectories) shown in Figure
B.1.2 NS-Gauss
Given two-dimensional velocity field u = (wz, Myy) its vorticity is given by the scalar = curl u U. Uz Note that, for any time the velocity can recovered from the vorticity using the so-called stream function [47).
22

---

For this dataset, we specify the initial conditions for the Navier-Stoker equations in terms of the vorticity given by,
ai o2
(x xz) 2
y/)
(2 v)
exp
(40)
20
where we chose = 100 Gaussians with O;; ~ 11-1,11; og ~ M1.01000 11, Hj ~ (0,1) and y ~ M0.1) Thus, the initial vorticity field is superposition of large number of Gaussians. The initial velocity field is then recovered from the vorticity.
The underlying solution operator S(t, is given by S(t, (.,,)) = (t, with U2 My solving the Navier-Stokes equations (31 with periodic boundary conditions.
We generated 20000 NS-Gauss trajectories with the same train/validatie chiinvalidation/eest split and time-stepping as for NS- Sines. A visualization of random sample and the predictions made by POSEIDON (128 training trajectories) are shown in Figure 56
B.1.3 CE-RP
The well-known four-quadrant Riemann problem is the generalization of the standard Sod shock tube defined by dividing the domain D [0, 1]2 into a grid of p p
two-space dimensions 22
square subdomains
Dig
| D
T
Su 1 1,-
1- P
ps
P
where T the 2d torus. We fix 2 for this problem.
The initial data on each of these subdomains is constant and given by,
(Pa, u2 by po) = (Pin (Va)i.j, (y))s..P.)))
By sampling Pi.j 1.). (Ur)ks 11111,= (vy)i, ~ U-111), and pi,j ~ 1.0.1.1). we obtain stochastic version of the four-quadrant Riemann problem, which also generalizes the stochastic shock tubes (5) to two-space dimensions.
The underlying solution operator S(t, -) is given by S(t, Po, " 22 Po) [p(t), un. p()) solving the compressible Euler equations with periodic boundary conditions.
We generated 10000 CE-RP trajectories where the first 9640 trajectories belong to the training set, the following 120 to the validation set, and the last 240 trajectories to the test set. The time see-stepping the same as for NS-Sines and NS-Gauss. A visualization of a random sample and the predictions made by POSEIDON (128 finetuning trajectories) are shown in Figure 57
B.1.4 CE-CRP
This dataset corresponds to curved and multi-partitiones version of the CE-RP dataset. To define it, we denote the fractional part of E R as {a} := x - -|xx| sgn x and define the functions
cx(o, v) =
Σ
azdy sin(22iiz + jyy + Be.
$1==
oy(2.V) Σ
agyi, sin(2miz + jy + By.a.)-
:1==1
where on,i. U..11,0.11 and BRiis ~ Mj.11.- These functions are then used to create partition of the domain into curved subdomains,
Dis (z,,)) 2 Imin I for + oz (x, + 13 Tmax. Miin S (y+ v) +1 ymax
with Jmin -+I. Imax
#+1 p+i
1+1 Umin -tt and Umax P+I
Finally, the initial conditions are given
by
(p, Vo, ky,p) (pi.s, Mig Bhg Pi.s) in Da.s
23

---

where Pr.j ~ 1... (Vz).,
~ ~ (Vy)as ~ 1.1), and Pi,j Mo. A visualization of random sample of the initial conditions shown in Figure (8 and illustrates how this problem is curved, multi-partitiones version of the standard stochastic four-quadrant Riemann problem (CE-RP).
The underlying solution operator S(t.)) .) is given by 8(1, Po, 22 Po) [p(t), Vz.yy (4), solving the compressible Euler equations 37 with periodic boundary conditions.
We generated 10000 CE-CRP trajectories with the same train/validdationtess split CE-RP. The time-stepping is the same as for NS-Sines and NS-Gauss. visualization random sample and the predictions made by POSEIDON (128 training trajectories) are shown in Figure
B.1.5 CE-KH
This is well-known benchmark of compressible fluid dynamics that corresponds to the well-known Kelvin- Delvin-Helmholer instability (29). modern version presented, for instance, in [16].
The underlying initial data is,
(1, 0.5, 2.5)
y < 0.25 + oo(x) or y > 0.75 01 (1())
(p, Uas Ug p= =
(2, -0.5, 2.5)
otherwise.
The perturbations and 0r are given by
p
o;(x) Σÿ=1 Dz,j
8
oi cos((2](( Bey))
j=1
where e = 0.05, oi, Mo and Brg ~ U|..1.-
The underlying solution operator ((. is given by 8(1, Po, "2 Po) [p(t), U2 (1,,((()) solving the compressible Fuler equations 337 with periodic boundary conditions.
We generated 10000 CE-KH trajectories with the same train/validationntess split as CE-RP. The time-stepping is the same as for NS-Sines and NS-Gauss. A visualization of a random sample and the predictions made by POSEIDON (128 training trajectories) are shown in Figure 59
B.1.6 CE-Gauss
As in the NS-Sines dataset, we initialize the curl w of the initial velocity with superposition of Gaussians,
(x
xz 2 -
vi)
wo wo v)
exp
20
where we chose p 100 Gaussians with a; U|-1,1] 0, Ug..1. o.1]. x; ~ Uj.11 and yi Uy. Then, the initial field is generated from the vorticity by using the incompressibility condition. The underlying density and pressure are initialized with constants, p 0.1 and p 2.5, respectively.
The underlying solution operator S(t,) -) is given by S(t. Pu. ub... Pu) = [p(t), U2.y (t). solving the compressible Euler equations 37 with periodic boundary conditions.
generated 10000 CE-Gauss trajectories with the same train/validationntest split as CE-RP. Time- stepping is the same as for NS-Sines and NS-Gauss. A visualization of a random sample and the predictions made by POSEIDON (128 training trajectories) are shown in Figure 600
We remark that out of the 6 operators that consitute the pretraining dataset, 2 of them (CE-KH and CE-RP) are well known in the literature where the other four (NS-Sines, NS-Gauss, CE-Gauss, CE-CRP) are novel to the best of our knowledge.
B.2 Downstream Tasks
Next, we describe the suite of downstream tasks on which POSEIDON and baselines are evaluated. The list of tasks is summarized in Table
24

---

Table Abbrevitiions/Sumarry for all the downstream tasks. IC and RP stand for initial conditions and Riemann problem, respectively. Datasets where (*) is checked mark datasets where solutions are learned depending on PDE prannetes/oounee/oofiiiiitts
Abbreviation
PDE
Defining Feature
Visualization
NS-PwC
Navier-Stokes
31
Piccewise constant vorticity IC
Fig. 61
NS-BB
Navier-Stokes
Brownian Bridge IC
Fig. 62
NS-SL.
Navier-Stokes
Shear Layer IC
Fig.
NS-SVS
Navier-Stokes
Sine Vortex sheet
Fig. 64
NS-Tracer-PwC
Navicr-Stokes + Transport
Scalar Advection
Fig. 65
FNS-KF
Forced Navier-Stokes (33
Kolmogorov Flow
Fig. 66
CE-RPUI
Euler
RP with uncertain interfaces
Fig. 67
CE-RM
Euler 37
Richtneyer-Meshkor
Fig. 68
GCE-RT
Gravity
57
Sayleigh-Taylor
Fig. 69
Wave-Gate
Wave Eqn.
64
Waves in Gaussian medium
Fig. 70
Wave-Layer
Wave Eqn
64
Waves in layered medium
Fig. 71
ACE
Allen-Cahh Eqn.
Reaction-Diffusion
Fig. 72
SE-AF
steady state of Fuler
Flow past airfoil
Fig. 73
Poisson Gauss
Poisson Eqn.
68
Stationary diffusion
Fig. 74
Helmholtz
Helmholtz Fqn 69
Waves frequency domain
Fig.
B.2.1 NS-PwC
This downstream task is based on the Navier-Stokes equations 331 on the space-time domain [0, 1]2 [0, with periodic boundary conditions. The initial conditions are based on the vorticity, which is assumed to be constant along a uniform (square) partition of the underlying domain. To be more specific, the initial vorticity is given by,
wo(2,y) Ci.j in 1,,-1, [y
(41)
for 3/4 = for i 0, 1, p. and C. 1111111 The number of squares in each direction was chosen be p 10. Thus, this problem is an analogue of multiple Riemann problems, but the vorticity. The underlying initial velocity field (0), (0) is then recovered from the vorticity by using the incompressibility condition.
The underlying solution operator 8(1,)) is given by 8(1, v. = (t,,y(), with Wy My solving the Navier-Stokes equations 31 with periodic boundary conditions.
We generated 20000 NS-PwC trajectories with the same train/validationntest split as NS-Sines. Note that we included 8 time steps in the training dataset, with every other time step selected, starting from step 0 up to step 14. The testing error is evaluated the 14th time step (i.e. 0.7). A visualization of a random sample and the predictions made by POSEIDON-B, CNO and FNO by (128 training trajectories) are shown in Figure 61
B.2.2 NS-BB
(Fractional) Brownian Bridges are widely used as an initial conditions for the Navier-Stokes equations to study statistical properties of turbulent flows in the computational physics literature, see for instance 31 and references therein.
We generate Brownian Bridges directly in Fourier space with the following method:
W(x)= = Σ
Σ
(mnne)
sc, (x)se,, sc
(I))
(42)
k
|
SN
m.n.ee (0.1)
25

---

where
sin(2))
stemming all for i 0 for 1
(c,(2)
(43)
cos(x)
(mmn)
and the O'k 11111 These Brownian Bridges are propagated through the discretized Navier- Stokes system (32) from time -05 to t 0. The resulting flow fields are then taken initial conditions for this dataset.
The underlying solution operator S(t, :) is given by 8(t, u. = (t), with May My solving the Navier-Stokes equations (31 with periodic boundary conditions.
We generated 20000 NS- BB-BBB trajectories with the same train/validationntess split as NS-Sines. The same time-stepping is used as for NS-PwC. The testing error evaluated at the 14th time step (i.e. t= 0.7). A visualization of random sample and the predictions made by POSEIDON, CNO and FNO (128 training trajectories) shown in Figure 62
B.2.3 NS-SL
The Shear Layer (SL) is well-known benchmark for the Navier-Stokes equations 31 the way from [5], not earlier, see [31 for modern (stochastic) version, whose variant we consider here.
We take initial conditions the shear layer,
tanh 2x V-2.5
for y + as(x) S
uo u0(2, tanh 2x 0.75-v y)
otherwise
(44)
0 y)
where as: [0, R is perturbation of the initial data given by
(b(x)= = E Σ aA sin(2 Bk)-
(45)
The parameters are chosen to be U17,8...127 an Uo.111 Br Mg.277 5 0.025, Uj..08.0.12- and &
The underlying solution operator 8(t,) :) is given by 8(1, v. = (t), with uz My solving the Navier-Stokes equations (31 with periodic boundary conditions.
We generated 40000 NS-SL trajectories of which the first 39640 are in the training split, the next 120 in the validation split, and the remaining 240 in the test split. The same time-stepping is used as for NS-PwC. The testing error is evaluated at the 14th time step (i.e. t 0.7). A visualization of random sample and the predictions made by POSEIDON, CNO and FNO (128 training trajectories) are shown Figure
B.2.4 NS-SVS
The Sinusoidal Vortex Sheet (SVS) is another classic numerical benchmark for the Navier-Stokes equations [61] and references therein. We consider a modern (stochastic) version from [1] here. The initial datum for this problem is specified terms of the vorticity, by setting,
wi 1hp wo
(46)
26

---

where
wo(x) = 6(x --)- dr
(47)
T2
07(1)) = p 2/ |
(48)
(brr = 80
1
1 2 (r 1)
(r 1)
4(r
Gr 4(r
(49)
7
2
r = ((z, v) E T 1y + (.T.sn(2TT.)
a4 sin + (,))).
(50)
2
We choose p- 10 and the random variables a, and 3 are given by a, M0.00003125) 9 U(. The parameter p is chosen to be p T28
We generated 20000 NS-SL trajectories with the same training/validation/ees split as NS-Sines. The same time-stepping is used as for NS-PwC. The testing error evaluated at the 14th time step (i.e. t = 0.7). A visualization of random sample and the predictions made by POSEIDON, CNO and FNO (128 training trajectories) are shown Figure
B.2.5 NS-Tracer-PwC
This downstream task is the first of our tasks, where the underlying physics has not been completely encountered in the pretraining dataset.
In this experiment, we focus on the important problem of transport a passive tracer, for instance pollutant in river. This tracer is carried along by the Navier-Stokes flow field without feeding back into the velocity. Let c c(2, y, t) be the concentration of the passive scalar the fluid. The equation that governs is given by
dc
Ve RÂC,
(51)
where u is the velocity field of the fluid, which in turn is governed by the Navier-Stokes equations (31), and is the diffusivity constant. We choose to be equal the the artificial viscosity term used in the simulation of the flow (see B.1 for clarification).
The fluid velocity field u has the exact same initial data in the NS-PwC experiment. The tracer concentration initialized as a sphere centered the center of the domain
co(2, y) = 1 B3 ( () y)-
(52)
Thus, the source of stochasticity in this problem purely the random initial condition driving the fluid flow.
The underlying solution operator S(1,-) is now given by S(1, ug uy cn) (1), uy ((), c(O)), with My solving the Navier-Stokes equations (31) with periodic boundary conditions and e solving the transport equation 551)
We generated 20000 NS- OS--racer-Pwe trajectories with the same train/validation/test split as NS-Sines. The same time-stepping is used as for NS-PwC. The testing error is evaluated for the 14th time step (i.e. t = 0.7). A visualization of a random sample and the predictions made by POSEIDON, CNO and FNO (128 training trajectories) are shown in Figure 65
B.2.6 FNS-KF
Another downstream task which introduces a physical process that has not been encountered in the pretraining dataset, two-dimensional version of the well-known Kolmogorov Flow [47 is modeled by Navier-Stokes equations with a forcing term, namely
Ut + (u V)u + Vp =Au S, div u = 0,
(53)
in the domain (0, 1]2 with periodic boundary conditions. The forcing term f is chosen to be constant
time and is equal
f(z, v) 0.1 sin (2 v)).
(54)
27

---

The fluid velocity field u is initialized in the exact same way as in the NS-PwC experiment. The data simulated with the same method as the other flows governed by Navier-Stokes equations (see B.] for clarification).
We also remark that this problem be readily recast in the generic form (T) by considering the augmented solution vector U = [Mzz My f] and augmenting the PDE ① with the trivial equation ft and augmenting the initial data with (54). The underlying solution operator S(t,-) is then given by S(t, U2 with periodic boundary conditions and f being given by (54).
(t), uy (f), f], with Ug My solving the forced Navier-Stokes equations (53)
We generated 20000 FNS-Kr KF trajectories with the same ttinn/vlldaaiion/estt split as NS-Sines. The same time-stepping is used as for NS-PwC. The testing error is evaluated at the 14th time step (i.e. t= 0.7). visualization of a random sample and the predictions made by POSEIDON, CNO and FNO (128 training trajectories) shown in Figure 66
B.2.7 CE-RPUI
This downstream task considers the compressible Euler equations and is a variant of the uncertain interface problem considered in [50] well (hard) perturbation of CE-RP, where not just the amplitude of the jumps for each Riemann problem randomly varied, but even the location and shape of the initial interfaces is randomly perturbed. To realize this construction, we denote the fractional part any x E R as {x} == x [|||] sgn x and define the functions
(x = Σ
anii sin + 2p²) + (j + 22" + B2.b.)
i==
0y 3)
Σ
agiin sin (in((2 + 2p") Jz + (j + +p" + By.,,))
i==1
where anii. ~ U1 0100001.0.01) and Bu.i. ~ Ug..1. These functions are then used to create a partitioning of the domain into subdomains
Di, {xx,)) E T Imin S far + 02 (x( + 11 Tmaxx Danin I fy + o, v) +1 < ymax)
with Tmin p+I. Tmax by
++1 p+i. Vmin -t., and Ymax 211. Finally, the initial conditions given
(p, Uz, Uv p)l== = (pi,5, vi oii Pt.j) in Dis
where Pig ~ U[13), vf U 10,10- un ~ 10, and pt,j ~ Up.77
The underlying solution operator S(t. :) is given by 8(t, Po, 22.g. Po) = (p(t), (1), solving the compressible Fuler equations with periodic boundary conditions.
We generated 10000 CE-RPUI trajectories with the train/validationnnest split being the same as for CE-RP. The same time-stepping is used as for NS-PwC. The testing error is evaluated at the 14th time step (i.e. = 0.7). A visualization random sample and the predictions made by POSEIDON, CNO and FNO (128 training trajectories) are shown in Figure 67
B.2.8 CE-RM
Another well-known benchmark for the compressible Euler equations (37 is the Richtmeyer-Meshie Meshkoy problem, see [29]. A modern (stochastic) version is provided in [6 The compressible Euler equations are considered with the initial data given by,
Po(x, v) - -120 20 if 12 3d < 0.1
po(a, f2 if |a| < [(x, y. w.
w
0
otherwise.
otherwise
(55)
We assign periodic boundary conditions on D [0, 11a. The interface between the two states is given as
[x w, = 0.25 + a (w) sin 2) + b, (,())),
(56)
28

---

where K = 10, e >> O, and a and by (for j = 1...., K) are uniform random variables on the interval [0, 1]. We normalize the a, such that Ej 0, 1. We simulate up 2.
The underlying solution operator S(,,,) -) is given by S(t, Po, v..g. po) - [o(t), U2 (1) p(t) solving
the compressible Euler equations 37) with periodic boundary conditions.
We generated 1260 CE-RM trajectories with a train/validaaionnteest split of 1030/100/130, The approximate solutions where generated with the FISH hydrodynamic code, see [16] and references therein which implements high-resolution finite volume schemes. We save 21 snapshots, evenly spaced in time. The testing error is evaluated at the 14th time step (i.e. t = 1.4). A visualization of random sample and the predictions made by POSEIDON, CNO and FNO (128 training trajectories) are shown in Figure 68
B.2.9 GCE-RT
The compressible Euler equations with gravitation (GCE) are given by,
U4 div F(u) S, [p. pv, -1-, F= lpv, po @ pl, (E + pilo]t
04 + -o, Byy
0g
(57)
S = 0, 10, pv
dx
with P3 Uz,y> P be defined in (37) and yp being the gravitational potential.
For this experiment, we follow a well-known benchmark in astrophysics, namely the Rayleigh-Tayloo (RT) instability
model neutron star, realized as a 7 2 polytrope in gravitational equilibrium. Our benchmark is a two-dimensional stochastic variant of the setup of (28), Section 3.2.4, with the only variation being provided by the random fields used to generate the initial conditions. The domain D 1-1/2, +1/2|² 11/22 and the pressure and gravitational potential are given by
sin(ar) ar
(p(r) = K2Po sin(ar) or
p()) Ko Po
(58)
where r /2² + 2² is the radius, Ku = po/ p the polytropic constant,
4TG
(59)
2kg
and G is the gravitational constant. The initial velocity is set to vanish. The density profile is set
as
Ko
sin(ar)
p(r) = - R())
Pu
(60)
or
where
Ko,
r < TRT
k(()
2
(61)
Ko.
r 2 TRT-
I+A
Here, A is the Atwood number which parameterizes the density jump between the heavier and lighter fluid, characterizing the Rayleigh-Taylor instability. The interface between the fluids is given
RRT 0.25(1 + a cos (alan2(y, .) + b)):
(62)
where the amplitude a and phase are uniform random variables on [-1, and [-,, respectively. Similarly, we perturb the central density Po. pressure and Atwood number as
Po + 0.2c, Po + 0.2d, A 0.05(1 + 0.2e),
(63)
where e, d, are uniform random variables on [-1,1] We evolve the initial state up to a final time of T 5 and save 11 snapshots, evenly spaced time.
The new physical phenomena that we add in this case is gravitational forcing and the underlying solu- tion operator 8(t, -) is given by S(t, Po. o2.y og.yyPP. Po: sp) = (p(t), (t,,(t), sp] solving the gravitational
Euler equations (57) with periodic boundary conditions.
We generated 1260 GCE-RT trajectories (with the same train/validation/tess split as CE-RM) with a well-balanced second-order finite volume method, as described in (28). at 256 resolution, then downsampled to 1282. The testing error is evaluated at the 7th time step, and we take every snapshot up to and including the 7th as training data, visualization of a random sample and the predictions made by POSEIDON, CNO and FNO (128 training trajectories) are shown in Figure 69
29

---

point and B.2.10 Wave-Gauss
We consider the wave equation with spatially varying propagation speed, i.e. (c(z))² = 0, in D x (0,T),
(64)
in order to model the propagation of acoustic waves in a spatially varying medium.
The initial condition uo is given by of several Gaussians whose parameters at random. First, we draw a random integer n from the set {2,3, 3, 4, 5, 6}. Then, for S I n. we draw two locations, Fe,t Ve, /66,5666 We fix the amplitude of the ith Gaussian to 1.0 and draw the ith standard deviation 8; ~ Vij. 030.0399.011 Note that we restrict any two centers of the Gaussians to be closer than 2 standard deviations from each other. If this happens, draw discard the old one. The ith Gaussian defined as
drawn uniformly
(Fee x) (ye,, ) I,y (0,
gx exp
28
Finally, the initial condition is defined by
wo 3) Σ g((,,,)
2,y K (0.1)
(65)
We use absorbing boundary conditions. The propagation speed e is spatially dependent and is generated sum of Gaussians in several steps. First, random base speed is generated such that co ~ 411000 Then, we select points the domain, namely, (x1, V1) = (0.25, 0.25), (I',, 1/2 (0.25, 0.75). (Iz, 8/)) (0.75, 0.25) and (I4: M4) (0.75, 0.75). For each point we define a random vector (dx;, dy;). where da;;, dy: 011003125,0003 3.312 We also draw an amplitude vs 11 21000.2500| of Gaussian that corresponds the i-th point, as well as its standard deviation a; U1 1//12 The ith Gaussian is defined by
f((x, v) = Us exp (-
+ dx; - x²2 + (yi + dy. -
x, E (0, 1).
20
Finally, the propagation speed is defined by
c(x, v) = cu + Σ f( 2, E E (0,1).
Trajectories are generated with a finite-difference method, similar to the DeVITO code [43 at 128² resolution. The final time of all the simulations is T 1. We save 15 snapshots, evenly spaced in time.
Thus, this benchmark models the propagation of acoustic waves, generated by seismic sources, which propagate in smoothly varying medium. The wave equation (64) can be readily recast into the u, and the coefficient into the solution vector
generic form (T) by adding the time-derivative U [u(x, (), o(z, (), c(x)). Thus, the differential operator in (T can be rewritten as,
V, = Au, 0,
(66)
and the resulting solution operator is S(t, Un) = (t()), ct),
We generated 10512 Wave-Gauss trajectories with train/validationnnesst split of 10212/60/240. The same time-stepping is used as for NS-PwC. The testing error is evaluated at the 14th time step (i.e. I 0.7). visualization of random sample and the predictions made by POSEIDON, CNO and FNO are shown in Figure 70
B.2.11 Wave-Layer
In the Wave-Layer experiment, expeeriinen we also consider the wave equation with spatially dependent propaga- tion speed (64, initial conditions given by (65). We use absorbing boundary conditions.
30

---

The propagation speed e varies spatially and is generated as (vertically) layered medium, with each layer having a constant propagation speed drawn uniformly random. To generate one instance of c, we first draw random integer n from {3, 4, 5, 6). where represents a number of layers in Then, for each S n, we generate x depeedeentt frontier, defined by
ax(x) co
ai sin(2rix)),
n
where, first, a, values are drawn uniformly at random from (0, 1) and then co is drawn uniformly at random from (0, 1) and it is rescaled by constant that depends are impossible to intersect. Finally, a point (z. y) E (0, is in i-th frontier if and only if (z(z) s Each layer i has a constant speed of propagation
that the adjacent frontiers
y S (z)1 with 01 0 and an-1 ~ L/2000.50001 Trajectories are generated by fimite- mifte-difference method at 128 resolution. The final time of all the simulations I 1. We save 21 snapshots, evenly spaced in time.
As in the Wave-Gause benchmark, the resulting solution operator is S(t, Un) = ju(t), o(t), cl. with Uy and coefficient The wave-layer experiment models the propagation of acoustic waves, generated by seismic sources, inside a layered subsurface medium.
We generated 10512 Wave-Layer trajectories with the same train/validation/test split as Wave-Gauss. The same time-stepping used as for NS-PwC. The testing error is evaluated at the 14th time step (i.e. t 0.7). A visualization of random sample and the predictions made by POSEIDON, CNO and FNO are shown in Figure 71
We remark that both the Wave-Gauss and Wave-Layer tasks are very different from the pretraining dataset as the wave equation linear second-order (in time and space) equation that is different from the compressible Euler and incompressible Navier-Stoker equations that form the pretraining dataset.
B.2.12 ACE
The Allen-Cahn equation for modeling phase transitions in material science is given by
14 Au è² 2u(uu² 1),
(67)
with a reaction rate of e = 220. We consider this equation with periodic boundary conditions and initial conditions given by
K
1
uo uo v)
Σ
aij (82 (²) (in(zzz) sin(njy), Va, y (0, 1),
K2
where K is random integer drawn uniformly random from [16, 32]. Up.771100f and a.j U[-1.1]-
Trajectories are generated by a finite-difference method at 1282 resolution. The final time of all simulations is T 0.0002. We save 20 snapshots, evenly spaced in time.
The corresponding solution operator is S(t, uo) u(t) and maps the initial concentration to the concentration at time t.
We generated 1,5000 ACE trajectories with train/validation/tes: split of 14700/60/240. The same time-stepping is used as for NS-PwC. The testing error evaluated at the 14th time step. A visualization of random sample and the predictions made by POSEIDON, CNO and FNO (128 training trajectories) are shown in Figure 72
Again, it is essential to emphasize that the Allen-Cahn equation is nonlinear parabolic reaction- diffusion equation that very different from the PDEs used constructing the pretraining dataset.
B.2.13 SE-AF
This dataset contains the samples that describe the classical computational physics benchmark of flow past airfoils, modeled by the compressible Fuler equations (37). The samples are not time-dependent, interested in the steady-state solution.
31

---

Figure 3: Elliptic mesh for the airfoil problem
We follow standard practice in aerodynamic shape optimization and consider reference airfoil shape with upper and lower surface of the airfoil are located at (x, (2/ and (x, (x/c)) where c the chord length and yier and y corresponding to the well-known RAE2822 airfoil 46). The reference shape is then perturbed by Hicks-Henne Bump functions [88
15
(E) vle((E) + B.),,, = =
Σ
a B.(E).
i=1
In2
B₁(5) sin (T) qi In14 Ini'
{
c
at = 2(th - 01 + 1) 10-3, a. = 2 0.5) - 1) x 10 -
i= 1,.,,.,,11 15
with vi E [0, 1]4. We now formally define the airfoil shape as S {(x, v) DD:a [0,c),y S ソ and accordingly the shape function = X|s) (2. v), with being the characteristic function. funciiooo
The underlying operator of interest maps the shape function f into the density of the flow p steady state the compressible Fuler equations.
The equations are solved with the solver NUWTUN, see 45 and references therein, on 243 x 43 elliptic mesh (see Figure 3 given the following free-stream boundary conditions,
T
A = 0.729, p =1 a = 2.31%
The data is ultimately interpolated onto a Cartesian grid of dimensions 128 x 128 on the underlying
domain D = |-0.75. 1.75|?, and unit values are assigned to the density p(z. y) for all (r. y) in the set S. The shapes of the training data samples correspond to 30 bump functions, with coefficients sampled uniformly from [0, 1 330 During the training and evaluation processes, the difference between the learned solution and the ground truth is exclusively calculated for the points (s, y) that do not belong to the airfoil shape S.
We generated 10869 SE-AF solutions with a train/validation/tess split of 10509/120/240. 120/240. A visual- ization of random sample and the predictions made by POSEIDON, CNO and FNO (128 training samples) are shown in Figure 73
We note here that this SE-AF benchmark differs from what has been seen during pretraining in many aspects, namely the problem is indeeeeeenndnt, in contrast to the time-dependent PDEs for pretraining, ii) the solution operator is very different as it maps (shape) coefficient into the steady state solution, and iii) the geometry of the underlying domain is non-Cartesian and the boundary conditions are very different from what was encountered during pretraining.
B.2.14 Poisson-Gauss
We consider the Poisson equation,
Au f, in (0. 1)
(68)
32

---

with homogeneous Dirichlet boundary conditions. The solution operator maps the source term f to the solution The source term f consists of superposition of random number of Gaussians
(z (I- pr,;)² + (y Hy.s)
f(z,y) = Σ exp
20
with being an integer drawn from a geometric distribution Geom(0.4)) Ma.i Py, ~ U1. and 3, Uio 025,0.1J- Thus, this experiment models the diffusion of input (source) which superposition of Gaussians.
We generated 20000 Poisson-Gaue solutions (with train//allidationneest split of 19640/120/240)
with a fimite element method based on FENICS [40]. A visualization of a random sample and the predictions made by POSEIDON, CNO and FNO (128 training samples) are shown in Figure 74
We note here that both the Poisson-Gaute Gauss and Helmholtz benchmarks differ from what has been seen during pretraining in many aspects, namely i) the problems are time-independene in contrast to the time-dependent PDEs for pretraining, ii) the solution operator very different as maps coefficients into the steady-state solution, and iii) the boundary conditions are very different from the periodic boundary conditions, seen during pretraining.
B.2.15 Helmholtz
The Helmholtz equation models wave propagation in the frequency domain. We consider variant of this equation given by
- - waa(x, O, 2, E D,
(69)
and Dirichlet boundary conditions
u( y) =b x.y E OD,
where 5 is the frequency, D (0, 1)2 is the domain, is the spatial dependent function that defines properties of the medium of the wave propagation and b is the fixed value of the solution at the boundary ØD.
The boundary value follows uniform ditribution, namely ~ 2.250.05 The function a is defined sum of random number of Gaussians and generated in several steps. First, we draw an integer from [2, 7] uniformly at random. This number represents the number of Gaussians that will be randomly generated. For 1 S S 7i, it holds that A; ~ 1... and a; 01..05...1)- Additionally, two numbers 24, that represent x and coordinates of the Gaussians are generated such that Te B4 M0.2.0.8. The unnormalized function is obtained by
a(x, v) = Σ A exp (-(m=- (x)) + (yi ))² x, 1).
20?
Function is obtained by normalizing i.e.
ā min(a)
a(x, max min(a)
The solution operator maps the tuple (a, 6) to the solution This problem is a steady-state problem.
Trajectories are generated by a finite-difference method at 128 resolution, similar to De Vito [3]
We generated 19675 Helmholtz solutions with train/validation/eest split of 19035/128/512. A visualization of random sample and the predictions made by POSEIDON, CNO and FNO (128 training trajectories) are shown in Figure 75
33

---

C Models and Baselines
We compare multiple versions of POSEIDON with foundation model baselines, namely MPP [49 and a CNO [60] foundation model that is trained in a similar manner as POSEIDON, In addition, we compare against state-of-the-ar neural operators trained from scratch, namely CNO and FNO [3 as well as scOT trained from scratch. All these models and their training strategies are described in the following. Their approximate model sizes can be read off from Table
Table 5: Approximate model sizes of all the models considered in this paper.
Model Number of parameters
POSEIDON-L
629M
POSEIDON-B 158M
POSEIDON-T
21M
CNO-FM
109M
MPP-B
116M
CNO 39M
scOT 40M
FNO
37M
We train all models a realization of Equation in the Main Text; in particular, set p 1. For each gradient step, we draw from the set of N avaiTable trajectory snapshots {u', |ul, E RoxJxJ ) where is the number of input/output functions, is the size of the computational grid, and t, the (lead) time from the initial condition to the k-th snapshot in the trajectory, i.e. we get a batch of size B of pairs {uu)) ( where i S j. The loss is then computed
ΣE
11 Ev (ur) D Sp((t,)² - (ta). (u.)) (..,))
C([(u)) uf (2)
(70)
c
ΣB Eu,==1 (17)) +
where c 10-10 for numerical stability, 8p is the model, (.) is the 1-th sample from the batch and (.),,,y denotes the value at indices (s, u, v).
During training, we create a checkpoint after every epoch, but only keep the checkpoint corresponding the lowest validation loss (evaluated at the end of the epoch) which is then also used for testing.
C.1 POSEIDON Models
In the following, we give thorough details for all POSEIDON models that we pretrained, as well as details on finetuning these pretrained models. All models are pretrained on the datasets introduced in Section[..r i.c. they expect four dimensional inputs and outputs, density p. velocities u and v, and pressure p. During pretraining, set p = 1 and mask for all pretraining datasets corresponding to incompressible flow. Further, we use the full set of 77840 pretraining trajectories, unless otherwise specified.
To finetune the pretrained model on tasks whose input/output functions are the set of pretraining input/outputs (p, u, or where there are additional inputs/outputs this corresponds to the tasks NS- PWCC FNS-KF, SE-AF, GCE-RT, Wave-Layer, Wave-Gauss, ACE, Coisson-Gauss, and Helmholtz we transfer all parameters from the pretrained model, except
the embedding weight We.
the patch recovery weight WR and bias bx, and
the mixup convolutional kernel.
We refer to Section A.2 for notation. This means that solely embedding/recovery is trained from scratch and just the parameters whose dimensions would not match, i.c. minimal set of parameters
34

---

is trained from random initialization. All other parameters are kept trainable and no parameter is frozen.
In general, we do not apply any weight decay on (time-conditioned) layer norm parameters. We finetune all parameters using the same optimizer with different learning rates, i.e. build two three parameter groups, depending on the finetuning task. In case the embedding and recovery do not have to be replaced, finetune all parameters, except parameters of the (lim-condittined) layer norm, with learning rate 1j, and parameters of the layer norm with learning rate TiN' If the embedding/recovery to replaced and trained from scratch, we finetune all embedding/recover parameters (including the embedding bias bg) with learning rate 17, layer norm parameters with learning rate IJ,, and all other parameters with learning rate 17.
C.1.1 POSEIDON-T
POSEIDON-Ti the smallest pretrained model, an instantiated scOT with the following hyperparame-
ters:
Embedding/latent dimension C: 48
Number of Swin V2 transformer blocks at each level (vi) &
This results in model with 21M parameters (for POSEIDON models, we exclude embedding and recovery parameters in this count).
The model is pretrained on 8 NVIDIA RTX 4090 GPUs using the following (data- parallel) training protocol:
Pretraining
Optimizer: Adam W 41
Scheduler: Cosine Decay with linear warmup of 2 epochs
Maximum learning rate: 10-3
Weight decay: 0.1
Effective batch size: 640, resulting in a per-device batch size of 80
Number of epochs: 40
Early stopping: No
Gradient clipping (maximal norm):
The pretrained model is finetuned on every task on single GPU following this finetuning protocol (r) only applicable to certain downstream tasks):
Finetuning
Optimizer: Adam
Scheduler: Cosine Decay
Initial learning rate îj: 5 1
Initial learning rate 7/: 5 10
Initial learning rate 7iv 5 10-4
Weight decay: 10-5
Batch size:
Number of epochs: 200
Early stopping: No
Gradient clipping (maximal norm):
C.1.2 POSEIDON-B
POSEIDON-B is the base model, an instantiated scOT with the following hyperparameters:
Embedding/latent dimension C:
Number of Swin V2 transformer blocks at each level (vi) &; 8
This results in model with 158M parameters.
35

---

The model is pretrained on 8 NVIDIA RTX 4090 GPUs using the following (data- parallel) training protocol:
Pretraining
Optimizer: Adam 41
Scheduler: Cosine Decay with linear warmup of 2 epochs
Maximum learning rate: 5 10-4
Weight decay: 0.1
Effective batch size: 320, resulting in a per-device batch size of
Number of epochs: 39 (40 were initially planned)
Early stopping: No
Gradient clipping (maximal norm):
The pretrained model finetuned on every task on a single GPU following this finetuning protocol (i) is only applicable to certain downstream tasks):
Finetuning
Optimizer: Adam W
Scheduler: Cosine Decay
Initial learning rate îý: 5 10
Initial learning rate îj: 5 10-4
Initial learning rate Tiv: 5 10
Weight decay: 10-6
Batch size: 40
Number of epochs: 200
Early stopping: No
Gradient clipping (maximal norm):
C.1.3 POSEIDON-L
POSEIDON-L is the largest model we trained, an instantiated scOT with the following hyperparame-
ters:
Embedding/latent dimension C: 192
Number of Swin V2 transformer blocks at each level (vi) t;:
This results a model with 629M parameters.
The model is pretrained on NVIDIA RTX 4090 GPUs using the following (data- parallel) training protocol:
Pretraining
Optimizer: Adam W
Scheduler: Cosine Decay with linear warmup of I epoch
Maximum learning rate: 2 10-4
Weight decay: 0.1
Effective batch size: 128, resulting in per-device batch size of 16
Number of epochs: 20
Early stopping: No
Gradient clipping (maximal norm):
36

---

The pretrained model is finetuned on every task on a single GPU following this finetuning protocol (i) only applicable certain downstream tasks):
Finetuning
Optimizer: Adam 41
Scheduler: Cosine Decay
Initial learning rate fy: 5 10
Initial learning rate 77: 10
Initial learning rate Niv 5 10-4
Weight decay: 10-6
Batch size: 16
Number of epochs: 200
Early stopping: No
Gradient clipping (maximal norm):
C.1.4 Models for Dataset Ablations (see Section D.3
For models used in the pretraining dataset ablations, utilize the pretraining and finetuning strategies as for POSEIDON-B BB. For the model trained on half of the pretraining dataset, we only train on the first half of each subset (NS-Sines, NS-Gaussia NS-Gaussians, CE-RP, CE-CRP, CE-KH, CE-Gauss); the same logic applies to the model trained on an eighth of the pretraining dataset. The model trained on less diverse pretraining dataset is not trained NS-Sines, CE-CRP, and CE-Gauss, such that the pretraining dataset size is directly comparable to the model trained on half of the pretraining dataset.
C.2 SCOT
We additionally train scOT from scratch every downstream task, to compare its performance to POSEIDON and other baselines. Its hyperparameters are as follows:
Embedding/latent dimension C: 48
Number of Swin V2 transformer blocks at each level (vi) 4; 8
This results a model with 40M parameters. is trained on one or multiple GPUs (depending on the dataset size) with the following parameters:
Optimizer: Adam W 41
Scheduler: Cosine Decay with linear warmup of 20 epochs
Maximum learning rate î): 5 10-4
Weight decay: 10-6
Batch size: 40 (on single GPU, else the effective batch size is larger)
Number of epochs: 400
Early stopping: If the validation loss does not improve for 40 epochs
Gradient clipping (maximal norm):
C33
CNO
A Convolational Neural Operator (CNO) is model that (approximately) maps bandlimited functions bandlimited functions (6). Let B., be the space of bandlimited functions with the bandlimit w.
CNO is compositional mapping between function spaces G B., (D) - B., (D) and is defined as
G:u P(u)
+
+ +((()) = u.
(71)
where
1 Pi o Ex KC(v)), 1515
(72)
where L is the number of CNO blocks and D (0, 1)a is the domain.
37

---

First, the imput function u E (D) is lifted to the latent space of bandlimited functions through lifting layer:
P: {u Bw (D, + {vo (D,
Here, do 2 dx is the number channels in the lifted, latent space. The lifting operation is performed by convolution operator and activation operator which will be defined below.
Then, the lifted function is processed through the composition of series of mappings between functions (layers), with each layer consisting of three elementary mappings, i.e., Pi is either the upsampling or downsampling operator, K2 is the convolution operator and E, is the activation operator.
Finally, the last output function in the iterative procedure is projected to the output space with projection operator Q, defined as
Q: {UL (D) + { (D,
The projection operation is also performed by a convolution operator and activation operator.
Upsampling and Downsampling Operators. For some w7 >> w, we can upsample function f E B the higher band B by simply setting.
Uw B (D) + (D( Uw) f(x), Vr E D.
On the other hand, for w w, downsample function B to the lower band B by setting Dw. (B((D) + (D), defined by
2 / ()drr-yf(l)dg) Vo D, w
D fw))(()
(h * (x)
where is the convolution operation on functions defined above and ha is the so-called interpolation sine filter:
fw (zo, x1) = sinc(2wwro) sinc(2wxx), (zo, E R..
(73)
Activation Operator.
First, the input function f E B. is upsampled to higher bandlimit w > w, then the activation function is applied and finally the result is downsampled back to the original bandlimit Implicitly assuming that w is large enough such that a (Bw) C Ba. define the activation operator in (7T) as,
En, (D) (D), Ex,zf(x) D o uu (x)), V D.
(74)
The above ingredients are assembled together in the form of an Operator U-Net architecture that has bandlimited functions as inputs and outputs. In addition to the blocks that have been defined above, one also needs additional ingredients, namely incorporate skip connections through ResNer blocks of the form, Rw,, Bi (D, Rd) (D), R.d) such that
(w) K. En (w)),, VU Bw (D,
(75)
Additionally, the so-called Invariant blocks of the form, Inn. B., (D, Rd) + B (D, is defined such that
(v) Ew,w (w), Ve E (.,
(76)
Finally, all these ingredients are assembled together in a modified Operator U-Net architecture which graphically depicted in Figure Note that instead of lead-time conditioned layer normalization we incorporate a lead-time conditioned instance normalization into CNO. A lead-time conditioned
instance normalization is applied to input by
(N (v) = a(t) O (N) + ((1)
(77)
where IN(v) is a regular instance normalization. the case of CNO, we use (small) MLPs to parametrize att) and (t)). This choice of conditional layer is similar to the FILM layer introduced in (55), applied on top of the instance normalization. Additionally, we observed that including time t
additional, constant input channel of the CNO slightly enhances its performance.
The specifications of the CNO model that we used and trained from scratch in all the experiments, as well as the training details are summarized in the following list:
38

---

Lifting dimension: 54
Number of up/downsampling layers: 4
Number of residual blocks in the bottleneck: 6
Number of residual blocks the middle layers: 6
Trainable parameters: 39.1M
Optimizer: AdamW [4]]
Scheduler: Linear with decreasing factor of 0.9 every 10 epochs
Initial learning rate: 5 10
Weight decay: 10-6
Number of epochs: 400
Batch size: 32
Early stopping: the validation loss does not improve for 40 epochs
Source code for CNO is available at httpss/gah.....oooooommc///o/oot
Scales:
Gaasa
Block
G(u)
Lio
www
Σ
20
Gooe
Block
Layer
Figure 4: Schematic representation of CNO 71 as a modified U-Net with a sequence of layers
mapping between bandlimited functions.
C.4 FNO
Fourier neural operator (FNO) 333 is composition
9: - G = Q Lr o... L1 R.
(78)
has a "lifting operator" u(x) -> R(u(z),)) æ), where R is represented by linear function R Rd Rdd where d, is the number of components the input function and d, is the "lifting dimension". The operator is non-linear projection, instantiated by shallow neural network with single hidden layer and leaky ReLU activation function, such that pl+ (x) + ((u)(z) = Q (g(+((x)
Each hidden layer Le : (z) + (x)) of the form (x) = INo (We (z)) + (Kro') (z))
with We Xd, trainable weight matrix (residual connection), an activation function, corre- sponding to leaky ReLU, IN standard instance normalization or time-conditioned instance normal- ization (see Equation 77) and the non-local Fourier layer,
Kgv = F (Pa(k) Fru
where Fk) denotes the (truncate))-Fouricc coefficients of the discrete Fourier transform (DFT) of of (zz), computed based on the given J grid values in each direction. Here, P;(k) E xd, is complex Fourier multiplication matrix indexed by & E zá, and g denotes the inverse DFT. As with CNO (Section C.3) we include time as additional channel in addition the time-conditioned instance normalization layers for all time-dependent problems.
We used the following hyperparameters and training details to train the FNO models:
39

---

Lifting dimension: 96
Number of Fourier layers: 5
Number of Fourier modes: 20
Trainable parameters: 37.0M
Optimizer: AdamW 41
Scheduler: Cosine Decay
Initial learning rate: 5 10
Weight decay: 10
Number of epochs: 400
Batch size:
Early stopping: the validation loss does not improve for 40 epochs
C.5
CNO-FM
In addition to the POSEIDON models, we also pretrain a CNO foundation model baseline. We use the same pretraining datasets as the POSEIDON models (see I.T). i.e. NS-Sines, NS-Gauss, CE-RP, CE-KH, CE-CRP and CE-Gauss datasets. The inputs and the outputs of the model have 4 channels, ie. P, u, and p. For the NS-Sines and NS-Gauss datasets, we mask out the pressure predictions during training, while predicting a constant value 1 for density.
The specifications of the CNO-FM model that pretrained, well as the training details summarized in the following list:
Lifting dimension: 82
Number of up'downsampling layers: 4
Number of residual blocks the bottleneck: 8
Number of residual blocks in the middle layers: 8
Trainable parameters: 109M
Optimizer: AdamW
Scheduler: Linear with decreasing factor of 0.9 every epoch
Initial learning rate: 5 10-4
Weight decay: 10
Effective batch size: 256, resulting in a per-device batch size of 32
Number of epochs: 40
Early stopping: No
To finetune the CNO-FM, we differentiate between two scenarios: one where the input and output share the same context as the pretrained models (comprising the variables p. u, v, and p. either masked or unmasked), and another where the downstream task is out-of-context (i.e. when the input and target variables differ from those used during pretraining).
To explain the finetuning technique, let us denote the pretrained CNO model by Grm and decompose it
GFM = Q GFMB P.
where the lifting layer, Q is the projection layer and Gras, the base part of the CNO-FM.
When the context of variables is retained the downstream task, we introduce an additional linear layer L that is applied prior to the lifting layer P. All other parameters from 9FM transferred over to the downstream task model. Hence, the model that is finetuned is
GFT Q GFMB Po L.
(79)
schematic representation of the CNO-FM finetuning procedure is shown in Figure When the downstream task is out-of-comtext, in addition to the linear layer L that is applied before P. the
40

---

projection layer is replaced by new. randomly initialized projection layer Q*. Other parameters are transferred over to the downstream task model. The model that is finetuned is
GFT Q* 9FMB P.
(80)
We set the number of epochs for the downstream tasks to 200. Since the loss converges significantly faster than when training from scratch, even 50- 100 epochs were sufficient effectively finetune the CNO-FM. Parameters of Grt are divided into three distinct groups
Group 1: Projection Q (or Q*), Lifting P and Linear layer C
Group 2: All the conditional instance normalization layers / (N(t))3())
Group 3: Other parameters in GF
We experimented with learning rates for each group of parameters, well as schedulers. An efficient way to finetune CNO-FM for in-context downstream tasks was to set the initial learning rates of the parameter groups to lr, 2.5 10-1 Ir2 = 5 10-1 and Ir3 10 10-1_ For out-of-context tasks, the learning rates that we used are Ir; = 7.5 10-1, Ir2 = 5 10-4 and Ir3 = 10-4 In both cases, the learning rate scheduler is linear with with decreasing factor of 0.9 every epochs.
The CNO codes are available at tttp://ithiccaooooooommtt//////m
CNO
Base
Lift Project
Input
output
PRE - TRAINING
CNO Base
Input
4d data
output
Train from scratch:
Fine Tune:
CNO
Base
Linear
New Project
OF OF- FINETUNING
-
Figure 5: Schematic representation of the finetuning procedure of CNO-FM.
C.6 MPP
Multiple physics pretraining (MPP) is pretraining approach for autoregressive physical surrogate modeling [49] MPP uses a scalable axial attention transformer backbone to reduce the quadratic complexity of the usual attention mechanism. Multiple input fields of MPP are projected onto single, shared embedding space. MPP also uses spatial and temporal attention blocks to capture spatial and temporal dependencies in the data. To train finetune MPP models, one uses the normalized MSE loss. We will finetune the MPP-AVIT-B foundation model for all our downstream tasks. The MPP-AVIT-B model has 116M trainable parameters.
MPP models are autoregressive models with fixed context size of S They predict the solution at a time step N of PDE of interest given the previous S time steps. Thus, they rely the history of the solution, encompassing multiple time steps, to forecast future time steps accurately. This differs interested in (i.e. OLT defined in the Main Text), which aims generate
from the the task that the entire solution trajectory given only the initial datum and boundary conditions.
Therefore, need to adjust the MPP finetuning strategy. We adapt the all2all strategy. Let = (Wo, ......, uT) be solution trajectory of length + 1. Let (i,j) be two integers such that j > i. We rely the fact that MPP predicts snapshot time and finetune MPP to predict u, based on the history uj-1, Ug-2 U, Since there not always 8 past time steps in the training
41

---

samples, we fill the remaining time steps with copies of u, (see Figure 6 We generate (T 1)/2 samples out of the trajectory U. For steady-state operators of the form f2......... fr. - u, the L channels are copied 7 times, and MPP is finetuned using these samples (see Figure 6 The inference strategy for the time-dependent problems is straightforward. Given the initial snapshot up. one autoregressively applies the finetuned model T times to predict uT (see Figure 6). During the finetuning of MPP, we do not predict dummy variables like the speed function the Wave equation or the forcing term in Kolmogorov flow, as the model had difficulties in predicting them, so the errors accumulated fast. This contrasts with other models that predict the dummy variables alongside the solution. Final testing errors for all the models are not calculated for these dummy variables.
For each downstream task. we finetuned MPP-AVIT-B model for 100 epochs. We did not use more than 100 epochs as the training usually converged after 10 to 50 epochs. We used the Adam [25] optimizer with a cosine annealing scheduler and linear warmup.
FINE - TUNING MPP (aliZall)
INFERENCE
Ho
Me
MPP
MPP
Me
...
M4
MPP
Mo
No
3
MPP
Mo
12:
22
12
FINE TUNING MPP (steady)
MPP
No
Hn
..0
F7
NT
MPP
context size T
Figure 6: Schematic representation of MPP all2all and steady finetuning and inference strategies.
42

---

D Results
D.1 Performance on Downstream Tasks
We evaluate all models on the median relative L' error at a certain snapshot in time on each solution function of interest. For vectorized functions such as the velocity fluid flow, we evaluate over the entire vector. Since all our downstream tasks have different solution functions of interest, we provide an overview over the actual functions of interest in Table T should there be list, we compute the mean over all metrics. It also depicts, how the rollout in time was done for each task i.e., either the solution at final time is computed directly with the final time as lead time or autoregressively (AR) as presented in Main Text (8). Within AR, POSEIDON models were always evaluated with uniform (homogeneous) rollout whereas in the case of CNO and CNO-FM, autoregressive (AR) rollout is heterogeneous; for MPP, it is always uniformly autoregressive.
Evaluation of Downstream Tasks with Scaling Plots. In Figures 7 to 21 we present the test errors (y-axis) for all the models (POSEIDON-L. POSEIDON-B, CNO-FM, MPP, FNO on the left sub-figure and POSEIDON-L, POSEIDON-B, CNO, ScOT and FNO on the right sub-figure of each figure) vs, the total number of trajectories (time-dependent PDEs) or total number of samples (for time-dependent PDEs) on the x-axis. As mentioned in the Main Text, the POSEIDON models clearly outperform all the baselines on most of the tasks as the corresponding test errors are significantly lower than the baselines for the same number of samples. Note that the metries EG and AG (TT) were computed
based on these plots. We do not include the POSEIDON-T results as they would further clutter the scaling plots. However, the EG and AG metrics for POSEIDON-T are presented in Table [8]
On Scaling Laws. If we denote the number of trajectories (samples) by M, we can fit power laws of the form,
(81) to the scaling plots in Figure to 21 Here, Cmodel denotes the model-specific scaling factor and the scaling exponent is Ommodel. The scaling exponents, resulting from these fits are presented in Table 7 We observe from this table, that all the models that we consider obey scaling laws of the form (81) with different scaling exponents for different problems (MPP-B does not converge in some cases). These include the POSEIDON foundation models which show consistent scaling laws. For instance, POSEIDON-L. has scaling exponent of approximately 0.5 higher in all the cases except for CE-RM (where all models converge very slowly). Nevertheless, we would like to emphasize that the scaling exponent alone does not govern the final error, except the asymptotic infinite data limit. Rather, the scaling factor Cmodel in (8 plays a decisive role in determining errors in the pre-asymptotic limited data regime that all downstream tasks correspond to.
Emodel (M) R Cuodee] -0model
Moreover, a closer analysis of the scaling plots reveals a more nuanced picture for the POSEIDON models. In some of the downstream tasks, for instance the Poisson Gauss benchmark, see from the scaling plot Figure 20 that both POSEIDON-L and POSEIDON-B display a biphasic behavior, with scaling law of the form,
E Emodel (M) = model M -ooddda if M S MP
Chu [ci C'modde M -aaaaaa Lnndda
model
(82)
if
M 2 M
model
with a <aa. Thus, the scaling behavior characterized by two phases, with different exponents. For instance, for the Poisson-Gauss benchmark (Figure 20), we find that M = 32 for both models. Moreover, for POSEIDON-B, Oh = 0.23 and az = 0.99 and for POSEIDON-L, Q'g: 0.33 and ag = 0.94. We speculate that these phase transitions separate two phases, a warmup phase where POSEIDON is slowly learning about an operator that is very different from those encountered in the pretraining dataset (as is the case with Poisson- Gauss) and learning phase, where fast learning takes place and the model is able to quickly learn the specifics of the downstream task.
Summarizing Downstream Task Performance. In Table 9
provide statistical summary of the performance of all models on all downstream tasks by presenting the (median) EG and the (mean) AG over all tasks. These statistics provide an (average) account of model performance over all downstream tasks and clearly quantify how the POSEIDON family of foundation models significantly outperforms all the baselines.
43

---

Table 6: The evaluation metrics are computed for each downstream task on different functions of interest, and rollout done differently.
Downstream Task
Functions of Interest
Rollout
NS-PwC
(Uay Myy
AR
NS-SVS
(Uny My)
AR
NS-BB
(Ug, Myy
AR
NS-SL
(Ug, Myy
AR
NS-Tracer-PwC
(Ugy UyJ+ C
AR
FNS-KF
(Uzy Uy)
direct
CE-RPUI
P. (Ux 4),
AR
CE-RM
P. (Vz, Uy)
direct
SE-AF
direct
GCE-RT
P. (Uz, Uy), P. Φ
direct
Wave-Layer
direct
Wave-Gauss
u
direct
Maison-Gauss "
ACE
direct
direct
Helmholtz
direct
Table 7: Scaling exponents with power law
(T)
Dataset
POSEIDON-B
POSEIDON-L
scOT
CNO-FM
CNO
MPP-B
FNO
NS-PwC
0.36
0.49
0.52
0.43
0.62
0.55
0.34
NS-SVS
0.49
0.48
0.77
0.56
0.55
0.53
0.02
NS-BB
0.32
0.51
0.57
0.47
0.64
0.52
0.40
NS-SL
0.36
0.47
0.48
0.46
0.45
0.45
0.59
NS-Tracer-PwC
0.78
0.66
0.59
0.41
0.56
0.43
0.44
FNS-KF
0.98
0.56
1.04
0.30
0.45
0.29
0.43
CE-RPUI
0.35
0.37
0.43
0.27
0.32
0.05
0.23
CE-RM
0.10
0.11
0.09
0.10
0.11
-0.20
0.11
SE-AF
0.30
0.32
0.27
0.35
0.13
0.31
0.24
GCE-RT
0.53
0.59
0.44
0.47
0.31
0.11
0.43
Wave-Layer
0.57
0.51
0.33
0.40
0.51
0.13
0.43
Wave-Gauss
0.59
0.50
0.45
0.37
0,46
0.07
0.34
ACE
0.74
0.85
0.77
0.72
0.47
0.46
0.88
Poisson-Gant
0.99
0.94
1.07
0.67
0.50
0.71
0.61
Helmholtz
0.38
0.43
0.68
0.42
0.54
0.27
0.31
44

---

Table 8: Efficiency gain (EG) and Accuracy Gain (AG) for the POSEIDON models all downstream tasks.
Pretrained Models
Scratch
POSEIDON-L
POSEIDON-B B
POSEIDON-T
FNO
EG
AG
EG
AG
EG
AG
EG
AG
NS-PwC
890.6
24.7
1024.0
19.7
1024.0
19.8
1
NS-SVS
502.9
7.3
518.9
7.9
212.0
6.
NS-BB
552.5
29.3
816.0
14.7
365.0
19.4
NS-SL
21.9
5.5
19.1
4.7
9.7
3.7
1
NS-Tracer-PwC
49.8
8.7
20.4
5.4
35.1
6.2
1
FNS-KF
62.5
7.4
16.1
4.7
77.9
5.9
1
CE-RPUI
352.2
6.5
370.8
6.2
909.7
5.8
1
CE-RM
4.6
1.2
3.1
1.1
2.8
1.1
1
SE-AF
3,4
1.2
2.9
1.2
2,4
1.1
1
GCE-RT
5.3
2.0
3.2
1.5
1.7
1.2
1
Wave-Layer
46.5
6.1
24.9
4.7
14.5
3.4
1
Wave-Gauss
62.1
29.3
4.3
19.5
3.1
ACE
17.0
8.7
6.5
9.8
7.2
1
Poisson- Gasson
42.5
20.5
24.4
13.0
18.2
8.4
1
Helmholtz
78.3
6./
64.7
5.0
64.7
4.9
1
Table 9: (Median) Efficiency gain (EG) and (Mean) Accuracy Gain (AG) over all downstream tasks for all models. We also present N(EG) as the number of tasks for which the EG of the model is greater than 10 and N(AG) as the number of tasks where the AG of the model is greater than 2.
Median
N(E))
N(AG)
POSEIDON-L
49.8
9.58
12
13
POSEIDON-B
24.4
6.71
11
12
POSEIDON-T
19.5
6.49
10
12
CNO-FM
10.6
2.91
10
MPP-B
2.0
1.82
3
CNO
4.6
2.61
ScOT
2.57
45

---

Poseidon-B
Poseidon-I.
CNO-FM
MPP-B
FNO
scOT
CNO
Median relative I! error
0.100
0.010
0.001
32
256
2048
32
256
2048
Number of trajectories
Figure NS-PwC. Number of trajectories vs. median relative L' error on the set.
Poseidon-B
Poseidon-L. scOT
CNO-FM
MPP-B
FNO
CNO
1.000
Median relative I error
0.010
0.001
32
256
2048
32
256
2048
Number of trajectories
Figure 8: NS-SVS. Number of trajectories median relative T error the test
Poseidon-B FNO
Poseidon-L scO'T
CNO-FM CNO
MPP-B
1.000
Median rolative [1] error
0.100
0.010
0.001
32
256
2048
32
256
2048
Number of trajectories
Figure 9: NS-BB. Number trajectories vs. median relative I,1 error on the test set.
46

---

Poseidon-B
Poseidon-I.
CNO-FM
MPP-B
FNO
scOT
CNO
Median relative I! error
0.100
0.010
0.001
32
256
2048
32
256
2018
Number of trajectories
Figure 10: NS-SL. Number of trajectories vs. median relative L' error on the test set.
Poseidon-B
Poseidon-L scOT
CNO-FM
MPP-B
FNO
CNO
1.000
Modian relative L' error
0.100
0.010
32
256
2048
32
256
2048
Number of trajectories
Figure 11: PW-.r Number of trajectories vs. median relative I error on the test set.
Poseidon-B
Poscidon-L scOT
CNO-FM
MPP-B
FNO
CNO
1.000
Median relative L error
0.100
0.010
32
256
2048
32
256
2048
Number trajectories
Figure 12: FNS-KF. Number of trajectories vs. median relative I,1 error on the test set.
47

---

Poseidon-B
Poseidon-I.
CNO-FM
MPP-B
FNO
scOT
CNO
Median relative I. error
0.100
0.010
256
2048
32
256
2048
Number of trajectories
Figure 13: CE-RPUL. Number of trajectories vs. median relative L' eπror on the test set.
Poscidon-B B
Poscidon-L scOT
CNO-FM
MPP-B
FNO
CNO
1.00
Median relative L? error
6 x 10
x 10
3 x 10
2 10
0.10
32
128
512
32
128
512
Number of trajectories
Figure 14: CE-RM. Number of trajectories median relative L error the test set.
Poseidon-B
Poseidon-L scO'T
CNO-FM CNO
MPP-B
FNO
1.000
Median relative Lt error
0.100
0.010
4
32
256
2048
32
256
2048
Number of trajectories
Figure 15: SE-AF. Number of samples vs. median relative I.1 error on the test set.
48

---

Poscidon-B
Poscidon-L
CNO-FM
MPP-B
FNO
scOT
CNO
Median relative L crror
0.100
0.010
2
128
512
32
128
512
Number of trajectories
Figure 16: GCE-RT. Number of trajectories vs. median relative L' error on the test set.
Poseidon-B FNO
Poseidon-L scOT
CNO-FM
MPP-B
CNO
1.00
Median relative L error
0.10
0.01
2
32
128
512
2048
8
32
128
512
2048
Number of trajectories
Figure 17: Wave-Layer. Number of trajectories median relative L error on the test set.
Poseidon-D B
Poseidon-L scOT
CNO-FM CNO
MPP-B
FNO
1.00
Median relative L' error
0.10
0.01
2
128
2048
512
2048
Number of trajectories
Figure 18: Wave-Gauss. Number of trajectories vs. median relative I error on the test set.
49

---

Poseidon-B
Poseidon-I.
CNO-FM
MPP-B
FNO
scOT
CNO
Median relative I. error
0.100
0.010
0.001
0.000
32
256
2018
32
256
2048
Number of trajectories
Figure 19; ACE. Number of trajectories vs. median relative L error on the test set.
Poscidon-B B
Poscidon-L scOT
CNO-FM
MPP-B
FNO
CNO
1.000
Median relative I
0.010
0.001
32
256
2048
32
256
2048
Number of trajectories
Figure 20: MaisoonGGause Number of samples median relative L error the test set.
Poseidon-B
Poseidon-L scO'T
CNO-FM CNO
MPP-B
FNO
1.000
Median relative Lt error
0.100
0.010
4
32
256
2048
32
256
2048
Number of trajectories
Figure 21: Helmholtz. Number samples vs. median relative I error on the test set.
50

---

D.2 Scaling with respect to Model Size
In Figure we plot how the training loss and evaluation (validation) loss during pretraining changes with model size for the POSEIDON models. We observe from this figure (bottom row) that there is consistent decay in losses with increasing model size. The role of model size vis vis downstream tasks has already been shown in the scaling plots of the previous subsection where we compared POSEIDON-L. with the smaller POSEIDON-B. The corresponding metries FG and AG are shown in Table [ We also see from the statistical summary Table 9 that there is gain, on average, in downstream performance with increasing model size for the POSEIDON family of models.
Poseidon-I.
Poseidon-B
Poseidon-T
10
10
10-1
10-1
Training/ /alidation Loss
10
20
15
Epoch
Epoch
0.20
0.20
0.18
0.18
0.16
0.16
0.14
0.14
0.12
0.12
0.10
0.10
0.08
0.08
0.06
0.04
21M
158M
629M
21M
158M
629M
Model Size
Model Size
Figure 22: (Top) Training (left) and evaluation (right) losses up to epoch 20 for different model sizes. (Bottom) Scaling at epoch 20 for training loss (left) and evaluation loss (right).
51

---

D.3 Scaling with respect to Pretraining Dataset Size and Quality
As mentioned in the Main Text, scaling of the POSEIDON models with dataset size of great interest. To that end, in Figure 23 we plot the training and evaluation losses, during pretraining, for the POSEIDON-B model, trained with one-eighth, one-half and full size of the pretraining dataset (the details of the corresponding setups are given in Section C). We see from this figure that POSFIDON-B scales with dataset size on the pretraining dataset. Moreover, in Figures 24 to 38 (left subfigure of each figure), we compare the performance of POSEIDON-B, B, trained on the full pretraining dataset with the same model trained on one-eighth of it.
Moreover, in Figures 24 to 38 (right subfigure of each figure), we compare the performance of POSEIDON-B, trained on the half the pretraining dataset with the same model pretrained on less diverse dataset (see Main Text and Section C for details of the setup).
full pretzaining data
half pretraining data
eighth pretraining data
10
10
Training/ Validation Loss
10-1
10
0
10
15
20
5
10
15
20
Epoch
Epoch
0.20
0.20
0.18
0.18
0.16
0.16
0.14
0.14
0.12
0.12
0.10
0.10
0.08
0.08
0.06
0.06
0.04
0.04
20000
10000
60000
80000
20000
10000
60000
80000
Number of trajectories
Number trajectorics
Figure 23: (Top) Training (left) and evaluation (right) losses up epoch 20 for different pretraining dataset sizes. (Bottom) Scaling epoch for training loss (left) and evaluation loss (right).
52

---

full dataset
eighth of dataset
half of dataset
less diverse dataset
0.100
Median relative L error
0.010
2
128
512
8
32
128
512
Number of trajectories
Figure 24: NS-PwC. Number of trajectories vs. median relative L error on the test set.
full dataset
... eighth of dataset
half of dataset
less diverse dataset
Median relative L, error
0.100
0.010
2
128
512
32
128
512
Number of trajectories
Figure 25: NS-SVS. Number of trajectories vs. median relative L error on the test set.
full dataset
eighth dataset.
half of dataser
less diverse dataset
Median relative Lt error
0.100
0.010
2
8
32
128
512
8
32
128
512
Number of trajectories
Figure 26: NS-BB. Number of trajectories median relative L error the test set.
53

---

full dataset
eighth of dataset
half of dataset
less diverse dataset
Median relative L error
0.100
0.010
2
128
512
32
128
512
Number of trajectories
Figure 27: NS-SL. Number of trajectories vs. median relative L' error on the test set.
full dataset
-..
eighth of dataset
half of dataset
less diverse dataset
1.00
Median relative L' error
0.10
128
32
512
Number of trajectories
Figure 28: NS-Tracer-PwC Number of trajectories vs. median relative L error on the test set.
full dataset
eighth dataset.
half of datases
less diverse dataset
1.00
Median relative L error
0.10
2
8
32
128
512
2
32
128
512
Number of trajectories
Figure 29: FNS-KF. Number of trajectories vs. median relative L error the test set.
54

---

full dataset
eighth of dataset
half of dataset
less diverse dataset
Median relative L error
0.100
2
128
512
8
32
128
512
Number of trajectories
Figure 30: CE-RPUL, Number of trajectories vs. median relative L' error on the test set.
full dataset
....... eighth of dataset
half of dataset
less diverse dataset
Median relative L' error
3 x 10
2 x 10-1
32
128
512
32
128
512
Number of trajectories
Figure 31: CE-RM. Number of trajectories vs. median relative L error on the test set.
fill dataset
eighth dataset.
half of dataser
less diverse dataset
Median relative L error
0.100
2
8
32
128
512
8
32
128
512
Number of trajectories
Figure 32: SE-AF. Number of samples median relative L error the test set.
55

---

full dataset
eighth of dataset
half of dataset
less diverse dataset
1.00
Median relative L error
0.10
128
512
32
512
Number of Urajectories
Figure 33: GCE-RT. Number of trajectories vs, median relative L error on the test set.
full dataset
-.e eighth of datasel.
half of clttaset
less diverse dataset
1.00
Modian relative L' error
0.10
128
32
512
Number of trajectories
Figure 34: Wave-Layer. Number trajectories vs. median relative L error on test set.
full dataset
eighth dataset.
half of dataser
less diverse dataset
1.00
Median relative L error
0.10
2
8
32
128
512
2
8
32
128
512
Number of trajectories
Figure 35: Wave-Gauss. Number of trajectories median relative L error the test set.
56

---

full dataset
eighth of dataset
half of dataset
less diverse dataset
1.00
Median relative !! error
0.10
0.01
0.00
2
8
32
128
512
32
128
512
Number of trajectories
Figure 36: ACE. Number of trajectories vs, median relative L error on the test set.
full dataset
-o eighth of dataset
half of dataset
less diverse dataset
1.00
Median relative L emor
0.10
128
32
512
Number of trajectories
Figure 37: Poisson-Gauss. Number of samples vs. median relative L error on the test set.
full dataset
eighth dataset.
half of dataser
less diverse dataset
Median relative 11 error
6 6 10
4.x 10
3 x 10
2 10
0.100
32
128
512
32
128
512
Number of trajectories
Figure 38: Helmholtz. Number of samples median relative L error on the test set.
57

---

D.4
Case Studies
Given the excellent performance of POSEIDON models across the board, including on tasks that involve PDEs (physical processes) not encountered during pretraining, it is important to understand what underpins this performance. To this end, we will present three case studies in order to explain POSEIDON's 's robust performance.
D.4.1 CE-RPUI
First we consider the CE-RPUI downstream task. Clearly, POSEIDON models perform very well on this task, as shown in Figure 13 as well as and Also, as seen from Figure [67 where we visualize single random sample for all the variables at time T = 0.7, POSEIDON-B is much more accurate, when finetuned on 128 trajectories, than CNO and FNO, which are trained from scratch with the same number of trajectories. Note that the underlying solution is very complex, with a mixture of shocks and roll-up vortices. While POSEIDON captures these shocks and vortices very sharply, CNO and (especially) FNO fail to do so. What explains this impressive performance of POSEIDON on this difficult downstream task?
We start by observing that, like all other downstream tasks, this task is out-of-distribution (o.o.d.) with respect to the pretraining dataset. Although the underlying PDE (compressible Euler equations (37) is present in the pretraining dataset, this data distribution has not been seen during pretraining. This o.o.d. nature of the task is clearly seen from Figure 39 where we plot how the same random sample (visualized in Figure (gguee670) is inferred with a POSEIDON-B B model zero-shor. We see from the figure (second column from the left) that the zero-shot results are rather poor. However, even with 1 task-specific example, we see from Figure (39 (third column from left) that at least, some large scale features (such shock locations) are approximated reasonably accurately. With just downstream trajectories, the quality of approximation improves dramatically and even the vortex roll-ups are captured accurately. The quality of approximation continues to improve with 32 and 128 downstream trajectories, as shown in the right-most columns of Figure 39 Thus, from this figure we conclude that a few task-specific samples suffice to accurately approximate the underlying solution operator. This is also evidenced in the scaling plot Figure 13
How does POSEIDON succeed in learning this complex solution operator with so few samples? We know that the pretraining dataset contains the CE-RP operator, where the initial condition (see Figure for sample) has similar four- four-quadrant Riemann problem structure as the intial conditions in the CE-RPUI benchmark, the main difference being that the interfaces, across which the initial data is discontinuous, are now perturbed sinusoidally, instead being axis-aligned. However, is precisely these perturbations that are responsible for the roll-up of small-scale vortices that are absent in the CE-RP operator. Thus, the model potentially needs to learn how to approximate small-scale vortices accurately from some other operator in the pretraining dataset, while learning how to propagate large-scale shocks from the CE-RP operator.
One would think that the CE-KH operator in the pretraining dataset provides the information about vortex roll-ups, see Figure 59 for visualizing a sample. However, the underlying vortices are much larger. So, where does this missing information come from? One possible source could be the CE- CRP operator (see Figure 58) where vortices of many different scales are being formed. Perhaps, the model leverages shock propagation from CE-RP, large vortex roll-ups from CE-KH and small-scale vortex dynamics, as well as curved shock propagation, from CE-CRP in order to provide very good approximation with a few training examples on CE-RPUL. A partial test of this hypothesis is to check the model, pretrained with the less-diverse dataset that excludes CE-CRP performs worse than the model pretrained with the full dataset. This is already shown in Figure 30 (Right) where the performance of the model, pretrained with the less-diverse dataset is worse than the model trained with the similarly sized but fully diverse dataset. This behavior is further reinforced from Figure 40) where the approximation of the same sample, considered in Figure 39 with these ablated models is shown. As predicted, the model pretrained on the less diverse dataset is clearly less accurate at resolving small-scale vortices than the competing one trained on the more-diverse dataset. It misses the input from the CE-CRP operator regarding small-scale vortex dynamics.
This qualitative analysis illustrates how the POSEIDON model leverages different operators from its pretraining dataset to amortize different aspects in order to construct accurate approximations during
58

---

Figure 39: How POSEIDON-B approximates random sample for the CE-RPUI task when trained with different numbers of task-specific trajectories.
finetuning with few task-specific examples and throws some light into how a foundation model for PDBs can learn effective representations from its pretraining phase.
D.4.2 ACE
Next, consider the ACE downstream task, where the underlying PDE is the nonlinear parabolic Allen-Cahn equation (67). which is clearly not included in the pretraining dataset for POSEIDON. More importantly, the type of physics that the Allen-Cahn Equation models is that of reaction- diffusion. On the other hand, the PDEs included in the pretraining dataset, Compressible Euler and Incompressible Navier-Stokes very high Reynolds number, are convection-dominatee one does not expect that the pretrained model has learned effective representations about reaction- diffusion. Yet, we see from Figure 19 that POSEIDON very effective at learning this solution operator from few training examples. This point is also reinforced from Figure 772, where we show how single randomly chosen sample is well-approximated by POSEIDON. How does POSEIDON learn these new physics?
Hence,
To understand the factors behind POSEIDON'S performance, we plot how the same random sample, visualized in Figure 72] is approximated by POSEIDON-B, when fine-tuned with different number of task-specific examples, ranging from 1 128, in Figure 41 We observe from this figure that already with just 1 task-specific trajectory, POSEIDON is able to learn the large-scale features of the solution of Allen-Cahn approximately. particular, it has learnt both front propagation (potentially from all propagating shock waves seen during pretraining) as well as diffusion (spreading) of localized features. With more downstream trajectories, it is able adjust local features quite well further approximate the diffuse fronts. This case study shows how POSEIDON can learn new features from few task-specific training examples.
Given these encouraging results on the ability of POSEIDON to generalize to the unseen physics underlying the Allen-Cahn equation, we investigate this ability further by freezing the latent space of POSEIDON during finetuning by setting 0, = 0., for all r. the gradient descent procedure (TO) for finetuning. Thus, only the embedding and recovery parameters are learned and the rest frozen. This results an extremely lighteiight model for training as less than 0.5% of the total parameters in POSEIDON are being retrained. Nevertheless, as shown in Figure 42 this very parsimonious
59

---

Groand Touth
half
lss divense dataset
Figure 40: sample of CE-RPUI when POSEIDON-B is pretrained on half of the pretraining dataset vs. less diverse pretraining dataset.
Figure 41: How POSEIDON-B approximates random sample for the ACE task when trained with different numbers task-specific trajectories.
60

---

Ground Toulk
Injeetsey
tniestedes
16 Enjontatics
Iniistodis
Figure 42: How POSEIDON-B with Frozen Latent Representation approximates the random sample as in Figure 41 for the ACE task when trained with different numbers of task-specific trajectories.
form of POSEIDON has already learned the solution of the Allen-Cahn equation qualitatively with only one training trajectory, although there is a quantitative mismatch. This mismatch is corrected when further samples are shown to the model. this model (0.031) is actually lower than FNO with 128 trajectories (0.037) although it is higher than the corresponding model where all the parameters of POSEIDON-B are finetuned (0.014). This experiment demonstrates that the latent representations learned from the equations of fluid dynamics during pretraining are very rich and can unexpectedly contain information about reaction-diffie -iiffusion equations, which are then leveraged by the frozen-laten latent model to learn the underlying solution operator.
particular, with 32 trajectories, the error with
D.4.3 Poisson-Gauss
In our final case study, we consider the Poisson-Gaus task. The underlying PDE is the Poisson equation (68) and the solution operator maps the coefficient, which is superposition of Gaussians, into the solution. A visualization of the solution operator for a single random sample shown in Figure and shows how the source is diffused and smoothed out.
We remark that this task is very different from the pretraining dataset in various ways. First, the underlying PDE is time-independent in contrast to the two time-dependent PDEs seen during pretraining. Second, the underlying physics of diffusionn of features and their smoothing out patently different from the physics seen in the pretraining dataset. Finally, the Dirichlet boundary conditions considered here are also different from the periodic boundary conditions of the pretraining dataset. Nevertheless, we see from the scaling plot Figure 20 and Table I and [ that POSEIDON models perform very well in this case. This also observed from Figure 74] where we observe that POSEIDON-B learns this particular random sample far better than CNO and FNO, with the same number (51 training samples. To understand the reasons behind POSEIDON'S performance, in Figure 43 again plot how this foundation model approximates this particular random sample, when trained with an increasing number of task-specific samples. We see from this figure that for 1 sample, the approximation is very poor, indicating how much out-of-distribution this task is, with reference to the pretraining dataset. In fact, the model simply learns to approximate the input. However, within few samples (16), it has learnt that the input needs to be spread (diffused) out. It takes about 128 samples for the model to realize that the input needs to be both spread out as well as smoothened and by 512 samples, the local adjustments needed to further smoothen the output have been made.
A few remarks are in order to explain this qualitative picture. First, POSEIDON could have used the first few samples in training to forget the information from the pretraining phase. Yet, it does not do that and uses the very first sample to already just output the identity operator. Then, there appears to be a warmup phase where the model slowly learns the underlying physics, for instance diffusion and smoothening and then fast learning phase where the operator can be better approximated. This qualitative picture is also consistent with the observed biphasic power scaling, see the subsection on scaling laws in section D.I and the fact that there is phase transition between the warmup and fast learning phases in the power law (22). This case study sheds further light into how POSEIDON can learn unseen physics from few task-specific training examples.
As with the Allen-Cahn equations of the previous section, we further study the factors underpinning the ability of POSEIDON to generalize to this PDE by freezing the latent space of POSEIDON during
61

---

Goond Trubh
stmple
sunckes
sumples
aunoles
0.0
-03
04
Figure 43: How POSEIDON-B approximates random sample for the Poisson-Gauss task when trained with different numbers of task-specific samples.
comples
04
405
Figure 44: How POSEIDON-B, with a Frozen Latent Representation, approximates a random sample for the Poisson-Gaus task when trained with different numbers of task-specific samples.
finetuning by setting 8, = 0r, for all т, in the gradient descent procedure 10 for finetuning. Thus, only the embedding and recovery parameters are learned and the rest frozen. shown in Figure [44] even this frozen-latent form of POSFIDON has already learned the basic features of the underlying solution operator, i.e., Diffusion and Smoothing, qualitatively with only a few training samples, although there is quantitative mismatch. This mismatch is corrected when further samples are shown to the model. In particular, with 512 trajectories, the error with this model (0.11) is significantly lower than FNO (0. (0.282) although it is higher than the corresponding model where all the parameters of POSEIDON-B are finetuned (0.022). This experiment further demonstrates that latent representations learned from the equations of fluid dynamics during pretraining rich enough to even contain information about the priori unrelated physics steady state diffusion, which are then leveraged by the frozen-laten model to learn the underlying solution operator.
62

---

D.5 Results with DPOT
Table 10: Efficiency gain EG 11 with S = 10224 for time-dependent and S 4096 for time-
independent PDES) and Accuracy Gain (AG) ((TT) with S = 128 for time-dependent and S 512 for time-independent PDEs) for DPOT and tested downstream tasks.
Finetuned DPOT
DPOT from Scratch
M
L
M
AG
EG
AG
EG
AG
NS-PwC
44.8
12.5
39.7
12.0
17.0
6.1
23.3
10.2
NS-SL
4.5
2.4
4.7
2.4
2.1
1.3
3.0
1.6
FNS-KF
0.0
1.0
0.0
0.9
0.0
0.8
0.0
0.8
CE-RPUI
53.5
3.7
53.6
3.6
26.1
2.5
31.2
2.9
SE-AF
3.5
1.2
4.7
1.4
4.4
1.3
5.1
1.4
Wave-Layer
23.5
5.5
28.9
6.0
14.1
3.6
17.8
4.2
Wave-Gauss
25.2
4.4
27.8
4.5
18.0
3.3
20.5
3.6
The DPOT foundation model [19] has been trained operators for the compressible and incom- pressible Navier-Stokes equations,
equations and Shallow- Water equations. The model has been setup to take sequence of time steps for a time-dependent PDE and output the next time step. However, we can modify it for finetuning for our OLT operator learning task by following exactly the same procedure for finetuning the MPP foundation model. For steady state problems, an identical procedure as with MPP is used. This allows us to perform a fair comparison between DPOT and the POSEIDON models proposed here.
To this end, we consider DPOT-M (with 120 M parameters) and DPOT-L (with 509 M parameters) which are comparable in size to the POSEIDON-D and POSEIDON-L models, respectively. Given compute constraints, we focus this comparison on a representative subset of 7 downstream tasks which are listed in Table 10 Moreover, FNO architecture, is also employed for each task to evaluate DPOT's model performance. For both finetuning and training models from scratch, we employed the Adam optimizer (25] with weight decay of 10 6 and a cycle learning rate policy. For finetuning DPOT models, the maximum learning rate set to 10-4 and training was conducted for 100 epochs. When training models from scratch, we used a maximum learning rate of 10-3 and trained for 200 epochs.
DPOT model, with the Adaptive
The resulting EG and AG scores presented in Table 10 These scores should be compared with the corresponding EG and AG scores of POSEIDON-L and scOT from Table and POSEIDON-B from Table 8 Comparing these results, we make the following observations,
For all these tasks except SE-AF, POSEIDON is significantly better, both in terms of efficiency and accuracy gains, the corresponding DPOT model. Even for SE-AF, the models are very comparable. The superiority in performance of POSEIDON is seen very clearly when we consider the mean AG scores over these downstream tasks which amount POSEIDON-L (8.14), POSEIDON-B (6.5), DPOT-M (4.39) and DPOT-L (4.4). Hence, the POSEIDON-L. model is almost /wice more accurate than both the DPOT models considered here. In fact, the DPOT models' performance lies in-between CNO-FM with an average AG score of 2.66 and the POSEIDON models. Similar results also hold for the efficiency gain score.
Surprisingly, DPOT foundation models do not seem to scale with model size, at least on this set of 7 representative downstream tasks seen from the mean AG scores of 4.4 for both the DPOT-M and DPOT-L models where an increase of the number of parameters by factor of 5 does not lead any noticeable increase in model performance on downstream tasks.
As surprisingly, the stand-alone DPOT neural operators performed well on this dataset. For instance, the average AG score of mained-from-scratte DPOT-L 3.53, which only 25% lower than the DPOT-L foundation model. On the other hand, POSEIDON-L is almost 5
63

---

times more accurate than the underlying scOT neural operator. These results indicate that DPOT foundation models do not harness latent representations as well as POSEIDON does and they rely on the capacity of the underlying neural operator to learn downstream tasks.
Taken together, our results indicate that (our modification of) DPOT performs better than the CNO- FM and MPP foundation models but is significantly inferior to the POSEIDON models. Moreover, the lack of scaling with model size for DPOT on downstream tasks and questions over how it uses latent representations further point to the advantages of POSEIDON over this competing model. Nevertheless, this comparison merits further study.
64

---

D.6 Further Ablations and Results
D.6.1 On all2all training
The all2all training strategy, described in the Main Text, aims to leverage the semi-group structure of the solution operator of the time-dependent PDE (I) to scale-up the training data per trajectory. use every possible pair of snapshots, per trajectory, in the learning
As shown in Figure 2 (d), process leading to the loss function (6) It is instructive to compare this strategy with the vanilla training strategy based on the loss function (5) As this strategy is applicable for any (time-dependent) operator learning algorithm, we study it for the CNO model [60] here. To this end, we consider the NS-SL task and compare the all2all and vanilla strategies and plot the results in Figure 44 to observe that the all2all training strategy significantly outperforms the vanilla training strategy for this task.
However, there is a caveat with the all2all strategy. It lies in the computational cost of training as the number of training pairs grows quadratically with the number of available time snapshots at which the trajectory is sampled. One option to reduce this cost is to select subser of snapshots from within all available snapshots per trajectory and apply all2all training to this subset, bringing down the computational cost proportionately by the relative reduction in the cardinality of the selected subset. Yet, there is the possibility that by sampling too few snapshots, the overall error will increase.
To investigate this trade-off, we consider the NS-PwC task and the CNO model. The data for this task available in the time-interval [0.0.7], sampled 14 time snapshots (excluding the initial time 0). Denoting the ith-snapshot by 4 with i 0, 1, ..... 14, We select the following subsets of time snapshots,
Thi Snapshots at to = 0 and f14 = 0.7. The training only considers learning the map between initial datum and solution at final time 414. Samples corresponding to identity function are also included.
Ty: Snapshots at lor 414
Ty: Snapshots at to, t2, t4, to, ts, t10, t12,
Tí: Snapshots at (j. for all 0 S S 14
For each of the above subsets of time snapshots, all2all training is used leading to 3, 6, 36 and 120 training pairs per trajectory for T14, 'T, T2, TT, respectively.
In Figure 46 we plot the test error number of trajectories. From the left panel of this figure, see that there is consistent gain in accuracy as a more dense sampling of the snapshots is performed. The models are monotonically more accurate as we go from T14 through Tr to Tz. However, observe from Figure 46 that going beyond To to Ti does not yield any further decrease in test error as the difference between the newly added snapshots and the existing ones in Tz is not statistically significant enough to aid the training process. Moreover, by choosing Tz over Ti, we reduce the computational cost of training by factor of 3.3. These considerations motivate to not too dense sampling strategy for pretraining (and finetuning) our foundation models.
also
D.6.2 Direct. vs. Autoregressive Inference
As mentioned in the Main Text, our time-conditioned models can either be directly evaluated at the time of interest, or an autoregressive rollout can be performed (see Equation g of the Main Text). This per be of any form that the user wants, i.c. with homogeneous step-sizes in time, or with heterogeneous step-sizes in time. For simplicity, we only consider homogeneous autoregressive rollouts for POSEIDON, scOT and FNO models, for the CNO models we find slight performance boost with heterogeneous rollout strategy.
Figure 47 shows for the NS-PwC and the Wave-Layer downstream task how the error behaves when using direct or (homogeneous) autoregressive rollouts. We can directly see that it depends very much on the task at hand, as autoregressive rollout works better for the NS-PwC task, whereas direct lead-time input works better for Wave-Layer; this seems to be very dataset- and dynamics-dependent. We therefore choose the best strategy for each task which is listed in Table
65

---

vanilla
allaal
10
Medion relative L' error
6 11
10-2
x 10""
128
256
512
1024
2048
Number of trajectories
Figure 45: NS-SL. Testing errors of the CNO models trained in an all2all and vanilla manner. Performance improves with all2all training.
Ti AR
Median relative Lt error
Ti direct
T2 AR
Ta direct
0.100
Tia direct
direct
T direct
128
256
512
1024
128
256
512
1024
Number of trajectories
Figure 46: NS-PwC. Testing errors of the CNO models trained in an all2all manner on different T. trajectories. (Left) Errors from directly evaluating the trained models. Performance improves as denser trajectories are incorporated. (Right) Saturation effect observed. Adding denser trajectories no longer enhances performance, as the additional samples are statistically less significant.
D.6.3 Error Growth over Time for POSEIDON-B
Autoregressive inference can only work better than direct lead-time input when the error that ac- cumulates at every step smaller than the error obtained by direct lead time input. In Figure 48 we can directly see that the error scales sub-linearly for the NS-PwC experiment and this is true in general for our downstream tasks. This leads to two observations. First, there is no blow-up (for instance exponential growth) of error in time with these models. Second, the fact that the error grows in time proves that it is harder to predict the solution at final time from initial data than predicting time-averaged quantities. In other words, the 1 in time will be greater than the -e-error. This justifies our choice of evaluating different models at the final lead time the underlying task.
To further demonstrate how POSEIDON compares with FNO over time, we plot errors for the NS-PwC and NS-SI, experiments as function of time with both models Figure 49 We observe from this figure that the difference in error between FNO and POSEIDON-B actually grows over time and is the highest the final time FNO has much larger rate of error growth over time than POSEIDON, justifying our decision evaluate models with respect to error at the final time.
66

---

Poseidon-B Autoregressive
FNO Autoregressive
Poseidon-B Direct
FNO
Direct
1.00
Median relative L error
0.10
x
0.01
2
32
128
512
2048
8
32
128
512
2048
Number of traeectorie
Figure 47: Homogenenous autoregressive rollout direct lead-time input NS-PwC (left) and Wave-Layer (right).
0.016
Median relative D! error
0.014
0.012
0.010
0.008
0.006
0.4
0.5
0.6
0.7
Time
Figure 48: Error accumulation for autoregressive rollout of POSEIDON-B finetuned on 128 trajectories of the NS-PwC dataset.
D.6.4 Out-of-distribation Time Extrapolation
Here, we consider the NS-SL downstream task. As mentioned before, FNO (and other neural operators) were trained from scratch as well as POSEIDON (and other foundation models) were finetuned to learn the solution up final lead time of T 0.7. We Poseidon foundation model and the neural operator baseline (relatively) perform when we consider an out-of-distribution time extrapolation the downstream task level. To this end, Figure 50 we plot the test еттоrs, with respect to increasing number of task-specific trajectories, for both FNO and POSEIDON-BI but evaluated final times of T 0.7 and the extrapolated final time of T 1.0. A homogeneous autoregressive rollout is used in all cases. We observe from this figure that both POSEIDON-B and FNO are worse at extrapolating in time than they are predicting within the time-period that they have been trained on. In addition to significantly outperforming FNO at both = 1.0, POSEIDON-B in fact performs relatively better
to investigate how the
time T 0.7 and at the extrapolated time at out-of out-of-distributtin than FNO. It is best seen from the EG metric [I]. where POSEIDON'S 's EG for time T 0.7 improved EG 30 time T 1.0. This gain can be attributed the fact that during pretraining, POSEIDON models have been trained for longer time horizon.
67

---

Poecido
FNO
Posiido
FNO
0.00
Mediia relative J1o00
0.20
Mediaa pelacive E eEror
0.00
0.20
000
0.15
0.04
0.00
0..00
0.00
0.00
0000
0.00
0.2
0.4
0.0
0.7
6.1
012
0.3
0.4
0.0
Time
Time
Figure 49: Error accumulation for the finetuned POSHIDON-B and FNO for 128 training trajectories on NS-PwC (left) and NS-SL (right).
Poseidon-B T 1.0
Poseidon-B T 0.7 FNO T = 0.7
FNO T 1.0
Median relative L' error
0.100
0.010
128
512
2048
Number of trajectories
Figure 50: Out-of-distributiin extrapolation in time for POSEIDON-B and FNO on NS-SL up to ==
D.6.5 Generalization of POSEIDON with respect to Changing PDE Parameters
Several of our downstream tasks such GCE-RT, Wave-Layer, Wave-Gaus and Helmholtz involve operators that map the coefficient in the PDE to its solution. This setup is very different from the pretraining dataset where the underlying solution operators only map the initial data solutions at later times and there is no PDE coefficient that is encountered. Nevertheless, from Tables [ and 88, observe that the POSEIDON models generalize very well these very different setups for the operators for downstream tasks. To further test the ability of POSEIDON to generalize for different PDE parameters, we consider the Navier-Stokes Equations (BI)) with a viscosity coefficient 4 x 10 The ground truth data is generated using the Azeban spectral hyper viscosity solver very different from the setup of the pretraining data and downstream tasks considered so far as in all of them, only hyperviscosity of 4 x 10 was applied to high-enough Fourier modes order model the incompressible Euler equations with zero viscosity. In this new task, the initial conditions are identical to the NS-PwC downstream task. We see from Figure 51 that Poseidon-B generalizes very well to this new viscosity coefficient and outperforms FNO readily, in terms of both sample efficiency and accuracy. In particular, the AG and EG scores of Poseidon-B are FG 925.5 and AG = 47.5, which are completely comparable to (even better than) the scores of EG 1024 and AG 19.7 (see Table for the original NS-PwC task). Taken
622 This new viscosity coefficient
68

---

Poseidon-B
FNO
1.000
Median relative L error
0.100
0.010
0.001
2
32
128
512
2048
Number of trajectories
Figure 51: Error for the NS-PwC downstream task, but with viscosity y = 4 10 10-3 (on all modes) instead of 10-4 applied only on high-enough Fourier modes simulate the inviscid limit
together with other downstream tasks involving different PDE coefficients, this experiment clearly demonstrates the ability of POSEIDON to generalize to different PDE parameters via finetuning.
D.6.6 POSEIDON Evaluated on Different Grids
As POSEIDON is based operator transformer (scOT), it can be evaluted grid resolutions, different from the underlying computational grid. Following B] we can simply downsample (upsample) the input function from the given grid to the computational grid, process the input with POSEIDON and upsample (downsample) the output from the computational grid to the given grid resolution. We perform this evaluation of POSEIDON-B task and present the result in Figure 52 to observe that the test error is (approximately) invariant to the grid resolution.
multiple grid resolutions for the NS-PwC
D.6.7 Robustness of Poseidon with respect to Noise
To study how robust POSEIDON is to noise, we consider the downstream CE-RPUI task and at inference time, we add Gaussian noise to the inputs (initial conditions) at different noise-to-signal ratios (NSRs) of 0.1%, 1% and 3% respectively. The resulting errors, computed with respect to Ground Truth where the outputs are noisy, for varying numbers of training trajectories. are shown in Figure 53 The errors in the zero noise (clean) case are also shown in this Figure. We observe from this figure that POSEIDON-L'S performance robust input noise and the error does not grow significantly even when the noise level is an appreciable 3%, demonstrating the robustness of this foundation model with respect to noise.
D.6.8 Histograms of Errors for Different Tasks
In Figure 54 we plot the distribution of emoor across the test set for all downstream tasks with the POSFIDON-B model, finetuned with 128 trajectories (samples).
69

---

0.020
training resolution
0.018
Median relative L error
0.016
0.014
0.012
0.010
60
70
90
100
110
120
130
Resolution
Figure 52: Test performance of POSEIDON-B finetuned on 128 trajectories of the NS-PwC dataset for multiple resolutions.
clean
0.1%
1%
3%
1.000
Median relative L error
0.100
0.010
32
128
Number of trajectories
Figure 53: Effect of injecting Gaussian noise in the initial condition CE-RPUI (before normalizing the data; normalization constants are as before) with POSEIDON-L.
70

---

NS-PwC
NS-SVS
150
30
40
100
20
20
10
0.02
0.08
0.000
0.050
0.100
NS-SL
NS- S5--Trccu-Pooe
FNS-KF
1010
30
80
30
200
60
40
10
20
0
0.01
0.02
0.03
0.01
0.05
0.10
0.15
0.05
0.10
0.15
REUT
CM-RM
SE-AF
10
12.5
10.0
Density
7.5
5.0
2-
2.5
0.0
0.1
0.2
0.3
0.4
0.176
0.200
0.225
0.250
0.275
0.0
0.1
0.2
GCE-RT
Wavo-Layer
Wave-Gauss
8
15
6
10
4
2
7- J. www............. Po
0,
0.10
0.20
0.4
0.2
0.1
ACE
Poisson Poisson-Gauss
Hehmholtz
30
125
100
20
76
50
10
2-
26
0.00
0.00
0.02
0.03
0.00
0.05
0.19
0.15
0.0
0.5
1.0
Relative
error (mesn over all QOLs)
Figure 54: Error distribution of POSEIDON-B finetuned on all downstream tasks (for 128 trajectories in the time- dependenendntt, and 512 in the time-inde iime-independent case). The kernel density estimate is done over the mean of functions/quantities interest.
71

---

E
Computational Resources
All experiments were run on different types of GPUs, on the Euler cluster of ETH Zurich. Depending on the experiment, we use between 8 and 128 CPU cores and up to 512GB of RAM, with pretrainings using the most CPU cores and RAM. However, we note that this more than is actually needed, as we tried to minimize being bottlenecked by dataloading. For all our models and baselines, we used consumer-grade GPUs with 24GB of VRAM. All our pretrainings were performed in (data-)parallel on 8 NVIDIA GeForce RTX 4090 GPUs. All finetuning experiments and most scratch trainings were performed on a single GPU, while some scratch training runs with lot of data were performed in (data-)parallel, Pretraining times be read off from Table I1
Table 11: Approximate pretraining times on 8 NVIDIA GeForce RTX 4090 GPUs. Batch sizes are given in parentheses.
POSEIDON-L
POSEIDON-B
POSEIDON-T
CNO-FM
165h (16)
118h (40)
22h (80)
178h (32)
we provide an overview over the inference times of each model for single call to it. We observe from this table that even the biggest POSEIDON-L. has an (average) inference time of less than 10 secs. This is contrast to the PDE solvers that were used to generate the data in this paper. Their run times, for a resolution 128 ranged from anywhere between 0.1 sec (for highly optimized GPU solver [62] for the NS datasets to 10 secs for FENICS [40] FEM solver for the Foisson-Gaus dataset to approx 100 secs for NEWTUN (55 for generating the airfoils datasets to 500 secs for the well-balanced scheme to generate the GCE-RT. Thus, observe gain in inference time from anywhere between 1 orders of magnitude.
In Table 12
Table 12: Approximate inference times (per call and normalized to single sample) for different models, all reported on a NVIDIA GeForce RTX 4090 GPU for the FNS-KF experiment. Batch sizes are given in parentheses. We note that the values given here are just proxies as this was not tested in controlled environment.
Model Approximate inference time
POSEIDON-L
4 ms (16)
POSEIDON-B
2.9ms (40)
POSEIDON-T
1.6ms (40)
CNO-FM
1.8ms (32)
MPP-B
10ms (4)
CNO 0.9ms (32)
ScOT
3ms (40)
FNO
2ms (40)
F Pretrained Models, Datasets, and Source Code
The source code corresponding to this work is available on Github httpss//gthub..oo//ama// ethz/poseidon), Everything is tightly integrated into Huggingface Transformers 773 and we make heavy use of Huggingface Accelerate for distributed training.
In addition to the code, we make (pretrained) models and datasets available on the Huggingface Hub httip::/hugginnfacc.oommam htt.s/accc.occamaaabbeeez) see the Poseidon collection for pretrained models and pretraining datasets, the Poseidon- Downstream Tasks collection for all downstream tasks, or the PDEGYM collection for all datasets PDEGYM.
72

---

G Visualizations
-2
0
2
(a) Inputs: horizontal velocity and vertical velocity v.
Groun Truth
Puseiton--B
(b) (Top) Ground truth. (Bottom) Samples predicted by POSEIDON B T Figure 55: NS-Sines. Visualization of random sample.
73

---

-2
(a) Inputs: horizontal velacity and vertical velacity
Ground Truth
Poscidonn-B B
2
(b) (Top) Ground truth. (Bottom) Samples predicted by POSEIDON-B B at 7
Figure 56: NS-Gauss. Visualization of a random sample.
74

---

0.0
0.4
00.5
0.0
1.0
-0.2
(a) Inputs: density p. horizontal velocity u, vertical velocity e and pressure p.
Ground
Pocilee
(b) (Top) Ground truth. (Bottom) Samples predicted by POSEIDON B at 7 1. Figure 57: CE-RP. Visualization of random sample.
75

---

0.5
(a) Inputs: density p. horizontal velocity u, vertical velocity and pressure p.
Groun
Poodium
(b) (Top) Ground truth. (Bottom) Samples predicted by POSHIDON-B 7 =1 1.
Figure 58: CE-CRP. Visualization of random sample.
76

---

-0.5
0.5
-0.2
0.2
-0.10
-0.05
0.05
0.10
6.0
7.0
x10
(a) Inputs: density horizontal velocity u, vertical velocity e and pressure p.
Poodium
(b) (Top) Ground truth. (Bottom) Samples predieted by POSHIDON-B 7
Figure 59: CE-KH. Visualization of a random sample.
77

---

Imputs: density P. horizontal velocity u, vertical velocity and pressure p.
Creund
Poseloo-
00
03
2.0
(b) (Top) Ground truth. (Bottom) Samples predicted by POSEIDON-B T
Figure 60: CE-Gauss. Visualization of random sample.
78

---

(a) Inputs: horizontal velocity and vertical velocity
Grone Thrrk
Prestde-B
CNNO
PSO
(b) (Top) Ground truth. (From second row onwards) Samples predicted by the finetuned POSEIDON-B. B, CNO. and FNO at 0.7.
Figure 61: NS-PwC. Visualization of a random sample.
79

---

0.5
0.0
0.5
1.0
0.5
0.0
0.5
(a) Inputs: horizontal velocity and vertical velocity
Goone Thrnks
ESO
(b) (Top) Ground truth. (From second row onwards) Samples predicted by the finetuned POSEIDON-B. B, CNO. and FNO at 0.7.
Figure 62: NS-BB. Visualization of a random sample.
80

---

(a) Inputs: horizontal velocity and vertical velocity
Gooned Turlk
Prestde-B D
CSVO
(b) (Top) Ground truth. (From second row onwards) Samples predicted by the finetuned POSEIDON-B. CNO. and FNO at 0.7.
Figure 63: NS-SL. Visualization of a random sample.
81

---

0.5
0.0
0.5
0.50
0.25
0.00
0.25
0.50
(a) Inputs: horizontal velocity and vertical velocity
Gooned Turlk
Prostde-B
CINO
(b) (Top) Ground truth. (From second row onwards) Samples predicted by the finetuned POSEIDON-B. B, CNO. and FNO at 0.7.
Figure 64: NS-SVS. Visualization of a random sample.
82

---

Tracer
Inputs: horizontal velocity u. vertical velocity and tracer concentration
Towoe
Ground Teub
Poecitonn 25
CNO
INO
(b) (Top) Ground truth. (From second onwards) Samples predicted by the finetuned POSEIDON-B, CNO. and FNO at 0.7.
Figure 65: C.... Visualization of random sample.
83

---

(a) Inputs: horizontal velocity u, vertical velocity a and forcing term f.
Creund Iruth.
CNO
FNO
1.0
05
00
(b) (Top) Ground truth. (From second row onwards) Samples predicted by the finetuned POSEIDON-B -B, CNO, and FNO at 0.7.
Figure 66: FNS-KF, Visualization of random sample.
84

---

-2.0
-0.50
-025
0.00
0.50
1.50-006
-0.1
0.0
-5.0
-15
-1.
(a) Inputs: density p. horizontal velocity a, vertical velocity e and pressure p.
Gronn
CND
FXO
(b) (Top) Ground truth. (From second row onwards) Samples predicted by the finetuned POSEIDON-B, B, CNO, and FNO at T = 0.7.
Figure 67: CE-RPUL. Visualization of random sample.
85

---

0.0
1.0
3.0
113
x10
x10
(a) Inputs: density horizontal velocity vertical velocity and pressure p.
Grand Trush
Powestrnh
FNO
(b) (Top) Ground truth. (From second row onwards) Samples predicted by the finetuned POSEIDON-B B, CNO. and FNO at 1.4.
Figure 68: CE-RM. Visualization of random sample.
86

---

0.50
X30
(a) Inputs: density horizontal velocity u, vertical velocity v, pressure p, and gravitational potential g.
Greal
(b) (Top) Ground truth. (From second row onwards) Samples predicted by the finetuned POSEIDON- B, CNO, and FNO at the seventh time step.
Figure 69: CE-RM. Visualization random sample.
87

---

0
2
0
(a) Inputs: displacement u and propagation field c.
Goond Turn
Prestbe-B B
CSV
(b) (Top) Ground truth. (From second row onwards) Samples predicted by the finetuned POSEIDON-B. BB CNO, and FNO at the 14-th time step.
Figure 70: Wave-Gauss. Visualization of random sample.
88

---

1.0
0.5
(a) Inputs: displacement a and propagation field c.
Cacond
Proot-oon
CNN
-15
-05
(b) (Top) Ground truth. (From second row onwards) Samples predicted by the finetuned POSEIDON-B, CNO. and FNO at T' = 0.7.
Figure 71: Wave-Layer. Visualization a random sample.
89

---

(a) Inputs: concentration u.
Grouni Trulh
Pomiraan-B
0.
0.0
1.0
(b) (Top) Ground truth. (From second row onwards) Samples predicted by the finetuned POSEIDON-B BB, CNO. and FNO at the 14-th time step.
Figure 72: ACE. Visualization of random sample.
90

---

Shape
0.0
0.2
0.4
0.6
0.
(a) Inputs: airfoil shape function
Groual Trubh
Powiide-
FNO
(b) (Top) Ground truth. (From second row onwards) Samples (density p) predicted by the finetuned POSEIDON-B, CNO, and FNO.
Figure 73: SE-AF Visualization of random sample.
91

---

0.00
0.05
0.10
0.15
0.20
(a) Inputs: source term f.
Ground Trolh
Powiidhu-
(b) (Top) Ground truth. (From second onwards) Samples (solution u) predicted by the finetuned POSEIDON-
B, CNO, FNO.
Figure 74: Poisson-Gaus. Visualization of a random sample.
92

---

0.0
022
0.4
0.6
08
(a) Inputs: propagation speed f.
Grounll Trull
Powiidoo-B
FNO
(b) (Top) Ground truth. (From second row onwards) Samples (solution u) predicted by the finetuned POSEIDON-
B. CNO, and FNO.
Figure 75: Helmholtz. Visualization of a random sample.
93

---

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: In the paper and SM, we thoroughly justified all our claims and contributions. We detailed all models, datasets and training, as well testing strategies used. Additionally, we supported our assertions with numerous experiments conducted throughout the study.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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
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 be expected to generalize other settings.
is fine include aspirational goals as motivation long 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: We dedicated an entire section to discussing the limitations and clearly outlined the next steps address them.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "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 the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or 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 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
94

---

Answer: [NA]
Justification: The purpose of this paper is to introduce a foundation model for learning the solution operators of PDBs, featuring novel benchmarks, training techniques, and new paradigm for foundation models for PDEs. Theoretical results will be addressed future work.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's 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 short proof sketch 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 ex- perimental results of the paper to the extent that affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We will make the weights our models, along with the codes and benchmarks, open source. We have clearly explained how the models were trained and finetuned, as well as which datasets were used. We also explained clearly how the datasets were generated. By following our instructions, end users will find it fairly easy reproduce our results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary 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 accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility.
95

---

In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: As explained in the previous answer, will make the weights of our models, along with the codes and benchmarks, open source. We explained clearly how the datasets were generated. We will not release the datasets and pretrained checkpoints during the review process as the files are very large and we are not aware of a file hosting service that enables anonymous release of this size. Code is available the reviewer's disposal.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.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 NeurlPS code and data submission guidelines (https
//iips.cc/ppll wttpcccppppppppeeeeeeeeeeeeeeeee
sonooooooc for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments 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) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We dedicated many sections in the main paper and the SM to explain the training and test details, along with all the relevant information necessary to understand our results.
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
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
96

---

Justification: In Section D.6.8 we presented histograms of emrrrs for various tasks, All the relevant statistical information can be inferred from these histograms. All other plots would get too cluttered by adding error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, 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.
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 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 out of range (e.g. negative eτror rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provided estimation about the training times and resources that we used in all the experiments.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate 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 the paper (e.g., preliminary or failed experiments that didn't '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
Answer: [Yes]
Justification:
Our research respects 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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- cration due to laws or regulations their jurisdiction).
10. Broader Impacts
97

---

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our paper is exclusively dedicated to advancing academic research in the area of Partial Differential Equations. Our models are tailored for utilization by researchers with an interest in this domain.
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 societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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 partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate 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 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 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 system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
To the best of our knowledge, our paper does not present any identifiable safety risks.
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do require this, but we encourage authors to take this into account and make 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]
98

---

The segments of code not authored by us, such as the CNO filtering, are explicitly acknowledged within the codebase, and due credit accorded the respective owners of these assets.
Justification:
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 dataset.
The authors should state which version of the asset used and, possible, include 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 (c.g.. website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, www.ssppppsssaaeee thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's 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: All the new assets are well documented. Detailed explanations of the datasets and codes are provided.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmdel 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 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:
This paper involves neither crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
99

---

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 equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve research with 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 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 the institution conducting the review.
100


Reviews:
[{"review_id": "GDmWS9TGbZ", "reviewer": "Reviewer_oEdC", "summary": "This paper proposes a PDE foundation model based on a multiscale Swin Transfromer backbone and a flexible pretraining strategy.", "questions": "See Weaknesses", "limitations": "See Weaknesses", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-organized and clearly written.\n2. The experimental results are comprehensive and solid, which is a valuable contribution to the research.\n3. The studied topic is very interesting.", "weaknesses": "1. Recently, several other PDE foundation models have been proposed. Could you compare your model with [1] to further demonstrate the capability of the scOT backbone?\n2. I am interested in understanding the purpose of adding the ConNeXt block. I think it seems to be somewhat over-designed. Could you conduct some ablation studies to justify its inclusion?\n3. The generalization ability to downstream tasks is crucial for foundation models. Aside from relying on dataset diversity, where does Poseidon's generalization ability come from? Do the scOT backbone and all2all pretraining method offer inherent advantages for enhancing generalization?\n4. PDEs are often solved on discretized meshes, which may be highly irregular [2]. How can Poseidon be extended to handle irregular mesh data?\n\n[1] DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training\n\n[2] Geometry-Informed Neural Operator for Large-Scale 3D PDEs", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We start by thanking the reviewer for your appreciation of the merits of our paper and your welcome suggestions to improve it. We address your detailed concerns below. \n\n[W1:] The reviewer's suggestion to compare with DPOT is very well-taken. As we had clearly stated in our paper (see line l347-348), the DPOT model was not publicly available when we submitted our paper, making it impossible to compare with it then. In the meanwhile, DPOT has been released publicly and we have followed the reviewer's suggestion to compare with it. However, we could like to state that i) DPOT takes a sequence of time steps and outputs the next step. Our objective (l86-87) is to learn the entire trajectory of the operator from only the initial data. Hence, DPOT needs to be extended to our setup -- we do so by following exactly the same modifications that we introduced for extending the MPP foundation model in our paper (see SM C.6 and Figure 6) and ii) DPOT only allows an input of 4 channels, which precludes us from testing on the GCE-RT downstream task (SM B2.9) with 5 input channels. Also, given the tight timeframe of the rebuttal, we could only compare with the DPOT-M model (with ca. 120 M parameters), which is comparable in size with our Poseidon-B model. As fine-tuning DPOT takes considerable compute on our GPUs, we had time to fine-tune it for 3 representative tasks, NS-SL (see SM B2.3), CE-RPUI (SM B2.7) and Wave-Layer (SM B2.11). The corresponding EG and AG scores (see Eqn. 11 for definition) for DPOT on these tasks are: NS-SL (EG$=3.9$, AG$=2.1$), CE-RPUI (EG$=40.8$, AG$=3.2$) and Wave-Layer (EG$=21.8$, AG$=4.5$). For your convenience, we reproduce the corresponding scores for the comparable (in size) Poseidon-B model from SM Table 8 as NS-SL (EG$=19.1$, AG$=4.7$), CE-RPUI (EG$=370.8$, AG$=6.2$) and Wave-Layer (EG$=24.9$, AG$=4.7$). Also, we can compare the scores of other models from Table 1 to conclude that DPOT-M 's performance lies between CNO-FM and Poseidon-B for these tasks. It is a strong baseline but preliminary results suggest that it is not as accurate or efficient as Poseidon-B and Poseidon-L as seen here. We plan to include the results for this extended DPOT in the  **CRV**, if accepted and thank the reviewer for suggesting its inclusion.\n\n[W2:] Regarding the reviewer's question about ConvNeXt layers, we would like to point out that some form of residual connections between the encoder and decoder are necessary for stable training of scOT.  To study the role of ConvNeXt in this context, we follow your suggestion and ablate scOT by replacing it therein with plain residual connections. This study is performed on 2 downstream tasks-- Poisson-Gauss (**SM** B2.14) and Helmholtz (**SM** B2.15), where we train the underlying scOT on 1024 samples to obtain the following test errors: Poisson Gauss: with ConvNeXt ($0.013$) vs. with plain Residual block ($0.017$) and  Helmholtz: with ConvNeXt ($0.068$) vs. with plain Residual block ($0.095$). Thus, in both cases, there is an advantage of using the ConvNeXt layers. We can check this fact for other tasks as well and report the results in a **CRV**, while making a comment on the utility of ConvNeXt as the reviewer has righty suggested. \n\n[W3:] The reviewer's suggestion about further highlighting the generalization ability of Poseidon is excellent. In the paper, we have already presented several factors affecting generalization. To recall, i) Model architecture does matter (see l261-277, Table 1 and **SM** Table 11) which clearly show that the scOT backbone generalizes much better than a CNO backbone, even when both models are trained on exactly the same pretraining data. ii) Data size is also key (see l297-303 and **SM** Figures 24-38) as we show that greater size of pretraining dataset enables more accuracy even on downstream tasks. It is in this context that all2all training plays a role as it increases training dataset size, leading to more accuracy (see **SM** Figure 43) iii) Diversity of pretraining data is absolutely crucial (l304-312) for generalization as you correctly point out and iv) the choice of the pretraining dataset which implicitly contains a rich variety of physics that is learnt by the foundation model, enabling it to generalize better. We have highlighted this point in **SM** D.4 with 3 case studies. In particular, in D.4.2 (with Allen-Cahn reaction-diffusion PDE) and in D.4.3 (with elliptic Poisson PDE), we have shown that the latent space of Poseidon is rich enough to learn unseen physics (reaction-diffusion, steady state diffusion etc) by fine-tuning with very few samples. These different factors affecting generalization will be further discussed and highlighted in the main text of a **CRV**, if accepted. \n\n[W4:] In response to the reviewer's question about Poseidon's ability to handle data on irregular grids coming from PDEs with complex geometries, we would like to point out the SE-AF downstream task (**SM** B2.13), where the underlying data is on an irregular grid (Figure 3), is precisely considered for this very purpose. We follow the protocol of Ref. [62] to process the data in order to feed it into Poseidon (and all the baselines) and find that the results, presented in **SM** Figure 15 , with Poseidon models is very good, even though this test case requires Poseidon on generalize on several fronts not encountered in the pre-training dataset, namely i) to steady states ii) to an operator mapping domain shape to density field and iii) to irregular grids and non-periodic boundary conditions. Needless to say, SE-AF is only 1 task with irregular grids and further evaluation of Poseidon on other such tasks, like some mentioned in the GINO paper could be interesting. However, the experience with SE-AF bodes well for Poseidon in this context.\n\nWe sincerely hope to have addressed  all your concerns and would kindly request the reviewer to update their assessment accordingly.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lRr79ZR5FS", "reviewer": "Reviewer_KcRd", "summary": "The paper introduces a foundation model for learning PDE solution operators, with a proposed architecture, training method and training dataset consisting of fluid dynamics PDEs. The foundation model was evaluated on various downstream tasks and was shown to outperform baselines in terms of sample efficiency and accuracy.", "questions": "- A fair comparison with FNO would involve fine-tuning an FNO that has already been trained with other PDEs (i.e. using meta-learning methods/transfer learning). What would be the sample efficiency for that, and evaluation comparison compared to Poiseidon? Especially for situations where the FNO has already been trained on a given PDE, and is fine-tuned for another initial condition.\n\n- The authors chose to use L1 error for the underlying final time, but in many settings the intermediate dynamics may be more important and challenging to model compared to the final state where T>>0. Please share empirical results where the evaluation is based on different time indices, and/or some averaged quantity. Ideally, some there should be some charts/visualization confirming that Poseidon's better performance is robust to the chosen time index.\n\n- Table 1 shows results comparing Poseidon-L against benchmarks and FNO. From the main paper, it is not clear if the comparisons are fair, as the benchmark model parameters sizes are not indicated. Please provide indication of the relative sizes of the various models, and whether it is fairer to use Poseidon-L or a smaller version as comparison. Key results from the appendix should be shifted to the main paper if they are needed to support a major claim.\n\n- For the SE-AF task, a model trained from scratch could even perform better than the largest Poseidon model. This is significant, as the CNO is much smaller and did not benefit from any amortized training at all. It would be useful for the authors to provide some analysis on why this is the case, and whether it reflects a significant weakness of the Poseidon model.\n\n- It is unclear what the key differences and contributions in terms of architecture innovation and training method are, compared to benchmarks. It would help if the main paper had a more explicit comparison of what are the new key components that this paper is contributing to the literature, as well as ablation results on how these components impact the final performance. \n\n- For the key results (e.g. Table 1), it would be useful to provide some clear indication of the statistical significance of the results explicitly (e.g. error bars)\n\n- Does Poseidon learn the PDE solution operator where PDE parameters (e.g. viscosity of the fluid) can vary? Some empirical evaluation of this would be useful.", "limitations": "In the main paper, it would be useful to more clearly highlight some of the limitations of the model such as its longer inference time, as well as the paper's limited scope to mainly PDEs governing fluid dynamics.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper addresses an important problem of developing practical foundation models for dynamical systems.\n- The paper also highlighted results of several experiments that demonstrate Poseidon's strong empirical performance compared to baselines in the 15 downstream tasks.\n- The pre-trained models, when released, could serve as a good base for future work to be built upon.", "weaknesses": "- The paper could improve in clarity. For example, the approach section (2) need to be made much clearer. Currently, many terms, operations and acronyms are left undefined, with too many references to the appendix for critical information, making it very hard to parse and follow. This is especially so for the model architecture section. It will be very useful to provide an overview of the key architecture components, and how intuitively these differ from standard model architectures. Some of the other sections, e.g. pretraining/finetuning, can likely be significantly simplified.\n\n- It would help a lot if the main paper contains some summary of what the training dataset and downstream task PDEs roughly corresponds to -- now they are just acronyms. Some of the tasks seem quite highly correlated to one another, e.g. solutions of the same PDE but with different initial condition.\n\n- The problem formulation seem to not explicitly consider the learning of solution operators that take into account different PDE parameters, but only different initial conditions\n\n- The Poiseidon models have a significantly higher inference time compared to benchmarks. This is an important weakness that should be highlighted in the main paper rather than placed in just the appendix, especially for foundation models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We start by thanking the reviewer for your appreciation of the merits of our paper and your welcome suggestions to improve it. We address your detailed concerns below. \n\n[W1/W2:] The reviewer's concern about clarity are well-taken as are your suggestions to improve it. Given the page limit, we had to make some choices regarding what to present in the main text and what to leave to the SM. We can certainly move more material from SM about methods to main text in a CRV as we can add an extra page there. A short description of the downstream tasks is in lines l191-206, with a detailed description in SM B.2 and a summary in Table 4. We will expand the description in the main text and move Table 4 there to add clarity. We would also like to point out that only 6 of the 15 tasks involve PDEs seen during pretraining and the majority of tasks (9/15) consider *new* physics in the form of either adding new terms to PDEs or new PDEs altogether. \n\n[W3:] Regarding the reviewer's point about the tasks not considering operators with different PDE parameters, we have clearly explained how PDE parameters (coefficients/forcings etc) are included in our problem formulation (see l79-80 Main Text, l1274-1277 on forcing in Navier-Stokes, l1334-1336 on gravity in Euler and l1363-1367 on coefficients/material properties in the Wave Eqn). In addition, there are 3 steady state PDE tasks  (steady Euler (SE-AF), Poisson, Helmholtz) where the operators again map PDE parameters (coefficients, sources) to the solution, see e.g. Eqn (69) and l1463-1470 for Helmholtz. Thus, almost half (7/15) of the downstream tasks actually involve operators which map PDE parameters, not just initial conditions.  \n\n[W4:] The inference times for all models (reported in SM Table 11) only show Poseidon-L (with 630M params) as having a larger inference time compared to baselines. This was because it was inferred on a RTX3090 GPU whereas other models were evaluated on RTX4090s as is clearly mentioned in the caption of Table 11. We reran inference for the setup in Table 11, on a RTX4090 to find that Poseidon-L has an inference time of only 4 ms, which is much more comparable to baselines. We apologize for the possible confusion and reiterate that inference time of Poseidon-L is only a factor of 2 over FNO whereas Poseidon-T (with 21M params) is actually faster than FNO at inference (Table 11), while at the same time being much more accurate and sample efficient (see SM Tables 8 and 9). Moreover, all ML models are *3-5 orders of magnitude* faster than the physics-based PDE solvers in terms of inference. Thus, we show that inference time is not a limitation for Poseidon models. \n\n[Q1:]  We already have an answer in our paper to the reviewer's question by replacing FNO with CNO in your suggested analysis.  As CNO (with same number of parameters) is shown to be outperforming FNO in 14 out of the 15 tasks (Table 1), it clearly constitutes a stronger baseline on this task set for your question. So, we also built a foundation model with a CNO backbone (CNO-FM) and pretrained it on exactly the same data as Poseidon, Yet, we find that (Tables 1,8 and 9) CNO-FM is significantly inferior in performance to the Poseidon models. We have highlighted this point in l261-267 of the main text. Hence, Poseidon scores over CNO in your suggested comparison.\n\n[Q2:] The reviewer has a valid point about comparing model performances at different time indices. On all our 12 time-dependent tasks (3 are steady-state), we have found that the maximum error (for all models) occurs at the final time. Hence, we used the final time error for comparison. Nevertheless, we follow your suggestion and plot errors at different times for Poseidon-B and FNO (see Fig 1 in the 1-page.pdf and pt 1 in the reply to all reviewers) for 2 tasks to illustrate that Poseidon's gain in performance is consistent over time. We will include figures like this for all 12 tasks and all models in the CRV, if accepted.\n\n[Q3:] The sizes for all models are shown in Table 5. Poseidon-L is an order of magnitude larger than FNO/CNO and Poseidon-B is comparable in size to other foundation models (CNO-FM,MPP) while Poseidon-T is actually smaller in size than FNO/CNO. Yet, from Table 8 and 9, we see that it still significantly outperforms FNO/CNO in terms of accuracy/sample efficiency. This issue of relative size has been discussed in l278-286 of main text and we are happy to highlight it further in a CRV. \n\n[Q4:] SE-AF (SM B 2.13) is a very challenging downstream task for Poseidon as it has to generalize on multiple fronts (see l1436-1440), when compared to pretraining, namely i) to steady states ii) to an operator mapping Domain shape to solution field iii) to irregular grids and non-periodic BCs. In spite of this challenge, Poseidon models did very well (see SM Figure 15) and scaled better to catch up with CNO with more training samples. Hence, we do not consider performance on this task as a limitation but rather as highlighting their potential for generalization. \n\n[Q5:] We can add further elements from SM D.4 and D.5 to the discussion in l261-312 on factors underpinning model performance.\n\n[Q6:] Full Error distributions (pdfs) for Poseidon-B have been provided in SM Fig. 49 and we plan to add figures like Fig. 2 (left) (1-page pdf) to compare error distributions of all models in the CRV. \n\n[Q7:] Please see pt [W3] regarding PDE parameters. Following the reviewer's excellent suggestion, we have considered an additional task where the viscosity of the NS Eqns is changed (see pt [3] in the reply to all reviewers) and Fig 2 (right) (1-page pdf) which shows that Poseidon is able to readily learn this solution operator with a new viscosity. We will add this test to the CRV \n\nWe sincerely hope to have addressed  all your concerns, particularly about possible limitations for dealing with PDE parameters and inference times of Poseidon, and would kindly request the reviewer to update their assessment accordingly.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5hYIEIgwid", "reviewer": "Reviewer_UfXd", "summary": "The paper introduces a new PDE foundation model, named Poseidon. The backbone of the model is a multiscale vision transformer. A data augmentation based on the semi-group property of time-dependent PDEs is also proposed to scale up the amount of training data. After pretraining, Poseidon has shown higher accuracy in a variety fluid dynamics benchmark problems than other task-specific neural operator, such as Fourier Neural Operator. The advantage of the proposed model architecture has been proven by comparing the accuracy with other baseline model of comparable size, such as the Multiple Physics Pretraining (MPP). In addition, the performance of Poseidon improves as the model size increases, verifying the scaling law in the literature of Large Language Models.", "questions": "1. How long does it take to pre-train the model? Is it efficient compared to existing PDE foundation models?", "limitations": "The authors did address the limitations.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors proposed Poseidon, a new PDE foundation model based on multi-scale vision transformer. It also incorporates the semi-group property of time-dependent PDEs to do data augmentation, scaling up the training data. Extensive experiments and evaluations are conducted across a suite of 15 challenging downstream tasks. The authors have shown that Poseidon outperforms existing baselines, both in terms of sample efficiency and accuracy. The use of a large-scale, diverse dataset for pretraining further underscores the robustness and reliability of the model. The paper has shown that the model architecture does matter by comparing its performance with other baseline models with comparable size. In addition, with the increase of the model size, Poseidon also shows better accuracy. The framework shows its potential as a general-purpose foundation model for PDEs, capable of generalizing to new and even unseen physics with minimal task-specific training.", "weaknesses": "The robustness of POSEIDON to noisy or incomplete training data is not thoroughly examined. This is important for real-world applications where data can often be imperfect. The paper also does not provide a comprehensive analysis of the computational cost and memory usage of POSEIDON compared to baseline models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We start by thanking the reviewer for your appreciation of the merits of our paper and your welcome suggestions to improve it. We address your detailed concerns below. \n\n[W1:] The reviewer's suggestion on evaluating the robustness of Poseidon to noisy data is excellent. We follow it up by considering one of our downstream tasks CE-RPUI (SM B2.7 for detailed description) and at inference time, we add Gaussian noise to the inputs (initial conditions) at different noise-to-signal ratios (NSRs) of $0.1$%, $1$% and $3$% respectively. The resulting errors, computed with respect to a Ground Truth where the outputs are not noisy, for varying numbers of training trajectories, are shown in Figure 3 of the 1-page rebuttal pdf. The errors in the zero noise (clean) case are also shown in this Figure. We observe from this figure that Poseidon-L's performance is robust to input noise and the error does not grow significantly even when the noise level is an appreciable $3$%, demonstrating the robustness of this foundation model. We observe similar behavior with other Poseidon models and would include a discussion on this topic in the CRV, if accepted. We thank the reviewer for this suggestion that further highlights the robustness of our model.\n\n[W2:] Regarding the reviewer's point about a comprehensive analysis regarding the computational costs of Poseidon, we would like to point out that the training costs have been provided in SM Table 10 (see SM sec E for further details). This table shows that the training, even for the biggest Poseidon-L model only took place on consumer grade GPUs with 24GB of VRAM. Similarly, the inference costs were provided in Table 11 where all the models barring Poseidon-L were inferred on a RTX-4090 GPU whereas Poseidon-L was inferred on a RTX-3090. In the meanwhile, we have also timed inference runs for Poseidon-L on a RTX-4090 too to find that the inference time is only 4 ms. Thus, all Poseidon models have inference times between 1.6-4 ms, which is comparable to the ML baselines and is between 3-5 *orders of magnitude* faster than physics-based PDE solvers. We will provide the new inference time of Poseidon-L and expand on the discussion in SM sec. E in a **CRV**, if accepted. \n\n[Q1:] The pretraining times on consumer GPUs (RTX4090s) are provided in Table 10. Regarding other foundation models, there are 2 models with a similar framework as ours in the literature (MPP of Ref. [50] and DPOT of Ref. [20]). The pretraining costs of these models *have not been transparently disclosed* to the best of our knowledge. Instead, Ref. [20] and [50] say that their models have been trained on professional GPUs (H100 or V100) but do not mention how much computational time it actually took for the pretraining. \n\nWe sincerely hope to have addressed  all your concerns and would kindly request the reviewer to update their assessment accordingly.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
At the outset, we would like to thank all three reviewers for their thorough and patient reading of our article. Their criticism and constructive suggestions will enable us to improve the quality of our article. If our paper is accepted, we will incorporate all the changes that we outline below in a camera-ready version (**CRV**) of our article. As allowed by the conference, we are uploading a one page pdf that contains figures with numerical results which support our arguments. These figures are described below. With this context, We also  proceed to answer the points raised by each of the reviewers individually in their respective rebuttal fields.

Yours Sincerely,  

Authors of *POSEIDON: Efficient Foundation Models for PDEs*.

Detailed Description of the 1-page Rebuttal pdf
-------------------------------------------------------
The 1-page pdf contains 5 figures arranged in 3 rows with the following content,

[1] Row 1: Figure 1: This figure shows how the test error grows with time when our foundation model (Poseidon-B) and the FNO baseline are evaluated on two downstream tasks: NS-PwC (Left) and NS-SL (Right) (See **SM** B.2.1 and B.2.3 for detailed description of the tasks). We observe from both figures that the gains in accuracy with Poseidon-B over FNO are consistently observed over time and are substantial for all the time indices that we consider. We plan to add similar figures for all time-dependent tasks and with all the baselines in a **CRV**, if accepted.  

[2] Row 2: Figure 2 (Left): shows empirical histograms representing the entire test error distributions for Poseidon-B and FNO for the NS-SL downstream task (Detailed description in **SM** B.2.3) when 128 trajectories are used to train FNO (from scratch) and fine-tune Poseidon-B. We plan to add similar figures for all tasks and all baselines in a **CRV**, if accepted. 

[3] Figure 2 (Right) shows how test errors scale for Poseidon-B and FNO for a new *Downstream Task*, suggested by Reviewer KcRd. In this task, we consider the Navier-Stokes Equations (**SM** Eqn. 31 of the submitted paper) with a viscosity coefficient $\nu=4\times 10^{-3}$. The ground truth data is generated using the Azeban spectral hyper viscosity solver (Ref. [64]). This new viscosity coefficient is very different from the setup of the pretraining data and downstream tasks in our original paper as there, only a hyperviscosity of $4 \times 10^{-4}$ was applied to high-enough Fourier modes in order to model the incompressible Euler equations with zero viscosity (see **SM** lines l1082-l1086). In this task, the initial conditions are identical to the NS-PwC downstream task (**SM** B2.1 for details). We see from Figure 2 (right) that Poseidon-B generalizes very well to this new viscosity coefficient, that was not seen during pretraining, and outperforms FNO readily, in terms of both sample efficiency and accuracy. In particular, the AG and EG scores of Poseidon-B (defined in Eqn. 11) are $EG=925.5$ and $AG=47.5$, which are completely comparable to (even better than) the scores of $EG=1024$ and $AG=19.7$ (see **SM** Table 8) for the NS-PwC task with much lower hyperviscosity that was reported in the main paper. We will add the results of other baselines to this figure in a **CRV**, if accepted. 

[4] Row 3: Figure 3 presents how the test error scales for the Poseidon-L model when noise in injected into the input at the level of inference. To study this question, we consider the CE-RPUI task (**SM** B.2.7 for details) and add Gaussian noise to the inputs (initial conditions) at different noise-to-signal ratios (NSRs) of $0.1$%, $1$% and $3$% respectively. The errors in the zero noise (clean) case are also shown in this Figure. The errors are computed with respect to a Ground Truth where the outputs are not noisy. We observe from this figure that Poseidon-L's performance is robust to input noise and the error does not grow significantly even when the noise level is an appreciable $3$%, demonstrating the robustness of the foundation model.

Meta-review: This paper presents a PDE foundation model based on a multiscale swin transformer architecture. It presents a training method and a training dataset consisting of fluid dynamics PDEs. The evaluation on various types of PDEs demonstrate the effectiveness of the foundation model. Though the reviews are generally consistently positive and I recommend to accept, various changes should be included in the final version, as promised during the rebuttal. In particular, the comparison with existing foundation models (e.g., DPOT) is too brief in the current Intro and experiments to fairly show the existing progress. A new paragraph should be included in the Intro to compare with the existing foundation models to provide a complete literature review.
Decision: Accept (poster)