Personalized Federated Learning for Cross-City Traffic Prediction

1 Introduction

Traffic prediction plays an important role in urban computing. The accuracy of traffic prediction depends on the availability of sufficient data, which is mostly collected by road sensors deployed by authorities or vehicle devices involved in spatial crowdsourcing tasks [Liu et al., 2022; Zhong et al., 2023]. However, due to diverse levels of urban development, some cities only have limited traffic data insufficient for accurate traffic prediction. To cope with this issue of unbalanced data availability, an intuitive solution is to share knowledge from data-rich cities with data-scarce cities, i.e., cross-city traffic knowledge transfer.

∗Corresponding authors.

Extensive research has been carried out on cross-city traffic knowledge transfer. Based on how urban areas are being divided, the research works fall into 1) Convolutional Neural Network (CNN)-based methods and 2) Graph Neural Network (GNN)-based methods. CNN-based methods divide a city into equal-size grid cells to learn region embeddings [Zhang et al., 2022b; Jin et al., 2022a; Wang et al., 2018; Yao et al., 2019a]. However, equal-size grid cells fail to reflect the actual urban structure. Neither can CNNs effectively model non-Euclidean road networks. In contrast, GNN-based methods can model irregular urban structures by capturing the spatial relevance of the regions [Yao et al., 2023; Ouyang et al., 2023b; Ouyang et al., 2023a].

Both CNN- and GNN-based methods may suffer from negative transfer, i.e., the performance of a target city degrades after domain knowledge is transferred from a source city [Zhang et al., 2022a]. Many graph-based cross-city transfer learning approaches have attempted to address this problem. Jin et al. [2022b] reweighted the source regions to transfer similar knowledge to the target city. However, external auxiliary data (e.g., POI, check-in data) is essential to measuring regional similarity. Domain adversarial learning is commonly used in transfer learning to solve the negative transfer problem caused by domain shift. Recent works [Ouyang et al., 2023b; Yao et al., 2023; Tang et al., 2022; Ouyang et al., 2023a] have adopted the domain classifier to align the distribution by distinguishing the source and target domain features. As the training objective of the domain classifier is to identify the source of samples, retraining is necessary when a new city joins. Lu et al. [2022] and Liu et al. [2023] focused on graph few-shot learning based on the Meta-Learning framework, using a reconstructed graph structure to constrain the shift of the original spatial structure. However, the aforementioned centralized methods do not uphold data privacy concerns by either the source or the target cities, as shown in Figure 1(a). To uphold privacy, we adopt a decentralized training strategy, as shown in Figure 1(b). Cities keep data private and transfer the traffic knowledge to a data-scarce target city to improve the performance of prediction.

Federated learning (FL) [Yang et al., 2020] is a decentralized collaborative learning paradigm designed to preserve data privacy. In the FL paradigm, each city works as a client, while an FL server orchestrates the collaborative training of an FL model across participating cities without exposing local data. Chen et al. [2022] proposed a Cross-city Federated Transfer Learning framework, CcFTL, to transfer knowledge from multi-source urban data (e.g., POIs and population density). However, the FedAvg [McMahan et al., 2017] algorithm used in CcFTL requires client data to be independent and identically distributed (IID), and thus encounters convergence issues on non-IID spatio-temporal traffic data. Such data leads to client drift, deteriorating the performance of FL.

Thus motivated, we propose a personalized Federated learning method for Cross-city Traffic Prediction (pFedCTP). We design a Spatio-Temporal Neural Network (ST-Net) to extract spatial structure features, spatio-temporal knowledge, and traffic patterns from raw traffic data. To overcome spatial heterogeneity caused by different city spatial structures, we decouple the components within ST-Net and selectively transfer space-independent traffic patterns to the FL server. On the other hand, each city's traffic patterns vary temporally, causing temporal heterogeneity in cross-city traffic prediction. To deal with temporal heterogeneity across cities, we propose an adaptive layer-wise model interpolation method to customize personalized local client models.

The contributions of this paper are as follows:

• We propose a novel personalized Federated learning method for Cross-city Traffic Prediction (pFedCTP). It learns traffic knowledge from multiple data-rich source cities and transfers the knowledge to a data-scarce target city while preserving inter-city data privacy.
• We design an ST-Net for cross-city traffic prediction under FL. To overcome spatial heterogeneity, we decouple the ST-Net and share space-independent traffic patterns with the server. Meanwhile, we propose an adaptive layer-wise model interpolation method to alleviate the effect of temporal heterogeneity.
• We conduct extensive experiments on four real-world

traffic speed datasets to verify that pFedCTP achieves superior performance over several representative stateof-the-art methods. It is capable of reducing average MAE and RMSE by 1.9% and 0.8% respectively compared to the best-performing baseline.

2 RelatedWork

2.1 Traffic Prediction

Traffic prediction [Tedjopurnomo et al., 2020] is an important area of research with real-world impact. Conventional approaches, like HA and ARIMA [Williams and Hoel, 2003], utilize statistical information from traffic data. Deep neural networks have further boosted this area of research. CNN-based methods use a grid to partition a city for traffic prediction [Zhang et al., 2017; Yao et al., 2018; Yao et al., 2019b], but they fail to account for spatial correlation among road networks. In contrast, graph neural networks perform better in a non-Euclidean space. A popular paradigm of spatio-temporal prediction combines GNN and Recurrent Neural Network (RNN) to capture spatio-temporal features simultaneously [Li et al., 2017a; Zhao et al., 2019; Bai et al., 2020].

2.2 Traffic Knowledge Transfer across Cities

Traffic knowledge transfer across cities aims to help datascarce cities improve traffic prediction by transferring knowledge from data-rich cities. Related work can be divided into three categories: 1) Similarity-based methods, 2) DAbased methods, and 3) ML-based methods. Similaritybased methods calculate the similarity scores between the source and the target regions based on the external auxiliary information (e.g., POI, check-in, weather) to avoid interference from irrelevant knowledge [Wang et al., 2018; Jin et al., 2022a]. Although urban knowledge mining from multi-source data has been widely verified [Li et al., 2023; Zhang et al., 2021; Huang et al., 2023], the data is not always available. DA-based methods refer to those that leverage domain adversarial learning to alleviate domain distribution discrepancies [Yao et al., 2023; Ouyang et al., 2023a; Ouyang et al., 2023b]. The training objective of the domain discriminator makes it essentially difficult to adapt to new cities. ML-based methods adopt the Meta-Learning framework Model-Agnostic Meta-Learning (MAML) [Yao et al., 2019a; Lu et al., 2022] and Reptile [Liu et al., 2023]. Such methods regard the target city as a new task and learn generalizable knowledge from multiple related tasks to adapt to new tasks rapidly. However, the aforementioned centralized methods fail to uphold data privacy concerns.

2.3 Personalized Federated Learning

As a new paradigm for distributed machine learning, FL aims to train a global model without collecting data to the server. However, FedAvg [McMahan et al., 2017] based approaches cannot deal with non-IID client data effectively. Personalized Federated Learning (PFL) methods have emerged to address this limitation. Previous research on PFL can be categorized into two types [Tan et al., 2022]: 1) global model personalization and 2) personalized local models. Global model personalization aims to learn a well-generalized global model

𝜃𝒮: Shared parameters 𝜃𝒫: Private parameters

Figure 2: The framework of pFedCTP. The left side corresponds to a local ST-Net composed of four modules, in which the spatial structure (SS), spatio-temporal feature (ST), and predictor are private, and their parameters are denoted as θP; traffic pattern (TP) is shared with other cities and the parameters are denoted as θS. The right side corresponds to Personalized Federated Learning, where cities act as clients participating in FL to train a well-generalized global model that is transferred to the data-scarce target city for fine-tuning on the local data.

to adapt rapidly to local models. Personalized local models customize model architecture for each client by modifying the FL model aggregation process. Model decoupling [Arivazhagan et al., 2019] and model interpolation [Hanzely and Richt´arik, 2020; Zhang et al., 2023] are two commonly used techniques to achieve local model personalization. Model decoupling keeps part of local parameters private to balance the performance between generalization and personalization. Model interpolation mixes the global and local models to save valuable local model parameter information. In this paper, we attempt to learn a generalized global model that absorbs traffic knowledge from multiple cities and can quickly adapt to a target city. Meanwhile, we also learn a personalized local model that customizes spatio-temporal features from the local traffic data of a city.

3 Preliminary

Traffic Spatio-Temporal Graph. Given a city, we define its traffic spatio-temporal graph as GST = (V, E, A, X). Specifically, V denotes the node set, N = |V| is the number of nodes, and E ⊆ (V×V) denotes the edge set. Moreover, A ∈ RN×N is the adjacency matrix of the traffic spatio-temporal graph, where aij = 1 indicates an edge exists between node i and j; otherwise, aij = 0. Furthermore, X ∈ RN×Ttotal×df is the traffic feature (e.g., traffic flow and traffic speed) matrix, Ttotal is the total time steps of traffic data, df is the node feature dimensionality, and Xt ∈ RN×df represents the traffic data at time t.

Traffic Spatio-Temporal Graph Prediction. Given the historical traffic data of T steps, the goal of traffic spatiotemporal graph prediction is to predict the traffic data at future T ′ steps by learning a function f(·) based on the traffic spatio-temporal graph GST :

Federated Cross-City Traffic Prediction. Given C - 1 source cities with abundant traffic data G S ST = {G 1ST , . . . , G C-1 ST } and a target city with scarce traffic data G T ST , the goal of federated cross-city traffic prediction is to learn traffic knowledge from the source cities without sharing their local data and transfer the traffic knowledge to the target city to improve the traffic prediction performance.

4 Methodology

We propose a novel personalized federated learning framework to transfer the traffic knowledge from the data-rich source cities to the data-scarce target city. As shown in Figure 2, the framework consists of two stages. Stage I: Local Spatio-Temporal Neural Network (ST-Net). We design an ST-Net to learn global shared traffic knowledge and local personalized spatio-temporal features for the cross-city traffic prediction task. Each city is seen as a client and deploys an ST-Net to train on the local traffic data. Stage II: Personalized Federated Learning. We decouple the components within ST-Net and selectively share space-independent traffic patterns to overcome spatial heterogeneity. Meanwhile, to alleviate temporal heterogeneity, we propose adaptive layerwise aggregation (model interpolation) to balance global generalization and local personalization.

4.1 ST-Net

In this section, we introduce the ST-Net designed for crosscity traffic prediction under FL. An ST-Net instance is deployed on each client to extract spatio-temporal features and predict traffic patterns. It has four modules: spatial structure, traffic pattern, spatio-temporal feature, and predictor.

Spatial Structure

City spatial structure is tightly connected to factors such as terrain, history and culture, and population density. Learning city spatial features helps understand spatial topological structure and improves the performance of spatio-temporal prediction tasks. Following a previous work [Tan et al., 2023], we extract the random walk information to generate spatial structure embedding Hrw, using the raw adjacent matrix. Then Hrw is fed into the Graph Convolutional Neural Network (GCN) [Kipf and Welling, 2016] to generate the spatial features Hss.

where σ is an activation function, H(l) ss represents the node features at layer l, H(0) ss = Hrw, ˜ A = ̂D- 12 ̂ ÂD- 12 , ̂ A = A+IN, ̂D is the degreematrix of ̂ A, IN is the identitymatrix, and N is the number of nodes.

Traffic Pattern

The current or future traffic condition is tightly correlated with its previous observations. We use Gated Recurrent Unit (GRU) [Chung et al., 2014] to generate temporal embedding Hgru from the historical traffic data X, for GRU has fewer parameters than RNN but a faster training speed.

Furthermore, we extract representative traffic knowledge, i.e., traffic patterns, from temporal embeddings. We feed the temporal embedding Hgru into a self-attention module to adaptively extract temporal traffic features.

whereWQ,WK,WV are learnable parameters, d′ is the dimension of the query, key, and value matrix. Then we extend the temporal attention mechanism to multi-heads to capture the temporal features Hatt. We further concatenate the temporal embedding Hgru and traffic features Hatt, and feed the concatenation result into Multi-Layer Perceptrons (MLP) to generate the temporal pattern Htp = MLP([Hgru||Hatt]), where || represents the concatenation operation.

Spatio-Temporal Feature

We combine GRU and Graph Attention Network (GAT) [Veliˇckovi´c et al., 2017] to learn spatio-temporal features from the observed historical data. We feed the adjacent matrix A and temporal embedding Hgru into GAT to generate spatio-temporal features Hgat.

whereW ′ is the weight matrix, eij represents the importance of node j to node i, Ni is node vi's neighbor node set, and softmax normalizes node vj's neighbor nodes. Finally, the output features Hgat are obtained by weighting the input features:

To further personalize a local ST-Net, we include a hypernetwork [Ha et al., 2016] to dynamically generate the weights for the network based on the input features. We pass the spatio-temporal features Hgat as input to the hypernetwork to produce personalized model weight W ∈ Rd×d and bias b ∈ Rd . Accordingly, we customize spatio-temporal features using the linear transformation Hst = W × Hgat + b.

Predictor

We design a predictor to predict the traffic situation of the future T ′ steps. We combine the private spatial structure embedding Hss, the shared traffic patterns Htp, and the private spatio-temporal features Hst to obtain the final representation for traffic prediction. Specifically, we use an MLP as the predictor.

where ̂ Y ∈ RN×T ′×df is the prediction result. Then we use Mean Squared Error (MSE) loss to measure the performance of our model in training. Given the ground truth Y ∈ RN×T ′×df , the loss function of ST-Net for traffic prediction is defined as:

4.2 Personalized Federated Learning

In this section, we describe the PFL designed for the crosscity traffic prediction task to handle client drift caused by spatio-temporal data heterogeneity.

Parameter Decoupling

To overcome spatial heterogeneity, we decouple ST-Net and only learn the shared module (traffic pattern), keeping specific spatial-related features private. The ST-Net parameters are divided into the shared model parameters θS and the private model parameters θP. For the four modules in ST-Net (spatial structure, traffic pattern, spatio-temporal feature, and predictor), only traffic pattern is shared during FL, and the rest modules contain city-specific spatial knowledge, which can be seen as noise to the target city and should be kept private. Parameter decoupling not only avoids unrelated traffic knowledge but also reduces communication costs between server and clients.

Adaptive Model Interpolation

To alleviate temporal heterogeneity, we adopt model interpolation to customize the local model parameters at the start of a communication round rather than overwriting them with the global model parameters. The original local model has effective information, and thus dropping themwill cause information loss. A trade-off parameter λ ∈ [0, 1] is generally used to determine the mixing degree of the local shared model and global mean model. λ = 1 means the global model replaces the local model completely; λ = 0 means the opposite. The model interpolation on client i at round r is denoted as:

However, it is hard to denote a constant λ to all clients because they have different contributions. Besides, it is infeasible to consider that each layer in the model has an equal mixing degree. Therefore, we propose an adaptive layer-wise model interpolation method for local ST-Net initialization in each round. For layer lr i ∈ θrSi in client i at round r, the model interpolation process works as follows:

where sim(·) captures the similarity of the global and the local shared layer, lr represents the corresponding layer in global model θrS at round r, and lr - lr-1 i denotes the update of the parameters. In this paper, we use the Cosine similarity function.

pFedCTP Learning Process

Algorithm 1 outlines the learning process of pFedCTP. After the initialization (lines 1-2), it adopts a workflow of FL (lines 3-10) followed by transfer (lines 11-14). In FL, clients receive a global module from the server (line 4) and adaptively aggregate with the local traffic pattern module (lines 5-7). Then clients train ST-Net on samples of their local dataset and upload the updated shared module to the server in parallel (lines 8-9). The server aggregates all modules from clients for the next round of communication (line 10). The FL ends after R rounds of communication, and the server stores a well-generalized global traffic pattern module. Subsequently, the target city receives the global module transferred from the server (line 11) and performs layer-wise traffic pattern module updates (lines 12-13). Finally, the target city fine-tunes the ST-Net with local scarce traffic data to improve traffic prediction performance (line 14).

Notably, all clients (source and target) participate in the first stage but only the target client participates in the second. Thus, the total communication cost between the server and clients is 2 ×C ×R+ 1 times. Moreover, the local training on each client only uses sampled data instead of the entire local dataset lest unbalanced client data deteriorates the performance of the target client. To this end, we introduce a hyperparameter B to denote the batch number in local training. Its effect will be analyzed in Section 5.3. Assuming that a datarich source city gets m samples from local data during ST-Net training, we have m << M 1 , where M is the total number of possible data samples. The computational complexity mainly includes the local model training cost O(m× C × R + p), where p is the number of samples in the target city, and the adaptive model interpolation cost O(L×C ×R+L), where L is the number of layers in the shared module. In contrast to the conventional FL approach that trains a model on the entire dataset and shares the entire model architecture, our model's complexity is significantly reduced.

5 Experiments

5.1 Experiment Settings

Datasets. We evaluate the performance of pFedCTP on four traffic speed datasets: PEMS-BAY, METR-LA [Li et al., 2017b], DiDi-Chengdu, and DiDi-Shenzhen. PEMS-BAY

1 Refer to the dataset Didi-Chengdu shown in Table 1. Suppose that the batch number B = 150 and the batch size = 32. We have m= 150 ×32 = 4, 800, andM ≈ 17, 280, i.e., the time span.

Algorithm 1: The pFedCTP Framework.

Input: Communication rounds R, client number C, ST-Net shared parameters θS, batch number B Output: The target client ST-Net ̂ θR+1 t // Inititialization 1 for client i from 1 to C in parallel do 2 Initialize client model θ0 i ; // Personalized Federated Learning 3 for round r = 1, 2, . . . ,R do 4 Server sends θrS to all clients; 5 for client i from 1 to C in parallel do 6 for layer lr i ∈ θrSi do 7 update lr i according to Eq. (11); 8 ̂ θr i ← Train ST-Net in B batches on samples; 9 update ̂ θrSi to the Server; 10 Server computes θr+1 S = ∑Ci=1 ̂ θrSi C ; // Fine-Tuning on Target City 11 Server sends θR+1 S to the target city; 12 for layer lR+1 t ∈ θR+1 St do 13 update lR+1 t according to Eq. (11); 14 ̂ θR+1 t ←Fine-tune ST-Net on target city's data;

and METR-LA include traffic information from the San Francisco Bay Area and Los Angeles County in the USA, respectively. DiDi-Chengdu and Didi-Shenzhen are provided by the Didi GAIA Initiative [DiDi, 2020]. The detailed information of each dataset is shown in Table 1. We alternately use three datasets as source cities with abundant traffic data and the fourth as the target city with relatively less data. To simulate data scarcity in the target city, we only use 3 days of traffic data as training data.

Metrics. We use two commonly used metrics to evaluate the performance of traffic prediction: Mean Absolute Error (MAE) and Root Mean Squared Error (RMSE).

Baselines. We compare pFedCTP with relevant methods of three types: Target-only methods, Cross-city methods, and Federated Cross-city methods.

Target-only methods only train a model on the datascarce target city, without involving any source cities or their data. (1) Spatio-Temporal Graph Convolution Network (STGCN) [Yu et al., 2017] integrates graph convolution and gated temporal convolution through spatio-temporal convolutional blocks for traffic prediction; (2) T-GCN [Zhao et al., 2019] uses temporal GCN to learn complex topological structures and GRU to obtain temporal correlation from dynamic traffic data; (3) We also train an ST-Net (cf. Section 4.1) on a target city only to validate its effectiveness.

From Cross-city methods, for fairness, we choose those that transfer knowledge from multiple source cities to the target city. (1) Domain Adversarial Spatial-Temporal Network (DastNet) [Tang et al., 2022] adopts domain adversarial learning to learn the domain-invariant node embedding; (2) Spatio-temporal few-shot traffic prediction model (STGFSL) [Lu et al., 2022] adopts MAML [Finn et al., 2017] as the base Meta-Learning framework.

We also compare with several FL-based methods. (1) FedAvg [McMahan et al., 2017]: We combine both source cities and the target city to train a global ST-Net, and we test the final global model on the target city data; (2) FedProx [Li et al., 2020] uses a preliminary term to guide local models to pay more attention to model weights that are close to the global model weights during updates; (3) Per-FedAvg [Fallah et al., 2020] is a combination of FL and Meta-Learning methods to optimize the global model for fast personalization. The source cities are trained in the meta-training phase and the target city is tested in the meta-testing phase.

Implement Details. For all the experiment methods, we set the history time steps T = 12 and the future time steps T ′ = 6. We simulate short-, medium-, and long-term prediction results using 1, 3, and 6 steps respectively. Other important hyperparameters are set as follows: the client number C = 4, the batch size = 32, the learning rate = 0.01, the number of GCN layers = 1, and the hidden dimensions = 32 for all methods. For the baselines, we train the target-only model for 100 epochs. The code is available at https://github.com/ZYuSdu/pFedCTP.

5.2 Results and Discussion

We show the traffic prediction results of all methods on the four datasets in Table 2. We have the following observations: (1) Target-only baselines mostly perform not well because they are trained on limited target city traffic data. Compared to T-GCN, STGCN struggles to predict short- and mid-term traffic using scarce data but has a significant advantage in long-term prediction in the Didi-Chengdu and Didi-Shenzhen datasets. (2) For Cross-city baselines, the state-of-the-art method ST-GFSL performs better on some long-term traffic prediction indicators, while DastNet is unable to handle cross-city traffic prediction well for scarce data in the target city. (3) In respect of FL-based baselines, Per-FedAvg outperforms the others. As a Meta-Learning based method, PerFedAvg samples tasks to learn a generalized global model, which can quickly adapt to the new target city. (4) Compared to the other target-only baselines, our ST-Net shows clear advantages in the prediction task on scarce traffic data. Overall, pFedCTP achieves the best results among all, reducing the average MAE and RMSE by 1.9% and 0.8% respectively compared to the best-performing baseline, i.e., ST-GFSL.

5.3 Hyperparameter Analysis

We investigate the effect of the communication rounds and the local batches in each round on Didi-Chengdu. The results are shown in Figure 3.

Communication Rounds. We tune R according to the set of {30, 60, 90, 120}. We can see that as the number of communication rounds increases, the performance of the model is improved. However, excessive communication rounds not only increase costs but also increase the over-fitting risk of the model. Our pFedCTP achieves good results when R = 90.

Local Batches. We vary the local batch number B according to the set of {50, 100, 150, 200} to find the suitable setting. Small B values require massive communication between the server and clients, causing a waste of resources, whereas large B values interfere with the prediction performance in long-term traffic prediction. When B = 150, pFedCTP makes better short- and long-term predictions.

5.4 Ablation Study

We also conduct ablation studies to verify the effectiveness of our framework designs.

Ablation Study of ST-Net. We compare the ST-Net (ST) with several variants: ST-w/o TP (temporal pattern), ST-w/o SS (spatial structure), and ST-w/o HN (hypernetwork). The ablation results are shown in Figure 4. We observe that the prediction performance deteriorates severely after we remove the traffic pattern module, meaning that the module is necessary for ST-Net and FL. Without the spatial structure module, the performance deteriorates clearly on the two Didi datasets, as the rich information in their complex spatial structure features is not used. Moreover, we observe that the hypernetwork might slightly enhance the prediction performance in some cases by customizing the network parameters.

Ablation Study of pFedCTP. The ablation study on pFedCTP includes pFedCTP-w/o FT (Fine-tuning on the target city), pFedCTP-All (sharing the entire ST-Net), and pFedCTP-T that only involves source cities in FL and trains each ST-Net on the entire local traffic dataset. The results are shown in Figure 5. We can see that pFedCTP-w/o FT is slightly worse than pFedCTP, verifying that the model transferred from the server can adapt to the target city through fine-tuning. Next, pFedCTP-All can result in higher errors since sharing the entire ST-Net with the server can also transfer irrelevant noise from other clients. Further, pFedCTP-T achieves comparable performance with pFedCTP. However, pFedCTP-T uses the entire dataset for local training and thus incurs much longer training time than pFedCTP that uses sampled data. For instance, on the Didi-Chengdu dataset, pFedCTP-T needs 72 minutes but pFedCTP only 9 minutes.

6 Conclusion

In this paper, we have proposed a personalized federated learning framework for cross-city traffic prediction, pFedCTP, that aims to help a data-scarce city improve traffic prediction performance without violating data privacy. Following the FL paradigm, pFedCTP treats cities as clients each of which trains a Spatio-Temporal Neural Network (STNet) on its local traffic data. To overcome spatial heterogeneity, pFedCTP decouples the components within ST-Net and shares space-independent traffic knowledge. Meanwhile, pFedCTP uses a layer-wise model interpolation method to deal with temporal heterogeneity. Experiments on four real traffic datasets demonstrate the effectiveness of pFedCTP.

Acknowledgments

This research is partially supported by the National Key R&D Program of China 2021YFF0900800 and NSFC No. 92367202.