FedSPU: Personalized Federated Learning for Resource-constrained Devices with Stochastic Parameter Update

TL;DR

FedSPU addresses personalization in personalized federated learning for resource-constrained IoT devices facing non-iid data. It maintains full local model architectures on clients while randomly freezing a subset of neurons, so only active parameters are updated and aggregated, which preserves personalization and reduces backpropagation cost. The authors prove convergence under standard assumptions and show that, with an early stopping mechanism using $L_t = \lambda L_{train} + (1-\lambda) L_{test}$, FedSPU achieves competitive or superior accuracy while significantly lowering training time and memory usage across EMNIST, CIFAR10, and Google Speech datasets. Empirically, FedSPU outperforms dropout baselines by about $4.45\%$ on average in final accuracy, with memory savings up to $54\%$ and substantial reductions in training rounds when combined with early stopping.

Abstract

Personalized Federated Learning (PFL) is widely employed in IoT applications to handle high-volume, non-iid client data while ensuring data privacy. However, heterogeneous edge devices owned by clients may impose varying degrees of resource constraints, causing computation and communication bottlenecks for PFL. Federated Dropout has emerged as a popular strategy to address this challenge, wherein only a subset of the global model, i.e. a sub-model, is trained on a client's device, thereby reducing computation and communication overheads. Nevertheless, the dropout-based model-pruning strategy may introduce bias, particularly towards non-iid local data. When biased sub-models absorb highly divergent parameters from other clients, performance degradation becomes inevitable. In response, we propose federated learning with stochastic parameter update (FedSPU). Unlike dropout that tailors the global model to small-size local sub-models, FedSPU maintains the full model architecture on each device but randomly freezes a certain percentage of neurons in the local model during training while updating the remaining neurons. This approach ensures that a portion of the local model remains personalized, thereby enhancing the model's robustness against biased parameters from other clients. Experimental results demonstrate that FedSPU outperforms federated dropout by 7.57% on average in terms of accuracy. Furthermore, an introduced early stopping scheme leads to a significant reduction of the training time by 24.8%-70.4% while maintaining high accuracy.

Figures (5)

• Figure 1: In dropout (a), clients train sub-models with fewer parameters. In FedSPU (b), clients train full models with partial parameters frozen.
• Figure 2: Demonstration of the FedSPU framework.
• Figure 3: The aggregation scheme in FedSPU.
• Figure 4: Comparison between replacing all parameters of a local sub-model and replacing a portion of the parameters for a local full model during global communication.
• Figure 5: Comparison of memory footprint (MB) with different $p_{k}$. "Others" stand for FedMP, Hermes, PruneFL and FedSelect.