Gradient Rotation Unit for Non-I.I.D. Federated Learning

Federated Learning (FL) enables collaborative training of a global model without exposing raw data by aggregating local updates from clients. However, the convergence efficiency on non-i.i.d. data remains challenging, leading to performance loss and resource bottleneck. Meanwhile, the nature of non-i.i.d. challenge are not yet fully understood. In this paper, we first reveal that non-i.i.d. data leads to server-side multi-objective aggregation conflict challenge which hampers the convergence efficiency. We then propose Federated Gradient Rotation Unit (FGRU), a simple yet general approach to mitigate this challenge by deliberately aligning optimization trajectories across clients. FGRU is a server-side plugin that rotates gradients to other before aggregation. On a series of challenging non-i.i.d. FL tasks, FGRU leads to significant gains in convergence efficiency and performance. Experimental results demonstrate that FGRU improves inference accuracy by approximately 3-5% and accelerates convergence by 2-3X on simulated non-i.i.d. data using MNIST, CIFAR-10, and Fashion-MNIST datasets. The results also show that FGRU is model agnostic and can be combined with existing non-i.i.d. FL frameworks such as SCAFFOLD and FedProx to further improve performance.