FedBABU: Towards Enhanced Representation for Federated Image Classification

TL;DR

FedBABU introduces a federated learning approach that enhances representation by training only the universal body of the model during federation and fine-tuning the personalized head during evaluation, improving both global and personalized performance.

Contribution

The paper proposes FedBABU, a novel federated learning algorithm that updates only the model's body during training and fine-tunes the head for personalization, addressing performance degradation issues.

Findings

FedBABU achieves consistent performance improvements.

The method enables efficient personalization.

Training only the body enhances universal representation.

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.