Adaptive Hybrid Model Pruning in Federated Learning through Loss Exploration

2024-05-16

AI Summary

To address the high communication overhead, client computational constraints, and slow convergence with degraded accuracy in federated learning—particularly under complex models and non-IID data—this paper proposes AutoFLIP. Methodologically, AutoFLIP introduces a novel knowledge distillation–driven pruning mechanism that models client-wise loss topology and gradient dynamics, enabling adaptive co-optimization of structured and unstructured pruning to automatically identify and compress low-contribution substructures. Crucially, lightweight pruning decisions are made locally during the federated loss exploration phase, requiring no additional communication or global model download. Experimental results demonstrate that AutoFLIP reduces average client computation cost by 48.8% and communication cost by 35.5%, accelerates global convergence, and improves both final model accuracy and distributional robustness.

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
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing

Abstract

The rapid proliferation of smart devices coupled with the advent of 6G networks has profoundly reshaped the domain of collaborative machine learning. Alongside growing privacy-security concerns in sensitive fields, these developments have positioned federated learning (FL) as a pivotal technology for decentralized model training. Despite its vast potential, specially in the age of complex foundation models, FL encounters challenges such as elevated communication costs, computational constraints, and the complexities of non-IID data distributions. We introduce AutoFLIP, an innovative approach that utilizes a federated loss exploration phase to drive adaptive hybrid pruning, operating in a structured and unstructured way. This innovative mechanism automatically identifies and prunes model substructure by distilling knowledge on model gradients behavior across different non-IID client losses topology, thereby optimizing computational efficiency and enhancing model performance on resource constrained scenarios. Extensive experiments on various datasets and FL tasks reveal that AutoFLIP not only efficiently accelerates global convergence, but also achieves superior accuracy and robustness compared to traditional methods. On average, AutoFLIP reduces computational overhead by 48.8% and communication costs by 35.5%, while improving global accuracy. By significantly reducing these overheads, AutoFLIP offer the way for efficient FL deployment in real-world applications for a scalable and broad applicability.

Problem

• Reduces communication costs in federated learning
• Optimizes model pruning for heterogeneous data distributions
• Improves computational efficiency in resource-limited devices

Innovation

• Adaptive hybrid pruning for federated learning optimization
• Federated loss exploration to identify pruning substructures
• Reduces computational and communication costs significantly