Robust federated learning under statistical heterogeneity via Hessian spectral decomposition

Abstract

Federated Learning (FL) is a collaborative machine learning paradigm in which a global model is learned via aggregating local ones. Although statistical heterogeneity of the local training data is necessary for the generalisability of the global model, it also introduces local model “drift” that slows down the convergence. Thus, how to optimally aggregate local models in FL remains an open problem. Recognising that training data lends varying evidential credence to different parts of a local model, we propose a novel approach to exploit such evidential asymmetry in FL aggregation in not independent and identically distributed (non-IID) data by applying a unique weight coefficient to each of the local parameter updates. To this end, we measure the parameter-level evidential credence making use of the eigenvalues of the Hessian of the local likelihood function, which are theoretically connected to the observed Fisher information. We employ these eigenpairs to propose a novel aggregation method, which we name FedHess. Our experiments show FedHess achieves smoother and faster convergence to a more accurate global model when compared with popular baselines such as Federated Average (FedAvg), FedProx, SCAFFOLD, Federated Curvature (FedCurv) and FedDF across different types of heterogeneous training data drawn from a number of benchmark datasets.