A Remedy for Heterogeneous Data: Clustered Federated Learning with Gradient Trajectory
Ruiqi Liu, Songcan Yu, Linsi Lan, Junbo Wang, Krishna Kant, Neville Calleja
Abstract
Federated Learning (FL) has recently attracted a lot of attention due to its ability to train a machine learning model using data from multiple clients without divulging their privacy. However, the training data across clients can be very heterogeneous in terms of quality, amount, occurrences of specific features, etc. In this paper, we demonstrate how the server can observe data heterogeneity by mining gradient trajectories that the clients compute from a two-dimensional mapping of high-dimensional gradients computed by each client from its bottom layer. Based on these ideas, we propose a new clustered federated learning with gradient trajectory method, called CFLGT, which dynamically clusters clients together based on the gradient trajectories. We analyze CFLGT both theoretically and experimentally to show that it overcomes several drawbacks of mainstream Clustered Federated Learning (CFL) methods and outperforms other baselines.