FedBABU: Towards Enhanced Representation for Federated Image Classification
arXiv:2106.06042
Abstract
Federated learning has evolved to improve a single global model under data heterogeneity (as a curse) or to develop multiple personalized models using data heterogeneity (as a blessing). However, little research has considered both directions simultaneously. In this paper, we first investigate the relationship between them by analyzing Federated Averaging at the client level and determine that a better federated global model performance does not constantly improve personalization. To elucidate the cause of this personalization performance degradation problem, we decompose the entire network into the body (extractor), which is related to universality, and the head (classifier), which is related to personalization. We then point out that this problem stems from training the head. Based on this observation, we propose a novel federated learning algorithm, coined FedBABU, which only updates the body of the model during federated training (i.e., the head is randomly initialized and never updated), and the head is fine-tuned for personalization during the evaluation process. Extensive experiments show consistent performance improvements and an efficient personalization of FedBABU. The code is available at https://github.com/jhoon-oh/FedBABU.
Published at ICLR 2022
References in corpus (6)
- PyTorch: An Imperative Style, High-Performance Deep Learning Library
- Federated Learning with Personalization Layers
- Federated Learning Based on Dynamic Regularization
- Exploiting Shared Representations for Personalized Federated Learning
- Personalized Federated Learning using Hypernetworks
- Personalized Federated Learning with First Order Model Optimization