SubFLOT: Submodel Extraction for Efficient and Personalized Federated Learning via Optimal Transport

TL;DR

SubFLOT introduces a novel server-side personalized federated pruning framework using optimal transport, enabling efficient, customized models for resource-limited devices while maintaining training stability.

Contribution

It proposes OTP and SAR modules for personalized pruning and divergence mitigation, advancing federated learning efficiency and personalization without raw data access.

Findings

SubFLOT outperforms existing methods in accuracy and efficiency.

The OTP module effectively generates personalized submodels.

SAR stabilizes training by penalizing divergence based on pruning rate.

Abstract

Federated Learning (FL) enables collaborative model training while preserving data privacy, but its practical deployment is hampered by system and statistical heterogeneity. While federated network pruning offers a path to mitigate these issues, existing methods face a critical dilemma: server-side pruning lacks personalization, whereas client-side pruning is computationally prohibitive for resource-constrained devices. Furthermore, the pruning process itself induces significant parametric divergence among heterogeneous submodels, destabilizing training and hindering global convergence. To address these challenges, we propose SubFLOT, a novel framework for server-side personalized federated pruning. SubFLOT introduces an Optimal Transport-enhanced Pruning (OTP) module that treats historical client models as proxies for local data distributions, formulating the pruning task as a Wasserstein distance minimization problem to generate customized submodels without accessing raw data. Concurrently, to counteract parametric divergence, our Scaling-based Adaptive Regularization (SAR) module adaptively penalizes a submodel's deviation from the global model, with the penalty's strength scaled by the client's pruning rate. Comprehensive experiments demonstrate that SubFLOT consistently and substantially outperforms state-of-the-art methods, underscoring its potential for deploying efficient and personalized models on resource-constrained edge devices.