Federated Optimization with Linear-Time Approximated Hessian Diagonal

Mrinmay Sen, C. Krishna Mohan, A. Kai Qin

Abstract: Federated learning (FL) or federated optimization is a type of distributed optimization where multiple clients collaboratively train a global model without sharing local data. One of the key challenge in FL is the communication overhead due to slow convergence of the global model. In this paper, we propose a federated learning algorithm to handle this slow convergence by incorporating Hessian diagonal while training client’s models. To reduce the computational and memory complexity in local clients, we introduce a linear time Hessian diagonal approximation technique by using only the first row of the Hessian. Our extensive experiments show that our proposed method outperforms state-of-the-art FL algorithms, FedAvg, FedProx, SCAFFOLD and DONE in terms of training loss, test loss and test accuracy.