FedBE: Making Bayesian Model Ensemble Applicable to Federated Learning

Abstract

Federated learning aims to collaboratively train a strong global model by accessing users' locally trained models but not their own data. A crucial step is therefore to aggregate local models into a global model, which has been shown challenging when users have non-i.i.d. data. In this paper, we propose a novel aggregation algorithm named FedBE, which takes a Bayesian inference perspective by sampling higher-quality global models and combining them via Bayesian model Ensemble, leading to much robust aggregation. We show that an effective model distribution can be constructed by simply fitting a Gaussian or Dirichlet distribution to the local models. Our empirical studies validate FedBE's superior performance, especially when users' data are not i.i.d. and when the neural networks go deeper. Moreover, FedBE is compatible with recent efforts in regularizing users' model training, making it an easily applicable module: you only need to replace the aggregation method but leave other parts of your federated learning algorithm intact. Our code is publicly available at https://github.com/hongyouc/FedBE.

arXiv commentAccepted to ICLR 2021; the camera-ready version with code URL

References in corpus (28)

• Distilling the Knowledge in a Neural Network
• MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications
• Communication-Efficient Learning of Deep Networks from Decentralized Data
• Federated Learning: Strategies for Improving Communication Efficiency
• MobileNetV2: Inverted Residuals and Linear Bottlenecks
• A Tutorial on Bayesian Optimization of Expensive Cost Functions, with Application to Active User Modeling and Hierarchical Reinforcement Learning
• Federated Learning with Non-IID Data
• On Calibration of Modern Neural Networks
• On the Convergence of FedAvg on Non-IID Data
• Towards Federated Learning at Scale: System Design
• Ensemble Distillation for Robust Model Fusion in Federated Learning
• FedMD: Heterogenous Federated Learning via Model Distillation
• Federated Optimization in Heterogeneous Networks
• Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization
• FedPAQ: A Communication-Efficient Federated Learning Method with Periodic Averaging and Quantization
• Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss
• Local SGD Converges Fast and Communicates Little
• Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs
• Group Knowledge Transfer: Federated Learning of Large CNNs at the Edge
• Essentially No Barriers in Neural Network Energy Landscape
• On the Convergence of Local Descent Methods in Federated Learning
• Large scale distributed neural network training through online distillation
• One-Shot Federated Learning
• Adaptive Federated Optimization
• Federated Multi-Task Learning
• Variance Reduced Local SGD with Lower Communication Complexity
• Federated Learning with Unbiased Gradient Aggregation and Controllable Meta Updating
• Private Query Release Assisted by Public Data

Cited by in corpus (6)

• Ensemble Distillation for Robust Model Fusion in Federated Learning
• Auto-FedAvg: Learnable Federated Averaging for Multi-Institutional Medical Image Segmentation
• Communication-Efficient Federated Distillation
• Distributed Computation for Marginal Likelihood based Model Choice
• Edge Bias in Federated Learning and its Solution by Buffered Knowledge Distillation
• FedRAD: Federated Robust Adaptive Distillation