Federated Inference through Aligning Local Representations and Learning a Consensus Graph

1 INTRODUCTION . Machine learning models become increasingly data hungry as the promise of deep learning continues to realize . More and more applications grow in scale thanks to the availability of distributed data across devices and organizations . Federated learning ( McMahan et al. , 2017 ; Yang et al. , 2019 ; Li et al. , 2019 ; Kairouz et al. , 2019 ) emerges as a formalism that allows data owners to collaboratively train a common model by using one ’ s own data without sharing . Such a formalism is poised to resolve the challenges of expensive data communication and the risk of privacy violation , in light of policies such as the General Data Protection Regulation ( Albrecht , 2016 ) . An issue less addressed by federated learning is the inference process . In fact , inference therein is trivial : once the common model is learned , each data owner retains a copy and applies it on local data , independently of other owners . However , such a scenario is not the only one how data are distributed in practice . In this work , we consider the following scenario : a datum has multiple parts , each of which belongs to a separate owner . Then , the inference must be collectively performed by all participating owners , since none of them alone possesses the entire information . Vertical federated learning studies such a scenario ( Hardy et al. , 2017 ; Hu et al. , 2019 ; Chen et al. , 2020 ) . This concept is figuratively named through cutting the data matrix vertically along the feature axis , rather than the data axis . From the sporadic literature addressing in this scenario , methods generally introduce model parameters distributed with data parts and optionally global parameters that reside in a central server . All parameters are learned jointly , causing however a practical drawback— expensive coordination ( even synchronization ) is required among data owners and the central server ( if present ) . Consider a live example—the national electricity grid , over which thousands of phasor measurement units ( PMUs ) have been deployed to monitor the grid condition ( Smartgrid.gov ) . PMU measurements , as time series data , are owned by several parties . These data may be used to train machine learning models that identify grid events ( e.g. , fault , oscillation , and generator trip ) . Such an event detection system relies on collective series measurements at the same time window but distributed across different data owners . To minimize coordination among owners and maximize autonomy , it is more desirable if each maintains a model of their own and does not participate joint training . In this work , we propose a local–global model framework that maintains data owner autonomy while staying effective in global inference . Therein , each data owner trains a local model with its data part ; the training is independent and incurs no coordination . In the deep learning terminology , the local models produce data representations for input data . Then , a central server takes these representations as input and trains a model for global inference . We term the scenario where inference is collectively performed by data owners but training incurs little coordination among them , federated inference , to distinguish federated learning and vertical federated learning . The local–global framework we propose for this scenario , however , bears two technical challenges . One is the ambiguity of local data representations , because feature dimensions can be arbitrarily permuted without changing the local model . Another challenge is how the global model leverages innate interactions of local data missed by independent local models . We resolve the first challenge through aligning the feature dimensions across all local representations . We resolve the second challenge through employing graph neural networks as the global inference model , where the graph corresponds to the explicit or implicit relational structure of the data owners . When such a graph is not present , we treat the combinatorial graph structure as a random variable of the Bernoulli distribution and optimize the distribution parameters as well . We summarize the contributions of this work as follows . 1 . We formalized federated inference , a less addressed scenario of machine learning with distributed data , where inference is conducted jointly by data owners without data sharing and coordinated training ( Section 2 ) . 2 . We propose an inference framework that consists of autonomous local models and a central model that digests local data representations and produces a global output ( Section 3 ) . 3 . We address the ambiguity challenge of this framework through aligning local representations ( Section 4 ) and address the missing of local model interactions through employing a graph neural network in the central global model . We further propose to simultaneously learn the graph structure if not present ( Section 5 ) . 4 . We study approaches to latent alignment and approaches to graph structure learning and develop theoretical insights into these approaches ( Theorems 1–3 ) . As a byproduct , a more efficient Bernoulli sampling method icdf is proposed to sample graphs for structure learning . 5 . We demonstrate experiments with four real-life data sets including power grids and traffic networks and show the effectiveness of the proposed framework ( Section 6 ) . 2 PROBLEM SETTING . Federated inference refers to a machine learning scenario where both training and inference are conducted on distributed data collectively by data owners . Each owner enjoys data and model autonomy but is subject to centralized coordination that produces a global prediction . This scenario stands in contrast to federated learning , whose inference process is local and separate among owners . Formally , we use a superscript i to index data owners . Let x be a datum with label y and let xi be the part of datum the ith owner possesses ; that is , x = ( x1 , x2 , . . . , xn ) with n owners . The problem is to learn a model y = f ( x ) = f ( x1 , x2 , . . . , xn ) ( 1 ) collectively with all data parts and to perform inference jointly by all owners . Additionally , owners share neither data nor models with each other for , e.g. , privacy reasons . Moreover , owners do not participate joint training , which often incurs expensive coordination . Similarity to federated learning . Federated inference shares the defining data characteristics of federated learning , first coined in McMahan et al . ( 2017 ) : distributed , non-IID , and unbalanced . Data are distributed among owners but not shared . In fact , data may even be heterogeneous . For example , time series measures of the power grid may have different attribute dimensions and may be under different sampling frequencies . As a result , data size may vary significantly among owners . Dissimilarity to federated learning . The root of the differences between these two concepts is the constituent of one datum ( data point ) . In federated learning , a data point is the basic unit of data and thus all owners learn a common model but use it separately . On the other hand , federated inference is concerned with data split in parts across owners . All parts of a data point contribute to the inference collectively . Similarity and dissimilarity to vertical federated learning . Both concepts are concerned with the split of a datum across owners . However , approaches taken for vertical federated learning differ substantially from ours because of joint training among owners . From the less prolific literature , two lines of work are noted . One takes the data matrix literally , by assuming tabular data and studying linear models , where model parameters have natural correspondence to the data parts ( Hardy et al. , 2017 ; Nock et al. , 2018 ; Heinze et al. , 2014 ; 2016 ) . Often , these approaches are hard to generalize to complex data and/or many owners . Another line of work uses a local–global model framework similarly as we do but jointly trains these two parts , incurring expensive communication and creating dependence of local models ( Hu et al. , 2019 ; Chen et al. , 2020 ) .1 In contrast , we allow data owners to train their models independently , maintaining local model autonomy . Data example . Let us consider the power grid . Figure 1 pictorially illustrates PMU measurements distributed across data owners . A panel of time series corresponds to a specific time window and the series collectively represent one data point , which the event detection system classifies . In this simplified illustration , each data owner possesses one series recorded by one PMU ; but in practice they may own different amounts of PMUs ( and thus series ) . Moreover , the series may differ in length because of varying sampling frequencies ; and the series are multivariate with possibly different number of variates . All these variations contribute to data heterogeneity , which necessitates the construction of separate local models . Note that if an event does not cascade over the entire grid , some local models may report event whereas others report normal , resulting in conflicting opinions . A consensus ( global ) model is responsible to resolve the conflict . Additionally , missing data may occur . 3 FEDERATED INFERENCE FRAMEWORK . As such , the proposed framework for federated inference consists of local models f i and a global model g , such that their composition is the sought f denoted in ( 1 ) . Each data owner i possesses a local model trained with its data , independently of other owners . This way , no data sharing is invoked and privacy is of minimal concern . However , because the local models lack a global vision and may be conflicting , a central ( global ) model is key to coordinating the local opinions for final prediction . To maintain autonomy , local models are frozen once pretrained and will not join the training of the global model . Data owners send local data representations to a centralized server for global model training ( and inference ) . In other words , the global model queries neither the raw data nor the local models from data owners . As long as owners agree to send the less decipherable representations to the central server , global inference can be made . Local models . We treat a neural network except the final output layer as a feature extractor , which produces the representation hi of an input xi ; and treat for simplicity the output layer as a logistic regression . That is , a local model f i reads : f i ( xi ) = softmax ( W ihi + bi ) where hi = embedding ( xi ) . ( 2 ) 1Note that the approach proposed by Hu et al . ( 2019 ) assumes no parameters for the global model . Were global parameters present , gradient communication is inevitable . We interchangeably use “ representation ” , “ embedding ” , and “ latent vector ” to mean hi . These hi ’ s are assumed to have the same shape across i , although xi can have different shapes and the embedding function can have different architectures to cope with data heterogeneity . A simple example of the embedding function is a fully connected layer hi = ReLU ( U ixi + ci ) ; but an arbitrarily complex function is applicable . Global model . The global model is a function g of all local representations : ŷ = g ( h1 , h2 , . . . , hn ) . ( 3 ) An example of g is a fully connected layer , followed by mean pooling and another fully connected layer : ŷ = softmax ( W1 · 1n ∑n i=1 ReLU ( W0h i + b0 ) + b1 ) . ( 4 ) Challenges . Two considerations are pertinent to this framework . First , when the latent dimensions have semantic meaning ( e.g. , when the local models are trained to yield disentangled representations ( Higgins et al. , 2018 ) ) , each latent feature of the local representations may not match , because an arbitrary permutation of the latent dimensions does not change a local model . Second , a naive mean pooling as in ( 4 ) may miss the interdependencies between local data , leading to a less well performing global model . Such interdependencies naturally occur in the power grid example because of the physics of an electricity network . Hence , in subsequent sections , we use latent alignment to address the first problem and graph neural network to address the second one . Incorporating these two components , we show the full , proposed framework in Figure 2 and Algorithm 1 ( supplement Section C ) .

Summary: This work is motivated by the problem of federated inference, which arises when clients own different parts of a datum, and learning must be collectively performed on the complete data representation. This application is presented as an extension of vertical federated learning. The proposed approach is based on the training of local models at the clients’ level. From the local models, an opportune data representation is extracted and shared in a centralized setting, which is then used to perform global inference. To cope with the mismatch of data representations across clients, a centralized alignment  procedure is proposed, which provides the input for a graph neural network encoding the relational structure among clients. Experiments are provided for the specific application of modeling real-life time series datasets from data owners of the US power grid, showing that the proposed local-global analysis pipeline outperforms both local and global training.