Probabilistic Fusion of Neural Networks that Incorporates Global Information

As one of the approaches in Federated Learning, model fusion distills models trained on local clients into a global model. The previous method, Probabilistic Federated Neural Matching (PFNM), can match and fuse local neural networks with varying global model sizes and data heterogeneity using the Bayesian nonparametric framework. However, the alternating optimization process applied by PFNM causes absence of global neuron information. In this paper, we propose a new method that extends PFNM by introducing a Kullback-Leibler (KL) divergence penalty, so that it can exploit information in both local and global neurons. We show theoretically that the extended PFNM with a penalty derived from KL divergence can fix the drawback of PFNM by making a balance between Euclidean distance and the prior probability of neurons. Experiments on deep fully-connected as well as deep convolutional neural networks demonstrate that our new method outperforms popular state-of-the-art federated learning methods in both image classification and semantic segmentation tasks.