Online Spatio-Temporal Correlation-Based Federated Learning for Traffic Flow Forecasting

A great variety of methods have been proposed to increase the accuracy of traffic flow forecasting. Overall, existing methods can be divided into two categories: parametric and non-parametric methods. AutoRegressive Integrated Moving Average (ARIMA) and its variants are classical parametric methods and have been researched extensively [17, 18]. However, this series of methods have limitations in regressing traffic flows with high fluctuation, because they are based on the assumption that traffic flows have stationary distribution, which is not the case in reality. Non-parametric methods are based on machine learning models, e.g., Graph Network (GN) [19] and RNN [20]. GN-based methods are devoted to evaluating spatial correlation among multiple nodes based on the constructed topology network, while RNN can capture temporal dependence inside traffic sequences. Therefore, many researches combine GN approaches and RNN to measure spatio-temporal correlation and have achieved satisfactory performance [21, 22].

In [21], Guo et al. proposed an Attention based Spatio-Temporal Graph Convolutional Network (ASTGCN), which contains a spatio-temporal attention mechanism for dynamically capturing the spatial patterns and temporal features. In [22], Zhu et al. proposed a traffic prediction method based on Graph Convolutional Network (GCN) and RNN, where the topology graph of road network is utilized for depicting the spatial correlation, aiming to increase forecasting accuracy.

However, these researches all adopt the centralized training strategy, where the central server updates prediction models directly based on traffic data transmitted from decentralized traffic nodes, thus resulting in huge communication overhead and privacy leakage.

Some recent researches pay attention to forecasting traffic flow using FL [23, 14, 15, 16]. In [23], Qi et al. treated the vehicles as clients and introduced the blockchain to avoid the single point of failure. This scenario is not similar to ours. In [14, 15], Liu et al. and Zhang et al. investigated a similar scenario where traffic sensors mounted along roads belong to some certain organizations. The traffic data detected by sensors should be directly transmitted to the organizational servers for training, also confronting with excessive communicating burden. In order to preserve privacy, multiple organizational servers adopt FL manner to collaboratively train the prediction model. However, the two proposed methods are far from reasonable, because only the spatial correlation among intra-organization nodes is evaluated and that of inter-organization ones is ignored, the lack of which leads to low forecasting accuracy.

Different from the above researches, Meng et al. proposed a model named cross-node federated graph neural network (CNFGNN) [16]. In this work, each traffic sensor serves as a client, and GNN is utilized to capture the spatial relation at the central server. However, in the backpropagation process of GNN, the central server has to communicate with clients for updating parameters, which incurs considerable communication overhead. Besides, since the transportation graph composed by clients at each round remains unchanged, the captured spatial correlation among multiple clients barely changes, thus failing to cater for traffic fluctuation.

Given the dynamic changes of traffic flow, our work mainly focuses on designing an online method for dynamically capturing spatio-temporal correlation among traffic nodes in FL framework, which is quite different from the above researches.

Suppose there is a central server and $N$ traffic nodes. All traffic nodes compose a transportation network, denoted as a directed graph $\mathcal{G}=(\mathcal{S},\mathcal{E})$. $\mathcal{S}=\{s_{1},s_{2},\cdots,s_{N}\}$ represents the node set and $s_{n}$ denotes the $n$-th traffic node. $\mathcal{E}=\left({{e_{m,n}}}\right)\in{\mathbb{R}^{N\times N}}$ denotes the adjacency matrix of these nodes. If $s_{m}$ is adjacent to $s_{n}$, $e_{m,n}=1$ (otherwise $e_{m,n}=0$). Let $\mathcal{A}_{n}$ represent the set of traffic nodes adjacent to $s_{n}$, and we define that $s_{n}\in\mathcal{A}_{n}$. Let $x_{t,n}$ denote the traffic speed observed by $s_{n}$ at the $t$-th time stamp. Primary descriptions are listed in Table I.

In this scenario, traffic nodes are regarded as clients to perform traffic flow forecasting using global model parameters from the central server and then incrementally update local models based on newly-observed speed data. The central server aggregates local models uploaded by clients to yield the fresh global model.

At the $t$-th round, $s_{n}$ receives the global model parameters $w_{t}$ from the central server and forecasts traffic speeds of future $F$ time stamps based on traffic speeds of $T$ historical time stamps, which can be formulated as

where $X_{t,n}^{T}=\left(x_{t-T+1,n},x_{t-T+2,n},\dots,x_{t,n}\right)$ denotes the sequence of traffic speeds from $t-T+1$ to $t$ observed by $s_{n}$, and $\hat{X}_{t,n}^{F}=\left(\hat{x}_{t+1,n},\hat{x}_{t+2,n},\dots,\hat{x}_{t+F,n}\right)$ represents the sequence of predicted traffic speeds from ${t+1}$ to ${t+F}$. $f(\cdot)$ denotes the chosen prediction model.

Let $X_{t,n}^{F}=\left(x_{t+1,n},x_{t+2,n},\dots,x_{t+F,n}\right)$ denote the true traffic speeds of the future time stamps $s_{n}$ observes. Then $s_{n}$ computes the prediction error as follows:

where $l$ denotes the loss function, and is adopted to evaluate the bias between the predicted and true values.

Let $\mathcal{R}$ represent the maximum number of global round. We denote $w^{*}$ as the optimal model parameters which can be obtained in hindsight after the server seeing the data sequences of all $N$ clients over $\mathcal{R}$ rounds [24]. The total prediction regret of $N$ clients can be denoted by $REG$, which evaluates the difference in prediction loss using the actual model parameters and $w^{*}$ [25]. $REG$ is formulated as

The objective of traffic flow forecasting adopting OL manner in FL is to minimize the global regrets over all $N$ clients, i.e., $\min REG.$ Since $w^{*}$ remains constant, we have to iteratively optimize the optimal model parameters $w_{t}(1\leq t\leq\mathcal{T})$, so as to generate the minimum regret.

We integrate the OL manner into the widespread used Federated Averaging (FedAvg) method [26], which is elaborated in Algorithm 1. Specifically, at the $t$-th round, each client concurrently optimizes its locally-owned prediction model via Online Gradient Descent (OGD) (Line 9-10) for $E$ epochs using the newly arriving traffic data. This update manner differs from traditional offline learning mechanism where the prediction model is pre-trained based on historical traffic data obtained in advance. After accomplishing the local training procedure, clients upload their updated local models to the central server. Then, the central server performs model aggregation and yields the fresh global model based on averaging mechanism (Line 5).

An efficient way of assessing spatio-temporal correlation is fundamental to satisfactory prediction performance. Therefore, in this subsection, we elaborates how to evaluate spatial and temporal correlation among clients. Firstly, clients capture the hidden temporal patterns inside speed sequences respectively with GRU-based encoders. Then the central server measures the spatial correlation among clients using GAT and returns the updated hidden temporal patterns to clients. Finally, clients input the updated hidden patterns into the decoders to generate the predicted results.

In this subsection, we illustrate the process of prediction model update, which includes incrementally optimizing local models at clients and the period-aware model aggregation mechanism.

Algorithm 2 elaborates the process of FedOSTC, where $w_{t,n,e}^{(e)}$ and $w_{t,n,d}^{(e)}$ denote the parameters of encoder and decoder on $s_{n}$ at the $e$-th epoch in the $t$-th round respectively. At each global round, the server transmits current global model to clients. Encoders at clients input traffic sequences and output the temporal patterns of these traffic sequences (Line 23-25). The server evaluates the spatial correlation among clients using GAT, and then updates hidden states from clients (Line 16). When clients receive the updated hidden states from the server, they make prediction, calculate prediction loss and optimize parameters of encoders and decoders via OGD, which is iterated for $E$ epochs (Line 17-20). Then clients transmit their updated model parameters to the central server for aggregation using the period-aware weighted aggregation mechanism (Line 5-9). The process continues until the maximum round $\mathcal{R}$ arrives.

In this subsection, the convergence analysis of FedOSTC is demonstrated theoretically. Based on the theory of Online Convex Optimization (OCO) [24], we obtain the upper bound of $REG$.

For simplicity, we reformulate ${l}(X_{t,n}^{F},\hat{X}_{t,n}^{F};{w_{t}})$ as $l(w_{t})$, since the loss and prediction results are dependent on the adopted model parameters. Firstly, some assumptions are introduced.

To Facilitate the analysis on the convergence rate of FedOSTC, the period-aware weighted aggregation mechanism described in (15) is transformed.

Therefore, the dynamic weighting aggregation mechanism can be considered as a relational expression of $w_{t}$ and $w_{t+1}$. The proof of Lemma 4.1 is elaborated as follows.

Theorem 4.2 indicates that the minimum regret of FedOSTC has an upper bound. Furthermore, $REG_{\min}=o(\mathcal{R})$ indicates that FedOSTC performs as well as the the optimal parameters $w^{*}$. Therefore, FedOSTC is efficient in tackling with the issue of traffic flow forecasting. The proof of Theorem 4.2 is detailed.

In the scenario of traffic flow forecasting, the observed traffic speeds from clients have different distribution. While in FL framework, all clients train a same prediction model cooperatively. We hope the global prediction model from the central server has good generalization ability so that all clients can achieve satisfactory performance. Therefore, in this subsection, we compare model generalization ability of the four methods based on FL framework. The prediction RMSEs of clients on two datasets with different forecasting steps are illustrated in Fig.4. By comparing FedAvgOff and CNFGNN with FedAvgOn and FedOSTC respectively, we can observe that RMSE variance of OL manner is much smaller than that of batch learning manner, which indicates that OL manner can yield global model with better generalization ability in FL. Furthermore, the prediction RMSE of FedOSTC is lower than that of FedAvgOn for almost all clients. The RMSE variance over all clients in FedOSTC is obviously lower than that in FedAvgOn. Therefore, by adopting the proposed spatio-temporal correlation evaluating mechanism and period-aware weighted aggregation mechanism, FedOSTC can increase the generalization ability of global model.

In the scenario of traffic flow forecasting, clients make prediction and perform incremental optimization based on only a speed sequence. It is quite different from those methods adopting batch learning manner, where the local batch size is set to 128. In this subsection, we explore the effect of local epoch on prediction performance and further consider whether 5 epochs are redundant for incremental optimization. The prediction RMSEs and MAEs on two datasets with different settings of $E$ are illustrated in Fig. 6.

We observe that different settings of $E$ yield various effects on prediction errors. On the two datasets, both prediction RMSEs and MAEs get lower with more local epochs, which indicates that the internal patterns of speed sequences can be obtained better with more local optimizations. On the other hand, larger local epoch means clients have to spare more computation resources for local optimization, which poses much burden on the resource-limited clients. Therefore, in the future research, we will explore the balance between local epoch and prediction performance.