pMixFed: Efficient Personalized Federated Learning through Adaptive Layer-Wise Mixup

TL;DR

pMixFed is a novel personalized federated learning method that adaptively combines shared global and local models using layer-wise mixup, improving robustness, personalization, and handling data heterogeneity.

Contribution

It introduces an adaptive, layer-wise mixup strategy with a gradual personalization transition to enhance model accuracy and robustness in heterogeneous federated learning scenarios.

Findings

Outperforms state-of-the-art PFL methods in accuracy.

Achieves faster training convergence.

Demonstrates robustness across various data heterogeneity settings.

Abstract

Traditional Federated Learning (FL) methods encounter significant challenges when dealing with heterogeneous data and providing personalized solutions for non-IID scenarios. Personalized Federated Learning (PFL) approaches aim to address these issues by balancing generalization and personalization, often through parameter decoupling or partial models that freeze some neural network layers for personalization while aggregating other layers globally. However, existing methods still face challenges of global-local model discrepancy, client drift, and catastrophic forgetting, which degrade model accuracy. To overcome these limitations, we propose $\textit{pMixFed}$, a dynamic, layer-wise PFL approach that integrates $\textit{mixup}$ between shared global and personalized local models. Our method introduces an adaptive strategy for partitioning between personalized and shared layers, a gradual transition of personalization degree to enhance local client adaptation, improved generalization across clients, and a novel aggregation mechanism to mitigate catastrophic forgetting. Extensive experiments demonstrate that pMixFed outperforms state-of-the-art PFL methods, showing faster model training, increased robustness, and improved handling of data heterogeneity under different heterogeneous settings.