ExtremeCast: Boosting Extreme Value Prediction for Global Weather Forecast

Medium-range weather forecast provides forecasts of weather conditions including temperature, wind speed, humidity, and more, for the next few days to two weeks. Traditional medium-range forecast models are mainly physics-based dynamic models such as the Integrated Forecasting System (IFS) from European Centre for Medium-Range Weather Forecasts (ECMWF)  (MODEL 2003), whose development is primarily limited by the computational cost of numerical models  (Ben-Bouallegue et al. 2023).

In recent years, data-driven models based on machine learning show significant potential for weather forecast  (de Burgh-Day and Leeuwenburg 2023). FourCastNet  (Kurth et al. 2023) utilizes Adaptive Fourier Neural Operator networks  (Guibas et al. 2021) for prediction and performs a two-step finetuning to improve the accuracy of autoregressive multi-step forecast, where the model’s output is fed back as input for predicting the next step. This was followed by the Pangu-Weather  (Bi et al. 2023), which uses 3D Swin-Transformer  (Liu et al. 2021) and proposes hierarchical temporal aggregation aimed at reducing iterations in autoregressive forecast through the integration of multiple models. Subsequently, GraphCast  (Lam et al. 2023) uses graph neural networks and performs a 12-step autoregressive finetuning. FengWu  (Chen et al. 2023a) treats the prediction problem as a multi-task optimization and introduces a novel finetune strategy named replay buffer. FuXi  (Chen et al. 2023b) reduces the cumulative error of the multi-step forecast by cascading three U-Transformer models  (Petit et al. 2021) optimized at different lead times.

Although many ML models outperform ECMWF-IFS on the root mean square error (RMSE) metric, most of them suffer from the prediction smoothing issue and lag behind ECMWF-IFS on the extreme value metrics.

The application of ML models in predicting extreme weather remains relatively unexplored. Most existing models for extreme weather prediction are limited to regional scale or low resolution ($>0.25^{\circ}$) rather than global scale and high resolution ($0.25^{\circ}$). Zhao et al.  (Zhao, Lu, and Kou 2003) use wavelet transform  (Farge 1992) to predict the occurrence of extreme weather, but cannot predict their actual values. Porto et al.  (Porto et al. 2022) combine multiple models to address the challenge of learning diverse extreme weather patterns, which introduces several times the additional training cost. Annau et al.  (Annau, Cannon, and Monahan 2023) propose to separate high and low-frequency in the data through convolution and enhance the learning of high-frequency to improve extreme value prediction. However, the learning of high-frequency is notably more challenging than low-frequency, which has not been effectively solved. Morozov et al.  (Morozov et al. 2023) apply bias correction on extreme values through quantile regression, whose main limitation is that it only performs regression on $7$ quantiles, which may not sufficiently capture the continuous distribution of the data.

Some work improves extreme values by adjusting loss. Lopez-Gomez et al.  (Lopez-Gomez et al. 2023) increase the learning weight of extreme values by designing an exponential-based loss function. However, this method may cause exponential explosion in practical applications (Wang et al. 2019). Ni et al.  (Ni 2023) identify that using GAN loss  (Creswell et al. 2018) can enhance the performance of the model in predicting extreme values, which is relatively difficult to optimize  (Berard et al. 2019).

For the global scale and high resolution extreme weather forecast, FuXi-Extreme  (Zhong et al. 2023) uses a diffusion model to improve prediction accuracy of extreme weather. Building on this previous work, we use a cascaded diffusion model as a baseline component of our model and design Exloss and ExBooster to achieve SOTA performance.

Formally, we denote the weather state at time period $i$ as a tensor $\mathcal{X}^{i}\in\mathbb{R}^{C\times H\times W}$ , where $C$ represents the number of atmospheric variables, $H$ and $W$ are the height and width. In this paper, we consider $C=69$ atmospheric variables, which will be introduced in detail in Table 1. To align the experimental setup with ML models like Pangu-Weather (Bi et al. 2023) and GraphCast (Lam et al. 2023), we adopt a latitude and longitude resolution of $0.25^{\circ}$, which means the global atmospheric data is projected into a two-dimensional plane that $H=721$ and $W=1440$ horizontally.

The model learns to map from $\mathcal{X}^{i}$ to $\mathcal{X}^{i+1}$ with a time interval of 6 hours. Formally, it can be written as:

where $\theta$ denotes the parameters of the model. By iterating the above process $n$ times, i.e. autoregressive prediction, the model can forecast any future time step $\mathcal{X}^{i+n}(n\in\mathbb{Z})$.

Our model consists of two trainable modules and one training-free module. They are, respectively, a deterministic model $M_{d}$, a probabilistic generation model $M_{g}$, and the ExBooster module. Specifically, $M_{d}$ learns the mapping from $\mathcal{X}^{i}$ to $\mathcal{X}^{i+1}$, followed by $M_{g}$, a conditional diffusion model that that refines the predictions of $M_{d}$ with additional details. Ultimately, ExBooster simulates atmospheric uncertainty through random sampling from the output of $M_{g}$ and amalgamating these samples to derive the final prediction.

In addition to one-step ($\mathcal{X}^{i}\to\mathcal{X}^{i+1}$) training, we also incorporate a finetuning of autoregressive regime to improve model’s performance in multi-step prediction. The training process of our model can be divided into four stages:

• •

  Stage 1: One-Step pretraining. Use MSE loss to train $M_{d}$ learning mapping from $\mathcal{X}^{i}$ to $\mathcal{X}^{i+1}$.

• •

  Stage 2: Muti-Step finetuning. Use the proposed Exloss to finetune $M_{d}$ learning autoregressive forecasting. This stage improves not only the accuracy of the model’s multi-step predictions, but also the ability to predict extreme weather due to the use of Exloss.

• •

  Stage 3: Diffusion training. Use Exloss to train $M_{g}$ to refine extreme values while freezing $M_{d}$’s parameters.

• •

  Stage 4: Inference. ExBooster that does not require training is added for complete weather forecast, enhancing the forecast hit rate of extreme events.

In the following, we will start by analyzing the inherent limitations of MSE, followed by describing the details of the proposed Exloss and the ExBooster.

We abbreviate Equation 1 as $F_{\theta}(X)=P(Y|X)$, where $X$ is the input of the model, and $Y$ is the target. According to the previous work (Aravkin et al. 2012; Kendall, Gal, and Cipolla 2018), when the data is normally distributed $Y\sim\mathcal{N}\left(\mu,\sigma^{2}\right)$, using MSE loss to predict $Y$ is equivalent to using maximum likelihood estimation (Jiang and Zhang 2009) to predict the mean $\mu$ of $Y$ given $X$, as shown in Equation 2:

However, this type of optimization has a significant negative impact on extreme values. Consider the introduction of a new variable, denoted as $Y_{M}=max\{Y_{1},Y_{2},...,Y_{n}\}$, where $Y_{i}(i=1,2..,n)$ represents different samples. In the context of weather forecasting, assuming that $Y_{1},Y_{2},...,Y_{n}$ correspond to the temperatures in different regions, then $Y_{M}$ represents the highest temperature in the world. According to the extreme value theory, $Y_{M}$ obeys the extreme value distribution (more specifically, maximum value distribution), that is, $Y_{M}\sim G(\widetilde{\mu},\widetilde{\sigma})$  (Kotz and Nadarajah 2000), whose probability density function is:

where $\widetilde{\mu}$ and $\widetilde{\sigma}$ are its position and scale parameter respectively. It is worth mentioning that, unlike the normal distribution, the extreme value distribution is asymmetric. When employing maximum likelihood estimation for predicting the extreme value $Y_{M}$, we can derive the optimization objective through a similar procedure:

Define the optimization function obtained by maximum likelihood estimation of this extreme distribution as $obj_{M}(\cdot)$:

It can be proved (in Appendix A) that the optimal solution of Equation 4, that is, the minimum of $obj_{M}(\cdot)$ is obtained if and only if $\widetilde{\mu}=Y_{M}$, which meets the optimization goal. But due to the asymmetry of the data distribution, $obj_{M}(\cdot)$ is also asymmetric, as depicted in the upper left corner of Figure 1, which ultimately results in biased estimates.

$obj_{M}(\cdot)$ can be viewed as a loss function that considers the data distribution. Under this assumption, the area enclosed by $obj_{M}(\cdot)$ and the x-coordinate axis can be regarded as the expectation of the total loss.

As illustrated in Figure 1, the shading on the left and the shading on the right represent the expectation of the total loss when model provides different estimates (underestimation and overestimation, respectively). Assuming the highest temperature in the world at a certain time is $Y_{M}=40^{\circ}C$, predicting $\widetilde{\mu}=39.5^{\circ}C$ will get a lower loss than predicting $\widetilde{\mu}=40.5^{\circ}C$ based on probability expectation, which encourages the model to predict a smaller $Y_{M}$, that is $\widetilde{\mu}=39.5^{\circ}C$, rather than $\widetilde{\mu}=40.5^{\circ}C$. This finally results in underestimating the degree of extremes. The same situation occurs when predicting minimum values, such as the global minimum temperature. For a comprehensive analysis of $obj_{M}(\cdot)$, please refer to Appendix A.

In practice, the highest temperature predicted by the MSE-based model is always smaller than the true value  (Ben-Bouallegue et al. 2023), which supports the above proof.

Section 3.3 proves from EVT that using MSE loss will underestimate the degree of extremes. To address this issue, a simple and straightforward idea is to make the originally asymmetric $obj_{M}(\cdot)$ symmetrical through scaling.

It is cumbersome and unnecessary to transform $obj_{M}(\cdot)$ into a completely symmetric function. In actual operation, we use linear scaling to make the optimization objective function $obj_{M}(\cdot)$ symmetrical within a certain range $\epsilon$, where $\epsilon$ is a hyperparameter that depends on both the data distribution and the accuracy of the model’s predictions. In simple terms, when the model’s predictions are highly accurate, the prediction error is smaller, allowing for the selection of a smaller value for $\epsilon$. For detailed guidelines on selecting an appropriate value for $\epsilon$, please refer to Appendix B. Mathematically, the process of achieving symmetry through scaling can be expressed as follows:

The purpose of the scaling factors $s_{1}$ and $s_{2}$ is to equalize the areas of the two shaded regions in Figure 1, which ensures that the expectations of the total loss are the same for underestimation and overestimation. Subsequently, $s_{1}$ and $s_{2}$ are combined into a function $S(\hat{\mathcal{Y}},\mathcal{Y})$. By applying scaling function to the loss, we propose Exloss:

where $\hat{\mathcal{Y}}$ and $\mathcal{Y}$ represent the model prediction and the target respectively, $S(\hat{\mathcal{Y}},\mathcal{Y})\in\mathbb{R}^{C\times H\times W}$ is asymmetrical scaling function. For the detailed form of $S(\hat{\mathcal{Y}},\mathcal{Y})$, please refer to Appendix B.

In fact, by calculating $S(\hat{\mathcal{Y}},\mathcal{Y})$, it can be observed that Exloss increases the loss weight when the model underestimates extreme value. This behavior is highly rational, as it adaptively increases the penalty for underestimations.

The atmospheric system has uncertainty (Thompson 1957), which means that for a given initial condition, there are multiple possible predictions. Some low-probability extreme events only exist in some prediction results. In such scenarios, multiple sampling of predictions aids in capturing a broader spectrum of potential extreme events.

In this work, we propose the ExBooster, drawing insights from ensemble algorithms (Gneiting and Raftery 2005). The process begins by introducing noise to the prediction, generating $m$ ensemble results. Due to the addition of noise, these different prediction results form a larger range of values than the original single prediction, which means that they contain more extreme events. Then, we use the ranking histogram method (Hamill 2001) to integrate these m prediction results. Specifically, all pixels in these samples are first sorted and grouped by the size of $m$, and then the original value is replaced with the median of the corresponding part according to the ranking of each element in the initial prediction.

As depicted in Figure 3, ExBooster adaptively expands the range of pixel values without altering the overall distribution, thereby producing more extreme predictions. In practice, we choose $m=50$, that is, randomly sample $50$ predictions, which is enough to mimic the uncertainty of atmospheric dynamics (Price et al. 2023). For detailed implementation, please refer to Appendix C.

We use the RQE to evaluate the models’ ability to predict extreme values on the global scale. As depicted in Figure 4, the RQE values of most ML models exhibit a tendency to become progressively more negative as the lead time increases, indicating an increasing underestimation of extreme values. On the contrary, the dynamic model has a relatively better RQE, that is, an RQE value closer to 0, and does not decrease as the lead time increases.

For our model, overall, the RQE does not decrease as the lead time increases, which is due to Exloss’s asymmetric optimization, avoiding the underestimation of extreme values. Especially for q500 (the channel of absolute humidity at the 500hPa pressure layer), u500, and u10, our model predicts with a RQE close to 0 consistently, surpassing both ML models and the dynamic model, showing the accurate and stable extreme value prediction capabilities.

Our model slightly overestimates (around $10^{-3}$) the extreme values in certain atmospheric variables like z500, which is very rare for ML models but is occasionally observed in dynamic models as well. Appropriate overestimation can provide a more adequate upper bound for the prediction of extreme events, which is beneficial for some applications like disaster early warning  (Li et al. 2021).

Different from RQE, SEDI focus more on extreme values in local areas, because each pixel uses its quantile in the time dimension as a threshold that varies from month to month. For instance, a temperature of $-10^{\circ}C$ may not be considered as an extreme event for high-latitude areas during winter, but it could be extreme for mid-latitude areas in summer.

As illustrated in Figure 5, for surface temperature, our model’s prediction accuracies at 90th, 95th, 98th extreme degrees are significantly higher than that of other models. Regarding the most extreme scenarios at the 99.5th quantile, our model demonstrates relatively better performance when the lead time is less than 24 hours, and for lead times exceeding 24 hours, our model has similar performance to FuXi-Extreme and ECMWF-IFS.

For surface wind speed, our model significantly outperforms other models across all extreme levels. It is worth noting that the SEDI of our model’s first step prediction (lead time=6h) is remarkably high, predominantly exceeding 0.95. This is particularly challenging for wind speed which exhibit rapid and diverse changes, highlighting the excellent performance of our model in forecasting extreme wind speeds.

To better showcase the performance of our model, we select two natural disasters, typhoon and heatwave, for visual analysis. Specifically, we choose Nanmadol typhoon  (Wu et al. 2023) at 6:00 on September 17, 2022, and the heatwave in the southeastern coastal area of China at 12:00 on August 14, 2022  (Jiang et al. 2023), for three-day forecasting. Figure 7 illustrates the results, with the two rows showing surface wind speed and surface temperature, respectively.

Regarding typhoons, the wind speeds of our prediction are the most extreme and closest to the real wind speeds. This aligns with the findings observed in the far-leading SEDI of wind speed, showing the exceptional performance of our model in predicting extreme wind speeds. This capability is of significant importance for estimating the intensity of typhoon landings and assessing the potential damage to coastal areas.

In the case of the heatwave, both of the predictions of our model and FuXi-Extreme are more extreme than other models. Although fuxi-extreme’s predictions are more extreme than ours, our model gets a prediction ($37.01^{\circ}C$) that is closer to the true value ($37.25^{\circ}C$) for the regional maximum temperature prediction. For more visualizations, please refer to Appendix E.

We assess the individual contributions of each module by conducting ablation experiments where we remove Exloss, and ExBooster. Additionally, as an effective implementation component, diffusion model is also considered in the ablation.

Through Table 2, we find that all three modules are beneficial in enhancing the prediction of extreme values. Specifically, we discover that the asymmetric optimization of Exloss plays a crucial role in enhancing the forecasting capabilities of extreme weather events. This is evident from the comparison between row ➀ and row ➁, where the use of Exloss leads to significant improvements in both RQE and SEDI metrics. Furthermore, comparing row ➂ and row ➅ (best version), the lack of Exloss results in a significant decrease in all metrics.

Moreover, Diffusion and ExBooster have great improvements in the prediction of extreme temperatures and extreme wind speeds respectively, which highlights the favorable decoupling properties of different modules in our model.

Although the focus of this paper is extreme weather prediction, the overall forecast accuracy of our model is still competitive. As shown in Figure 6, while our model demonstrates outstanding performance in predicting extreme wind speeds, it also achieves a high accuracy in overall wind speed prediction, as indicated by RMSE u10. Additionally, the RMSE metrics of our model on other variables are also comparable to top ML models and surpass ECMWF-IFS.