Non-Convex Optimization in Federated Learning via Variance Reduction and Adaptive Learning

This paper proposes a novel federated algorithm that leverages momentum-basedvariance reduction with adaptive learning to address non-convex settings acrossheterogeneous data. We intend to minimize communication and computationoverhead, thereby fostering a sustainable federated learning system. We aim toovercome challenges related to gradient variance, which hinders the model'sefficiency, and the slow convergence resulting from learning rate adjustmentswith heterogeneous data. The experimental results on the image classificationtasks with heterogeneous data reveal the effectiveness of our suggestedalgorithms in non-convex settings with an improved communication complexity of$\mathcal{O}(\epsilon^{-1})$ to converge to an $\epsilon$-stationary point -compared to the existing communication complexity $\mathcal{O}(\epsilon^{-2})$of most prior works. The proposed federated version maintains the trade-offbetween the convergence rate, number of communication rounds, and test accuracywhile mitigating the client drift in heterogeneous settings. The experimentalresults demonstrate the efficiency of our algorithms in image classificationtasks (MNIST, CIFAR-10) with heterogeneous data.