FedGSE: Gradient-Based Submodel Extraction for Resource-Constrained Federated Learning

Federated Learning (FL) has emerged as a pivotal paradigm for multi-client collaborative learning, primarily due to its inherent capability to safeguard privacy. Nonetheless, the heterogeneity among FL clients, characterized by their disparate resource capabilities and not independent and identical (NonIID) local datasets, presents a significant challenge. Specifically, low-resource clients, such as edge devices, grapple with the inadequacy to accommodate the entire model parameter set for training purposes. To mitigate this issue, preceding research has ventured into devising methodologies that entail extracting sub-models from the overarching global model, tailored to the specific communication, computational, and memory constraints of individual clients. Despite these advancements, prevailing sub-model extraction techniques, which predominantly hinge on pre-established rules, overlook a crucial factor: the impact of NonIID local data on the trajectory of neuron update dynamics. This oversight can amplify discrepancies between the practical local updates inferred by the sub-model and those anticipated via the entire model, thereby undermining overall performance. In this paper, we introduce FedGSE, an innovative Gradient-based Neuron Selection methodology designed explicitly for FL environments. This methodology aims to curate sub-models that significantly reduce discrepancies in local updates, enhancing alignment with the global model's learning trajectory. Central to the FedGSE approach is a sophisticated algorithm that handpicks critical neurons for sub-model construction. These neurons are identified through their pronounced gradient magnitudes, resulting from the training of the global model on a dataset mirroring the client's data distribution. Consequently, the sub-model's induced local gradient updates closely emulate those derived from directly training the client's data on the full global model, fostering enhanced alignment and performance. Extensive experiments over diverse datasets and tasks demonstrate the superiority of FedGSE over existing baselines.