Earth System Foundation Model (ESFM): A unified framework for heterogeneous data integration and forecasting
Firat Ozdemir, Yun Cheng, Salman Mohebi, Fanny Lehmann, Simon Adamov, Zhenyi Zhang, Leonardo Trentini, Dana Grund, Oliver Fuhrer, Torsten Hoefler, Siddhartha Mishra, Sebastian Schemm, Benedikt Soja, Mathieu Salzmann
Why It Matters
What makes this one worth your time
This work is relevant for AI researchers and engineers focused on climate science, as it provides a versatile tool for integrating and forecasting diverse environmental data, potentially improving predictions of extreme weather events.
ESFM offers a unified model for integrating diverse Earth system data and forecasting with improved accuracy.
Summary
The paper introduces the Earth System Foundation Model (ESFM), an open framework for integrating and forecasting heterogeneous Earth system data. It extends the 3D Swin UNet backbone to handle diverse datasets, including those with missing values, and introduces axial attention for capturing inter-variable dependencies. The model can predict variables in data-sparse regions and transform into a probabilistic model using adaptive layer norm-based ensembles. The ESFM demonstrates competitive performance on various datasets and accurately estimates extreme weather events.
Key contributions
- Introduction of a unified framework for integrating heterogeneous Earth system data.
- Implementation of axial attention for improved inter-variable dependency capture.
- Development of adaptive layer norm-based ensembles for probabilistic forecasting.
Notable insights
- Axial attention is used to capture inter-variable dependencies, enhancing prediction accuracy in data-sparse regions.
- Adaptive layer norm-based ensembles allow the deterministic model to be transformed into a probabilistic one, increasing its versatility.
Possible limitations
- Not stated in the abstract
Abstract
arXiv:2605.00850v1 Announce Type: cross Abstract: Foundation models (FMs) for the Earth system learn statistical relationships between physical variables across massive datasets to enable versatile downstream applications through finetuning, separating them from task-specific weather models. Here, we introduce Earth System Foundation Model (ESFM), a fully open model building on the 3D Swin UNet backbone of the pioneering Aurora model. ESFM introduces extensions that increase functionality and foster adoption in climate sciences. First, the encoding scheme and training protocols have been extended to handle diverse datasets, including those containing missing values across all spatio-temporal dimensions such as satellite data, as well as station data, all under one backbone. Axial attention is introduced to capture inter-variable dependencies. As a result ESFM skillfully predicts variables in regions or on pressure levels where no data is present at the initial time, while preserving inter-variable relationships, for example between temperature, pressure, and humidity. Individual variable tokenization enables different sets of variables to be shuffled during training and simplifies the process of building extensions for new downstream tasks. Adaptive layer norm-based ensembles allow for a simple yet effective way to transform deterministic ESFM to a probabilistic FM. We present findings using dense gridded data (ERA5, CMIP6), regionally masked dense data, sparse gridded MODIS satellite data, and station data. Results demonstrate competitive or superior performance relative to state-of-the-art benchmarks. Case studies of Super Typhoon Doksuri (2023) and 2024 sudden stratospheric warming events show accurate positional and magnitude estimations of extreme weather. ESFM retains the strengths of previous foundation models, such as long-term stability, but facilitates application to a variety of downstream tasks.