This is the second post in a series on ECMWF's ML-based weather forecasting models. In the last post, we looked at the first, deterministic version of the data-driven, transformer-based weather forecasting system. In this post, we're looking at the follow-up paper (Lang et al. 2024) that was released a few months after the first one.
In the first post, we already mentioned the double-penalty problem, which leads the model to predict increasingly blurred features as the forecast horizon (and uncertainty) increases. This phenomenon occurred in all deterministic forecasting models up to this point, including GraphCast, Pangu-Weather, and FourCastNet.
While there are likely a number of ways to resolve this issue, the preferred approach for modelling teams (including the ECMWF team) has been to extend these deterministic models into probabilistic forecasts, which yield entire ensembles of predictions rather than just one value.
The reason this approach is so attractive is because it solves not only the blurring problem, but more importantly also yields uncertainty estimates on the forecast. This has been integral to the usefulness of physics-based forecasts for a long time. And as a proud Bayesian and advocate for probabilistic machine learning, that is of course immensely satisfying to me.
The ECMWF modelling team decided to implement the probabilistic component by keeping most of the architecture of the deterministic AIFS model unchanged, and instead introducing a new training objective: the Continuous Ranked Probability Score (CRPS). We'll go through the implementation details in this post, and throughout I will assume that you have read the first post beforehand.
Probabilistic Forecasts
Before getting into the technical details of AIFS-CRPS, it's worth clarifying what an ensemble actually is. In weather forecasting, an ensemble is a set of plausible future trajectories of the atmosphere, each sharp individually, which together represent the probability distribution for the future atmospheric state. ECMWF has been running physics-based ensembles for a long time. Their operational system, IFS-ENS, runs 50 perturbed copies of the same model, each starting from a slightly different initial condition and/or subject to stochastic perturbations in their parametrisations. The output is 50 predictions, which can be interpreted as a distribution over possible atmospheric states.
For ML-based models, there are a number of ways you could implement this probabilistic component. You could simply run the deterministic model multiple times with perturbed initial conditions, but this doesn't fix the blurring issue, and you're unlikely to get a well-calibrated ensemble that does not underestimate uncertainties. A more promising, and I'd say the most hyped, approach is using diffusion models. The flagship example is Google DeepMind's GenCast (Price et al. 2023), and ECMWF experimented with this approach as well (Enter the Ensembles). Diffusion models have been quite popular in the image and video generation fields, and their rapid progress makes them a natural choice for the architecture. They are also by nature probabilistic: you start with random noise and iteratively shape it using the denoising network, meaning different input noise maps produce different predictions. In image generation, diffusion models also stand out for producing sharper, less blurry images compared to previous methods. So they were an obvious candidate for dealing with the issues of deterministic forecasting models, and both GenCast and the ECMWF experiments produced promising results.
In the end, the AIFS-CRPS model in the paper takes a different approach though. It uses an architecture setup that is almost identical to the deterministic AIFS from the first paper, but is trained with a distribution-aware objective instead of the MSE loss. Thanks to this, each ensemble member is produced in a single forward pass through the network rather than a complicated denoising procedure. This has a bunch of advantages: it's faster and easier to run, architectural improvements of the deterministic model translate directly to the ensemble model, and ultimately, these models even tend to produce better predictions.
Updates To The AIFS Model
Before we get to the scoring-rule and probabilistic stuff, it's worth mentioning a few non-probabilistic upgrades that ECMWF quietly slipped into the pipeline between the original AIFS and AIFS-CRPS papers. The model backbone is still the same encoder-processor-decoder we covered in post 1, with the same graph-transformer mappers and the same 16-layer sliding-window-attention processor. But a few ingredients of the training pipeline are genuinely different, and would still make sense if you stripped the ensemble training away. So they're worth flagging on their own.
From Full States To Tendencies
The original AIFS predicted the full atmospheric state at $t+6$h directly, not a tendency. The AIFS authors flagged it as a known weak spot in the first paper:
AIFS currently uses "full" normalised states in its training loss. This stands in contrast to predicting normalised forecast tendencies as done by, e.g., Keisler [2022] or Lam et al. [2023]. The potential downside here is that the output of the last linear layer of AIFS that maps latent space variables to physical space may not be adequately normalised. This will likely be revised in a future AIFS upgrade.
That revision is implemented in the AIFS-CRPS architecture. I actually found the exact moment of this
change in the
codebase: issue #39 on
ecmwf/anemoi-core. The AIFS-CRPS paper itself sweeps this change under the rug
a little and only
mentions it in passing in Section 2.1, which confused me quite a bit when reading the paper for the
first time.
In practice, this means we change the prediction equation from
$$\hat{\mathbf{x}}_{t+6} = f_\theta\big(\mathbf{x}_{t-6},\, \mathbf{x}_t\big)$$
in the first paper to
$$\hat{\mathbf{x}}_{t+6} = \mathbf{x}_{t} + f_\theta\big(\mathbf{x}_{t-6},\, \mathbf{x}_t\big).$$
Low-Pass Filter On The Input State
The second change is closely related to the first. Since the AIFS-CRPS model is able to predict sharp features over the course of the rollout, rather than producing increasingly smoothed successive steps, small errors in the tendency predictions tend to pile up over the course of the rollout. This leads to implausible high-frequency features in the forecast.
The ECMWF team mitigates these errors by effectively applying a low-pass filter to the input state. In this way, the network has to redraw all the sharp features from scratch for every prediction, and doesn't have to fix any of its own leftover noise from the previous step.
They implement this low-pass filter by using down- and upsampling matrices $D$ and $U$. These are not neural networks, but sparse interpolation matrices generated by ECMWF's MIR (Meteorological Interpolation and Regridding) software, without any learnable parameters. The state update equation thus becomes $$\hat{\mathbf{x}}_{t+6} = U(D(\mathbf{x}_{t})) + f_\theta\big(\mathbf{x}_{t-6},\, \mathbf{x}_t\big).$$
Changes To The Training Curriculum
Improvements have also been made to the training procedure. The first paper had a three-phase training schedule: pretrain on single 6 h steps, then a rollout fine-tuning phase that smoothly grew the forecast window from 6 h to 72 h, then an operational-analysis fine-tuning phase.
AIFS-CRPS splits this into four phases. Stage 1 is the same single-step pretrain. Stage 2 is a dedicated 2-step rollout phase with the learning rate reset, which the original didn't have. Stage 3 grows the rollout window from 3 to 12 steps, but the learning rate is held constant through this phase. Stage 4 is the operational-analysis fine-tune. Furthermore, the amount of data for the final stage has been increased from two years (2019–2020) to eight years (2016–2023).
| Phase | Data | Forecast horizon | Iterations | Learning rate |
|---|---|---|---|---|
| 1. Pretrain | ERA5 1979–2017 | single 6 h step | 300,000 | $1\times 10^{-3} \to 0$, cosine |
| 2. Rollout-2 | ERA5 1979–2017 | 2 steps (12 h) | 60,000 | $1\times 10^{-5} \to 0$, cosine |
| 3. Long rollout | ERA5 1979–2017 | 3 → 12 steps (18 → 72 h) | ~45,000 | $1\times 10^{-6}$, constant |
| 4. Operational fine-tune | IFS analyses 2016–2023 | up to 12 steps | not reported | $5\times 10^{-7}$, constant |
Floor On The Upper-Atmosphere Loss Weight
The first paper mentioned that the accuracy of the prediction is reduced in the upper atmosphere, since the loss has a pressure-dependent weighting. AIFS-CRPS introduces an optional minimum floor to this weight of 0.2, in order to not fully neglect the upper-air variables.
The Continuous Ranked Probability Score
Now we are finally getting to the probabilistic part! As mentioned above, architecturally, the AIFS and AIFS-CRPS models are nearly identical. The big difference is the loss function. For the deterministic AIFS variant, the typical Mean Squared Error (MSE) loss was used. But the MSE only optimises for a point prediction (getting the prediction as close to this reference value as possible).
To move to probabilistic modelling, we usually do not look at the distance between two points but at the distance between two distributions. Typical keywords for a Bayesian modeller here are KL divergence, ELBO, variational inference, and all those fun topics. ML practitioners may have heard about the Wasserstein distance, which is used in GANs. Another distance metric, and the one that is used here, is the energy distance. Between two distributions $P$ and $Q$, it is defined as $$D^2_\mathrm{Energy}(P, Q) = 2\,\mathbb{E}\|\mathbf{X} - \mathbf{Y}\| \;-\; \mathbb{E}\|\mathbf{X} - \mathbf{X}'\| \;-\; \mathbb{E}\|\mathbf{Y} - \mathbf{Y}'\|,$$ where $\mathbf{X}$ and $\mathbf{X}'$ are independent random variables drawn from $P$, and $\mathbf{Y}$ and $\mathbf{Y}'$ are independent random variables drawn from $Q$. So if $P$ is the distribution of atmospheric states predicted by our probabilistic model, and $Q$ is the distribution of the true atmospheric state, then all we have to do is minimise this loss.
But this comes with a slight problem: we don't know the real distribution of atmospheric states. All we have is a single best-estimate value from the ERA5 dataset. So we have to use a small trick. We simply say the distribution of the true state is a point mass at $\mathbf{y}$ — all of the probability sits on the random variable $\mathbf{Y}$ being equal to $\mathbf{y}$, and nothing on anything else. In physics and math, we call this the Dirac delta distribution $\delta(\mathbf{y})$.
If we assume $Q = \delta(\mathbf{y})$, then every sample of $\mathbf{Y}$ is just $\mathbf{y}$, so $$\mathbb{E}\|\mathbf{Y} - \mathbf{Y}'\| = 0$$ and $$\mathbb{E}\|\mathbf{X} - \mathbf{Y}\| = \mathbb{E}\|\mathbf{X} - \mathbf{y}\|.$$ The energy distance therefore simplifies to $$D^2_\mathrm{Energy}\big(P, \delta(\mathbf{y})\big) = 2\,\mathbb{E}\|\mathbf{X} - \mathbf{y}\| \;-\; \mathbb{E}\|\mathbf{X} - \mathbf{X}'\|.$$ If we divide this expression by 2, we get the so-called Continuous Ranked Probability Score (CRPS), $$\mathrm{CRPS}(P, \mathbf{y}) = \mathbb{E}\|\mathbf{X} - \mathbf{y}\| \;-\; \tfrac{1}{2}\,\mathbb{E}\|\mathbf{X} - \mathbf{X}'\|.$$
The CRPS is our new loss function, and a so-called proper scoring rule. A scoring rule $S(P, \mathbf{y})$ takes a predicted distribution $P$ and an observation $\mathbf{y}$, and returns a number; lower is (usually) better. It's proper if, on average, the lowest score is achieved by predicting the true distribution. That is, if the data really comes from some distribution $Q$, then
$$ \mathbb{E}_{\mathbf{y} \sim Q}\,[S(Q, \mathbf{y})] \;\le\; \mathbb{E}_{\mathbf{y} \sim Q}\,[S(P, \mathbf{y})] \quad \text{for any } P. $$
Now, if our model predicts $M$ ensemble members $\{x_j\}_{j=1}^M$ for a single processor grid point $x$, we can estimate the CRPS as
$$ \mathrm{CRPS}\big(\{x_j\}_{j=1}^M, y\big) \;=\; \underbrace{\frac{1}{M}\sum_{j=1}^M |x_j - y|}_{\text{average distance to the observation}} \;-\; \underbrace{\frac{1}{2 M^2}\sum_{j=1}^M\sum_{k=1}^M |x_j - x_k|}_{\text{inter-ensemble spread}}. $$
This is a nicely interpretable expression. The first term is the average absolute error of the ensemble members against the observation — basically the Mean Absolute Error (MAE). The model tries to minimise this term. The second term, on the other hand, is half the average pairwise distance between ensemble members. This term contributes negatively to the loss, meaning the model tries to keep this average distance as large as possible. Both terms play against one another, with the optimum sitting where the two terms balance out, which is exactly when the ensemble spread matches the true forecast uncertainty.
The CRPS expression replaces the MSE in the loss function for a single processor grid point $x$, calculated over all ensemble members $\{x_j\}_{j=1}^M$. Just like in the first paper, the total loss is an area-weighted mean over all grid points. The individual variables are scaled in the same way as in the AIFS paper, so that they contribute roughly equally to the loss. Additionally, the pressure-dependent weighting is used to decrease the effect of upper-air variables (with the optional weight floor discussed above).
It is not surprising that the CRPS is used for ML-based weather prediction. The score itself originates in meteorology, where it is commonly used to assess the quality of (classical, physics-based) ensemble forecasts. So using it as a loss for a probabilistic forecasting model is quite natural.
The Almost-Fair CRPS
There is a small subtlety in the second term of the estimator above. If $j = k$, the pairwise sum always evaluates to zero. This means we are effectively losing one degree of freedom in the estimator, which makes the above expression biased. Ferro (2013) showed that this bias is corrected by switching the denominator from $\frac{1}{2M^2}$ to $\frac{1}{2M(M-1)}$, leading to the fair CRPS:
$$ \mathrm{fCRPS}\big(\{x_j\}, y\big) \;=\; \frac{1}{M}\sum_{j=1}^M |x_j - y| \;-\; \frac{1}{2 M(M-1)}\sum_{j=1}^M\sum_{k=1}^M |x_j - x_k|. $$
The fCRPS is unbiased in expectation at any ensemble size $M$. However, in the case when all $x_j = y$ except for one, $x_l$, the fair CRPS simplifies to $$ \mathrm{fCRPS}\big(\{x_j\}, y\big) \;=\; \frac{1}{M} |x_l - y| \;-\; \frac{2 (M-1)}{2 M(M-1)} |x_l - y| \;=\; 0, $$ and therefore becomes independent of the model predictions. Neural networks love to cheat if they can, so the authors anticipate this would become a problem at some point. The paper's solution is to introduce a composite score, which they call the almost-fair CRPS, given by $$ \mathrm{afCRPS}_\alpha \;:=\; \alpha\,\mathrm{fCRPS} + (1-\alpha)\,\mathrm{CRPS}, $$ with $\alpha = 0.95$. Rearranged for numerical stability, it reads $$ \mathrm{afCRPS}_\alpha \;=\; \frac{1}{2 M(M-1)} \sum_{j=1}^M \sum_{\substack{k=1 \\ k \ne j}}^M \Big(|x_j - y| + |x_k - y| - (1-\epsilon)|x_j - x_k|\Big), $$ with $\epsilon = (1-\alpha)/M$.
Making The Model Stochastic
Besides the changes mentioned above, the model architecture is essentially the same encoder-processor-decoder setup from the first paper. The only architectural addition to make the model stochastic is adding a noise input to the processor blocks, so that for a fixed input state the model can produce a distribution of forecasts conditioned on input noise.
For each ensemble member, and at each forecast step, a random noise vector is drawn from a standard Gaussian $\boldsymbol{\xi} \sim \mathcal{N}(0, \mathbf{I})$. It is passed through a small two-layer MLP followed by a layer normalisation (LN) to produce a noise embedding, $$\nu = \mathrm{LN}\big(\mathrm{MLP}(\boldsymbol{\xi})\big),$$ which is then used to condition the forward pass using Conditional Layer Normalisation (CLN).
In a standard LN, a grid point $x$ gets normalised $x^\prime = \frac{x - \mu}{\sigma}$, where $\mu$ and $\sigma$ are calculated across the hidden dimension, and passed through a learned affine transformation, $$\mathrm{LN}(x) = \gamma \odot x^\prime + \beta,$$ where $\gamma, \beta \in \mathbb{R}^{d}$ are per-layer learnable parameters, with $d$ being the size of the hidden dimension. In standard LN, these parameters are fixed after training and applied to inputs indiscriminately.
For conditional layer normalisation, we replace the two LN layers in the transformer by CLN layers. We use the noise embedding to calculate the noise-dependent transformation parameters $\gamma$ and $\beta$ using $$\gamma(\boldsymbol{\xi}) = W_\gamma\,\nu(\boldsymbol{\xi}) + b_\gamma, \qquad \beta(\boldsymbol{\xi}) = W_\beta\,\nu(\boldsymbol{\xi}) + b_\beta,$$ where $W_\gamma, b_\gamma, W_\beta, b_\beta$ are per-layer learnable parameters. In this way, the same input state can induce differing post-norm activations and ultimately output states, based on the input noise.
Training
The model training follows the four-stage curriculum laid out above.
For each forecast step, we start with the same ERA5 initial state $\mathbf{x}_t$. For every ensemble member, a unique noise sample $\boldsymbol{\xi}$ is drawn, and a forecast produced. The $M$ forecasts are then scored jointly under the afCRPS. For rollout training, each state is initialised with a fresh noise sample.
Similarly to the deterministic version, the training is distributed via data and tensor sharding. In addition, ensemble parallelism is used. Each ensemble member runs its own copy of the model on its own set of GPUs. After the forward pass, a differentiable all-gather collects the per-member forecasts onto every device, so the loss and gradients can be calculated.
The paper trains two different versions of the model:
- AIFS-CRPS O96 (~1°): O96 input grid, O48 processor grid, ensemble size $M = 4$, with the linear pressure weighting and its 0.2 floor. Trained for about 4 days on 64 H100 GPUs.
- AIFS-CRPS N320 (~0.25°, ERA5's native resolution): N320 input grid, O96 processor grid, ensemble size $M = 2$, no minimum-pressure floor. Trained for about 7 days on 128 H100 GPUs.
Inference And Operational Use
At inference time, the individual ensemble members are basically fully independent, so that one can draw as many as desired. ECMWF has made the ensemble model operational last year as AIFS ENS (separate from the AIFS Single deterministic model). The operational model includes 50 ensemble members plus a control run.
Importantly, one change is made from the training procedure: In training, every member starts from the same ERA5 state, so all of the ensemble spread is generated internally, by the noise channel. In operation, the initial states for AIFS ENS are the same initial states that are used for the physics-based IFS ensemble. These initial conditions are themselves perturbed using singular-vector perturbations. The operational forecast therefore carries two sources of uncertainty: initial-condition uncertainty from the perturbed states, and forecast-model uncertainty from the noise. Since these perturbations are specifically tuned to improve the forecast skill of the IFS model, they appear to degrade the performance of the AIFS ENS system slightly. ECMWF flags this as an area for future improvements.
While the authors argue there is still lots of room for improvement and experimentation, it is pretty amazing how fast they were able to build an ML-based operational system that can compete with established physics-based methods, all while being much cheaper and faster to run. In the paper, they especially flag the good performance of the model over longer time horizons as a promising avenue for further research.
There have already been some follow-up papers on this work: AIFS 1.1, which introduces additional variables; a preprint that modifies the CRPS loss to explicitly include multi-scale spatial features; and the updated AIFS v2. There has also been recent work on combining ECMWF's ML-based and physics-based models. So there is still a bunch of things we can dig into in the future.