Fed-ensemble: Ensemble Models in Federated Learning for Improved Generalization and Uncertainty Quantification
Naichen Shi, Fan Lai, Raed Al Kontar, Mosharaf Chowdhury · IEEE Transactions on Automation Science and Engineering · 2023
The increase in the computational power of edge devices has opened up the possibility of processing some of the data at the edge and distributing model learning. This paradigm is often called federated learning (FL), where edge devices exploit their local computational resources to train models collaboratively. Though FL has seen recent success, it is unclear how to characterize uncertainties in FL predictions. In this paper, we proposeFed-ensemble: a simple approach that brings model ensembling to FL. Instead of aggregating local models to update a single global model,Fed-ensembleuses random permutations to update a group of$K$models and then obtains predictions through model averaging.Fed-ensemblecan be readily utilized within established FL methods and does not impose a computational overhead compared with single-model methods. Empirical results show that our model has superior performance over several FL algorithms on a wide range of data sets and excels in heterogeneous settings often encountered in FL applications. Also, by carefully choosing client-dependent weights in the inference stage,Fed-ensemblebecomes personalized and yields even better performance. Theoretically, we show that predictions on new data from all$K$models belong to the same predictive posterior distribution under a neural tangent kernel regime. This result, in turn, sheds light on the generalization advantages of model averaging and justifies the uncertainty quantification capability. We also illustrate thatFed-ensemblehas an elegant Bayesian interpretation.Note to Practitioners—provides an algorithm that extracts a set of$K$solutions without imposing any additional communication overhead in FL. Given multiple solutions,Fed-ensemblecan be exploited to personalize inference as well as quantify uncertainty. Such capabilities may be beneficial within multiple practical systems that require uncertainty-aware decision-making. Further,Fed-ensemblemay be useful for model validation and hypothesis testing.