Bayesian Coreset Optimization for Personalized Federated Learning

AI summary

Overview

Research area: Federated learning (FL), Bayesian coreset selection, and personalized/distributed Bayesian inference.

Technical level: Advanced. The paper is heavily theoretical (Hölder-smooth function classes, minimax lower bounds, KL-divergence-based variational objectives) and assumes familiarity with variational inference, Bayesian neural networks, and federated optimization.

Scope: The paper proposes CoreSet-PFedBayes, a personalized federated learning architecture in which each client uploads updates derived from a weighted coreset of its own data rather than its full dataset, supported by generalization-error bounds and benchmark experiments.

What This Paper Is About

In federated learning, each client must train on its own local data and send updates to a central server, but training on every client's entire dataset is cumbersome and communication-heavy. This paper asks whether each client can instead learn from a small, intelligently weighted subset (a Bayesian coreset) of its own data while still achieving near-optimal accuracy. The authors propose CoreSet-PFedBayes, which learns personalized coreset weights per client so that the model trained on the subset behaves like the model trained on the full data.

Key Contributions

1. A new architecture that incorporates Bayesian coreset optimization into the federated learning setting, with each client's data weighted by a per-client coreset weight vector.
2. Novel objective functions for the federated Bayesian coreset problem, including one that combines the KL divergence between the coreset-weighted client distribution and the full-data client distribution with a likelihood-matching term, plus a sparsity constraint on the coreset weights.
3. Theoretical analysis showing the convergence rate of CoreSet-PFedBayes falls within logarithmic bounds, with matching upper and lower rates in the coreset size n_k.
4. Experiments on multiple benchmark datasets and against a range of existing baselines, showing accuracy gains from coreset selection over random subsampling, and additional gains over submodular subset-selection methods on medical datasets.

Main Findings

• Minimax optimality up to a logarithm. The generalization error of CoreSet-PFedBayes under the L2 norm is upper bounded by O(n_k^{-2β/(2β+Λ)} log^{2δ'}(n_k)) and lower bounded by O(n_k^{-2β/(2β+Λ)}), where n_k is the coreset size per client, β is the Hölder-smoothness parameter, Λ is the intrinsic dimension of each client's data, and C_2, C_3, δ' are constants. This means the rate is minimax optimal up to a logarithmic term in n_k.

• The coreset penalty is a closed-form function. Theorem 1 states that the difference between the upper bound on the overall generalization error of the full-data setup and that of the coreset-weighted setup is always bounded above by a positive closed-form function ℑ(w, n_k) depending on the coreset weights w and the coreset size n_k.

• Coreset weighting widens the generalization-error spread. Theorems 3 and 4 show the lower bound on the generalization error of the coreset-weighted CoreSet-PFedBayes objective is always higher than the lower bound for the vanilla PFedBayes objective, with the delta difference scaling as O(n_k^{-2β/(2β+Λ)}) — the error increases because of limited coreset samples, but remains bounded in closed form with respect to coreset size.

• Accuracy gains over random 50% sampling. In Table 1, comparing against PFedBayes trained on a randomly sampled 50% of the data, CoreSet-PFedBayes (k = 50%) improves by +4.87% on MNIST, +8.61% on FashionMNIST, and +9.71% on the remaining dataset reported in the text. The paper notes the method also beats some baselines that were trained on the full dataset (for example, outperforming pFedMe on FashionMNIST).

• Reported accuracies. CoreSet-PFedBayes (k = 50%) reaches 92.48% personal and 96.3% global on MNIST, 89.55% personal and 92.7% global on FashionMNIST, and 69.66% personal and 71.5% global on CIFAR. A RandomSubset (50%) control reaches 80.2/88.4 on MNIST, 87.12/90.75 on FashionMNIST, and 48.31/61.35 on CIFAR. PFedBayes (Full) reaches 98.79/97.21 on MNIST, 93.01/93.30 on FashionMNIST, and 83.46/64.40 on CIFAR. Note that Table 1 uses "-" where no accuracy is reported because of very slow convergence of the corresponding algorithm.

• KL divergence converges faster with coresets. Figure 3 shows that, on MNIST, the optimal distribution parameters learned from the coreset approach those learned on the full dataset within a few epochs (decreasing KL divergence), and that training on a randomly selected subset takes more epochs to converge than the coreset-selected subset.

• Accelerated Iterative Thresholding recovers the true posterior. On the Housing Prices 2018 dataset, using Bayesian Riemann linear regression with coreset sizes k = 220, 260, and 300, the radius capturing the weights of the coreset points matches the true posterior most closely at k = 300.

• Medical dataset comparison. The paper reports that on medical datasets the method shows some gains compared with submodular optimization approaches for subset selection on client data. The medical datasets used are described as consisting of 3 classes each, with one class (Normal) shared between 2 clients and the other two classes distributed separately to the two clients. No specific accuracy figures for the medical experiments appear in the provided content.

Methodology in Plain English

The authors set up a standard federated setting: one server and N clients, each holding n samples of the same size. Each client's data is assumed to follow a regression model with Gaussian noise, the unknown functions are assumed to be β-Hölder-smooth, and each client has the same fully connected deep neural network with its own parameters.

On top of this, they borrow the idea of a Bayesian coreset: instead of using all n of a client's data points, learn a weight vector w_i that selects and weights a small subset of size n_k (with n_k much smaller than n). The coreset is chosen so that the weighted data likelihood stays close to the full-data likelihood, measured in an L2 norm over a weighting distribution that shares support with the true posterior.

The client-side objective is rewritten to include these weights, and the server aggregates client distributions as before, with the family of client and server distributions assumed Gaussian so the KL divergence has a closed form. The key new objective picks the coreset weights w that both keep the coreset-trained client distribution close to the full-data client distribution, and keep the coreset-weighted likelihood close to the original likelihood, subject to a sparsity constraint on w.

For optimization, the algorithm uses a reparameterization trick (writing θ in terms of μ and ρ) and mini-batch stochastic gradient descent with Monte Carlo sampling to estimate the client objective, plus an extra convergence-accelerating parameter. The sparse coreset subproblem is solved with Accelerated Iterative Thresholding (A-IHT). Clients update over several local iterations before uploading to the server, which then updates the global distribution.

The theoretical work then compares the average generalization error of CoreSet-PFedBayes against the PFedBayes baseline using Hellinger distance and the L2 norm, deriving upper bounds, lower bounds, and the closed-form error gap.

Why This Matters

This work connects two areas that have largely been studied separately — Bayesian coreset selection and federated learning — and shows that the loss in generalization from training on a small subset can be characterized in closed form rather than being a vague empirical trade-off. It also argues that model-centric subset selection, which accounts for how each client's model distribution changes, can outperform data-centric or submodular subset selection.

Real-world applications:

• Recommendation systems, where the paper notes personalization and privacy-preserving recommendation are a major focus, and where inference over large datasets with only a subset of representative points is common.
• Healthcare and medical deployments, the setting the authors specifically test against submodular methods, where data is partitioned across institutions and sharing full datasets is impractical.
• Mobile and edge machine learning, where each user's device holds only a small amount of local data and must still deliver a personalized experience.
• Resource-constrained distributed training, where reducing the data each client trains on can lower communication and computation overhead between clients and the server.

Industry relevance: federated learning is deployed in production systems such as recommendation and mobile app models, both of which the paper cites. A method that preserves accuracy while training each client on a fraction of its data is directly relevant to communication cost and on-device resource limits. The authors note that privacy mechanisms such as differential privacy and homomorphic encryption are not the focus of this work but can be integrated with their approach.

Future Directions

• Integrating privacy mechanisms. The paper explicitly states that differential privacy and homomorphic encryption can be combined with the approach but leaves this to future work; evaluating CoreSet-PFedBayes under a privacy budget is a natural next step.
• Filling out the medical and submodular comparison. The abstract claims gains over submodular subset selection on medical datasets, but the provided content describes the setup without reporting specific accuracy numbers, so a fuller empirical comparison remains open.
• Scaling beyond the tested configuration. Experiments use 10 clients on MNIST, FashionMNIST, and CIFAR-10 under a non-i.i.d. setting; whether the logarithmic-bound behavior holds for many more clients and heterogeneous data sizes is not established here.
• Tightening the error gap. Theorems 3 and 4 establish that the coreset-weighted lower bound always exceeds the vanilla lower bound; quantifying and reducing the closed-form function ℑ(w, n_k) in practice is an open direction.

Target Audience

Researchers and graduate students working on federated learning, Bayesian deep learning, variational inference, and coreset or subset-selection methods. It is also relevant to practitioners building federated systems where communication cost and on-device data limits matter, though the theoretical content requires a background in statistical learning theory to follow in full.

Authors’ 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.