Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Probabilistic Federated Learning on Uncertain and Heterogeneous Data with Model Personalization

Ratun Rahman, Dinh C. Nguyen

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.

Probabilistic Federated Learning on Uncertain and Heterogeneous Data with Model Personalization

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.
Paper Structure (24 sections, 9 theorems, 82 equations, 3 figures, 6 tables, 1 algorithm)

This paper contains 24 sections, 9 theorems, 82 equations, 3 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Works
  3. Our Key Contributions
  4. Paper Structure
  5. Preliminaries
  6. Bayesian Neural Networks (BNNs)
  7. Federated Learning
  8. Personalized Federated Learning
  9. Proposed Probabilistic FL Framework with Model Personalization
  10. Model Training Concept
  11. Convergence Analysis
  12. Detailed Proofs of Convergence
  13. Proof of Lemma \ref{['lemma2']}
  14. Proof of Lemma \ref{['lemma3']}
  15. Proof of Lemma \ref{['lemma5']}
  16. ...and 9 more sections

Key Result

Lemma 1

Under Assumption assump:4, there exist a constant $C_1 > 0$ such that for any $\delta > 0$, with probability $1-\delta$ we have

Figures (3)

  • Figure 1: A simple probabilistic model structure in which probabilistic properties link every layer to every layer after it. The likelihood graph shows the likelihood of moving from the current layer to the next. It shows that probabilistic (distributional) parameters spread uncertainty through depth.
  • Figure 2: Our proposed Meta-BayFL FL algorithm where there are $K$ clients connected to the global server with adaptive model aggregation using a personalized learning approach. Each client creates local models using BNN and their local data and different temporary learning rates, and chooses the best rates for local model training. The server aggregates the local models and distributes the updated global model to all the clients.
  • Figure 3: Performance comparison on the CIFAR-10 non-IID dataset between basic CNN, FedAVG (FL with CNN), and BayFL (FL with BNN) approaches. A total of 5 clients with various data distributions and batch sizes are used for FL approaches with 10 local epochs in every global round.

Theorems & Definitions (11)

  • Lemma 1
  • Lemma 2
  • Lemma 3
  • Lemma 4
  • Lemma 5
  • Lemma 6
  • Lemma 7
  • Theorem 1
  • Remark 1
  • Theorem 2
  • ...and 1 more