Personalized Federated Learning for Gradient Alignment

2026-05-03

AI Summary

This work addresses the challenges in personalized federated learning arising from high local gradient variance due to heterogeneous and limited client data, as well as the distortion of personalized optimization directions during model aggregation. The authors propose pFLAlign, a novel framework that, for the first time, derives a gradient alignment mechanism from a PAC-Bayesian perspective. pFLAlign employs a two-stage strategy: it adaptively adjusts gradient directions during local training and re-aligns personalized directions after global aggregation to preserve client-specific information. This approach significantly enhances both personalization performance and training stability, achieving state-of-the-art results across multiple benchmarks. Ablation studies further confirm the effectiveness of each component in the proposed framework.

Abstract

Personalized federated learning (pFL) aims to adapt models to client specific data distributions, yet it often fails to reliably preserve personalized information. Local training is hindered by high variance gradients induced by limited and heterogeneous client data, while aggregation further distorts client specific optimization directions. To address these challenges, we propose pFLAlign, a gradient alignment framework to maintain client specific information during both local training and aggregation. pFLAlign consists of two complementary mechanisms: one adapts local gradient directions to reduce variance during client side optimization, and the other mitigates aggregation induced distortion by realigning the global model with each client's personalized direction. Theoretically, we derive pFLAlign from a PAC Bayesian analysis, which reveals how personalized gradient alignment preserves client specific information. Our experiments and ablation studies show that pFLAlign consistently improves personalization performance and training stability, achieving state of the art results.

Problem

• personalized federated learning
• gradient variance
• data heterogeneity
• model aggregation
• personalization preservation

Innovation

• personalized federated learning
• gradient alignment
• PAC-Bayesian analysis
• local optimization
• model aggregation