Empowering Federated Learning: Tackling Client Drift and Resource Constraints with Knowledge Distillation by Combining Losses

Federated Learning is an emerging approach to Machine Learning which allows for decentralised model training which safeguards privacy. Its potential applications, particularly in Medicine, Smart Manufacturing, Finance, and the Internet of Things, hold significant promise. However, it faces hurdles due to resource constraints and the diverse nature of data and devices at the client end. This paper highlights the critical challenge of client drift and its effects on Machine Learning model performance across various architectural configurations. Furthermore, our findings reveal that the use of pretrained models such as ResNet offers a compelling solution to mitigate the impact of client drift to some extent. Nonetheless, it is worth noting that leveraging pretrained models necessitates substantial client-side resources. In response to the dual challenges of client drift and resource constraints, we propose an innovative approach involving Knowledge Distillation, namely combining distillation loss and classification loss while using knowledge distillation at the client. Here, the teacher model is trained on a more compact dataset, while the student model undertakes training on a larger, more diverse dataset. This approach not only improves robustness but also enhances privacy. The outcomes of our experiments substantiate the efficacy of this technique, showcasing an approximate improvement of 50% in the accuracy and loss of the student model.