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

2026-04-07

AI Summary

This work addresses the challenges of system and statistical heterogeneity in federated learning, which often lead to inefficient models, insufficient personalization, and parameter divergence among submodels—issues exacerbated by existing pruning methods that struggle to balance server-side deployability with client-side resource constraints. To overcome these limitations, the paper proposes SubFLOT, a novel framework that, for the first time, integrates optimal transport theory into federated pruning. By minimizing Wasserstein distance, SubFLOT enables the server to generate personalized submodels aligned with each client’s data distribution. Additionally, it introduces a sparsity-aware adaptive regularization mechanism that effectively mitigates parameter divergence. Extensive experiments demonstrate that SubFLOT consistently outperforms state-of-the-art methods across diverse heterogeneous settings, significantly enhancing both training stability and performance of efficient personalized models on edge devices.

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
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Federated Learning
• Model Pruning
• Personalization
• System Heterogeneity
• Parametric Divergence

Innovation

• Optimal Transport
• Federated Pruning
• Personalized Federated Learning
• Wasserstein Distance
• Model Heterogeneity