Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

PDE foundation models are skillful AI weather emulators for the Martian atmosphere

Johannes Schmude, Sujit Roy, Liping Wang, Theodore van Kessel, Levente Klein, Marcus Freitag, Eloisa Bentivegna, Robert Manson-Sawko, Bjorn Lutjens, Manil Maskey, Campbell Watson, Rahul Ramachandran, Juan Bernabe-Moreno

TL;DR

This work demonstrates that PDE foundation models pretrained on a diverse set of PDE solutions can be adapted to weather emulation for the Martian atmosphere, addressing data scarcity typical of planetary contexts. By extending a 2D PDE-FM (Poseidon) to 3D via axial vertical attention and preserving pretraining, the authors achieve a 34.4% improvement on a held-out Martian year, even with sparse initial conditions and modest compute (~13 GPU hours). The OpenMARS dataset provides Martian atmospheric fields used for training and evaluation, and the approach shows robust performance under data sparsity and across vertical levels, highlighting data-efficient pathways for real-world, data-limited atmospheric modeling. The results point to broader applicability of PDE-FMs to Earth regional weather in data-poor regimes and suggest architectural directions (e.g., multi-scale graph models) to further enhance performance when forcings and curved geometries are considered.

Abstract

We show that AI foundation models that are pretrained on numerical solutions to a diverse corpus of partial differential equations can be adapted and fine-tuned to obtain skillful predictive weather emulators for the Martian atmosphere. We base our work on the Poseidon PDE foundation model for two-dimensional systems. We develop a method to extend Poseidon from two to three dimensions while keeping the pretraining information. Moreover, we investigate the performance of the model in the presence of sparse initial conditions. Our results make use of four Martian years (approx.~34 GB) of training data and a median compute budget of 13 GPU hours. We find that the combination of pretraining and model extension yields a performance increase of 34.4\% on a held-out year. This shows that PDEs-FMs can not only approximate solutions to (other) PDEs but also anchor models for real-world problems with complex interactions that lack a sufficient amount of training data or a suitable compute budget.

PDE foundation models are skillful AI weather emulators for the Martian atmosphere

TL;DR

This work demonstrates that PDE foundation models pretrained on a diverse set of PDE solutions can be adapted to weather emulation for the Martian atmosphere, addressing data scarcity typical of planetary contexts. By extending a 2D PDE-FM (Poseidon) to 3D via axial vertical attention and preserving pretraining, the authors achieve a 34.4% improvement on a held-out Martian year, even with sparse initial conditions and modest compute (~13 GPU hours). The OpenMARS dataset provides Martian atmospheric fields used for training and evaluation, and the approach shows robust performance under data sparsity and across vertical levels, highlighting data-efficient pathways for real-world, data-limited atmospheric modeling. The results point to broader applicability of PDE-FMs to Earth regional weather in data-poor regimes and suggest architectural directions (e.g., multi-scale graph models) to further enhance performance when forcings and curved geometries are considered.

Abstract

We show that AI foundation models that are pretrained on numerical solutions to a diverse corpus of partial differential equations can be adapted and fine-tuned to obtain skillful predictive weather emulators for the Martian atmosphere. We base our work on the Poseidon PDE foundation model for two-dimensional systems. We develop a method to extend Poseidon from two to three dimensions while keeping the pretraining information. Moreover, we investigate the performance of the model in the presence of sparse initial conditions. Our results make use of four Martian years (approx.~34 GB) of training data and a median compute budget of 13 GPU hours. We find that the combination of pretraining and model extension yields a performance increase of 34.4\% on a held-out year. This shows that PDEs-FMs can not only approximate solutions to (other) PDEs but also anchor models for real-world problems with complex interactions that lack a sufficient amount of training data or a suitable compute budget.
Paper Structure (22 sections, 1 equation, 9 figures, 3 tables)

This paper contains 22 sections, 1 equation, 9 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Contributions
  3. Background
  4. AI models for atmospheric physics
  5. PDE FMs
  6. Poseidon
  7. Methods
  8. The OpenMARS dataset
  9. Data processing
  10. Training and fine-tuning
  11. Extension to three dimensions
  12. Sparsity
  13. Baselines
  14. Experiments
  15. Efficacy of multi-PDE pretraining
  16. ...and 7 more sections

Figures (9)

  • Figure 1: Extending scOT to three dimension
  • Figure 2: When extending the training to 40,000 steps, the randomly initialized model not only does not close the performance gap, but rather shows clear signs of overfitting. The differences between the 20,000 and 40,000 step runs stem from the learning rate scheduler.
  • Figure 3: Model performance after 20,000 training steps by level and masking ratio.
  • Figure 4: Impact of pretraining and dimensional extension on physical quantities.
  • Figure 5: Sample prediction for temperature at $\sigma=0.87007517$.
  • ...and 4 more figures