pFLAIR: personalized Federated Learning using Attention-Informed model Recombination

Personalized federated learning (pFL) has emerged as a key approach for addressing statistical heterogeneity, while retaining the privacy and collaboration benefits of federated learning. However, most existing pFL methods place the bulk of personalization on the client side, often adding extra local modules, inner-loop optimization, or per-client adaptation steps, which can be computationally and memory intensive, and contrary to a core practical motivation of FL: keeping client workloads and on-device model footprints small. This is especially limiting under label skew, the most realistic heterogeneity regime, where clients observe different subsets or frequencies of classes reflecting natural specialization in real-world deployments. We propose pFLAIR, a pFL framework that shifts a substantial part of personalization to the server by learning an explicit, trainable mixing mechanism that determines how client knowledge should be combined. Crucially, unlike several competitive pFL methods that rely on communicating explicit label frequencies or class-conditioned statistics to achieve their best performance, pFLAIR operates without requiring explicit label-distribution statistics. This avoids exposing such statistics to the server, which is the very premise of federated learning. Across heterogeneous settings, pFLAIR consistently outperforms strong pFL baselines, with particularly large gains under challenging label skew, and degrades far more gracefully than competing methods as the number of clients grows, demonstrating a better scalability for federated deployments.