FedRPN: An Efficient Framework for Optimizing System Heterogeneity in Federated Learning

Abstract

Federated Learning (FL) enables machine learning tasks to be performed on distributed data in a privacy-preserving manner, but faces challenges related to the heterogeneity of device systems. This necessitates the customization of resource requirements to accommodate the diverse capacities of participating clients. However, existing approaches struggle to generate resource-customized, ready-to-use inference models while still incurring substantial resource consumption throughout the system workflow. In response, this paper introduces a novel framework, FedRPN, which incorporates Resource-customized Prototypical Networks. RPN leverages resource-customized pre-trained models as the initial point and employs unbiased proto-typical classification, enabling rapid convergence, resource efficiency, and robustness to non-IID data. Additionally, we propose a globally-aware training strategy that produces deployable inference models of varying capacities. Building on RPN, we propose a two-stage process comprising prototype construction followed by model fine-tuning, which further enhances performance. Experimental results demonstrate that FedRPN reduces computational and communication resource consumption by 48% and 43%, respectively, while delivering improved performance.