Bayesian Coreset Optimization for Personalized Federated Learning
Prateek Chanda, Shrey Modi, Ganesh Ramakrishnan
TL;DR
This work presents CoreSet-PFedBayes, a personalized federated learning framework that optimizes per-client coresets via Bayesian coreset weights to train on compact, representative data rather than full local datasets. By integrating coreset weighting into the FL objective and adding a likelihood-consistency term, the method maintains alignment with full-data distributions while achieving strong generalization bounds, shown to be minimax-optimal up to logarithmic factors. The approach uses Accelerated Iterative Hard Thresholding (AIHT) to efficiently learn sparse per-client coresets and delivers consistent accuracy gains over random-sample baselines and submodular subset selectors on benchmarks like MNIST, Fashion-MNIST, CIFAR-10, and several medical datasets. Theoretical results quantify the generalization-error trade-offs with coresets, and experiments demonstrate reduced communication rounds and data requirements without sacrificing performance, highlighting the method’s practical impact for privacy-aware, heterogeneous FL.
Abstract
In a distributed machine learning setting like Federated Learning where there are multiple clients involved which update their individual weights to a single central server, often training on the entire individual client's dataset for each client becomes cumbersome. To address this issue we propose $\methodprop$: a personalized coreset weighted federated learning setup where the training updates for each individual clients are forwarded to the central server based on only individual client coreset based representative data points instead of the entire client data. Through theoretical analysis we present how the average generalization error is minimax optimal up to logarithm bounds (upper bounded by $\mathcal{O}(n_k^{-\frac{2 β}{2 β+\boldsymbolΛ}} \log ^{2 δ^{\prime}}(n_k))$) and lower bounds of $\mathcal{O}(n_k^{-\frac{2 β}{2 β+\boldsymbolΛ}})$, and how the overall generalization error on the data likelihood differs from a vanilla Federated Learning setup as a closed form function ${\boldsymbol{\Im}}(\boldsymbol{w}, n_k)$ of the coreset weights $\boldsymbol{w}$ and coreset sample size $n_k$. Our experiments on different benchmark datasets based on a variety of recent personalized federated learning architectures show significant gains as compared to random sampling on the training data followed by federated learning, thereby indicating how intelligently selecting such training samples can help in performance. Additionally, through experiments on medical datasets our proposed method showcases some gains as compared to other submodular optimization based approaches used for subset selection on client's data.