Personalized Federated Learning with Multi-branch Architecture

Federated learning (FL) is a decentralized machine learning technique that enables multiple clients to collaboratively train models without revealing the raw data to each other. Although the traditional FL trains a single global model with average performance among clients, the statistical data heterogeneity across clients motivates personalized FL (PFL) which learns personalized models with good performance on each client's data. A key challenge in PFL is how to promote clients with similar data to collaborate more in a situation where each client has data from complex distribution and does not know each other's distribution. In this paper, we propose a new PFL method, personalized federated learning with multi-branch architecture (pFedMB), which achieves personalization by splitting each layer of neural networks into multiple branches and assigning client-specific weights to each branch. pFedMB is simple but effective to facilitate each client to share the knowledge with similar clients by adjusting the weights assigned to each branch. We experimentally show that pFedMB performs better than the state-of-the-art PFL methods using CIFAR10 dataset.