WAFFLE: Weighted Averaging for Personalized Federated Learning

In federated learning, model personalization can be a very effective strategy to deal with statistical heterogeneity across clients. We introduce WAFFLE (Weighted Averaging For Federated LEarning): a personalized collaborative machine learning algorithm based on SCAFFOLD. SCAFFOLD uses stochastic control variates to converge towards a model close to the globally optimal model even in classification tasks where the marginal distribution of labels across clients is highly skewed. However, WAFFLE uses the Euclidean distance between clients’ updates to weigh their contributions and thus minimize the trained model’s loss on one specific agent. Through a series of experiments, we compare our proposed new method to two recent personalized federated learning methods, Weight Erosion and APFL, as well as two global methods, federated averaging and SCAFFOLD. We evaluate our method using two categories of non-identical client distributions (concept shift and label skew) on two benchmarked image data sets, MNIST and CIFAR10. Our experiments demonstrate the effectiveness of WAFFLE compared with other methods, as it achieves or improves accuracy with faster convergence.