W-MAE: Pre-trained weather model with masked autoencoder for multi-variable weather forecasting

Numerical Weather Prediction (NWP) involves using mathematical models, equations, and algorithms to simulate and forecast atmospheric conditions and weather patterns DBLP:journals/nature/BauerTB15 . NWP models generally use physical laws and empirical relationships such as convection, advection, and radiation to predict future weather states DBLP:journals/msj/Arakawa97 . The accuracy of NWP models depends on many factors, such as the quality of the initial conditions, the accuracy of the physical parameterizations Parameterization , the resolution of the grid, and the computational resources available. In recent years, many efforts have been made to improve the accuracy and efficiency of NWP models DBLP:journals/pt/SchultzBGKLLMS21 ; DBLP:journals/natmi/IrrgangBSBKSS21 . For example, some researchers DBLP:journals/jc/Arakawa04 ; DBLP:journals/mwr/GrenierB01 have proposed using adaptive grids to better capture the features of the atmosphere, such as the boundary layer and the convective clouds. Others have proposed using hybrid models DBLP:journals/ae/HanMSCLLX22 that combine the strengths of different models, such as the global and regional models DBLP:journals/geoinformatica/ChenWZZLY19 .

In recent years, there has been an increasing interest in the development of AI-based weather forecasting models. These models use deep learning algorithms to analyze vast amounts of meteorological data and learn patterns that can be used to make accurate weather forecasting. Compared with the traditional NWP models DBLP:journals/pt/SchultzBGKLLMS21 ; DBLP:journals/natmi/IrrgangBSBKSS21 , AI-based models DBLP:journals/ames/RaspDSWMT20 ; DBLP:journals/corr/abs-2202-07575 have the potential to produce more accurate and timely weather forecasts, especially for extreme weather events such as hurricanes and heat waves. Albu et al.DBLP:journals/remotesensing/AlbuCMCBM22 combine Convolutional Neural Network (CNN) with meteorological data and propose NeXtNow, a model with encoder-decoder convolutional architecture. NeXtNow is designed to analyze spatiotemporal features in meteorological data and learn patterns that can be used for accurate weather forecasting. Karevan and Suykens DBLP:journals/nn/KarevanS20 explore the use of Long Short-Term Memory (LSTM) network for weather forecasting, which can capture temporal dependencies of meteorological variables and are suitable for time series forecasting, but may struggle to capture spatial features. However, both LSTM-based and CNN-based models suffer from high computational costs, limiting their ability to handle large amounts of meteorological data in real-time applications.

Taking into account the computational costs and the need for timely forecasting, Pathak et al.DBLP:journals/corr/abs-2202-11214 propose FourCastNet, an AI-based weather forecasting model that employs adaptive Fourier neural operators to achieve high-resolution forecasting and fast computation speeds. FourCastNet represents a promising solution for real-time weather forecasting, but it requires significant amounts of training data to achieve optimal performance and may have limited accuracy in certain extreme weather events. In this work, we apply our pre-training method to FourCastNet to explore its impact on model performance. The results show that pre-trained FourCastNet achieves nearly 20% performance improvement in precipitation forecasting. This suggests that pre-training can be a feasible strategy to enhance the performance of FourCastNet and other weather forecasting models.

Self-supervised learning enables pre-training rich features without human annotations, which has made significant strides in recent years. In particular, Masked AutoEncoder (MAE) DBLP:conf/cvpr/HeCXLDG22 , a recent state-of-the-art self-supervised pre-training scheme, pre-trains a ViT encoder DBLP:conf/iclr/DosovitskiyB0WZ21 by masking an image, feeding the unmasked portion into a Transformer-based encoder, and then tasking the decoder with reconstructing the masked pixels. MAE adopts an asymmetric design that allows the encoder to operate only on the partial, observed signal (i.e., without mask tokens) and a lightweight decoder that reconstructs the full signal from the latent representation and mask tokens. This design achieves a significant reduction in computation by shifting the mask tokens to the small decoder. Our W-MAE model is built upon the MAE architecture and also utilizes the VIT encoder to process unmasked image patches. However, we employ a modified decoder specifically designed to reconstruct pixels for meteorological data to reduce computational overhead.

ERA5 DBLP:journals/rms/HersbachBBH20 is a publicly available atmospheric reanalysis dataset provided by the European Centre for Medium-Range Weather Forecasts (ECMWF). The ERA5 reanalysis data combines the latest forecasting models from the Integrated Forecasting System (IFS) DBLP:journals/ams/BougeaultTBBB10 with available observational data (e.g., pressure, temperature, humidity) to provide the best estimates of the state of the atmosphere, ocean-wave, and land-surface quantities at any point in time. The current ERA5 dataset comprises data from 1979 to the present, covering a global latitude-longitude grid of the Earth’s surface at a resolution of $0.25^{\circ}\times 0.25^{\circ}$ and hourly intervals, with various climate variables available at 37 different altitude levels, as well as at the Earth’s surface. Our experiments are conducted on the ERA5 dataset across two distinct data partition schemes. The first partition manifests a division wherein the training, validation, and test sets are composed of data from 2015, 2016 and 2017, and 2018, respectively. The second partition arrangement aligns with the division used in FourCastNet DBLP:journals/corr/abs-2202-11214 whereby the training, validation, and test sets are comprised of data from 1979 to 2015, 2016 and 2017, and 2018, respectively. We select data samples at six-hour intervals, where each sample comprises a collection of twenty atmospheric variables across five vertical levels (see Table 1 for details).

We focus on forecasting two important and challenging atmospheric variables (consistent with the work done on FourCastNet DBLP:journals/corr/abs-2202-11214 ): 1) the wind velocities at a distance of 10m from the surface of the earth and 2) the 6-hourly total precipitation. The variables are selected for the following reasons: 1) predicting near-surface wind speeds is of tremendous practical value as they play a critical role in planning energy storage and grid transmission for onshore and offshore wind farms, among other operational considerations; 2) neural networks are particularly well-suited to precipitation prediction tasks due to their impressive ability to infer parameterizations from high-resolution observational data. Additionally, our model reports forecasting results for several other variables, including geopotential height, temperature, and wind speed.

Vision Transformer (ViT) DBLP:conf/iclr/DosovitskiyB0WZ21 , with its remarkable spatial modeling capabilities, scalability, and computational efficiency, has emerged as one of the most popular neural architectures. Let $I\in\mathbb{R}^{H\times W\times C}$ denote an input image with height $H$, width $W$, and $C$ channels. In our study, we treat each sample of ERA5 data as an image, with different channels representing the different atmospheric variables contained in each sample. ViT differs from the standard Transformer in that it divides an image into a sequence of two-dimensional patches $\mathbf{x}_{p}\in\mathbb{R}^{N\times(P^{2}\cdot C)}$, where each patch has a resolution of $(P,P)$, resulting in $N=HW/P^{2}$ patches. The patches are flattened and mapped to a $D$-dimensional feature space using a trainable linear projection (Eq. 1). The output of this projection, referred to as patch embeddings, is then concatenated with a learnable embedding, which serves as the input sequence to the Transformer encoder. The encoder consists of alternating layers of Multi-headed Self-Attention (MSA) and Multi-Layer Perceptron (MLP) blocks (Eq. 2 and Eq. 3). The MLP block contains two fully connected layers with a GELU activation function. Layer Normalization (LN) is applied before each block, and residual connections are added. The overall procedure can be formalized as follows:

where $\mathbf{E}\in\mathbb{R}^{(P^{2}\cdot C)\times D}$, $\mathbf{E}_{pos}\in\mathbb{R}^{(N+1)\times D}$ and $L$ is the number of layers.

MAE is a self-supervised pre-training scheme that masks random patches of the input image and reconstructs the missing pixels. Specifically, a masking ratio of $m$ is selected, and $m$ percentage of the patches is randomly removed from the input image. The remaining patches are then passed through a projection function (e.g., a linear layer) to project each patch (contained in the remaining patch sequence $S$) into a $D$-dimensional embedding. A positional encoding vector is then added to the embedded patches to preserve their spatial information, and the function is defined as:

where $pos$ is the position of the patch along the given axis and $i$ is the feature index. Subsequently, the resulting sequence is fed into a Transformer-based encoder. The encoder processes the sequence of patches and produces a set of latent representations. The removed patches are then placed back into their original locations in the sequence of patches, and another positional encoding vector is added to each patch to preserve the spatial information. After that, all patch embeddings are passed to the decoder to reconstruct the original input image. The objective function is to minimize the difference between the input image and the reconstructed image. Our proposed W-MAE architecture utilizes MAE to reconstruct the spatial correlations within meteorological variables.

Our W-MAE uses the standard VIT encoder to process the unmasked patches, as shown in Figure 2. Subsequently, all patches are fed into the W-MAE decoder in their original order. The standard MAE learns representations by reconstructing an image after masking out most of its pixels, and its decoder uses self-attention for weight matrix computation over all input patches. However, since meteorological images are approximately three times larger than typical computer vision images, using the original MAE decoder would result in a significant increase in computational resources. Therefore, our W-MAE decoder splits the input one-dimensional patch sequence into a two-dimensional patch matrix, following the shape of the two-dimensional meteorological image. As a result, our decoder only considers the relationships between patch blocks along the second dimension when performing weight computation via self-attention. Specifically, for an input image $I\in\mathbb{R}^{H\times W\times C}$ with height, width, and channel as $H$, $W$, and $C$, respectively, we divide it into patches with each patch resolution set to $(P,P)$. As a result, we get $N=HW/P^{2}$ one-dimensional patches. To ensure the original image proportions are preserved, we then resize the individual $N$ patches into a $(X,Y)$ two-dimensional patch matrix, where $X=H/P$ and $Y=W/P$. Consequently, we can apply the standard self-attention formula as shown below:

The query, key, and value vectors generated by the transformer block are represented by Q, K, and V, respectively, where the feature dimensionality of Q and K is $d$. The original MAE decoder adopts a set of $N$ patch vectors as values for Q, K, and V. We replace these with an $(X,Y)$ two-dimensional patch vector matrix. To evaluate the effectiveness of our improved structure, we employ a patch size of $(32,32)$ and measure the resulting performance. Notably, using our W-MAE decoder achieves a 1.2 $\times$ speedup (for each training epoch) compared with the vanilla MAE decoder, while reducing memory overheads by up to 8%. This significantly lowers the computational resources demanded while still retaining high-quality pixel reconstruction capabilities.

Given data samples from the ERA5 dataset, each sample contains twenty atmospheric variables and is represented as a 20-channel meteorological image. Our W-MAE first partitions meteorological images into regular non-overlapping patches. Next, we randomly sample patches to be masked, with a mask ratio of 0.75. The visible patches are then fed into the W-MAE encoder, while the W-MAE decoder reconstructs the missing pixels based on the visible ones. The mean squared error between the reconstructed and original images is computed in the pixel space and averaged over the masked patches. There are some differences in model architecture for two dataset division settings (see Section 3.1 for details). For clarity, we denote the models trained on two years of data and thirty-seven years of data as W1 and W2, respectively. For W1, we set the patch size as $16\times 16$, while for W2, it is $8\times 8$. The W1 model consists of encoders with depth=16, dim=768 and decoders with depth=12, dim=512. The W2 model employs encoders with depth=12, dim=768 and decoders with depth=6, dim=512. To reduce the memory overhead and save on training time, we employ a smaller model architecture for W2, given that the training data volume for W2 is nearly 20 times that of W1. Furthermore, with a patch size of $8\times 8$, the number of generated patches reaches the maximum limit beyond which the self-attention network will be effective. Therefore, in the W2 model, we replace the self-attention blocks used in both the encoder and decoder with adaptive Fourier neural operator (AFNO) blocks DBLP:journals/corr/abs-2111-13587 obtained from FourCastNet. This process is equivalent to applying the pre-training architecture of our W-MAE to FourCastNet. We employ the AdamW optimizer DBLP:conf/iclr/LoshchilovH19 with two momentum parameters $\beta_{1}$=0.9 and $\beta_{2}$=0.95, and set the weight decay to 0.05. The pre-training process takes 600 epochs for W1 and 200 epochs for W2. Note that W2 is currently still in training and has not yet reached its optimal performance.

We perform multi-variable forecast fine-tuning on the ERA5 dataset, with the task of predicting future states of meteorological variables based on the meteorological data from the previous time-step. We denote the modeled variables as a tensor $X(k\Delta t)$, where $k$ represents the time index and $\Delta t$ is the temporal spacing between consecutive time-steps in the dataset. Throughout this work, we consider the ERA5 dataset as the ground-truth and denote the true variables as $X_{true}(k\Delta t)$. For simplicity, we omit $\Delta t$ in our notation, and $\Delta t$ is fixed at 6 hours. After pre-training, we fine-tune our pre-trained W-MAE still using the ERA5 dataset to learn the mapping from $X(k)$ to $X(k+1)$. The fine-tuning process for multi-variable forecasting takes a further 120 epochs and 100 epochs for W1 and W2, respectively.

To evaluate the performance of our W-MAE model on multi-variable forecasting, we perform autoregressive inference to predict the future states of multiple meteorological variables. Specifically, the fine-tuned W-MAE is initialized with $N_{f}$ different initial conditions, taken from the testing split (i.e., data samples from the year 2018). $N_{f}$ varies depending on the forecasting days $d$, which differs for each forecast variable, as shown in Table 2. Subsequently, the fine-tuned W-MAE is allowed to run iteratively for $t$ time-steps to generate future states at time-step $i+j\Delta t$. For each forecast step, we evaluate the latitude-weighted Anomaly Correlation Coefficient (ACC) and Root Mean Squared Error (RMSE) DBLP:journals/ames/RaspDSWMT20 for all forecast variables. The latitude-weighted ACC for a forecast variable $v$ at forecast time-step $l$ is defined as follows:

where $\tilde{\mathbf{X}}_{\mathrm{pred/true}}(l)\left[v,m,n\right]$ represents the long-term-mean-subtracted value of predicted or true variable $v$ at the location denoted by the grid coordinates $(m,n)$ at the forecast time-step $l$. The long-term mean of a variable refers to the average value of that variable over a large number of historical samples. The long-term mean-subtracted variables $\tilde{\mathbf{X}}_{\mathrm{pred/true}}$ represent the anomalies of those variables that are not captured by the long-term mean values. $L(m)$ is the latitude weighting factor at the coordinate $m$ and $\mathrm{lat}(m)$ denotes the latitude value. $N_{lat}$ provides the total number of latitude locations in the dataset, which is used to normalize the latitude weighting factor to maintain the overall weight balance across all the locations on the grid. The latitude-weighted RMSE for a forecast variable $v$ at forecast time-step $l$ is defined as follows:

where ${\mathbf{X}}_{\mathrm{pred/true}}$ represents the value of predicted or true variable $v$ at the location denoted by the grid coordinates $(j,k)$ at the forecast time-step $l$. $N_{lat}$ and $N_{lon}$ denote the total number of latitude locations and longitude locations, respectively.

We report the mean ACC and RMSE for each of the variables at each forecast time-step and compare them with the corresponding FourCastNet forecasts that use time-matched initial conditions. We present the performance comparisons of W-MAE and FourCastNet under two different training data durations, namely, two years and thirty-seven years. Figure 3 displays the latitude-weighted ACC and RMSE values for our W-MAE model forecasts (indicated by the red line with markers) and the corresponding FourCastNet forecasts (indicated by the blue line with markers).

As depicted in Figure 3, with two years of training data, the W-MAE pre-trained delivers stable and satisfactory performance in forecasting the future states of multiple meteorological variables. In contrast, FourCastNet exhibits a rapid drop in performance at very short lead times and becomes unable to predict beyond 24 hours. The reason for this is that our W-MAE model, through pre-training, can learn the basic structure and patterns within the data from the limited training data. Subsequently, the pre-trained W-MAE model utilizes the learned fundamental data features which serve as prior knowledge for the model, making it more robust in downstream task fine-tuning. However, traditional models in the case of inadequate training data encounter two limitations: (1) limited capacity to capture latent patterns and features, leading to overfitting or poor performance, and (2) more prone to be affected by noise and outliers, reducing predictive robustness. These limitations may account for the underperformance of FourCastNet.

Under the setting of thirty-seven years of training data, our W-MAE continues to outperform FourCastNet, maintaining stable prediction performance. It is important to note that FourCastNet incorporates a distinct approach to mitigate accumulation errors (see DBLP:journals/corr/abs-2202-11214 for more details), necessitating a second round of training for multi-variable forecasting. To ensure fairness and enable a straightforward comparison, we adopt a first-round-training model of FourCastNet with a training epoch of 150 (matching the count mentioned in FourCastNet). It is observed that the performance of FourCastNet, which lacks the integration of a tailored training process for error accumulation suppression, experiences a significant decline beyond a certain forecast lead time. These findings underscore the superiority of our W-MAE framework in maintaining performance stability and robustness of results. Additionally, in Figure 3, we observe that beyond a certain lead time (approximately 130 hours), the performance of our W-MAE model with thirty-seven years of training data is not as strong as that achieved with two years of training data. This may be attributable to the fact that our W-MAE model with thirty-seven years of training data has not yet converged to its optimal performance. As mentioned in Section 5.1, our W-MAE model with two years of training data undergoes 600 epochs of pre-training, while the model with thirty-seven years of training data is only pre-trained for 200 epochs.

Moreover, we conduct an analysis of the time overhead for fine-tuning our pre-trained W-MAE model in multi-variable forecasting tasks. The fine-tuning process takes 100 epochs, with each epoch lasting an average of 63 minutes on 8 NVIDIA Tesla A800 GPU, resulting in a total fine-tuning time of nearly 105 hours. In comparison, FourCastNet requires 150 epochs of training, averaging 55 minutes per epoch on the same platform resulting in a total training time of around 137 hours. The results indicate that our W-MAE achieves a nearly 23% speedup compared with FourCastNet. This demonstrates that by leveraging the powerful initial parameters provided by pre-training, our W-MAE can achieve fast convergence during the fine-tuning process, which saves training time and computational resources.

Total Precipitation (TP) in the ERA5 dataset is a variable that represents the cumulative liquid and frozen water that falls to the Earth’s surface through rainfall and snowfall. Compared with other meteorological variables, TP presents sparser spatial characteristics. It is calculated from other variables, requiring an intermediate calculation in the derivation process, making it more challenging. For the precipitation forecasting task, we aim to predict the cumulative total precipitation in the next 6 hours using multiple meteorological variables from the previous time-step. We compare the performance of our W-MAE model and FourCastNet using both ACC and RMSE as the evaluation metrics.

Experimental results for different training data settings, i.e. using two years or thirty-seven years of training data, are shown in Figure 5¹¹1Note that, only a portion of FourCastNet’s prediction results are displayed in Figure 5. This is because the inability of FourCastNet to make predictions at longer lead times of more than 108 hours for two years of training data. Specifically, our experimental results show that in precipitation forecasting, FourCastNet produces considerable prediction errors beyond a certain lead time, resulting in NAN (not a number) predictions that render the results infeasible.. When using two years of training data, our W-MAE model outperforms FourCastNet in predicting the total precipitation at the 6-hour result. When the training data spans thirty-seven years, at the first forecast time-step (i.e., with a forecast lead time of 6 hours), our model achieves a 20% improvement in terms of the ACC metric compared with FourCastNet (0.996 vs. 0.804). At the 20$th$ forecast time-step (i.e., with a forecast lead time of 5 days), our model significantly outperforms FourCastNet (0.886 vs. 0.234). It is worth noting that despite incorporating the error accumulation suppression optimization techniques and undergoing a second round of training (see DBLP:journals/corr/abs-2202-11214 for detailed training process), the yielding FourCastNet model still exhibits inferior performance in precipitation forecasting compared to our W-MAE model. Moreover, concerning the time costs associated with precipitation forecasting tasks, our W-MAE framework outperforms FourCastNet, with a total training time of 20.6 hours and 23 hours, respectively. These observations suggest that our W-MAE not only achieves superior performance but also reduces the training time, resulting in efficient model training.

The methodology of ensemble forecast mainly involves adding noise to perturb initial weather states and observing the change in forecast results. In this paper, we follow FourCastNet DBLP:journals/corr/abs-2202-11214 to use Gaussian noise to perturb the initial conditions and generate an ensemble of 100 perturbed initial condition sets (represented as $E$=100). The perturbations are scaled by a factor $\sigma$=0.3. We apply large-member ensemble on multi-variable forecasting tasks for $Z_{500}$. As shown in Figure 6, we compare control forecast (i.e., the unperturbed forecasts) with the mean of a 100-member ensemble forecast using W-MAE. For the short-term forecast time, such as for 1 day, 100-member ensemble forecasts exhibit a slightly lower accuracy than control forecasts by a single model. However, when predicting beyond 5 days, the accuracy of 100-member ensemble forecasts displays a marked improvement compared with control forecasts. This observation affirms that the large-member ensemble forecast is valuable when the precision of individual models is compromised. Moreover, our W-MAE takes an average of 320 ms to infer a 24-hour forecast on one NVIDIA Tesla A800 GPU, which is orders of magnitude faster than traditional NWP methods in ensemble forecast. The results prove the efficiency of our W-MAE on ensemble forecast.

Effect of pre-training. To fully validate the effectiveness of our method, we apply our pre-training technique to FourCastNet. During the pre-training stage, we set the patch size to $8\times 8$ (same as that used in FourCastNet), the mask ratio to 0.75, and trained the model for 250 epochs. Then, we fine-tune the pre-trained FourCastNet for the precipitation forecasting task. We compared the performance of the pre-trained FourCastNet with that of the vanilla FourCastNet without pre-training. Our ablation results, as shown in Figure 7, indicate that the pre-trained FourCastNet outperformed FourCastNet without pre-training by nearly 30% in precipitation forecasting when both were trained on two years of data. Moreover, the pre-trained FourCastnet outperformed FourCastNet without pre-training, which was trained on thirty-seven years of data, by nearly 20%. Furthermore, in both training data settings, i.e., using two years and thirty-seven years of training data, FourCastnet with pre-training shows a significant performance improvement compared with that without pre-training. This suggests that applying pre-training to weather forecasting is promising.

Effect of masking ratio. Given the disparity in information density between meteorological images and images in computer vision, we conducted an investigation on the effect of mask ratios on meteorological image reconstruction. Specifically, the training loss under different mask ratios are compared and we set the maximum pre-training epoch to 1000. Figure 8 displays the effect of masking ratios on pixel reconstruction, with the optimal ratio found at approximately 0.6. However, when considering the time overhead of pre-training, we discover that setting the mask ratio to 0.75, compared with 0.6, saves 8% of the computational memory usage and reduces the training time by nearly 4 minutes per epoch. Moreover, the performance of the model at a 0.75 mask ratio can be comparable to that of 0.6. Therefore, to strike a balance between computational efficiency and model performance, we select a mask ratio of 0.75 for model pre-training. Furthermore, Figure 9 visualizes the pixel reconstruction results under varying mask ratio conditions.