Cross-Node Federated Graph Neural Network for Spatio-Temporal Data Modeling

1 INTRODUCTION . Modeling the dynamics of spatio-temporal data generated from networks of edge devices or nodes ( e.g . sensors , wearable devices and the Internet of Things ( IoT ) devices ) is critical for various applications including traffic flow prediction ( Li et al. , 2018 ; Yu et al. , 2018 ) , forecasting ( Seo et al. , 2019 ; Azencot et al. , 2020 ) , and user activity detection ( Yan et al. , 2018 ; Liu et al. , 2020 ) . While existing works on spatio-temporal dynamics modeling ( Battaglia et al. , 2016 ; Kipf et al. , 2018 ; Battaglia et al. , 2018 ) assume that the model is trained with centralized data gathered from all devices , the volume of data generated at these edge devices precludes the use of such centralized data processing , and calls for decentralized processing where computations on the edge can lead to significant gains in improving the latency . In addition , in case of spatio-temporal forecasting , the edge devices need to leverage the complex inter-dependencies to improve the prediction performance . Moreover , with increasing concerns about data privacy and its access restrictions due to existing licensing agreements , it is critical for spatio-temporal modeling to utilize decentralized data , yet leveraging the underlying relationships for improved performance . Although recent works in federated learning ( FL ) ( Kairouz et al. , 2019 ) provides a solution for training a model with decentralized data on multiple devices , these works either do not consider the inherent spatio-temporal dependencies ( McMahan et al. , 2017 ; Li et al. , 2020b ; Karimireddy et al. , 2020 ) or only model it implicitly by imposing the graph structure in the regularization on model weights ( Smith et al. , 2017 ) , the latter of which suffers from the limitation of regularization based methods due to the assumption that graphs only encode similarity of nodes ( Kipf & Welling , 2017 ) , and can not operate in settings where only a fraction of devices are observed during training ( inductive learning setting ) . As a result , there is a need for an architecture for spatio-temporal data modeling which enables reliable computation on the edge , while maintaining the data decentralized . To this end , leveraging recent works on federated learning ( Kairouz et al. , 2019 ) , we introduce the cross-node federated learning requirement to ensure that data generated locally at a node remains decentralized . Specifically , our architecture – Cross-Node Federated Graph Neural Network ( CNFGNN ) , aims to effectively model the complex spatio-temporal dependencies under the cross-node federated learning constraint . For this , CNFGNN decomposes the modeling of temporal and spatial dependencies using an encoder-decoder model on each device to extract the temporal features with local data , and a Graph Neural Network ( GNN ) based model on the server to capture spatial dependencies among devices . As compared to existing federated learning techniques that rely on regularization to incorporate spatial relationships , CNFGNN leverages an explicit graph structure using a graph neural networkbased ( GNNs ) architecture , which leads to performance gains . However , the federated learning ( data sharing ) constraint means that the GNN can not be trained in a centralized manner , since each node can only access the data stored on itself . To address this , CNFGNN employs Split Learning ( Singh et al. , 2019 ) to train the spatial and temporal modules . Further , to alleviate the associated high communication cost incurred by Split Learning , we propose an alternating optimization-based training procedure of these modules , which incurs only half the communication overhead as compared to a comparable Split Learning architecture . Here , we also use Federated Averaging ( FedAvg ) ( McMahan et al. , 2017 ) to train a shared temporal feature extractor for all nodes , which leads to improved empirical performance . Our main contributions are as follows : 1 . We propose Cross-Node Federated Graph Neural Network ( CNFGNN ) , a GNN-based federated learning architecture that captures complex spatio-temporal relationships among multiple nodes while ensuring that the data generated locally remains decentralized at no extra computation cost at the edge devices . 2 . Our modeling and training procedure enables GNN-based architectures to be used in federated learning settings . We achieve this by disentangling the modeling of local temporal dynamics on edge devices and spatial dynamics on the central server , and leverage an alternating optimization-based procedure for updating the spatial and temporal modules using Split Learning and Federated Averaging to enable effective GNN-based federated learning . 3 . We demonstrate that CNFGNN achieves the best prediction performance ( both in transductive and inductive settings ) at no extra computation cost on edge devices with modest communication cost , as compared to the related techniques on a traffic flow prediction task . 2 RELATED WORK . Our method derives elements from graph neural networks , federated learning and privacy-preserving graph learning , we now discuss related works in these areas in relation to our work . Graph Neural Networks ( GNNs ) . GNNs have shown their superior performance on various learning tasks with graph-structured data , including graph embedding ( Hamilton et al. , 2017 ) , node classification ( Kipf & Welling , 2017 ) , spatio-temporal data modeling ( Yan et al. , 2018 ; Li et al. , 2018 ; Yu et al. , 2018 ) and multi-agent trajectory prediction ( Battaglia et al. , 2016 ; Kipf et al. , 2018 ; Li et al. , 2020a ) . Recent GNN models ( Hamilton et al. , 2017 ; Ying et al. , 2018 ; You et al. , 2019 ; Huang et al. , 2018 ) also have sampling strategies and are able to scale on large graphs . While GNNs enjoy the benefit from strong inductive bias ( Battaglia et al. , 2018 ; Xu et al. , 2019 ) , most works require centralized data during the training and the inference processes . Federated Learning ( FL ) . Federated learning is a machine learning setting where multiple clients train a model in collaboration with decentralized training data ( Kairouz et al. , 2019 ) . It requires that the raw data of each client is stored locally without any exchange or transfer . However , the decentralized training data comes at the cost of less utilization due to the heterogeneous distributions of data on clients and the lack of information exchange among clients . Various optimization algorithms have been developed for federated learning on non-IID and unbalanced data ( McMahan et al. , 2017 ; Li et al. , 2020b ; Karimireddy et al. , 2020 ) . Smith et al . ( 2017 ) propose a multi-task learning framework that captures relationships amongst data . While the above works mitigate the caveat of missing neighbors ’ information to some extent , they are not as effective as GNN models and still suffer from the absence of feature exchange and aggregation . Alternating Optimization . Alternating optimization is a popular choice in non-convex optimization ( Agarwal et al. , 2014 ; Arora et al. , 2014 ; 2015 ; Jain & Kar , 2017 ) . In the context of Federated Learning , Liang et al . ( 2020 ) uses alternating optimization for learning a simple global model and reduces the number of communicated parameters , and He et al . ( 2020 ) uses alternating optimization for knowledge distillation from server models to edge models . In our work , we utilize alternating optimization to effectively train on-device modules and the server module jointly , which captures temporal and spatial relationships respectively . Privacy-Preserving Graph Learning . Suzumura et al . ( 2019 ) and Mei et al . ( 2019 ) use statistics of graph structures instead of node information exchange and aggregation to avoid the leakage of node information . Recent works have also incorporated graph learning models with privacypreserving techniques such as Differential Privacy ( DP ) , Secure Multi-Party Computation ( MPC ) and Homomorphic Encryption ( HE ) . Zhou et al . ( 2020 ) utilize MPC and HE when learning a GNN model for node classification with vertically split data to preserve silo-level privacy instead of nodelevel privacy . Sajadmanesh & Gatica-Perez ( 2020 ) preprocesses the input raw data with DP before feeding it into a GNN model . Composing privacy-preserving techniques for graph learning can help build federated learning systems following the privacy-in-depth principle , wherein the privacy properties degrade as gracefully as possible if one technique fails ( Kairouz et al. , 2019 ) . 3 CROSS-NODE FEDERATED GRAPH NEURAL NETWORK . 3.1 PROBLEM FORMULATION . Given a dataset with a graph G = ( V , E ) , a feature tensor X ∈ R|V|× ... and a label tensor Y ∈ R|V|× ... , we consider learning a model under the cross-node federated learning constraint : node feature xi = Xi , ... , node label yi = Yi , ... , and model output ŷi are only visible to the node i . One typical task that requires the cross-node federated learning constraint is the prediction of spatiotemporal data generated by a network of sensors . In such a scenario , V is the set of sensors and E describes relations among sensors ( e.g . eij ∈ E if and only if the distance between vi and vj is below some threshold ) . The feature tensor xi ∈ Rm×D represents the i-th sensor ’ s records in the D-dim space during the past m time steps , and the label yi ∈ Rn×D represents the i-th sensor ’ s records in the future n time steps . Since records collected on different sensors owned by different users/organizations may not be allowed to be shared due to the need for edge computation or licensing issues on data access , it is necessary to design an algorithm modeling the spatio-temporal relation without any direct exchange of node-level data . 3.2 PROPOSED METHOD . We now introduce our proposed Cross-Node Federated Graph Neural Network ( CNFGNN ) model . Here , we begin by disentangling the modeling of node-level temporal dynamics and server-level spatial dynamics as follows : ( i ) ( Figure 1c ) on each node , an encoder-decoder model extracts temporal features from data on the node and makes predictions ; ( ii ) ( Figure 1b ) on the central server , a Graph Network ( GN ) ( Battaglia et al. , 2018 ) propagates extracted node temporal features and outputs node embeddings , which incorporate the relationship information amongst nodes . ( i ) has access to the not shareable node data and is executed on each node locally . ( ii ) only involves the upload and download of smashed features and gradients instead of the raw data on nodes . This decomposition enables the exchange and aggregation of node information under the cross-node federated learning constraint . 3.2.1 MODELING OF NODE-LEVEL TEMPORAL DYNAMICS . We modify the Gated Recurrent Unit ( GRU ) based encoder-decoder architecture in ( Cho et al. , 2014 ) for the modeling of node-level temporal dynamics on each node . Given an input sequence xi ∈ Rm×D on the i-th node , an encoder sequentially reads the whole sequence and outputs the hidden state hc , i as the summary of the input sequence according to Equation 1. hc , i = Encoderi ( xi , h ( 0 ) c , i ) , ( 1 ) where h ( 0 ) c , i is a zero-valued initial hidden state vector . To incorporate the spatial dynamics into the prediction model of each node , we concatenate hc , i with the node embedding hG , c , i generated from the procedure described in 3.2.2 , which contains spatial information , as the initial state vector of the decoder . The decoder generates the prediction ŷi in an auto-regressive way starting from the last frame of the input sequence xi , m with the concatenated hidden state vector . ŷi = Decoderi ( xi , m , [ hc , i ; hG , c , i ] ) . ( 2 ) We choose the mean squared error ( MSE ) between the prediction and the ground truth values as the loss function , which is evaluated on each node locally .

Summary: Graph neural networks and federated learning are both promising directions of works individually. This papers is apparently one of the first few attempts to combine them for spatio-temporal data modeling. The time series data in the local nodes is modelled by an Encoder-Decoder architecture and spatial locality property of various nodes is captured by the server. The Encoder at each node projects the time series data into an embedding space. This embedding is used by the GNN at the server as node features. The server side GNN outputs node embeddings. The Encoder embeddings and the GNN embeddings are then concatenated and fed to the decoder that predicts the outputs for the subsequent time steps. To ensure that all the nodes encode their temporal data in a common space, the encoders are shared by the clients. Overall, the results look promising.