[Paper Review] FedBABU: Towards Enhanced Representation for Federated Image Classification

Jaehoon Oh, Sangmook Kim+1 more|arXiv (Cornell University)|Jun 4, 2021

Privacy-Preserving Technologies in DataComputer Science48 references48 citations

TL;DR

FedBABU decouples the model into a trainable body and a fixed head during federated training, enabling efficient personalization by fine-tuning the head at inference; it improves representation power and personalization over FedAvg in heterogeneous data settings.

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.

Motivation & Objective

• Motivate understanding of how single global models interact with personalized models under data heterogeneity.
• Decouple neural networks into a universal body and a personalized head to identify sources of personalization degradation.
• Propose and validate FedBABU, which updates only the body during federated training and freezes the head.
• Demonstrate that a fixed head across clients yields strong representation power and accelerates on-device personalization.

Proposed method

• Decouple model parameters into body (extractor) and head (classifier).
• During FedBABU, update only the body on each client; the head uses a fixed, randomly initialized global head.
• Aggregate only the body parameters across clients; do not update or aggregate the head during federated training.
• Evaluate personalization by comparing initial global-model performance and client-specific fine-tuning performance.
• Show that a fixed head can provide comparable centralized performance to a jointly trained head.
• Optionally replace trained heads with templates to assess representation power without the head.

Experimental results

Research questions

• RQ1 Can decoupling body and head mitigate personalization degradation in federated learning with heterogeneous data?
• RQ2 Does freezing the head during federated training improve the quality of universal representations for subsequent personalized fine-tuning?
• RQ3 How does FedBABU compare to FedAvg and other personalization methods in terms of initial and personalized accuracy under varying data heterogeneity?
• RQ4 Is the representation power of FedBABU’s global model robust when evaluated without the head, using class templates?
• RQ5 Can FedBABU generalize to regularized FL frameworks like FedProx?

Key findings

• Training-only-the-body in FedBABU yields stronger or comparable representation power versus FedAvg, especially under high heterogeneity.
• Fixed head across clients provides robust guidance for representation learning, improving personalization efficiency.
• FedBABU enables rapid personalization, often with only a few fine-tuning epochs, and outperforms FedAvg in several settings.
• When the head is not trained, FedBABU achieves higher or competitive initial accuracy without sacrificing personalization.
• Fine-tuning the head or full model yields better personalization for FedAvg, while FedBABU can achieve strong results by fine-tuning the head alone.