🤖 machine learning

Federated Learning with Uncertainty and Personalization via Efficient Second-order Optimization

This paper proposes a novel, computationally efficient second-order optimization method for Federated Learning that achieves the uncertainty quantification and personalization benefits of Bayesian approaches while significantly outperforming existing state-of-the-art methods in accuracy and efficiency.

Original authors: Shivam Pal, Aishwarya Gupta, Saqib Sarwar, Piyush Rai

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Shivam Pal, Aishwarya Gupta, Saqib Sarwar, Piyush Rai

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a group of students (the clients) who all live in different houses and have their own unique sets of homework problems. They want to learn a subject together, but there's a strict rule: no one can leave their house or share their actual homework papers. They can only send their answers to a teacher (the server) who sits in the middle.

This is Federated Learning. The teacher collects the answers, averages them out to create a "Master Guide," and sends it back to the students.

The Problem with the Old Way

Usually, the teacher just takes a simple average of the answers (like FedAvg). This works great if everyone has similar homework. But what if Student A only has math problems, Student B only has history, and Student C has a mix? A single "Master Guide" won't help anyone much because it's a blurry mix of everything.

Also, the old methods don't tell you how sure they are about an answer. If a student guesses, the teacher doesn't know if it's a lucky guess or a confident calculation.

Some researchers tried to fix this using Bayesian Learning. Instead of just sending one answer, students would send a whole "cloud of possibilities" (a probability distribution) to show how confident they are. This is great for personalization and knowing when to be unsure. But, calculating and sending these "clouds" is like trying to mail a library of books instead of a single postcard. It's too heavy, too slow, and too expensive for students with weak internet or old computers.

The New Solution: FedIvon

The authors of this paper created a new method called FedIvon. Think of it as a clever way to send a "smart summary" that feels like a heavy library but weighs as little as a postcard.

Here is how it works, using a simple analogy:

1. The "Second-Order" Shortcut (The Smart Map)

Imagine you are hiking down a mountain to find the lowest valley (the best answer).

Old methods (like Adam) are like a hiker who only looks at the slope directly under their feet. They take small steps based on how steep the ground is right now. It works, but it can be slow and might get stuck in small dips.
Traditional Bayesian methods try to map the entire mountain range to understand the shape of the valley perfectly. This is accurate but takes forever to draw the map.
FedIvon is like a hiker who uses a smart compass. It doesn't map the whole mountain, but it quickly estimates the curvature of the ground (is it a sharp cliff or a gentle slope?) using a trick called IVON. This allows it to take bigger, smarter steps without needing to draw the whole map. It gets the benefits of the "full map" (uncertainty and precision) without the heavy workload.

2. The "Personalized" Prior (The Teacher's Hint)

In this new system, the teacher sends back a "hint" (a prior distribution) based on what the whole class knows.

When a student learns, they start with the teacher's hint but then adjust it heavily based on their own specific homework.
If a student has very little homework (data), they rely more on the teacher's hint.
If a student has lots of unique homework, they trust their own data more.
This creates a Personalized Model for every student that is still connected to the group.

3. The Efficient Exchange

Instead of sending a heavy "cloud" of data, FedIvon sends just two small numbers for every part of the model:

The best guess (the average).
The confidence level (how much the answer might vary).

Because the math behind FedIvon is so efficient (it calculates these numbers implicitly while learning, rather than doing a separate, expensive calculation), it runs almost as fast as the simple, non-Bayesian methods.

What Did They Find?

The authors tested this on three different "homework sets" (datasets):

EMNIST (handwritten letters/numbers).
SVHN (house numbers from street signs).
CIFAR-10 (pictures of objects like cats, cars, etc.).

They simulated a scenario where 200 students each had very few examples (less than 100) and very different types of data.

The Results:

Better Accuracy: FedIvon got higher scores on tests than the old "average" methods and even better than other complex Bayesian methods.
Better Confidence: It was much better at knowing when it was unsure. In a test where they tried to spot "strange" images (Out-of-Distribution), FedIvon was the best at saying, "I don't know what this is," rather than guessing wrong confidently.
Speed: Despite doing all this fancy math, it didn't slow down the process. It was as fast as the simple methods.

The Bottom Line

FedIvon is a new way for computers to learn together without sharing private data. It manages to be smart (knowing how sure it is), personal (adapting to each user's unique data), and fast (not slowing down the network). It solves the problem of "Bayesian learning is too heavy" by using a clever mathematical shortcut that makes the heavy lifting feel light.

Technical Summary: Federated Learning with Uncertainty and Personalization via Efficient Second-order Optimization

Problem Statement

Federated Learning (FL) enables collaborative model training across decentralized clients without sharing raw data. However, standard FL approaches face two primary challenges:

Data Heterogeneity: When client data is non-IID (non-independent and identically distributed), a single global model often fails to represent individual client distributions effectively, leading to suboptimal performance.
Lack of Uncertainty Quantification: Standard FL methods typically learn point estimates of model parameters. This makes them susceptible to overfitting in low-data regimes and unable to provide predictive uncertainty, which is crucial for tasks like Out-of-Distribution (OOD) detection.

While Bayesian FL offers a principled solution by learning posterior distributions to capture uncertainty and enable personalization (via hierarchical frameworks), existing methods suffer from significant bottlenecks. Current Bayesian FL approaches often rely on computationally expensive routines (e.g., MCMC sampling, Laplace's approximation requiring explicit Hessian computation, or deep ensembles) and incur high communication costs by transmitting full posterior distributions. These factors limit scalability, particularly for resource-constrained clients.

Methodology: FedIvon

The authors propose FedIvon, a novel Bayesian FL algorithm that leverages Improved Variational Online Newton (IVON) to perform efficient second-order variational inference. The method balances the benefits of Bayesian inference with computational and communication efficiency comparable to first-order optimizers like Adam.

Core Components

Client-Side Variational Inference (IVON):
- Instead of learning a point estimate, each client approximates its local posterior distribution $q_k(\theta)$ using a Gaussian with a diagonal covariance matrix: $N(\theta | m_k, \text{diag}(\sigma_k^2))$ .
- Second-Order Efficiency: Unlike Adam, which adjusts learning rates based on squared gradients, IVON uses a reparameterization to implicitly estimate the diagonal Hessian ( $h$ ) without explicit computation. The variance is derived as $\sigma^2 = 1/(h + \delta)$ .
- Optimization: The algorithm updates both the mean vector ( $m$ ) and the variances ( $\sigma^2$ ) using natural gradients. This provides uncertainty estimates naturally and captures the geometry of the loss function for faster convergence, all while maintaining computational costs similar to Adam.
Server-Side Posterior Aggregation:
- The server aggregates client posteriors to compute a global posterior. Since the product of Gaussian distributions is also Gaussian, the global posterior is approximated as a Gaussian $N(\theta | \mu, \Lambda^{-1})$ .
- Aggregation Strategy: The global mean ( $\mu$ ) and precision matrix ( $\Lambda$ ) are computed as weighted sums of local means and precisions, respectively, where weights are proportional to client dataset sizes ( $N_k$ ). This strategy is computationally efficient as it only requires transmitting the mean and diagonal variance (or Hessian) rather than full model weights or complex distributions.
Personalization via Hierarchical Bayesian Framework:
- FedIvon naturally supports personalized FL. The global posterior from the server serves as a prior ( $p_k(\theta)$ ) for the client's local optimization.
- The client minimizes a personalized loss function: $L_k(q) = \mathbb{E}_{q_k}[\bar{\ell}_k(\theta)] + \beta D_{KL}(q_k(\theta) \| p_k(\theta))$ .
- The hyperparameter $\beta$ controls the trade-off between local adaptation and global knowledge. When $\beta=0$ , the model is fully personalized; higher values incorporate more global information.

Key Contributions

Efficient Second-Order Bayesian FL: Introduction of FedIvon, which utilizes IVON to perform variational inference with computational costs similar to first-order methods (like Adam) but provides second-order information (Hessian estimates) implicitly.
Improved Performance and Uncertainty: Demonstration that FedIvon achieves superior predictive accuracy and better uncertainty estimates compared to state-of-the-art (SOTA) Bayesian and non-Bayesian FL methods.
Native Personalization: A framework that naturally enables client-level model personalization by leveraging the server's posterior as a prior, effectively balancing local data characteristics with global shared knowledge.

Experimental Results

The authors evaluated FedIvon on three datasets (EMNIST, SVHN, CIFAR-10) under non-IID settings with 200 clients, each having limited data (<100 examples).

Standard FL Results

Accuracy & Calibration: FedIvon outperformed baselines (FedAvg, FedLaplace) in test accuracy, Expected Calibration Error (ECE), Negative Log Likelihood (NLL), and Brier Score across all datasets.
OOD Detection: In Out-of-Distribution detection tasks, FedIvon achieved competitive or superior AUROC scores compared to baselines, validating its ability to quantify predictive uncertainty effectively.
Convergence: While FedIvon showed slightly slower initial convergence due to stochastic sampling of weights for gradient estimation, it surpassed other methods as training progressed due to improved Hessian approximation.

Personalized FL Results

Experiments on CIFAR-10 (class distribution skew) and CIFAR-100 (class concept drift) compared FedIvon against pFedME, pFedBayes, and pFedVEM.

CIFAR-10: FedIvon achieved client accuracy comparable to pFedVEM but demonstrated significant improvements in global server accuracy, indicating better aggregation of heterogeneous data.
CIFAR-100: FedIvon showed significant improvements over all baselines in both personalized client accuracy and global server accuracy, highlighting its robustness in handling concept drift.

Significance and Claims

The paper claims that FedIvon bridges the gap between the theoretical benefits of Bayesian FL (uncertainty quantification, personalization) and the practical constraints of federated systems (computation, communication). By using an efficient second-order optimization approach, the method avoids the heavy computational overhead of MCMC or explicit Hessian calculations found in prior works.

The authors assert that FedIvon is the first method to offer:

Computational efficiency comparable to first-order optimizers (like Adam).
Full Bayesian benefits (uncertainty and personalization) without the scalability issues of existing Bayesian FL methods.
Superior performance in both standard and personalized FL settings, particularly in scenarios with high data heterogeneity and limited client data.

The work positions FedIvon as a scalable, practical solution for deploying robust, uncertainty-aware, and personalized machine learning models in real-world federated environments.