Probabilistic Federated Learning on Uncertain and Heterogeneous Data with Model Personalization

2026-03-18

AI Summary

This work addresses the performance degradation of conventional federated learning under client data heterogeneity and uncertainty. To this end, we propose Meta-BayFL, a unified probabilistic personalized federated learning framework that, for the first time, integrates Bayesian neural networks with meta-learning. Our approach models inter-layer uncertainty via Bayesian neural networks, incorporates a meta-learning mechanism with adaptive learning rates, and refines the global model aggregation strategy. Theoretical analysis establishes the convergence of the proposed method, and extensive experiments demonstrate its superiority over state-of-the-art approaches such as pFedMe and Ditto, achieving up to a 7.42% accuracy gain on CIFAR-10, CIFAR-100, and Tiny-ImageNet. Furthermore, we validate the practical feasibility of deploying Meta-BayFL on edge devices.

Technology Category

• Machine Learning: Bayesian Learning
• Reasoning under Uncertainty: Probabilistic Programming
• Multiagent Systems: Multiagent Learning

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Conventional federated learning (FL) frameworks often suffer from training degradation due to data uncertainty and heterogeneity across local clients. Probabilistic approaches such as Bayesian neural networks (BNNs) can mitigate this issue by explicitly modeling uncertainty, but they introduce additional runtime, latency, and bandwidth overhead that has rarely been studied in federated settings. To address these challenges, we propose Meta-BayFL, a personalized probabilistic FL method that combines meta-learning with BNNs to improve training under uncertain and heterogeneous data. The framework is characterized by three main features: (1) BNN-based client models incorporate uncertainty across hidden layers to stabilize training on small and noisy datasets, (2) meta-learning with adaptive learning rates enables personalized updates that enhance local training under non-IID conditions, and (3) a unified probabilistic and personalized design improves the robustness of global model aggregation. We provide a theoretical convergence analysis and characterize the upper bound of the global model over communication rounds. In addition, we evaluate computational costs (runtime, latency, and communication) and discuss the feasibility of deployment on resource-constrained devices such as edge nodes and IoT systems. Extensive experiments on CIFAR-10, CIFAR-100, and Tiny-ImageNet show that Meta-BayFL consistently outperforms state-of-the-art methods, including both standard and personalized FL approaches (e.g., pFedMe, Ditto, FedFomo), with up to 7.42\% higher test accuracy.

Problem

• federated learning
• data uncertainty
• data heterogeneity
• non-IID data
• model personalization

Innovation

• Bayesian Neural Networks
• Federated Learning
• Meta-Learning
• Model Personalization
• Data Heterogeneity