FedPeWS: Personalized Warmup via Subnetworks for Enhanced Heterogeneous Federated Learning

2024-10-03

arXiv.org

AI Summary

To address parameter conflict and slow convergence in federated learning under extreme data heterogeneity, this paper proposes a personalized subnet warm-up mechanism. During the initial training phase, clients learn only lightweight binary masks to activate and update sparse subnetworks tailored to their local data distributions; subsequently, they transition to full-model aggregation. This approach is the first to integrate personalized sparse training with a two-stage optimization paradigm—warm-up followed by standard aggregation—and introduces gradient masking alongside joint mask-parameter updates. Evaluated on multiple highly heterogeneous benchmarks, the method achieves average accuracy improvements of 3.2–7.8% over baselines including FedAvg and FedProx, accelerates convergence by 1.8–2.5×, and incurs no additional communication overhead.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing

Abstract

Statistical data heterogeneity is a significant barrier to convergence in federated learning (FL). While prior work has advanced heterogeneous FL through better optimization objectives, these methods fall short when there is extreme data heterogeneity among collaborating participants. We hypothesize that convergence under extreme data heterogeneity is primarily hindered due to the aggregation of conflicting updates from the participants in the initial collaboration rounds. To overcome this problem, we propose a warmup phase where each participant learns a personalized mask and updates only a subnetwork of the full model. This personalized warmup allows the participants to focus initially on learning specific subnetworks tailored to the heterogeneity of their data. After the warmup phase, the participants revert to standard federated optimization, where all parameters are communicated. We empirically demonstrate that the proposed personalized warmup via subnetworks (FedPeWS) approach improves accuracy and convergence speed over standard federated optimization methods.

Problem

• Addresses convergence issues in federated learning due to extreme data heterogeneity
• Proposes personalized subnetworks during warmup to handle conflicting updates
• Enhances accuracy and convergence speed in heterogeneous federated learning

Innovation

• Personalized warmup phase via subnetworks
• Focus on learning specific subnetworks initially
• Revert to standard federated optimization after warmup