Federated Hybrid Model Pruning through Loss Landscape Exploration

TL;DR

AutoFLIP tackles the challenge of efficient federated learning under non-IID data by introducing a federated loss exploration phase that informs a hybrid pruning strategy combining unstructured and structured pruning. The global pruning mask, updated each FL round, reduces both communication and computation while aligning client updates to improve convergence. Across multiple datasets, architectures, and non-IID partitions, AutoFLIP achieves substantial FLOPs and communication reductions while maintaining or boosting global accuracy relative to FedAvg and state-of-the-art pruning baselines. The approach offers practical gains for deploying scalable, resource-efficient FL in edge environments, with future work focusing on personalization, privacy, and multi-server settings.

Abstract

As the era of connectivity and unprecedented data generation expands, collaborative intelligence emerges as a key driver for machine learning, encouraging global-scale model development. Federated learning (FL) stands at the heart of this transformation, enabling distributed systems to work collectively on complex tasks while respecting strict constraints on privacy and security. Despite its vast potential, specially in the age of complex models, FL encounters challenges such as elevated communication costs, computational constraints, and the heterogeneous data distributions. In this context, we present AutoFLIP, a novel framework that optimizes FL through an adaptive hybrid pruning approach, grounded in a federated loss exploration phase. By jointly analyzing diverse non-IID client loss landscapes, AutoFLIP efficiently identifies model substructures for pruning both at structured and unstructured levels. This targeted optimization fosters a symbiotic intelligence loop, reducing computational burdens and boosting model performance on resource-limited devices for a more inclusive and democratized model usage. Our extensive experiments across multiple datasets and FL tasks show that AutoFLIP delivers quantifiable benefits: a 48.8% reduction in computational overhead, a 35.5% decrease in communication costs, and a notable improvement in 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.

Figures (9)

• Figure 1: FL optimization process. At each communication round, participant clients perform a local update to then send the new parameters to the Server for the global aggregation.
• Figure 2: AutoFLIP pruning procedure. The local guidance matrices are computed a priori through the federated exploration phase. The global guidance matrix is computed by the server by aggregating the elements of the local guidance matrices corresponding to the participant clients. The pruning mask is received by the participant clients. All steps preliminary to the FL procedure are denoted in gray, while the steps intrinsic to the FL procedure with pruning are denoted in red.
• Figure 3: Hybrid pruning of both individual synapses (dotted lines) and entire structural units, e.g, neurons (highlighted nodes), based on the pruning mask $PG_{\text{global}}$ strategy.
• Figure 4: Variance reduction of weight updates $\sigma^2_{\Delta W}$ over FL rounds for AutoFLIP and FedAvg.
• Figure 5: Distribution of parameter deviations in $G_{\text{global}}$ after exploration. The absolute frequency (in log-scale) is shown for each normalized deviation. Higher frequencies are recorded for smaller deviation values, indicating that many parameters are irrelevant for loss improvement.