Federated Prompt Learning for Weather Foundation Models on Devices

1 Introduction

Climate change has a profound impact on both natural ecosystems and human societies [Karl et al., 2009; Kjellstrom et al., 2016]. It leads to higher temperatures, sea level changes and more frequent extreme weather events [Hagemann et al., 2013]. As a result, precise weather forecasting is becoming increasingly important. Data from meteorological devices in various regions is vital. However, analyzing this data with deep learning through centralized cloud computing presents challenges such as network dependence and privacy concerns [Chakraborty and Rodrigues, 2020]. First, sending large volumes of data to centralized system places a heavy burden on communication networks, which is impractical for low-resource weather devices. Second, data from sensitive locations is subject to privacy laws, restricting sharing across devices [Chen et al., 2023a]. To address these issues, ondevice intelligence for analyzing data directly on the devices is crucial, reduces the need for data transfers, protects privacy, and decreases reliance on networks.

Federated Learning (FL) [McMahan et al., 2017] is a promising method for on-device intelligence that trains a uniform model collaboratively across multiple devices without exchanging data. However, the models often underperform due to statistical heterogeneity among clients and data homogeneity on within individual clients' data. Personalized FL (PFL) offers new insights by developing specialized models for each device, enabling tailored on-device intelligence [Paulik et al., 2021]. Recent PFL methods have introduced various methods to improve personalization [Chen et al., 2022; Tan et al., 2022; Li et al., 2021b]. Despite these advances, two significant challenges remain. First, there is often inadequate consideration of the impact of physical geographic location on local models. For example, devices on seashores and hilltops may collect different data types even if they are geographically close. Second, the substantial communication demands of large neural networks burden both clients and servers. Edge devices with limited resources may struggle to process the necessary updates for these complex models [Xiong et al., 2023]. Moreover, the transfer of entire model parameters hampers communication efficiency.

To tackle the above issues, this paper introduces Federated Prompt Learning for Weather Foundation Models on Devices (FedPoD), which allows devices to obtains high customized models with efficient communication. FedPoD comprising two pivotal components: (1) Adaptive Prompt Tuning and (2) Dynamic Graph Modeling. Adaptive Prompt Tuning against data homogeneity and reduces communication load via updating local prompts based on the frozen foundation model (FM) to capture local information and guide FM to generate accurate prediction, coupled with multi-level communication. Additionally, FedPoD uses Dynamic Graph Modeling on the server to manage prompts from clients and to build multiple graphs dynamically, considering various perspectives. This process takes geographic features into account and promotes priority collaborative learning among clients with similar data, mitigating the effects of data heterogeneity. As shown in Table 1, using a pre-trained foundation model leads to fewer parameters and higher performance compared to starting from scratch with FedAvg [McMahan et al., 2017]. Furthermore, FedPoD achieves the best results with the proposed adaptive prompt tuning and dynamic graph modeling.

Main Contributions. With extensive experiments across datasets including real-world on-deivce weather series datasets on various setting, we show that our FedPoD consistently outperforms state-of-the-art baselines. Besides, we conduct further analysis to provide more insights in FedPoD from the perspective of ablation, hyperparameter sensitivity, and privacy. The main contributions is as follows:

• We present FedPoD, a communication-efficient framework for on-device weather forecasting that addresses the challenges of data heterogeneity among devices and data homogeneity within individual clients during federated learning.
• We show Adaptive Prompt Tuning that uses prompts to represent information and guide generation. These prompts enable multi-level communication and knowledge sharing, reducing the impact of data homogeneity.
• We introduce Dynamic Graph Modeling to create dynamic links between participants' prompts. This prioritizes collaborative optimization for clients with similar representations, enhancing personalization.
• With extensive experiments, we show FedPoD consistently achieves the best and help improve the communication efficiency while keeping privacy, and adaptive prompt tuning also benefits baselines.

2 RelatedWork

Weather Forecasting. Weather forecasting is a crucial tool that analyzes the variations in weather patterns. Recently, weather forecasting has made significant strides by incorporating data-driven approaches [Chen and Lai, 2011; Sapankevych and Sankar, 2009; Voyant et al., 2012]. RNNs have shown promising in weather forecasting [Shi et al., 2015; Grover et al., 2015]. Besides, Transformers [Zhou et al., 2021; Zhou et al., 2022; Wu et al., 2021; Chen et al., 2023c] can capture non-stationary changes, which have contributed to their widespread use in weather analysis. To overcome challenges caused by intricate spatial-temporal correlation, spatial-temporal modeling methods [Yu et al., 2017] can be an effective solution. However, these methods all focus on data-intensive centralized training, which poses a challenge to weather forecasting practices.

Personalized Federated Learning. Weather forecasting involves significant communication loads and raises privacy issues due to the large volume of data processes on parallel [Chavan and Momin, 2017]. Federated learning (FL) [McMahan et al., 2017] offers a way to perform ondevice intelligence but is often hampered by data heterogeneity and homogeneity. Personalized FL (PFL) seeks to overcome these issues by training customized models for each device, providing fresh insights. For example, [T Dinh et al., 2020; Hanzely et al., 2020; Li et al., 2021a] add a regularization that decomposes the personalized model optimization from the global. [Li et al., 2021b; Collins et al., 2021] share part of the model and keep personalized layers private. [Zhang et al., 2020] enables a flexible method by adaptively weighted aggregation. [Fallah et al., 2020] start from a Model-Agnostic Meta-Learning where a meta-model is learned to generate the initialized local model for each client. In addition, [Chen et al., 2022] utilize structure information to explore the topological relations among clients.

3 Preliminaries and ProblemFormulation

Weather Forecasting. A multivariate weather time series represented by Xi ∈ Rm×n , where m and n is the series length and the number of variables, respectively. Each data point is shown as xt ∈ R1×n . The weather forecasting task can be divided into two categories below:

• Task 1-Multivariate to Univariate Forecasting: Predicting a single variable for future Q periods using all variables from the past P periods.
• Task 2-Multivariate toMultivariate Forecasting: Predicting all variables for future Q periods from all variables in the past P periods.

These tasks can be defined as follows:

where f denotes the learning system, xT1 t ∈ R1×1 is the predicted variable at the t-th step, and xT2 t ∈ R1×n is the predicted variable at the t-th step.

On-device Weather Forecasting based on FL. Each device 1 possesses a local data varying location pattern, leading to statistical heterogeneity. Thus, we can define the task ondevice weather forecasting as:

where the Dk and D′k denote the input dataset and prediction in k-th client, respectively, and fk is the personalized model for k-th client. This makes vanilla FL that train an uniform model unsuitable, and the task is updated to the PFL problem that solves below bi-level optimization.

1 We take "device(s)" and "client(s)" to mean the same one.

where each client hold a customized model parameterized by vi, w denotes the global model. R(·) is a regularization term, G(·) is the aggregation strategy. Previous studies have had difficulty managing the non.iid of geographic data, often overlooking how spatial-temporal correlation is affected by more than just location [Chen et al., 2023b]. In this work, we aim to address two main challenges: (1) How can we ensure efficient communication between clients and servers while guaranteeing the framework's high performance? (2) How can we minimize the heterogeneity caused by complex geographic features in the most cost-effective way?

4 Methodology

In this section, we detail our FedPoD, illustrated in Fig. 1. Each client hold a pre-trained FM (PFM) 2 and three types of prompts that updated locally. In each round, we introduce multi-level communication based on prompts uploaded by participants, including inter-clients and client-server. In the server, we present a novel aggregation method, Dynamic Graph Modeling, to building dynamic graphs based on structural information from prompts, reducing influence of data heterogeneity. With the updated prompts from the server, clients perform local optimization with a specialized promptwise loss. We'll describe them in more detail below.

Adaptive Prompt Tuning. We introduceAdaptive Prompts Tuning for local updating to minimize the effects of data homogeneity within devices while keeping computational loads low. Unlike prompt tuning in Natural Language Processing, which simply adjusts inputs to guide a pre-trained large language model (LLM) to produce outputs [White et al., 2023]. It involves using lightweight prompts that dynamically represent local knowledge and act as information carriers in multilevel communication. This helps lessen the overall impact of both global data heterogeneity and local data homogeneity during collaborative training. Specifically, we use trainable parameters as prompts, including TEMPORAL PROMPTS (PT ) and INTER-VARIABLES PROMPTS (PV ), to capture the local temporal dynamics and the relationships among variables. These prompts are incorporated into the time series and are refined during the local training phase. The updating process of PT and PV is shown in Alg. 1.

After updating the TEMPORAL PROMPTS (PT ) and INTER-VARIABLES PROMPTS (PV ), we apply two learnable matrices, Wt and Wv, to them. This is represented as X = PT ⊙Wt +PV ⊙Wv. These matrices help to adjust the significance of the prompts, ensuring they contribute to our optimization goal without straying off course. Furthermore, we introduce SPATIAL PROMPT (PS), to encode local geographic pattern for comprehensive modeling, via updating with original input Xipt and X. The final prediction X is then derived using the following formula:

2 Detailed information about the utilized PFM in Appendix B. Complete version: https://arxiv.org/abs/2305.14244.

Algorithm 1 Implementation of PT and PV Updating

Initialize Original input series Xipt, frozen PFM FM, Temporal/Variable updating steps Kt and Ks.

end for

▷ ∥.∥V : concat along variable dimension PV ← ∥PV ,P ′V ∈ Rm×Kv ∥V

end for

▷ P ′V : Next inter-variable prompt block

where Xgeo denotes the client's geographic location represented by (ϕ, λ), ϕ and λ is the latitude and longitude coordinates, respectively, for simultaneous updating of PS and X to adjust the parameters of PS.

Local Optimization from Multi-Task Perspective. For each client's local optimization, we focus on two key elements: (1) multi-level communication regularization and (2) a multi-task perspective. The first involves interactions among clients and between clients and the server, aiming to mitigate the effects of data homogeneity. The second treats the optimization of various prompts as separate tasks, helping to lessen the unpredictability that comes with mixed updates. Consequently, we suggest a prompt-based optimization objective for local updates, as follows:

where MSE(·) denotes the mean square error loss that evaluate the distance between ground-truth y and output y′ , R({Pi}; {Pj}l ; {Pi}l ; {P}∗ ) is the regularization term utilized to measure the distance between prompts, including personalized prompts {Pi}l of the i-th client, neighboring j-th client's prompts {Pj}l , and global prompts {P}∗ obtained by averaging all client's prompts. The underlying motivations is allowing local client to craft highly customized models via decomposing prompt parameters from the neighboring and global prompts while keeping the comprehensive knowledge. Then, inspired by [Kendall et al., 2018], we conceptualize the optimization from a multi-task view, can be formulated as:

Here, the ξ and τ are importance coefficients that obey λ, τ ∈ (0, 1), the L2 is L2 regularization (e.g., Euclidean distance, Cosine Similarity, etc.). SG represents the subgraph step used to adjust the scope of interaction between clients. The inter-client regularization term 1 τ 2 · 1 (|N |/SG)-1 ∑j∈N L2 ({Pi}, {Pj}), drives the local up-

▷ P ′T : Next temporal prompt block dating process towards a more comprehensive representation via considering neighboring clients with distinct feature distributions within a given range. Regulation terms 1 ξ2L2 ({Pi}, {P}∗ ) and 1 ξ2L2 ({Pi}, {Pi}l ) are employed with the purpose of prompting the local clients to attain a more personalized representation.

Dynamic Graph Modeling for Global Aggregation. We introduce Dynamic Graph Modeing (DGM) on the server to boost personalization by constructing spatial-temporal correlations among clients. This promotes collaborative optimization among clients with similar local knowledge representation. DGM uses the prompts shared by clients and their geographic data, like latitude and longitude, to form graphs. These graphs reveal possible relationships among clients, leading to a more customized optimization process. Specifically, we divide local prompts into three classes: (1) Temporal and Inter-Variables Prompts {PT,i,PV,i}Ni=1 ; (2) Spatial Prompts {PS,i}Ni=1 and (3) Full Prompts {Pi}Ni=1 , where N is the number of clients. First, the server generates a static graph Ageo according to the location information based on Haversine formula [Robusto, 1957] as:

where i, j ∈ N, i ̸= j, ϕi and ϕj are the latitude coordinates of client i and j, respectively, ∆ϕ = ϕi -ϕj is the difference in latitude between the two points in radians, ∆λ = λi - λj is the difference in longitude between the client i and client j, D is the calibre of the Earth.

To grasp the potential correlations between clients dynamically, we use two matrices, Wi are Wj , to apply linear transformations to the prompt vectors Pi and Pj of two different clients. The relation of the i-th client to the j-th client is calculated using the formula ei,j = α(WiPi,WjPj), where α(·) denotes a shared attention mechanism that operates in the space RF ′ × RF ′ → R. Here, W ∈ RF ′×F helps determine the attention coefficients. We then introduce another matrix W to calculate the weight of the connection (edge)

and construct an adjacency matrix as follows:

For three types of prompts, we create three corresponding adjacency matrices (graphs), denoted as ATV , AS, and A, via Eq. 8. We then merge these with AGeo (from Eq.7) using an attention mechanism to capture more precise correlation representations. Based on these matrices, we reconstruct prompts to deliver personalized prompts {Pi}l for each client:

where √dk is the dimension of adjacent matrix, and α is importance coefficient. The term [AGeo - AS] highlights the discrepancy between the actual geographic correlation and the encoded spatial correlation, enabling the dynamic adjustment of spatial-temporal correlation among clients to achieve a more precise potential correlation graph modeling.

Optimization for FedPoD. The overall optimization objective of FedPoD is to solve a bi-level optimization problem, as below:

where {P } denotes local prompts including PT , PV , and PS , {P}∗ is global prompts, the local model was parameterized by {P} after receiving the pre-trained FM. The {Pj}l is personalized local models from other clients that achieve by the additional regularization term G(·) that is a graph-based constraint that ensures each client aggregates with similar neighbor nodes. The learned graph with the adjacent matrix A

Algorithm 2 Implementation of FedPoD

1: Initialize local data {Di}Ni=1 , foundation model FM, prompts {PT,i,PV,i,PS,i}Ni=1 2: Initialize {PT,i,PV,i,PS,i}Ni=1 as {Pi}Ni=1 , Wbt,i,Wbs,i 3: Server-side: Broadcast frozen FM to each clients ▷ Model communication 4: for rounds R = 1, 2, 3... do ▷ FL rounds in sequence 5: Client-side: 6: Download {P}l (personalized prompts), {P}∗ (global prompts) from the server 7: for each client i in parallel do ▷ Clients in parallel 8: {Pi} ← LOCALUPDATE(FM , Di, {Pi}Ni=1 ) ▷ Prompt-based training 9: Upload {Pi} to the server ▷ Model communication 10: end for 11: Server-side: 12: Ageo ← HAVERSINE FORMULA(ϕ, λ) (Eq. 7) ▷ Generate the static graph 13: A ← DYNAMIC GRAPH MODELING({PT}i}Ni=1 , {PV }i}Ni=1 , {PS}i}Ni=1 ) (Eq. 8) ▷ Generate the dynamic graph 14: ATV ← DYNAMIC GRAPH MODELING({PT}i}Ni=1 , {PV }i}Ni=1 ) (Eq. 8) ▷ Generate the dynamic graph 15: AS ← DYNAMIC GRAPH MODELING({PT}i}Ni=1 , {PV }i}Ni=1 ) (Eq. 8) ▷ Generate the dynamic graph 16: A′ ← ATTENTION(A,ATV,AS,Ageo) (Eq. 9) ▷ Attention for filtering 17: {Pi}l,N i=1 ← αA{Pi}Ni=1 + (1 - α)A′ {Pi}Ni=1 (According to Eq. 9) ▷ Update personalized prompts 18: {Pi}∗ ← n nk ∑Ni=1 P s , wr ← n nk ∑Ni=1 wr,i ▷ Update global prompts and layers 19: end for 20: LocalUpdate(FM,D,PT,PV,PS,Flayer) 21: for local epoch e = 1, 2, ... do 22: Update PT,PV (Algorithm 1) 23: Update PS (Eq. 4) 24: Update rest of trainable parameters (Eq. 4) 25: Compute local loss (Eq. 6) ▷ Optimization from multi-task view 26: end for

(computed by A′ ,A) is expected to be sparse and able to preserve proximity relationships among clients. Algorithm 2 shows the detailed implementation of our FedPoD.

5 Theorems and Proofs

Theorem 1. Consider a on-device weather forecasting system with m clients. Let D1, D2, ..., Dm be the true data distribution and ˆ D1, ˆ D2, ..., ˆ Dm be the empirical data distribution. Denote the head h as the hypothesis from H and d be the VCdimension of H. The total number of samples over all clients is N. Then with probability at least 1 - δ:

Theorem 2 (Transmitting Prompts Ensure Privacy). Consider a device with a frozen pre-trained foundation model parameterized by θf , and trainable prompts parameterized by θp but initialized before updates. Transmitting these prompts can ensure privacy in multi-level communication.

Proof. Detailed proofs can be found at Appendix C.

6 Experiments

Datasets. Three weather multivariate time-series datasets from [Chen et al., 2023b], including AvePRE, SurTEMP, and SurUPS collected by 88, 525, and 238 devices, respectively. Detailed information can be found at Appendix A.

Baselines. We compare with competitive FL methods, including FedAvg [McMahan et al., 2017], FedProx [Li et al., 2020], pFedMe [T Dinh et al., 2020], Per-FedAvg [Fallah et al., 2020], FedATT [Jiang et al., 2020], APFL [Deng et al., 2020], FedAMP [Huang et al., 2021], and SFL [Chen et al., 2022], while keeping the local foundation model consistent. Details about baselines can be found at Appendix A. In addition, we adapt two fine-tuning methods for each baseline for evaluate our method's effectiveness, as below:

• Conventional Fine-tuning: Update local FM with an FFN as the fine-tune head.
• Adaptive Prompts Tuning (Ours): Update prompts with the frozen FM and multi-level communication.
• Other Prompt Tuning: Add parameters to input to updating models [Chen et al., 2023b; Guo et al., 2023].

Implementation. The task of on-device weather forecasting is to predict the next 12 hours using the data from the previous 12 hours. Main experiments are conducted in 25 local epoch within 50 communication round. Following [Chen et al., 2022], Mean Absolute Error (MAE) and Root Mean Squared Error (RMSE) are used as evaluation metrics. All results are in 100× the original value for a clearer comparison. Detailed information about the implementation, local updating process and the aggregation can be found at Appendix B.

6.1 Main Results

Table 2 presents main results, showing that FedPoD outperforms baselines in most scenarios, often by a significant margin, across various tuning strategies. Notably, our adaptive prompt tuning outperforms conventional fine-tuning while using about 74% parameters (see Table 1). This method enhances baseline models by enabling them to learn fewer parameters for considerable performance boosts. FedPoD records an average performance increase of 23.6%/12.9%, 11.7%/19.7%, and 12.6%/4.3% over FedAvg, FedProx, and Per-FedAvg, respectively. These percentages reflect MAE improvements for Task1/Task2. The gains are particularly striking against SFL, which employs graph-based aggregation [Chen et al., 2022]. With adaptive prompt tuning, FedPoD improves by 9.8% and 6.7% on average. These figures rise to 15.9% for Task1 and 11.1% or Task2 with conventional fine-tuning. In addition, FedPoD can show a superior performance relative to FL-based prompt methods, PromptFL [Guo et al., 2023] and MetePFL [Chen et al., 2023b]. We credit these benefits to two main strategies: (1) adaptive prompt tuning guides the PFM to generate more accurate prediction based on lightweight prompts with multi-level communication, and (2) dynamic graph modeling encourages collaborative optimization among clients with similarly distribution to mitigate data heterogeneity. These components effectively address data heterogeneity and homogeneity issues through a lightweight plug-and-play means.

6.2 Framework Analysis

Ablation Study. We present the ablation results from two angles: (1) examining local prompts and their aggregation method, and (2) assessing the local optimization objective. This helps confirm the effectiveness of our Adaptive Prompt Tuning and Dynamic Graph Modeling. For (1), Table 3 re- Hyper-parameter Sensitivity. We examine the impact of hyper-parameters from two angles: the prompt updating step and the subgraph step. Our configuration is as follows: 5 local epochs and 10 communication rounds, while other settings follow main experiments. Table 6 shows that the best performance for Task 2 is achieved with a step of 1, and for Task 1 with a step of 6. This inconsistency is due to the variable nature of weather patterns. Additionally, setting the step to 12 results in the poorest performance for both Task 1 and Task 2. This is because a single-step update does not account for the erratic periodicity of weather patterns, leading to inflexibility. Our findings on the impact of SG are shown in Table 7, where SG ∈ {1, 2, 4, 6, 8, 10}. The results suggests that FedPoD achieve the suboptimal when SG = 1 across different tasks, while optimal results are achieved for Task 1/2 when SG = 10/2. Bigger SG means more knowledge will be involved in local optimization. In our experiment, not all clients train in each round for training due to the considerable overhead. With a large SG, clients are optimized locally with a restricted range of client knowledge, potentially overlooking valuable input from other participants and negatively affecting performance. We set SG = 1 as the default because it considers all clients and allows for flexibility in specific scenarios. As only prompts PT ,PS,PV , which have fewer parameters, are involved, SG = 1 does not significantly increase communication costs.

Table 4: Ablation results about the multi-task optimization objective (MAE/RMSE report). Bold: the best, Underline: the second best.

Table 5: Differential privacy experiment results (MAE/RMSE report), FedPoD-DP denotes FedPoD with DP.

7 Conclusion and FutureWorks

In this paper, we seek to tackle data heterogeneity among devices and data homogeneity within individual devices in on-device weather forecasting. To achieve this, we propose FedPoD, which is built on Adaptive Prompt Tuning and Dynamic Graph Modeling. The former aims to mitigate data homogeneity via extracting latent knowledge with the frozen foundation model, alongside multi-level communication, and the last deals with data heterogeneity by prioritizing devices with similar distribution for aggregation and collaborative training based on prompt-related graphs. Extensive experiments on real-world on-device weather forecasting datasets shows FedPoD consistently outperforms stateof-the-art methods. However, FedPoD can struggle with long-term predictions due to the gradual updating of prompts. We plan to address this limitation in future research and extend our approach to more on-device spatiotemporal reasoning challenges, such as forecasting and imputation.