Bayesian Neural Network For Personalized Federated Learning Parameter Selection
Mengen Luo, Ercan Engin Kuruoglu
TL;DR
This work tackles non-IID data in federated learning by introducing element-wise personalization guided by Bayesian uncertainty. It uses a diagonal Gaussian posterior estimated via Laplace approximation to identify parameters with high uncertainty as candidates for personalization, applying a KL-divergence-based aggregation to reconcile global and client distributions. Personalization is implemented through a mask that selects a subset of parameters to personalize based on a threshold on parameter variance, enabling fine-grained customization. Empirical results on MNIST, Fashion-MNIST, and CIFAR-10 show the approach often surpasses layer-based baselines, with performance gains varying by dataset and model depth. The method offers a principled, scalable path to personalized federated learning in heterogeneous environments.
Abstract
Federated learning's poor performance in the presence of heterogeneous data remains one of the most pressing issues in the field. Personalized federated learning departs from the conventional paradigm in which all clients employ the same model, instead striving to discover an individualized model for each client to address the heterogeneity in the data. One of such approach involves personalizing specific layers of neural networks. However, prior endeavors have not provided a dependable rationale, and some have selected personalized layers that are entirely distinct and conflicting. In this work, we take a step further by proposing personalization at the elemental level, rather than the traditional layer-level personalization. To select personalized parameters, we introduce Bayesian neural networks and rely on the uncertainty they offer to guide our selection of personalized parameters. Finally, we validate our algorithm's efficacy on several real-world datasets, demonstrating that our proposed approach outperforms existing baselines.