Prompt Federated Learning for Weather Forecasting: Toward Foundation Models on Meteorological Data

1 Introduction

Climate change will significantly impact all regions; however, the specific effects will vary [Kjellstrom et al., 2016]. Increasing global temperatures and melting ice will lead to alterations in sea levels, ocean currents, weather patterns, and cloud cover [Hagemann et al., 2013]. To effectively tackle the challenge of global climate change, the implementation of a large-scale collaborative data-sharing platform is essential. Although this work is labor-intensive and demands a multitude of skilled experts, the utilization of machine learning techniques can enhance efficiency in addressing this problem. Nonetheless, the machine learning domain faces a challenge when attempting to employ a centralized uniform model to serve all regions due to their heterogeneity. An effective solution involves pre-training a foundational model using extensive weather data and enabling each region to fine-tune the model using a relatively small data to enhance its ability to capture local weather patterns.

Weather forecasting is a fundamental analytical task aimed at modeling the dynamic changes of weather on both a global and regional scale. Multi-sensor weather forecasting serves as a critical tool in mitigating the loss of human lives and property by providing early warnings for extreme weather events resulting from global climate change [Chattopadhyay et al., 2020]. The objective of this approach is to capture potential correlations between multiple meteorological factors and the tendency of weather variations in order to gain a comprehensive understanding of specific regions. Unlike conventional time-series data, weather time-series data in meteorology are gathered from sensing devices distributed across diverse geographical locations [Campbell and Diebold, 2005].

One-step forecasting [Chen et al., 2022b; Chen et al., 2023], empowered by recent advancements in deep learning, have garnered considerable attention due to their high efficiency. However, the performance of these strategies is hindered by the non-stationary nature of weather changes. This limitation arises from their reliance on fixed patterns derived from prior knowledge. To address this issue and capture temporal information, other studies [Karevan and Suykens, 2020; All´eon et al., 2020] suggest formulating the tasks as autoregression problems. These studies utilize preceding time step variables to predict variables at the subsequent time step [Chen and Lai, 2011]. This technique is commonly implemented using RNN or Transformer models. The choice of these models is based on their superior performance in time series analysis [Shi et al., 2015].

However, previous works focus on using an uniform model to serve all regions regardless of heterogeneity. In contrast, foundation model (FM) represents a novel service architecture that aims to pre-train a large model with extensive data. Subsequently, this model can be fine-tuned for specific tasks using relevant data, such as weather forecasting in a particular region. The FM has the capability to capture common knowledge shared among multiple tasks or participants. Further refinements can then enhance its alignment with the specific requirements of a given task. The FM has demonstrated remarkable success in Natural Language Processing (NLP), exemplified by ChatGPT. Notably, recent progress in foundation models has been observed across diverse domains, including ViT [Xu et al., 2022], BERT [Yates et al., 2021], and CLIP [Radford et al., 2021].

In contrast to existing FMs, training a FM on weather forecasting tasks must tackle the following challenges. First, sharing raw data across countries/regions will not be easy. Second, transmitting and processing the continuously collected data is a challenge for low-resourced sensors or devices. Third, real-time forecasting is critically important. In summary, we need a solution to tackle data security, communication, and computation efficiency issues and provide on-device decisions independently.

This paper will design a novel machine-learning approach to train foundation models on weather forecasting tasks. The model will be capable of understanding and constructing the complex spatiotemporal relationship of meteorological data to provide reliable analysis support on weather forecasting and global climate challenges. Specifically, we propose a novel Meteorological Prompt Federated Learning (MetePFL) approach to collaboratively learn a Transformerbased foundation model (FM) across devices with multivariate time-series data (see Figure 1). The MetePFL only considers the model parameters exchange among devices rather than direct data sharing. Considering the low-resourced sensors' communication efficiency constraint, a brand-new prompt learning mechanism is introduced upon a pre-trained FM to comprehensively explore the correlation among weatherrelated variables while computing a few parameters.

Three weather forecasting datasets based on multivariate time series with multiple meteorological factors, i.e., precipitation, temperature, upstream, are leveraged to verify the effectiveness of MetePFL. The main contributions of this work are summarized in four-fold:

• This is the first work to explore a foundation modelbased solution to enhance weather forecasting tasks towards a global scale.
• The proposed prompt federated learning approach is a novel mechanism to collaboratively learns a foundation model for the applications with many satellite sites or stations across regions.
• A spatio-temporal prompt learning mechanism has been designed to efficiently tackle multivariable time series.
• Experiments on three real datasets have demonstrated the effectiveness and superior of our proposed solution. It is worth noting that we obtain excellent performance with only 2.38% of the model's parameters trained.

2 RelatedWork

Weather Forecasting. Weather forecasting plays a crucial role in the global climate analysis system. Conventional forecasting methods utilize numerical weather prediction (NWP) models, which incorporate physical constraints to simulate weather phenomena [Bauer et al., 2015]. However, with the emergence of data-driven approaches, weather forecasting has shifted towards approaches driven by data, such as ARIMA [Chen and Lai, 2011], SVM [Sapankevych and Sankar, 2009], and NNs [Voyant et al., 2012]. While these basic models exhibit potential, they encounter difficulties in comprehending nonlinear temporal dynamics. Deep Learning methods, especially models based on Recurrent Neural Networks (RNNs), have exhibited promising outcomes in weather forecasting [Shi et al., 2015]. Recently, Transformers have demonstrated superior performance compared to RNN-based models in time series analysis [Bojesomo et al., 2021], thereby gaining popularity for weather-related tasks [Chen et al., 2023]. However, these models neglect the data exposure concerns when utilizing multi-sensor data for real-world tasks and modeling spatio-temporal correlations.

Foundation Model. A fully supervised learning paradigm needs a large scale of data. The foundation model provides a practical solution for scenario-specific tasks, aiming to pretrain a model using extensive prior knowledge. The FM has widespread applications in NLP [Yates et al., 2021] and CV [Xu et al., 2022; Radford et al., 2021], providing an effective cross-task learning strategy. For example, ChatGPT can be used as a baseline model for researchers to fine-tune it to achieve more accurate responses for downstream tasks.

Federated Learning. Federated learning (FL) is a new learning paradigm to embody the collaborative training of models without requiring data exposure from each participant (e.g., meteorological sensors) [McMahan et al., 2017; Long et al., 2021; Long et al., 2020; Jiang et al., 2020]. Vanilla FL suffer from the heterogeneity of data and devices. Personalized FL aims to solve the problem by multiple techniques [Tan et al., 2022a; Chen et al., 2022a; Wang et al., 2022; Ma et al., 2022; Tan et al., 2021; Zhang et al., 2023; Li et al., 2023; Long et al., 2023; Li et al., 2019; Li et al., 2020; Gao et al., 2022] via training better personalized model for each client. However, these methods are not suitable for weather forecasting due to all parameters must be considered during communication so that hinder the realtime forecasting. pre-train-based strategy can mitigate the problem [Tan et al., 2022b] but can not explore the spatiotemporal correlations. Different from the above methods, this paper focuses on establishing a high-efficiency FL approach that provides analytical support to across-regional weather forecasting systems with heterogeneous meteorological data.

Prompt Learning. Prompt learning as a lightweight mechanism is widely used in NLP [Li and Liang, 2021; Liu et al.,

2021], which requires fewer parameters and is more adaptive than fine-tuned pre-trained models by represented by several prompt tuning strategies in different applications [Zhou et al., 2022a]. Different from language data, understanding multidimensional correlation among multivariate data in the weather forecasting task is critical. However, the key point is often ignored by the federated prompt learning method [Guo et al., 2022]. The paper introduces a novel prompt mechanism within the FL framework based on pre-trained FM to explore the temporal dynamics and the potential correlation among clients while computing only a few parameters.

3 ProblemFormulation

Given N clients that possess individual local private datasets D, each client has a multivariate time series denoted as Xi ∈ Rm×n . In this notation, each sample at a specific time step t is represented as xt ∈ R1×n . Weather forecasting using multivariate time series can be defined as the process of utilizing historical values of all variables for a duration of P periods to predict the values of a specific variable in the future over Q periods, can be defined below:

where f is a learning system, and x′t ∈ R1×1 is the value of the variable to be forecasting at the t-th time step. Valida FL system aims to minimize the average loss of the global model w on all clients' local dataset:

where nk is the number of samples hold by the k-th client. n is the number of samples held by all clients. Fk(wk) denotes the local objective function of k-th client. The distinguishing factor is that each client possesses a unique pattern, and sensors are deployed in specific locations, resulting in a statistically heterogeneous environment. To address this challenge, PFL is typically modeled as a bi-level optimization problem.

where each client hold a personalized model parameterized by vi, w denotes the global model. R(·) is the regularization term to control model update, via avoiding the local model updating be far away to the optimal global model.

4 Meteorological Prompt Federated Learning

The framework of MetePFL is depicted in Figure 1. In contrast to conventional Federated Learning (FL) where random global parameters are broadcasted to each client, MetePFL employs a fixed FM, thereby reducing computation costs and improving performance without requiring extensive backpropagation. During each round, only the prompt parameters of the clients are taken into consideration. The MetePFL framework consists of the Spatial-Temporal Prompt (STP) and the optimization process.

4.1 Spatial-Temporal Prompt

The Spatial-Temporal Prompt (STP) shown in Figure 2 can be divided into Temporal prompt learning (TPL) and Spatial prompt learning (SPL).

Temporal prompt learning (TPL). To effectively understand the temporal dynamics under multivariate interactions, we use a multi-step incremental learning mechanism for learning the prompt parameters along the time dimension. The TPL comprise four phases, is shown in Figure 3.

Given a time-series X ∈ Rm×n , the m and n represent the number of time steps and variables, respectively. Define a initial step l, the four temporal prompt format: (1) initial prompt ˆ P ; (2) Temporal prompt-I: P1; (3) Temporal promptII: P2 ∈ R2l×n ; (4) Temporal prompt-III: P3 ∈ R2l×n .

The first p steps are fixed as Seq ∈ Rp×n . In addition, the data within the p to k steps are spliced with the ˆ P ∈ R(k-p+l)×n to generate P1 by

where ∥∥ denote the concat operation. To enhance the accuracy of forecasting performance, we adopt a two-stage objective in the initial learning phase: "X(0∼p) 1 -→X(0∼k) 2 -→ X(0∼k+l) ". During this phase, we encourage the FM to enhance prompt parameters associated with prior knowledge in X(p∼k) in order to establish a foundation for the subsequent stage. In the next stage, the FM learns the remaining PT,1 parameters in X(k∼k+l) . This approach enables the FM to develop a comprehensive understanding of the relationships among variables across different time periods, rather than establishing rigid back-and-forth associations. The formulation of the initial learning phase is:

where the embed(x) represents the learnable position encoder, and the FM is the pre-trained foundation model.

Splicing discrete prompts to form a complete series is challenging because their parameters cannot match the FM simultaneously. To establish the relationship between two prompts, a dual-correct strategy is adopted in the second learning phase. This strategy encourages the FM to correct PT,1 based on the previous PT,1 while learning PT,2. The objectives of the second learning phase are "X(0∼k+l) -→ X(0∼k+3l) ." This process can be formulated as mapping the input data fromX(0∼k+l) toX(0∼k+3l) during the second phase.

The objective of the third learning phase can be expressed as "X(0∼k+3l) -→X(0∼k+5l) ", for further correcting these previously learned prompts to improving the continuity of several prompts and providing smooth transitions between them. The third learning phase can be formulated as:

To prevent the prompt parameters from being overly biased toward expressing short-term over long-term dependence, the final stage of learning intends to uniformly adjust these parameters. The corresponding values from the three previous learning phases are concatenated along the time dimension and then multiplied by the uniform prompt PT,w. The final learning can be expressed as follows.

Spatial Prompt Learning (SPL). We regard multiple meteorological factors within a specific space. This enables the establishment of correlations between these factors from a spatial perspective on the local client. The SPL are shown in Figure 4.

The trainable parameters that serve as spatial prompts can be represented as PS ∈ R(m-k)×1 , where m - k represents the length of the forecasting period. Prior to the initial learning phase, the first p hours of the data are considered fixed and denoted as Seq ∈ Rk×n . Subsequently, Seq are combined with the initial spatial prompts along the temporal dimension, and any gaps in the spatial dimension are filled with zero-valued parameters, resulting in ˆ PS ∈ R(m-k)×(n-1) .

The first learning phase comprises:

In learning spatial prompts for the ith variable, the ith learning phase of SPL can be formulated as follow:

One advantage of SPL over auto-regression is its continuous correction of learned prompt parameters from the previous iteration. This correction utilizes all the previously predicted values, leading to improved forecast accuracy. Moreover, SPL improves the model's perception of correlations among multiple variables by considering adjacent variables together. Through multiple iterations of regression, SPL enhances the model's ability to account for these spatial correlations.

To prevent isolating the TPL and SPL effects within the STP, we combine the result of TPL and SPL empirically using a Gate operation. The output of the STP is formulated as R = [(1 - sigmoid(RS)) ∗ tanh(RT )] ∗ W .

4.2 Optimization of MetePFL

The optimization objective of MetePFL is

where {Pk} is prompts including PS, PT , {P g } is the global prompt parameters, G(A) is a regularization term used to represent the correlation among client according the adjacent matrixA. specifically, the termFk({P})+λR({P}, {P g }) can be formulated as Lk = LMSE + λ∥{Pk} - {P g }∥22 . The optimization objective on i-th client as:

Algorithm 1 MetePFL algorithm.

Initialized PT , PS . for each communication round t = 0, 1, 2, · · · , T do Local model initialize: for each client i = 0, 1, 2, ..., N in parallel do 0 ← PT , PS Local model update: for each client i = 0, 1, 2, · · · , N in parallel do Update PT , PS for e local steps: Train PT , PS with loss function L in Eq.(11) end for Each selected client sends PT , PS to the server Aggregation: A←GraphGenerator({P1}, {P2}, · · · , {PN}) A′ ← GAT (A) Update P for r steps GCN(A′ , {Pi}Ni=1): {P s i }Ni=1 =A′ {Pi}Ni=1 {P s i }Ni=1 = α{P s i }Ni=1 + (1 - α){Pi}Ni=1 Get Global Prompts {P g } ← G({P s 1 }, {P s 2 }, ..., {P sN}) end for

where the A ∈ {0, 1}, S({Pj}, {Pi}) is the similarity of prompt parameters of client i and client j measured by cosine or distance, G(·) is the average operation. During the optimization of MetePFL, each client update their model via solving the local objective function as Lk after them receive the fixed foundation model at first. Then each client upload their prompts P rather than complete model to the server that conduct graph-based aggregation, which significantly reduce the communication overhead while exploring the potential correlations among clients. The aggregation including two steps: Graph Attention Network (GAT) [Veliˇckovi´c et al., 2017]-based graph structure learning that explore the dynamic correlations among clients and Graph Convolution Network (GCN) [Kipf and Welling, 2016] that utilized to parameters reconstruction using learned adjacent matrix A and the prompt parameters uploaded by clients. The GCN automatically updates the parameters of each node by aggregating the models of its neighbors in the graph.

5 Experiments

Baselines. We compare our MetePFL with STGCN [Yu et al., 2017], LSTM [Graves, 2012], ConvLSTM [Shi et al., 2015], Transformer [Zerveas et al., 2021], Informer [Zhou et al., 2021], Autoformer [Wu et al., 2021], and Fedformer [Zhou et al., 2022b]. The LSTM-based models have four layers. The Transformer consists of an eight-layer Encoder, while the Informer, Autoformer, and FEDformer consist of two encoders and a decoder. The Transforme models are trained by data-centric and FL setting, respectively.

Datasets. We compiled three multivariate time series datasets from NASA 1 , Average Precipitation (AvePRE), Surface Temperature (SurTEMP), and Surface Upstream (SurUPS) collected by 88, 525, and 238 devices, respectively. All three datasets cover the hour-by-hour variability of 12 different weather-related meteorological variables.

Experimental Setups. Models' input and output dimensions are C = 12 and C = 1, respectively. The dataset is split into training, validation, and testing in a 6:2:2 ratio. For pretraining, we use 2/3 of the training set (i.e., 50% of the entire dataset) for training and 1/6 as the validation set (i.e., 10% of the complete dataset) based on the above partition, following the pre-training strategy from Zerveas et al. [Zerveas et al., 2021]. For fine-tuning and prompt learning, the last 1/6 of the training set is used for training, while the validation and test sets remain unchanged. In the federated training process, we set k to control the number of clients participating in training per round, and we use k = 0.1 and k = 0.2 in the experiments. The forecasting uses 12 time steps in history (P=12, i.e., the past twelve hours) to predict 15 time steps in the future (Q=15, i.e., fifty hours in the future) with a time window length of 27 h. We set l to 3, considering the validity time and trigger threshold of weather events. All models were trained on an NVIDIA Tesla V100 GPU using an initial learning rate of 1e-3 and a batch size of 128, ADAM [Kingma and Ba, 2014] as the optimizer. The algorithm used in the FL of the Transformer-based network is FedAvg. The communication round was set to 20, α = 0.5, and the early stopping strategy was applied. Mean absolute error (MAE) and root mean absolute error (RMSE) as evaluation metrics. The code is avaliable at https:// github.com/ shengchaochen82/MetePFL.

5.1 Overall Comparison

Table 1 presents a performance comparison between our MetePFL and baselines. Results indicate the superiority of the Transformer-based model over the LSTM-based model and STGCN when trained centrally across all three datasets. However, when utilized as a FM within the FL framework, the performance of the Transformer-based model is diminished compared to the trained FM. This reduction can be attributed to the heterogeneity of weather data collected from multiple sensors. Notably, while the Transformer may be less effective than other similar models in centralized training, it demonstrates a significant performance advantage over FedFEDformer, Fed-Autoformer, and Fed-Informer when employed as an FM within the MetePFL. This observation suggests that STP enhances the Transformer's capability to comprehend spatiotemporal data. Furthermore, our reliable FM and aggregation algorithms (Transformer and FedAvg) surpass PromptFL [Guo et al., 2022]. This outcome provides additional validation of the effectiveness and superiority of our proposed MetePFL.

1 https://disc.gsfc.nasa.gov/

To evaluate the effectiveness and superiority of MetePFL over fine-tuning, we implement different MetePFL version based on four FL algoirthms: FedAtt [Jiang et al., 2020], FedProx [Sahu et al., 2018], Scaffold [Karimireddy et al., 2020], and FedAvg [McMahan et al., 2017]. We maintained the prompt setting and compared the results against fine-tuning. The results are presented in Table 2. Our experiments reveal the following findings: (1) MetePFL outperforms the finetuning method, highlighting the sensitivity of data correlation to the STP approach; (2) the graph-based aggregation used in MetePFL surpasses other algorithms, indicating its effectiveness in mitigating the negative impact of Non-IID.

5.2 Framework Applicability

To determine the applicability of MetePFL, we replaced its FM with pre-trained Informer and Autoformer using the same pre-training strategy. Additionally, we used fine-tuning and PromptFL as reference strategies. The results as shown in Table 3. Our proposed MetePFL remains valid for other Transformer FMs and significantly outperforms fine-tuning and PromptFL. By comparing the performance of MetePFL under different FL algorithm, we demonstrate the effectiveness of the graph-based aggregation once again.

5.3 Parameter Utilization

Table 4 compares parameter utilization for different federated prompt learning strategies. PromptFL and MetePFL are significantly more advantageous than regular training (train from scratch) and fine-tuning, with parameter utilization of 2.22% and 2.38%, respectively - nearly 28% lower than fine-tuning. While PromptFL is better than MetePFL in parameter utilization (-0.16%), it performs nearly 20% worse than MetePFL (see Table 1). Despite considering only 2.38% parameters, MetePFL achieves excellent performance and significantly improves inter-device communication efficiency.

5.4 Ablation Study

To evaluate the effectiveness of STP, we conducted ablation studies under general scenarios rather than under the FL framework. The results are presented in Table 5. The Encoder-only and Decoder-only Transformer underwent pretraining with one-step forecasting and auto-regressive mechanisms similar to previous experiments. The Encoder-only Transformer trained with STP outperformed the Decoderonly Transformer, indicating the superior learning mechanism of the proposed STP over conventional auto-regression. Moreover, the performance gap between the pre-trained Encoder with and without STP confirms the efficacy of the proposed STP in general scenarios.

The effectiveness of SPL, Gate, and their advantage over position-aware embedding (PE) was verified using the Encoder-only Transformer. The results are presented in Table 6. SPL demonstrates a greater improvement compared to PE. Additionally, SPL enhances the model's performance, while the model without Gate experiences a significant decline. TPL does not perform as effectively as desired in generic scenarios. Consequently, we conducted ablation experiments in the FL setting, as shown in Table 7. The results indicate that both TPL and SPL enhance the model's forecasting performance. In conclusion, the effectiveness of the Gate operation is demonstrated, and both TPL and SPL can improve the model's performance in the FL setting.

6 Conclusion

This paper proposes a novel machine learning approach to train foundation models for weather forecasting tasks, capable of capturing the spatiotemporal relationships of meteorological data based on multivariate time series. To enhance the performance while keeping data secure and reducing communication overhead, we introduce a prompt learning mechanism based on the fixed foundation model within the FL framework. Additionally, we utilize a graph-based approach to mitigate the impact of data heterogeneity on model effectiveness. Extensive experiments on three real-world weather datasets confirm the effectiveness of our proposed method.