MoFedAGR: Mitigating client drift with adaptive gradient regularization and global momentum in federated learning

Federated learning is a novel distributed machine learning framework with privacy-protection, yet it is vulnerable to the effects of heterogeneous data. Heterogeneous data drive client models that overfit local datasets and depart from the global optimum during local training, which is termed client drift. To address the impact of client drift, we approach this issue from the perspectives of optimization and generalization. We comprehensively considering the effects of client drift during the training process, and quantifying it as the aggregation error. We first propose adaptive gradient regularization, which is based on gradient regularization and further and applies different regularization strengths to each parameter based on the magnitude of the parameter variance between the local model and the global model, thereby mitigating the performance degradation caused by aggregation error and helping model converge to a flatter minimum. In order to obtain the variance between local and global models to compute our adaptive gradient regularization term, we introduce global momentum from the server side as the approximation of global gradient and further utilize it as a gradient correction term. Next, we propose MoFedAGR, which combines gradient correction term and adaptive gradient regularization term, helping client models converge to a consistent flat minimum. We have provided the theoretical convergence bounds of the algorithm we proposed. Furthermore, experiments on several image classification datasets demonstrate that our algorithm significantly improves model performance while exhibiting strong generalization capabilities.