FedAli: Personalized Federated Learning Alignment with Prototype Layers for Generalized Mobile Services

TL;DR

FedAli is a novel personalized federated learning method that uses prototype-based regularization to improve client generalization and reduce model divergence in mobile systems with heterogeneous data.

Contribution

Introduces FedAli, a prototype-based regularization technique with an ALP layer, to enhance inter-client alignment and robustness in personalized federated learning.

Findings

Significantly improves client generalization in heterogeneous settings.

Effectively reduces client drift through prototype alignment.

Accelerates convergence with prototype pre-training.

Abstract

Personalized Federated Learning (PFL) enables distributed training on edge devices, allowing models to collaboratively learn global patterns while tailoring their parameters to better fit each client's local data, all while preserving data privacy. However, PFL faces two key challenges in mobile systems: client drift, where heterogeneous data cause model divergence, and the overlooked need for client generalization, as the dynamic of mobile sensing demands adaptation beyond local environments. To overcome these limitations, we introduce Federated Alignment (FedAli), a prototype-based regularization technique that enhances inter-client alignment while strengthening the robustness of personalized adaptations. At its core, FedAli introduces the ALignment with Prototypes (ALP) layer, inspired by human memory, to enhance generalization by guiding inference embeddings toward personalized prototypes while reducing client drift through alignment with shared prototypes during training. By leveraging an optimal transport plan to compute prototype-embedding assignments, our approach allows pre-training the prototypes without any class labels to further accelerate convergence and improve performance. Our extensive experiments show that FedAli significantly enhances client generalization while preserving strong personalization in heterogeneous settings.