PrithviWxC Foundation Model Validation on Weather Downscaling for Cross Domain Learning

AI foundation models hold considerable promise for leveraging the vast and diverse datasets available in atmospheric and geoscientific research. These models have the potential to advance scientific discovery by capturing complex spatial and temporal relationships inherent in earth system processes. However, the development and deployment of such models is often hindered by limited computational resources.

Accurate reconstruction of fine-scale atmospheric features from coarse-resolution data is a critical challenge in geoscientific modeling, as well as a benchmark for understanding the performance of climate-related models. High-resolution atmospheric data are essential for capturing localized phenomena, such as convective systems, topographic effects, and land-atmosphere interactions, that influence weather patterns and climate processes. However, the generation and storage of high-resolution datasets are computationally expensive, necessitating methods that can infer fine-scale structures from lower-resolution observations.

The primary objective of this study is to validate PrithviWxC [4], a ViT-based foundation model [1], for the task of downscaled image reconstruction. While Prithvi WxC was trained on 160 atmospheric variables from the Modern-Era Retrospective Analysis for Research and Applications, Version 2 (MERRA-2) dataset [2], we implement a smaller version of the initial model, with 21 million parameters, and pretrained on the same dataset with a set of six variables. 

The evaluation process involves the assessment of the model's capacity to reconstruct fine-scale features through the process of downscaling atmospheric data. In this procedure, inputs at a high spatial resolution of 1 km from [5] are first coarsened to 25 km resolution, the same as European Centre of Medium-range Weather Forecasts Reanalysis v.5 (ERA5) [3], and subsequently upscaled to recover the original fine-grained structure. This process serves as a benchmark for assessing the model's capacity to learn and preserve spatial details during resolution transformations, which is an essential requirement for geoscientific modeling tasks.

We fine-tune the model for downscaled image reconstruction on a set of 30784 128x128 patches, and validate its output, produced after learning on a limited temporal period, on tiles coming from the ERA5 dataset, which encompass all seasonality. In particular, we aim at highlighting the model's ability to generalize to data domains beyond its pretraining distribution, demonstrating its adaptability and the transferability of knowledge embedded within ViT architectures. By applying PrithviWxC to a knowledge domain that is distinct from its original training context, we demonstrate the potential for cross-domain learning in geoscientific applications.

 

REFERENCES

[1] Dosovitskiy, Alexey. "An image is worth 16x16 words: Transformers for image recognition at scale." arXiv preprint arXiv:2010.11929 (2020).

[2] Gelaro, Ronald, et al. "The modern-era retrospective analysis for research and applications, version 2 (MERRA-2)." Journal of climate 30.14 (2017): 5419-5454.

[3] Hersbach, Hans, et al. "The ERA5 global reanalysis." Quarterly Journal of the Royal Meteorological Society 146.730 (2020): 1999-2049.

[4] Schmude, Johannes, et al. "Prithvi wxc: Foundation model for weather and climate." arXiv preprint arXiv:2409.13598 (2024).

[5] Wedi, Nils P., et al. "A baseline for global weather and climate simulations at 1 km resolution." Journal of Advances in Modeling Earth Systems 12.11 (2020): e2020MS002192.