A Foundation Model for Weather and Climate

## Abstract

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. Source code is available at https://github.com/microsoft/ClimaX .

## 1 Introduction

Modeling weather and climate is an omnipresent challenge for science and society. With rising concerns around extreme weather events and climate change, there is a growing need for both improved weather forecasts for disaster mitigation and climate projections for long-term policy making and adaptation efforts []. Currently, numerical methods for global modeling of weather and climate are parameterized via various general circulation models (GCM) []. GCMs represent system of differential equations relating the flow of energy and matter in the atmosphere, land, and ocean that can be integrated over time to obtain forecasts for relevant atmospheric variables [, ].
While extremely useful in practice, GCMs also suffer from many challenges, such as accurately representing physical processes and initial conditions at fine resolutions, as well as technological challenges in large-scale data assimilation and computational simulations [].
These factors limit their use in many scenarios, especially in simulating atmospheric variables quickly at very short time scales (e.g., a few hours) or accurately at long time scales (e.g., beyond 5-7 days) [].

In contrast, there has been a steady rise in data-driven approaches for forecasting of atmospheric variables, especially for meteorological applications [, , , , , , , , , ].
The key idea here is to train deep neural networks to predict the target atmospheric variables using decades of historical global datasets, such as the ERA-5 reanalysis dataset [].
Unlike GCMs, these networks are not explicitly grounded in physics, and lack general-purpose utility for Earth system sciences as they are trained for a specific predictive modeling task. Yet, with growing compute and datasets, there is emerging evidence that these models can achieve accuracies competitive with state-of-the-art numerical models in many scenarios, such as nowcasting of precipitation [, ] and medium-range forecasting of variables like temperature, wind and humidity [, , , , , ].
While these trends are encouraging, there remain concerns regarding the generality of such data-driven methods to diverse real-world scenarios, such as forecasting of extreme weather events and longer-term climate projections, especially under limited spatiotemporal supervision and computational budgets.

Variants of the aforementioned challenges apply broadly throughout machine learning (ML).
In disciplines such as natural language processing and computer vision, it is well acknowledged that ML models trained to solve a single task using supervised learning are label-hungry during training and brittle when deployed outside their training distribution [].
Recent works have shown that it is possible to mitigate the supervision bottleneck by pretraining [, ] large unsupervised “foundation” models [] on huge passive datasets, such as text and images scraped from the internet [, , , ].
Post pretraining, there are many ways to finetune the same model on arbitrary target task(s) with little to none (i.e., zero-shot) additional supervision.
Besides low target supervision, these models also generalize better to shifts outside their training distribution [, ], improving their reliability.

Inspired by the above successes, this work studies the question: how do we design and train a foundation model for weather and climate that can be efficiently adapted for general-purpose tasks concerning the Earth’s atmosphere?
We propose ClimaX, a foundation model for weather and climate.
For pretraining any foundation model, the key recipe is to train a deep architecture on a large dataset using an unsupervised objective.
For example, many foundation models for language and vision train large transformers on Internet-scale datasets using generative modeling.
While conceptually simple, this scaling recipe is riddled with challenges for weather and climate domains, that we discuss below and propose to resolve with ClimaX.

First, it is unclear what constitutes an Internet-scale passive dataset for pretraining ClimaX.
The size of historical weather and climate datasets at any given time is fixed and increases at an almost constant rate everyday, as it corresponds to processed sensor measurements of naturally occurring phenomena.
Our first key proposal is to go beyond these datasets to explicitly utilize physics-informed climate simulation models.
Many such models are in use today, for example, the CMIP6 collection [] of climate modeling simulations
consists of runs of $\sim$100 distinct climate models from 49 different climate modeling groups.
We show that the heterogeneity in these simulation datasets serves as a source of rich and plentiful data for pretraining ClimaX.

Second, we need a model architecture that can aptly embrace the heterogeneity of the above climate datasets.
Climate data is highly multimodal, as observations typically correspond to many different, unbounded variables with varying datatypes (e.g., pressure, temperature, humidity).
Moreover, many observational datasets are irregular in the sense that they differ in their spatiotemporal coverage and might correspond to different subsets of atmospheric variables.
We resolve the above challenges in ClimaX by repurposing the vision transformer [, ].
In contrast to earlier work where the input data is represented as an image with different atmospheric variables treated as the channels thereof [, ], we treat them as separate modalities to enable more flexible training even with irregular datasets.
This has the side-effect of drastically increasing the sequence length, which we propose to resolve via a cross-attention style channel aggregation scheme prior to the self-attention layers.

Third and last, we need a pretraining objective that can learn complex relationships between the atmospheric variables and permit effective finetuning for downstream tasks.
Given the spatiotemporal nature of climate data, we propose a randomized forecasting objective for pretraining ClimaX. Here, the goal of the model is to forecast an arbitrary set of input variables at an arbitrary time into the future.
While simple and intuitive, we show that such a pretraining objective aids finetuning to novel tasks and timescales even beyond the pretraining window, such as sub-seasonal to seasonal cumulative predictions, climate projections, and downscaling of climate models.
See Figure [1](#S0.F1) for a list of tasks considered in this work.

Empirically, we demonstrate that a single pretrained model can be finetuned for many tasks (e.g., multi-scale weather forecasting, climate projections, downscaling) under a range of operating conditions involving different spatiotemporal resolutions, geographical regions, and target prediction variables, including those unseen during training.
Notably, our benchmark results are state-of-the-art on ClimateBench [] and competitive with the operational Integrated Forecasting System (IFS) [] on WeatherBench [], even when our model is trained on moderate resolutions using only a maximum of 80 NVIDIA V100 GPUs.

Finally, we show promising scaling laws of ClimaX
with natural axes of performance improvements for larger number of pre-training datasets,
larger models, and scaling to higher resolution gridded datasets.
While especially the last is in line with recent and concurrent works on data-driven weather forecasting [, , ],
to the best of our knowledge, ClimaX is the first of its kind data-driven model that can effectively scale using heterogeneous climate datasets during pretraining, and generalize to diverse downstream tasks during finteuning, paving the way for a new generation of data-driven models for Earth systems science.

## 2 Background and Related Work

Current weather and climate models in use today rely extensively on numerical methods and computational simulations to predict and understand the Earth’s weather and climate systems.
These tasks include various *numerical weather prediction* (NWP) systems which use computer simulations to make short-term forecasts of weather conditions as well as climate models which use similar techniques to simulate and predict the long-term changes in the Earth’s climate.
Most notably, at the core of both weather and climate models lie the same set of primitive equations.

For climate modeling, earth system models (ESM) [], or “coupled models”, that couple together simulations which govern the atmosphere, cryosphere, land, and ocean processes are considered the state-of-the-art. Primarily these simulations are based on general circulation models (GCMs) [, , , ] which date back to the works of [, ] solving Navier-Stokes equations on a rotation sphere to model fluid circulation.
These models are often used to perform various *factor sensitivity* studies to examine how the changes in certain forcing factors like greenhouse gas concentrations can affect the global or regional climate and help in *climate projections* to help understand future conditions.

Numerical Weather Prediction (NWP) models share many components of GCMs, especially the atmospheric components [, , ].
However, incorporating *data assimilation* [, ] which involves combining observations and various measurements of the atmosphere and oceans together with these numerical models is important for accurate forecasts and simulations.
Another significant distinction between weather and climate models is the framing of the solution for underlying equations: *initial value problem* for weather, while *boundary value problem* for climate [].
Different difficulty levels of these solution approaches results in the fact where climate models tend to be global often at coarser spatio-temporal resolutions while weather models can range from global to local and regional models of very high spatio-temporal resolutions [].

Despite their noted success, including the recent 2021 Nobel Prize in Physics [], there is considerable debate around the limitations of general circulation models (GCMs), particularly structural errors across models and the fact that current GCMs are designed to reproduce observed climate [].
The climate science community has been aware of these challenges which resulted in the creation of Coupled Model Intercomparison Project (CMIP) as a standardized protocol for evaluating and comparing the performance of different climate models [].
As we will see in the following sections, not only has CMIP been playing a crucial role in the advancement of our understanding of climate change and its potential impacts, its evaluation procedure has resulted in enormous quantity of data making modern deep learning based approaches quite attractive for many tasks.
Notably, encoding this knowledge into a “foundation” machine learning model with much faster inference and data assimilation capabilities can pave the way for a much wider impact.

### 2.1 Data sources

Unlike data in computer vision or natural language processing, weather and climate data is not solely based on sensed data, instead incorporates information from a diverse range of sources.
For example, *reanalysis* weather data
blends meteorological observations with past short-range
weather forecasts via data assimilation [].
The data measurements themselves are highly heterogeneous, representing various physical variables with different data types (e.g. pressure, temperature, humidity) that are recorded at different, relatively sparse, spatial locations at different temporal frequencies.
These measurements can be integrated together with known physics inform the design of climate simulations, which again produce data with different variables at different scales.
From a machine learning perspective, the plethora of available data thus spans multiple axes: from direct weather measurements at land, sea, or atmosphere, over multiple decades of re-analyzed weather data at different spatial scales, to physics-informed climate projections for various scenarios.
Most notably, the data shares the same set of primitive equations, but with fairly different characteristics.
Below we describe two of the most commonly used data sources for weather and climate modeling.

#### 2.1.1 CMIP6

The Coupled Model Intercomparison Project (CMIP) [] is an international effort across different individual climate modeling groups to come together to compare and evaluate their global climate models.
While the main goal of CMIP is to improve the understanding of Earth’s climate system and improve the accuracy of its simulations, the recent data from their experimental runs is easily accessible on the CMIP6 [] archive.
In CMIP6,
where “6” refers to the most recent phase of the project,
$49$ groups are involved with their experiments covering wide range of climate variables including temperature, precipitation, sea level and others from hundreds of models.
This results in global projections of various climate scenarios from as early as 1850 onwards,
all following similar governing equations, but with different *forcings*, e.g., greenhouse gas emissions that affect the climate.

#### 2.1.2 ERA5

The ERA5 reanalysis archive [, ]
of the European Center for Medium-Range Weather Forecasting (ECMWF)
is the predominant data source for learning and benchmarking weather forecasting systems.
Once completed, the ERA5 reanalysis is set to embody a detailed record of the global atmosphere, land surface and ocean waves from 1950 onwards.
The currently available ERA5 reanalysis data combines the state of the art forecasting model called Integrated Forecasting System (IFS) []
of ECMWF with available observations to provide the best guess of the state of the atmosphere, ocean-wave and land-surface quantities at any point in time.
In its raw form, the available reanalyzed data is huge: 40 years, from 1979 to 2018, on a $0.25\degree\times 0.25\degree$ global latitude-longitude grid of the Earth’s sphere, at hourly intervals with different climate variables at 37 different altitude levels plus the Earth’s surface. The grid overall contains $721\times 1440$ grid points for latitude and longitude, respectively. The altitude levels are presented as pressure levels.

### 2.2 Tasks

Given the scale of data availability, increasing compute requirements of current numerical methods despite it being difficult to incorporate real observational data into them, machine learning is increasingly finding applications in many of the tasks related to weather and climate modeling.
When it comes to weather, the main task of interest is *forecasting* the future values of key weather variables.
These tasks can take the following forms depending on temporal and spatial horizons of interest:

- •
Global forecasting tasks that range from a few hours (i.e., nowcasting) to days and weeks in lead time (i.e., short and medium range forecasting). Often these tasks are evaluated on the ERA5 reanalysis dataset (see [Section 2.1.2](#S2.SS1.SSS2)) with Operational IFS []
of the European Center for Medium-Range Weather Forecasting (ECMWF) being the current state-of-the-art NWP baselines.
- •
Regional forecasting tasks which could range from weather forecasting in continental North America or Europe to individual state, county or city.
- •
Sub-seasonal to seasonal prediction (S2S) [, ] which is the task of forecasting the weather with lead times between 2 weeks and 2 months. S2S bridges the gap between weather forecasting and seasonal climate prediction, and is critical to disaster mitigation. Often at such long horizons, predicting instantaneous values of key weather variables can be a difficult task and therefore the focus is often on averaged value of key weather variables over a certain time horizon, e.g. weekly average precipitation.

Whereas deep learning approaches for regional or S2S tasks are scarce, most of the recent and concurrent work focuses on global forecasting tasks.
[] were the first to use pretraining on climate simulations to achieve good data-driven medium-range weather prediction with a ResNet [], []
used CNNs on a cubed sphere for global weather prediction, [] forecast weather sub-seasonally with a large ensemble of deep-learning weather prediction models, [] applied a graph neural network based approach to weather forecasting,
[] use deep generative models of radar for precipitation nowcasting,
[] build a reservoir computing-based, low-resolution, global prediction model, and MetNet [] takes as input radar and satellite data to forecast probabilistic precipitation maps.
These approaches are complemented by general machine learning models for fluid dynamics [, , , , ].
Finally, recent state-of-the-art neural weather models such as FourCastNet [], Pangu-weather [], or GraphCast [], which also perform global forecasting tasks, use the highest resolution $0.25\degree$ ERA5 data, and are optimized on the respective hardware resources.

On the other hand, climate tasks have to deal with much longer time horizons.
Possible categories of tasks where machine learning can help include climate projection and climate model downscaling:

- •
Climate projection is the task of
generating estimates of climate change under different future socio-economic scenarios.
Usually, this takes the form of figuring out the response of the climate system to different forcing factors such as greenhouse gases and aerosol emissions.
Climate projection is a crucial task in understanding and preparing for the potential impacts of climate change.
While the application of machine learning in this field is still in its early stages, recent efforts have been made to standardize evaluation in this domain. One example of this is ClimateBench [], which is a benchmark dataset drawing on CMIP6 to provide an evaluation framework for machine learning models that aim to improve the accuracy of climate projections. This benchmark aims to provide a consistent and reliable evaluation method for various machine learning models that are applied to climate projections.
- •
A more popular application of ideas in machine learning is towards downscaling of climate model.
Global climate models typically have a coarse spatial resolution, which means that they can only provide a rough estimate of climate conditions at a local or regional scale.
Moreover, the simulations often reflect systematic biases that deviate from trends in the observation data.
The aim of climate model downscaling is to create locally accurate climate information from global climate projections by relating those to observed local climatological conditions.
This process improves the spatial and temporal resolution of the data, making it more suitable for use in local and regional analyses.
Downscaling methods can be divided into dynamic approaches that relate outputs of global climate models with those of regional climate models, and statistical approaches that infer the desired transformations using data-driven approaches [].
Dynamic approaches are physically consistent, but can be slow and have large biases, whereas statistical approaches need large amounts of data to learn expressive mappings that hold for target output scenarios.

Similar to weather forecasting, deep learning has emerged as appealing alternative in climate science as well. Recent approaches comprise surrogate models to emulate climate projections
[, , , , ], extract contextual cues from existing datasets or simulations [, , ],
and perform climate model downscaling [, , ].
Climate model downscaling usually inputs low-resolution reanalysis data and local orographic information to obtain high-resolution local information.
Many recent approaches are based on convolutional architectures [, , ].

### 2.3 Foundation models

[] gave the term “foundation models” to the emerging paradigm of training scalable deep learning models on broad data via self-supervision which could then be adapted (often via finetuning) to a wide range of downstream tasks. Current notable examples include BERT [], GPT [] and PaLM [], in language, CLIP [], Florence [], BEiT [] for vision-language.
Outside applications on data crawled from web, this paradigm has also started finding success in various scientific domains like protein design [].
Key significance of such models has been identified as *emergence* with respect to model capabilities and *homogenization* with respect to methodologies for different tasks, domains, and modalities, enabled by the principles of transfer learning [] at scale.
While a foundation model itself should be considered incomplete, it can provide a common basis from which various task-specific models can be derived.
Current research at the intersection of weather and climate science and ML has largely focused on designing separate models for every task of interest despite potential availability of fairly diverse large scale data with shared underlying physics and geology across these tasks.
A few recent works have proposed pretraining techniques for satellite imagery and remote sensing [, , ] but they have so far not been applied to multi-sensory data and variables in weather and climate.

## 3 Approach

Given the availability of large scale data sources, together with shared physics and geology between various weather and climate tasks, we aim to build a generalizable deep learning foundation model. The model needs to be able to input heterogeneous datasets of different variables, and provide spatio-temporal coverage based on physical groundings. We, therefore, first take a closer look at input representations, and next design a model to cope with their heterogeneity - local, global, and across variables.

### 3.1 Input representation

We are interested in gridded prediction tasks, in which the model takes an input of shape $V\times H\times W$ and predicts an output of shape $V^{\prime}\times H^{\prime}\times W^{\prime}$. $V$ refers to the number of input variables, which can be weather conditions such as geopotential and temperature, or climate forcing factors such as CO${}_{2}$ and SO${}_{2}$. $H$ and $W$ refer to the spatial resolution of the input data, which depends on how densely we grid the globe.
This general representation captures a broad variety of downstream tasks in Earth systems science.
Similarly, $V^{\prime},H^{\prime},W^{\prime}$ refer to the variables and spatial resolution of the predicted outputs.
We mainly work with two spatial resolutions: $5.625\degree$ ($32\times 64$ grid points) and $1.40625\degree$ ($128\times 256$ grid points). Semantically, a $H\times W$ map can represent the entire globe or a specific region such as North America.

### 3.2 Model architecture

We aim to design a foundation model that we can pretrain on heterogeneous data sources and then finetune to solve various downstream weather and climate tasks. From [Section 3.1](#S3.SS1), one could think of the tasks as image-to-image translation problems with $V$ input channels and $V^{\prime}$ output channels. This makes any image architecture a natural fit, such as UNet [], ResNet [], or Vision Transformers (ViT) [].
However, the settings of climate and weather tasks are much broader, where we may want to make predictions for regional or even spatially incomplete data, forecast unseen climate variables, or finetune the model on data at different resolutions from pretraining. Current CNN-based architectures are not applicable in these scenarios, as they require the input to be perfectly gridded, contain a fixed set of variables, and have a fixed spatial resolution. Transformers-based architectures, on the other hand, provide much better flexibility by treating the image-like data as a set of tokens. Therefore, we build ClimaX architecture upon Vision Transformers (ViT) [, ], and propose two major architectural changes, namely *variable tokenization* and *variable aggregation* to further improve the flexibility and generality, which we will describe next.

Figure: Figure 2: Pretraining phase of ClimaX. Variables are encoded using variable-separate tokenization, and subsequently aggregated using variable aggregation. Together with position embedding and lead time embedding those are fed to the ViT backbone.
Refer to caption: /html/2301.10343/assets/x2.png

#### 3.2.1 Variable tokenization

Given an input of shape $V\times H\times W$, ViT tokenizes the input into a sequence of $(H/p)\times(W/p)=h\times w$ patches, with each patch having a size of $V\times p^{2}$, where $p$ is the patch size. This tokenization scheme works well for image data, as $V$ is always the RGB channels, which is the same for all datasets. However, this is not true for climate and weather data, where the number of physical variables can vary between different datasets.
For example, in the CMIP6 project [], each dataset contains simulated data of a different climate model, and thus has a different set of underlying variables. Therefore, we propose *variable tokenization*, a novel tokenization scheme that tokenizes each variable in the input separately. Specifically, each input variable as a spatial map of shape $H\times W$ is tokenized into a sequence of $h\times w$ patches, which results in $V\times h\times w$ patches in total. Finally, each input patch of size $p^{2}$ is linearly embedded to a vector of dimension $D$, where $D$ is the chosen embedding size. The output of the variable tokenization module therefore has a dimension of $V\times h\times w\times D$.
Figure [3](#S3.F3) illustrates our proposed tokenization scheme.

Figure: Figure 3: Variable tokenization. Each variable is independently tokenized.
Refer to caption: /html/2301.10343/assets/x3.png

#### 3.2.2 Variable aggregation

While variable tokenization allows ClimaX to learn from datasets with varying numbers of input variables, it has two inherent problems. First, it results in a sequence of length $V\times h\times w$ which increases linearly with the number of variables. Since we use attention to model the sequence, the memory complexity scales quadratically with the number of variables. This is computationally expensive, as we can have up to $48$ input variables in our experiments. Moreover, because we tokenize each variable separately, the input sequence will contain tokens of different variables with very different physical groundings,
which can create difficulties for the attention layers to learn from.
We therefore propose *variable aggregation* to solve the two mentioned challenges. For each spatial position in the $h\times w$ map, we perform a cross-attention operation, in which the query is a learnable vector, and the keys and values are the $V$ embedding vectors of $V$ variables at that position. The cross-attention module outputs a single vector for each spatial position, thus reducing the sequence length to $h\times w$, significantly lowering the computational cost. Moreover, the sequence now contains unified tokens with universal semantics, creating an easier task for the attention layers. Figure [4](#S3.F4) shows our proposed variable aggregation.

Figure: Figure 4: Position-based variable aggregation reduces a sequence of length $V\times h\times w$ to $h\times w$.
Refer to caption: /html/2301.10343/assets/x4.png

#### 3.2.3 Transformer

Post variable aggregation, we need a sequence model for generating the output tokens.
While in principle, one could use any general sequence model, we propose to extend a standard Vision Transformer (ViT).
Moreover, since the standard ViT treats image modeling as pure sequence-to-sequence problems, it can perform tasks that some other variations cannot [, ], such as learning from spatially incomplete data, where the input does not necessarily form a complete grid. This is useful in the regional forecasting task we consider in [Section 4.2.2](#S4.SS2.SSS2).
In the experiments, we report results with $8$ attention layers, an embedding size of $1024$, and a hidden dimension of $1024\times 4$. After the attention layers, we employ a prediction head that takes a token and outputs a vector of size $V^{\prime}\times p^{2}$. The prediction head is a 2-layer MLP with a hidden dimension of $1024$. We provide more details in [Appendix A](#A1).

### 3.3 Datasets

#### 3.3.1 Pretraining

We believe that CMIP6’s diversity and scale presents an attractive opportunity for pretraining large-scale foundation models.
However, handling the inconsistent set of variables across different data sources can be a challenge.
In this work we only use a subset of variables from five different data sources (MPI-ESM, TaiESM, AWI-ESM, HAMMOZ, CMCC) containing global projections of climate scenarios from 1850 to 2015 with the time delta of $6$ hours as described in [Table 8](#A3.T8).
Due to variable original resolution, we choose to simplify our data-loading by regridding them to commonly used resolutions [, ] of $5.625\degree$ ($32\times 64$ grid points) and $1.40625\degree$ ($128\times 256$ grid points)(^1^11Regridding was done using the xesmf Python package [] using bilinear interpolation.).

#### 3.3.2 Finetuning and evaluation

We use the ERA5 reanalysis data as described in [Section C.2](#A3.SS2), as the source of datasets for finetuning and evaluation for various weather related downstream tasks.
Due to its large size, it is common to regrid [, ] the high-resolution data to lower resolutions like $5.625\degree$ ($32\times 64$ grid points) and $1.40625\degree$ ($128\times 256$ grid points) to fit within the available computational constraints(^2^22Regridding was done using the xesmf Python package [] using bilinear interpolation.).
We follow the evaluation procedure by [] and use this data to assess the forecasting performance of our ML models at different lead time horizons.
More details about the individual datasets are in their appropriate experiment sections.

### 3.4 Training

#### 3.4.1 Pretraining

We pretrain ClimaX on CMIP6 data to predict future weather conditions given the current conditions. That is, given the weather snapshot $X_{t}$ of shape $V\times H\times W$ at a particular time $t$, ClimaX learns to predict the future weather scenario $X_{t+\Delta t}$ of the same shape at lead time $\Delta t$. To obtain a pretrained model that is generally applicable to various temporal forecasting tasks, we randomize the lead time from $6$ hours to $168$ hours (i.e., 1 week) during pretraining. We add the lead time embedding to the tokens to inform the model of how long it is forecasting into the future. The lead time embedding module is a single-layer MLP that maps a scalar to a vector of the embedding size $D$. [Figure 2](#S3.F2) depicts the forward pass of ClimaX in pretraining. For an input $X_{t}$, we sample a lead time $\Delta t\sim\mathcal{U}[6,168]$ and get the corresponding ground truth $X_{t+\Delta t}$. Input variables are tokenized separately using variable tokenization, and are subsequently aggregated at each spatial location, resulting in a sequence of $h\times w$ unified tokens. We add the tokens with the lead time embedding and positional embedding before feeding the sequence to the ViT backbone. The output of the last attention layer is fed to a prediction head, which transforms the sequence back to the original shape of $V\times H\times W$.

We employ the latitude-weighted mean squared error [] as our objective function. Given the prediction $\tilde{X}_{t+\Delta t}$ and the ground truth $X_{t+\Delta t}$, the loss is computed as:

$$ $\mathcal{L}=\frac{1}{V\times H\times W}\sum_{v=1}^{V}\sum_{i=1}^{H}\sum_{j=1}^{W}L(i)(\tilde{X}_{t+\Delta t}^{v,i,j}-X_{t+\Delta t}^{v,i,j})^{2},$ (1) $$

in which $L(i)$ is the latitude weighting factor:

$$ $L(i)=\frac{\cos(\text{lat}(i))}{\frac{1}{H}\sum_{i^{\prime}=1}^{H}\cos(\text{lat}(i^{\prime}))},$ (2) $$

where $\text{lat}(i)$ is the latitude of the corresponding $i\text{th}$ row of the grid. The latitude weighting term accounts for the non-uniformity in areas when we grid the round globe. Grid cells toward the equator have larger areas than the cells near the pole, and thus should be assigned more weights.

#### 3.4.2 Finetuning

ClimaX has four learnable components, including the token embedding layers, the variable aggregation module, the attention blocks, and the prediction head. We evaluate the performance of ClimaX on various downstream tasks, which we categorize into two finetuning scenarios: one in which the downstream variables belong to the set of pretraining variables, and the other with variables unseen during pretraining. In the first case, we finetune the entire model, and in the latter, we replace the embedding layers and the prediction head with newly initialized networks, and either finetune or freeze the other two components. We present more details of each downstream task in [Section 4](#S4).

## 4 Experiments

We finetune ClimaX on a diverse set of downstream tasks to evaluate its performance and generality. We categorize the tasks into forecasting, climate projection, and climate downscaling. The experiments aim to answer the following questions:

- •
How does ClimaX perform on global forecasting compared to the current state-of-the-art NWP system?
- •
Can we finetune ClimaX to make forecasts for a specific region or at different temporal horizons from pretraining?
- •
How well does ClimaX perform on climate tasks that are completely different from pretraining?

In addition to the main experiments, we analyze the scaling property of ClimaX, i.e., how the performance of ClimaX improves with increasing data size, model capacity, and data resolution. Finally, we perform comprehensive ablation studies to understand the trade-off between computation and performance when finetuning ClimaX.

### 4.1 Neural baselines

In global forecasting, we compare ClimaX with IFS [], the current gold standard in weather forecasting. In tasks we do not have a baseline, we compare with UNet [, ] and ResNet [], two CNN baselines commonly used in vision tasks. We borrow the ResNet architecture from Weatherbench []. The exact architectural details of these baselines are in [Section A.2](#A1.SS2).

### 4.2 Forecasting

#### 4.2.1 Global forecasting

Given global weather conditions $X_{t}$ at a particular time $t$, we want to forecast the weather at a future time $X_{t+\Delta t}$, in which $\Delta t$ is the lead time. The input variables include $6$ atmospheric variables at $7$ vertical levels, $3$ surface variables, and $3$ constant fields, resulting in $48$ input variables in total. The details of the variables are in [Table 9](#A3.T9). We evaluate ClimaX on predicting four target variables: geopotential at $500$hPa (Z500), the temperature at $850$hPa (T850), the temperature at $2$ meters from the ground (T2m), and zonal wind speed at $10$ meters from the ground (U10). Z500 and T850 are the two standard verification variables for most medium-range NWP models and are often used for benchmarking in previous deep learning works, while the two surface variables, T2m and U10, are relevant to human activities. We consider seven lead times: $6$ hours, $\{1,3,5,7\}$ days, 2 weeks, and 1 month, which range from nowcasting to short and medium-range forecasting and beyond. We consider predicting each target variable at each lead time a separate task, and finetune a separate model for each task. We discuss alternative finetuning protocols in Section [4.6](#S4.SS6).

We compare ClimaX with IFS and the two CNN baselines on the ERA5 dataset at both $5.625\degree$ and $1.40625\degree$ resolutions. Following [], we split the data into three sets, in which the training data is from $1979$ to $2015$, the validation data is in $2016$, and the test data is in $2017$ and $2018$. We finetune ClimaX and train the other deep learning baselines using the latitude-weighted MSE loss in [Equation 1](#S3.E1). We perform early stopping on the validation loss for all deep learning models, and evaluate the best checkpoint on the test set. For IFS, we download the predictions from the TIGGE archive [] for the year $2018$(^3^33We were not able to download IFS predictions for $2017$ due to some server issues.). We compare all methods on latitude-weighted root mean squared error (RMSE) and latitude-weighted anomaly correlation coefficient (ACC), two commonly used metrics in previous works. The formulations of the two metrics are in [Section D.1](#A4.SS1). Lower RMSE and higher ACC indicates better performance.

[Figures 5](#S4.F5) and [6](#S4.F6) show the performance of ClimaX and the baselines at $5.625\degree$ and $1.40625\degree$, respectively. At low resolution, IFS outperforms ClimaX on 6-hour to 5-day prediction tasks. On longer horizons, however, ClimaX performs comparably to or slightly better than IFS, especially on 14-day prediction. At higher resolution, the performance of ClimaX closely matches that of IFS even for short horizons, and is superior in forecasting at $7$ days and beyond. The trends are similar for both RMSE and ACC. The two CNN baselines perform similarly and achieve reasonable performance, but lag behind ClimaX and IFS on all tasks.
We include other additional task-specific baselines [, , ] in [Section D.2](#A4.SS2).
These baselines are trained on higher-resolution ERA5 ($0.25\degree$) so are not directly comparable.

Figure: Figure 5: Performance on global forecasting on ERA5 at $5.625\degree$.
Refer to caption: /html/2301.10343/assets/x5.png

Figure: Figure 6: Performance on global forecasting on ERA5 at $1.40625\degree$.
Refer to caption: /html/2301.10343/assets/x6.png

#### 4.2.2 Regional forecasting

It is not always possible to make global predictions, especially when we only have access to regional data In this section, we evaluate ClimaX on *regional forecasting* of the relevant variables in North America, where the task is to forecast the future weather in North America given the current weather condition in the same region.
We create a new dataset from the ERA5 data at $1.40625\degree$ that has the same set of variables but just focuses on the North America region. We call this dataset ERA5-NA and present details of how to construct it in [Section C.2](#A3.SS2). Training, validation, and test splits are done similarly to [Section 4.2.1](#S4.SS2.SSS1).
[Figure 7](#S4.F7) illustrates the finetuning process of ClimaX on this task, where the only difference from global forecasting is the input now only contains tokens that belong to North America.

Since the task has not been considered in previous works, we compare ClimaX with the two CNN baselines ResNet and UNet, and the scratch-trained version of ClimaX, which we refer to as Cli-ViT.
In addition, we finetune two ClimaX models, in which one was pretrained on CMIP6 at $1.40625\degree$, and the other was pretrained on $5.625\degree$ data. To finetune the low-resolution model on higher-resolution data, we follow the common practice of interpolating the positional embedding [, ]. We denote this model as ClimaX-pos-interp.
We evaluate all methods on predicting Z500, T2m, and T850 at lead times of $3$, $5$, and $7$ days. Latitude-weighted RMSE is used as the evaluation metric.

[Figure 8](#S4.F8) compares the performance of ClimaX and the baselines. ClimaX is the best performing method among different target variables and lead times. Interestingly, even though pretrained on data at a lower resolution, ClimaX-pos-interp achieves the second best performance in predicting Z500 and T850, and only underperforms ResNet in predicting T2m at 3-day lead time. This result shows that ClimaX can gain strong performance on tasks that have different spatial coverage or even different spatial resolution from pretraining.

Figure: Figure 7: Finetuning setup for Regional Forecasting in North America.
Refer to caption: /html/2301.10343/assets/x7.png

Figure: Figure 8: Performance on Regional (North America) forecasting for key variables.
Refer to caption: /html/2301.10343/assets/x8.png

#### 4.2.3 Sub-seasonal to seasonal cumulative prediction

Sub-seasonal to seasonal (S2S) prediction is the task of forecasting at a time range between $2$ weeks and $2$ months [], which bridges the gap between weather forecasting and climate projection. Compared to the other two well-established tasks, S2S prediction has received much less attention, despite having a significant socioeconomic value in disaster mitigation.
Recent works have proposed data-driven approaches based on both traditional machine learning [, , ] and deep learning [, , ], but their performances often lag behind adaptive bias correction methods [] on standard benchmarks [].
Here, following the S2S competition ([https://s2s-ai-challenge.github.io/](https://s2s-ai-challenge.github.io/)), we aim to predict the biweekly average statistics of weeks 3-4 and weeks 5-6, which correspond to lead times of $2$ weeks and $4$ weeks, respectively. We construct ERA5-S2S, a new dataset from $5.625\degree$ ERA5 that has the same input variables, but the output variables are averaged from the lead time to $2$ weeks ahead into the future.

We compare ClimaX with ResNet, UNet, and Cli-ViT on the S2S prediction of four target variables: T850, T2m, U10, and V10.
[Table 1](#S4.T1) compares the RMSE of ClimaX and the baselines. ClimaX achieves the lowest error for all variables, and the performance gap with the best baseline UNet is larger at increasing lead times. ClimaX also has significant performance gains over its scratch-trained counterpart Cli-ViT, showing the effectiveness of our pretraining procedure in capturing features that are generally useful for various temporal prediction tasks.

**Table 1: RMSE of ClimaX and baselines on $5.625\degree$ ERA5-S2S prediction tasks.**
|  | T850 | T2m | U10 | V10 |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Weeks 3-4 | Weeks 5-6 | Weeks 3-4 | Weeks 5-6 | Weeks 3-4 | Weeks 5-6 | Weeks 3-4 | Weeks 5-6 |
| Resnet | 2.12 | 2.13 | 1.88 | 2.16 | 1.91 | 1.94 | 1.52 | 1.59 |
| Unet | 1.91 | 1.95 | 1.67 | 1.79 | 1.85 | 1.90 | 1.52 | 1.57 |
| Cli-ViT | 1.96 | 1.96 | 1.79 | 1.90 | 1.83 | 1.92 | 1.51 | 1.56 |
| ClimaX | 1.89 | 1.92 | 1.66 | 1.70 | 1.81 | 1.86 | 1.50 | 1.54 |

### 4.3 Climate projection

To further test the generality of ClimaX, we evaluate the model on ClimateBench [], a recent benchmark designed for testing machine learning models for climate projections. The goal of ClimateBench is to predict the annual mean global distributions of surface temperature, diurnal temperature range, precipitation, and the $90$th percentile of precipitation, given the four anthropogenic forcing factors: carbon dioxide (CO${}_{2}$), sulfur dioxide (SO${}_{2}$), black carbon (BC), and methane (CH${}_{4}$). We note that this is not a temporal modeling task, as we do not predict the future given the past. Instead, we answer questions like *what will be the annual mean temperature for a specified CO${}_{2}$ level?* In particular, note that the input variables and the task itself are completely different from pretraining.

[Figure 9](#S4.F9) illustrates the finetuning pipeline of ClimaX for ClimateBench. As the input and output variables are unseen during pretraining, we replace the pretrained embedding layers and prediction heads with newly initialized networks, while keeping the attention layers and the variable aggregation module. We consider two finetuning protocols, in which we either freeze(^4^44We finetune the LayerNorm in ClimaX${}_{\text{frozen}}$, as suggested by [].) (ClimaX${}_{\text{frozen}}$) or finetune (ClimaX) the attention layers. In addition, we introduce two components to the pipeline in [Figure 2](#S3.F2). We use a history of the preceding ten years of the forcing factors to make predictions for a particular year, creating an input of shape $T\times V\times H\times W$.
Each time slice of the input goes through variable tokenization, variable aggregation, and the attention layers as usual, which output a feature tensor of shape $T\times h\times w\times D$, where $D$ is the embedding size. The feature tensor then goes through a global average pooling layer, reducing the dimension to $T\times D$. Finally, the $10$-year history is aggregated using a cross-attention layer before being fed to the prediction head, which linearly transforms the $D$-dimensional feature vector to a $H\times W$ map. The history aggregation and the global pooling modules are the two additions to the original ClimaX architecture. These architectural designs are inspired by the neural network baseline in [].

We compare ClimaX with ClimaX${}_{\text{frozen}}$, Cli-ViT, and the best baseline from ClimateBench. Following [], we use the standard mean squared error ([Equation 1](#S3.E1) without the weighting term) as the loss function. We evaluate all methods on RMSE, NRMSE${}_{s}$ (Spatial), NRMSE${}_{g}$ (Global), and Total = NRMSE${}_{s}$ + 5 $\times$ NRMSE${}_{g}$ []. Details of the metrics are in [Section D.1](#A4.SS1). [Table 2](#S4.T2) shows the results. ClimaX${}_{\text{frozen}}$ performs the best in predicting two temperature-related variables, followed by ClimaX. This shows that the pretrained attention layers can serve as a strong feature extractor in seemingly unrelated tasks. Where downstream data is scarce (ClimateBench has only $754$ data points), further finetuning the attention layer can lead to overfitting and thus slightly hurt the performance. In two precipitation-related tasks, ClimaX${}_{\text{frozen}}$ slightly underperforms ClimateBench baseline in terms of NRMSE${}_{s}$ and NRMSE${}_{g}$ but outperforms on RMSE. We hypothesize that this was because ClimaX did not observe the precipitation variable during pretraining, which has very different behaviors from other variables.

Figure: Figure 9: Finetuning pipeline for ClimateBench. A different set of input and output variables requires different embedding layers and prediction heads. Attention layers can be frozen or finetuned.
Refer to caption: /html/2301.10343/assets/x9.png

**Table 2: Performance of ClimaX and the baselines on ClimateBench. Spatial and Global denote the normalized root mean squared error NRMSE${}_{s}$ and the NRMSE of the global mean NRMSE${}_{g}$, respectively. Total is a weighted combination of Spatial and Global.**
|  | Surface temperature | Diurnal temperature range | Precipitation | 90th percentile precipitation |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Spatial | Global | Total | RMSE | Spatial | Global | Total | RMSE | Spatial | Global | Total | RMSE | Spatial | Global | Total | RMSE |
| ClimateBench-NN (reproduced) | 0.123 | 0.080 | 0.524 | 0.404 | 7.465 | 1.233 | 13.632 | 0.150 | 2.349 | 0.151 | 3.104 | 0.553 | 3.108 | 0.282 | 4.517 | 1.594 |
| ClimateBench-NN (paper) | 0.107 | 0.044 | 0.327 | N/A | 9.917 | 1.372 | 16.778 | N/A | 2.128 | 0.209 | 3.175 | N/A | 2.610 | 0.346 | 4.339 | N/A |
| Cli-ViT | 0.086 | 0.044 | 0.305 | 0.362 | 6.997 | 1.759 | 15.792 | 0.146 | 2.224 | 0.241 | 3.430 | 0.550 | 2.800 | 0.329 | 4.447 | 1.579 |
| ClimaX | 0.086 | 0.043 | 0.300 | 0.362 | 7.148 | 0.961 | 11.952 | 0.147 | 2.360 | 0.206 | 3.390 | 0.554 | 2.739 | 0.332 | 4.397 | 1.575 |
| ClimaX${}_{\text{frozen}}$ | 0.085 | 0.043 | 0.297 | 0.360 | 6.688 | 0.810 | 10.739 | 0.144 | 2.193 | 0.183 | 3.110 | 0.549 | 2.681 | 0.342 | 4.389 | 1.572 |

### 4.4 Climate model downscaling

Climate models are often run at coarse grids due to their high computational cost. Although these predictions are useful in understanding large-scale climate trends, they do not provide sufficient detail to analyze regional and local phenomena. Downscaling aims to obtain higher-resolution projections and reduce biases from the outputs of these models.
To evaluate the applicability of ClimaX to the task of climate model downscaling, we construct a new dataset based on CMIP6 and ERA5 data sources for coarse inputs and higher resolution targets.
Specifically, we use all MPI-ESM, a dataset from CMIP6, and its variables listed in [Table 8](#A3.T8) at $5.625\degree$ as input, and train separate models to downscale to each ERA5 target variable at $1.40625\degree$. We compare ClimaX with Cli-ViT and the two CNN baselines, UNet and ResNet, as most recent deep downscaling methods [, , , , ] are based on convolution. We were not able to compare with YNet [], the current best method on deep downscaling as we did not have access to high-resolution auxiliary data such as elevation and topographical information.
For all methods, we first bilinearly interpolate the input to match the resolution of the desired output before feeding it to the model. We evaluate all methods on RMSE, Pearson correlation, and Mean bias, which were commonly used in existing deep downscaling works [, ].
Details of the metrics are in [Section D.1](#A4.SS1).

[Table 3](#S4.T3) compares ClimaX and the baselines quantitatively. ClimaX achieves the lowest RMSE and a mean bias closest to 0 0 for all three target variables, and performs similarly to the baselines in terms of Pearson correlation.
While pretrained to perform forecasting, ClimaX has successfully captured the spatial structure of weather data, which helps in downstream tasks like downscaling.
[Figure 10](#S4.F10) visualizes the downscaled predictions of ClimaX for the three target variables. The input is at a much lower resolution and contains a lot of bias compared to the ground truth. While the prediction is missing some fine details, it has successfully captured the general structure of the ERA5 data and removed input biases.

**Table 3: Performance of ClimaX and the baselines on downscaling from MPI-ESM ($5.625\degree$) to ERA5 ($1.40625\degree$).**
|  | Z500 | T850 | T2m | U10 | V10 |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | RMSE | Pearson | Mean bias | RMSE | Pearson | Mean bias | RMSE | Pearson | Mean bias | RMSE | Pearson | Mean bias | RMSE | Pearson | Mean bias |
| ResNet | $825.75$ | $0.96$ | $-108.54$ | $3.60$ | $0.96$ | $0.19$ | $2.89$ | $0.98$ | $0.14$ | $4.05$ | $0.65$ | $0.06$ | $4.11$ | $0.45$ | $0.09$ |
| UNet | $858.35$ | $0.95$ | $35.10$ | $3.66$ | $0.96$ | $-0.34$ | $2.95$ | $0.98$ | $0.16$ | $4.09$ | $0.64$ | $-0.06$ | $4.13$ | $0.44$ | $0.08$ |
| Cli-ViT | $811.61$ | $0.96$ | $-54.32$ | $3.58$ | $0.97$ | $-0.29$ | $2.80$ | $0.99$ | $-0.06$ | $4.01$ | $0.66$ | $-0.08$ | $4.07$ | $0.47$ | $0.01$ |
| ClimaX | $807.43$ | $0.96$ | 2.70 | 3.49 | $0.97$ | $-0.11$ | $2.79$ | $0.99$ | $-0.06$ | $3.99$ | $0.66$ | $0.04$ | $4.06$ | $0.47$ | $-0.02$ |

Figure: Figure 10: Example visualizations of downscaled prediction of key variables by ClimaX.
Refer to caption: /html/2301.10343/assets/x10.png

### 4.5 Scaling laws analysis

Transformers have shown favorable scaling properties for language [, ], vision [], or even multi-modal tasks [, , ]. That is, their performance improves with respect to data size and model capacity given sufficient compute.
In this section, we study the scaling laws of ClimaX in weather forecasting. [Figure 11](#S4.F11) presents the performance of ClimaX as a function of data size and model capacity. The $x$-axis is the pretraining data size measured in Gigabytes, which corresponds to $1$ to $5$ CMIP6 datasets, and the $y$-axis shows the RMSE of ClimaX on the 3-day forecasting task. We compare four ClimaX models with different capacities by varying the embedding dimension from $128$ to $1024$. All experiments are conducted on the $5.625\degree$ data. The error rate of the two biggest models decreases consistently as we increase the data and model size. This highlights the unique ability of ClimaX in learning from diverse and heterogeneous data sources, which allows us to further improve the performance by simply pretraining on more data. However, the two smaller models do not scale as well as the bigger ones, where increasing data size does not gain much improvement or can sometimes hurt performance. This result shows that larger models not only perform better but are also more data efficient.

In addition to data size and model capacity, data resolution is another important scaling dimension in the context of weather and climate. In many vision tasks such as classification, understanding the general, high-level structure of the image is sufficient to make accurate predictions. To model the underlying complex physical processes that govern weather and climate, however, it is important for a model to look at fine-grained details of the input in order to understand the spatial and temporal structure of data as well as the interactions between different variables. High-resolution data contains finer details and local processes of weather conditions that are not present in the low-resolution data, and thus provides stronger signals for training deep learning models. [Figure 12](#S4.F12) compares the performance of ClimaX pretrained and finetuned on $5.625\degree$ and $1.40625\degree$ data on global forecasting. Except for T2m at $1$ day and $3$ days lead times, ClimaX ($1.40625\degree$) consistently achieves lower RMSE and higher ACC than the low-resolution model. We note that for the high-resolution data we have to use a larger patch size ($4$ compared to $2$ for low-resolution data) due to lack of memory issue. We can further improve the performance of ClimaX on the $1.40625\degree$ data by reducing the patch size, as the model is able to capture better details.

Figure: Figure 11: Error on ERA5 3-day forecasting for different variables with respect to CMIP6 5.625$\degree$ data seen during pre-training. Bigger models are more sample efficient.
Refer to caption: /html/2301.10343/assets/x13.png

Figure: Figure 12: Scaling performance with respect to data resolution. Despite a larger patch size, ClimaX ($1.40625\degree$) achieves consistently better performance than the low-resolution model on almost all tasks, except for T2m forecast at 1 day and 3 days lead times.
Refer to caption: /html/2301.10343/assets/x14.png

### 4.6 Ablation studies

In the main forecasting results, we finetune a separate ClimaX model for each target variable at each lead time, as we found this protocol led to the best performance. However, this can be computationally expensive, as finetuning cost scales linearly with respect to the number of target variables and lead times. In this section, we consider different finetuning alternatives to investigate the trade-off between computation and performance.

Figure: Figure 13: Performance of ClimaX and its variations on weather forecasting. ClimaX-cont is a lead-time-conditioned model that we finetune to make predictions at 6 hours to 7 days. ClimaX-iter forecasts at a $6$-hour lead time and rolls out the predictions to forecast at longer horizons. ClimaX-all-vars predicts the future conditions of all variables in the input at particular lead-times.
Refer to caption: /html/2301.10343/assets/x16.png

#### 4.6.1 Should we finetune ClimaX for each variable separately or all at once?

Instead of finetuning ClimaX for each target variable separately, we could alternatively finetune once to predict all variables in the input simultaneously, which we denote as ClimaX-all-vars. [Figure 13](#S4.F13) shows that ClimaX-all-vars achieves comparable performance to ClimaX in most of the tasks and only underperforms for forecasting T2m. This suggests that with a limited budget, one can finetune ClimaX to predict all target variables at the same time without losing much performance.

#### 4.6.2 Should we do iterative forecast or direct forecast?

To avoid finetuning a different model for each lead time, we can finetune ClimaX to make predictions at a short horizon such as $6$ hours, and roll out the predictions during inference to make forecasts at longer horizons. We call this model ClimaX-iter, where *iter* stands for iterative prediction []. We note that in order to roll out more than one step, ClimaX-iter must predict for all input variables, or in other words. This provides the benefit of finetuning a single model that can predict for any target variable at any lead time. [Figure 13](#S4.F13) shows that ClimaX-iter works reasonably well up to 1-day prediction, but the performance degrades significantly at longer lead times. This is not surprising, because ClimaX-iter is not finetuned to predict multiple steps into the future, leading to quick error accumulation. One can employ a multi-step objective for finetuning as in [] to achieve better results.

#### 4.6.3 Can we finetune ClimaX to work for all lead times?

Another way to avoid finetuning for each lead time separately is to finetune a lead-time-conditioned model. Specifically, during finetuning, we randomize the lead time from $6$ hours to $7$ days, resembling the pretraining setting. Note that unlike ClimaX-iter, we still have to finetune a separate model for each target variable. We call this model ClimaX-cont, wherein *cont* stands for *continuous*, a standard term used in previous works []. [Figure 13](#S4.F13) shows that ClimaX-cont performs competitively on $6$-hour to $7$-day forecasting, but fails to extrapolate to $2$ weeks and $1$ month lead times that are unseen during training.
One can also randomize the lead time from $6$ hours to $1$ month, but that means the model sees much fewer data points for each target lead time, potentially hurting the performance.

The cost for finetuning each set of weights is a constant $C$, which is about $15$ hours on an $8\times\text{V}100s$. Among different finetuning protocols, ClimaX is the most expensive, whose total cost is $C\times\#variables\times\#lead_times$, scaling linearly with the number of target variables and lead times. Following ClimaX are ClimaX-all-vars and ClimaX-cont, whose total costs are $C\times\#lead_times$ and $C\times\#variables$, respectively. Finally, ClimaX-iter is the cheapest finetuning protocol, where we only have to finetune a single model that works for all target variables and at all lead times. The performance is proportional to the computational cost, as ClimaX is the best performing model, while ClimaX-iter is the worst.

## 5 Discussion and Future Work

The scaling of datasets, model architectures, and computation has resulted in a transformative impact in various subdisciplines of artificial intelligence, from natural language and speech processing to computer vision, as well as scientific applications in biology and chemistry. In particular, it has led to the emergence of general-purpose foundation models that are trained on large datasets and compute clusters, and can be easily adapted to a variety of downstream tasks efficiently, both in terms of compute and data supervision.
Our work represents a pioneering effort to enable such broad scaling and generality in data-driven models for weather and climate. This approach goes beyond the limitations of both traditional numerical modeling and existing data-driven forecasting methods. Unlike ClimaX, numerical models scale only in terms of computation and not in terms of dataset size, whereas existing data-driven models are typically limited to specific tasks and lack general-purpose applicability across a wide range of tasks.

In addition to traditional considerations in language and vision, foundation models like ClimaX open up new opportunities for scaling through the use of simulation datasets and grid resolutions. To simplify our approach, we chose to use pretraining datasets that include standard variables that have been benchmarked in previous research on data-driven forecasting [, ]. Additionally, we avoided datasets that simulate future scenarios under different forcings to prevent any potential leakage for the climate projection task. Future research could explore incorporating both observational and simulated datasets that include a wider range of climate variables, higher spatiotemporal resolutions, and even extend into future scenarios.
Further, we showed that resolution plays a crucial role in scaling of ClimaX. Due to our compute restrictions, we trained ClimaX on low to moderate resolutions.
Nevertheless, our empirical trends suggest that scaling to higher resolutions ($0.25\degree$) is likely to lead to even better results.

Scaling efforts in the future can benefit from better sequence modeling architectures, especially those designed for multimodal spatiotemporal inputs.
As we saw in ClimaX, the number of channels for climate datasets is much greater than those handled for standard multimodal settings (e.g., audio-video, vision-language models).
Moreover, in practice, there is also a significant range of resolutions across different climate datasets.
This heterogeneity drastically increases the raw length of input sequences for standard architectures such as ViT.
In the future, we believe that investigating single multi-scale architectures (e.g., []) can potentially aid in scaling to such diverse multi-resolution and multi-modal datasets by learning to infer features relevant to atmospheric phenomena at increasing spatial resolutions.

In conclusion, we believe that the generality of our approach has potential applications beyond the tasks considered in this work. It would be interesting to explore the generalization of a pretrained ClimaX backbone to other Earth systems science tasks, such as predicting extreme weather events [, ] and assessing anthropogenic contributions to climate change [, ], as well as broader domains that are closely tied to weather and climate conditions, such as agriculture, demography, and actuarial sciences.

We would like to thank ECMWF for enabling this line of research with accessible public datasets, contributors of WebPlotDigitizer [] for making it easier to build LABEL:tab:rmse-compare and LABEL:tab:acc-compare, and numerous other open-source libraries, notably numpy [] and PyTorch [].
Some icons in [Fig. 1](#S0.F1) by Freepik, smalllikeart, and GOWI from [flaticon.com](flaticon.com).

## Appendix A Model

This section presents the implementation details and hyperparameters of ClimaX and the two CNN baselines UNet and ResNet.

### A.1 ClimaX

#### A.1.1 Implementation details

ClimaX receives a tensor of shape $V\times H\times W$ and outputs a tensor of shape $V^{\prime}\times H\times W$, where the number of input and output variables $V$ and $V^{\prime}$ can vary between different datasets(^5^55The spatial resolution $H\times W$ can also vary. In that case, we employ the common practice of interpolating the positional embedding, and everything else remains the same [, ].). To do that, we assume a set $\mathcal{V}$ that contains all possible variables we could encounter during pretraining and finetuning. Each variable in $\mathcal{V}$ has a separate token embedding layer.

The variable tokenization module tokenizes the input to a sequence of $V\times h\times w$ tokens, with each token being a vector of size $p^{2}$. After that, for each token, we extract the corresponding embedding layer that transforms the token to a vector of dimension $D$. Each embedding layer is a single convolution layer with $in_channels=1,out_channels=D,kernel_size=p,stride=p$. This results in a tensor of shape $V\times h\times w\times D$.

To differentiate between tokens of different input variables, we add the sequence with a *variable positional embedding*, which is a tensor of shape $|\mathcal{V}|\times D$. For each input variable, we extract the corresponding variable positional embedding to add to its tokens. After that, all tokens go through the variable aggregation module, which outputs a tensor of shape $h\times w\times D$.

The tokens are then fed to the attention layers, which output a tensor of the same shape $h\times w\times D$. The prediction head takes each token of dimension $D$ and maps it to a vector of dimension $|\mathcal{V}|\times p^{2}$, and the output is reshaped to $|\mathcal{V}|\times H\times W$. Finally, we extract predictions of $V^{\prime}$ target variables and compute the loss.

#### A.1.2 Hyperparameters

**Table 4: Default hyperparameters of ClimaX**
| Hyperparameter | Meaning | Value |
| --- | --- | --- |
| $\mathcal{V}$ | Default variables | All ERA5 variables in Table [9](#A3.T9) |
| $|\mathcal{V}|$ | Number of default variables | 48 |
| $p$ | Patch size | $2$ for $5.625\degree$<br>$4$ for $1.40625\degree$ |
| $D$ | Embedding dimension | $1024$ |
| Depth | Number of ViT blocks | $8$ |
| # heads | Number of attention heads | $16$ |
| MLP ratio | Determine the hidden dimension of<br>the MLP layer in a ViT block | $4$ |
| Prediction depth | Number of layers of the prediction head | $2$ |
| Hidden dimension | Hidden dimension of the prediction head | $1024$ |
| Drop path | For stochastic depth [] | $0.1$ |
| Dropout | Dropout rate | $0.1$ |

### A.2 CNN Baselines

#### A.2.1 ResNet Hyperparameters

We use the following hyperparameters for ResNet in all of our experiments.

**Table 5: Default hyperparameters of ResNet**
| Hyperparameter | Meaning | Value |
| --- | --- | --- |
| Padding size | Padding size of each convolution layer | $1$ |
| Kernel size | Kernel size of each convolution layer | $3$ |
| Stride | Stride of each convolution layer | $1$ |
| Hidden dimension | Number of output channels of each residual block | $128$ |
| Residual blocks | Number of residual blocks | $28$ |
| Dropout | Dropout rate | $0.1$ |

#### A.2.2 UNet Hyperparameters

We borrow our UNet implementation from [PDEArena](https://github.com/microsoft/pdearena/blob/main/pdearena/modules/twod_unet.py) [].
We use the following hyperparameters for UNet in all of our experiments.

**Table 6: Default hyperparameters of UNet**
| Hyperparameter | Meaning | Value |
| --- | --- | --- |
| Padding size | Padding size of each convolution layer | $1$ |
| Kernel size | Kernel size of each convolution layer | $3$ |
| Stride | Stride of each convolution layer | $1$ |
| Channel multiplications | Determine the number of output channels<br>for Down and Up blocks | $[1,2,2,4]$ |
| Blocks | Number of blocks | $2$ |
| Use attention | If use attention in Down and Up blocks | False |
| Dropout | Dropout rate | $0.1$ |

#### A.2.3 Other implementation details

Following the implementation of ResNet in [, , ], we found the following details important for the performance of both CNN baselines:

- •
Use Batch normalization
- •
Use Leakyrelu with a slope of $0.3$ as the activation function
- •
Postnorm instead of Prenorm
- •
Use periodic convolutions in the longitude direction but not the latitude direction.
- •
Use a kernel size of $7$ in the first CNN layer.

## Appendix B Training details

##### Data normalization

We normalized all inputs during pre-training as well as fine-tuning.
For each variable, at each pressure level (for atmospheric variables), we compute the mean and standard deviation to normalize them to zero mean and unit variance. We de-normalize the predictions to get back to the original range before computing evaluation metrics.

##### Software and hardware stack

We use PyTorch [], timm [], numpy [] and xarray [] to manage our data and model training.
We used 32GB NVIDIA V100 devices for training. For pretraining we distribute the batch across 80 V100s on AzureML. We leverage fp16 floating point precision in our model.

### B.1 Pretraining

#### B.1.1 Objective

We use the loss function in [Equation 1](#S3.E1) for pretraining.

#### B.1.2 Optimization

We used the AdamW optimizer [, ] with parameters ($\beta_{1}=0.9,\beta_{2}=0.95$). We used weight decay of $1e-5$ for all parameters except for the positional embedding. We used a learning rate of $5e-4$, with a linear warmup schedule for $10000$ steps ($5$ epochs), followed by a cosine-annealing schedule for $190000$ steps ($95$ epochs).

### B.2 Finetuning

#### B.2.1 Objective

We use lat-weighted MSE in [Equation 1](#S3.E1) for finetuning ClimaX in temporal forecasting and downscaling tasks. In ClimateBench, we finetune using standard MSE without the weighting term, as this led to better results and was suggested by [].

#### B.2.2 Optimization

For all tasks, we used AdamW with parameters ($\beta_{1}=0.9,\beta_{2}=0.999$). We used weight decay of $1e-5$ for all parameters except for the positional embedding. We used a linear warmup schedule for $10000$ steps ($5$ epochs), followed by a cosine-annealing schedule for $90000$ steps ($45$ epochs). The learning rate for each task is as follows:

**Table 7: Learning rate for finetuning ClimaX in different downstream tasks**
| Task | Learning rate |
| --- | --- |
| Weather forecasting | $5e-7$ |
| Climate projection | $5e-4$ |
| Climate downscaling | $5e-5$ |

We used a small learning rate for weather forecasting as the task resembles pretraining. For downscaling, we used a larger learning rate, as the nature of the task is different from pretraining, even though the input variables are similar. In climate projection, we needed to initialize new weights for the embedding layers and prediction heads, and thus used a similar learning rate to training from scratch.

## Appendix C Datasets

### C.1 CMIP6-ClimaX

We created CMIP6-ClimaX for pretraining ClimaX, which consists of $5$ datasets from the CMIP6 project. We downloaded the datasets from the official CMIP6 search interface at [https://esgf-data.dkrz.de/search/cmip6-dkrz/](https://esgf-data.dkrz.de/search/cmip6-dkrz/). These datasets share the following attributes:

- •
Experiment ID: historical
- •
Table ID: 6hrPlevPt, i.e., 6-hourly data on pressure levels.
- •
Variant label: r1i1p1f1. The variant label distinguishes among closely related simulations by a single model, in which “r” specifies the initial condition, “i” specifies the observational dataset and initialization method used for determining the initial condition, “p” specifies the perturbed physics version of the model, and “f” specifies the forcing index.

All datasets have a temporal coverage from $1850$ to $2015$ and a temporal resolution of $6$ hours. We chose these datasets as they contain similar climate variables at similar vertical levels to ERA5.
We note that there are more than $5$ datasets from CMIP6 that suit our selection criteria, but we were not able to download others due to some issues on the data servers. We regridded these datasets to $5.625\degree$ and $1.40625\degree$ using the xesmf Python package [] using bilinear interpolation.
We provide a detailed description of these $5$ data sources and the available variables we used to construct CMIP6-ClimaX in [Table 8](#A3.T8).

**Table 8: Resolution and variables of CMIP6-ClimaX dataset used for pretraining. *Static* represents variables don’t depend on time, *Single* represents surface variables, and *Atmospheric* represents time-varying atmospheric properties at the chosen altitudes.**
| Data Source | Original resolution | Variables |  |  |
| --- | --- | --- | --- | --- |
| Type | Abbrev. | Levels |  |  |
| MPI | 100km | Single | t2m |  |
| Single | u10 |  |  |  |
| Single | v10 |  |  |  |
| Atmospheric | z | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | u | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | v | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | t | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | q | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Tai | 100km | Single | t2m |  |
| Atmospheric | z | 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | u | 250, 500, 850 |  |  |
| Atmospheric | v | 250, 500, 850 |  |  |
| Atmospheric | t | 250, 500, 850 |  |  |
| Atmopheric | q | 250, 500, 600, 700, 850, 925 |  |  |
| AWI | 250km | Single | t2m |  |
| Single | u10 |  |  |  |
| Single | v10 |  |  |  |
| Atmospheric | z | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | u | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | v | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | t | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | q | 50, 250, 500, 600, 700, 850, 925 |  |  |
| HAMMOZ | 250km | Single | t2m |  |
| Single | u10 |  |  |  |
| Single | v10 |  |  |  |
| Atmospheric | z | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | u | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | v | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | t | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | q | 50, 250, 500, 600, 700, 850, 925 |  |  |
| CMCC | 100km | Atmospheric | z | 50, 250, 500, 600, 700, 850, 925 |
| Atmospheric | u | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | v | 50, 250, 500, 600, 700, 850, 925 |  |  |
| Atmospheric | t | 250, 500, 850 |  |  |

We note that AWI and HAMMOZ are not the best data sources for higher resolution $1.40625\degree$ training, because their original resolution at $250$ km is lower than $1.40625\degree$, which is about $156$ km. We wanted to use other higher-resolution datasets but were not able to download them. We believe pretraining on other high-resolution datasets would lead to better performance.

### C.2 ERA5

We use the preprocessed version of ERA5 from WeatherBench [] for finetuning ClimaX. WeatherBench was created as a standard benchmark data and evaluation framework for comparing data-driven weather forecasting models. WeatherBench regridded the original ERA5 at $0.25\degree$ to three lower resolutions: $5.625\degree$, $2.8125\degree$, and $1.40625\degree$. See [https://confluence.ecmwf.int/display/CKB/ERA5%3A+data+documentation](https://confluence.ecmwf.int/display/CKB/ERA5%3A+data+documentation) for more details of the raw ERA5 data.
Table [9](#A3.T9) summarizes the variables we use for finetuning ClimaX.

**Table 9: ECMWF variables used in our ERA5 dataset. *Static* represents variables don’t depend on time, *Single* represents surface variables, and *Atmospheric* represents time-varying atmospheric properties at the chosen altitudes.**
| Type | Variable name | Abbrev. | ECMWF ID | Levels |
| --- | --- | --- | --- | --- |
| Static | Land-sea mask | LSM | 172 |  |
| Static | Orography |  |  |  |
| Single | 2 metre temperature | T2m | 167 |  |
| Single | 10 metre U wind component | U10 | 165 |  |
| Single | 10 metre V wind component | V10 | 166 |  |
| Atmospheric | Geopotential | Z | 129 | 50, 250, 500, 600, 700, 850, 925 |
| Atmospheric | U wind component | U | 131 | 50, 250, 500, 600, 700, 850, 925 |
| Atmospheric | V wind component | V | 132 | 50, 250, 500, 600, 700, 850, 925 |
| Atmospheric | Temperature | T | 130 | 50, 250, 500, 600, 700, 850, 925 |
| Atmospheric | Specific humidity | Q | 133 | 50, 250, 500, 600, 700, 850, 925 |
| Atmospheric | Relative humidity | R | 157 | 50, 250, 500, 600, 700, 850, 925 |

#### C.2.1 ERA5-NA

We constructed ERA5-NA from ERA5 to evaluate ClimaX and the baselines on regional forecasting. ERA-NA has the same set of variables as in Table [9](#A3.T9), but only contains data that belongs to the North America region. To do this, we first identified the latitude and longitude range to form a rectangular area that encapsulates North America, using the standard CORDEX domains [https://cordex.org/wp-content/uploads/2012/11/CORDEX-domain-description_231015.pdf](https://cordex.org/wp-content/uploads/2012/11/CORDEX-domain-description_231015.pdf). For each data sample, we then extracted the spatial positions that fall into this range, forming in ERA5-NA.

#### C.2.2 ERA-S2S

We built ERA5-S2S from ERA5 to serve as a benchmark dataset for sub-seasonal to seasonal prediction. ERA5-S2S consists of two sub-datasets, whose the goals are to predict the biweekly average statistics of target variables in weeks $3$ and $4$, and weeks $5$ and $6$, respectively. The input includes all variables in Table [9](#A3.T9), while the output variables are are averaged over two weeks, starting from the start of week 3 (5) and to the end of week 4 (6).

### C.3 ClimateBench

We refer to [] for complete details of ClimateBench.

## Appendix D Quantitative evaluation

### D.1 Metrics

This section presents all evaluation metrics we use in [Section 4](#S4). For all metrics, we denote $\tilde{X}$ and $X$ as the prediction and ground truth, which have a shape of $N\times H\times W$, where $N$ is the number of forecasts, or the number of test samples, $H\times W$ is the spatial resolution. $L(i)$ is the latitude weighting term to account for the non-uniformity in areas of the grid cells. We have removed the time notation for simplicity.

#### D.1.1 Weather forecasting metrics

##### Root mean square error (RMSE)

$$ $\text{RMSE}=\frac{1}{N}\sum_{k=1}^{N}\sqrt{\frac{1}{H\times W}\sum_{i=1}^{H}\sum_{j=1}^{W}L(i)(\tilde{X}_{k,i,j}-X_{k,i,j})^{2}}.$ (3) $$

##### Anomaly correlation coefficient (ACC)

Anomaly correlation coefficient (ACC) is the spatial correlation between prediction anomalies $\tilde{X}^{{}^{\prime}}$ relative to climatology and ground truth anomalies $X^{{}^{\prime}}$ relative to climatology:

$$ $\displaystyle\text{ACC}=\frac{\sum_{k,i,j}L(i)\tilde{X}^{{}^{\prime}}_{k,i,j}X^{{}^{\prime}}_{k,i,j}}{\sqrt{\sum_{k,i,j}L(i)\tilde{X}^{{}^{\prime}2}_{k,i,j}\sum_{k,i,j}L(i)X^{{}^{\prime}2}_{k,i,j}}},$ (4) $\displaystyle\tilde{X}^{{}^{\prime}}=\tilde{X}^{{}^{\prime}}-C,X^{{}^{\prime}}=X^{{}^{\prime}}-C,$ (5) $$

in which climatology $C$ is the temporal mean of the ground truth data over the entire test set $C=\frac{1}{N}\sum_{k}X$.

#### D.1.2 Climate projection metrics

##### Normalized spatial root mean square error (NRMSE s 𝑠 {}_{s} )

Normalized spatial root mean square error (NRMSE${}_{s}$) measures the spatial discrepancy between the temporal mean of the prediction and the temporal mean of the ground truth:

$$ $\text{NRMSE}_{s}=\sqrt{\left\langle\left(\frac{1}{N}\sum_{k=1}^{N}\tilde{X}-\frac{1}{N}\sum_{k=1}^{N}X\right)^{2}\right\rangle}\bigg{/}\frac{1}{N}\sum_{k=1}^{N}\left\langle X\right\rangle,$ (6) $$

in which $\langle A\rangle$ is the global mean of $A$:

$$ $\langle A\rangle=\frac{1}{H\times W}\sum_{i=1}^{H}\sum_{j=1}^{W}L(i)A_{i,j}$ (7) $$

##### Normalized global root mean square error (NRMSE g 𝑔 {}_{g} )

Normalized global root mean square error (NRMSE${}_{g}$) measures the discrepancy between the global mean of the prediction and the global mean of the ground truth:

$$ $\text{NRMSE}_{g}=\sqrt{\frac{1}{N}\sum_{k=1}^{N}\left(\langle\tilde{X}\rangle-\langle X\rangle\right)^{2}}\bigg{/}\frac{1}{N}\sum_{k=1}^{N}\left\langle X\right\rangle.$ (8) $$

##### Total normalized root mean square error (TRMSE)

Total normalized root mean square error (TRMSE) is the weighted sum of NRMSE${}_{s}$ and NRMSE${}_{g}$:

$$ $\text{TRMSE}=\text{NRMSE}_{s}+\alpha\cdot\text{NRMSE}_{g},$ (9) $$

where $\alpha$ is chosen to be $5$ as suggested by [].

#### D.1.3 Climate downscaling metrics

##### Root mean square error (RMSE)

This is the same as [Equation 3](#A4.E3).

##### Mean bias

Mean bias measures the difference between the spatial mean of the prediction and the spatial mean of the ground truth. A positive mean bias shows an overestimation, while a negative mean bias shows an underestimation of the mean value.

$$ $\text{Mean bias}=\frac{1}{N\times H\times W}\sum_{k=1}^{N}\sum_{i=1}^{H}\sum_{j=1}^{W}\tilde{X}-\frac{1}{N\times H\times W}\sum_{k=1}^{N}\sum_{i=1}^{H}\sum_{j=1}^{W}X$ (10) $$

##### Pearson coefficient

Pearson coefficient measures the correlation between the prediction and the ground truth. We first flatten the prediction and ground truth, and compute the metric as follows:

$$ $\rho_{\tilde{X},X}=\frac{\operatorname{cov}(\tilde{X},X)}{\sigma_{\tilde{X}}\sigma_{X}}$ (11) $$

### D.2 Results summary

Table LABEL:tab:rmse-compare and LABEL:tab:acc-compare summarize the global forecasting results of ClimaX and the baselines for all target variables and at all lead times. In addition to IFS and the two CNN-based baselines in the main text, we include FourCastNet [], PanguWeather [], and GraphCast [] for comprehensiveness. We want to emphasize that the results obtained by these methods are not comparable with ClimaX, as they were trained on ERA5 at $0.25\degree$, a much higher resolution compared to $5.625\degree$ and $1.40625\degree$ data used to train ClimaX. In [Section 4.5](#S4.SS5), we had a discussion on how the performance of ClimaX scales favorably with respect to data resolution. We hope this summary will provide future works with an easier comparison with existing baselines.

In spite of being trained on much lower resolutions, ClimaX outperforms FourCastNet in forecasting Z500, T850, and U10 at lead times from 3 days and beyond, in terms of both RMSE and ACC. For T2m, ClimaX achieves better results at horizons longer than 3 days. PanguWeather performs better than ClimaX on most of the tasks, but the gap between the two methods shrinks and becomes negligible as the lead time increases. ClimaX even outperforms PanguWeather in predicting U10 at 7 days lead times. This is because ClimaX is finetuned to perform direct prediction, which mitigates error accumulation for long horizon prediction. GraphCast achieves the lowest RMSE among all methods, but performs worse in terms of ACC compared to ClimaX and PanguWeather.

Table: Table 10: RMSE on global forecasting for different target variables at different lead times. Lower is better.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: Major: 
- Do you compute area-weighted RMSEs? If not, I strongly encourage fixing this (especially for the weather forecasting experiment; see e.g. Weatherbench2).
- For the fine-tuning experiments, do you start them based on the weights resulting from the first or second pre-training stage? If the former, then the second pre-training seems to also be a form of fine-tuning and I would suggest avoiding using the term "zero-shot" for reporting the forecasting results.
- From looking at Figure 3, I don't think that I agree with the author's interpretation that *"It is interesting that reconstruction performance is relatively little affected by lead time at the lower end of masking ratios"*. There's a clear cap between "0h, global" and "6h, global", especially for the lower end of masking ratios. Do I miss something?
- Can you include downscaling results on more variables than only T2M?

Minor:
- What's the difference between the encoder and decoder blocks? Fig. 1 suggests these are identical... Is the difference some densification by the "reconstruct batch" module in Fig. 1? If so, can you explain this more and make it clearer?
- Line 174: "reduce the masking ratio to 50%"... is this a typo (you use the same rate for the first stage)? What's correct?
- Why is there no reference line in figure 5b)?
- Please define what's meant by spatial and temporal RMSEs.
- Can you include snapshots like Fig. 9 but at the native temporal resolution used for prediction (i.e. not a monthly mean)?
- How does using different patch sizes (larger than the used size of 1) impact downscaling performance? 
- Why do you call the model Prithvi WxC?
soundness: 2
strengths: 1. Prithvi WxC is quite flexible as it can be used for a broad range of downstream applications, as convincingly shown in the experiments. 
2. The method contains original ideas such as the pre-training objective and using climatology-derived anomalies as targets, which I found interesting to read about.
3. The paper is generally clearly written and easy to read.
summary: This paper introduces a new foundation model, Prithvi WxC, for atmospheric modeling applications in weather and climate.  Prithvi WxC was trained on 3-hourly data from 1980 to 2019 from the MERRA-2 reanalysis dataset based on a masked reconstruction/forecasting pre-training objective. 
The model follows a transformer-based encoder-decoder architecture inspired by Hiera and MaxViT.
Afterward, the model is fine-tuned for various downstream tasks: Medium-range weather forecasting, global and regional downscaling, and learning a gravity wave flux parametrization. These tasks have different sets of spatial resolutions, variables, and datasets, showing the flexibility of the foundation model.
weaknesses: 1. The paper falls short of establishing a compelling case for Prithvi WxC as a foundation model for weather or climate. The practical significance and advantages of this approach remain inadequately demonstrated:

a.) While foundation models typically excel at zero-shot performance and data-efficient fine-tuning across diverse tasks, the evidence presented for Prithvi WxC's capabilities in these areas is not convincing. Baselines for the non-forecasting experiments are either very weak (interpolation-based downscaling) or non-existent (gravity wave experiments). Some highly relevant and simple baselines are: 
- How much worse(?) does Prithvi WxC perform on these tasks if you omit the pre-training stage (i.e. initialize with random weights instead of the frozen pre-trained ones, and train all parameters jointly from scratch on the tasks)? 
- How about completely removing the pre-trained transformer backbone (i.e. removing the Prithvi WxC block from Figures 12 & 13)? 
- For the latter, it would be also good to run an experiment where you replace the pre-trained Prithvi WxC backbone with some "lightweight" blocks (e.g. a (deeper) U-Net), trained in a task-specific way from scratch, to account for the huge difference in parameter counts if you completely remove Prithvi WxC. 

These ablations would immensely help in understanding how useful the pre-training stage is for these downstream applications (e.g. does using pre-trained Prithvi WxC improve performance over such simple baselines? Is it more data-efficient?). Besides, otherwise, it is hard to see evidence for the claim in the conclusion that *"Instead of building task-specific ML-models from scratch, these pretrained encoders can be used to develop more precise data-driven models of atmospheric processes"*.

b.) No ablations are included. I understand that training such a huge model is expensive but having a few ablations would have been very appreciated (perhaps, with a smaller-scale version of the model). For example:

- How crucial is it to predict climatology-normalized targets as opposed to normal per-variable means/stds? 
- What's the forecasting performance of Prithvi WxC after the first pre-training phase?
- How important is local vs global masking? What about the masking rates?
- What's the line of thought behind randomizing the distance between input timesteps? Can the model only use one input timestep? I presume this is possible by masking the corresponding snapshot by 100%, but no experiments with this setting are shown. 

c.) The weather forecasting results seem lukewarm, albeit it is hard to judge because the comparison is not apples-to-apples.
- Prithvi WxC is trained and evaluated on Merra-2. The baselines are evaluated on ERA5. These reanalysis datasets have different spatial resolutions. The evaluation years seem to be different too (correct me if I'm wrong). It would help to fix this mismatch. For example,  given the foundational nature of Prithvi WxC... why not fine-tune it on ERA5 directly? Showing that it can be competitive to these baselines in an apples-to-apples comparison would be a very strong result.
- Based on the mismatched comparison, Prithvi WxC seems to be competitive on 6h to 12h forecasts but it's quite notable that its performance implodes compared to the baselines for longer lead times. It is very unclear why. I wouldn't necessarily expect this version of Prithvi WxC to be state-of-the-art, but the performance does seem underwhelming. Especially given that the authors did "several things" to tune these results (i.e. a second forecasting-specific pre-training stage and autoregressive rollout fine-tuning).
- The hurricane evaluation includes hurricanes from 2017 to 2023. This seems to overlap with the training data period (up to 2019). 
- Either Figure 6 or its analysis in the main body of the text (lines 251-253) is wrong because I see all of the three models do best on exactly one of the three RMSE figures.
- For the hurricane forecasting experiments, I would appreciate a comparison to the state-of-the-art models included in the weather forecasting experiments (e.g. GraphCast) which have shown to be better than FourcastNet.

d.) The downscaling problem setup is artificial. Downscaling coarsened of existing reanalysis/model outputs is not of much use in practice. A realistic and important downscaling application, as discussed in the Appendix, would be to downscale coarse-resolution model outputs to high-resolution outputs (either of a different model, observations, or the same model run at higher resolution).
 
e.) The climate model parameterization experiments should be more carefully interpreted. 
- The model predicts outputs that are normalized by the 1980-2019 climatology. Unfortunately, decadal or centennial simulations of the future under a changing climate are inherently a non-stationary problem. It is highly unclear if Prithvi WxC would remain stable, let alone effective, under this highly relevant use case. This is particularly so as the in-the-loop (coupled to a running climate model) stability of ML-based climate model parameterizations is a well-known issue.
- The selling point for ML-based emulators of climate model parametrizations is often their computational cheapness. Thus, the runtime of Prithvi WxC should be discussed. Given the large parameter count of Prithvi WxC it might be important to note its runtime as a limitation for these kinds of applications.
- Line 461 claims that Prithvi WxC "outperforms" task-specific baselines but no baselines whatsoever are included in the manuscript for this experiment.
- Are the inputs a global map? I am not familiar with gravity waves, but I believe that most physics parameterizations in climate models are modeled column-wise (i.e. across atmospheric height but ignoring lat/lon interactions). This is surely a simplification of these parameterizations, but it seems to indicate that they're highly local problems. What's the motivation for using global context then?
- The end of the section should be worded more carefully, clearly stating the aforementioned limitations.

f.) No scaling experiments are included. Thus, it is unclear how important its 2.3 billion parameter size is, how well the model scales, and how its size impacts performance on the downstream applications. Besides, vision and language models are usually released with multiple model sizes that cover different use cases (e.g. balancing inference speed with accuracy). It would be really useful to get these (and carefully compare them) for Prithvi WxC.

2. Related work is insufficiently discussed. Please include an explicit section discussing it, focusing on:
- Carefully comparing similarities/differences to existing weather foundation models (e.g. architectures, pre-training objectives, downstream applications etc.). Besides, ClimaX is not properly discussed in the paper. Given that it's also a transformer-based foundation model, validated on forecasting, downscaling, and climate emulation, it is very important to include it in the comparison. 
- Similarly, please discuss how exactly the masking technique in this paper relates to the ones proposed in Vandal et al. and McNally et al..
- Carefully discuss how the architecture is derived from Hiera and/or MaxViT (and other papers of which components were derived, if any).

3. While the authors transparently discuss some issues/limitations with their experiments (e.g. the evaluation data mismatches), it would be nice to also include an explicit paragraph or section on this (and include aforementioned things like the issues with the climate model parameterization experiments).

Minor:
- Can you properly discuss, and include a reference to, what a Swin-shift is?
- Similarly, for the "pixel shuffle layers"
- Line 39: Pangu -> Pangu-Weather
- Line 48: Nowcasting should be lower-case
- Equation 1: Consider reformulating this as an objective/loss function.
- Also Eq. 1: What is $\hat{X}_t$? What is $\sigma_C$?
- Line 93: $\sigma^2_C = \sigma^2_C(X_t - C_t)$ doesn't make sense to me.
- Line 104: *" same 20 year period that we used for pretraining."* .... Do you mean 40 year period? If not, which 20-year period from the 40-year training period did you use?
- Line 157: Multiple symbols are undefined (e.g. $V_S$).
- Line 169: It's not entirely clear what "alternates" means in this context.
- Line 429: "baseline"... do you mean Prithvi WxC? 
- Line 507: "improved"... improved compared to what?
- Figure 12: Do you mean 'downscale' on the right "upscale" block?
- Sections D. 2.3 and D.2.4 in the appendix are literal copies of the corresponding paragraphs on pages 8 and 9. Please remove.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: ### Minor issue:  Weak baselines for downscaling
The comparison primarily involves interpolation-based methods and does not consider more advanced, AI-driven downscaling models or domain specific statistical/dynamical downscaling.
soundness: 2
strengths: ### Originality and significance 
Prithvi WxC pushes the foundation model concept in atmospheric science further by expanding beyond just forecasting— seen in earlier models foundation model Aurora. Its architecture and training approach enable it to tackle a variety of downstream tasks, such as downscaling and parameterization, making it a valuable tool for both short-term weather predictions and long-term climate modeling. With its modular design, Prithvi WxC can be adapted flexibly to new tasks that combines AI with physical climate models.

       
### Quality
The authors have thoroughly evaluated Prithvi WxC across a range of tasks, including zero-shot reconstruction, downscaling, and extreme event forecasting. The extensive validations across different downstream tasks support Prithvi WxC’s adaptability and effectiveness in diverse weather and climate applications.



### Clarity and open research
The paper is organized in a clear, logical flow, moving smoothly from motivation and background to model architecture, objectives, and results. Key ideas, like the mixed masking and forecasting objective, are presented in a way that makes the technical contributions accessible to both AI and climate science audiences. Code and comprehensive supplementary materials are provided.
summary: The paper presents Prithvi WxC, a new foundation model designed to support a wide range of weather and climate applications. Built on a large, transformer-based architecture, Prithvi WxC is trained on the extensive MERRA-2 dataset, which covers 160 variables capturing atmospheric data. The model is unique in its ability to address multiple tasks—including forecasting, downscaling, and parameterization—making it versatile in handling both regional and global weather patterns.
weaknesses: ### Reliance on a single reanalysis dataset (MERRA-2)
MERRA-2, with its relatively lower spatial resolution, is not commonly used in AI-driven weather and climate research, where higher-resolution datasets like ERA-5 are preferred for their superior predictive accuracy. The authors themselves acknowledge this limitation, citing the weaker performance of Prithvi WxC in hurricane track forecasting compared to the ERA-5-trained FourCastNet, attributing this discrepancy to MERRA-2's lower spatial resolution.      

Prithvi WxC’s exclusive training on the MERRA-2 reanalysis dataset raises questions about whether it truly qualifies as a foundation model. The narrower training base implies that the model may be learning a representation more specific to MERRA-2 characteristics, along with its biases and errors, rather than capturing a broader, more generalized understanding of weather and climate dynamics. By contrast, Aurora foundation model are pretrained on six diverse weather and climate datasets, including ERA-5, CMCC, IFS-HR, HRES Forecast, GFS Analysis, and GFS Forecasts, which span various sources like forecasts, analyses, reanalyses, and climate simulations [1]. This multi-source approach ensures a broader, more representative foundation that enhances versatility across diverse applications.Prithvi WxC would benefit from a similar multi-dataset training approach to strengthen its robustness and generalizability, which would then qualify it as a foundation model.

[1] Bodnar, C., Bruinsma, W. P., Lucic, A., Stanley, M., Brandstetter, J., Garvan, P., ... & Perdikaris, P. (2024). Aurora: A foundation model of the atmosphere. arXiv preprint arXiv:2405.13063.

### Ablation Study
While the authors propose a novel objective function, they only hypothetically attribute Prithvi WxC's strong short-term forecasting performance to its masking objective, without providing empirical evidence. This lack of testing weakens the claims about the model’s unique architecture and objective function. The observed performance could be influenced by several factors: the mixed objective itself, specific network structures or attention mechanisms, and choices made in the pretraining setup.

Without ablation experiments, the paper's assertions about the effectiveness of these innovations remain speculative, leaving readers uncertain about the impact of each component. An ablation study could isolate the contributions of these elements and would strengthen the paper by making its claims more concrete and providing clearer insights into Prithvi WxC’s architectural and training contributions.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: In addition to the weaknesses above, can the authors clarify the following:

1. How is the masking pre-training strategies better than just using variable lead-time that includes delta_t = 0? 
2. How is the local-global attention better than existing pre-training strategies of e.g., patch-variable-level tokenization employed in e.g., ClimaX? 
3. With masking as an additional pre-training strategy, what is the cost-performance tradeoff? 
4. Has the authors measure Prithvi's parameterization stability in an online-coupled setting?
soundness: 3
strengths: Operationally, the large-scale pre-training of 2.3B FM is impressive in the emerging field of AI4Science. The use of large-scale dataset is also noteworthy. The writing is clear and the downstream tasks are clearly defined, and touch upon some of the hardest challenges facing the field. The use of beyond-atmospheric variables (coupled ocean + land) is also welcomed to build a full Earth system FM.
summary: The paper introduces a new foundation model for weather and climate, called Prithvi WxC, with 2.3 billion parameters, and trained on relatively unconventional yet interesting reanalysis products of MERRA-2. The authors use a relatively novel pre-training strategy in this field, where in addition to forecasting-only pre-training, they also combine masking for reconstruction in the hope of better self-supervision, among several upsides (e.g., natural extension to data assimilation with sparsely-gridded observation). The FM is then evaluated on several downstream tasks, including forecasting, downscaling, and parameterization.
weaknesses: There are several major weaknesses in the paper, including non-existent/weak baselines and misleading claims, summarized below:

Major weaknesses
1. The term foundation model for this work is misleading, since in contrast to other similar works e.g., ClimaX, Aurora, the model is only trained on one data source which is MERRA-2. The problem: reanalysis products are by construction applicable for short-medium range applications as they are used either to evaluate NWPs or used as ICs for the next forecasting window. The climate dataset being missing, makes the claim of Prithvi being also a climate FM is an overreach at best. 

2. Related to the first point, the paper does not evaluate on any climate-related tests despite claiming it to be a climate FM too: (a) is the model stable over 50-100 years climate rollout? (b) what is the climate drift/bias compared to SOTA climate emulator? The paper applies downscaling to a climate dataset CORDEX, but this is less of a climate question and just a general downscaling problem since the former is more concerned about getting the long-term statistics, rather than a high-resolution state realization, correct (which is near impossible to nonlinear chaotic systems such as the Earth system). 

3. The downscaling benchmarking is also unacceptable as the baseline is too weak (e.g., bilinear, nearest neighbor). The author should either remove this part or add stronger baselines where the SOTA is at least a deep learning based method. Also, there is no benchmarking on the parameterization downstream task. At least use existing DL-based models from recent works. 

4. The forecasting performance appears to be unconvincing at best: with 2.3B parameter, it performs worse than e.g., <50M parameter GraphCast (~40x smaller) even at short lead time of 66 hours (<3 days). Even when the authors mention this result is "zero-shot" (which I find it unconvincing since there is rollout fine-tuning still) and the target is different (MERRA-2 vs ERA5), the obvious larger error growth (Figure 4) is alarming as it may not be useful for long-range climate forecasting. Also, why not benchmark against MERRA-FourCastNet in the forecasting task since this provides for a fairer comparison as both are trained with MERRA (as in the case for hurricane prediction). Finally, figure 4c (single line: forecasting cloud) is a case-in-point: the lack of sufficient apple-to-apple baselines where training/eval is done on MERRA-2. 

Overall, I find the results unconvincing given the lack of data sources, proper baselines, and inferior performance gain despite Prithvi being orders-of-magnitude larger than any SOTA model. As a side note: I believe the task of downscaling and parameterization is similar in that both attempts to resolve small-scale physics in an otherwise coarse-resolution model. I suggest the authors combine or use different downstream tasks e.g., climate projection.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Architectural flexibility: Can you provide a detailed explanation of the adaptability of the architecture when applied to non rectangular grid systems, such as those used for ocean simulation or polar regions?
Generalization strategy: What steps are taken to ensure the performance of the model when training or fine-tuning on datasets outside of MERRA-2, such as ERA5 or even higher resolution datasets?
Hurricane forecast: Can you provide more details on how the hurricane trajectory prediction of this model compares to specific task models such as FourCastNet, especially in areas with sparse data coverage?
soundness: 3
strengths: Originality: The introduction of basic models for weather and climate applications is novel and important. Unlike specific task models such as FourCastNet and GraphCast (Lam et al., 2022), Prithvi WxC addresses a wide range of tasks and effectively narrows the gap between artificial intelligence models for specific weather tasks and general artificial intelligence base models.
Quality: The mixed pre training objective of this model combines masking and prediction, which is robust, especially because it uses climate bias rather than just future state prediction, which enhances the adaptability of the model. The model also showed impressive results in zero sample evaluation of reconstruction and autoregressive prediction, outperforming the baseline in a short delivery cycle.
Meaning: The ability to generalize to multiple downstream tasks, such as downscaling and gravity wave parameterization, suggests that this model has the potential to have a significant impact on weather and climate modeling.
summary: The paper titled "Basic Models for Weather and Climate" introduces Prithvi WxC, a 2.3 billion parameter basic model for various weather and climate tasks. These include downscaling, autoregressive prediction, and extreme event estimation. The model is trained on 160 variables from the MERRA-2 dataset and combines mask reconstruction with prediction tasks to learn from various atmospheric data. Its encoder decoder architecture and ability to work with different spatial topologies make it suitable for global and regional weather modeling.
weaknesses: Generalization to other datasets: Although the model performs well on MERRA-2 data, its generalization ability to other datasets such as ERA5 or CMIP has not been fully explored. Validation on different datasets will better demonstrate its robustness.
Long term prediction: The accuracy of Prithvi WxC decreases with the extension of the prediction window, especially beyond 66 hours, and its performance is poor compared to models such as Pangu. A deeper investigation into how to maintain performance within an extended time frame can improve its utility in medium - and long-term forecasting.

Meta-review: The paper proposes a foundation model for weather and climate. While the foundation model backbone is similar to many other works in the realm of foundation models, the paper proposes interesting objectives based on the weather domain. The biggest weakness of the paper is in experiments that one of the reviewers discusses at length.

Decision: Reject