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 paper introduces a new problem, called federated inference, and proposed a framework of solutions including local representation alignment and learning a consensus graph. The technical contribution mainly comes from a theorem (Theorem 1) on the convergence of learning permutation matrices, and a new approach (icdf) for parameter inference of matrix Bernoulli distribution. The results demonstrate the proposed approaches are useful in improving the prediction accuracy.