ESFM - A foundation model framework for heterogeneous data integration

With increased availability of high quality diverse weather data, including reanalysis, satellite, surface stations, climate model data, the amount of data-driven foundation models (FM) in the environmental field has increased significantly over the past years with forecasting performances matching and sometimes exceeding physics-based numerical model predictions.  However, most FMs are trained with one dataset or a few datasets with similar sampling and/or resolution properties. While the proposed models achieve remarkable results with the datasets and variables they are trained on; it would be hard to anticipate similar performance under partially missing observations across different dimensions at test time. Similarly, typical design considerations risk limiting usage of these FMs to other heterogeneous datasets concerning the broader Earth sciences community.

We propose Earth System Foundation Model (ESFM), an FM capable of handling heterogeneous observations (i) across different resolutions, (ii) with spatially gridded and non-gridded nature, and (iii) with little to extreme sparsity. We achieve this through simple architectural design considerations and a masked training protocol. Namely, we bin similar ranges of grid resolutions together, while optimizing a different set of tokenizers for significantly different resolution bins to accommodate a single FM for observations across different resolutions. Similarly, we tokenize non-gridded data (i.e., station) separately with a single pixel patch size. Finally, we use variable specific tokenizers, coupled with learnable missing observation tokens, that allow ESFM to naturally accommodate for various subsets of available variables across different spatiotemporal positions. 

In this exploratory study, we show that ESFM is a flexible FM that can achieve impressive forecasting performance under different adverse setups with missing test data across any dimension on ERA5; spatio-temporal and inter-variable. We further test forecasting performance of ESFM in very sparse satellite imagery (3% pixel occupancy) data as well as station data. 

The proposed framework; also compatible for different backbone architectures than the one we experimented with; provides a general approach for integrating diverse Earth system data sources with varying resolutions, sampling patterns, and availability. This makes ESFM particularly relevant for the broader environmental sciences and Earth and space sciences, where challenges related to data heterogeneity and missing observations are central to the development of next-generation data-driven environmental modeling systems.