FedBranch: Heterogeneous Federated Learning via Multi-Branch Neural Network

As a privacy-preserving paradigm of decentralized machine learning, federated learning (FL) has become a hot spot in the field of machine learning. Existing FL approaches generally assume that the global model can be deployed and trained on any client. However, in practical applications, the devices participated in FL are often heterogeneous and have different computation capacities, resulting in the difficulty of large neural network model training. The current solutions, such as reducing the scale of the global model to fit all clients or removing weak devices to deploy a larger model, will lead to model accuracy degradation, owing to the limitation of model scale or the loss of data on weak clients. To address the device heterogeneity issue inherent in FL, we propose FedBranch, a heterogeneous FL framework based on multi-branch neural network model. Its core idea is to assign a proper branch model to each client according to their computation capacity. In FedBranch, a layer-wise aggregation method is designed to address aggregation of different branches. Meanwhile, we introduce a model regularization method to improve the convergence efficiency and model performance of FedBranch. Besides, we propose a training task offloading algorithm based on Split Learning to safely and effectively share training tasks among different branch models. Extensive experiments conducted on different datasets demonstrate that our FedBranch method has higher convergence efficiency and model accuracy than existing federated learning methods in various heterogeneous scenarios.