LightWeather: Harnessing Absolute Positional Encoding for Efficient and Scalable Global Weather Forecasting

We consider $N$ weather stations and each station collects $C$ meteorological variables (e.g., temperature). Then the observed data at time $t$ can be denoted as $\mathbf{X}_{t}\in\mathbb{R}^{N\times C}$. The 3D geographical coordinates of stations are organized as a matrix $\mathbf{\Theta}\in\mathbb{R}^{3\times N}$, which is naturally accessible in station-based forecasting. Given the historical observation of all stations from the past $T_{h}$ time steps and optional spatial and temporal information, we aim to learn a function $\mathcal{F}(\cdot)$ to forecast the values of future $T_{f}$ time steps :

where $\mathbf{X}_{t-T_{h}:t}\in\mathbb{R}^{T_{h}\times N\times C}$ is the historical data, and $\mathbf{Y}_{t:t+T_{f}}\in\mathbb{R}^{T_{f}\times N\times C}$ is the future data.

Spatial encoding provides the geographical knowledge of stations to the model, which can explicitly model the spatial correlations among worldwide stations. Specifically, we encode the geographical coordinates of the station into latent space by a simple fully connected layer, thus spatial encoding $\mathbf{S}^{i}\in\mathbb{R}^{d}$ can be denoted as:

where $\mathbf{\Theta}^{i}\in\mathbb{R}^{3}$ represents the coordinates of the station $i$.

Temporal encoding provides real-world temporal knowledge to the model. We utilize three learnable embedding matrices $\mathbf{T}\in\mathbb{R}^{24\times d}$, $\mathbf{D}\in\mathbb{R}^{31\times d}$ and $\mathbf{M}\in\mathbb{R}^{12\times d}$ to save the temporal encodings of all time steps (Shao et al. 2022). They represent the patterns of weather in three scales ( $\mathbf{T}$ denotes hours in a day, $\mathbf{D}$ denotes days in a month and $\mathbf{M}$ denotes the months in a year), contributing to model the multi-scale temporal correlations of weather. We add them together with data embedding to obtain $\mathbf{H}^{i,j}_{t}$:

The effectiveness of LightWeather lies in the fact that absolute positional encoding integrates geographical coordinates and real-world time features into the model, which are intrinsically linked to the evolution of atmospheric states in global weather system. Here we theoretically demonstrate this relationship.

Considering the use of historical data spanning $T_{h}$ steps for prediction, it is not difficult to draw the corollary:

The detailed proof is provided in Appendix A.1. According to Eq. (11), we conclude that predictions for the future are bifurcated into two components: the fitting to historical observations and the modeling of the function $G(\lambda,\varphi,z,\tau)$, which represents the spatial-temporal correlations. When the scale is small, e.g., in single-station forecasting, even a simple linear model can achieve great performances (Zeng et al. 2023). However, as the scale expands to global level, modeling $G$ becomes the key bottleneck of forecasting.

The majority of prior models are designed to fit historical observations more accurately by employing increasingly complex structures. As shown in Figure 3 (a) (b), they simplistically regard $G$ as a function of historical values, and the complex structures may lead to over-fitting of it. In comparison, LightWeather can explicitly model $G$ with $\lambda,\varphi,z,\tau$ introduced by absolute positional encoding, as shown in Figure 3 (c), thereby enhancing the predictive performance.

We theoretically analyze the efficiency of LightWeather from the perspectives of parameter volume and computational complexity.

The proof is provided in Appendix A.2. According to Theorem 2, we conclude that the parameter count of LightWeather is independent of the number of stations $N$. In addition, the computational complexity of LightWeather grows linearly as the increase of $N$. In canonical spatial correlation modeling methods (Wu et al. 2019), it will cause quadratic complexity and linear increase of parameters with respect to $N$, which is unafforable with large-scale stations. However, LightWeather can effectively model the spatial correlations among worldwide stations with $N$-independent parameters and linear complexity, scaling perfectly to fine-grained data.

We conduct extensive experiments on 5 datasets including worldwide and nationwide:

• •

  GlobalWind and GlobalTemp (Wu et al. 2023b) contains the hourly averaged wind speed and temperature of 3,850 stations around the world, spanning 2 years with 17,544 time steps.

• •

  Wind_CN and Temp_CN contains the daily averaged wind speed and temperature of 396 stations in China, spanning 10 years with 3,652 time steps.

• •

  Wind_US (Wang et al. 2019) contains the hourly averaged wind speed of 27 stations in US, spanning 62 months with 45,252 time steps.

We partition all datasets into training, validation and test sets in a ratio of 7:1:2. More details of the datasets are provided in Appendix B.1.

We compare our LightWeather with the following three categories of baselines:

• •

  Classic methods: HI (Cui, Xie, and Zheng 2021), ARIMA (Shumway, Stoffer, and Stoffer 2000).

• •

  Universal DL methods: Informer (Zhou et al. 2021), FEDformer (Zhou et al. 2022), DSformer (Yu et al. 2023), PatchTST (Nie et al. 2023), TimesNet (Wu et al. 2023a), GPT4TS (Zhou et al. 2023), Time-LLM (Jin et al. 2024), DLinear (Zeng et al. 2023).

• •

  Weather forecasting specialized DL methods: AirFormer (Liang et al. 2023), Corrformer (Wu et al. 2023b), MRIformer (Yu et al. 2024b).

In Appendix B.2, we provide a detailed introduction to the baselines.

We evaluate the performances of all baselines by two commonly used metrics: Mean Absolute Error (MAE) and Mean Squared Error (MSE).

Consistent with the prior studies (Wu et al. 2023b), we set the input length to 48 and the predicted length to 24. Our model can support larger input and output lengths, whereas numerous models would encounter out-of-memory errors, making comparison infeasible. We adopt the Adam optimizer (Kingma and Ba 2014) to train our model. The number of layers in MLP is 2, and the hidden dimensions are contingent upon datasets, ranging from 64 to 2048. The batch size is set to 32 and the learning rate to 5e-4. All models are implemented with PyTorch 1.10.0 and tested on a single NVIDIA RTX 3090 24GB GPU.

Absolute positional encoding is the key component of LightWeather. To study the effects of absolute positional encoding, we first conduct experiments on models that are removed the spatial encoding and the temporal encoding respectively. The results are shown in Exp. 1 and 2 depicted in Table 3, where we find that the removal of each encoding component leads a decrease on MSE. This indicates that both spatial and temporal encodings are beneficial.

Then we make a comparison between relative and absolute positional encoding. Specifically, relative spatial encoding embeds the indices of stations instead of the geographical coordinates. Kindly note that we project the temporal dimension of data into hidden space, thus there is no need for relative temporal encoding. The results are shown Exp. 3 of Table 3. It is notable that the adoption of relative positional encoding causes a degradation in performance. Moreover, comparison between Exp. 1 and 3 reveals that the use of relative positional encoding even results in poorer performance than the absence of positional encoding, which further substantiates the effectiveness of absolute positional encoding.

We investigate the effects of two important hyperparameters: the number of layers $L$ in MLP and the hidden dimension $d$. As illustrated in Figure 4 (a), LightWeather achieves the best performance when $L=2$, whereas an increase in $L$ beyond $2$ results in over-fitting and a consequent decline in model performance. Figure 4 (b) shows that the metrics decrease with the increment of hidden dimension and begin to converge when $d$ exceeds $1024$. For reasons of efficiency, we chose $d=64$ in our previous experiments, but this selection did not yield the peak performance of our model. Moreover, it should be emphasized that LightWeather can outperform other Transformer-based models even when $d=32$ and the parameters is less than 10k. This further substantiates that absolute positional encoding is more effective than the complex architectures of Transformer-based models.

To intuitively comprehend the collaborative forecasting capabilities of LightWeather for worldwide stations, we present a visualization of the forecasting results here. Kriging is employed to interpolate discrete points into a continuous surface, facilitating a clearer observation of the variations.

The forecasting results are shown in the right column of Figure 5, while the ground-truth values are in the left. Overall, the forecasting results closely align with the ground-truth values, indicating that LightWeather can effectively capture the spatial-temporal patterns of global weather data and make accurate predictions.

To better interpret the effectiveness of our model, we visualize the learned embeddings of absolute positional encoding. Due to the high dimension of the embeddings, t-SNE (Van der Maaten and Hinton 2008) is employed to visualize them on two-dimensional planes. The results are shown in Figure 6.

Figure 6 (a) shows that the embeddings of spatial encoding tend to cluster. Similar embeddings are either spatially proximate or exhibit analogous climatic patterns. Besides, the embeddings in Figure 6 (b) and (d) form ring-like structures in temporal order, revealing the distinct daily and annually periodicities of weather, which is consistent with our common understanding.