PaperlayerView exact revision on arXiv

Source-linked AI summary

ClimaX: A foundation model for weather and climate

Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K. Gupta, Aditya Grover

arXiv:2301.10343v5cs.LGcs.AI

TL;DR

Weather and climate modeling needs models that are both computationally practical and general across heterogeneous data and tasks. ClimaX addresses this need by pretraining a Transformer-based foundation model on CMIP6-derived simulations and fine-tuning it across weather and climate applications. It achieves strong benchmark performance and generalizes across operating conditions, while the authors identify broader variables, resolutions, and future scenarios as directions for further work.

  • Problem

    Existing numerical models are computationally intensive, while data-driven models trained for specific tasks and homogeneous datasets lack generality across weather and climate applications.

  • Method

    ClimaX pretrains a Transformer-based foundation model with a self-supervised forecasting objective on heterogeneous CMIP6-derived climate simulations, then fine-tunes it for diverse tasks.

  • Results

    ClimaX supports one pretrained model across forecasting, projection, and downscaling tasks, achieving state-of-the-art ClimateBench results and competitive WeatherBench performance with at most 80 NVIDIA V100 GPUs.

  • Takeaways & Limitations

    ClimaX demonstrates that heterogeneous climate simulations can support a general-purpose data-driven model across varied spatiotemporal resolutions, regions, and target variables.

  • Takeaways & Limitations

    The study restricts pretraining to standard benchmarked variables and excludes future-forcing scenarios, leaving broader variables, higher resolutions, and future scenarios for future work.

Abstract

from arXiv · show

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.

1. Introduction

Weather and climate modeling needs accurate, efficient forecasts and projections, but numerical models are computationally intensive while existing data-driven models are often task-specific. ClimaX addresses this gap with heterogeneous-data pretraining and a generalizable architecture that transfers across diverse tasks and conditions.

  • Motivation: Numerical weather and climate models face challenges representing fine-resolution physical processes and initial conditions while remaining computationally efficient.These limitations affect rapid short-term simulation and accurate forecasts beyond 5–7 days.
  • Motivation: Existing data-driven networks are commonly trained on curated, homogeneous datasets for specific predictive tasks, limiting their general-purpose utility.Their deployment can involve scenarios such as extreme-weather forecasting and longer-term climate projection.
  • Approach: ClimaX is designed as a weather-and-climate foundation model that can be efficiently adapted to general-purpose atmospheric tasks.The work explicitly studies how to design and train such a model.
  • Approach: ClimaX uses heterogeneous climate simulations, including CMIP6, whose diversity provides rich and plentiful pretraining data.CMIP6 includes simulations from approximately 100 climate models across 49 modeling groups.
  • Approach: Its architecture is designed to handle climate data with many variables, datatypes, and irregular spatiotemporal coverage.The approach repurposes the Vision Transformer and addresses variable heterogeneity through dedicated modeling choices.
  • Approach: A randomized forecasting objective trains ClimaX to predict arbitrary input variables at arbitrary future times, supporting downstream fine-tuning.The pretraining design targets relationships among atmospheric variables and transfer across tasks.
  • Results: A single pretrained ClimaX model can be fine-tuned across weather forecasting, climate projection, and downscaling tasks under varied operating conditions.The evaluated settings vary in resolution, geography, target variables, and include variables unseen during training.
  • Results: ClimaX achieves state-of-the-art ClimateBench results and competitive WeatherBench performance against operational IFS using at most 80 NVIDIA V100 GPUs.These results hold despite training at moderate resolutions and compute budgets.

2. Background and Related Work

Weather and climate modeling spans numerical simulation, data assimilation, forecasting, projection, and downscaling across distinct spatial and temporal scales. Related work shows both the strengths of established numerical systems and the emerging role of machine learning benchmarks and data sources.

  • Numerical Modeling: Numerical weather and climate systems use related primitive equations, with weather models emphasizing forecasts and climate models emphasizing long-term change.Earth system models additionally couple atmosphere, cryosphere, land, and ocean processes.
  • Numerical Modeling: Data assimilation combines atmospheric and ocean observations with numerical models to improve forecasts and simulations.This is an important component of numerical weather prediction systems.
  • Numerical Modeling: Climate models are generally global and coarser, whereas weather models range from global to local and regional settings with higher spatiotemporal resolutions.The different resolutions reflect distinct modeling and prediction requirements.
  • Numerical Modeling: GCMs remain successful but face debated limitations including structural errors across models and their design around reproducing observed climate.CMIP was created as a standardized protocol for evaluating and comparing climate models.
  • Data Sources: Weather and climate data combine diverse sources, including reanalysis that blends observations with past short-range forecasts through data assimilation.Measurements span heterogeneous physical variables and data types.
  • Data Sources: CMIP6 provides accessible simulations from 49 groups covering many climate variables, models, forcings, and climate scenarios from 1850 onward.The archive follows similar governing equations while varying model forcings such as greenhouse gas emissions.
  • Data Sources: ERA5 combines ECMWF’s Integrated Forecasting System with observations to estimate atmospheric, ocean-wave, and land-surface quantities from 1950 onward.It is a predominant source for training and benchmarking weather forecasting systems.
  • Tasks: Machine-learning weather tasks include global, regional, and sub-seasonal-to-seasonal forecasting across different spatial and temporal horizons.Operational IFS is a state-of-the-art numerical baseline for global forecasting, while S2S spans 2 weeks to 2 months.

3. Approach

ClimaX is designed as a generalizable foundation model for heterogeneous weather and climate data, supporting varied variables and spatial resolutions. Its approach combines variable-aware ViT processing, CMIP6 pretraining, and fine-tuning across downstream tasks.

  • Design goal: The model is designed to pretrain on heterogeneous data sources and then fine-tune for diverse weather and climate tasks.The intended coverage includes differing variables, spatial resolutions, geographical regions, and physical groundings.
  • Problem setting: ClimaX targets gridded prediction tasks mapping V × H × W inputs to V′ × H′ × W′ outputs across Earth-system applications.Variables may represent atmospheric conditions or climate forcing factors, while spatial dimensions encode grid resolution.
  • Model architecture: ClimaX extends Vision Transformers with variable tokenization and variable aggregation to handle flexible climate-data inputs.These components address variable heterogeneity and reduce the sequence length created by independently tokenizing variables.
  • Model architecture: Variable tokenization independently divides each input variable into spatial patches, producing V × h × w embedded tokens of dimension D.Each variable map has shape H × W, and each patch has spatial size p^2 before linear embedding.
  • Model architecture: Variable aggregation applies position-wise cross-attention over the V variable embeddings, reducing the sequence from V × h × w to h × w.A learnable query produces one aggregated vector for each spatial position.
  • Data and evaluation: The model uses CMIP6 data from five climate-model sources regridded to 5.625° and 1.40625° resolutions, then evaluates fine-tuned models on ERA5-based tasks.The reported downstream evaluation spans forecasting and other weather-related tasks under computational constraints.
  • Pretraining: ClimaX is pretrained on CMIP6 forecasting snapshots with randomized lead times from 6 to 168 hours and a latitude-weighted mean squared error objective.The model predicts future weather conditions from current conditions and adds lead-time embeddings to the tokens.

4. Experiments

ClimaX is evaluated across global, regional, sub-seasonal, climate, and downscaling tasks, with analyses of resolution, scaling, finetuning protocols, and computational trade-offs. Across these settings, pretraining supports strong performance and generalization beyond the original training conditions.

  • Experimental scope: Experiments span global and regional forecasting, sub-seasonal prediction, climate projection, climate downscaling, scaling laws, and finetuning trade-offs.The study compares ClimaX with operational and task-specific baselines across diverse variables, regions, lead times, and resolutions.
  • Global forecasting: At higher resolution, ClimaX closely matches IFS at short horizons and outperforms it at 7 days and beyond, while low-resolution performance is comparable or slightly better mainly at longer horizons.At low resolution, IFS leads from 6 hours to 5 days; at higher resolution, ClimaX is superior from 7 days onward, with similar RMSE and ACC trends.
  • Regional forecasting: ClimaX is the best-performing method across North American target variables and lead times, while interpolating positional embeddings preserves strong performance for a lower-resolution pretrained model.The interpolated model ranks second for Z500 and T850 and is beaten only by ResNet for T2m at a 3-day lead time.
  • Sub-seasonal prediction: ClimaX achieves the lowest RMSE for all four S2S target variables, with a larger advantage over UNet at increasing lead times and gains over scratch-trained Cli-ViT.The evaluation covers T850, T2m, U10, and V10 at sub-seasonal lead times.
  • Climate projection: On ClimateBench, frozen ClimaX attention performs best for two temperature-related variables, while scarce downstream data can make further attention-layer finetuning overfit.For precipitation tasks, the frozen model slightly trails the ClimateBench baseline on NRMSE metrics but outperforms it on RMSE.
  • Scaling and finetuning: Larger ClimaX models scale more consistently with additional pretraining data and are more data-efficient, while higher-resolution pretraining generally improves forecasting except for two short T2m horizons.ClimaX at 1.40625° uses a larger patch size because of memory constraints but achieves lower RMSE and higher ACC on almost all tasks.

5. Discussion and Future Work

ClimaX extends foundation-model scaling ideas to weather and climate, with opportunities spanning datasets, resolutions, and broader Earth-system applications. The discussion also identifies higher-resolution training and multimodal sequence modeling as future directions.

  • Discussion: ClimaX brings foundation-model scaling and generality to data-driven weather and climate modeling.Unlike numerical models that primarily scale computation and task-specific data-driven models, ClimaX is presented as enabling broader scaling and applicability.
  • Future work: Future work could combine observational and simulated datasets with more climate variables, higher spatiotemporal resolutions, and future scenarios.The authors avoided future-forcing simulations to prevent leakage in climate projection tasks.
  • Limitations and scaling: Training ClimaX at low to moderate resolutions was driven by compute restrictions, while empirical trends suggest 0.25° resolution may improve results.The passage frames resolution as important for scaling performance.
  • Future work: Better sequence-modeling architectures are needed for climate inputs with many channels and substantial resolution heterogeneity.The authors highlight multimodal spatiotemporal architectures as a promising direction.
  • Future applications: The pretrained ClimaX backbone could be explored for extreme-weather prediction, climate-change attribution, and domains such as agriculture.These applications extend beyond the tasks evaluated in the work.

A.1.1. Implementation details

ClimaX handles variable-dependent inputs and outputs through token embeddings, aggregation, attention, and a flexible prediction head. Its default configuration is documented alongside implementation details for the ClimaX, ResNet, and UNet models.

  • ClimaX architecture: ClimaX accepts V × H × W inputs and produces V′ × H × W outputs, allowing input and output variable counts to vary across datasets.A set of possible variables provides separate token-embedding layers for pretraining and fine-tuning.
  • ClimaX architecture: Variable tokenization converts each variable into spatial tokens, applies a variable-specific convolutional embedding, and produces V × h × w × D representations.The convolution uses one input channel, D output channels, kernel size p, and stride p.
  • ClimaX architecture: Variable positional embeddings distinguish tokens from different input variables before aggregation reduces them to an h × w × D tensor.The aggregated tokens then enter the attention layers.
  • ClimaX architecture: Attention preserves the h × w × D token shape, while the prediction head maps each token to |𝒱| × p^2 values and reconstructs spatial outputs.Predictions for the V′ target variables are extracted before computing the loss.
  • ClimaX architecture: Spatial resolution can vary by interpolating positional embeddings while leaving the rest of the architecture unchanged.This follows the cited common practice for variable-resolution inputs.

A.2.3. Other implementation details

The CNN baselines use several implementation choices identified as important for performance, including normalization, activation, normalization order, directional periodicity, and an initial large kernel.

  • CNN baseline details: The CNN baselines use batch normalization, LeakyReLU with slope 0.3, post-normalization, and a 7-sized first-layer kernel.These settings are listed as important performance details for ResNet and UNet baselines.
  • CNN baseline details: Periodic convolutions are applied in longitude but not latitude.This directional treatment is part of the reported baseline implementation.

B. Training details

Training normalizes variables, uses mixed-precision distributed hardware, and applies task-specific losses, optimizers, learning rates, and schedules. Pretraining and fine-tuning differ in both objectives and optimization settings.

  • Data and infrastructure: Each variable is normalized to zero mean and unit variance, then predictions are de-normalized before evaluation metrics are computed.For atmospheric variables, normalization is performed separately at each pressure level.
  • Data and infrastructure: Training uses 32GB NVIDIA V100 devices, distributes pretraining across 80 V100s, and uses fp16 precision.The software stack includes PyTorch, timm, NumPy, and xarray.
  • Pretraining: Pretraining uses a loss specified in Equation (1), AdamW with β1 = 0.9 and β2 = 0.95, and a 5e−4 learning rate with warmup followed by cosine annealing.The schedule uses 10,000 warmup steps and 190,000 cosine-annealing steps.
  • Fine-tuning: Fine-tuning uses latitude-weighted MSE for temporal forecasting and downscaling, but standard MSE for ClimateBench.Standard MSE performed better for ClimateBench and was recommended by the cited work.
  • Fine-tuning: All downstream tasks use AdamW with β1 = 0.9 and β2 = 0.999, 10,000 warmup steps, and 90,000 cosine-annealing steps.Weight decay is 1e−5 except for positional embeddings, with learning rates varying by task.
  • Fine-tuning: Weather forecasting uses a small learning rate, downscaling a larger one, and climate projection a scratch-like rate because new embedding and prediction-head weights are initialized.The choices reflect task similarity or difference from pretraining and the need to train new weights.

C. Datasets

The study uses CMIP6-derived datasets for pretraining and ERA5-based datasets for fine-tuning and evaluation across global, regional, and sub-seasonal-to-seasonal tasks.

  • Pretraining data: CMIP6-ClimaX comprises five CMIP6 datasets selected for ClimaX pretraining.The datasets use historical 6-hourly pressure-level simulations with the specified variant label.
  • Pretraining data: The CMIP6 datasets cover 1850–2015 at 6-hour resolution and were regridded to 5.625° and 1.40625° using bilinear interpolation.The datasets were chosen for similarity to ERA5 variables and vertical levels.
  • Pretraining data: AWI and HAMMOZ are unsuitable sources for 1.40625° training because their original 250 km resolution is coarser than the approximately 156 km target grid.The authors were unable to download other higher-resolution datasets and expect them to improve pretraining performance.
  • Fine-tuning and evaluation data: ERA5 from WeatherBench supplies fine-tuning data at 5.625°, 2.8125°, and 1.40625° resolutions.WeatherBench provides a standard benchmark and evaluation framework for data-driven weather forecasting.
  • Fine-tuning and evaluation data: ERA5-NA supports regional North American forecasting, while ERA5-S2S evaluates biweekly prediction for weeks 3–4 and weeks 5–6.ERA5-NA extracts samples within a rectangular North America domain, and ERA5-S2S averages target variables over two-week periods.

D.1. Metrics

The evaluation uses spatially weighted error, correlation, bias, and normalized discrepancy metrics for comparing predictions with ground truth across weather and climate tasks.

  • Metric setup: The metrics compare prediction and ground-truth fields shaped as N × H × W, with latitude weighting correcting for non-uniform grid-cell areas.Time notation is omitted for simplicity in the metric definitions.
  • Climate projection metrics: Anomaly correlation coefficient measures spatial correlation between prediction anomalies relative to climatology.Climatology is defined from the temporal mean of the ground-truth data over the test set.
  • Climate projection metrics: NRMSE_s measures spatial discrepancy between temporal means, while NRMSE_g measures discrepancy between global means.The global mean uses latitude-weighted averaging.
  • Climate projection metrics: TRMSE combines NRMSE_s and NRMSE_g as a weighted sum with α = 5.The weighting follows the value recommended by Watson-Parris, Rao, et al.
  • Downscaling metrics: RMSE measures prediction error using the same formulation as Equation (3), while mean bias measures the difference between predicted and ground-truth spatial means.Positive bias indicates overestimation and negative bias indicates underestimation.
  • Downscaling metrics: The Pearson coefficient measures correlation between prediction and ground truth after flattening both fields.The flattened arrays are then used to compute the metric.

D.2. Results summary

Global forecasting results compare ClimaX with IFS and CNN-based and foundation-model baselines across target variables and lead times, with resolution differences affecting comparability.

  • Global forecasting results: Tables 10 and 11 report global forecasting RMSE and ACC for all target variables and lead times, alongside IFS, CNN baselines, FourCastNet, PanguWeather, and GraphCast.RMSE favors lower values, while ACC favors higher values.
  • Global forecasting results: ClimaX outperforms FourCastNet for Z500, T850, and U10 from 3 days onward in both RMSE and ACC despite lower training resolutions.For T2m, ClimaX performs better at horizons longer than 3 days.
  • Global forecasting results: PanguWeather performs better on most tasks, but its advantage shrinks with increasing lead time and disappears for U10 at 7 days.The authors attribute long-horizon behavior to ClimaX’s direct-prediction fine-tuning, which mitigates error accumulation.
  • Global forecasting results: GraphCast has the lowest RMSE among the compared methods but lower ACC than ClimaX and PanguWeather.The comparison should account for GraphCast, FourCastNet, and PanguWeather being trained at 0.25°, unlike ClimaX’s lower-resolution training.

E. Qualitative evaluation

The qualitative evaluation presents ClimaX forecasts against ERA5 ground truth across lead times ranging from 6 hours to 1 month.

  • Figure layout: Each forecast figure shows the initial condition, ground truth at the selected lead time, ClimaX’s prediction, and the prediction bias.Bias is defined as the difference between prediction and ground truth.
Loading 2301.10343v5…