Research map / Federated learning and privacy / Personalized and heterogeneous FL

Improving Generalization in Federated Learning with Model-Data Mutual Information Regularization: A Posterior Inference Approach

NeurIPS 2024 · Federated learning and privacy › Personalized and heterogeneous FL · approach: Non-IID federated learning · 2 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes FedMDMI, a federated posterior inference method that regularizes a decomposed global model-data mutual information, samples local Gibbs posteriors with SGLD, and averages them at the server to improve generalization and calibration.

Problem

Point-estimate federated learning overfits scarce, heterogeneous client data and is overconfident; direct local Bayesian posterior inference under data heterogeneity yields biased local posteriors and a sub-optimal global posterior.

Approach

Add a global model-data mutual information regularizer, decompose it into local MI terms computed at clients, show the optimal local objective yields a Gibbs posterior sampled by SGLD, and approximate the global Gibbs posterior by averaging local posterior samples at the server, with the global model as a moving-average prior.

Key results

• FedMDMI reaches 83.42% on CIFAR-10 Dir(0.2)-L versus 82.63% for the best baseline FedEP (I), and 49.46% on CIFAR-100 Dir(0.2)-L versus 48.89%.
• FedMDMI attains the lowest ECE in most settings, e.g., 0.115 on CIFAR-10 Dir(0.2)-L versus 0.121 for FedEP.
• On Shakespeare non-iid-L, FedMDMI achieves 52.93% versus 51.39% for FedEP (I).

Measured effects reported in the paper

Surprising findings

• Gains over the best Bayesian baselines are modest (often under 1 point accuracy), while the FedMDMI convergence is slower than SCAFFOLD.
• β-PredBayes achieves lower ECE than FedMDMI in some CIFAR-100 settings, as reported in the appendix.

Limitations

• Assumes conditional independence of the global likelihood given w, which may fail in extreme FL settings.
• Convergence of SGLD for non-convex FL objectives is not theoretically established.
• Convergence is slower than optimization-based methods such as SCAFFOLD.

Experimental setup

Datasets: CIFAR-10 CIFAR-100 Shakespeare Stack Overflow

Benchmarks: CIFAR-10 CIFAR-100 Shakespeare Stack Overflow

Metrics: top-1 test accuracy expected calibration error train/test error

Baselines: FedAvg FedM MimeLite SCAFFOLD FedBE FedPA FedEP FALD β-PredBayes FedProx FedHB FedRep

Builds on: FedPA FedEP FALD SGLD FedAvg

Closest papers by meaning

• Calibrated One Round Federated Learning with Bayesian Inference in the Predictive Space (AAAI 2024 · similarity 0.90)
• Federated Generalised Variational Inference: A Robust Probabilistic Federated Learning Framework (ICML 2025 · similarity 0.89)
• FedSMU: Communication-Efficient and Generalization-Enhanced Federated Learning through Symbolic Model Updates (ICML 2025 · similarity 0.89)
• Improving Generalization in Federated Learning with Highly Heterogeneous Data via Momentum-Based Stochastic Controlled Weight Averaging (ICML 2025 · similarity 0.89)
• A New Theoretical Perspective on Data Heterogeneity in Federated Optimization (ICML 2024 · similarity 0.89)
• FedGPS: Statistical Rectification Against Data Heterogeneity in Federated Learning (NeurIPS 2025 · similarity 0.89)