PaperlayerView exact revision on arXiv

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.19101v2cs.LG

TL;DR

PDE operator learning is limited by expensive task-specific data and the difficulty of generalizing across diverse, unrelated equations. Poseidon addresses this with a multiscale, time-conditioned operator transformer and semi-group-based all2all pretraining on fluid-dynamics operators. Across 15 out-of-distribution tasks, it outperforms baselines in accuracy and sample efficiency, including tasks involving unseen physics, while its scope remains limited by training and evaluation on only a few PDEs and data distributions.

  • Problem

    Existing PDE operator-learning methods require many costly task-specific examples, while it remains unclear whether pretraining on few PDEs can generalize to unseen and unrelated physics.

  • Method

    Poseidon combines a multiscale operator transformer with time conditioning and semi-group-based all2all training, pretrained on diverse compressible Euler and incompressible Navier–Stokes operators.

  • Results

    Across 15 out-of-distribution downstream tasks, Poseidon outperforms FNO and other baselines in accuracy and sample efficiency, requiring a median 20 samples versus FNO’s 1024 for comparable error.

  • Takeaways & Limitations

    Poseidon demonstrates that a foundation model pretrained on a small set of fluid-dynamics PDEs can generalize to unseen and unrelated physical processes with few task-specific examples.

  • Takeaways & Limitations

    Poseidon was trained and evaluated on only a few PDEs and data distributions, with further diversity, elliptic equations, and broader time scales identified as important extensions.

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

Existing PDE operator-learning methods require expensive task-specific datasets, motivating a foundation-model approach that learns reusable representations across diverse PDEs. Poseidon addresses this challenge by pretraining on fluid-dynamics operators and evaluating transfer to diverse, unseen downstream physics.

  • Existing operator-learning methods can require many training examples, making PDE-specific datasets expensive to generate through simulation or physical measurement.
  • Foundation models offer a model for reducing task-specific data requirements by pretraining on diverse distributions and finetuning with few downstream samples.
  • Designing PDE foundation models is difficult because PDEs span diverse types, data distributions, spatio-temporal scales, and limited public datasets.
  • Poseidon combines a multiscale operator transformer, all2all trajectory training, and large-scale fluid-dynamics pretraining, then tests 15 out-of-distribution downstream tasks.
  • Poseidon outperforms downstream baselines in accuracy and sample efficiency, including gains on PDEs and physical processes absent from pretraining.

2 Approach

Poseidon approximates PDE solution operators with a time-conditioned, multiscale transformer and trains on trajectory pairs generated through the semi-group property. Its pretraining framework spans multiple PDEs and data distributions while supporting continuous-time inference.

  • Problem Formulation: The operator-learning task approximates S so that an initial datum produces the entire PDE solution trajectory over time.
  • Model Architecture: Poseidon’s scOT backbone processes lead time and function-valued initial data with a hierarchical multiscale vision transformer to approximate S(t, a).
  • Model Architecture: Patch-embedded inputs pass through shifted-window SwinV2 blocks, whose local attention improves efficiency while successive shifts provide domain-wide interaction.
  • Model Architecture: Lead-time-conditioned layer normalization enables continuous-in-time evaluations.
  • Model Architecture: The U-Net-style encoder-decoder arranges transformer blocks hierarchically with patch merging, patch expansion, and same-scale ConvNeXt connections.
  • Training and Inference: All2all training uses the semi-group property to include every ordered pair of states within a trajectory, increasing usable samples from O(K) to O(K^2) per trajectory.
  • Training and Inference: The trained operator can evaluate arbitrary initial conditions and times directly, or approximate a target time through autoregressive rollouts.
  • Pretraining: Pretraining indexes PDE types and data distributions, using trajectories from selected fluid-dynamics operators and shared input channels across operators.

3 Experiments

Experiments evaluate POSEIDON on 15 diverse out-of-distribution PDE tasks against neural-operator and foundation-model baselines. POSEIDON consistently improves accuracy and sample efficiency, including on unseen physics, while benefiting from larger models and pretraining datasets.

  • Downstream Tasks: The 15 downstream tasks span diverse PDE types, physical processes, spatio-temporal scales, and out-of-distribution settings.Six tasks use familiar Euler or Navier–Stokes equations with different data distributions, while nine involve PDEs not seen during pretraining.
  • Downstream Performance: POSEIDON requires only 20 samples on average to match FNO’s error with 1024 samples, and needs an order of magnitude fewer samples on 13 of 15 tasks.Four tasks require only 3 task-specific samples to reach FNO’s error with 1024 samples.
  • Baseline Comparisons: POSEIDON exceeds trained-from-scratch neural operators by a median order-of-magnitude gain in sample efficiency and an average fourfold gain in accuracy.CNO and scOT are comparable to each other and both generally outperform FNO, but POSEIDON performs substantially better than both.
  • Generalization: POSEIDON generalizes to unseen physics and remains effective when its latent space is frozen during finetuning.It performs best on 8 of 9 tasks involving unseen PDEs, while frozen-latent experiments indicate that pretraining learned rich physical representations.
  • Foundation-Model Baselines: POSEIDON outperforms CNO-FM on 14 of 15 tasks despite identical pretraining data, implicating the multiscale transformer architecture alongside the dataset.CNO-FM uses multiscale CNNs, whereas POSEIDON uses a multiscale vision transformer.
  • Scaling: POSEIDON’s advantage is not explained by model size alone: the 21M-parameter POSEIDON-T outperforms larger and comparable baselines, while larger POSEIDON models improve with scale.Training and validation errors decrease with increasing POSEIDON model size, and POSEIDON-L consistently outperforms POSEIDON-B downstream.
  • Scaling: Increasing pretraining dataset size lowers pretrained-model losses and generally improves downstream performance, while dataset diversity also affects results.The study compares random trajectory removal with removing all trajectories from three of six operators.

4 Discussion

The discussion presents POSEIDON as a broadly capable PDE foundation-model family whose architecture, pretraining, and training strategy support strong out-of-distribution performance. Its current scope remains limited by the small range of PDEs, geometries, and data distributions used for training and evaluation.

  • Contributions: POSEIDON combines a multiscale transformer, lead-time conditioning, all2all training, and fluid-dynamics pretraining to learn PDE solution trajectories.Lead-time conditioning supports continuous-in-time evaluation, while all2all training exploits the semi-group structure of time-dependent PDE solutions.
  • Main Findings: Across 15 out-of-distribution tasks, POSEIDON is the best-performing model on 14 and reaches FNO’s error with a median of 50 times fewer task-specific samples.The evaluation covers varied PDEs and data distributions, including physics not present during pretraining.
  • Related Work: Prior PDE foundation models are narrower in scope, covering specific elliptic equations, small one-dimensional datasets, or focused incompressible-flow settings.The discussion contrasts these models with POSEIDON’s broader evaluation and pretraining scope.
  • Limitations: POSEIDON was trained and evaluated on only a few PDEs and data distributions, especially lacking time-independent and non-Cartesian pretraining examples.The authors identify broader PDE datasets, more time scales, and more general geometries or boundary conditions as important extensions.

A Architecture of the scalable Operator Transformer (scOT)

scOT is a multiscale encoder-decoder operator transformer that converts discretized function inputs into outputs through patching, embedding, SwinV2 processing, and reconstruction. Its U-Net-style design combines hierarchical resolution changes, skip connections, windowed attention, and time-conditioned normalization.

  • Input representation: scOT begins by partitioning discretized input functions into non-overlapping patches and linearly embedding them into a shared latent representation.The embedding is a discretization of the corresponding continuous operator and is followed by time-conditioned normalization.
  • Transformer processing: The embedded tokens pass through multiple SwinV2 stages with windowed multi-head self-attention and shifted windows that progressively connect spatial regions.Relative logarithmic position encodings are used, and shifted windows enable attention across points over successive blocks.
  • Encoder-decoder connections: Encoder levels connect to the decoder through ConvNeXt blocks, while the bottleneck is convolution-free in the U-Net-style architecture.The ConvNeXt blocks operate on reshaped two-dimensional token grids and include time-conditioned normalization and residual connections.
  • Output reconstruction: After the final decoder stage, patch recovery linearly maps latent tokens back into discretized output patches, which are assembled through mixup.This completes the transformation from tokenized function inputs to function outputs.
  • Time conditioning: For time-dependent PDEs, scOT replaces ordinary layer normalization with time-conditioned normalization whose learnable gain and bias depend on continuous time t.This design is used in Poseidon and supports conditioning the token transformations on t.
  • Multiscale hierarchy: The encoder applies patch merging after stages to halve resolution while increasing latent dimension, while the decoder uses patch expansion to restore resolution.Patch merging includes an additional residual connection; patch expansion halves the latent dimension while doubling resolution.

B.1 Pretraining Datasets

The pretraining datasets comprise diverse incompressible Navier–Stokes and compressible Euler solution operators, using varied initial conditions and fluid-dynamics phenomena. Each dataset generates simulated trajectories for learning time-dependent PDE evolution.

  • Compressible Euler datasets: Compressible Euler trajectories are simulated to time T = 1 using a high-resolution finite-volume solver with WENO reconstructions and HLLC Riemann solvers.The pretraining target contains density, horizontal velocity, vertical velocity, and pressure snapshots.
  • Navier–Stokes datasets: Navier–Stokes datasets use either trigonometric initial conditions or vorticity fields formed from many Gaussian components.NS-Sines uses 10 modes, whereas NS-Gauss uses 100 Gaussian components and recovers velocity from vorticity.
  • Compressible Euler datasets: The four-quadrant Riemann problem samples piecewise-constant density, velocity, and pressure over a 2 × 2 partition of the domain.The stochastic initial values are sampled independently over specified uniform ranges.
  • Compressible Euler datasets: CE-CRP extends the four-quadrant Riemann problem to curved, multip partitioned subdomains while retaining stochastic fluid-state initialization.Its initial conditions use uniformly sampled density, velocity, and pressure values.

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

The ablation models use POSEIDON-B training settings while comparing POSEIDON with CNO and FNO operator-learning baselines. CNO is a modified U-Net over bandlimited functions, whereas FNO uses Fourier layers; both incorporate time information for time-dependent problems.

  • Dataset ablations: The dataset-ablation models retain POSEIDON-B’s pretraining and finetuning strategies, with variants trained on reduced or less diverse pretraining subsets.The less diverse variant excludes NS-Sines, CE-CRP, and CE-Gauss; half- and eighth-size variants use the first portions of each dataset subset.
  • Dataset ablations: A scOT model trained from scratch on each downstream task provides an additional comparison against POSEIDON and other baselines.
  • CNO: CNO maps bandlimited functions through a modified Operator U-Net with lifting, multiscale layers, activation operators, skip connections, and projection.Its reported configuration has 39.1M trainable parameters and four up/downsampling layers.
  • CNO: CNO uses lead-time-conditioned instance normalization, with MLPs parameterizing α(t) and β(t), and also receives time as a constant input channel.
  • FNO: FNO combines lifting, Fourier layers, residual connections, and nonlinear projection, using truncated Fourier coefficients and time as an additional channel for time-dependent problems.The trained FNO configuration has 37.0M parameters, five Fourier layers, and 20 Fourier modes.

D.1 Performance on Downstream Tasks

Across downstream tasks, POSEIDON models are evaluated using median relative L1 error and generally achieve lower errors than baselines at the same data volume. Their performance follows consistent, sometimes biphasic, scaling behavior and improves with model size.

  • Scaling plots: POSEIDON models clearly outperform baselines on most downstream tasks, achieving significantly lower test errors for the same number of trajectories or samples.The scaling plots compare POSEIDON variants with CNO-FM, MPP, FNO, CNO, and scOT.
  • Scaling laws: All considered models generally obey power-law scaling with different exponents across problems, while POSEIDON foundation models show consistent scaling laws.The fits use the number of trajectories or samples M and a model-specific scaling factor and exponent.
  • Scaling laws: Poisson-Gauss exhibits biphasic scaling for POSEIDON-B and POSEIDON-L, with a warmup phase followed by faster learning after M_pt=32.For POSEIDON-B, αw=0.23 and αℓ=0.99; for POSEIDON-L, αw=0.33 and αℓ=0.94.
  • Aggregate performance: Aggregate efficiency-gain and accuracy-gain statistics quantify POSEIDON’s advantage over baselines across all downstream tasks.
  • Model scaling: Increasing POSEIDON model size consistently decreases both pretraining training loss and evaluation loss.Figure 22 examines losses through epoch 20 and scaling at epoch 20.

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

POSEIDON-B scales with the size of its pretraining dataset, and comparisons examine both reduced data volume and reduced diversity. These ablations are evaluated across the downstream task suite.

  • Pretraining dataset size: POSEIDON-B’s pretraining training and evaluation losses improve as the dataset grows from one-eighth to one-half and full size.
  • Pretraining dataset size: Downstream comparisons contrast POSEIDON-B pretrained on the full dataset with the same model pretrained on one-eighth of it.Figures 24 to 38 report these comparisons across the task suite.
  • Pretraining dataset quality: Additional comparisons use POSEIDON-B pretrained on half the dataset versus the same model pretrained on a less diverse dataset.
  • Downstream coverage: The dataset-size and diversity experiments extend across NS, CE, SE, GCE, Wave, ACE, Poisson-Gauss, and Helmholtz tasks.

D.4.1 CE-RPUI

Case studies examine how POSEIDON transfers representations to difficult or unseen PDE operators. The analyses connect few-shot accuracy with reusable pretrained structures and rapid adaptation phases.

  • CE-RPUI: On CE-RPUI, POSEIDON-B with 128 trajectories captures shocks and roll-up vortices more sharply than CNO and FNO trained from scratch with the same data.The task is out-of-distribution because its data distribution was not seen during pretraining.
  • CE-RPUI: The CE-RPUI analysis suggests that POSEIDON combines information associated with shock propagation, large vortex roll-ups, and multiscale vortex dynamics.
  • CE-RPUI: This case study attributes few-shot CE-RPUI accuracy to representations learned from multiple pretraining operators and reused during finetuning.
  • ACE: On the unseen reaction-diffusion Allen-Cahn task, POSEIDON learns the solution operator from few examples despite pretraining on convection-dominated fluid equations.
  • ACE: With one task-specific trajectory, POSEIDON-B approximately learns Allen-Cahn’s large-scale solution features, including front propagation.
  • ACE: A frozen-latent POSEIDON model reaches error 0.031 with 32 trajectories, compared with FNO’s 0.037 using 128 trajectories, while full finetuning reaches 0.014.
  • Poisson-Gauss: For Poisson-Gauss, POSEIDON adapts to time-independent diffusion, smoothing, and Dirichlet boundaries absent from the pretraining setting.
  • Poisson-Gauss: Poisson-Gauss adaptation shows an initial warmup followed by faster learning, consistent with the observed biphasic power scaling.The model takes about 128 samples to learn spreading and smoothing, with further local adjustments by 512 samples.

D.5 Results with DPOT

POSEIDON substantially outperforms DPOT on representative downstream tasks, while DPOT shows little benefit from increasing model size. The comparison also suggests stronger use of pretrained representations by POSEIDON.

  • Overall comparison: POSEIDON outperforms corresponding DPOT models on all evaluated tasks except SE-AF, where performance is comparable.Across seven tasks, POSEIDON shows gains in both efficiency and accuracy.
  • Overall comparison: 8.14, 6.5, 4.39, and 4.4 are the mean AG scores for POSEIDON-L, POSEIDON-B, DPOT-M, and DPOT-L, respectively.POSEIDON-L is almost twice as accurate as both DPOT models under this comparison.
  • Scaling: A fivefold parameter increase from DPOT-M to DPOT-L does not noticeably improve downstream mean AG, which remains 4.4 for both models.The reported comparison covers seven representative downstream tasks.
  • Standalone baseline: The trained-from-scratch DPOT-L neural operator reaches an average AG score of 3.53, only 25% below the DPOT-L foundation model.This indicates strong standalone performance for the underlying DPOT neural operator on the evaluated dataset.
  • Representation use: POSEIDON-L is almost five times more accurate than the underlying scOT neural operator, suggesting stronger use of latent representations.The authors interpret the results as indicating that DPOT relies more on the downstream learning capacity of its underlying operator.
  • Overall comparison: Overall, the modified DPOT performs better than CNO-FM and MPP but remains significantly inferior to POSEIDON, although the comparison merits further study.The authors also highlight DPOT’s lack of model-size scaling and questions about latent-representation use.

D.6.1 On all2all training

All2all training uses the semi-group structure of time-dependent PDE solution operators to create training pairs from trajectory snapshots. Denser snapshot sampling improves accuracy until additional samples become statistically less informative, while inference strategy remains task-dependent.

  • All2all strategy: All2all training uses every possible pair of snapshots within each trajectory, scaling the training data through the solution operator’s semi-group structure.The strategy applies to time-dependent operator-learning algorithms, not only POSEIDON.
  • Cost trade-off: All2all training cost grows quadratically with the number of trajectory snapshots, creating a trade-off between computational cost and potential error reduction.Selecting a snapshot subset reduces cost proportionately but may increase overall error if too few snapshots are retained.
  • Snapshot subsets: 3, 6, 36, and 120 training pairs per trajectory result from the T14, T7, T2, and T1 snapshot subsets, respectively.These subsets progressively include denser time sampling for the NS-PwC task.
  • Accuracy and saturation: Denser snapshot sampling consistently improves test accuracy from T14 through T7 to T2, but adding the still-denser T1 subset yields no further decrease in test error.The saturation is attributed to newly added snapshots being statistically less significant.
  • Inference strategy: Autoregressive rollout performs better on NS-PwC, whereas direct lead-time input performs better on Wave-Layer, so the preferred strategy depends on task dynamics.The authors select the best strategy separately for each task.
  • CNO comparison: All2all training improves testing errors for CNO models compared with vanilla training on NS-SL.

D.6.3 Error Growth over Time for POSEIDON-B

POSEIDON-B’s rollout error grows sub-linearly over time, while its advantage over FNO increases toward the final time. The model also retains strong performance under time extrapolation, different resolutions, coefficient shifts, and input noise.

  • Error growth: POSEIDON-B rollout error grows sub-linearly over time on NS-PwC, with no observed exponential error blow-up.Autoregressive inference is advantageous only when per-step accumulated error remains below direct lead-time error.
  • Comparison with FNO: POSEIDON-B’s error advantage over FNO increases over time on NS-PwC and NS-SL because FNO has a larger error-growth rate.The difference is highest at the final time.
  • Time extrapolation: POSEIDON-B and FNO both extrapolate worse beyond their training interval, but POSEIDON-B significantly outperforms FNO at T = 0.7 and T = 1.0.POSEIDON’s EG improves from approximately 20 at T = 0.7 to approximately 30 at T = 1.0.
  • Operator generalization: POSEIDON generalizes to downstream operators mapping PDE coefficients to solutions, despite pretraining operators mapping initial data to later-time solutions.The authors report generalization to different PDE parameters through finetuning.
  • Resolution robustness: POSEIDON-B’s test error is approximately invariant across the evaluated grid resolutions on NS-PwC.Inputs and outputs are resampled to the model’s computational grid before and after processing.
  • Noise robustness: POSEIDON-L’s error does not grow significantly under Gaussian input noise up to a 3% noise-to-signal ratio on CE-RPUI.Errors are measured against clean ground-truth outputs for NSRs of 0.1%, 1%, and 3%.

E Computational Resources

The experiments used substantial CPU and memory resources alongside consumer-grade GPUs, while inference remained much faster than the PDE solvers used to generate data. Reported inference timings are approximate proxies.

  • Hardware: Experiments used 8–128 CPU cores, up to 512GB of RAM, and consumer-grade GPUs with 24GB of VRAM.Pretraining used the most CPU cores and RAM and was performed in parallel on eight NVIDIA GeForce RTX 4090 GPUs.
  • Solver comparison: At 128^2 resolution, the PDE solvers used for data generation range from 0.1 seconds for an optimized GPU solver to 10 seconds for a FENICS FEM solver.These solver runtimes are contrasted with the much shorter reported model inference time.
  • Measurement caveat: Inference times in Table 12 are approximate proxies because they were not measured in a controlled environment.

F Pretrained Models, Datasets, and Source Code

The section documents publicly available Poseidon code, pretrained models, datasets, and visualizations spanning the evaluated PDE tasks. The figures show random samples with task-specific inputs, including velocity, displacement, concentration, source, forcing, density, pressure, and propagation fields.

  • Source Code: Poseidon source code is available on GitHub and integrates with Huggingface Transformers and Accelerate for distributed training.
  • Pretrained Models and Datasets: Pretrained Poseidon models and datasets are available through the Huggingface Hub, including downstream-task and PDEGYM collections.
  • Visualizations: The visualizations cover random samples from Navier–Stokes, convection–diffusion, wave, acoustic, Poisson, and Helmholtz task families.
  • Visualizations: Task inputs include velocity fields, displacement and propagation fields, concentration, airfoil shape, source or forcing terms, and fluid density and pressure.
  • Visualizations: Additional Navier–Stokes visualizations include tracer concentration, while fluid task inputs also include forcing terms and density, velocity, and pressure.
  • Visualizations: The collection also includes repeated convection–diffusion, wave, and Helmholtz visualizations with inputs such as displacement, propagation field, concentration, and propagation speed.
Loading 2405.19101v2…