Poseidon: Efficient Foundation Models for PDEs

Maximilian Herde; Bogdan Raonić; Tobias Rohner; Roger Käppeli; Roberto Molinaro; Emmanuel de Bézenac; Siddhartha Mishra

Source-linked AI summary

Poseidon: Efficient Foundation Models for PDEs

Maximilian Herde, Bogdan Raonić, Tobias Rohner, Roger Käppeli, Roberto Molinaro, Emmanuel de Bézenac, Siddhartha Mishra

arXiv:2405.19101v2May 29, 2024cs.LG

TL;DR

Existing PDE operator-learning methods require many costly training examples, while it remains unclear whether foundation-model pretraining can generalize across diverse, unseen PDEs. Poseidon addresses this with a multiscale operator transformer and scalable trajectory-based pretraining, outperforming baselines across 15 downstream tasks while achieving comparable error with far fewer samples.

Problem
Existing PDE operator-learning methods are not sample efficient, and it remains unclear whether pretraining on few PDEs can generalize to unseen and unrelated equations.
Source: §1 · p. 2
Method
Poseidon combines a multiscale operator transformer, time-conditioned normalization, semi-group-based all2all training, and large-scale fluid-dynamics pretraining.
Source: §1 · p. 2, §4 · p. 10
Results
20 samples suffice on average to match FNO’s error with 1024 samples, while Poseidon outperforms FNO across 15 diverse downstream tasks.
Source: §1 · p. 3, §3 · p. 7
Takeaways & Limitations
Poseidon generalizes to unseen and a priori unrelated physical processes using only a few task-specific training examples.
Source: §1 · p. 3
Poseidon was trained and evaluated on only a few PDEs and data distributions, with pretraining focused primarily on Cartesian geometries and time-dependent fluid dynamics.
Source: §4 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

POSEIDON addresses the sample inefficiency and generalization challenges of PDE operator learning through foundation-model pretraining, a multiscale operator transformer, and trajectory-based data scaling. Across 15 out-of-distribution tasks, it improves accuracy and sample efficiency, generalizes to unseen physics, and scales with model and dataset size.

Method

POSEIDON combines a multiscale Operator Transformer, all2all trajectory training, and an open-source fluid-dynamics pretraining dataset.The model uses shifted-window attention adapted for operator learning and solution trajectories from time-dependent PDEs to scale training data.Source: §1 · p. 2

Results

POSEIDON requires 20 samples to attain the same error level as FNO with 1024 samples on median across downstream tasks.The model is reported as up to multiple orders of magnitude more sample efficient than task-specific neural operators.Source: Figure 1 · p. 2, §1 · p. 3

Evaluation

POSEIDON outperforms baselines in accuracy and sample efficiency across 15 downstream tasks spanning diverse PDE types, all out of distribution from pretraining.The tasks include linear and nonlinear, time-dependent and independent, elliptic, parabolic, hyperbolic, and mixed-type PDEs.Source: §1 · pp. 2–3

Generalization

POSEIDON generalizes with few task-specific examples to unseen and a priori unrelated physical processes, including nine tasks involving PDEs absent from pretraining.Pretraining uses compressible Euler and incompressible Navier–Stokes fluid-dynamics operators, while all downstream tasks are out of distribution.Source: §1 · pp. 2–3

Scaling

POSEIDON scales with model and dataset size during both pretraining and downstream adaptation, supporting its use as a general-purpose PDE foundation model.The paper attributes this potential to effective representations learned during pretraining and leveraged on unrelated downstream PDEs.Source: §1 · p. 3

2 Approach

Poseidon learns time-dependent PDE solution operators that map initial conditions to complete trajectories, using a multiscale, time-conditioned operator transformer. Its all2all training exploits the semi-group property to scale trajectory data quadratically, while pretraining transfers across PDEs and data distributions.

Problem Formulation

The operator-learning task approximates S(t, a) from an initial datum so the model generates the entire solution trajectory over t ∈ [0, T].The learned operator must produce the trajectory from only the initial datum and boundary conditions.Source: §2 · p. 3

Model Architecture

scOT is a hierarchical multiscale SwinV2 transformer with patch embeddings, shifted window attention, and a U-Net-style encoder-decoder for approximating S(t, a).Windowed attention improves efficiency, while shifting windows across layers enables domain-wide interactions; patch merging and expansion provide multiscale processing.Source: §2 · pp. 3–4
Lead-time-conditioned layer normalization enables scOT to evaluate the learned operator continuously in time.The conditioning uses learnable affine functions of time, α(t) and β(t).Source: §2 · p. 4

Training and Inference

Training fits scOT to sampled solution trajectories, and inference can directly evaluate any time or use autoregressive rollouts through successive applications of the trained model.Direct evaluation produces continuous-time trajectories, while κ successive applications approximate the operator at a target time.Source: §2 · pp. 4–5
All2all training uses the PDE solution operator’s semi-group property to provide O(K^2) training pairs per trajectory instead of O(K) under vanilla training.The method uses pairs (u(t_k), u(t̄_k)) with k ≤ k̄ and a relative loss to balance differing operator-output scales.Source: §2 · p. 5

Pretraining

Pretraining spans indexed PDE types and initial-data distributions, then transfers most parameters while selectively fine-tuning task-specific components for downstream operators.Inputs are channel-aligned across pretraining operators; newly initialized parameters and time embeddings use higher learning rates during fine-tuning.Source: §2 · pp. 5–6

3 Experiments

POSEIDON is evaluated on a diverse, out-of-distribution suite of 15 PDE operator-learning tasks and consistently achieves superior accuracy and sample efficiency. Its performance reflects the combined importance of pretraining data quality, model architecture and scale, while enabling strong generalization to unseen physics.

Experimental setup

The experiments pretrain POSEIDON on 77840 trajectories from six Euler and Navier–Stokes operators, then evaluate it on 15 diverse tasks spanning unseen PDEs, physical processes, operator types and boundary conditions.The downstream suite contains 11 nonlinear and 4 linear tasks, including 9 PDEs not seen during pretraining and 7 operators unlike the initial-condition-to-later-solution mapping used for pretraining.Source: §3 · p. 6

Main results

20 task-specific samples suffice on average for POSEIDON-L to match FNO’s error with 1024 samples, while 13 of 15 tasks require an order of magnitude fewer samples.At matched sample counts, POSEIDON-L’s accuracy gains over FNO range from 10% to a factor of 25, averaging an order of magnitude; POSEIDON also outperforms CNO and scOT in sample efficiency and accuracy.Source: §3 · p. 7, Table 1 · p. 8

Generalization to unseen physics

POSEIDON is the best-performing model on 8 of 9 tasks involving PDEs absent from pretraining, including every time-dependent task, and is only marginally surpassed on one time-independent task.All downstream tasks are out-of-distribution, and the time-independent tasks are treated as long-time limits of the time-dependent formulation.Source: §3 · pp. 6–7

Architecture and baselines

POSEIDON’s architecture matters beyond pretraining and model size: it beats CNO-FM on 14 of 15 tasks, while even the smallest POSEIDON-T outperforms comparable trained-from-scratch and foundation-model baselines.CNO-FM needs approximately 100 task-specific examples versus approximately 20 for POSEIDON to match FNO’s 1024-sample error; finetuned MPP modestly beats FNO on 8 tasks but fails on the remainder.Source: §3 · pp. 7–8

Scaling and pretraining data

Increasing POSEIDON’s model size generally improves downstream performance, while increasing pretraining-data size significantly improves accuracy on 9 of 15 downstream tasks; remaining tasks are on par or marginally inferior.At matched dataset size, less-diverse pretraining data performs worse on 10 of 15 tasks and is on par on 4; POSEIDON-L generally outperforms POSEIDON-B on downstream tasks.Source: §3 · pp. 8–9

Generalization mechanisms

POSEIDON can adapt to radically different physics with few examples, and a frozen latent space remains effective on the elliptic, time-independent Poisson-Gauss task.For Poisson-Gauss, it approximates the input after one example and learns diffusion and smoothing from 16 examples; frozen-latent finetuning still learns the solution operator accurately.Source: §3 · p. 9

4 Discussion

POSEIDON combines multiscale transformer architectures, time conditioning, and semi-group-based training to achieve strong, sample-efficient generalization across diverse PDE tasks, including unseen physics. Its scope remains limited by the narrow pretraining coverage, but broader datasets and extensions to new geometries and downstream problems are identified as promising directions.

Summary

POSEIDON is best on 14 of 15 downstream tasks, uses a median 50-fold fewer task-specific samples than FNO at equal error, and generalizes to unseen physics.It also scales with model and dataset size, while all models and datasets are publicly released.Source: §4 · p. 10
POSEIDON combines a multiscale SwinV2 transformer, time-modulated layer normalization for continuous-time evaluation, and all2all training based on the semi-group structure.It is pretrained on compressible Euler and incompressible Navier–Stokes operators and evaluated on 15 out-of-distribution tasks.Source: §4 · p. 10

Related Work

POSEIDON significantly outperforms DPOT in accuracy and sample efficiency, while comparisons with [67], [74], and are limited by differences in scope or task.MPP [49] and DPOT [19] are identified as comparable general-purpose PDE foundation models; the paper reports extensive prior comparisons with MPP.Source: §4 · p. 10

Limitations

POSEIDON was trained and evaluated on only a small subset of PDEs and data distributions, especially lacking time-independent equations, broader time scales, and diverse non-Cartesian geometries.Future work targets expanded pretraining datasets, uncertainty quantification [45], inverse problems, and PDE-constrained optimization [46].Source: §4 · p. 10

Supplementary Material for: POSEIDON: Efficient Foundation Models for PDEs. … B Datasets

The supplement specifies scOT’s operator-learning formulation and computational U-Net realization, including multiscale SwinV2 processing and time-conditioned normalization. It also notes that pretraining and downstream datasets are publicly available through PDEGYM.

A.1 Operator Learning with scOT

scOT partitions function-space inputs into patches, embeds them into latent representations, and processes outputs with SwinV2 transformer blocks using shifted windows for cross-window attention.The operator formulation uses shared learnable weights for patch embedding and shifted windows across layers so points can interact over multiple blocks.Source: §A.1 · p. 17

A.2 Computational Realization of scOT

The computational scOT architecture is an encoder-decoder with SwinV2 stages, patch merging, ConvNeXt skip connections, a convolution-free bottleneck, patch recovery, and mixup.This realization implements the continuous operators through discrete patch, attention, convolution, and reconstruction operations.Source: §A.2 · p. 18
SwinV2 attention operates within M × M windows with cyclic shifts between blocks, using cosine similarity and logarithmic relative-position biases.The attention formulation includes learnable head-specific parameters and shared positional-bias processing across heads.Source: §A.2 · pp. 18–19
Time-conditioned layer normalization enables scOT to process time-dependent PDEs, while time-independent PDEs use ordinary layer normalization.The time-dependent normalization applies learnable time-varying gain and bias to attention outputs and is also used in other architectural components.Source: §A.2 · pp. 18–19
The model’s multiscale hierarchy downsamples encoder tokens by halving resolution and doubling latent dimension, then reverses this through decoder patch expansion.Encoder stages also connect to time-conditioned ConvNeXt blocks, while decoder expansion doubles resolution and halves latent dimension.Source: §A.2 · p. 19
The implementation fixes key scOT hyperparameters at p = 4, M = 16, L = 4, [h1, h2, h3, h4] = [3, 6, 12, 24], and nc = 2.These fixed choices reduce the number of hyperparameters used to instantiate scOT.Source: §A.2 · p. 20

B Datasets

The supplement describes datasets used for pretraining and downstream evaluation, all of which are publicly available through the PDEGYM collection.The collection is hosted on Hugging Face.Source: B Datasets · p. 21

B.1 Pretraining Datasets … C.1.3 POSEIDON-L

This section describes the pretraining and downstream datasets spanning diverse fluid-dynamics operators and additional physical processes, with variation in initial conditions, geometries, boundary conditions, and temporal dependence. It also details the data generation, model comparisons, transfer strategy, and configurations used to evaluate Poseidon across pretrained model sizes.

B.1.1 NS-Sines / B.1.2 NS-Gauss

NS-Sines uses p = 10 random sine-and-cosine modes for its initial conditions, whereas NS-Gauss uses a superposition of p = 100 random Gaussians to initialize vorticity and recovers velocity from it.Each dataset contains 20000 trajectories with the same train/validation/test split and periodic Navier–Stokes evolution.Source: §B.1.1 · p. 22, §B.1.2 · p. 23

B.1.3 CE-RP / B.1.4 CE-CRP

CE-RP samples constant states on a fixed 2 × 2 partition, while CE-CRP generalizes it to curved, multi-partitioned subdomains with randomly sampled density, velocity, and pressure.Both datasets solve the compressible Euler equations with periodic boundary conditions and contain 10000 trajectories using the CE-RP split and shared time-stepping.Source: §B.1.3 · p. 23, §B.1.4 · pp. 23–24

B.1.5 CE-KH / B.1.6 CE-Gauss

CE-KH represents Kelvin–Helmholtz instability with piecewise compressible-flow initial data, while CE-Gauss initializes vorticity from 100 Gaussians and uses constant density ρ = 0.1 and pressure p = 2.5.Both datasets contain 10000 compressible-Euler trajectories with the same split as CE-RP.Source: §B.1.5 · p. 24, §B.1.6 · p. 24

B.1 Pretraining Datasets

The pretraining dataset contains six operators: NS-Sines, NS-Gauss, CE-RP, CE-CRP, CE-KH, and CE-Gauss; CE-KH and CE-RP are well known, while the other four are novel to the authors’ knowledge.Pretraining predicts four variables—density, horizontal velocity, vertical velocity, and pressure—over T = 1, storing 21 uniformly spaced snapshots at 128 × 128 resolution generated on 512 × 512 and downsampled.Source: §B.1 · p. 22, §B.1.6 · p. 24

B.2.1 NS-PwC / B.2.2 NS-BB / B.2.3 NS-SL / B.2.4 NS-SVS

The downstream suite includes Navier–Stokes benchmarks with piecewise-constant vorticity, Brownian-bridge initial conditions, shear layers, and sinusoidal vortex sheets.NS-PwC, NS-BB, NS-SL, and NS-SVS evaluate periodic Navier–Stokes solution operators, generally at t = 0.7; NS-SL uses 40000 trajectories, while the other stated datasets use 20000.Source: §B.2.1 · p. 25, §B.2.2 · p. 26, §B.2.3 · p. 26

B.2.5 NS-Tracer-PwC / B.2.6 FNS-KF

NS-Tracer-PwC adds a passive diffusive tracer transported by the Navier–Stokes flow, while FNS-KF adds a time-independent forcing term modeling two-dimensional Kolmogorov flow.These tasks introduce physical processes not completely encountered during pretraining and retain the NS-PwC initial velocity setup; NS-Tracer-PwC evaluates at t = 0.7.Source: §B.2.5 · p. 27, §B.2.6 · p. 27

B.2.6 FNS-KF

FNS-KF generates 20,000 forced Navier–Stokes trajectories and evaluates predictions at t = 0.7 using the 14th time step.The solution augments the state with a forcing field satisfying f_t = 0, and compares POSEIDON, CNO, and FNO using 128 training trajectories.Source: §B.2.6 · p. 28

B.2.7 CE-RPUI

CE-RPUI perturbs compressible-Euler interface amplitudes, locations, and shapes, generating 10,000 trajectories evaluated at t = 0.7.The task uses periodic boundaries and the same train/validation/test split and time-stepping as CE-RP, with comparisons against CNO and FNO.Source: §B.2.7 · p. 28

B.2.8 CE-RM

CE-RM models a stochastic Richtmeyer–Meshkov compressible-Euler problem, generating 1,260 trajectories with 21 snapshots and testing at t = 1.4.Approximate solutions come from the FISH high-resolution finite-volume code, using a 1030/100/130 train/validation/test split.Source: §B.2.8 · pp. 28–29

B.2.9 GCE-RT

GCE-RT adds gravitational forcing to a stochastic Rayleigh–Taylor instability benchmark for compressible Euler equations on a model neutron star.It generates 1,260 trajectories at 256^2 resolution, downsampled to 128^2, and evaluates at the seventh time step using earlier snapshots for training.Source: §B.2.9 · p. 29

B.2.10 Wave-Gauss / B.2.11 Wave-Layer

Each experiment generated 10,512 trajectories using finite-difference simulations at 128^2 resolution, with testing error evaluated at the 14th time step, t = 0.7.The simulations run to T = 1; Wave-Gauss saves 15 snapshots, while Wave-Layer saves 21.Source: §B.2.10 · p. 30, §B.2.11 · p. 31

B.2.12 ACE

ACE studies a nonlinear parabolic reaction-diffusion equation with periodic boundaries, generating 15,000 trajectories and testing at the 14th time step.The Allen–Cahn equation is explicitly characterized as very different from the PDEs used to construct the pretraining dataset.Source: §B.2.12 · p. 31

B.2.13 SE-AF

SE-AF maps perturbed RAE2822 airfoil shapes to steady compressible-Euler flow density on a non-Cartesian aerodynamic domain, using 10,869 solutions.The benchmark is time-independent, uses a shape-to-steady-state operator, and has geometry and boundary conditions unlike pretraining.Source: §B.2.13 · pp. 31–32

B.2.14 Poisson-Gauss / B.2.15 Helmholtz

Poisson-Gauss maps Gaussian-mixture source terms to steady solutions, while Helmholtz models frequency-domain wave propagation with spatially varying media and Dirichlet boundaries.Poisson-Gauss uses 20,000 finite-element solutions; Helmholtz sets ω = 5π/2 and constructs its medium coefficient from randomly parameterized Gaussians.Source: §B.2.14 · p. 33, §B.2.15 · p. 33

B.2.15 Helmholtz

The Helmholtz task is a steady-state solution-operator problem, with 19,675 finite-difference solutions at 128^2 resolution split into 19,035/128/512 train/validation/test examples.Predictions from POSEIDON, CNO, and FNO are visualized using 128 training trajectories.Source: §B.2.15 · p. 33

C Models and Baselines

POSEIDON models are compared with MPP [49], a CNO foundation model [60], and scratch-trained CNO, FNO [33], and scOT, with model sizes summarized in Table 5.The comparisons include foundation-model and state-of-the-art neural-operator baselines.Source: C Models and Baselines · p. 34, Table 5 · p. 34

C.1 POSEIDON Models

For mismatched downstream input/output functions, POSEIDON transfers all pretrained parameters except embedding, patch-recovery, and mixup-convolution parameters, which are trained from scratch.All other parameters remain trainable; finetuning uses task-dependent learning rates and generally excludes weight decay on time-conditioned layer norms.Source: §C.1 · pp. 34–35

C.1.1 POSEIDON-T

POSEIDON-T uses 48-dimensional embeddings, four SwinV2 blocks per level, and 21M parameters, with pretraining for 40 epochs at effective batch size 640.It uses AdamW, cosine decay with two warmup epochs, maximum learning rate 10^-3, weight decay 0.1, and gradient clipping at 5; downstream finetuning runs for 200 epochs.Source: §C.1.1 · p. 35

C.1.2 POSEIDON-B

POSEIDON-B uses 96-dimensional embeddings, eight SwinV2 transformer blocks at each level, 158M parameters, effective batch size 320, per-device batch size 40, and 39 epochs.The cited configuration specifies cosine decay with linear warmup of 2 epochs and gradient clipping with maximal norm 5.Source: §C.1.2 · pp. 35–36

C.1.3 POSEIDON-L

The cited POSEIDON-L configuration uses 192-dimensional embeddings, eight SwinV2 transformer blocks at each level, 629M parameters, effective batch size 128, 20 epochs, and one warmup epoch.Source: §C.1.3 · p. 36, §C.2 · p. 37

C.1.4 Models for Dataset Ablations (see Section D.3) … D Results

The paper specifies ablation models and several operator-learning baselines, including scOT, CNO, FNO, CNO-FM, and MPP, with distinct architectures and finetuning strategies. These baselines are configured for both in-context and out-of-context downstream PDE tasks, while MPP is adapted from autoregressive forecasting to trajectory generation.

C.1.4 Models for Dataset Ablations (see Section D.3)

Dataset-ablation models reuse POSEIDON-B’s pretraining and finetuning strategies while varying dataset size and diversity under directly comparable sampling schemes.Half-data models use the first half of every subset, eighth-data models use the corresponding reduced subsets, and less-diverse models omit NS-Sines, CE-CRP, and CE-Gauss.Source: §C.1.4 · p. 37

C.2 scOT

scOT is trained from scratch on every downstream task as a comparison baseline, using a 40M-parameter configuration.The stated configuration uses embedding dimension C = 48 and eight SwinV2 transformer blocks at each level.Source: §C.2 · p. 37

C.3 CNO

CNO maps bandlimited functions through a modified Operator U-Net composed of lifting, multiscale convolutional layers, activation operators, and projection.Its time conditioning uses instance normalization parameterized by small MLPs, with time also supplied as an additional constant input channel.Source: §C.3 · pp. 37–38

C.4 FNO

FNO combines lifting and projection operators with residual Fourier layers, and includes both time-conditioned instance normalization and time as an input channel for time-dependent problems.The trained FNO configuration has 37.0M parameters, five Fourier layers, and 20 Fourier modes.Source: §C.4 · pp. 39–40

C.5 CNO-FM

CNO-FM is pretrained on the same six datasets as POSEIDON and finetuned by transferring the base model while adapting input and output projections to downstream variable contexts.In-context tasks add a linear input layer, whereas out-of-context tasks additionally replace the projection layer with a randomly initialized one.Source: §C.5 · pp. 40–41

C.6 MPP

MPP is an autoregressive foundation model adapted to whole-trajectory PDE tasks through an all2all finetuning strategy that creates T(T + 1)/2 training samples per trajectory.Missing history is filled by copies of the earlier state; inference autoregressively predicts future snapshots, while dummy variables are excluded from prediction and testing errors.Source: §C.6 · pp. 41–42, Figure 6 · p. 42

D.1 Performance on Downstream Tasks · D.2 Scaling with respect to Model Size

Poseidon foundation models significantly outperform baselines across downstream tasks, with consistent data-scaling behavior and improved average performance as model size increases. Their scaling can include task-dependent biphasic learning dynamics when downstream operators differ substantially from pretraining.

D.1 Performance on Downstream Tasks

Downstream evaluation uses median relative L1 error on each function of interest, with task-dependent direct or autoregressive rollout procedures.Vector-valued functions are evaluated over the entire vector, and list-valued metrics are averaged.Source: §D.1 · p. 43, Table 6 · p. 44
POSEIDON models generally achieve significantly lower test errors than baselines at the same number of trajectories or samples across the downstream tasks.Scaling plots compare POSEIDON-L and POSEIDON-B with CNO-FM, MPP, FNO, CNO, and scOT, while POSEIDON-T is reported separately in Table 8.Source: §D.1 · p. 43, Figure 7 · p. 46, Figure 8 · p. 46
All considered models follow power-law scaling with data, while POSEIDON-L has an exponent of approximately 0.5 or higher except on CE-RM.The scaling factor, not only the exponent, strongly affects errors in the limited-data regime.Source: §D.1 · p. 43, Table 7 · p. 44
POSEIDON significantly outperforms all baselines on average across downstream tasks in efficiency and accuracy gains.Table 9 summarizes median efficiency gain and mean accuracy gain over all tasks.Source: §D.1 · p. 43, Table 9 · p. 45
On Poisson-Gauss, POSEIDON-B and POSEIDON-L show biphasic scaling after M_pt = 32, with exponents 0.23 then 0.99 and 0.33 then 0.94, respectively.The slower warmup phase is speculated to reflect learning an operator unlike those in pretraining, followed by faster task-specific learning.Source: §D.1 · p. 43, Figure 20 · p. 50

D.2 Scaling with respect to Model Size

Larger POSEIDON models also improve average downstream performance, as reflected in the comparisons between POSEIDON-L and POSEIDON-B.The corresponding efficiency and accuracy gains are reported in Tables 8 and 9.Source: Table 8 · p. 45, Table 9 · p. 45, §D.2 · p. 51
Increasing model size consistently decreases pretraining training and validation losses.Figure 22 reports loss curves and scaling at epoch 20 for different POSEIDON model sizes.Source: Figure 22 · p. 51, §D.2 · p. 51

D.3 Scaling with respect to Pretraining Dataset Size and Quality · D.4 Case Studies

POSEIDON-B scales with pretraining dataset size, while Figures 24–38 compare the effects of dataset quantity and diversity across downstream tasks. Because POSEIDON performs robustly even on unseen physical processes, the paper introduces three case studies to investigate the basis of this performance.

D.3 Scaling with respect to Pretraining Dataset Size and Quality

POSEIDON-B scales with pretraining dataset size, as shown by training and evaluation losses for one-eighth, one-half, and full datasets.Figure 23 reports losses through epoch 20 and scaling at epoch 20 for each dataset size.Source: §D.3 · p. 52, Figure 23 · p. 52
Figures 24–38 compare full-data and one-eighth-data POSEIDON-B models across downstream tasks using test-set median relative L1 error.The comparisons cover fluid, chemical, wave, elliptic, and related tasks, with figures specifying trajectories or samples as the data-axis quantity.Source: §D.3 · p. 52, Figure 24 · p. 53, Figure 25 · p. 53
The downstream comparisons also include FNS-KF, CE-RPUI, CE-RM, SE-AF, and GCE-RT, each evaluated through median relative L1 error.Figures 29–33 report these task-specific comparisons against the number of trajectories or samples.Source: §D.3 · p. 52, Figure 29 · p. 54, Figure 30 · p. 55
POSEIDON-B trained on half the pretraining dataset is compared with the same model pretrained on a less diverse dataset across downstream tasks.Figures 24–38 place this quantity-versus-diversity comparison in the right subfigure of each task result.Source: §D.3 · p. 52, Figure 24 · p. 53, Figure 25 · p. 53
Additional quantity-versus-diversity comparisons cover Wave-Layer, Wave-Gauss, ACE, Poisson-Gauss, and Helmholtz.Figures 34–38 report median relative L1 error against trajectories or samples for these tasks.Source: §D.3 · p. 52, Figure 34 · p. 56, Figure 35 · p. 56

D.4 Case Studies

The case studies are motivated by POSEIDON’s strong performance across tasks, including PDEs representing physical processes absent from pretraining.They are intended to explain the model’s robust performance.Source: §D.4 · p. 58
D.4 presents three case studies to investigate what underpins POSEIDON’s robust performance.The supplied passage introduces the case-study analysis but does not specify the individual studies’ contents.Source: §D.4 · p. 58

D.4.1 CE-RPUI · D.4.2 ACE · D.4.3 Poisson-Gauss

Across CE-RPUI, Allen–Cahn, and Poisson-Gauss, POSEIDON adapts accurately from few task-specific examples, including PDEs and physics absent from pretraining. Frozen latent representations retain substantial transferability, while diverse pretraining supports multiscale shock and vortex resolution.

D.4.1 CE-RPUI

POSEIDON-B outperforms CNO and FNO on complex CE-RPUI solutions, sharply resolving shocks and roll-up vortices with 128 finetuning trajectories.The benchmark combines large-scale shocks with small-scale vortices, which CNO and especially FNO fail to capture accurately.Source: §D.4.1 · p. 58
With 4 CE-RPUI trajectories, POSEIDON captures vortex roll-ups accurately, and approximation quality continues improving through 32 and 128 trajectories after poor zero-shot results.One trajectory recovers some large-scale features such as shock locations, showing that a few task-specific samples suffice to approximate the solution operator.Source: §D.4.1 · p. 58, Figure 39 · p. 59
Diverse pretraining improves CE-RPUI resolution of small-scale vortices, consistent with combining shock propagation from CE-RP, large vortices from CE-KH, and multiscale vortex dynamics from CE-CRP.The less-diverse pretrained model is less accurate and misses the CE-CRP contribution to small-scale vortex dynamics.Source: §D.4.1 · p. 58, Figure 40 · p. 60

D.4.2 ACE

POSEIDON learns the unseen reaction-diffusion physics of Allen–Cahn from few examples, despite pretraining only on convection-dominated Compressible Euler and high-Reynolds-number Navier–Stokes equations.With one trajectory it captures large-scale front propagation and diffusion, while additional trajectories refine local diffuse-front features.Source: §D.4.2 · p. 59, Figure 41 · p. 60
On Allen–Cahn, frozen-latent POSEIDON reaches error 0.031 with 32 trajectories, below FNO’s 0.037 with 128 trajectories but above fully finetuned POSEIDON-B’s 0.014.Only the embedding and recovery parameters are trained, retraining less than 0.5% of total parameters.Source: §D.4.2 · pp. 59, 61, Figure 42 · p. 61

D.4.3 Poisson-Gauss

POSEIDON-B learns the highly out-of-distribution Poisson-Gauss operator progressively: diffusion emerges by 16 samples, smoothing by 128, and local refinements by 512.The task differs from pretraining in being time-independent, diffusion-based, and governed by Dirichlet rather than periodic boundary conditions.Source: §D.4.3 · p. 61, Figure 43 · p. 62
For Poisson-Gauss, frozen-latent POSEIDON achieves error 0.11 with 512 trajectories, below FNO’s 0.282 but above fully finetuned POSEIDON-B’s 0.022.The frozen model qualitatively learns diffusion and smoothing from only a few samples, supporting transfer from fluid-dynamics representations to steady-state diffusion.Source: §D.4.3 · pp. 61–62, Figure 44 · p. 62

D.5 Results with DPOT · D.6 Further Ablations and Results

Across seven representative downstream tasks, POSEIDON substantially outperforms DPOT in efficiency and accuracy gains, while DPOT shows little benefit from scaling model size. The comparison also suggests POSEIDON leverages pretrained representations more effectively than DPOT.

D.5 Results with DPOT

POSEIDON-L achieves the highest mean AG score at 8.14, ahead of POSEIDON-B (6.5), DPOT-L (4.4), and DPOT-M (4.39) across seven tasks.POSEIDON is significantly better than DPOT on every task except SE-AF, where performance is comparable; similar conclusions hold for efficiency gain.Source: Table 10 · p. 63, §D.5 · p. 63
The DPOT comparison uses models pretrained on several fluid-dynamics equations and adapts them to the OLT task using the same finetuning procedure as MPP.DPOT-M and DPOT-L are selected to match POSEIDON-B and POSEIDON-L in scale, enabling a fair comparison.Source: §D.5 · p. 63
DPOT-M and DPOT-L obtain identical mean AG scores of 4.4 despite DPOT-L having 5× more parameters, indicating no observed downstream scaling benefit.The comparison uses DPOT-M (120 M parameters) and DPOT-L (509 M parameters) on seven representative downstream tasks.Source: §D.5 · p. 63
Trained-from-scratch DPOT-L reaches an average AG score of 3.53, only 25% below the DPOT-L foundation model.This strong stand-alone performance suggests DPOT relies substantially on the capacity of its underlying neural operator for downstream learning.Source: §D.5 · pp. 63–64
POSEIDON-L is almost 5 times more accurate than the underlying scOT neural operator, indicating stronger use of latent representations.The authors interpret this contrast with DPOT’s stand-alone performance as evidence that POSEIDON harnesses latent representations more effectively.Source: §D.5 · p. 64
Overall, the modified DPOT model outperforms CNO-FM and MPP but remains significantly inferior to POSEIDON on the evaluated downstream tasks.The authors note that the comparison merits further study.Source: §D.5 · p. 64

D.6.1 On all2all training · D.6.2 Direct. vs. Autoregressive Inference

All2all training improves operator-learning accuracy by exploiting every snapshot pair, but its quadratic pair growth motivates moderately dense sampling. Direct versus autoregressive inference is task-dependent, so Poseidon uses the better strategy for each downstream task.

D.6.1 On all2all training

All2all training significantly outperforms vanilla training for the NS-SL task with CNO models.The strategy uses every possible pair of snapshots per trajectory, exploiting the semi-group structure of time-dependent solution operators.Source: §D.6.1 · p. 65, Figure 45 · p. 66
All2all training has quadratic computational cost because training pairs grow quadratically with the number of trajectory snapshots.Selecting a subset reduces cost proportionately, but overly sparse sampling can increase overall error.Source: §D.6.1 · p. 65
Denser snapshot sampling consistently improves accuracy from T14 through T7 to T2, but T1 adds no further test-error reduction.The T14, T7, T2, and T1 subsets yield 3, 6, 36, and 120 training pairs per trajectory, respectively.Source: §D.6.1 · p. 65, Figure 46 · p. 66
Choosing T2 over T1 reduces training cost by a factor of 3.3 while preserving the observed accuracy gains.This motivates using a not-too-dense snapshot sampling strategy for pretraining and finetuning.Source: §D.6.1 · p. 65

D.6.2 Direct. vs. Autoregressive Inference

Poseidon, scOT, and FNO use homogeneous autoregressive rollouts, while CNO obtains a slight performance boost from heterogeneous rollout steps.Time-conditioned models can alternatively be evaluated directly at the requested time.Source: §D.6.2 · p. 65
Autoregressive rollout performs better for NS-PwC, whereas direct lead-time input performs better for Wave-Layer.The comparison uses homogeneous autoregressive rollouts and indicates that the preferred inference strategy depends on the dataset and dynamics.Source: §D.6.2 · p. 65
The models therefore select the best direct or autoregressive inference strategy separately for each task, as reported in Table 6.This choice reflects the task-dependent error behavior observed for NS-PwC and Wave-Layer.Source: §D.6.2 · p. 65

D.6.3 Error Growth over Time for POSEIDON-B … D.6.8 Histograms of Errors for Different Tasks

Across downstream evaluations, POSEIDON shows sub-linear error growth, strong extrapolation and parameter generalization, grid-resolution invariance, and robustness to input noise. Error distributions are also reported across all downstream tasks.

D.6.3 Error Growth over Time for POSEIDON-B

POSEIDON-B’s error grows sub-linearly over time without exponential blow-up, while its advantage over FNO increases toward the final time.For NS-PwC, autoregressive error accumulation is shown directly, and FNO has a larger error-growth rate than POSEIDON on NS-PwC and NS-SL.Source: §D.6.3 · p. 66, Figure 48 · p. 67, Figure 49 · p. 68

D.6.4 Out-of-distribution Time Extrapolation

POSEIDON-B significantly outperforms FNO at both T = 0.7 and extrapolated T = 1.0, with EG improving from approximately 20 to 30 under extrapolation.Both models perform worse outside the training time range, but POSEIDON-B degrades relatively less, attributed to longer time horizons during pretraining.Source: §D.6.4 · p. 67, Figure 50 · p. 68

D.6.5 Generalization of POSEIDON with respect to Changing PDE Parameters

POSEIDON generalizes through finetuning to coefficient-to-solution operators and altered PDE parameters, outperforming FNO for viscosity ν = 4 × 10^-3.For the changed-viscosity task, POSEIDON-B obtains EG = 925.5 and AG = 47.5, compared with EG = 1024 and AG = 19.7 on the original NS-PwC task.Source: §D.6.5 · pp. 68–69, Figure 51 · p. 69

D.6.6 POSEIDON Evaluated on Different Grids

POSEIDON-B’s test error is approximately invariant across the multiple grid resolutions evaluated for NS-PwC.The operator-transformer design enables evaluation on grids differing from the underlying computational grid by resampling inputs and outputs.Source: §D.6.6 · p. 69, Figure 52 · p. 70

D.6.7 Robustness of Poseidon with respect to Noise

POSEIDON-L remains robust to Gaussian input noise, with error not growing significantly even at a 3% noise-to-signal ratio on CE-RPUI.The evaluation compares clean inputs with NSRs of 0.1%, 1%, and 3%, using noise-free ground truth outputs.Source: §D.6.7 · p. 69, Figure 53 · p. 70

D.6.8 Histograms of Errors for Different Tasks

Error distributions are plotted across all downstream tasks using POSEIDON-B finetuned with 128 trajectories for time-dependent tasks and 512 for time-independent tasks.The histograms use kernel density estimates over the mean of all functions or quantities of interest.Source: §D.6.8 · p. 69, Figure 54 · p. 71

E Computational Resources

Experiments used consumer-grade GPUs with parallelized pretraining, while Poseidon inference was substantially faster than the PDE solvers used for data generation. Even POSEIDON-L required less than 10^-2 seconds per inference call.

Hardware and Training Setup

Pretraining used eight NVIDIA GeForce RTX 4090 GPUs in parallel, whereas finetuning and most scratch training ran on a single 24GB-VRAM GPU.Experiments used up to 128 CPU cores and 512GB RAM, primarily to avoid dataloading bottlenecks.Source: E Computational Resources · p. 72

Pretraining Times

Pretraining times are reported in Table 11 for eight NVIDIA GeForce RTX 4090 GPUs, with batch sizes provided in parentheses.The table gives approximate pretraining times for the reported configurations.Source: Table 11 · p. 72, E Computational Resources · p. 72

Inference Times

Inference time was less than 10^-2 seconds even for POSEIDON-L, yielding a 1–5-order-of-magnitude gain over the PDE solvers used for data generation.Solver runtimes ranged from 0.1 seconds to approximately 500 seconds at resolution 128^2, depending on the dataset and solver.Source: Table 12 · p. 72, E Computational Resources · p. 72

F Pretrained Models, Datasets, and Source Code · G Visualizations

Poseidon releases its source code through GitHub with Hugging Face integration, alongside pretrained models and datasets hosted on the Hugging Face Hub. The paper also visualizes random samples from its downstream PDE datasets, including fluid, wave, reaction, airfoil, Poisson, and Helmholtz tasks.

F Pretrained Models, Datasets, and Source Code

Poseidon’s source code is available on GitHub and tightly integrated with Hugging Face Transformers, using Hugging Face Accelerate for distributed training.The implementation makes heavy use of the Hugging Face ecosystem.Source: F Pretrained Models, Datasets, and… · p. 72
Pretrained models, pretraining datasets, downstream tasks, and PDEGYM datasets are available through collections on the Hugging Face Hub.The listed collections include Poseidon, Poseidon – Downstream Tasks, and PDEGYM.Source: F Pretrained Models, Datasets, and… · p. 72

G Visualizations

Random-sample visualizations cover Navier–Stokes and compressible-Euler datasets, including NS-Sines, NS-Gauss, CE-RP, CE-CRP, CE-KH, and CE-Gauss.The listed visualizations show examples from these fluid-dynamics task datasets.Source: Figure 55 · p. 73, Figure 56 · p. 74, Figure 57 · p. 75
Additional Navier–Stokes visualizations include NS-PwC, NS-BB, NS-SL, NS-SVS, NS-Tracer-PwC, and FNS-KF, with inputs such as velocity, tracer concentration, forcing, density, and pressure.The corresponding captions identify velocity inputs, optionally alongside tracer concentration, forcing term, density, and pressure.Source: Figure 61 · p. 79, G Visualizations · p. 79, Figure 62 · p. 80
The visualizations also include CE-RPUI, CE-RM, Wave-Gauss, and Wave-Layer, with inputs involving density, velocity, pressure, displacement, and propagation fields.The captions identify random samples for these compressible-Euler and wave datasets and specify their input fields.Source: Figure 67 · p. 85, Figure 68 · p. 86, Figure 69 · p. 87
Further random-sample visualizations cover ACE, SE-AF, Poisson-Gauss, and Helmholtz, using concentration, airfoil shape, source-term, and propagation-speed inputs.The corresponding captions identify these downstream PDE datasets and their stated inputs.Source: Figure 72 · p. 90, G Visualizations · pp. 90–91, Figure 73 · p. 91