FedAH: Aggregated Head for Personalized Federated Learning

TL;DR

The paper tackles statistical heterogeneity in Federated Learning by proposing FedAH, a personalized-head-based PFL method that preserves a global feature extractor while introducing an Aggregated Head to fuse global information into each client's head. FedAH computes the Aggregated Head via element-wise mixing with learnable weights and trains the head and feature extractor in alternating steps, with the server aggregating updated client models to form a new global state. Across five CV/NLP benchmarks and two heterogeneity scenarios, FedAH outperforms ten SOTA FL methods by about 2.87% on CIFAR-100 and demonstrates robustness to client dropouts, scalability to many clients, and stability under varying local epochs. The approach advances practical PFL by enabling finer-grained global knowledge transfer to personalized components, and its open-source implementation facilitates adoption and further research.

Abstract

Recently, Federated Learning (FL) has gained popularity for its privacy-preserving and collaborative learning capabilities. Personalized Federated Learning (PFL), building upon FL, aims to address the issue of statistical heterogeneity and achieve personalization. Personalized-head-based PFL is a common and effective PFL method that splits the model into a feature extractor and a head, where the feature extractor is collaboratively trained and shared, while the head is locally trained and not shared. However, retaining the head locally, although achieving personalization, prevents the model from learning global knowledge in the head, thus affecting the performance of the personalized model. To solve this problem, we propose a novel PFL method called Federated Learning with Aggregated Head (FedAH), which initializes the head with an Aggregated Head at each iteration. The key feature of FedAH is to perform element-level aggregation between the local model head and the global model head to introduce global information from the global model head. To evaluate the effectiveness of FedAH, we conduct extensive experiments on five benchmark datasets in the fields of computer vision and natural language processing. FedAH outperforms ten state-of-the-art FL methods in terms of test accuracy by 2.87%. Additionally, FedAH maintains its advantage even in scenarios where some clients drop out unexpectedly. Our code is open-accessed at https://github.com/heyuepeng/FedAH.

Figures (4)

• Figure 1: An example for FedAH. $r^t$ : global feature extractor; $h^t$ : global heads; $h^{t-1}_{1},\ldots,h^{t-1}_{n}$: local head of clients $1,\ldots, n$ in the previous iteration; $\hat{h}^t_{1},\ldots,\hat{h}^t_{n}$ : Aggregated Heads of clients $1,\ldots, n$ in the current iteration.
• Figure 2: Local learning process of FedAH on client $i$ in the $t$-th iteration.
• Figure 3: The inputs and outputs of the feature extractor and head in personalized-head-based PFL methods.
• Figure 4: Test accuracy (%) curves of different methods on Tiny-ImageNet using ResNet-18.