[Paper Review] ClimaX: A foundation model for weather and climate

Tung Nguyen, J. Brandstetter+3 more|arXiv (Cornell University)|Jan 24, 2023

Meteorological Phenomena and Simulations168 citations

TL;DR

ClimaX is a flexible transformer-based foundation model pretrained on heterogeneous CMIP6 climate data, enabling fine-tuning for diverse weather and climate tasks with superior generalization and competitive benchmarks.

ABSTRACT

Most state-of-the-art approaches for weather and climate modeling are based on physics-informed numerical models of the atmosphere. These approaches aim to model the non-linear dynamics and complex interactions between multiple variables, which are challenging to approximate. Additionally, many such numerical models are computationally intensive, especially when modeling the atmospheric phenomenon at a fine-grained spatial and temporal resolution. Recent data-driven approaches based on machine learning instead aim to directly solve a downstream forecasting or projection task by learning a data-driven functional mapping using deep neural networks. However, these networks are trained using curated and homogeneous climate datasets for specific spatiotemporal tasks, and thus lack the generality of numerical models. We develop and demonstrate ClimaX, a flexible and generalizable deep learning model for weather and climate science that can be trained using heterogeneous datasets spanning different variables, spatio-temporal coverage, and physical groundings. ClimaX extends the Transformer architecture with novel encoding and aggregation blocks that allow effective use of available compute while maintaining general utility. ClimaX is pre-trained with a self-supervised learning objective on climate datasets derived from CMIP6. The pre-trained ClimaX can then be fine-tuned to address a breadth of climate and weather tasks, including those that involve atmospheric variables and spatio-temporal scales unseen during pretraining. Compared to existing data-driven baselines, we show that this generality in ClimaX results in superior performance on benchmarks for weather forecasting and climate projections, even when pretrained at lower resolutions and compute budgets. The source code is available at https://github.com/microsoft/ClimaX.

Motivation & Objective

• Motivate the creation of a general, data-driven foundation model for weather and climate to overcome task-specific limitations.
• Leverage heterogeneous, physics-informed climate datasets (CMIP6) for pretraining to enhance generality.
• Develop architectural innovations enabling variable modalities and irregular spatiotemporal coverage.
• Demonstrate fine-tuning versatility across forecasting, climate projection, and downscaling tasks.
• Assess scaling behavior with data, model size, and resolution to guide future research.

Proposed method

• Extend Vision Transformer (ViT) with two key modules: variable tokenization (tokenize each input variable separately) and variable aggregation (cross-attention at each spatial position to fuse variables).
• Add position and lead-time embeddings to guide temporal forecasting across arbitrary horizons.
• Pretrain ClimaX on CMIP6 data with a randomized forecasting objective to predict future states from current states across lead times 6h–168h.
• Use a latitude-weighted mean squared error to account for spherical area differences in global gridding.
• Fine-tune ClimaX to downstream tasks, including variables seen during pretraining and unseen variables with modular fine-tuning strategies.
• Evaluate on forecasting, climate projection, and downscaling benchmarks, comparing to neural baselines and NWP baselines.

Experimental results

Research questions

• RQ1 Can a single pretrained ClimaX model adapt to multiple weather and climate tasks with heterogeneous data?
• RQ2 How does ClimaX perform on global/regional forecasting, sub-seasonal to seasonal predictions, climate projection, and downscaling compared to specialized baselines?
• RQ3 What is the impact of dataset size, model capacity, and resolution on ClimaX’s performance (scaling laws)?
• RQ4 What finetuning strategies best transfer ClimaX to unseen variables and varying spatiotemporal resolutions?

Key findings

• A single pretrained ClimaX can be finetuned for a broad set of tasks across different lead times, resolutions, and regions.
• ClimaX achieves state-of-the-art performance on ClimateBench and is competitive with the operational IFS on WeatherBench, even with moderate compute.
• Scaling analyses show performance gains with more pretraining data, larger models, and higher resolution inputs.
• Ablation studies support effective finetuning strategies, including handling unseen variables and multiple tasks with shared components.
• The architecture maintains robustness to irregular data coverage and variable sets thanks to variable tokenization and aggregation.