Covariances for Free: Exploiting Mean Distributions for Federated Learning with Pre-Trained Models

Goswami, Dipam, Magistri, Simone, Wang, Kai, Twardowski, Bartłomiej, Bagdanov, Andrew D., van de Weijer, Joost

arXiv.org Artificial Intelligence 

Using pre-trained models has been found to reduce the effect of data heterogeneity and speed up federated learning algorithms. Recent works have investigated the use of first-order statistics and second-order statistics to aggregate local client data distributions at the server and achieve very high performance without any training. In this work we propose a training-free method based on an unbiased estimator of class covariance matrices. Our method, which only uses first-order statistics in the form of class means communicated by clients to the server, incurs only a fraction of the communication costs required by methods based on communicating second-order statistics. We show how these estimated class covariances can be used to initialize a linear classifier, thus exploiting the covariances without actually sharing them. When compared to state-of-the-art methods which also share only class means, our approach improves performance in the range of 4-26% with exactly the same communication cost. Moreover, our method achieves performance competitive or superior to sharing second-order statistics with dramatically less communication overhead. Federated learning (FL) is a widely used paradigm for distributed learning from multiple clients or participants. In FL, each client trains their local model on their private data and then send model updates to a common global server that aggregates this information into a global model. The objective is to learn a global model that performs similarly to a model jointly trained on all the client data. A major concern in existing federated optimization algorithms (McMahan et al., 2017) is the poor performance when the client data is not identically and independently distributed (iid) or when classes are imbalanced between clients (Zhao et al., 2018; Li et al., 2019; Acar et al., 2021; Karimireddy et al., 2020a). Luo et al. (2021) showed that client drift in FL is mainly due to the drift in client classifiers which optimize to the local data distribution, resulting in forgetting knowledge from other clients from previous rounds (Legate et al., 2023b; Caldarola et al., 2022).