PaperlayerView exact revision on arXiv

Source-linked AI summary

Bayesian Nonparametric Federated Learning of Neural Networks

Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Trong Nghia Hoang, Yasaman Khazaeni

arXiv:1905.12022v1stat.MLcs.LG

TL;DR

Federated learning must operate when data are siloed, private, or unavailable for pooling. The paper develops a Bayesian nonparametric neural-network matching framework that aggregates local weights into compressed global models, including from pre-trained networks, with as few as one communication round. Experiments on image-classification federated-learning problems show favorable performance and compression properties.

  • Problem

    Federated learning needs methods that can learn across siloed sources when data pooling is impractical or pre-trained local models are available without their data.

  • Method

    A Beta-Bernoulli process matches local neural-network parameters into a global network, decoupling local learning from aggregation and requiring no additional data for pre-trained models.

  • Results

    The method empirically demonstrates favorable properties, including compressed federated networks that can outperform state-of-the-art neural-network federated-learning algorithms with a modest communication budget.

  • Takeaways & Limitations

    The framework supports expressive, compressed global neural networks from local models without pooled data, additional supervision, or repeated local-model inference.

  • Takeaways & Limitations

    The approach assumes practical consistency in input-dimension ordering and applies permutation invariance on the first hidden-layer side.

Abstract

from arXiv · show

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 provide 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, data pooling and with as few as a single communication round. We then demonstrate the efficacy of our approach on federated learning problems simulated from two popular image classification datasets.

1. Introduction

The paper addresses federated learning when data remain siloed or unavailable, developing a Bayesian nonparametric framework to aggregate local neural networks. Its approach matches local parameters into a global network without pooling data and supports combining pre-trained models without additional data or learning-algorithm knowledge.

  • Data silos make pooling impractical because of communication costs, time sensitivity, computational burden, and privacy concerns.
  • Federated learning trains local models and distills them into a global model without explicitly combining source data.
  • The paper targets the more extreme setting where pre-trained local models are available but their training data are inaccessible.
  • A Beta-Bernoulli process posterior matches local neural-network parameters across sources, allowing matches to existing global parameters or creation of new ones.
  • The framework combines pre-trained models without additional data or knowledge of the algorithms that produced them.
  • Unlike naive ensembles, its Bayesian nonparametric construction produces compressed global models requiring one forward pass instead of J local-model passes.

2. Background and Related Works

The background situates the approach within Bayesian nonparametrics and prior federated and distributed-learning methods. The Beta-Bernoulli process models shared, sparsely selected latent components, while the proposed framework is nonparametric and communication-conscious.

  • Beta-Bernoulli Process: The Beta-Bernoulli process builds subsets by Bernoulli-selecting atoms from a shared Beta-process random measure.
  • Beta-Bernoulli Process: Its subsets are conditionally independent given the shared measure, while marginalizing that measure induces dependencies among them.
  • Beta-Bernoulli Process: The Indian Buffet Process represents a distribution over sparse binary matrices with infinitely many columns and is closely related to the Beta-Bernoulli process.
  • Federated and Distributed Learning: Earlier federated-learning work includes multitask SVM formulations, but that approach does not extend to neural-network structures.
  • Federated and Distributed Learning: The proposed framework is nonparametric, allowing federated-model complexity to grow or shrink with varying data complexity.
  • Federated and Distributed Learning: Distributed deep-learning methods emphasize scalable training and typically require frequent communication between distributed nodes.

3. Probabilistic Federated Neural Matching

Probabilistic Federated Neural Matching models local neural-network neurons as shared or distinct global atoms, matches them across models, and aggregates matched parameters into a global network. Its Bayesian nonparametric construction supports compressed models and extension across communication rounds.

  • Single-layer matching: PFNM identifies subsets of neurons across J local models and combines matched neurons to form a global model.The framework treats matching as the central inference problem for constructing the global network.
  • Model size: The global network has an unknown hidden-layer width L that is inferred to be much smaller than the total number of local hidden units.This Bayesian nonparametric construction can recover compressed global models with fewer parameters than the collection of all local parameters.
  • Single-layer matching: Neural-network hidden-layer neurons are permutation invariant, so matching must compare feature extractors rather than their local indices.Reordering hidden-layer columns and corresponding output-layer rows leaves the network outputs unchanged.
  • Bayesian nonparametric model: The model represents global hidden neurons as Beta-Bernoulli-process atoms, selects atom subsets for each batch, and models local atoms as noisy measurements.Binary assignments identify which global atoms are used by each batch, while Gaussian observations connect local weight vectors to those atoms.
  • Inference: The inference procedure forms assignment costs from local weight vectors, solves neuron assignments with the Hungarian algorithm, and estimates global weight vectors from matched instances.The resulting global neurons and inferred parameters are concatenated into a new global hidden layer.
  • Additional communications: PFNM extends naturally to additional communication rounds by reinitializing local models from subsets of the global model, with global model size expected to shrink as local models improve.The procedure first aggregates local models in a single round and can then repeat matching after local updates.

4. Experiments

Experiments evaluate PFNM for single- and limited-communication federated learning of neural networks on MNIST and CIFAR10. PFNM generally matches or exceeds strong alternatives while producing substantially smaller global models and requiring fewer communication rounds.

  • Experimental setup: Experiments simulate homogeneous and heterogeneous federated learning on MNIST and CIFAR10, using randomly partitioned batches with varying class proportions and sizes.Heterogeneous partitions are generated with unbalanced batch sizes and class proportions.
  • Single communication: Single-communication experiments compare PFNM with local models, ensembles, federated averaging, and k-means when only pre-trained local models are available.Ensembles provide a strong accuracy reference but require storing and evaluating every local network.
  • Single communication: PFNM consistently outperforms local models and federated averaging, often matches or exceeds ensembles, and achieves this with fewer stored parameters.For multi-layer networks, PFNM can perform marginally worse than ensembles on deeper CIFAR10 models while retaining significant compression.
  • Single communication: PFNM’s performance remains strong under heterogeneous data, whereas local neural networks degrade noticeably as batch sizes and class proportions become unbalanced.The gap between PFNM and ensembles increases with the number of batches on CIFAR10 but remains constant or favors PFNM on MNIST.
  • Limited communication: With J = 25 batches and additional communication, PFNM outperforms ensembles in all reported one- and two-layer scenarios given sufficient rounds.Additional communication steadily improves accuracy and reduces global model size, although two-layer CIFAR10 PFNM temporarily drops in performance during sharp compression.
  • Limited communication: PFNM requires significantly fewer communication rounds than federated averaging and D-SGD to reach a given performance level while also shrinking the global model.The experiments allow up to twenty rounds for homogeneous partitions and up to fifty for heterogeneous partitions.

5. Discussion

The discussion concludes that the methods effectively learn compressed federated networks from pre-trained local models and can outperform state-of-the-art neural-network federated-learning algorithms with modest communication. It also identifies unsupervised matching, poor local models, and extension to CNNs and RNNs as directions and boundaries for future work.

  • Conclusions: The methods are particularly effective for learning compressed federated networks from pre-trained local networks with a modest communication budget.The authors report favorable empirical properties and performance against state-of-the-art federated-learning algorithms for neural networks.
  • Limitations and future work: The current matching approach is completely unsupervised, and supervision may improve global-network performance when local networks are poor.The authors propose more sophisticated combinations of local networks for regimes with very few training instances.
  • Limitations and future work: The framework is proposed for extension to CNNs and RNNs, where permutation invariance also requires matching but pooling operations add bookkeeping.The stated CNN/RNN extension is a future direction rather than an evaluated result in this discussion passage.

A. Single Hidden Layer Inference

Single-hidden-layer inference alternates between estimating global atoms and optimizing local-to-global neuron assignments. The assignment step is reduced to a linear sum assignment problem, while matched weights determine MAP global-atom estimates.

  • MAP estimation maximizes the posterior probability of global atoms and local assignment variables.
  • Given assignments, Gaussian-Gaussian conjugacy provides straightforward MAP estimates for the global atoms.
  • The method infers assignments by optimizing the posterior with respect to the batch-specific matching matrices.
  • Inference iteratively fixes all but one batch assignment, optimizes the remaining assignment, and repeats with a randomly selected batch until convergence.
  • The assignment objective is rearranged as a linear sum assignment problem and solved using the Hungarian algorithm.
  • The resulting procedure completes the proof of Proposition 1.

B. Multilayer Inference Details

Multilayer inference proceeds greedily from the top layer downward. Each inferred global layer supplies the matching reference for the next lower layer, ending with the bottom hidden layer.

  • The multilayer procedure matches hidden layers in a top-down loop from layer C through layer 2.Each iteration collects the corresponding hidden layer from all batches and invokes single-layer neural matching.
  • After matching a layer, the algorithm forms the global neuron layer and updates its inferred width from the union of matched neurons.
  • The bottom layer is matched using weights connected to both the input and the layer above, after which the global multilayer model is returned.
  • At each step, the resulting global layer is used to match the next lower layer until the bottom hidden layer is reached.

C. Complexity Analysis

The matching algorithm’s cost depends on model dimensions and batch count rather than dataset size. In the reported experiments, runtime varied with input dimension and network depth.

  • The worst-case per-layer complexity is O(D(JL_j)^2) for cost-matrix construction plus O((JL_j)^3) for the Hungarian algorithm.This case occurs when no neurons are matched.
  • The best-case per-layer complexity is O(DL^2_j) when all neurons are matched.
  • The complexity is independent of data size and is closer in practice to the best case when the global model is moderate relative to local models.
  • At J = 30 groups, reported runtimes were 40 seconds for Figure 2a,b and 500 seconds for Figure 2c,d.The text attributes the larger latter runtime to the higher CIFAR-10 input dimension.
  • For six-layer experiments with J = 10, reported runtimes were 60 seconds for Figure 2e,f and 150 seconds for Figure 2g,h.The computations used two CPU cores and 4GB memory.

D. Experimental Details and Additional Results

Experiments simulate federated learning by partitioning MNIST and CIFAR-10 across non-overlapping batches, then training local neural networks under specified settings. Results are summarized across repeated trials using mean accuracies and standard deviations.

  • Data partitioning: The experiments partition MNIST and CIFAR-10 across multiple non-overlapping batches that may have different data distributions.Two partition strategies are considered, including homogeneous partitioning with similar class distributions across batches.
  • Data partitioning: Each partition-strategy and dataset combination is evaluated over 10 trials to obtain mean accuracies and standard deviations.
  • Batch networks training: The framework and ensemble methods operate on collections of neural-network weights trained independently on the batch datasets.
  • Batch networks training: Local networks are implemented and trained in PyTorch using the AMSGrad optimizer with default parameters unless otherwise specified.Parameter settings are summarized in Table 1.

D.1. Parameter Settings for the Baselines

The baseline experiments specify ensemble, downpour SGD, and Federated Averaging procedures, alongside sensitivity analyses for selected parameters.

  • Ensemble predictions use the class with the highest aggregated probability across local neural networks.
  • Downpour SGD updates the master network with accumulated local gradients after every mini-batch.The local learners use PyTorch and communicate with npush = nfetch = 1.
  • Federated Averaging initializes local networks with the same seed and uses all local networks for each central update.Local training runs for 10 epochs initially and 5 epochs after the first communication round.
  • Figure 5 analyzes parameter sensitivity at J = 25, while Figure 6 analyzes σ2 sensitivity for fixed σ2.
  • The experiments include the settings σ0^2 = 10 and γ0 = 1 for varying J.

D.2. Parameter Settings for Matching with Additional Communications

Matching with additional communications trains local networks repeatedly using AMSGrad and decays the learning rate after each communication.

  • Local networks train for 10 epochs in the first round and 5 epochs thereafter, using AMSGrad with optimizer resets after every communication.The initial learning rate is decayed by a factor of 0.99 after each communication.

D.3. Parameter Sensitivity Analysis for PFNM

PFNM is empirically evaluated for sensitivity to its prior variance, neuron-discovery parameter, and local-weight variance. Its performance is generally robust, with σ2 having greater effects under heterogeneous partitioning.

  • PFNM uses σ0^2 as the global-network weight prior variance, γ0 to control new-neuron discovery, and σ2 as local-weight variance around global weights.Increasing γ0 increases the learned global network size.
  • Performance does not fluctuate significantly across the parameter values tested for a single-hidden-layer model with J = 25 batches.The analysis measures accuracy on the training data.
  • The experiments set σ0^2 = 10 and γ0 = 1 for single-communication-round experiments.
  • σ2 has slightly higher performance impact than the other parameters and is selected using training data.
  • Under homogeneous partitioning, σ2 has almost no effect, whereas heterogeneous partitioning makes its effect more noticeable while all tested values remain competitive.
Loading 1905.12022v1…