Model-heterogeneous Federated Learning with Partial Model Training

Federated Learning (FL) enables a large number of resource-limited devices to train a model collaboratively without data sharing. However, many works focus on model-homogeneous FL where the global and local models are the same size and ignore the inherently heterogeneous computation capabilities on different devices, which restricts resource-constrained devices from contributing to FL. In this paper, we consider model-heterogeneous FL and propose Federated Partial Model Training (FedPMT), where devices with small computation capabilities work on partial models (subsets of the global model) and contribute to the global model. Different from Dropout based partial model generation that removes neurons in hidden layers at random, model training in FedPMT is achieved from the back-propagation perspective. As such, all devices in FedPMT prioritize the most crucial parts of the global model. Empirical results show that FedPMT outperforms the existing benchmark FedDrop to a large extent. Meanwhile, compared to the popular model-homogeneous benchmark, FedAvg, FedPMT reaches the learning target with a shorter completion time and achieves a better trade-off between learning accuracy and completion time.