Heterogeneous Federated Learning Through Multi-Branch Network

Abstract

Recently, federated learning has gained increasing attention for privacy-preserving computation since the learning paradigm allows to train models without the need for exchanging the data across different institutions distributively. However, heterogeneity of computational capabilities of edge devices is seldom discussed and analyzed in the current literature for heterogeneous federated learning. To address this issue, we propose a novel heterogeneous federated learning framework based on multi-branch deep neural network models which enable the selection of a proper sub-branch model for the client devices according to their computational capabilities. Meanwhile, we also present an aggregation method for model training, MFedAvg, that performs branch-wise averaging-based aggregation. With extensive experiments on MNIST, FashionMNIST, MedMNIST, and CIFAR-10, it demonstrates that our proposed approaches can achieve satisfactory performance with guaranteed convergence and effectively utilize all the available resources for training across different devices with lower communication cost than its homogeneous counterpart.