λ$\mathtt{\lambda}$ μ$\mathtt{\mu}$

# CaFA: Global Weather ForeCasting
with Factorized Attention on Sphere

Zijie Li†, Anthony Zhou†, Saurabh Patil†, Amir Barati
Farimani†111Correspondence<EMAIL_ADDRESS>
† Mechanical Engineering Department, Carnegie Mellon University

###### Abstract

Accurate weather forecasting is crucial in various sectors, impacting
decision-making processes and societal events. Data-driven approaches based on
machine learning models have recently emerged as a promising alternative to
numerical weather prediction models given their potential to capture physics
of different scales from historical data and the significantly lower
computational cost during the prediction stage. Renowned for its state-of-the-
art performance across diverse domains, the Transformer model has also gained
popularity in machine learning weather prediction. Yet applying Transformer
architectures to weather forecasting, particularly on a global scale is
computationally challenging due to the quadratic complexity of attention and
the quadratic increase in spatial points as resolution increases. In this
work, we propose a factorized-attention-based model tailored for spherical
geometries to mitigate this issue. More specifically, it utilizes multi-
dimensional factorized kernels that convolve over different axes where the
computational complexity of the kernel is only quadratic to the axial
resolution instead of overall resolution. The deterministic forecasting
accuracy of the proposed model on $1.5^{\circ}$ and 0-7 days’ lead time is on
par with state-of-the-art purely data-driven machine learning weather
prediction models. We also showcase the proposed model holds great potential
to push forward the Pareto front of accuracy-efficiency for Transformer
weather models, where it can achieve better accuracy with less computational
cost compared to Transformer based models with standard attention.

##  1 Introduction

Weather forecasting plays a pivotal role in numerous aspects of modern
society, spanning from agriculture and transportation to disaster preparedness
and public safety. The ability to predict weather conditions accurately allows
informed decision making in a wide array of engineering and social events.
Numerical weather prediction (NWP) (Bauer et al., 2015), a cornerstone of
modern meteorology, utilizes numerical models to simulate the Earth’s
atmosphere and predict the future weather. The simulation process consists of
constructing partial differential equations to describe the global weather
dynamics and then solving the equation via numerical algorithms. Resolving
finer physical scales, such as simulating microphysics phenomena or using a
finer spatio-temporal discretization, allows greater NWP accuracy; however,
this greatly increases the computational cost. Performing a 10-day forecast
can take an hour on a supercomputer (Bauer et al., 2020; Benjamin et al.,
2018). Additionally, accurately resolving small scale physics requires a well-
tailored parametrization of the numerical model that approximates the
underlying physical process.

The recent progress in machine learning and increasing amount of available
data have facilitated the emergence of machine-learning-based weather
prediction (MLWP) models. Compared to NWP, MLWP offers the potential to
capture and learn finer-scale physics from data without describing the physics
in an analytical mathematical model or searching for precise parameterization
fits. By ingesting datasets encompassing large amounts of historical weather
records, these models can learn and fit intricate patterns and correlations in
the global weather dynamics. Recently, MLWP models, particularly those based
on deep neural networks, have shown promising potential in global weather
prediction (Bi et al., 2023; Lam et al., 2023; Kurth et al., 2023; Chen et
al., 2023c, a; Nguyen et al., 2023a, b; Verma et al., 2024; Bonev et al.,
2023; Keisler, 2022). They have demonstrated competitive accuracy compared to
physics-based numerical models in medium-range weather forecasting. Once
trained, they can conduct prediction tasks very efficiently on computing
devices (e.g. graphic processing units) specialized for neural network
computation. For instance, a 10-day forecast at a scale of around 30 km can be
completed within a minute (Bi et al., 2023; Lam et al., 2023; Kurth et al.,
2023).

Among different types of neural network architectures, Transformer models
based on the attention mechanism (Vaswani et al., 2017) stands out as a
notable class. Initially designed for processing sequences in natural
language, Transformer models have found successful applications in a wide
array of engineering and scientific domains, such as image classification
(Dosovitskiy et al., 2021) and protein folding (Jumper et al., 2021). At its
essence, the Transformer’s attention mechanism provides the ability to capture
non-local patterns at all length scales within the context and employs data-
controlled weights to aggregate different features dynamically, with its
computational cost scaling quadratically with respect to the context size;
this can range from one to the full length of the input data. Prior efforts to
reduce the computational cost include linearizing the computational complexity
via kernel approximation (Choromanski et al., 2020), hardware-aware
optimization of the algorithm implementation (Dao et al., 2022), or
sparsifying the attention computation by restricting the context size through
partitioning the input into different windows (Beltagy et al., 2020; Liu et
al., 2021). For global weather forecasting tasks, previous Transformer-based
models apply attention to the input mostly with the spatial structure
flattened (Nguyen et al., 2023a, b). Sliding window attention can be used
afterwards to restrict the attention context (Chen et al., 2023c; Bi et al.,
2023; Chen et al., 2023a) and therefore reduce the computational cost. In
these cases, the multi-dimensional (longitude-latitude) spatial structure of
the input is not taken into account explicitly. While attending between every
pair of tokens that represent embedded features at different locations can
exploit all possible interactions in the data, its computation is
prohibitively expensive. Another drawback of discarding spatial structure is
that useful inductive biases about the underlying physics such as boundary
condition and geometry information (the target domain is on a sphere) can not
be easily applied, although some of this information can be implicitly encoded
to the model, e.g. via positional encoding.

In this work, we explore a different perspective of designing the attention
layer, where the multi-dimensional structure is preserved during the attention
calculation. We build our model with the low-rank attention kernel operator
(Li et al., 2024) that decomposes the multi-dimensional attention kernel into
a set of single-dimensional sub-kernels. The sub-kernels are parameterized
using dot-product attention with angular distance modulation which facilitates
a smooth response of the attention score under varying distances. We
numerically validate the deterministic forecasting capability of the proposed
model on WeatherBench2 (Rasp et al., 2023)’s data which has a resolution of
$1.5^{\circ}$. The proposed model - CaFA (ForeCast with Factorized Attention)
outperforms the deterministic operational Integrated Forecast System (IFS) on
most investigated variables and shows competitive performance over other
state-of-the-art end-to-end MLWPs in 0-7 days’ global weather forecasting. Our
study represents a stride towards reducing the computational cost of
Transformer-based weather models without sacrificing accuracy, thereby
enhancing their accessibility and efficiency.

##  2 Related Works

The advances in computing and general numerical methods for partial
differential equations have revolutionized numerical weather prediction over
the last century. The core of numerical modeling of weather lies in describing
the complex Earth dynamics including land, sea and atmosphere, as coupled
equation systems, or earth system models (ESM) (Hurrell et al., 2013). These
equation systems are often simulated via general circulation models (GCM)
(Phillips, 1956; Lynch, 2008; Satoh, 2013) numerically. Under-resolved physics
in GCMs (sub-grid physics that are beyond the scale of finest discretization)
are usually approximated with empirical parameterization schemes (Stensrud,
2007) and adjusted with domain expertise (Sherwood et al., 2014; Webb et al.,
2013). Not only have the numerical methods of climate/weather prediction been
greatly advanced, the amount of high-quality observed and analyzed data has
also increased vastly over the last century. One of the most notable instances
is historical data from the European Center for Medium-Range Weather Forecasts
(ECMWF) reanalysis v5 (ERA5) (Hersbach et al., 2020), which combines
historical observations with results from high-fidelity integrated Forecasting
System (IFS) (Wedi et al., 2015). These datasets offer an alternative for
advancements in climate and weather modeling, that is using data-driven models
instead of manually adjusting parameterizations to better capture underlying
physics. Various machine learning models have been employed to enhance the
predictive capabilities of complex atmospheric processes, including earlier
models that target specific variables/regions (Gagne et al., 2014; Lagerquist
et al., 2017) or global weather variables at relatively coarser resolutions
(Dueben and Bauer, 2018; Scher, 2018; Scher and Messori, 2019; Weyn et al.,
2019; Rasp and Thuerey, 2021; Weyn et al., 2020). More recently, Keisler
(2022) proposes a graph-neural-network-based model for forecasting multiple
surface and upper-air variables on $1^{\circ}$ resolution with accuracy
surpassing many prior MLWPs, which underscores the vast potential of MLWP in
enhancing weather prediction accuracy and efficiency. FourCastNet (Kurth et
al., 2023) demonstrates the promising forecasting capability and superior
inference efficiency of MLWP on $0.25^{\circ}$ resolution. Pangu-Weather (Bi
et al., 2023) and then GraphCast (Lam et al., 2023) further elevate the
accuracy of MLWP to the level such that their medium-range forecast (within 10
days) surpasses the state-of-the-art deterministic NWP on many weather
variables. They marked the first milestone where MLWP can outperform NWP in
terms of both deterministic forecasting accuracy and inference efficiency at
such a high resolution.

Each of these advancements are built upon a specific neural network
architecture. FourCastNet, which is built upon Adaptive Fourier Neural
Operator (Guibas et al., 2022; Li et al., 2021), employs spectral convolution
in the Fourier space to capture interactions in the spatial domain. Followup
work with the Spherical Fourier Neural Operator (Bonev et al., 2023) improves
and extends this framework by doing the convolution in spherical harmonic
space to better account for the underlying spherical geometry. GraphCast
scales up the graph neural network in Keisler (2022) and leverages multi-scale
message passing to better capture both local and non-local interaction between
different spatial points. Pangu-Weather is based on the Swin-Transformer (Liu
et al., 2021) that leverages sliding window attention to capture spatial
interactions. FuXi (Chen et al., 2023c) and Fengwu (Chen et al., 2023a), which
are also based on the Swin Transformer, improves the model’s long-term
forecasting accuracy with an improved training strategy and model
architectures. In parallel, ClimaX (Nguyen et al., 2023a) proposes a versatile
climate modelling framework with a modified Vision Transformer (ViT)
(Dosovitskiy et al., 2021) backbone and showcases it’s ability to address
various climate and weather tasks including global weather prediction. Stormer
(Nguyen et al., 2023b) scales up ClimaX’s Transformer architecture and
proposes to pretrain the model with dynamic lead times to further enhance the
long-term forecasting accuracy.

Beyond deterministic regression, the integration of deep-learning-based
generative models into weather modeling has also gained traction, including
ensemble prediction (Price et al., 2023; Li et al., 2023a), probabilistic
forecasting and nowcast (Leinonen et al., 2023; Sønderby et al., 2020; Nath et
al., 2023) and downscaling (Leinonen et al., 2021; Mardani et al., 2023; Chen
et al., 2023b). Moreover, machine learning models can also be combined with
physics priors to create a hybrid model, which can potentially further improve
the robustness and accuracy. A canonical paradigm is using machine learning
model to predict the under-resolved processes (Rasp et al., 2018) such as
various subgrid physics while using a numerical solver built upon governing
equation to solve the overall dynamics propagation. NeuralGCM (Kochkov et al.,
2023) presents a pioneering realization of this paradigm with performance on
par with state-of-the-art MLWP and NWP in terms of both deterministic and
ensemble forecast. Additionally, physics can be incorporated by formulating
the system propagation as a specific ordinary differential equation (Verma et
al., 2024) to guarantee certain physics constraints like momentum
conservation.

Figure 2.1: Main schematic of the proposed Transformer-based weather forecast
model - CaFA. The model approximates a Markovian mapping that forwards the
last system state to the next system state with a fixed time interval $\Delta
t$.

Figure 2.2: Main architecture of the model. (a) Each Transformer block in the
Processor comprises of a feed-forward network (abbreviated as "FFN", which is
a two-layer MLP) and a factorized attention layer. The factorized attention
layer contains three major components: projection, attention and tensor-matrix
product. (b) Cross-attention is used to project the latent embedding on a
coarse grid to the high-resolution grid (i.e. upsample). The query location
(at longitude $\theta$ and latitude $\varphi$) can be arbitrary points in the
domain. (c) Dot-product attention with relative distance encoding. (d) The
Encoder encodes the input surface variables and upper-air variables into
latent embeddings and projects them to a coarser latent grid via patchifying.
(e) The Decoder projects the latent grid back to the original grid and decodes
the latent embedding into surface and upper-air variables’ residual.

##  3 Methodology

###  3.1 Attention mechanism

The attention mechanism (Bahdanau et al., 2015; Vaswani et al., 2017) has
become the core component of modern neural network architectures,
revolutionizing various fields including natural language processing (Devlin
et al., 2019; Brown et al., 2020) and computer vision (Dosovitskiy et al.,
2021). Attention operates on three sets of vectors: queries
$\\{\mathbf{q}_{i}\\}_{i=1}^{N_{q}}$, keys
$\\{\mathbf{k}_{i}\\}_{i=1}^{N_{k}}$, and values
$\\{\mathbf{v}_{i}\\}_{i=1}^{N_{v}}$ ($N_{k}=N_{v}$) and computes the weighted
average of values:

$\mathbf{z}_{i}=\sum_{j=1}^{N_{v}}h(\mathbf{q}_{i},\mathbf{k}_{j})\mathbf{v}_{j},$
(1)

where $\mathbf{q}_{i},\mathbf{k}_{i},\mathbf{v}_{i}\in\mathbb{R}^{d}$,
$h:\mathbb{R}^{d}\times\mathbb{R}^{d}\mapsto\mathbb{R}$ is a suitable weight
function such as softmax with a scaling factor (Vaswani et al., 2017):
$h(\mathbf{q}_{i},\mathbf{k}_{j})=\exp(\mathbf{q}_{i}\cdot\mathbf{k}_{j}/\tau)/\sum_{s}\exp(\mathbf{q}_{i}\cdot\mathbf{k}_{s}/\tau)$
with scaling $\tau=\sqrt{d}$. For self-attention, the query/key/value is
transformed from the same input $\mathbf{u}_{i}$ through a learnable linear
transform:

$\mathbf{q}_{i}=\mathbf{u}_{i}W_{q},\mathbf{k}_{i}=\mathbf{u}_{i}W_{k},\mathbf{v}_{i}=\mathbf{u}_{i}W_{v},$
(2)

where $\mathbf{u}_{i}\in\mathbb{R}^{d_{\text{in}}}$ is the input vector and
$W_{q},W_{k},W_{v}\in\mathbb{R}^{d_{\text{in}}\times d}$ are learnable weight
matrices. In cross-attention, queries are usually obtained from a different
input source from the keys and values. In its continuous limit, attention in
Equation (1) parametrizes a kernel integral transform operator (Kovachki et
al., 2023; Cao, 2021). With suitable mesh weights it becomes a Riemann sum
approximation of the kernel integral transform:

$\mathbf{z}_{i}=\sum_{j=1}^{N_{v}}h(\mathbf{q}_{i}\cdot\mathbf{k}_{j})\mathbf{v}_{j}\mu_{j}\approx\int_{\Omega}\kappa\left(x_{i},\xi\right)v(\xi)d\xi,$
(3)

where $\mu_{j}$ is the mesh-based weight for grid point $x_{j}$, and the
weight function $h$ together with dot product of
$\mathbf{q}_{i},\mathbf{k}_{j}$ paramterizes the kernel function $\kappa$. The
application of the attention in modeling a diverse range of PDE-related
problems (Cao, 2021; Guo et al., 2023; Li et al., 2023b, 2024; Hao et al.,
2023; Liu et al., 2024; Hemmasian and Barati Farimani, 2023; Kissas et al.,
2022) suggests its potential for learning the complex global weather dynamics.
It is also observed that in Li et al. (2024) the attention matrices of a
Transformer trained on PDE problems often have a low-rank structure which
motivates us to explore a low-rank parametrization.

###  3.2 Axial factorized attention on sphere

#### Factorized attention

The original formulation of attention has a complexity $O(N^{2})$ (with
respect to the total number of grid points $N$) in terms of its floating point
operation. As the number of grid points grows exponentially with respect to
the dimension of the problem, the computational cost of applying attention to
higher dimensional problems is relatively expensive. To mitigate this issue,
Li et al. (2024) proposes an axial-factorized attention leveraging the
tensorized structure of data, which comprises two learnable integral
operators: axial projection and axial factorized attention. The projection
operator projects input functions with a high-dimensional domain $\Omega$ into
$n$ sub-functions with a single-dimensional domain
($\Omega=\Omega_{1}\times\Omega_{2}\ldots\times\Omega_{n}$), and then the
attention kernels are computed based on projected single-dimensional
functions. Below we will introduce their discretized forms.

The projection operator is defined as:

$\text{{Projection}:}\quad
U^{(m)}_{i^{(m)}}=\gamma\left(\sum_{j^{(1)}=1}^{S_{1}}\mu^{(1)}_{j^{(1)}}\ldots\sum_{j^{(m-1)}}^{S_{m-1}}\mu^{(m-1)}_{j^{(m-1)}}\sum_{j^{(m+1)}}^{S_{m+1}}\mu^{(m+1)}_{j^{(m+1)}}\ldots\sum_{j^{(n)}}^{S_{n}}\mu^{(n)}_{j^{(n)}}U_{j^{(1)}j^{(2)}\ldots
i^{(m)}\ldots j^{(n)}}\right),$ (4)

where $U^{(m)}\in\mathbb{R}^{S_{m}\times d}$ is the projected feature along
the $m$-th axis, $d$ is the number of feature channels, $U$ is obtained from
input tensor $Z\in\mathbb{R}^{S_{1}\times\ldots\times S_{n}\times d}$ via a
pointwise linear transformation, $S_{m}$ denotes the size of the grid that
discretizes sub-domain $\Omega_{m}$ with $\mu^{(m)}_{j^{(m)}}$ being the
corresponding mesh weights for grid point $j^{(m)}$, and
$\gamma:\mathbb{R}^{d}\mapsto\mathbb{R}^{d}$ is a pointwise two-layer multi-
layer perceptron.

Based on the projected functions with a single-dimensional domain, the (self)
axial factorized attention is computed as follows:

Axial attention kernel: $\displaystyle\quad
A^{(m)}=Q^{(m)}(K^{(m)})^{T},\text{where:
}Q^{(m)}=U^{(m)}W_{q}^{(m)},K^{(m)}=U^{(m)}W_{k}^{(m)},$ (5) Tensor-matrix
product: $\displaystyle\quad\left(Z_{\text{out}}\right)_{i^{(1)}i^{(2)}\ldots
i^{(n)}}$ $\displaystyle=$
$\displaystyle\sum_{i^{(1)}=1}^{S_{1}}\mu^{(1)}_{j^{(1)}}A^{(1)}_{i^{(1)}j^{(1)}}\sum_{i^{(2)}=1}^{S_{2}}\mu^{(2)}_{j^{(2)}}A^{(2)}_{i^{(2)}j^{(2)}}\ldots\sum_{i^{(n)}=1}^{S_{n}}\mu^{(n)}_{j^{(n)}}A^{(n)}_{i^{(n)}j^{(n)}}~{}V_{j^{(1)}j^{(2)}\ldots
j^{(n)}},$ (6)

where $A^{(m)}\in\mathbb{R}^{S_{m}\times S_{m}}$ is the $m$-th axial kernel
computed with dot product attention, $V$ is the value tensor which is derived
from the input via a learnable linear transformation following standard
attention (Equation (2)), and $W_{q}^{(m)},W_{k}^{(m)}$ are learnable linear
transformation matrices for computing queries and keys. In Li et al. (2024),
all the problems considered have a uniform equi-spaced grid, thus the mesh
weight on each grid point along $m$-th axis is $\mu^{(m)}_{j}=1/S_{m},\forall
j\in\\{1,2,\ldots,S_{m}\\}$ and the weighted sum in Equation (4) simply
amounts to average pooling over all but the $m$-th axis. With an equi-angular
grid on the spherical domain $\mathbb{S}^{2}$ and assuming uniform quadrature
rule, the mesh weight on each latitude grid point is non-uniform. Given a
latitude coordinate $\varphi\in[-\pi/2,\pi/2]$ and longitude coordinate
$\theta\in[0,2\pi)$ (and assuming unit radius), the integration of function
$f:\mathbb{S}^{2}\mapsto\mathbb{R}$ on the sphere takes the form:

$\int_{\mathbb{S}^{2}}fdS=\int_{0}^{2\pi}\int_{-\pi/2}^{\pi/2}f\cos\varphi~{}d\varphi
d\theta,$ (7)

where a suitable mesh weight for each latitude grid point (cell) $j$ is
$\mu^{(\text{lat})}_{j}=\frac{\pi}{N_{\text{lat}}}\cos\varphi_{j}$ and
$\mu^{(\text{lon})}_{j}=\frac{2\pi}{N_{\text{lon}}}$ for a longitude grid
point on an equi-angular grid. In summary, the vanilla attention-based kernel
integral in Equation (3) has now been modified to:

$z(\theta_{i},\varphi_{i})=\int_{0}^{2\pi}\kappa_{\text{lon}}(\theta_{i},\theta)\int_{-\pi/2}^{\pi/2}\kappa_{\text{lat}}(\varphi_{i},\varphi)v(\varphi)\cos\varphi~{}d\varphi
d\theta,$ (8)

with latitudinal kernel and longitudinal kernel
$\kappa_{\text{lon}},\kappa_{\text{lat}}$ parameterized using dot product
attention with angular distance modulation which will be described in the next
paragraph. Illustrative figures comparing factorized attention and standard
attention are provided in Appendix Figure C.3, C.4, C.5.

#### Distance encoding and positional encoding

The original dot-product attention mechanism does not explicitly take the
relative position between different points into account, but it is found that
modulating the attention with relative position information between tokens is
often beneficial (Su et al., 2024; Li et al., 2023b). In this work we propose
to use a learnable distance encoding function to directly modulate the dot
product between queries and keys:

$A^{(m)}_{ij}=\sum_{c=1}^{d}\psi^{(m)}_{c}(e^{(m)}_{ij})Q^{(m)}_{ic}K^{(m)}_{jc}$
(9)

where $Q^{(m)},K^{(m)}\in\mathbb{R}^{S_{m}\times d}$ are query and key
matrices of $m$-th axis (Equation (5)), $e^{(m)}_{ij}$ is the absolute angular
distance along $m$-th axis between point $i$ and $j$, and
$\psi^{(m)}_{c}:\mathbb{R}\mapsto\mathbb{R}$ is a learnable function. To
parameterize $\psi^{(m)}_{c}$, we use a Bessel basis function with learnable
coefficients (Gasteiger et al., 2019), which previously have shown to be very
effective for modelling atomic interactions (Gasteiger et al., 2019, 2021)
(here we omit superscript $(m)$ for better readability):

$\psi_{c}(e_{ij})=b_{c}+\sum_{n=1}^{N_{\text{basis}}}W_{nc}\sqrt{\frac{2}{\pi}}\frac{\sin(ne_{ij})}{e_{ij}},$
(10)

where $e_{ij}\in[0,\pi]$ (in practice we add a small number to the left
endpoint to avoid overflow), in which the longitude difference
$\theta_{ij}=|\theta_{i}-\theta_{j}|$ is wrapped into $[0,\pi]$:
$e_{ij}=\min(\theta_{ij},2\pi-\theta_{ij})$, and $W_{nc},b_{c}$ are a
learnable scalar coefficient and offset. For all the experiments, we set
$N_{\text{basis}}=64$ along the longitude direction and $N_{\text{basis}}=32$
along the latitude direction. After relative distance encoding, we propose to
apply a element-wise non-linearity function- leakyReLU (leaky rectified linear
unit) to the attention matrix in Equation (9), which suppresses the magnitude
of the negative part in the matrix and in practice we find this slightly
improves the performance. This also coincides with the observation in
Choromanski et al. (2020), where they use ReLU to replace Softmax in the
attention kernel and observe performance improvement.

In addition to the relative distance encoding, we also apply learnable
absolute positional encoding to the feature embedding before every Transformer
block. The learnable positional encoding is computed by first wrapping the
longitude/latitude coordinates with real-valued spherical harmonics and then
feeding them to a two-layer MLP. The spherical harmonics can be viewed as an
extension of the popular Fourier positional encoding (Vaswani et al., 2017) on
flat Euclidean space to the spherical domain $\mathbb{S}^{2}$, which generally
improves the performance of coordinate-based networks on spheres (Rußwurm et
al., 2024; Elhag et al., 2024). Example visualization of learned attention
kernels, distance and positional encoding are shown in Appendix C.

###  3.3 Grid projection

To reduce the computational cost and improve model scalability, the majority
of the proposed model (Processor) operates on a latent grid with reduced
resolution. For the height dimension (different pressure levels), the Encoder
will project 3D variables onto a 2D surface via a height compression module
and then in the Decoder the 3D structure of upper-air variables is recovered
via a height recovery module. For the spatial dimension (longitudinal and
latitudinal), the 2D grid is uniformly downsampled in the Encoder and then re-
projected to the original grid via cross-attention in the Decoder.

#### Height compression/recovery

The height compression is achieved via a 1D height-wise convolution filter,
with kernel size equaling the number of pressure levels. To recover the
height, we first expand the number of channels of the latent embedding by the
number of pressure levels via a linear transformation: $d\mapsto
N_{\text{levels}}d$, and then reshape the tensor from $N_{\text{lat}}\times
N_{\text{lon}}\times(N_{\text{levels}}d)$ to $N_{\text{lat}}\times
N_{\text{lon}}\times N_{\text{levels}}\times d$. After recovering the 3D
structure, we apply height-wise learnable linear modulation (Perez et al.,
2018) to the features: $\tilde{Z}_{\cdots lc}=\alpha_{lc}Z_{\cdots
lc}+\beta_{lc},\quad
l\in\\{1,2,\ldots,N_{\text{levels}}\\},c\in\\{1,2,\ldots,d\\}$, where
$\tilde{Z},Z\in\mathbb{R}^{N_{\text{lat}}\times N_{\text{lon}}\times
N_{\text{levels}}\times d}$ are 3D latent embeddings, and
$\alpha_{lc},\beta_{lc}$ are a learnable scalar scaling and offset.

#### Spatial downsample/upsample

We use patchification to downsample the spatial resolution, which is a widely
adopted technique in the vision Transformer (Dosovitskiy et al., 2021) to
downsample the data resolution and is also used in prior weather forecasting
Transformer models. The patchification amounts to applying non-overlapping
convolution to the input, with kernel size and stride equal to the patch size
$p$ (single-sided zero padding is used when the resolution is not divisible by
$p$). This process downsamples the grid resolution from $N_{\text{lat}}\times
N_{\text{lon}}$ to $(N_{\text{lat}}/p)\times(N_{\text{lon}}/p)$.

To recover the original resolution, we use factorized cross-attention (Figure
2.2(b)) to project the low-resolution latent grid back to the original
resolution, which is different from previous works that first apply linear
transformation to the channel dimension and then reshape it to "un-patchify"
the features. The proposed cross-attention module learns a continuous
representation (Mildenhall et al., 2020) for the target variables and allows
querying at arbitrary locations instead of a fixed resolution grid, which can
be useful for downstream applications such as data downscaling and regional
forecasting. The longitudinal and latitudinal basis functions are paramterized
by SIREN (Sitzmann et al., 2020) \- a MLP with a sine activation function,
which works well in approximating and representing continuous signals,
especially their high-frequency components.

###  3.4 Model overview

#### Architecture

We use an Encoder-Processor-Decoder (EPD) model architecture (Sanchez-Gonzalez
et al., 2020; Lam et al., 2023) (Figure 2.2). The model $F_{\theta}(\cdot)$
predicts the change of the variable $\Delta X$ over a fixed time interval
$\Delta t$ given $X_{t}$, i.e. $\hat{X}_{t+\Delta
t}=f(X_{t})=X_{t}+F_{\theta}(X_{t})$. Long-term forecasting is generated by an
auto-regressive rollout for multiple steps:

$\hat{X}_{t+N_{T}\Delta t}=\underbrace{f\circ f\circ\ldots\circ f}_{\times
N_{T}}(X_{t}).$ (11)

The Encoder takes the surface variables and upper-air variables as input and
fuses them together. The 3D upper-air variables are first projected onto a 2D
surface via the height compression operation and then concatenated with the
surface variable embedding. After concatenation, the surface and upper-air
variable embeddings are fused together via a two-layer MLP and then
downsampled to the reduced-resolution grid via patchfication.

The Processor is a stack of Transformer blocks (we use 6 blocks across all the
experiments) operating on the latent grid, where each block comprises of a
two-layer MLP (also known as feed forward network, FFN) followed by the
factorized attention layer. We place FFN in front of attention layer contrary
to many existing Transformer architectures that place FFN after the attetion
layer. The main reason is that we aim to fuse positional encoding with
features via non-linear transformation before doing the attention, which we
found performs slightly better. We use residual connections (He et al., 2016)
between different Transformer blocks and inside the attention layer (Figure
2.2 (a)). Layer-normalization (Ba et al., 2016) is applied after the residual
connection between each Transformer block (which is known as a post-norm
scheme in Transformer literatures). We add layer-normalization after
projecting features into queries and keys.

The Decoder takes the latent embedding output from the processor (represented
on a reduced resolution grid) as input and decodes it into the 2D surface and
3D upper-air variables. The latent embedding is first upsampled via cross-
attention, and then passed into two separate MLPs, where the first one is used
to decode the embedding to 2D surface variables and the second is used to
project and decode the 2D embedding to 3D upper-air variables.

#### Training

We train the model by employing a curriculum that first minimizes the single-
step prediction loss and then multi-step prediction loss, similar to prior
works (Lam et al., 2023; Chen et al., 2023a, c; Nguyen et al., 2023b; Kurth et
al., 2023). In the first stage, we train the model for 160k gradient steps,
where in the first 50k steps the loss is computed on a single step of
prediction-target pairs $(\hat{X}_{t+\Delta t},X_{t+\Delta t})$, and for the
remainng 110k steps the loss is computed based on a two-step rollout of
predictions $([\hat{X}_{t+\Delta t},\hat{X}_{t+2\Delta t}],[X_{t+\Delta
t},X_{t+2\Delta t}])$. During the second stage, we train the model for 50k
gradient steps and increase the model training rollout steps from $4$ to $20$
with a curriculum. The training loss function is a latitude and variable
weighted $L1$ loss. The motivation for choosing $L1$ loss over $L2$ loss is to
reduce the blurriness of the model’s prediction. However, it can make long-
term predictions to have higher MSE which we mitigate with multi-step loss
during training. The latitude weight and variable weight are computed similar
to Lam et al. (2023). The training of the final model on the $1.5^{\circ}$
resolution grid is carried out with 4 A6000 GPUs using Distributed Data
Parallel and float32 precision, which takes roughly a week. It takes around 6
seconds to generate the prediction up to a week (28 steps) on a single A6000
using float32 precision, with model wrapped by torch.compile. We provide a
brief summary of the model’s major hyperparameters in Table 3.1, and the model
has around 200 M trainable parameters in total. Further details about model
training are deferred to the Appendix A.1.

The proposed model is developed and validated based on the ERA5 reanalysis
data (Hersbach et al., 2020) from the European Centre for Medium-Range Weather
Forecasts (ECMWF), where we use the preprocessed version from WeatherBench2
(Rasp et al., 2023). Similar to GraphCast (Lam et al., 2023), during the model
development stage we train the model on $1979-2015$ and validate on $2016$.
The final model (on the $1.5^{\circ}$ grid) is trained on $1979-2019$ with
data from four additional recent years added to the training dataset.

Hyperparameter | Value
---|---
Base hidden dimension | 256
Processor hidden dimension | 768
Number of Transformer blocks | 6
Self-attention heads | 16
Self-attention head dimension | 128
Cross-attention heads | 32
Cross-attention head dimension | 128
Latent grid downsample ratio (patch size) | 2$\times$2
Table 3.1:  Major hyperparameter choices of the model. Base hidden dimension
denotes the width of MLPs and number of channels of the latent embedding in
Encoder and Decoder. In the patchification layer, the latent embedding’s
channel size is projected to Processor hidden dimension and re-projected back
in the Decoder’s upsampling (cross-attention) layer. For the FFN inside each
Transformer block, the expansion ratio is set to $6$, i.e. the hidden units
number is $6$ times the input channel number. Gaussian Linear Error Units
(GELUs)(Hendrycks and Gimpel, 2023) is used as the non-linear activation
inside every Transformer block.

##  4 Experiments and Results

###  4.1 Evaluation setting

Constants | Surface variables | Upper-air variables
---|---|---
Land sea mask |  $2$m temperature ($t2$m) | Geopotential ($zX$)
Soil type |  $10$m $u$ of wind ($u10$m) | Temperature ($tX$)
Angle of sub-gridscale orography |  $10$m $v$ of wind ($v10$m) |  $u$ of wind ($uX$)
Anisotropy of sub-gridscale orography | Mean sea level pressure (MSLP) |  $v$ of wind ($vX$)
| Total precipitation 6hr ($\text{TP}6$) | Specific humidity ($qX$)
| Total precipitation 24hr ($\text{TP}24$) |
Table 4.1: Weather variables used by the model. Surface and upper-air
variables are time-dependent variables predicted by the model given an initial
condition. $13$ pressure levels of upper-air variables are considered:
$50,100,150,200,250,300,400,500,600,700,850,925,1000~{}\text{hPa}$. $X$
denotes the pressure level in the abbreviation of upper-air variables, e.g.
$z500$ denotes geopotential at $500$hPa.

Variables and their corresponding abbreviations are listed in Table 4.1. The
time resolution is $\Delta t=6\text{hr}$. The major deterministic metrics are
root mean squared error (RMSE), anomaly correlated coefficient (ACC), and bias
(see Appendix A.2 for definitions). All the metrics are weighted by the
corresponding cell area based on latitude (Lam et al., 2023; Rasp et al.,
2023):

$w(i)=\frac{(\sin\varphi^{u}_{i}-\sin\varphi^{l}_{i})}{{\frac{1}{N_{\text{lat}}}}\sum_{j=1}^{N_{\text{lat}}}(\sin\varphi^{u}_{j}-\sin\varphi^{l}_{j})},$
(12)

where $\varphi_{u},\varphi_{l}$ are the corresponding upper/lower latitudinal
boundaries for each cell. We exclude the evaluation of total precipitation, as
prior works (Lam et al., 2023; Lavers et al., 2022) have suggested the bias of
ERA5 precipitation data and its skewed distribution makes deterministic metric
like RMSE less meaningful. For MLWP, all the metrics are computed by comparing
model’s prediction against ERA5 data. For IFS HRES and IFS ENS (ensemble
mean), their evaluation metrics are computed by comparing model’s prediction
against IFS analysis data following WeatherBench 2.

The major model is trained on a $1.5^{\circ}$ resolution equi-angular grid
($240\times 121$) and evaluated on the same grid following the the
deterministic evaluation protocol in WeatherBench 2 (Rasp et al., 2023). The
initial time for each prediction rollout is 00:00/12:00 UTC of every day in
year 2020. To compare our model to the open-sourced weather-forecast
Transformer - ClimaX (Nguyen et al., 2023a), a ViT (Dosovitskiy et al., 2021)
based model that is customized for climate related tasks with a competitive
accuracy on coarse resolution. We also train our model on a $64\times 32$
($5.625^{\circ}$) grid following the splitting and testing strategy in Nguyen
et al. (2023a). Other MLWP baseline models considered are Pangu-Weather
(Pangu) (Bi et al., 2023), FuXi (Chen et al., 2023c) and GraphCast (Lam et
al., 2023), which represent the state-of-the-art purely data-driven MLWP. All
three models are trained on a $0.25$ resolution grid with evaluation results
projected to a $1.5^{\circ}$ grid (Rasp et al., 2023). Pangu trains different
lead time models ($1h,3h,6h,24h$) and cascade them during inference to
generate a final prediction. FuXi scales up the number of Swin Transformer
(Liu et al., 2021) blocks compared to Pangu (48 vs 16). Instead of cascading
models with different lead times, FuXi uses a fixed $6h$ lead time but fine-
tunes separate models for different time scales and cascade them during
prediction. GraphCast is built upon a multi-resolution graph on the spherical
geometry and uses a model with lead time $6h$ to autogressively generate
predictions.

###  4.2 Discussion

The quantitative results for different models’ prediction accuracies are
presented in Table 4.2 and 4.2. For pressure variables ($z500$, MSLP), we
observe that CaFA consistently outperforms other MLWP baseline within the
range of 5 days. For other target variables, CaFA falls behind GraphCast and
FuXi for 1 day prediction but gradually catches up. We hypothesize the primary
reason is that compared to these models, CaFA is trained on lower resolution
data ($1.5^{\circ}$ vs $0.25^{\circ}$) and therefore does not have access to
finer scale physics which can be important for features like the wind velocity
field. The influence of finer-scale physics will gradually diminish as time
evolves. This is due to the chaotic nature of weather dynamics and
deterministic ML models will tend to approximate the statistical mean which
usually dampens the high frequency components (Bonavita, 2023). In general,
CaFA is on par with other MLWP baselines for prediction from 3 days up to a
week and performs particularly well in variables that features a decaying
spectrum such as geopotential. It is also a more lightweight Transformer-based
model compared to FuXi, which uses 48 Transformer blocks with hidden dimension
$1536$ (each block comprises an attention layer and a FFN) in total whereas
CaFA uses 6 Transformer blocks plus 1 cross-attention layer with hidden
dimension $768$.

lcccccccccccccccc Variables z500 [kg${}^{2}/$m2] t850 [K] t2m [K] MSLP [Pa]

Lead hours 24 72 120 168 24 72 120 168 24 72 120 168 24 72 120 168
IFS HRES 42 135 304 521 0.62 1.15 1.79 2.59 0.51 0.87 1.33 1.88

60 149 309 506

IFS ENS 42 132 277 439 0.63 1.09 1.59 2.14 0.56 0.86 1.20 1.58

62 146 280 423
CaFA(Ours) 39 123 269 449

0.60 0.99 1.56 2.28

0.55 0.84 1.24 1.76

48 130 271 438

Pangu 44 133 294 501

0.61 1.03 1.68 2.47

0.56 0.91 1.39 1.97

55 143 297 486
FuXi

40 125 276 433

0.54 0.95 1.56 2.11

0.53 0.84 1.27 1.65

49 132 276 416

GraphCast

39 124 274 468

0.51 0.93 1.53 2.29

0.51 0.82 1.25 1.81

48 132 276 454

Climatology 820 820 820 820 3.41 3.41 3.41 3.41 2.61 2.61 2.61 2.61 725 725
725 725

Table 4.2: Deterministic forecast RMSE on year $2020$’s data for temperature
and pressure at different lead times. Darker colors indicate better
performance. The evaluation results for baseline models are taken from
WeatherBench2 (Rasp et al., 2023).

lcccccccccccccccccccc Variables q700 [g/kg] u850 [m/s] v850 [m/s] u10m [m/s]
v10m [m/s]

Lead hours 24 72 120 168 24 72 120 168 24 72 120 168 24 72 120 168 24 72 120
168
IFS HRES 0.56 0.97 1.28 1.54

1.18 2.32 3.65 4.99

1.20 2.34 3.69 5.07

0.83 1.61 2.56 3.50

0.87 1.67 2.66 3.66
IFS ENS 0.51 0.84 1.06 1.22

1.14 2.10 3.13 4.03

1.16 2.12 3.16 4.08

0.81 1.47 2.19 2.82

0.84 1.51 2.27 2.95

CaFA(Ours) 0.55 0.84 1.07 1.30

1.11 1.98 3.09 4.24

1.12 2.00 3.13 4.30

0.69 1.34 2.16 2.98

0.72 1.39 2.24 3.12

Pangu 0.54 0.88 1.20 1.48

1.17 2.11 3.38 4.70

1.19 2.15 3.44 4.79

0.73 1.43 2.36 3.28

0.76 1.49 2.45 3.44

FuXi \- - - - 1.03 1.96 3.15 3.96

1.05 2.00 3.21 4.03

0.66 1.33 2.19 2.74

0.69 1.38 2.28 2.87

GraphCast 0.47 0.80 1.06 1.31

1.00 1.94 3.12 4.36

1.02 1.97 3.17 4.43

0.65 1.32 2.18 3.05

0.68 1.37 2.26 3.19

Climatology 1.61 1.61 1.61 1.61

5.66 5.66 5.66 5.66

5.47 5.47 5.47 5.47

3.83 3.83 3.83 3.83

3.96 3.96 3.96 3.96

Table 4.3: Deterministic forecast RMSE for specific humidity and wind speed
components at different lead times and pressure levels on year 2020’s data.

As shown in Figure 4.2 and 4.3, CaFA has lower RMSE and higher ACC than IFS
HRES from $1$ to $7$ days on most of the target variables. Comparing to the
ensemble mean of IFS forecasts, CaFA’s performance is on par or better up to 5
days, but after 5 days the ensemble mean generally has better accuracy. We
find that for CaFA, using a larger time step size can improve forecasting
accuracy beyond 7 days by trading off accuracy within 7 days range (the
comparison is shown in Appendix C.2). We also analyze the upper-air variables’
prediction of different pressure levels by comparing it to IFS HRES
prediction. Figure 4.1 shows that CaFA has better RMSE across most pressure
levels except for $50$hPa, where HRES consistently performs better. Such a
trend is also observed in GraphCast (Lam et al., 2023) and one of the reason
is that both CaFA and GraphCast are trained on loss that outweighs near-
surface (higher pressure) accuracy over smaller pressure levels’ accuracy. An
example rollout visualization of 3 key variables is shown in Figure 4.6,
qualitatively the models’ prediction agrees well with reference ERA5 analysis
data in terms of the general pattern. On the wind velocity field, it is
observed that the high-frequency details are blurrier as lead time increases.
Visualizations of other variables are shown in Appendix C.

We also compare our model to ClimaX (Nguyen et al., 2023a) 222ClimaX’s
official code for model and its training is publicly available (with MIT
license) at: https://github.com/microsoft/ClimaX/tree/main. to better
understand the influence of different attention mechanisms. To factor out the
influence of discretization (in space and time) and feature selection, we
slightly modify the original ViT-based architecture in ClimaX to train an
autoregressive model (the resulting model is called ClimaX-AR) with lead time
equal to $6$ hours using the weather variables specified in Table 4.1 (which
uses more atmospheric variables than original ClimaX). This ensures an exact
same training pipeline between CaFA and ClimaX-AR. We observe Climax-AR
performs notably better than the direct prediction version used in original
ClimaX, particularly at shorter lead times. As shown in Table 4.2, equipped
with factorized attention, CaFA performs consistently better on four key
variables for lead time 1-7 days, using significantly smaller amount of FLOPs.
We also provide a computational benchmark comparing factorized attention and
standard attention on different resolution grid in the Appendix B.3, which
shows that factorized attention greatly reduces the computational cost
associated with higher resolution grid. This highlights that the potential of
proposed factorized attention in improving the efficiency/accuracy balance of
Transformer-based weather models.

Figure 4.1:  Normalized RMSE difference:
$(\text{RMSE}_{\text{HRES}}-\text{RMSE}_{\text{CaFA}})/\text{RMSE}_{\text{HRES}}$.
Blue colors indicate IFS HRES has larger RMSE while red colors indicate CaFA
has larger RMSE. Darker colors indicate a larger normalized difference. The
plotting style follows GraphCast (Lam et al., 2023).

Figure 4.2: Comparison of CaFA and NWP in temperature and pressure’s
prediction of year 2020.  Figure 4.3: Comparison of CaFA and NWP in specific
humidity and wind velocity prediction of year 2020.

lccccccccccccccccc Variables z500 [kg${}^{2}/$m2] t850 [K] t2m [K] u10m [m/s]
TFLOPs

Lead hours 24 72 120 168 24 72 120 168 24 72 120 168 24 72 120 168

CaFA*(Ours) 51 170 348 544

0.59 1.04 1.71 2.47

0.49 0.82 1.28 1.79

0.58 1.28 2.15 2.97 0.61

ClimaX-AR 71 215 404 585

0.73 1.25 1.93 2.62

0.56 0.95 1.42 1.90

0.69 1.50 2.37 3.08 2.22

ClimaX* 96 244 440 599

1.11 1.59 2.23 2.77

1.10 1.43 1.83 2.18

1.41 2.18 2.94 3.43 \-

Table 4.4:  Comparison of accuracy and computational cost against ClimaX
(Nguyen et al., 2023a) on a $5.625^{\circ}$ resolution grid. The data
splitting follows the strategy used in ClimaX, where the test years are
$2017,2018$. ClimaX* is the direct prediction model reported in the original
paper (Table 10) and uses fewer atmospheric variables than our framework. T
(tera) FLOPs (number of floating point operations) of models’ forward
propagation is measured using the DeepSpeed profiling utility (Rasley et al.,
2020) with PyTorch 2.1 and a batch size of 8. CaFA* is a smaller version of
the major model, with a processor hidden dimension of $512$ ($2/3$ of the
number specified in Table 3.1) and the total trainable amount of parameters
therefore roughly matches ClimaX ($\sim$ 100M). The number of attention heads
is $16$ for both ClimaX and CaFA, while the hidden dimension of each attention
head is $128$ for CaFA and $64$ for ClimaX.

(a) $z500$ example rollout visualization

(a) $t2$m example rollout visualization

(a) $u10$m example rollout visualization Figure 4.6:  Example rollout
visualizations of the model’s prediction versus reference ERA5 reanalysis data
at different lead times. The initialization time is 00:00 UTC on August 11,
2020.

##  5 Conclusion

In this work we propose an end-to-end machine learning weather prediction
model based on factorized attention on a sphere. Our work presents a
computationally efficient methodology for computing a data-controlled
attention kernel that is capable of capturing long-range dependencies in the
domain. The factorized axial kernels only have a single-dimensional domain and
thus their computation is very efficient. The proposed attention mechanism
respects the underlying periodic boundary along longitudinal direction and
responds smoothly with respect to the angular distance between spatial points.
Our work signifies an advancement in enhancing the efficiency and accuracy of
Transformer-based weather prediction models.

Due to the limited computational budget, the finest resolution we have studied
in this work is a spatial resolution of $1.5^{\circ}$ and temporal resolution
of $6$ hours. Scaling the model to higher resolutions in both space and time
can potentially allow the model to capture finer scale physics and provide
better predictions. In parallel, it is interesting to study the scaling of
model itself to test how much accuracy improvement can be achieved by using
more compute. In addition, we have only explored using factorized attention to
model spatial interactions, which works particularly well for physical
processes that have a decaying spectrum such as geopotential. To further
improve the learning capability of higher-frequency components, it is also an
interesting direction to combine the proposed model with other localized
operators such as message passing around local neighborhoods.

In the current study, we only focus on investigating the short to medium range
deterministic forecast capability of the proposed model. However, due to the
chaotic nature of weather dynamics, deterministic and purely data-driven
models tend to generate blurry predictions in the long run as a result of
increasing uncertainty. The IFS ensemble prediction is generally observed to
have superior accuracy in longer range forecasting compared to deterministic
MLWP. We anticipate that the integration of probabilistic frameworks and
ensemble prediction that accounts for the inherent uncertainty will be an
important next step. It is also worth pointing out the current model is a
purely data-driven approach and it is not guaranteed to satisfy physics
constraints. Hence, incorporating more physics priors and symmetries into the
model can also be a meaningful future direction.

## References

  * Ba et al. (2016) J. L. Ba, J. R. Kiros, and G. E. Hinton.  Layer normalization.  _arXiv preprint arXiv:1607.06450_ , 2016.
  * Bahdanau et al. (2015) D. Bahdanau, K. Cho, and Y. Bengio.  Neural machine translation by jointly learning to align and translate.  In Y. Bengio and Y. LeCun, editors, _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_ , 2015.  URL http://arxiv.org/abs/1409.0473.
  * Bauer et al. (2015) P. Bauer, A. J. Thorpe, and G. Brunet.  The quiet revolution of numerical weather prediction.  _Nature_ , 525:47–55, 2015.  URL https://api.semanticscholar.org/CorpusID:4451289.
  * Bauer et al. (2020) P. Bauer, T. Quintino, N. Wedi, A. Bonanni, M. Chrust, W. Deconinck, M. Diamantakis, P. Düben, S. English, J. Flemming, P. Gillies, I. Hadade, J. Hawkes, M. Hawkins, O. Iffrig, C. Kühnlein, M. Lange, P. Lean, O. Marsden, A. Müller, S. Saarinen, D. Sarmany, M. Sleigh, S. Smart, P. Smolarkiewicz, D. Thiemert, G. Tumolo, C. Weihrauch, C. Zanna, and P. Maciel.  The ecmwf scalability programme: Progress and plans, 02/2020 2020.  URL https://www.ecmwf.int/node/19380.
  * Beltagy et al. (2020) I. Beltagy, M. E. Peters, and A. Cohan.  Longformer: The long-document transformer.  _arXiv preprint arXiv:2004.05150_ , 2020.
  * Benjamin et al. (2018) S. G. Benjamin, J. M. Brown, G. Brunet, P. Lynch, K. Saito, and T. W. Schlatter.  100 years of progress in forecasting and nwp applications.  _Meteorological Monographs_ , 59:13.1 – 13.67, 2018.  doi: 10.1175/AMSMONOGRAPHS-D-18-0020.1.  URL https://journals.ametsoc.org/view/journals/amsm/59/1/amsmonographs-d-18-0020.1.xml.
  * Bi et al. (2023) K. Bi, L. Xie, H. Zhang, X. Chen, X. Gu, and Q. Tian.  Accurate medium-range global weather forecasting with 3d neural networks.  _Nature_ , 619(7970):533–538, 2023.
  * Bonavita (2023) M. Bonavita.  On some limitations of data-driven weather forecasting models, 2023.
  * Bonev et al. (2023) B. Bonev, T. Kurth, C. Hundt, J. Pathak, M. Baust, K. Kashinath, and A. Anandkumar.  Spherical Fourier neural operators: Learning stable dynamics on the sphere.  In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 2806–2823. PMLR, 23–29 Jul 2023.  URL https://proceedings.mlr.press/v202/bonev23a.html.
  * Brown et al. (2020) T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei.  Language models are few-shot learners.  In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 1877–1901. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf.
  * Cao (2021) S. Cao.  Choose a transformer: Fourier or galerkin.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, _Advances in Neural Information Processing Systems_ , 2021.  URL https://openreview.net/forum?id=ssohLcmn4-r.
  * Chen et al. (2023a) K. Chen, T. Han, J. Gong, L. Bai, F. Ling, J.-J. Luo, X. Chen, L. Ma, T. Zhang, R. Su, et al.  Fengwu: Pushing the skillful global medium-range weather forecast beyond 10 days lead.  _arXiv preprint arXiv:2304.02948_ , 2023a.
  * Chen et al. (2023b) L. Chen, F. Du, Y. Hu, Z. Wang, and F. Wang.  Swinrdm: integrate swinrnn with diffusion model towards high-resolution and high-quality weather forecasting.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 37, pages 322–330, 2023b.
  * Chen et al. (2023c) L. Chen, X. Zhong, F. Zhang, Y. Cheng, Y. Xu, Y. Qi, and H. Li.  Fuxi: a cascade machine learning forecasting system for 15-day global weather forecast.  _npj Climate and Atmospheric Science_ , 6(1):190, 2023c.
  * Choromanski et al. (2020) K. M. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Q. Davis, A. Mohiuddin, L. Kaiser, et al.  Rethinking attention with performers.  In _International Conference on Learning Representations_ , 2020.
  * Dao et al. (2022) T. Dao, D. Fu, S. Ermon, A. Rudra, and C. Ré.  Flashattention: Fast and memory-efficient exact attention with io-awareness.  _Advances in Neural Information Processing Systems_ , 35:16344–16359, 2022.
  * Devlin et al. (2019) J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In J. Burstein, C. Doran, and T. Solorio, editors, _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.  doi: 10.18653/v1/N19-1423.  URL https://aclanthology.org/N19-1423.
  * Dosovitskiy et al. (2021) A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=YicbFdNTTy.
  * Dueben and Bauer (2018) P. D. Dueben and P. Bauer.  Challenges and design choices for global weather and climate models based on machine learning.  _Geosci. Model Dev._ , 2018.  doi: 10.5194/gmd-2018-148.  URL https://www.geosci-model-dev-discuss.net/gmd-2018-148/gmd-2018-148.pdf.
  * Elhag et al. (2024) A. A. A. Elhag, Y. Wang, J. M. Susskind, and M. Á. Bautista.  Manifold diffusion fields.  In _The Twelfth International Conference on Learning Representations_ , 2024.  URL https://openreview.net/forum?id=BZtEthuXRF.
  * Gagne et al. (2014) D. J. Gagne, A. McGovern, and M. Xue.  Machine learning enhancement of storm-scale ensemble probabilistic quantitative precipitation forecasts.  _Weather and Forecasting_ , 29(4):1024 – 1043, 2014.  doi: 10.1175/WAF-D-13-00108.1.  URL https://journals.ametsoc.org/view/journals/wefo/29/4/waf-d-13-00108_1.xml.
  * Gasteiger et al. (2019) J. Gasteiger, J. Groß, and S. Günnemann.  Directional message passing for molecular graphs.  In _International Conference on Learning Representations_ , 2019.
  * Gasteiger et al. (2021) J. Gasteiger, F. Becker, and S. Günnemann.  Gemnet: Universal directional graph neural networks for molecules.  _Advances in Neural Information Processing Systems_ , 34:6790–6802, 2021.
  * Guibas et al. (2022) J. Guibas, M. Mardani, Z. Li, A. Tao, A. Anandkumar, and B. Catanzaro.  Efficient token mixing for transformers via adaptive fourier neural operators.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=EXHG-A3jlM.
  * Guo et al. (2023) R. Guo, S. Cao, and L. Chen.  Transformer meets boundary value inverse problems.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=HnlCZATopvr.
  * Hao et al. (2023) Z. Hao, Z. Wang, H. Su, C. Ying, Y. Dong, S. Liu, Z. Cheng, J. Song, and J. Zhu.  Gnot: A general neural operator transformer for operator learning.  In _International Conference on Machine Learning_ , pages 12556–12569. PMLR, 2023.
  * He et al. (2016) K. He, X. Zhang, S. Ren, and J. Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778, 2016.
  * Hemmasian and Barati Farimani (2023) A. Hemmasian and A. Barati Farimani.  Reduced-order modeling of fluid flows with transformers.  _Physics of Fluids_ , 35(5):057126, 05 2023.  ISSN 1070-6631.  doi: 10.1063/5.0151515.  URL https://doi.org/10.1063/5.0151515.
  * Hendrycks and Gimpel (2023) D. Hendrycks and K. Gimpel.  Gaussian error linear units (gelus), 2023.
  * Hersbach et al. (2020) H. Hersbach, B. Bell, P. Berrisford, S. Hirahara, A. Horányi, J. Muñoz-Sabater, J. Nicolas, C. Peubey, R. Radu, D. Schepers, A. Simmons, C. Soci, S. Abdalla, X. Abellan, G. Balsamo, P. Bechtold, G. Biavati, J. Bidlot, M. Bonavita, G. De Chiara, P. Dahlgren, D. Dee, M. Diamantakis, R. Dragani, J. Flemming, R. Forbes, M. Fuentes, A. Geer, L. Haimberger, S. Healy, R. J. Hogan, E. Hólm, M. Janisková, S. Keeley, P. Laloyaux, P. Lopez, C. Lupu, G. Radnoti, P. de Rosnay, I. Rozum, F. Vamborg, S. Villaume, and J.-N. Thépaut.  The era5 global reanalysis.  _Quarterly Journal of the Royal Meteorological Society_ , 146(730):1999–2049, 2020.  doi: https://doi.org/10.1002/qj.3803.  URL https://rmets.onlinelibrary.wiley.com/doi/abs/10.1002/qj.3803.
  * Ho et al. (2020) J. Ho, N. Kalchbrenner, D. Weissenborn, and T. Salimans.  Axial attention in multidimensional transformers, 2020.  URL https://openreview.net/forum?id=H1e5GJBtDr.
  * Hurrell et al. (2013) J. W. Hurrell, M. M. Holland, P. R. Gent, S. Ghan, J. E. Kay, P. J. Kushner, J.-F. Lamarque, W. G. Large, D. Lawrence, K. Lindsay, W. H. Lipscomb, M. C. Long, N. Mahowald, D. R. Marsh, R. B. Neale, P. Rasch, S. Vavrus, M. Vertenstein, D. Bader, W. D. Collins, J. J. Hack, J. Kiehl, and S. Marshall.  The community earth system model: A framework for collaborative research.  _Bulletin of the American Meteorological Society_ , 94(9):1339 – 1360, 2013.  doi: 10.1175/BAMS-D-12-00121.1.  URL https://journals.ametsoc.org/view/journals/bams/94/9/bams-d-12-00121.1.xml.
  * Jumper et al. (2021) J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakool, R. Bates, A. Žídek, A. Potapenko, et al.  Highly accurate protein structure prediction with alphafold.  _Nature_ , 596(7873):583–589, 2021.
  * Keisler (2022) R. Keisler.  Forecasting global weather with graph neural networks.  _arXiv preprint arXiv:2202.07575_ , 2022.
  * Kingma and Ba (2015) D. Kingma and J. Ba.  Adam: A method for stochastic optimization.  In _International Conference on Learning Representations (ICLR)_ , San Diega, CA, USA, 2015.
  * Kissas et al. (2022) G. Kissas, J. H. Seidman, L. F. Guilhoto, V. M. Preciado, G. J. Pappas, and P. Perdikaris.  Learning operators with coupled attention.  _Journal of Machine Learning Research_ , 23(215):1–63, 2022.
  * Kochkov et al. (2023) D. Kochkov, J. Yuval, I. Langmore, P. Norgaard, J. Smith, G. Mooers, J. Lottes, S. Rasp, P. Düben, M. Klöwer, et al.  Neural general circulation models.  _arXiv preprint arXiv:2311.07222_ , 2023.
  * Kovachki et al. (2023) N. Kovachki, Z. Li, B. Liu, K. Azizzadenesheli, K. Bhattacharya, A. Stuart, and A. Anandkumar.  Neural operator: Learning maps between function spaces with applications to pdes.  _Journal of Machine Learning Research_ , 24(89):1–97, 2023.
  * Kurth et al. (2023) T. Kurth, S. Subramanian, P. Harrington, J. Pathak, M. Mardani, D. Hall, A. Miele, K. Kashinath, and A. Anandkumar.  Fourcastnet: Accelerating global high-resolution weather forecasting using adaptive fourier neural operators.  In _Proceedings of the platform for advanced scientific computing conference_ , pages 1–11, 2023.
  * Lagerquist et al. (2017) R. Lagerquist, A. McGovern, and T. Smith.  Machine learning for real-time prediction of damaging straight-line convective wind.  _Weather and Forecasting_ , 32(6):2175–2193, 12 2017.  ISSN 15200434.  doi: 10.1175/WAF-D-17-0038.1.
  * Lam et al. (2023) R. Lam, A. Sanchez-Gonzalez, M. Willson, P. Wirnsberger, M. Fortunato, F. Alet, S. Ravuri, T. Ewalds, Z. Eaton-Rosen, W. Hu, et al.  Learning skillful medium-range global weather forecasting.  _Science_ , 382(6677):1416–1421, 2023.
  * Larsson et al. (2016) G. Larsson, M. Maire, and G. Shakhnarovich.  Fractalnet: Ultra-deep neural networks without residuals.  In _International Conference on Learning Representations_ , 2016.
  * Lavers et al. (2022) D. A. Lavers, A. Simmons, F. Vamborg, and M. J. Rodwell.  An evaluation of era5 precipitation for climate monitoring.  _Quarterly Journal of the Royal Meteorological Society_ , 148(748):3152–3165, 2022.  doi: https://doi.org/10.1002/qj.4351.  URL https://rmets.onlinelibrary.wiley.com/doi/abs/10.1002/qj.4351.
  * Leinonen et al. (2021) J. Leinonen, D. Nerini, and A. Berne.  Stochastic super-resolution for downscaling time-evolving atmospheric fields with a generative adversarial network.  _IEEE Transactions on Geoscience and Remote Sensing_ , 59(9):7211–7223, 2021.  doi: 10.1109/TGRS.2020.3032790.
  * Leinonen et al. (2023) J. Leinonen, U. Hamann, D. Nerini, U. Germann, and G. Franch.  Latent diffusion models for generative precipitation nowcasting with accurate uncertainty quantification.  _arXiv preprint arXiv:2304.12891_ , 2023.
  * Li et al. (2023a) L. Li, R. Carver, I. Lopez-Gomez, F. Sha, and J. Anderson.  Seeds: Emulation of weather forecast ensembles with diffusion models.  _arXiv preprint arXiv:2306.14066_ , 2023a.
  * Li et al. (2021) Z. Li, N. B. Kovachki, K. Azizzadenesheli, B. liu, K. Bhattacharya, A. Stuart, and A. Anandkumar.  Fourier neural operator for parametric partial differential equations.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=c8P9NQVtmnO.
  * Li et al. (2023b) Z. Li, K. Meidani, and A. B. Farimani.  Transformer for partial differential equations’ operator learning.  _Transactions on Machine Learning Research_ , 2023b.  ISSN 2835-8856.  URL https://openreview.net/forum?id=EPPqt3uERT.
  * Li et al. (2024) Z. Li, D. Shu, and A. Barati Farimani.  Scalable transformer for pde surrogate modeling.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Liu et al. (2024) X. Liu, B. Xu, S. Cao, and L. Zhang.  Mitigating spectral bias for the multiscale operator learning.  _Journal of Computational Physics_ , 506:112944, 2024.  ISSN 0021-9991.  doi: https://doi.org/10.1016/j.jcp.2024.112944.  URL https://www.sciencedirect.com/science/article/pii/S0021999124001931.
  * Liu et al. (2021) Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, and B. Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pages 10012–10022, 2021.
  * Loshchilov and Hutter (2019) I. Loshchilov and F. Hutter.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=Bkg6RiCqY7.
  * Lynch (2008) P. Lynch.  The origins of computer weather prediction and climate modeling.  _Journal of Computational Physics_ , 227(7):3431–3444, 2008.  ISSN 0021-9991.  doi: https://doi.org/10.1016/j.jcp.2007.02.034.  URL https://www.sciencedirect.com/science/article/pii/S0021999107000952.  Predicting weather, climate and extreme events.
  * Mardani et al. (2023) M. Mardani, N. Brenowitz, Y. Cohen, J. Pathak, C.-Y. Chen, C.-C. Liu, A. Vahdat, K. Kashinath, J. Kautz, and M. Pritchard.  Generative residual diffusion modeling for km-scale atmospheric downscaling.  _arXiv preprint arXiv:2309.15214_ , 2023.
  * Mildenhall et al. (2020) B. Mildenhall, P. P. Srinivasan, M. Tancik, J. T. Barron, R. Ramamoorthi, and R. Ng.  Nerf: Representing scenes as neural radiance fields for view synthesis.  In _European Conference on Computer Vision_ , pages 405–421. Springer, 2020.
  * Nath et al. (2023) P. Nath, P. Shukla, and C. Quilodrán-Casas.  Forecasting tropical cyclones with cascaded diffusion models.  _arXiv preprint arXiv:2310.01690_ , 2023.
  * Nguyen et al. (2023a) T. Nguyen, J. Brandstetter, A. Kapoor, J. K. Gupta, and A. Grover.  ClimaX: A foundation model for weather and climate.  In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 25904–25938. PMLR, 23–29 Jul 2023a.  URL https://proceedings.mlr.press/v202/nguyen23a.html.
  * Nguyen et al. (2023b) T. Nguyen, R. Shah, H. Bansal, T. Arcomano, S. Madireddy, R. Maulik, V. Kotamarthi, I. Foster, and A. Grover.  Scaling transformer neural networks for skillful and reliable medium-range weather forecasting.  _arXiv preprint arXiv:2312.03876_ , 2023b.
  * Paszke et al. (2019) A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ , 32, 2019.
  * Perez et al. (2018) E. Perez, F. Strub, H. De Vries, V. Dumoulin, and A. Courville.  Film: Visual reasoning with a general conditioning layer.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 32, 2018.
  * Phillips (1956) N. A. Phillips.  The general circulation of the atmosphere: A numerical experiment.  _Quarterly Journal of the Royal Meteorological Society_ , 82(354):535–539, Oct. 1956.  doi: 10.1002/qj.49708235422.
  * Price et al. (2023) I. Price, A. Sanchez-Gonzalez, F. Alet, T. Ewalds, A. El-Kadi, J. Stott, S. Mohamed, P. Battaglia, R. Lam, and M. Willson.  Gencast: Diffusion-based ensemble forecasting for medium-range weather.  _arXiv preprint arXiv:2312.15796_ , 2023.
  * Rabe and Staats (2021) M. N. Rabe and C. Staats.  Self-attention does not need $o(n^{2})$ memory.  _arXiv preprint arXiv:2112.05682_ , 2021.
  * Rasley et al. (2020) J. Rasley, S. Rajbhandari, O. Ruwase, and Y. He.  Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters.  In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, KDD ’20, page 3505–3506, New York, NY, USA, 2020. Association for Computing Machinery.  ISBN 9781450379984.  doi: 10.1145/3394486.3406703.  URL https://doi.org/10.1145/3394486.3406703.
  * Rasp and Thuerey (2021) S. Rasp and N. Thuerey.  Data-driven medium-range weather prediction with a resnet pretrained on climate simulations: A new model for weatherbench.  _Journal of Advances in Modeling Earth Systems_ , 13(2):e2020MS002405, 2021.  doi: https://doi.org/10.1029/2020MS002405.  URL https://agupubs.onlinelibrary.wiley.com/doi/abs/10.1029/2020MS002405.  e2020MS002405 2020MS002405.
  * Rasp et al. (2018) S. Rasp, M. S. Pritchard, and P. Gentine.  Deep learning to represent subgrid processes in climate models.  _Proceedings of the National Academy of Sciences_ , 115(39):9684–9689, 2018.  doi: 10.1073/pnas.1810286115.  URL https://www.pnas.org/doi/abs/10.1073/pnas.1810286115.
  * Rasp et al. (2023) S. Rasp, S. Hoyer, A. Merose, I. Langmore, P. Battaglia, T. Russel, A. Sanchez-Gonzalez, V. Yang, R. Carver, S. Agrawal, M. Chantry, Z. B. Bouallegue, P. Dueben, C. Bromberg, J. Sisk, L. Barrington, A. Bell, and F. Sha.  Weatherbench 2: A benchmark for the next generation of data-driven global weather models, 2023.
  * Rußwurm et al. (2024) M. Rußwurm, K. Klemmer, E. Rolf, R. Zbinden, and D. Tuia.  Geographic location encoding with spherical harmonics and sinusoidal representation networks.  In _The Twelfth International Conference on Learning Representations_ , 2024.  URL https://openreview.net/forum?id=PudduufFLa.
  * Sanchez-Gonzalez et al. (2020) A. Sanchez-Gonzalez, J. Godwin, T. Pfaff, R. Ying, J. Leskovec, and P. Battaglia.  Learning to simulate complex physics with graph networks.  In _International conference on machine learning_ , pages 8459–8468. PMLR, 2020.
  * Satoh (2013) M. Satoh.  _Atmospheric circulation dynamics and general circulation models_.  Springer Science & Business Media, 2013.
  * Scher (2018) S. Scher.  Toward data-driven weather and climate forecasting: Approximating a simple general circulation model with deep learning.  _Geophysical Research Letters_ , 45(22):12,616–12,622, 2018.  doi: https://doi.org/10.1029/2018GL080704.  URL https://agupubs.onlinelibrary.wiley.com/doi/abs/10.1029/2018GL080704.
  * Scher and Messori (2019) S. Scher and G. Messori.  Generalization properties of neural networks trained on Lorenzsystems.  _Nonlinear Processes in Geophysics Discussions_ , pages 1–19, 6 2019.  ISSN 2198-5634.  doi: 10.5194/npg-2019-23.  URL https://www.nonlin-processes-geophys-discuss.net/npg-2019-23/.
  * Sherwood et al. (2014) S. C. Sherwood, S. Bony, and J.-L. Dufresne.  Spread in model climate sensitivity traced to atmospheric convective mixing.  _Nature_ , 505(7481):37–42, 2014.
  * Sitzmann et al. (2020) V. Sitzmann, J. Martel, A. Bergman, D. Lindell, and G. Wetzstein.  Implicit neural representations with periodic activation functions.  _Advances in neural information processing systems_ , 33:7462–7473, 2020.
  * Sønderby et al. (2020) C. K. Sønderby, L. Espeholt, J. Heek, M. Dehghani, A. Oliver, T. Salimans, S. Agrawal, J. Hickey, and N. Kalchbrenner.  Metnet: A neural weather model for precipitation forecasting.  _arXiv preprint arXiv:2003.12140_ , 2020.
  * Srivastava et al. (2014) N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov.  Dropout: A simple way to prevent neural networks from overfitting.  _Journal of Machine Learning Research_ , 15(56):1929–1958, 2014.  URL http://jmlr.org/papers/v15/srivastava14a.html.
  * Stensrud (2007) D. J. Stensrud.  _Parameterization Schemes: Keys to Understanding Numerical Weather Prediction Models_.  Cambridge University Press, 2007.
  * Su et al. (2024) J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Y. Liu.  Roformer: Enhanced transformer with rotary position embedding.  _Neurocomputing_ , 568:127063, 2024.
  * Vaswani et al. (2017) A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polosukhin.  Attention is all you need.  In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.
  * Verma et al. (2024) Y. Verma, M. Heinonen, and V. Garg.  ClimODE: Climate and weather forecasting with physics-informed neural ODEs.  In _The Twelfth International Conference on Learning Representations_ , 2024.  URL https://openreview.net/forum?id=xuY33XhEGR.
  * Webb et al. (2013) M. J. Webb, F. H. Lambert, and J. M. Gregory.  Origins of differences in climate sensitivity, forcing and feedback in climate models.  _Climate Dynamics_ , 40:677–707, 2013.
  * Wedi et al. (2015) N. Wedi, P. Bauer, W. Denoninck, M. Diamantakis, M. Hamrud, C. Kuhnlein, S. Malardel, K. Mogensen, G. Mozdzynski, and P. Smolarkiewicz.  _The modelling infrastructure of the Integrated Forecasting System: Recent advances and future challenges_.  European Centre for Medium-Range Weather Forecasts, 2015.
  * Weyn et al. (2019) J. A. Weyn, D. R. Durran, and R. Caruana.  Can machines learn to predict weather? Using deep learning to predict gridded 500-hPa geopotential height from historical weather data.  _Journal of Advances in Modeling Earth Systems_ , page 2019MS001705, 7 2019.  ISSN 1942-2466.  doi: 10.1029/2019MS001705.  URL https://onlinelibrary.wiley.com/doi/abs/10.1029/2019MS001705.
  * Weyn et al. (2020) J. A. Weyn, D. R. Durran, and R. Caruana.  Improving data-driven global weather prediction using deep convolutional neural networks on a cubed sphere.  _Journal of Advances in Modeling Earth Systems_ , 12(9):e2020MS002109, 2020.

_Supplementary Material_

###### Contents

  1. 1 Introduction
  2. 2 Related Works
  3. 3 Methodology
    1. 3.1 Attention mechanism
    2. 3.2 Axial factorized attention on sphere
    3. 3.3 Grid projection
    4. 3.4 Model overview
  4. 4 Experiments and Results
    1. 4.1 Evaluation setting
    2. 4.2 Discussion
      1. 5 Conclusion
        1. A Model implementation
          1. A.1 Training details
            1. A.2 Metrics
              1. B Comparison against standard attention
              2. C Further results

##  Appendix A Model implementation

###  A.1 Training details

We specify the detailed hyperparamters for the model training in Table A.1 and
A.2. We implement the model and the experiment using PyTorch (Paszke et al.,
2019). We use AdamW optimizer (Kingma and Ba, 2015; Loshchilov and Hutter,
2019) and the majority of the hyperparameter setting follows ClimaX (Nguyen et
al., 2023a). Regularization techniques including dropout (Srivastava et al.,
2014) and droppath (Larsson et al., 2016) are not used in the proposed model
(which is used in ClimaX baseline). For testing, we evaluate the final model
checkpoint and no early stopping is applied. During second-stage training,
gradient checkpointing is used to reduce the memory cost of backpropagating
through time. In addition, only the most recent 20 years’ (up to the
validation years) training data is used for second stage training.

Training Hyperparameter | Value
---|---
Optimizer | AdamW
Learning rate | $1e-8\mapsto 3e-4\mapsto 1e-7$
Batch size | 8
Gradient steps | 160k
Learning rate scheduler | Cosine Annealing with
| Linear Warmup (for 1600 steps)
Weight decay | 1e-6
Rollout step curriculum | $[1,2]$
Curriculum milestone | $[0,50\text{k}]$
EMA | None
Table A.1:  Training hyperparameter for stage 1. On $64\times 32$ grid we adjust the batch size to $16$ and peak learning rate to $5e-4$. The curriculum milestone determines how many steps to rollout after specific milestone. In the above setting, after 50k gradient steps, the model will rollout for 2 steps. Training Hyperparameter | Value
---|---
Optimizer | AdamW
Learning rate | $1e-8\mapsto 3e-7\mapsto 1e-7$
Batch size | 8
Gradient steps | 50k
Learning rate scheduler | Cosine Annealing with
| Linear Warmup (for 500 steps)
Weight decay | 1e-6
Rollout step curriculum | $[4,8,12,16,20]$
Curriculum milestone | $[0,10\text{k},20\text{k},30\text{k},40\text{k}]$
EMA | Deacy=0.999
Table A.2:  Training hyperparameter for stage 2. We save the Exponential
Moving Average (EMA) of model weights during the second stage.

The training loss is a latitude-weighted and variable-weighted $L1$ loss:

$\text{Loss}(\hat{X},X)=\frac{1}{N_{\text{lat}}N_{\text{lon}}N_{\text{t}}}\sum_{c=1}^{N_{c}}\lambda_{c}\sum_{l=1}^{N_{\text{level}}}\lambda_{l}\sum_{t=1}^{N_{\text{t}}}\sum_{i=1}^{N_{\text{lat}}}\sum_{j=1}^{N_{\text{lon}}}w(i)||\hat{X}_{tijl}^{c}-X_{tijl}^{c}||,$
(13)

where $w(i)$ is the latitude-based weight (see Equation 12), $\lambda_{l}$ is
the level-based weight, $\lambda_{c}$ is the variable-based weight, $\hat{X}$
is model’s prediction and $X$ is reference data, $N_{t}$ denotes total number
of rollout timesteps, $N_{\text{lat}}$ denotes number of latitude grid points,
$N_{\text{lon}}$ denotes number of longitude grid points. The level-based
weight is a simple linear interpolation between $0.05$ and $0.065$:
$\lambda_{l}=\text{Index}(l)/12\times 0.015+0.050$, where the index of
pressure level $l$ ranges from 0 to 12 corresponds to 50hPa to 1000 hPa. We
list the specification of the variable-based weight below (see Table 4.1 for
the meaning of the variable abbreviation):

lccccccccccc Variable z t u v q t2m MSLP u10 v10 TP6 TP24
Weight 1.0 1.0 0.5 0.5 0.1 1.0 0.1 0.1 0.1 0.05 0.05

Table A.3: Variable-based weight.

We didn’t perform search over the best combination of variable-based weights
and level-based weights. We find uniform level-based weights perform almost
similar to current level-based weights in preliminary experiments. The
variable weights we’ve used is similar to the variable weights used in
GraphCast(Lam et al., 2023). Recent work Stormer (Nguyen et al., 2023b)
showcase that using variable-based weight is beneficial for variables that are
assigned higher loss weight like 2m temperature, geopotential.

###  A.2 Metrics

The following metrics are defined on a surface variable or a upper-air
variable at a specific level.

#### Root mean squared error (RMSE)

The RMSE for time step $t$ is defined as:

$\text{RMSE}_{t}(\hat{X},X)=\sqrt{\frac{1}{N_{\text{lat}}N_{\text{lon}}}\sum_{i=1}^{N_{\text{lat}}}\sum_{j=1}^{N_{\text{lon}}}w(i)(\hat{X}_{tij}-X_{tij})}.$
(14)

#### Anomaly correlation coefficient (ACC)

The ACC is defined as the Pearson correlation coefficient of the anomalies
with respect to the climatology:

$\displaystyle\text{Anomaly:}\quad\hat{X}^{\prime}_{tij}=\hat{X}_{tij}-C_{\tilde{t}ij},\quad
X^{\prime}_{tij}=\hat{X}_{tij}-C_{\tilde{t}ij},$ (15)
$\displaystyle\text{ACC}_{t}(\hat{X}^{\prime},X^{\prime})=\frac{\sum_{i=1}^{N_{\text{lat}}}\sum_{j=1}^{N_{\text{lon}}}w(i)\hat{X}^{\prime}_{tij}X^{\prime}_{tij}}{\sqrt{\left(\sum_{i=1}^{N_{\text{lat}}}\sum_{j=1}^{N_{\text{lon}}}w(i)\hat{X^{\prime}}^{2}_{tij}\right)\left(\sum_{i=1}^{N_{\text{lat}}}\sum_{j=1}^{N_{\text{lon}}}w(i)X^{2}_{tij}\right)}},$
(16)

where we use the climatology statistics $C$ provided in WeatherBench2 (Rasp et
al., 2023) (computed on ERA5 data from $1990$ to $2019$), $\tilde{t}$ denotes
the day of year and time of day of the corresponding time index $t$.

#### Bias

The bias is defined as the point-wise difference between prediction and ground
truth:

$\text{Bias}_{t}(\hat{X},X)=\hat{X}_{tij}-X_{tij}.$ (17)

##  Appendix B Comparison against standard attention

Figure B.1: Axial factorized attention. The multi-dimensional spatial
structure of value is preserved, each axial attention kernel matrix’s size is
$S_{m}^{2}$ for $n$-dimensional grid with total size $S=S_{1}\times
S_{2}\times\ldots\times S_{n}$.

Figure B.2: Standard attention for sequence. The multi-dimensional spatial
structure is not considered explicitly during the attention (it can be
implicitly encoded via positional encoding). The full attention kernel
matrix’s size is quadratic to the total grid size $S$.

Two illustrative diagrams of factorized attention and standard attention are
shown above (Figure B.2 and B.2). Instandard attention, the spatial structure
is not preserved during the contraction of value and attention kernel.
Compared to axial Transformer (Ho et al., 2020), the axial attention kernel in
CaFA is still dependent on global information while in axial Transformer its
context is limited to that specific row/column.

We compare factorized attention to two different implementations of standard
attention in PyTorch’s attention function:
`torch.nn.functional.scaled_dot_product_attention` \- a standard one without
memory optimization and another one with memory optimization (Rabe and Staats,
2021) 333FlashAttention (Dao et al., 2022) is currently not available for
float32 computation on A6000 in PyTorch 2.1.. The results are shown in Figure
3(a), 3(b), 3(c). We observe that factorized attention has significantly
better scaling efficiency in terms of runtime and FLOPs. At the lowest
resolution, the runtime of factorized attention is larger than standard
attention. This is because in factorized attention layer the input is first
processed by two projection layer and then sent to compute the axial attention
kernels, with all computation done sequentially 444For projection and
attention kernel’s computation, each axis’s computation are independent to
each other and the memory cost is small, so potentially they can be optimized
to run in parallel.. Compared to memory efficient attention, factorized
attention has larger memory footprint at current stage. The memory cost of
factorized attention can also be optimized by adopting a similar chunking
strategy in Rabe and Staats (2021) during the contraction of value and
attention kernels.

(a) Forward latency

(b) Theoretical FLOPs

(c) Peak memory usage

Figure B.3: Computational benchmark between factorized attention and different
implementation of the standard softmax attention (Vaswani et al., 2017). The
number of attention head is 16, each attention head’s dimension is 128, batch
size is set to 1, input and output channel dimension is 512. The benchmark is
carried out on A6000 with CUDA 12, float32 precision is used. The runtime
statistics are profiled using DeepSpeed. The factorized attention layer
contains following learnable modules that are not in standard attention layer:
distance encoding (Equation (10)) and projection layer (Equation (2)).

##  Appendix C Further results

In this section we provide more quantitative and qualitative analysis of
model’s prediction capability.

We analyze how the model perform in terms of $L1$ and $L2$ (RMSE) loss. The
normalized difference with respect to IFS HRES is shown in Figure C.1. We
observe that the model performs relatively better in terms of $L1$ norm in the
long run, which is possibly because that the model is trained on $L1$ norm and
penalizes outlier less than $L2$ norm. We also compare models trained with
different time step sizes (lead time 6 hour vs 12 hour). We find that model
with larger time step size performs better in terms of forecasting accuracy
beyond 7 days but performs worse in shorter range, while the ACC of both model
variants fall below $0.6$ on day 10.

Example visualization of the learned positional encoding, distance encoding
and attention kernels are shown in Figure C.3, C.4, C.5. Interestingly, the
attention kernel at earlier layer exhibit sharper patterns than deeper layer,
where attention scores are mostly concentrated around the diagonal.

Figure C.1: Relative $L1$ and $L2$ error difference with respect to IFS HRES
of 4 selected key variables. Figure C.2: Normalized RMSE difference and ACC
difference comparison between CaFA trained with $6$ hour interval and CaFA
trained with $12$ hour interval. Negative RMSE difference and positive ACC
difference indicates better performance. CaFA trained with $12$ hour interval
uses less gradient steps for the second stage (36k) and fewer rollout steps
(up to 16). Figure C.3: Example visualization of learned spherical harmonic
based positional encoding. "Layer" denotes their corresponding attention layer
number, "c" denotes channel number. Figure C.4: Example visualization of
learned self-attention kernel along different axes. The kernel matrices are
selected from random batch and head. For better clarity all the kernel
matrices shown are normalized such that all the elements fall into the range
of $[-1,1]$. Figure C.5: Example visualization of learned distance encoding
along different axes. The distance encoding are selected from random batch,
head and channel. For better clarity all the kernel matrices shown are
normalized such that all the elements fall into the range of $[-1,1]$.

(a) $v10$m example rollout visualization

(b) $u850$ example rollout visualization

(c) $v850$m example rollout visualization

(a) $t850$ example rollout visualization

(b) $q700$m example rollout visualization

(c) MSLP example rollout visualization

Figure C.7:  More example rollout visualization of model’s prediction versus
reference ERA5 reanalysis data at different lead times. The initialization
time is 00:00 UTC on August 11, 2020.