On the Convergence and Stability of Distributed Sub-model Training

TL;DR

This paper introduces a distributed shuffled sub-model training method for federated learning, demonstrating improved convergence and generalization stability through theoretical analysis and extensive experiments.

Contribution

It proposes a novel shuffled sub-model training approach with convergence guarantees and stability-based generalization improvements.

Findings

The algorithm converges at a proven rate.

Shuffling enhances training stability and generalization.

Experimental results validate theoretical claims.

Abstract

As learning models continue to grow in size, enabling on-device local training of these models has emerged as a critical challenge in federated learning. A popular solution is sub-model training, where the server only distributes randomly sampled sub-models to the edge clients, and clients only update these small models. However, those random sampling of sub-models may not give satisfying convergence performance. In this paper, observing the success of SGD with shuffling, we propose a distributed shuffled sub-model training, where the full model is partitioned into several sub-models in advance, and the server shuffles those sub-models, sends each of them to clients at each round, and by the end of local updating period, clients send back the updated sub-models, and server averages them. We establish the convergence rate of this algorithm. We also study the generalization of distributed sub-model training via stability analysis, and find that the sub-model training can improve the generalization via amplifying the stability of training process. The extensive experiments also validate our theoretical findings.