In federated learning problems, data is scattered across different servers and exchanging or pooling it is often impractical or prohibited. We develop a Bayesian nonparametric framework for federated learning with neural networks. Each data server is assumed to train local neural network weights, which are modeled through our framework. We then develop an inference approach that allows us to synthesize a more expressive global network without additional supervision or data pooling. We then demonstrate the efficacy of our approach on federated learning problems simulated from two popular image classification datasets. The standard machine learning paradigm involves algorithms that learn from centralized data, possibly pooled together from multiple data sources. The computations involved may be done on a single machine or farmed out to a cluster of machines. However, in the real world, data often lives in silos and amalgamating them may be rendered prohibitively expensive by communication costs, time sensitivity, or privacy concerns. Consider, for instance, data recorded from sensors embedded in wearable devices. Such data is inherently private, can be voluminous depending on the sampling rate of the sensing modality, and may be time sensitive depending on the analysis of interest. Pooling data from many users is technically challenging owing to the severe computational burden of moving large amounts of data, and fraught with privacy concerns stemming from potential data breaches that may expose the user's protected health information (PHI).Federated learning avoids these pitfalls by obviating the need for centralized data and instead designs algorithms that learn from sequestered data sources with different data distributions. To be effective , such algorithms must be able to extract and distill important statistical patterns from various independent local learners coherently into an effective global model without centralizing data. This will allow us to avoid the prohibitively expensive cost of data communication. To achieve this , we develop and investigate a probabilistic federated learning framework with a particular emphasis on training and aggregating neural network models on siloed data.We proceed by training local models for each data source, in parallel. We then match the estimated local model parameters (groups of weight vectors in the case of neural networks) across data sources to construct a global network. The matching, to be formally defined later, is governed by the posterior of a Beta-Bernoulli process (BBP) (Thibaux & Jordan, 2007; Yurochkin et al., 2018) , a Bayesian nonparametric model that allows the local parameters to either match existing global ones or create a new global parameter if existing ones are poor matches. Our construction allows the size of the global network to flexibly grow or shrink as needed to best explain the observed data. Crucially, we make no assumptions about how the data is distributed between the different sources or even about the local learning algorithms. These may be adapted as necessary, for instance to account for non-identically distributed data. Further, we only require communication after the local algorithms have converged. This is in contrast with popular distributed training algorithms that rely on frequent communication between the local machines. Our construction also leads to compressed global models with fewer parameters than the set of all local parameters. Unlike naive ensembles of local models, this allows us to store fewer parameters and leads to more efficient inference at test time, requiring only a single forward pass through the compressed model as opposed to J forward passes, once for each local model. While techniques such as distillation allow for the cost of multiple forward passes to be amortized, training the distilled model itself requires access to data pooled across all sources, a luxury unavailable in our federated learning scenario. In summary, the key question we seek to answer in this paper is the following: given pre-trained neural networks trained locally on non-centralized data, can we learn a compressed federated model without accessing the original data, while improving on the performance of the local networks?The remainder of the paper is organized as follows. We briefly introduce the Beta-Bernoulli process in Section 2 before describing our model for federated learning in Section 3. We thoroughly vet the proposed models and demonstrate the utility of the proposed approach in Section 4. Finally, Section 5 discusses limitations and open questions. In this work we have developed models for matching fully connected networks, and experimentally demonstrated the capabilities of our methodology, particularly when prediction time is limited and communication is expensive. We also observed the importance of convergent local neural networks that serve as inputs to our matching algorithms. Poor quality local neural network weights will affect the quality of the master network. In future work we plan to explore more sophisticated ways to account for uncertainty in the weights of small batches. Additionally, our matching approach is completely unsupervised -incorporating some form of supervised signal may help to improve the performance of the global network when local networks are low quality. Finally, it is of interest to extend our modeling framework to other architectures such as Convolutional Neural Networks (CNNs) and Recurrent Neural Networks (RNNs). The permutation invariance necessitating matching inference arises in CNNs too -any permutation of the filters results in same output, however additional bookkeeping is needed due to pooling operations.Ohad Shamir, Nati Srebro, and Tong Zhang. Communication-efficient distributed optimization using an approximate newton-type method. In International conference on machine learning, pp. The goal of maximum a posteriori (MAP) estimation is to maximize posterior probability of the latent variables: global atoms DISPLAYFORM0 and assignments of observed neural network weight estimates to global atoms {B j } J j=1 , given estimates of the batch weights DISPLAYFORM1 arg max DISPLAYFORM2 MAP estimates given matching (Proposition 1 in the main text) First we note that given {B j } it is straightforward to find MAP estimates of {✓ i } based on Gaussian-Gaussian conjugacy: DISPLAYFORM3 where L = max{i : DISPLAYFORM4 . . , J} is the number of active global atoms, which is an (unknown) latent random variable identified by {B j }. For simplicity we assume ⌃ 0 = I 2 0 , ⌃ j = I 2 j and µ 0 = 0.Inference of atom assignment. We can now cast optimization corresponding to (1) with respect to only {B j } J j=1 . Taking natural logarithm we obtain: DISPLAYFORM5 Let us first simplify the first term of (3):1 2 DISPLAYFORM6 We consider an iterative optimization approach: fixing all but one B j we find corresponding optimal assignment, then pick a new j at random and proceed until convergence. In the following we will use notation j to say "all but j". Let L j = max{i : B j i,l = 1} denote number of active global weights outside of group j. We now rearrange (4) by partitioning it into i = 1, . . . , L j and i = L j + 1, . . . , L j + L j . We are interested in solving for B j , hence we can modify objective function by subtracting terms independent of B j : DISPLAYFORM7 Now observe that P l B j i,l 2 {0, 1}, i.e. it is 1 if some neuron from batch j is matched to global neuron i and 0 otherwise. Due to this we can rewrite (5) as a linear sum assignment problem: DISPLAYFORM8 Now we consider second term of FORMULA12 : DISPLAYFORM9 First, because we are optimizing for B j , we can ignore log P (B j ). Second, due to exchangeability of batches (i.e. customers of the IBP), we can always consider B j to be the last batch (i.e. last customer of the IBP). Let m j i = P j,l B j i,l denote number of times batch weights were assigned to global atom i outside of group j. We now obtain the following: DISPLAYFORM10 We now rearrange (7) as linear sum assignment problem: DISPLAYFORM11 Combining FORMULA20 and FORMULA25 we arrive at the cost specification for finding B j as minimizer of DISPLAYFORM12 , where: DISPLAYFORM13 This completes the proof of Proposition 2 in the main text. FIG3 illustrates the overall multilayer inference procedure visually, and Algorithm 1 provides the details. Nodes in the graphs indicate neurons, neurons of the same color have been matched. On the left, the individual layer matching approach is shown, consisting of using the matching assignments of the next highest layer to convert the neurons in each of the J servers to weight vectors referencing the global previous layer. These weight vectors are then used to form a cost matrix, which the Hungarian algorithm then uses to do the matching. Finally, the matched neurons are then aggregated and averaged to form the new layer of the global model. As shown on the right, in the multilayer setting the resulting global layer is then used to match the next lower layer, etc. until the bottom hidden layer is reached FIG3 ,... in order).
We propose a Bayesian nonparametric model for federated learning with neural networks. Uses beta process to do federated neural matching. The paper considers federate learning of neural networks, where data is distributed on multiple machines and the allocation of data points is potentially inhomogenous and unbalanced.