Cost-Free Personalization via Information-Geometric Projection in Bayesian Federated Learning
Authors: Nour Jamoussi
Giuseppe Serra
Photios A. Stavrou
Marios Kountouris
Venue: Accepted by TMLR ()

Abstract: Bayesian Federated Learning (BFL) combines uncertainty modeling with decentralized training, enabling the development of personalized and reliable models in the presence of data heterogeneity and privacy constraints. Existing approaches typically rely on Markov Chain Monte Carlo (MCMC) sampling or variational inference, often incorporating personalization mechanisms to better adapt to the local data distributions. In this work, we propose an information-geometric projection framework for personalization in parametric BFL. By projecting the global model onto a neighborhood of the user's local model, our method enables a tunable trade-off between global generalization and local specialization. Under mild assumptions, we show that this projection step is equivalent to computing a barycenter in the statistical manifold, allowing us to derive closed-form solutions and achieve cost-free personalization. We apply the proposed approach within a variational learning setup using the Improved Variational Online Newton (IVON) optimizer and extend it to general aggregation schemes in BFL. Empirical evaluations under heterogeneous data distributions confirm that our method effectively balances global and local performance with minimal computational overhead.

Published in Transactions on Machine Learning Research (01/2026)
Cost-Free Personalization via Information-Geometri Projection in Bayesian Federated Learning
Nour Jamoussi
nour-jomoussiBeureomm.f
Communication Systems Department
EURECOM, Fronce
Giuseppe Serra Communication Systems Department EURECOM, France
giuseppe..crruteeerccoo.
Photios A. Stavrou Commanication Systems Department EURECOM, Fronce
Jolios.tuurouuCeercoom.ff
Marios Kountouris
mariosk@ugr.es
Department of Computer Science and Artificial Intelligence, University of Granada, Spašn Communication Systems Departmend, EURECOM, France
Reviewed on OpenReview: https: // openreview. net/ forum? id= 9y0jCrrajDRR
Abstract
Bayesian Foderated Learning (BFL) combines uncertainty modeling with decentralized train- ing, enabling the development of personalized and reliable models in the presence of data heterogeneity and privacy constraints. Existing approaches typically rely on Markov Chain Monte Carlo (MCMC) sampling or variational inference, often incorporating personaliza- tion mechanisms to better adapt to the local data distributions. In this work, propose an information--eommetii projection framework for personalization in parametric BFL. By projecting the global model onto ncighborhood of the user's local model, our method en- ables a tunable trade-off between global generalization and local specialization. Under mild assumptions, we show that this projection step equivalent to computing a barycenter the statistical manifold, allowing us derive closod-form solutions and achieve cost- free personalization. We apply the proposed approach within variational learning setup using the Improved Variational Online Newton (IVON) optimizer and extend it to gen- eral aggregation schemes in BFL. Empirical evaluations under heterogeneous data distribu- tions confirm that our method effectively balances global and local performance with min- imal computational overhead. Code is available at https:///ithuut: com//NourJamousss// Learnnin
Inf
1 Introduction
Federated learning (FL) is collaborative machine learning paradigm designed to preserve data privacy, By connecting central server to multiple participants, commonly referred to as clients or end-users, FL enables distributed model training while keeping data local to each client. In a typical FL selting, clients train local models on their private datasets, and the server aggregates these locally updated models into global model. After each round of local updates, clients share their updated model parameters with the server, which aggregates them refine the global model. The updated global model is then sent back the clients for further local training McMahan et al. (017

---

Published in Transactions on Machine Learning Research (01/2026)
In real-world federated learning scenarios, statistical heterogencity among clients is a pervasive challenge. For instance, variations in Magnetic resonance imaging (MRT) scanners across hospitals can induce feature skewness, where the input data distributions differ substantially between clients. Likewise, label distribution shifts may occur when certain clients have disproportionately more samples from specific classes than others, such as specialized hospitals locusing on particular diseases compared to general hospitals. This hetero- gencity violates the standard assumption of independent and identically distributed (i.i.d.) data, leading to degraded performance and increased uncertainty of the global model when evaluated on local test data.
Personalized Federated Learning (PFL) is paradigm in which clients adapt the global model to their local data, i.c., perform model personalization, with the goal of mitigating the impact of statistical heterogeneity. In Tan et al. (2023). the authors propose a comprehensive laxonomy of PFL approaches, classifying them into two main categories: global model personalization techniques and personalized model learning techniques. The former category includes methods such as client selection Yang et al. (2021) , meta-learning Fallah et al.. (2000), and transfer learning Chen et al. (2020), while the latter encompasses strategics such as parameter decoupling Arivazhagan al. (2019). knowledge distillation [Jeong & Kountouris (2023). and clustering-
based approaches [Sattler et al. 2020). Furthermore, several studies Zhang et al. (202bb); Boroujeni et al.)
2024 | Kotelevskii et al. [2022] Zhu et al. (2023) have applied Bayesian methods within the PFL framework,
demonstrating improved model calibration and uncertainty quantification. The transition from deterministic FI. to Bayesian FI. (BFT.) naturally motivates an investigation of the manifolds to which clients' and global posterior distributions belong. This perspective enables the exploration of their geometric properties and the definition of meaningful operations on these manifolds, as discussed in Jamoussi et al. (2024).
In this work, we study the PFL paradigm through the lens of Bayesian learning, leveraging principles from Information Geometry to develop a novel personalization method for BFL. Unlike existing approaches, the proposed method does not require fine-tuning, additional training, or access to either local or shared global data, making it fully private and computationally cost-free personalization technique. By specifying a divergence metric, we obtain personalized model posteriors by projecting the global posterior distribution onlo a sphere centered at the local posterior distribution, where the radius encodes the desired degree of personalization specified by the end-user. This formulation enables unrestricted personalization flexibility, allowing users to scamlessly control the trade-off between local adaptation and global generalization without incurring additional computational or privacy costs.
We show that, for any divergence function that is convex in its first argument, the projection problem is equivalent to computing the weighted barycenter between the local and global posterior distributions. This equivalence establishes a conceptual bridge between two fundamental concepts in Information Geometry, geometric projection and geometric barycenter, thereby providing an interpretable relationship between the radius of the projection sphere and the weights in the barycentric aggregation. Consequently, the degree of personalization becomes both intuitive and directly controllable by end-users.
The paper is organized as follows. Section 2 reviews the related work. Section B introduuess the parametric BFL framework based on variational learning for local training and posterior aggregation to estimate the global posterior model. We exlend the approach of Pal et al. (2024) by employing the Improved Variational Online Newton (IVÓN) optimizer Shen et al. (2024) for multiple aggregation methods. In Section 4 we present a novel personalization strategy for global BFL models using information-goometrie projection and show its equivalence to barycentric aggregation for divergence functions convex in their first argument. The proposed method is general and applicable beyond parametric BFT, including non-parametric settings. For the parametric case, we exploit closed-form barycentric solutions for specific divergences, yielding significant computational gains. Section reports the experimental results, comparing aggregation techniques under the variational learning framework and benchmarking our method against state-of-the-art approaches. Finally, Section G discusses the broader applicability of the proposed method to continual learning.
2 Background and Related Work
Bayesian Federated Learning. BFL. socks to integrate the advantages of Bayesian Deep Learning, such as improved model calibration and uncertainty quantification, within the FL framework. Following the classification proposed in Cao et al. (2023), BFL approaches can be broadly categorized into two main types:
2

---

Published in Transactions on Machine Learning Research (01/2026)
client-side BFL Zhang et al. (2022b); Zhu et al. (2023): Liu et al.
20233 Boroujeni et al. 2024 Hasan et al.
(2024b; Bhatt et al. 2024) and server-side BFI. Corinzia et al.
(2019 Chen & Chao (2020); A1-Shedivat
et al. (020); Guo et al. (2023), with certain overlaps between these groups, reflecting the diversity in model
formulations and inference techniques. In Jamoussi et al. (2024). the authors extend the discussion of the
parametric client-side BFL framework by categorizing met hods according to the techniques used to aggregate
local posteriors and propose a unified framework for aggregation rooted in information geometry.
However, within the scope of this work, it is useful to introduce an additional categorization of BFL based on the adopted modeling assumptions, distinguishing between parametrie and nonparametric Bayesian methods.
Parametric BFL: Parametric Bayesian learning provides principled framework within Bayesian statistics, wherein the data generation process is assumed to follow a parametric model. Specifically, this approach assumes that the data distribution can be fully characterized by finite set of param- eters. The parametric family most commonly considered in BFL is the Gaussian family, in which both client-specific posteriors and the global model posterior are approximated using Gaussian dis- tributions Zhang et al. (2022b); Ozer et al. (0022;; Guo et al. (2023) : Kim & Hospedales (2023); Pal
ct al. (2024): Swaroop al. (2025). This Gaussian assumption facilitates tractable inference and
eflicient parameter aggregation.
Nonparametric BFL: Unlike parametric BFL, nonparametric Bayesian learning does not restrict the model specific family of distributions. Instead, it provides the flexibility to learn directly from the data, allowing the structure and complexity to emerge naturally rather than being predefined. This approach leverages distributions over infnite-dimeesiooall function spaces, such Gaussian Processes or Dirichlet Processes, as employed in the context of BFL in Yurochkin et al. (2019). Additionally, particle-based variational inference, used in Kassab & Simeone (202)) for BFT, aligns with the nonparametric perspective by using a set particles to approximate posteriors without assuming a predefined parametric family.
Personalized Bayesian Federated Learning. Personalized BFL. extends the standard BFL framework by incorporating cpient-specific adaptations into the shared global model. FedPop Kotelevskii et al. (2022) achieves personalization by modeling each client's data generation process as a combination of fixed shared population parameters, which describe the common data model, and client-specific random effects, which capture heterogencity in client data, providing personalization and uncertainty estimation capabilities. To elficiently infer these parameters, FedPop employs MCMC methods to approximate local posterior distri- butions and perform stochastic optimization in a federated setting. Although MCMC enables flexible and asymptotically exact Bayesian inference, it can be computationally expensive in large-scale scenarios. To address this, paedBaaes Zhang et al. (022b) adopts variational inference to approximate client-specific posteriors, optimizing a tractable Evidence Tower Bound (FLBO) to achieve computational efficiency in BFL. Although paye presents a promising approach personalization in BFL, certain aspects of its methodology introduce notable limitations. First, the degree of personalization is uniformly applied across intuitive rationale for tailoring
all clients and determined empirically, without an adaptive mechanism personalized based on the use of it to individual client characteristics. Second, it defines the local model as variational inference, where the global model serves as a prior within the KL divergence term of the ELBO. This formulation may be viewed as controversial within the Bayesian community, since the classical interpre- tation of a prior belief assumes independence from the observed data. More recently, PFedVEM Zhu et al. (2023) has leveraged variational inference estimate the client-specific uncertainty and model deviation by modeling client parameters Gaussian distributions centered around global latent variable. It employs confidence-basee based aggregation strategy that ensures that clients with lower uncertainty and smaller deviations contribute more to the global model.
Unlike the aforermentooeed methods, our approach introduces personalization as an additional step indepen- dent of the training process, requiring only access to the global and local posteriors, with no need for extra data or fino-tuning. Morcover, our method provides access to three different variants of the model, i.c., local, personalized, and global, as suggested in Divi et al. (2021), with the ability continually adjust the personalized variant without incurring additional training costs.
3

---

Published in Transactions on Machine Learning Research (01/2026)
Information Geometry and Optimal Transport in FL. Information geometry Amari
2016 provides
general framework for studying the geometric properties of statistical manifolds through divergence fune- tions, allowing distributions to be intuitively treated as elements of a manifold. This framework naturally extends fundamental geometric concepts, such as projections onto constraint sets Csiszár (1975): Csiszár & Matus (2003) and the identification of barycenters, which serve as statistical centers of mass for a given set of probability distributions Nielsen & Nock (2009;; D'Ortenzio et al. (2022).
In contrast, optimal transport Villani (2009) focuses on identifying the most efficient transport plan between two probability distributions with respect to given cost function. Under specific assumptions on the cost function, optimal transport enables the definition of distance functions on the statistical manifold, as exem- plified by the family of Wasserstein-p distances Villani (2009). The Sinkhorn algorithm Cuturi (0133) serves as a natural bridge between information geomelry and optimal transport, enabling efficient compulation of transport plans by incorporating entropy regularization into the optimal transport problem.
Multi-Marginal Optimal Transport (MMOT) Gangbo & Swiçch (1998); Pass (2015) extends the classical two- two-marginal optinal transport problem to multiple distributions by seeking an optimal joint coupling that minimizes given n-ary cost function over all marginal distributions. Furthermore, for specific forms of the T-ary cost, the MMOT problem can be shown to be equivalent to computing the Wasserstein barycenter (WB) of the set of marginals. This equivalence highlights MMOT as more general framework within which the WB emerges as a special case, characterizing the optimal interpolation of multiple distributions in Wasserstein space. Recently, applications of concepts from information geometry and optimal transport have been explored in the FL setting. In Farnia et al. (2022), the authors introduce FedOT, a PFL algo- rithm that integrates optimal transport with model training. FedOT employs deterministic FL, leveraging MMOT to align data distributions. Specifically, it learns transport maps that transform data points from heterogeneous distributions into shared domain while simultaneously training a predictive model the mapped data. Similarly, FedDRO Li et al. (2024) aggregates client data distributions via a Wasserstein barycenter and trains the global model against worst-case perturbations within Wasserstein ball centered at this barycenter. In the context of BFL, Hassan et al. (2023) considers structured latent variable models in which local latent variables are kepl private, formulating a decentralized variational inference problem and proposing a communicttion-fficceet aggregation scheme based on Wasserstein barycenters. Similarly, Jamoussi et al. (2024) introduces unifying framework for aggregation in BFL grounded in barycentric aggregation.
3 Parametric Bayesian Federated Learning
3.1 Learning Phase
The primary objective of BFT, is to estimate the posterior distribution of the global model parameters, denoted as ((0"1D), using the posterior distributions of local models, p(() However, exact posterior inference is typically computationally intractable, necessitating the use of approximate inference methods. In this study, we adopt variational learning to approximate local posterior distributions based on a shared prior distribution p(0) and client-specific likelihoods p(Dx) Given a parametric distribution family Q, optimization seeks a distribution g Q that minimizes the KL divergence from the true posterior distribution p(O)D), i..,
min (kL(((()))<())))) (00)0
However, direct optimization of equation 1 is generally intractable, motivating the use of the Negative Evidence Lower Bound (ELBO) as a surrogate objective:
EE (ogp + (kr.(a(0)llp(0). 2)
min
(2)
We approximate the posterior distributions of the local models p(0x Dz), Vk E [I....,N} by optimizing the objective in equation using the IVON optimizer [hhen et al. (2024), which is grounded in the Bayesian learning rule Khan & Rue (2021). Unlike classical variational inference, IVON integrates variational learning directly into the optimization process, without requiring modifications to the model architecture or loss
4

---

Published in Transactions on Machine Learning Research (01/2026)
function. IVON maintains Gaussian posterior over the weights through efficient second-order updates based reparameterized Hessian estimates, enabling uncertainty-aware learning at a computational cost comparable to deterministic training with Adam. This makes Bayesian deep learning more scalable in practice. Subsequently, the local posteriors are aggregated to obtaín the global posterior distribution p(0 (D).
Based on this formulation, we introduce the following assumptions regarding the common prior p(0) and the variational family 2, which will hold throughout the experimental setting described in Section 5
Assumption (Mean-field Model). The voriational family Q consists of d-dimensional Gaussian distri- butions with independent marginals. Specifically, ~ N(,, where μ E Rd is the mean vector and Σ diag(of......of is a diagonal covoriance matriz.
3.2 Aggregation Phase
The aggregation phase refers to the process combining locally trained model updates from multiple clients into a single global model. In what follows, we outline and discuss several aggregation techniques for the N local distributions, each consistent with Assumption and parametrized with { ))))
Empirical Arithmetic Aggregation, also known as naive aggregation, is employed in Zhang et al. (2022b); Ozer et al. (2022): Bhatt et al. [202]] Fischer et al. (2024). computes the weighted average of the distribution statistics:
N Σ
ΣEAA
W
MEAA
wk/4k
(3)
Barycentric Aggregation, adopted for BFL in Jamoussi et al. 2024 uses the barycenter of the clients' posteriors as the aggregated model, minimizing the average discrepancy among client distributions. Some examples include:
Wasserstein-2 Barycenter: minimizes the average discrepancy in terms of the Wasserstein-2 distance between the clients' distributions. Under Assumption [ closed-form solution for the Wasserstein barycenter is derived in Álvarez-Estee Esteban et al. (0166):
N WyHk
Ew3
Σ
Ww3
(4)
Reverse KT. Barycenter: minimizes the average reverse KT. divergences between the local distri- butions and coincides with the multiplicative aggregation of posteriors proposed in Al-Shedivat
ct al. 2020): Liu et al. (2023); Guo et al. (2023 Pal ct al. (2024). Its closed-form expressions
are given in Koliander et al. (2022):
N ERKL
PRKL ERKL
Σ
E pk.
To support the different aggregation strategies that operate on covariance matrices, and to avoid arbitrarily setting the effective sample size parameter in IVON, which ideally corresponds to the total dataset size but is often unknown at the server, we adopt a subclass of IVON that explicitly stores the covariance matrix and performs sampling directly from it, rather than relying the Hessian matrix. This reformulation is made
possible by the relation derived in Shen et al. (2044)
1
0²
(6)
N(h
which links the variance o² to the Hessian approximation R the dataset size N, and the weight decay term d. In practice, the covariance matrix is first computed locally after estimating the Hossian. These
For multivariate Gaussian, the Hcssian of the negative log-likelihood is proportional the inverse covariance matrix. When the covariance is diagonal, the Hessian also diagonal, so Hessian approximation reduces per paramaaeeter entries.
5

---

Published in Transactions on Machine Learning Research (01/2026)
Pa
Pu 1
Pg
Pg.k
Pg.k
P g.k
Pk
11 0
Figure Personalization through cnermaatin-gcometiin projection. The figure presents two projection scenarios illustrated with two local spheres St and Sj of increasing radius r and r highlighting the impact of the radius the closeness of the projected distribution to the global local distribution.
local covariance matrices are then aggregated at the server, and the resulting global covariance is distributed back to the clients, where it is used to reconstruct the updated local Hessians. Unlike prior work [Pal et al. (2024), which employed IVON but was limited to RKI.B-based aggregation due to its direct dependence on the Hessian, our formulation enables a broader range of covariance- based aggregation strategies.
4 Personalization via Informatiin-Geometriic Projection
The main contribution of this work lies in interpreting the personalization problem as an information- geometric projection problem. Given a divergence function D, the objective is to project the global posterior Pg onto the projection set of the 1th client, defined as a sphere centered at the local posterior Pk with radius Tk- The projection is performed after FL training has been completed. Conceptually, this projection can be viewed identifying the distribution within the *local neighborhood" of the client's posterior that best aligns with the global model. The radius тк quantifies the degree of personalization required by client ke where smaller values of Tæ correspond to stronger adherence to the local posterior, while larger values allow greater influence from the global model, as shown in Figure In the following, we formally define the local sphere S₁ induced by the divergence function D.
Definition 1. (Local Sphere Sp(p,r) r)) Given statistical manifold M and divergence function D, the local sphere Sp(p, r) centered at E M with radius r E [0, 00) is defined as the set
Splp.
= =
{ D(ali) CM
To simplify the notation, we denote by S₁ = Sp(pk, rK) the local sphere associated with the ttth client, centered at the local postcrior pk with radius rk- We are now ready to formally define the personalization problem as a projection problem.
Problem 1. (Projection onlo Sx) Given a statistical manifold M, a divergence function D, global posterior Pg E M, and a local sphere Sk, we define the projection of Pg onto Sy as the optimization problem
D(pllyg) PES&
min
(7)
The solution to this projection problem is identified as the personalized posterior distribution of client denoted by Pg,k arg minpss. D(pllyg)-
Remark 1. By varying the radius T E D(pylle)) the solutions of the associated projection problem trace the geodesic between pk and Pg. i.e., the shortest possible path connecting the two distributions under the geometry induced by the divergence D.
To support our derivation, we first introduce the definition of barycenter with respect to divergence D.

---

Published in Transactions on Machine Learning Research (01/2026)
Definition 2. (D-barycenter) Given statistical manifold M, divergence function D and set of dis- Iribulions {pA I S .M with associaled normalized weights {wE))== the barycenter of the set {pk) is defined as
N
ph Phplpal fur arg min 9EM
wy ((glpp))))
(8)
The following mild assumption essential for establishing the equivalence between the projection and barycenter formulations, shown in Theorem
Assumption The divergence metric convex function in its first argument.
Remark 2. Most divergences used in practice, such as the family of f-divergences and the Wasserstein-p distances, are conver in both arguments, and therefore satisfy Assumption
Theorem 1. Under Assumplion 2 the solution of the projection problem is equivalent to that of the weighted barycenter problem i.e.,
Pg.k Ph Pk), fwgg wk)),
(9)
where the weights wa and we are given by
1
À
Wg A A+1"
Wk =
(10)
1
A +1
for some À E [0)
The proof of Theorem
I is provided in Appendix
We highlight the following obscrvations regarding the relationship between To and X.
Remark 3. There erisls an inverse proportional relationship between the Lagrangian mulliplier X and the radius Tk. Specifically, as À + 0, the radius тк - oc, corresponding to the case where the personalized posterior coincides with the global posterior. Conversely, the limät X -> so, the radius rk vanishes (rk + 0), implying that the personalized posterior collapses to the local posterior. Hence, by selecting a value of X, implicitly determine the personalization radius Tk-
The main advantage of the equivalence between projections and barycenters lics the improved tractability of the barycentic formulation. For instance, under Assumption [ analytical solutions are available for both the RKL divergence and the Wasserstein-2 distance, shown in Koliander et al. (2022); Alvarez- Esteban et al. 2016). These closed-form expressions enable straightforward and computationally efficient personalization procedure, incurring virtually no additional cost.
5 Experiments
5.1
Experimental Setting
Datasets and Heterogeneity Simulation. Interpretability and reliable uncertainty quantification are crucial in high-stakes domains such banking and healthcare. When these institutions participate in FL, they typically do so in cross-silo setting, where the number of clients is small but each holds large local dataset. Reflecting this application focus, concentrate our empirical evaluation cross-silo simulations rather than cross-device benchmarks Caldas et al. (2018). To simulate these conditions, we evaluate our approach on three widely used image classification benchmarks: FashionMNIST,, SVHN, and CIFAR-10 100 and we induce label shift across 10 clients using Dirichlet- based partitioning. Following prior works Li et al. (2020); Yurochkin et al. (2019): Wang et al. (2020a bb: Lin et al. (2020); Ozer et al. (2022), we sample client-specffic label distributions from Dirichlet distribution with concentration parameter 0.5, which induces label shift and yields non-i.i.d. data across clients.

---

Published in Transactions on Machine Learning Research (01/2026)
Metrics. Following the taxonomy of evaluation metrics in Table IV of Tan et al 2023 our evaluation focuses mainly model performance and trustworthy AT criteria. For model performance, report overall global-model accuracy as well as the average and variance of personalized-modde accuracies, together with the worst-performing 10% of clients' accuracy. The latter two also connect to the fairness column in Table IV, and we elaborate on this Section 5.5 addition, our Bayesian formulation explicitly largets uncer- tainty quantification and model calibration, which are not explicitly listed among the model-performance brustworthy-Al metries in Tan et al. (2023). We therefore report Expected Calibration Error (ECE) and Negative Log- Likelihood (NLL) alongside accuracy-based measures to assess both the quality and the reliability of the predictive distributions. In contrast, we do not report system-performance metrics such as fault tolerance or system scalability, since our method has similar per-round complexity to FedAvg (see Section 3 and our focus here is on the statistical behavior of the personalized posteriors rather than on system- level aspects.
Models Architecture and Hyperparameters. Our model architecture consists of two convolutional layers with 5x5 kernels and ReLU activations, each followed by 2x2 max-pooling layer. The extracted features are flattened and passed through three fully connected layers of sizes 120, 84, and 10, respectively, produce the final class logits. For all state-of-the-art methods we compare against, we closely follow the implementation details provided in the original papers. The complexity of our architecture is comparable that of the models used in these methods, ensuring fair and consistent comparison. Whenever possible, use the official codebases for implementation; otherwise, carefully reproduce the setups following the authors' guidelines to maintain alignment with their reported settings. We make exception for pFedBayes on FashionMNIST, where, instead of the original multi-layer perceptron with one hidden layer, we use our convolutional architecture to enable a fair comparison.
It worth noting that training with IVON can be sensitive hyperparameter scttings, particularly to the initialization of the Hessian. Therefore, we report in Table | the hyperparameters used for training on each dataset. Here, Ng denotes the size of the training dataset for the 1tth client.
Table 1: IVON Hyperparameters.
params
FashionMNIST SVHN
CIFAR-10
initial learning rale
0.1
0.1
final learning rate
0.01
0.01
0.01
weight decay
2o-4
2c-4
2o-4
batch size
64
64
64
ESS
N₂
N₂
N
initial hessian (ho)
5
2
I
MC sample while training
1
1
MC sample while testing
10
10
10
Personalization Step. For personalization, we restrict our experimental evaluation RKLB and WB, as both barycenters preserve the Gaussian structure of the distributions. This preservation critical design choice, it facilitates the adaptability of the personalized models. In particular, ensures that the posterior distribution of the personalized model remains consistent with those of both the global and local models, which are assumed to be Gaussian under the variational learning framework.
5.2 Comparison between the Different Aggregation Methods Considered
Our experiments investigate various combinations of global update methods (EAA, RKLB, WB) and per- sonalization techniques (RKLB, WB). Statistical tests, presented in Appendix B support our observation that these configurations yield comparable performance. To reduce redundancy and enhance readability, we thercfore restrict the results reported in the subsequent experiments the WB method for both global updates and personalization.
8

---

Published in Transactions on Machine Learning Research (01/2026)
5.3 Effect of À the Trade-off off between Performance Local and Global Data
We analyze the impact of the personalization parameter on model performance across three datasets under heterogeneity simulated using the Dirichlet distribution. Across all datasets, we observe consistent trade-off between performance on global and local data. Results on CIFAR-10 are reported in Figure and the results on MNSST and SVHN are reported in Appendix C The global data, i.e., the union of all client test sets, are approximately uniform across classes while the local data follow distinct Dirichlet distributions. Notably, À = corresponds the global model, whereas À oo represents the local model. As À incrcases, performance on the global distribution deteriorates in terms of accuracy, calibration (ECE), and uncertainly quantification (measured by NLL.), whereas local performance improves and remains stable within certain range of A. The increased ECE and NLL for local models (A oo) global data suggest that over- personalization may reduce generalization and model confidence, particularly for underrepresented classes, Conversely, at lower values of X, the model achieves improved global performance but fails to effectively capture specifc distributions, showing reduced confidence on local data. These results underscore the importance of in controlling the trade-off between generalization and personalization, allowing the model adapt effectively heterogeneous data distributions in non-i.idd federated settings.
Average per ferfrmaance
Average performance
Clients performance on data
performance
35
Accuracy
0.35
231
0.33
0.32
0.55
05
Figure 2: Effect of on performance across local and global data distributions. Results are reported for the CIFAR-10 dataset. Notably, À corresponds to the global model, whereas + ao corresponds to the local model.
5.4 Comparison with State-of-the-Art
In this section, we provide a detailed analysis of the results presented in Table 2 and Figure 3 All reported values correspond to the mean and standard deviation computed over three independent runs.
Accuracy on Global Data
80
Methods Datasets
ECE on Global Data
NLL on Global Data
10.0
FedAVG
pFedBayes
7.5
FedPop
FFedVem
40
5.0
Proposed
FashionMNIST
0.2
2.5
SVHN
CIFAR-10
80
90
0.1
0.2
0
Accuracy on Local Data
ECE on Local Data
NLL on Local Data
Figure Trade-offs between local and global performance for personalized models. Each subplot presents results for different evaluation metric: Accuracy (left), RCE (center), and NLL (right). Points represent method dataset pairs. For Accuracy, the top-right region indicates better performance trade-off, whereas for ECE and NLL, the bottom-left region is preferable. Our method (with A 1) consistently achieves a favorable balance across all metrics and datasets.

---

Published in Transactions on Machine Learning Research (01/2026)
Comparison of all methods across four evaluation settings: personalized models on local data (PM on LD), personalized models on global data (PM on GD), global models on local data (GM on LD), and global models on global data (GM on GD). Results are reported for three datasets (FashionMNISTT, SVHN, and CIFAR-10) using Accuracy, NLL, and ECE evaluation metrics.
Table 2:
MNSHS
SVII
CIFAR-10
Sotting
Method
Are
ECE
NLL
Aoc
ECE
NL.
Ace
ECR
NI.L.
FodAvg
92.28 + 3.16
0.07 0.03
0.50 + 0.26
91.17 3.38
0.08 0.03
0.74 + 0.20
74:80 4.89
0.13 + 0.02
0.91 0.16
Bayde
91.42 + 3.35
0.05 - 0.01
0.25 + 0.00
80.27 = 3.27
0.07 + 0.02
0.53 + 0.19
71.84
6.25
0.17 + 0.04
1.10 + 0.25
PM on LD
FodPop
98.35 = 0.37 0.02
0.004
0.09 = 0.02
89.77
3.49
0.08 + 0.03
0.74 4 0.28
80.51 + 4.66
0.07 = 0.01
0.58 1: 0.12
PFodVem
93,00 + 2.80
0.06 = 0.02
0.34 + 0.14
91.69 = 3.08
0.07 # 0.022
0.35 + 0.20
72.61
5.09
0.23 = 0.04
1.91 4 0.12
Ours (A
92.22 + 3:23
0,06 0.02
0.34 = 0.14
91.87
3.03
0.05 # 0.02
0.40 = 0.14
76.82
1.01
0.08 = 0.01
0.66 + 0.10
FodAvg
79.12 +: 5.56
0.18 = 0.05
1.80 1: 0.68
68.67 8.54
0.27 1 0
3.11 4: 1.13
45.89 6.78
0.36 + 0.07
3.65 1:24
PPedBayes
77.51 +: 5.11
0.12 0.05
0.78 +: 0.25
66.20 + 10.71
0.26 0.10
2.28 1.04
44.48
6.50
0.38 + 0.088
3.12 0.90
PM on GD
FedPop
6.53 0.49
0.02
11.82 0.78
60.00 + 11.71
0.34 + 0.12
4.92 + 2.28
52.05
8.00
0.25 + 0.08
2.86 1.54
PeeeVVem
80.74 + 5.17
0.15
0.05
1.15 + 0.46
67.02 + 10:31
0.26 + 0.10
2.55 + 1.10
45.88 + 6.05
0.46 + 0.06
5.15 + 1.22
Ours (A
84.61 + 3.41
0.11 0.03
0.73 0.18
79.87
0.12 0.04
1.00
0.31
57.77
6.98
0.15
0.07
1.38
0.33
FedAvg
87.89 + 4.83
0.10
0.04
0.73 + 0.29
86.02
3.97
0.11 + 0.03
0.98 + 0.31
61.24
7.67
0.16 + 0.05
1.20 + 0.25
Bayd
88.00 1: 0.06
0.02
0.34 1: 0.13
85.76 6.11
0.09 0.04
0.67
0.32
5.17
1 0.033
1.25 0.21
GM on LD FedPop
VeedVem
89.49
3.95
0.16
4.74
0.10 4 0.03
4.50
= 2:43 0.35
Ours
86.45 + 4.88
0.08 - 0.03
0.47 + 0.20
87.18
4.63
0.07 + 0.03
0.58
+ 0.20
65.39
6.74
0.11 + 0.02 0.99 + 0.18
FedAng
87.88 1:
1:
86.06
0.55
0.01
61.63
3.81
1: 1.18 0.13
BaddBaye
86.02 + 0.39
0.03 + 0.005
0.34 + 0.02
86.03 0.41
0.08
0.005
0.06 + 0.04
63.86 + 1.58
0.16 + 0.01
1.25 + 0.06
GM on GD FedPop
pPld Vem
89.50 = 0.23
DJ07 0.0112
0.45 1: 0.02
86.32 = 0.22
0.09
0.004
0.80 - 0.03
60.88
1.44
0.29 = 0.01
2.41 1 0.10
Ours
86.14 + 0.60)
0.07 0.014
0.19 +: 0.02
86.54 1.05
0.06 + 0.01
0.60 + 0.06 65.05
3.57 0.06 = 0.01 0.99 0.09
Results of Personalized Models. It is worth noting that all methods in the comparison, except FedAvg, compare the local models of FedAvg, i.e., the models sent by
employ personalized models. Consequently, the clients before aggregation, with the personalized models produced by the other methods, including ours. The results are presented in the first two groups of rows in Table 2 and Figure 3
On Local Dala: Personalized models are designed to adapt to client-specific distributions, which is clearly reflected in their local accuracy scores, shown in the first group of rows in Table 2 Our method achieves performance comparable to or better than the baselines; for example, it attains the highest accuracy on SVHN and the second-highest on CIFAR-10. In terms of calibration (ECE) and uncertainty quantification (NLL), our method consistently yields low error rates on local data, achieving the best performance on SVHN and second- best on CIFAR-10.
On Global Data: A key limitation of many personalized methods is the degradation in global per- formance due to overfitting to client-specific distributions. This trade-off is particularly evident in FedPop on FNSSi, which achieves the best local accuracy but performs poorly on global evaluations. In contrast, our method maintains strong generalization, achieving the highest global accuracy across all datasets: 84.61% on FashionMNIST, 79.87% on SVHN, and 57.75% on CIFAR- 10. These represent improvements of 3.87%, 11. 20%, and 5.70%, respectively, over the second-best method on each dataset. Furthermore, our method consistently achieves the lowest ECE and NLL on global data, demonstrating both well-calibrated predictions and effective uncertainty quantification.
Trade-off Analysis: Figure 3 visualizes the tradc-off between local and global performance for per- sonalized models across the three metrics. For accuracy, the top-right region indicates lavorable performance; for ECE and NLL, better performance lies in the bottom-left region. Our method con- sistently appears closest to the optimal region in all three plots, demonstrating strong local accuracy without sacrificing global generalization. The plots confirm that our method achieves superior balance across all datasets, combining high accuracy with well-calibrated and confident predictions.
10

---

Published in Transactions on Machine Learning Research (01/2026)
Results of Global Models. We compare the performance of global models across all methods, excluding FedPop, which by design does not maintain global model. The results discussed below correspond to the last two groups of rows in Table 2 The results for global models on both local and global data are closely aligned, as the GM-on-LD setting reports the average performance the global model evaluated across all clients' local data. Our method performs strongly in the global setting, achieving the best results on SVHN and C-FAR-10 across all three metrics: accuracy, ECE, and NLL. This demonstrates not only high predictive performance but also well-calibrated and confident uncertainty estimates, highlighting the robustness of our approach even in non-personalized settings.
5.5
Client-Level Fairness Analysis
FashionMNIST
SVHN
CIFAR-10
100
Mean Accuracy
90
50
Ours (A
Fea/G
pFedBayes
Worst 10% Accuracy
100
PFodVom
so
70
30
Figure 4: Performance of the global model on clients' local data. Top row: average accuracy. Bottom row: worst-case performance. Fach column corresponds different dataset.
Despite the growing interest in BFL, its fairness implications remain largely overlooked aspect. Although not the main focus of this work, take a first step toward systematic investigation of client-level fairness from an accuracy-parity perspective. In Appendix D.1 we provide an extended literature review that motivates our choice of fairness metrics.
As a comprehensive fairness evaluation requires comparing shared models across clients under consistent conditions, we focus on the performance of the global model. It is important note that FedPop does not produce global model. Although achieves strong performance in terms of personalized models, the absence of single shared model renders it unsuitable for evaluating fairness from global perspective. As result, the notion of fairness across all clients is not applicable in its context.
Our results, presented in Figure show that the global model produced by our method is comparable to state- of-the-art approaches in terms of both mean accuracy and worst 10% accuracy across clients. Additionally, we conduct a study to examine the effect of Bayesian layers on algorithmic fairness in federated learning, and report our findings in Appendix D.2
6 Broader Applicability of the Proposed Method
Although the use of information geometry in BFL is natural, method also applies to any optimization problem defined over probability spaces rather than point estimates. The approach naturally extend to domain adaptation scenarios: when model is trained on souree data A and new target data B become available, retraining only on B may lead to catastrophic forgetting of the source domain. Our projection framework provides a principled trade-off, retaining prior knowledge while integrating new information. This flexibility further enhanced by the use of IVON, which combines the efficiency of deterministic learning with the ability to perform inference and adaptation in the posterior space. However, is important note that approach fundamentally differs from standard domain adaptation methods based optimal
11

---

Published in Transactions on Machine Learning Research (01/2026)
transport Courty et al. (2016). which operate on data distributions rather than model posteriors. Appendix provides additional insights and discussion.
7
Conclusions
In this work, we proposed personalized Bayesian Federated Learning method that balances local adapta- tion and global generalization while explicitly accounting for uncertainty quantification and model calibra- tion. Our approach combines variational learning with informttion-gooeee geometric tools, namely projection and barycenters, to produce client-specific models that remain robust under data heterogeneity with minimal additional cost. Through experiments on three benchmark datasets, FashionMNIST, SVHN, and CIFAR-10, we demonstrate that our method consistently achieves favorable trade-off between performance on local and global data distributions. In global evaluations, our personalized models outperform existing baselines in terms of accuracy, calibration, and uncertainty quantification, while maintaining highly competitive per- formance local data. This contrasis with other methods, which often lend to overfit underfit across clients. Furthermore, we show that the global model produced by our approach generalizes better than those of state-of-the-art alternatives. These empirical results highlight the effectiveness principled, information- geometric personalization combined with variational learning federated settings. Future work will explore adaptive mechanisms to control the degree of personalization, extend our personalization framework non- parametric BFL, and evaluate its performance under additional heterogeneity scenarios, such as feature and quantity shifts, as well as real-world cross-device FL benchmarks with naturally occurring client partilions.
Acknowledgments
The work has been supported by the EC through the Horizon Europe/JU SNS project, ROBUST-6G (Grant Agreement 101139068), the IMT "Futur, Ruptures & Impacts" program, and the European Research Council (ERC) under the European Union's Horizon 2020 Research and Innovation programme (Grant agreement No.101003431).
References
Maruan Al-Shedivat, Jennifer Gillenwater, Eric Xing, and Afshin Rostamizadeh. Federated Learning via Posterior Averaging: A New Perspective and Practical Algorithms. arXiv preprint arXiv:2010.05773, 2020.
Pedro C Asteran, E Del Barrio, JA Cuesta-Albertos, and C Matrán. A Fixed-point. Approach to Barycenters in Wasserstein Space. Journal of Mathematical Analysis and Applications, 441(2):744-702 2016.
Shun-ichi Amari. Information Geometry and Its Applications, volume 194. Springer, 2016.
Manoj Ghuhan Arivazhagan, Vinay Aggarwal, Aaditya Kumar Singh, and Sunav Choudhary. Federated Learning with Porsonalization Layers. arXio preprint 1012: 00818, 2019.
Julio Backhoff- Veraguas, Joaquin Fontbona, Gonzalo Rios, and Felipe Tobar. Bayesian Learning with Wasser- stein Barycenters. ESAIM: Probability and Statistics, 26:436-472, 2022.
Jeremy Bentham. Utilitarianism. Progressive Publishing Company, 1890.
Jeremy Bentham and John Stuart Mill. Utilitorianism and Other Essays. Penguin UK, 2004.
Shrey Bhatt, Aishwarya Gupta, and Piyush Rai. Federated Learning with Uncertainty via Distilled Predictive Distributions. In Asian Conference on Machine Learning, pp. 153 168. PMLR, 2024.
Rishi Bommasani and et al. On the Opportunities and Risks of Foundation Models. orXiv preprint arXiv:2108.072558 2021.
12

---

Published in Transactions on Machine Learning Research (01/2026)
Mahrokh Ghoddousi Boroujeni, Andreas Krause, and Giancarlo Ferrari Trecate. Personalized Federated Learning of Probabilistic Models: A PAC-Bayesian Approach. arXiv preprinl or Xix:2444 00.4400008551, 2024.
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wu, Tian Li, Jakub Konečný, H Brendan McMahan, Virginia Smith, and Ameet Talwalkar. LEAF: A Benchmark for Federated Settings, arXiv preprint arXiv:1812.01097 2018.
Longbing Cao, Hui Chen, Xuhui Fan, Joao Gama, Yew-Soon Ong, and Vipin Kumar. Bayesian Federated Learning: A Survey. arXw preprint arX::22304.133 13267, 2023.
Hong-You Chen and Wei-Lum Chao. Fedbe: Making Bayesian Model Ensemble Applicable to Federated Learning. arXiv preprinl arXiv:2009.0197, 2020.
Yiqiang Chen, Xin Qin, Jindong Wang, Chaohui Yu, and Wen Gao. Fedhealth: A Federated Transfer Learning Framework for Wearable Healthcare. IEEE Intelligent Systems, 35(4):83 100-933,- 2020.
Evgenii Chzhen, Christophe Denis, Mohamed Hebiri, Luca Oneto, and Massimiliano Pontil. Fair Regression with Wasserstein Barycenters. Advances in Neural Information Processing Systems, 33:7321-7331, 2020.
Luca Corinzia, Ami Beuret, and Joachim M. Buhmann. Variational Federated Multi-task Learning. arXiv preprint Xiv: 06268, 2019.
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal Transport for Domain Adaptation. IEEE transactions on pattern analysis and machine intelligence, 3999)11883-1865, 2016.
Csiszár. I-Divergence Geometry of Probability Distributions and Minimization Problems. The Annals of Probability, 3(1):146 158, 1975.
I. Csiszár and F. Matus. Information Projections Revisited. IEFE Transactions Information Theory, 49 (6):1474 1490, 2003.
Marco Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. Advances in Newral Information Processing Systems, 26, 2013.
Nico Daheim, Thomas Möllenhoff, Edoardo Maria Ponti, Iryna Gurevych, and Mohammad Emtiyaz Khan. Model Merging by Uncortainty-based Gradient Matching. arXiv preprint arXiv:2310. 1.0000880, 2023.
Julie Delon, Agnes Desolneux, and Antoine Salmona. Distributions. Journal of Applied Probubilily, 59(4):1178 1198, 2022.
Distances Between Gaussian
Janez Demšar. Statistical Comparisons of Classifiers Over Multiple Data Sets. Journal of Machine Learning Research, 7(Jan): 30, 2006.
Siddharth Divi, Yi-Shan Lin, Habiba Farrukh, and Z Berkay Celik. New Metrics to Evaluate the Performance and Fairness of Personalized Federated Learning. arXiv preprint 3rXiv:2107.13111 2021.
Alessandro D' D'Ortenzio, Costanzo Manes, and Umut Orguner. Fixed-point Iterative Computation of Gaussian Barycenters for Some Dissimilarity Mcasures. In 2022 International Conference Computational Science and Compuational Intelligence (CSCI), pp. 1422 1428, 2022.
Ivar Ekeland and Roger Temam. Conver Analysis and Varialional Problems. SIAM, 1999.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized Federated Learning: A Meta-Learning Approach. arXiv preprint arXiv:2002.00 07948, 2020.
Farzan Farnia, Amirhosscin Reisizadeh, Ramtin Podarsani, and Ali Jadbabaic. An Optimal Transport Approach to Personalized Federated Learning. IEEE Journal on Selected Areas in Information Theory, 3 (2):162-177) 171, 2022.
13

---

Published in Transactions on Machine Learning Research (01/2026)
John Fischer, Marko Orescanin, Justin Loomis, and Patrick McClure. Federated Bayesian Deep Learning: The Application of Statistical Aggregation Methods to Bayesian Models. ar preprint orXin:2403. 15263, 2024.
Wilfrid Gangbo and Andrzej Swigch. Optimal Maps for the Multidimensional Monge-kantorovich Prob- lem. Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences, 51(1):23 45, 1998.
Jiashi Gao, Ziwei Wang, Xiangyu Zhao, Xin Yao, and Xuetao Wei. Does Egalitarian Fairness Lead to Instabilit?? The Fairness Bounds in Stable Federated Learning under Altruistic Behaviors. Advances in Newral Information Processing Systems, 37:47849 77747749-47875, 2024.
Solenne Gaucher, Nicolas Schreuder, and Evgenii Chzhen. Fair Learning with Wasserstein Baryeenters for Non-decomposable Performance Measures. In International Conference on Arlificial Intelligence and Statistics, pp. 1430-2455 PMLR, 2023.
Charles Goddard, Shamane Siriwardhana, Malikeh Ehghaghi, Luke Meyers, Vlad Karpukhin, Brian Benedict, Mark McQuade, and Jacob Solawetz. Arcoe's Mergekit: A Toolkit for Merging Large Language Models. arXiv preprint arXiv:2403. 13257, 2024.
Han Guo, Philip Greengard, Hongyi Wang, Andrew Gelman, Yoon Kim, and Eric P Xing Federated Learning As Variational Inference: A Scalable Expectation Propagation Approach. arXiv preprint arXiv:2302. 0et223220.222. 2023.
Mohsin Hasan, Guojun Zhang, Kaiyang Guo, Xi Chen, and Pascal Poupart. Calibrated One Round Federated Learning with Bayesian Inference in the Predictive Space. In Proceedings of the AAAI conference on artificial intelligence, volume 38, pp. 12313-12321, 2024.
Conor Hassan, Robert Salomone, and Kerrie Mengersen. Federated Variational Inference Methods for Structured Latent Variable Models. arxiv preprint Xri.:23002 033114 2023.
Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. LoRA: Low-Rank Adaptation of Large Language Models. ICLR, 1(2):3, 2022a.
Zeou Hu, Kiarash Shaloudegi, Guojun Zhang, and Yaoliang Yu. Federated Learning Meets Aulti-Objective Optimization. JEEE Transactions on Network Science and Engincering, 9(4):2039-205) 2022b.
Nour Jamoussi, Giuseppe Serra, Photios A Stavrou, and Marios Kountouris. Information- Barycen- ters for Bayesian Federated Learning. arXiv preprint arXic:2412. 11646, 2024.
Owng-Hwan Jang, Sangdoo Yun, and Dongyoon Han. Model Stock: All We Need Is Just a Few Fine-tuned Models. In European Conference on Computer Vision, pp. 207-223. Springer, 2024.
Eunjeong Jeong and Marios Kountouris. Personalized Decentralized Federated Learning with Knowledge Distillation. In ICC 2023 IEEE International Conference on Communications, pp. 1982-1987, 2023.
Rahif Kassab and Osvaldo Simeone. Federated Generalized Bayesian Learning via. Distributed Stein Varia- tional Gradient Descent. IEEE Transactions on Signal Processing, 70:2180 2192, 2022.
Zixuan Ke, Yijia Shao, Haowei Lin, Tatsuya Konishi, Gyuhak Kim, and Bing Liu. Continual Pre-training of Language Models. arXiv preprint arXiv:2302.0322 2023.
Mohammad Emtiyaz Khan and Håvard Rue. The Bayesian Learning Rule. arXiv preprint arXiv:2107.04562 2021.
Minyoung Kim and Timothy Hospodales. Fodhb: Hicrarchical Bayesian Federated Learning. arXie preprint 044909930 2023.
James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska- Barwinska, et al. Overcoming Catastrophic Forgetting in Neural Networks. Proceedings of the national academy of sciences, 114(13): 3521333552351666 2017.
14

---

Published in Transactions on Machine Learning Research (01/2026)
Günther Koliander, Yousef El-Laham, Petar M Djurié, and Franz Hlawatsch. Fusion of Probability Density Functions. Proceedings of the IEEE, 110(4))404-453, 2022.
Nikita Kotelevskii, Maxine Vono, Alain Durmus, and Eric Moulines. Fedpop: A Bayesian Approach for Personalised Federated Learning. Advonces in Neural Information Processing Systems, 35:8687-8701, 2022.
Khang Le, Dung Q Le, Huy Nguyen, Dat Do, Tung Pham, and Nhat Ho. Eutropic Gromov- W Wasserstein between Gaussian Distributions. In Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Reseurch, pp. 12164 12203. PMLR, 17 23 Jul 2022.
Qinbin Li, Bingsheng He, and Dawn Song. Practical One-shot Federated Learning for Cross-silo Setting. arXiv preprint ar (xXt:20100000011 2020.
Wenqian Li, Shuran Fu, and Yan Pang. Distributionally Robust Federated Learning with Wasserstein Barycenter. In The Second Tiny Popers Track at ICLR 2024, 2024.
Tao Lin, Lingjin Kong, Sebastian U Stich, and Martin Jaggi. Ensemble Distillation for Robust Model Fusion in Federated Learning. Advances in Neural Informalion Processing Systems, 33:2351 2363, 2020.
Liangxi Liu, Xi Jiang, Feng Zheng, Hong Chen, Guo-Jun Qi, Heng Huang, and Ling Shao. A Bayesian Federated Learning Framework with Online Laplace Approximation. IEFE Transactions on Pattern Analysis and Machine Intelligence, 2023.
Eric Maskin. A Theorem on Utilitarianism. The Review of Economic Studies, 45(1):93 96, 1978.
Michael S Matena and Colin A Raffel. Merging Models with Fisher Fisher-weighted Averaging. Advances in Neural Information Processing Systems, 35:17703-17716, 2022.
Michael McCloskey and Neal J. Cohen. Catastrophic Interference in Connectionist Networks: The Sequential Learning Problem. In Psychology of Learning and Motivation, volume 24 of Psy. of Le. and Mot., pp. 109 165. Academic Press, 1989. doi: https://doi.org/ 1orii/orr/10.111666
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Commmicationn-fficient Learning of Deep Networks From Decentralized Data. In Artificial intelligence and statistics, pp. 1273-1282. PMLR, 2017.
Ninarch Mehrabi, Fred Morstatter, Nripsuta Saxena, Kristina Lerman, and Aram Galstyan. A Survey on Bias and Fairness in Machine Learning. ACM computing surveys (CSUR), 54(6):1 35, 2021.
Frank Nielsen and Richard Nock. Sided and Symmetrized Bregman Centroids. IEEE Transactions on Information Theory, 55(6))2882 2904, 2009.
Atahan Ozer, Kadir Burak Buldu, Abdullah Akgül, and Gozde Unal. How to Combine Variational Bayesian Networks in Federated Learning. arXiv preprint arXi::2206.1088 2022.
Shivam Pal, Aishwarya Gupta, Saqib Sarwar, and Piyush Rai. Simple and Scalable Federated Learning with Uncertainty via Improved Variational Online Newton. In OPT 2024: Optimization for Machine Learning, 2024.
Brendan Pass. Multi-Marginal Optimal Transport: Theory and Applications. ESAIM: Mathematical Mod-
elling and Numerical Analysis, 49(6):17711 1177771799,,,,,11111111 2015.
Peijie Qiu, Wenhui Zhu, Sayantan Kumar, Xiwen Chen, Xiaotong Sun, Jin Yang, Abolfazl Razi, Yalin Wang, and Aristeidis Sotiras. Multimodal Variational Autoencoder; a Barycentric View. arXiv preprint 2xXiu:2412.20455 20487 2024.
John Rawls. Some Reasons for the Maximin Criterion. The American Economic Review, 64(2):141114 146, 1974.
John Rawls. A Theory of Justice. Revised Edition, 1999.
15

---

Published in Transactions on Machine Learning Research (01/2026)
Felix Sattler, Klaus-Robert Müller, and Wojeiech Samek. Clustered Federated Learning: Model-Agnosts Distributed Multitask Optimization Under Privacy Constraints. IEFF Transactions on Newral Nebworks and Learning Systems, 32(8))37100 3722, 2020.
Nripsuta Ani Saxena, Karen Huang, Evan DeFilippis, Goran Radanovic, David C Parkes, and Yang Liu. How Do Fairness Definitions Fare? Examining Public Attitudes Towards Algorithmic Definitions of Fairness. In Proceedings of the 2019 AAAI/ACM Conference on AI, Ethics, and Society, pp. 99-106, 2019.
Yuesong Shen, Nico Daheim, Bai Cong, Peter Nickl, Gian Maria Marconi, Clement Bazan, Rio Yokota, Iryna Gurevych, Daniel Cremers, Mohammad Emtiyaz Khan, et al. Variational Learning Is Effective for Large Deep Networks. arXiv preprint orXiv:2402. 1 1741, 2024.
Haizhou Shi, Zihao Xu, Hengyi Wang, Weiyi Qin, Wenyuan Wang, Yibin Wang, Zifeng Wang. Sayna Ebrahimi, and Hao Wang. Continual Learning of Large Language Models: A Comprehensive Survey. ACM Computing Surveys, 2024.
Yuxin Shi, Han Yu, and Cyril Leung. Towards Fairness-aware Federated Learning. IEEE Transactions on Newral Nelworks and Learning Systems, 2023.
Siddharth Swaroop, Mohammad Emtiyaz Khan, and Finale Doshi- Connecting Federated ADMM to Bayes. Xiv preprint or arxi:22001. 1100001..732 17325, 2025.
Alysa Ziying Tan, Han Yu, Lizhen Cui, and Qiang Yang. Towards Personalized Federated Learning. IEEE Transactions on Neural Networks and Learning Systems, 34(12):9587 9603, 2023.
Cédric Villani. Optimal Transport: Old and New, volume 338. Springer, 2009.
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated Learning with Matched Averaging. arXiv preprint arXiv:2002.06; 06440, 2020a.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H. Vincent Poor. Tackling the Objective Incon-
sistency Problem Heterogeneous Federated Optimization. Advances in Neural Information Processing Systems, 33:7611-7623, 2020b.
Zheng Wang, Xiaoliang Fan, Jianzhong Qi, Chenglu Wen, Cheng Wang, and Rongshan Yu. Federated Learning with Fair Averaging. arXiu preprint ar 1tt:2104.14990 2021.
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Contijo-Lopes, Ari S Mor- cos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model Soups: Averaging Weights of Multiple Fine-tuned Models Improves Accuracy Without Increasing Inference Time. In Inler- national conference on machine learning, pp. 23965 23998. PMLR, 2022.
Miao Yang, Ximin Wang, Hongbin Zhu, Haifeng Wang, and Hua Qian. Federated Learning with Class Imbalance Reduction. In 2021 29th European Signal Processing Conference (EUSIPCO), pp.
IEEE, 2021.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Bayesian Nonparametric Federated Learning of Neural Networks. In International Conference on Machine Learning, pp.
PMLR, 2019.
Guojun Zhang, Saber Malekmohammadi, Xi Chen, and Yaoliang Yu. Proportional Fairness in Federated Learning. arXiv preprint arXit:2202.01666 2022a.
Xu Zhang, Yinchuan Li, Wenpeng Li, Kaiyang Guo, and Yunfeng Shao. Personalized Federated Learning via Variational Bayesian Inference. In International Conference on Machine Learning, pp. 26293 26310. PMLR, 2022b.
Junyi Zhu, Xingchen Ma, and Matthew Blaschko. Confidenco-aware Personalized Federated Learning via Variational Expectation Maximization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pallern Recognition, pp. 25542-24511, 2023.
16

---

Published in Transactions on Machine Learning Research (01/2026)
A Proof of Theorem
Proof. In our proof, we directly address the problem through the KKT conditions, leveraging the well-known
variational structure of projection problems (see Eq. 5.14, Ekeland & Temam (1999)). ) This approach allows
us highlight. the key concepts and ideas underlying the proof without unnecessary technical overhead. All cited Theorems refer Ekeland & Temam (999)..
Weak or Strong Duality: We first apply Proposition 5.1 (generalization of Slater's condition) show that the projection problem is stable. Subsequently, Proposition 2.2 then implies that it is also normal. Finally, Lemma 5.2 verifies the hypotheses of Proposition 2.1, which yields zero duality gap. Consequently, the max-min and min-max formulations attain the same finite optimal value, establishing strong duality.
Relations between the min-max and max-min solutions: Theorem 5.1 ensures that, for fixed Ta- grangian multipliers, if one can identify a primal variable that satisfics the saddle-point condition, then this variable solves the primal problem for an appropriate constraint value. As shown in the proof, this is equiv- alent to minimizing the Lagrangian with respect to the primal variable while keeping the multipliers fixed. This guarantees that the solution obtained via the barycenter coincides with the solution of the original projection problem.
We restate the problem for clarity. The projection problem is defined as:
min D(plpg) s.t. D(plp)) S rk- PEM
(11)
Assumption 2 ensures that (111p)) and (-IIp) are convex in their first argument. Consequently, the problem is a convex optimization problem with convex incquality constraint.
Because E Sg and the interior of Sg is nonempty (e.g., any sufficienly close to ps lies in the interior). the constraint admits a Slater point. Therefore, the KKT conditions are both necessary and sufficient for optimality.
The Lagrangian for problem 11
C(p, A) D(plpg) + X(D(plpa) rk).
A 2 0.
(12)
Let p* = Py be the minimizer of The KKT conditions are:
Primal feasability: D(p* S rk.
2. Dual feasibility: X 2
3. Complementary slackness:
x*(((p"pp)) -k) =
(13)
Since varying the radius F& parametrizes the geodesic between Pg and pk, and Py & Se in any nontrivial personalization setting, the solution necessarily lies on the boundary, i.e., D(p"lpo Tk which implics X > 0.
4. Stationarity:
+'V,",")+++, = O.
(14)
Divide equation
14 by (1+ ++):
I+ X- VpD(p))))) + 1+ vpD(pTpp)) = O.
(15)
Define normalized nonnegative weights
1
X
Wg
Wy
(16)
1 X
1
17

---

Published in Transactions on Machine Learning Research (01/2026)
Then equation 14
becomes:
Wi (yDP((p')pp) + we D(p 0.
(17)
The barycenter of {pg, pk with weights wy wy is defined as the minimizer of:
min (wplo)) + (pppp))) PEM
(5)
Since the objective convex, its minimizer must satisfy the first-order optimality condition:
Wig v + us v D(p = 0.
(18)
This exactly equation 17 obtained from the KKT stationarity condition. Thus, the minimizer of the constrained projection problem coincides with the barycenter:
p" - Pb pat, (Wg wx)).
(19)
This concludes the proof.
18

---

Published in Transactions on Machine Learning Research (01/2026)
B Comparison between the Different Aggregation Methods Considered
We conduct experiments with various combinations of global update and personalization methods, including EAA, RKLB, and WB for global aggregation, and RKLB and WB for personalization. Across all config- urations, we observe consistent trends and comparable results. To confirm these observations, we apply the Wilcoxon signed-rank test Demšar (2006) to compare the performance of different aggregation methods across multiple runs (random sceds) and datasets. We compute pairwise Wilcoxon tests between all methods to assess whether their performance differences indicating, for each pair of methods, whether one consistently outperforms the other. The results of the Wilcoxon signed-rank tests, presented in Figure 5 show that in most cases the p-values are high, indicat- ing no statistically significant difference between the aggregation methods. This suggests that, despite small fluctuations in performance across datasets, the aggregation methods perform similarly overall, and no single aggregation approach consistently outperforms the others in a statistically meaningful way across the con- sidered metrics. To further illustrate the variability and stability the aggregation methods across different datasets, we provide violin plots in Figure [ These plots show the distribution of global model accuracies across multiple random sceds for cach aggregation method on FashionMNIST,, SVHN, and CIFAR-10. The shape and spread of the distributions oller insight into both the typical performance (median) and variabil- ity (spread) of each method. Based on these results, which show that the WB-based aggregation achieves competitive and consistently stable performance across datasets, we restrict the experiments reported in the main text to this method for both the global update and personalization steps.
statistically significant. This yields malrix of p-values
1.00
WB RKLB EAA
WB RKLB EAA
WB RKLB CAA
0.75
0.75
0.75
0.50
0.50
0..50
0.25
0.25
0.25
EAA
RKLB WB
0.00
EAA RKLB
WB
0.00
EAA
RKLB WB
0.00
(a) Accuracy
(b) ECE
(c) NLL
Figure 5: Wilcoxon signed-rank test p-values comparing aggregation methods across all datasets for three evaluation metrics: (a) accuracy, (b) ECE, and (c) NLL. Lower p-values indicate statistically significant
differences between methods. Only the lower triangle of each matrix shown to avoid redundancy.
90
Accuracy
Accuracy
Accuracy
85
86
60-
86
84-
50-
RKLB WB
EAA
RKLB
Aggregation Method
Aggregation Method
Aggregation Method
(a) FashionMNNST
(b) SVHN
(c) CIFAR-10
Figure Accuracy distributions across random seeds for different aggregation methods on: (a) FashionM- NIST, (b) SVHN, and (c) CIFAR-10 datascts. lower spread and higher median indicate better and more stable performance.
19

---

Published in Transactions on Machine Learning Research (01/2026)
C Effect of À on the Trade-off off between Performance on Local and Global Data
To complement Section 5.3 in the main text, report in Figures 7 and the results illustrating the effect of À on the trade-off wetween performance on local and global data for the PashionMNIST and SVHN datasets.
Averge prffrmance
Average performnce
performance
prrfermacce
Accoratt
030
0.13
0.13
0.5
0109
100 1433 105
1.5
1.5
0.00
325 0.
1.0
100 11.5 163
10
1.2
20
Figure 7: Effect of À on the trade-off between performance on local and global data distributions. Results are reported for the FashionMNIST dataset. Notably, À = 0 corresponds to the global model, whereas + oc corresponds to the local model.
Average performance
Average performnnce
performance
performance
201
Accurary
0.15
0.09
100 1655 055
1.5
1.5
2.5
0.0 3255 0.5
1.0
100 11.5
18
14
25
Figure 8: Effect of on the trade-off between performance on local and global data distributions. Results are reported for the SVHN dataset. Notably, 0 corresponds to the global model, whereas + 00 corresponds to the local model.
20

---

Published in Transactions on Machine Learning Research (01/2026)
D Client-Level Fairness Analysis
Client heterogeneity introduces significant challenges in distributed learning architectures, one of which is ensuring algorithmic fairness. Algorithmic fairness refers to the impartial treatment of individuals or groups in automated decision-making processes, without bias, discrimination, or favoritism based on innate or ac- quired characteristics Saxena al. (2019) : Mehrabi et al. (2021). In the context of FL, where models are collaboratively trained data from multiple sources representing diverse populations, overlooking fairness considerations can amplify perpetuate existing societal biases in the global model, leading to unfair out- comes for certain subgroups. To mitigate this issue, growing body of research has proposed fairness-aware FL algorithms, including approaches based on personalized FL Tan et al. (2023) ; Jeong & Kountouris (2023). which tailor models to individual client data distributions. [Shi et al. (2023) presents a comprehensive taxon- omy of faireess-aware FT, strategies, covering key aspects such as client selection, optimization, contribution assessment, and incentive allocation.
D.1 Related Work
Fairness in Federated Learning. Fairness is a multifaceted concept that spans several disciplines, includ- ing the social sciences, law, machine learning, and statistics, each offering distinct perspectives and implicitly different definitions. In the context of machine learning, Mehrabi et al. (2021) provides comprehensive overview of key fairness notions, categorizing them into three primary types: individual fairness, which seeks to ensure similar outcomes for similar individuals; group fairness, which focuses on equal treatment across predefined demographic groups; and subgroup fairness, which combines elements of both individual- and group-based approaches to better capture fairness across a broader range of population segments.
In FL, fairness extends beyond the behavior of a single predictive model encompass the equitable treatment of clients participating in the distributed training process. Because clients contribute heterogeneous data and resources, an FL system is considered fair if it avoids systematically privileging disadvantaging certain participants. To formalize fairness in this context, we draw on two long-standing philosophical traditions, utilitarianism and egalitarianism, which have recently been adapted to the federated setting Zhang et al. (022a)); Hu et al. (202bb) Wang et al. (2021 0 Gao et al. (2024). yet remain underexplored within the BFL. paradigm.
Utilitarianism, rooted in the works of Bentham and Mill Bentham 18900) Bentham & Mill (0044) and formalized by Maskin Maskin (1978). evaluates the fairness of a system through the aggregate welfare it produces. In FL, this corresponds to optimizing global utility, for example, the average or mean accuracy across clients.
Egalitarianism, inspired by Rawls' difference principle Rawls 1974 1999). instead focuses on protecting the worst-off participants. Translated to FL, this notion aligns with max max-min fairness, which aims to maximize the minimum (worst-case) performance across clients, typically evaluated through the worst-10% accuracy metric.
Wasserstein Barycenters for Fairness. In the domain of fair learning, recent work has highlighted the effectiveness of WBs as unifying framework for promoting fairness in ML while preserving strong predictive performance. Chzhen et al. (2020) investigated fair regression by establishing a direct connection between the demographic parity constraint and the WB problem. They showed that the optimal fair predictor, min- imizing the squared error while ensuring independence from sensitive attributes, can be dcrived as the WB of the conditional distributions corresponding to different sensitive groups. This approach leads to simple yet powerful post-processing technique that transforms any regression model into fair one without requir- ing retraining. Importantly, their method provides robust, distribution-fres fairness guarantees, improving fairness metrics with only minimal reductions in predictive accuracy.
Building on this foundation, Gaucher et al. (0223) extended the use of WBs to classification tasks. They demonstrated that demographic parity classification can be achieved by solving fair regression problem, followed by appropriate thresholding. Their approach underscores the role of the WB in aligning group- wise distributions, thereby reducing disparities across sensitive attributes. In settings with binary sensitive
21

---

Published in Transactions on Machine Learning Research (01/2026)
attributes, the baryeenter plays central role in determining optimal classification thresholds, enabling lavorable trade-off between fairness and performance. Together, these works provide both theoretical justification and practical algorithms for transport-based fairness, offering cohesive framework applicable to both regression and classification.
Our work connects with these findings by incorporating WBs during the aggregation phase. We extend this line of research by analyzing WBs through the lens of client-level fairness definitions, offering complemen- tary perspective to the discussions presented in this section.
D.2 Effect of Variational Inference Layers on Client-Level Fairness
In this subsection, we investigate how increasing the number of variational inference layers (i.e.; Bayesian layers implemented via variational inference) in the model architecture affects fairness outcomes. To this end, we focus on the BA-BFL setting implemented within the Hybrid Bayesian Deep Learning framework Jamoussi et al. 2024)) alongside the standard FedAvg baseline.
It is important to note that both the WB and RKLB aggregation strategies reduce to simple arithmetic avoraging in the deterministic limit (i.c., when posterior variances vanish). Consequently, BA-BFL with zero Bayesian layers is equivalent to FedAvg, enabling a unified comparison Figure To ensure a fair compar- ison, we restrict our analysis to BA-BFL and FedAvg, deliberately excluding additional mechanisms specific IVON-based approach, since the differing optimizers prevent
to other Bayesian methods, including direct comparison with FedAvg in this context. This controlled setup allows for focused examination of how the degree of Bayesian modeling, quantified by the number of variational inference layers, influences the fairness of the global model across different clients. However, our analysis does not reveal a consistent trend across experiments, suggesting that the relationship between the number of Bayesian layers and client-level fairness may depend on other factors, such as task complexity.
40
Worst 10% Accuracy
Accracy
SVINN
70-
RCLB SUIN
RCBB -3AR-100
WR FIRIST
RCLD TVBJST
10
Number Bayosian
Numbbe Bayesian Layers
Figure 9: Impact of Variational Inference Layers on Fairness.
22

---

Published in Transactions on Machine Learning Research (01/2026)
E Beyond BFL: Broader Applicability of the Proposed Method
E.1 Merging Variational Foundation Models
Foundation Models (FMs) have emerged as powerful gencral-purpose learners, capable of adapting to a wide range of rownstreea tasks after large-scale pretraining. However, as data distributions shift and new domains emerge, keeping these models up to date without retraining from scratch remains major
challenge Bommasani & et al. (2021). Continual pretraining Ke et al. (2023): Shi et al. (2024), fine-tuning
Hu et. al. (2022a). and model merging McMahan et al. (2017); Matena & Rallel 2022 Daheim et al. (2023)
Goddard el al. (024)) offer promising paths forward, enabling FMs to integrate new know ledge while retaining
broad generalization. However, updating such large systems at scale faces key challenges: computational
constraints, catastrophic forgetting McCloskkyy & Cohon (199)), and potential misalignment in uncortainty
quantification. Addressing these challenges calls for principled and efficient updale rules that incorporate domain-specific adaptations into a global FM while preserving statistical rigor and scalability.
Within the spectrum of adaptation strategies, variational approaches provide rigorous framework for rep- resenting model uncertainly, making them particularly suitable in contexts where both reliability and in- terpretability are critical. We focus on variational FMs whose parameters encode posterior distributions, enabling updates to be expressed as operations on the statistical manifold of distributions. We assume these models are pretrained or fine-tuned using the IVON optimizer Shen et al. (2024), chosen for its ability to match Adam's computational cost while delivering strong Bayesian performance in large-scale settings. IVON has been shown to pretrain GPT-2 on OpenWebText and ResNet on ImageNet from scratch, as well as to fine-tune large masked language models (e-g., DeBERTaV3). Within this setting, we formu- late FM merging an onto a sphere centered at a specialized model, i.c, a fine-tuned model. This formulation naturally extends our earlier notion of specialization, which we showed to be equivalent to computing a barycenter in the variational parameter space, toward multi-model aggregation via barycentric averaging that minimizes the average iotormaaiionggometiin discrepancy across multiple fine-tuned models. Our approach thus provides an interpretable and theoretically grounded merging mechanism that generalizes existing techniques such as Fisher-weighted averaging Matena & Raffel (2022), as well as mixture and product-of-expers methods.
projection from a global model, i.e, the pretrained model,
Related Work: Model Merging. Originally proposed in the context of FL to mitigate communication overhead and enhance privacy McMahan et. al. (2017), model merging has since been adopted in both computer vision and large language models Goddard et al. (2024). Wortsman et al. (0022) demonstrated that averaging the weights of models fino-tuned under varied hyperparameters can improve both accuracy and
robustness. Using only a few fine-tuned models, Stock Jang et al. (2024) achieved robust merges via layer-wise linear interpolations that explicitly operate in the Euclidean geometry of the parameter space. In contrast, our method focuses the manifold geometry of the variational posteriors.
Generalization via D-Barycenters. Let {p denote the variational posteriors (e.g., IVON-trained posteriors of fine-tuned FMs), and let px denote the predictive distribution of the 4th FM.
Forward KL. With D KL(pxll),, the minimizer is the mizlure in posterior space: (0) IN We (0). After marginalizing over 0, the resulting predictive distribution also mixes pointwise as Pb x)) = E wk (2 x), a construction commonly referred to as the Miature of Experts.
Reverse KL. With D = (gllpp)), the solution is the log-opinion pool or Product of Experts: (0) x IT (0) In exponential families, this corresponds to natural-parumettr averaging. For Ganssians posteriors, A = Σ wn Aş and μ* = (At)- "' W/ AgHle where A₁ denotes the precision matrix of the tth model. This formulation connects directly to Fisher merging Matena &
Raffel (2022).
Wasserstein-2. With D = W3(pxlle), the minimizer is the Wasserslein-2 burycenter. For Gaussians
distributions, the barycenter remains Gaussian with Vwz = (I wxlit) Mw3 = Σi= Wgfli often yielding more robust. summaries than naive parameter averaging.
23

---

Published in Transactions on Machine Learning Research (01/2026)
} Practical implications. Barycentric merging provides single, interpretable control (the weights {wx) to balance global and domain-specific knowledge. It recovers popular FM-merging schemes as special cases and admits closed-form solutions for common variational families (e.g., diagonal Gaussians) under widely used divergences. When complemented by the IVON training regime, the weights can be instantiated from curvalure estimates, such that higher-curvature models receive grealer weight Daheii el. al. (2023) Shen ct al. 2024). Consoquently, high-uncertainty (low-curvature) models are down-weighted during aggregation.
Limitations and Possible Extensions. Like most merging methods in distribution space (Bayesian) parameter space (deterministic), our approach assumes architecturally aligned models, i.c., compatible layers and widths, to enable layer-wise aggregation. This constraint is particularly limiting for FMs, where specialized models are often smaller than the pretrained backbone. As a next step, we aim to relax this assumption through
iassert optimal transport maps (e.g., see Delon et al. (2022); Le et al. (2022)), which enable mappings between spaces of different dimensionalities. Furthermore, we plan conduct FM-scale experiments to assess the scalability and efficiency of the method in large-scale set sings.
E.2 Incremental Learning
We design a toy example to simulate a continual learning scenario involving two sequential tasks: Task A, consisting of the first five classes of MNIST, and Task B, corresponding the remaining five. We first train a model (Model A) on Task A. Subsequently, Task B becomes available, while access Task A is revoked, mimicking lypical continual learning setup where previously seen data due to constraints such data retention policies regulations like the General Data Protection Regulation (GDPR). We evaluate multiple strategics: training new model (Model B) solely on Task B. fine-tuning Model A on Task B with and without Elastic Weight Consolidation (EWC) Kirkpatrick et al. (2017), and computing the barycenter of Model and Model B. As shown in Figure 10 the barycentric combinations offer a favorable trade-off between performance on both tasks, preserving accuracy on Task while adapting effectively to Task B, thereby mitigating the catastrophic forgetting observed in the other scenarios.
longer accessible
1.0
0.8
Accuracy task
Model A
0.6
Model B
Model A finetuned on task B
0.4
Model A finetuned on task B with EWC
WB( Model A, Model B, [1/2,1/2]]
0.2
RKLB( Model A, Model B, (1/2,1/21)
0.0
0.8
Accuracy task
Figure 10: Accuracy Trade-off in an Incremental Learning Setting
24

---

Published in Transactions on Machine Learning Research (01/2026)
F Additional Discussion on Barycenters
Wasserstein Barycenters in Machine Learning. The Wasserstein Barycenter (WB) has emerged as a powerful tool in a wide range of machine learning applications, including Bayesian learning, multimodal representation learning, and fair learning.
The Bayesian Wasscrstein Barycenter (BWB), introduced in Bacckoof-Veraguss et al. (2022): minimizes the Wasserstein Bayes risk corresponding to the p- D-Wasserstein distance the loss, yielding a predictive posterior distribution with lower variance compared to Bayesian Model Averaging (BMA), i.e., the Bayes estimator that minimizes the Bayes risk under the L2 distance the KL divergence. In addition, the BWB enhances interpretability by respecting the geometric structure of the model space. It extends naturally to nonparametric model spaces and can be computed efficiently using stochastic gradient descent.
the context of multimodal learning, Qiu et al. (2024) redefines the aggregation unimodal inference distribut in Variational Autoencoders (VAEs) using the WB. Here, the WB is formulated as the central distribution that minimizes the average discrepancy in terms of the squared 2- Wasserstein distance:
N
1/" arg min vEP(X)
wk W2 (2R,,,), V),
wk 1
with W2 denotes the 2- Wasssrrseni distance. Unlike traditional aggregation methods such as the Product of Experts Mixture of Experts, which rely the asymmetric KL divergence, the WB better preserves the geometric structure of distributions and enables smooth interpolation across modalities. In multimodal VAEs, WB-based models such as WB-VAE and its variant MWB- VAE achieve superior classification accuracy and conditional generation performance, particularly on datasets with missing modalities.
The works Backhooff Vercauuf-Veeaguss el al. 20222 Qiu el al. (2024) collectively highlight the flexibility and scalability of WB-based methods. Their results demonstrate the WB's potential to unify diverse distributions
while maintaining geometric interpretability and robustness.
Forward KL Barycenter. As noted in Jamoussi al. (2024). the Forward Kullback-Leible (FKL) barycenter of Gaussian distributions, one of the a-divergence barycenters is not itself Gaussian. While it is possible to approximate the FKL barycenter with a Gaussian distribution through projection onto the space of Gaussians D'Ortenzio et al.] [2022), this approach fails to preserve parameter independence, as shown Figure [1 making it a suboptimal choice for practical implementation. The following equations characterize the parameters of the Gaussian approximation to the FKL Barycenter for a set of N Gaussian distributions with parameters {pk, Ex)
EPLL Σ
wh (Ex + (px MFKL) ))
PFKL
Wkfli-
Wasserstein-32 Barycenter
Approximated Barycenter
Barycenter
1000
L.25
0.875
0.750
0.750
0.625
0.15
0.625
0.500 8
0.12
0.500
0.375
0.09
0.50
0.08
0.125
0.03
0.125
0.000
3.00
0.000
Figure 11: Example of barycenters between two multivariate Gaussian distributions with independent pa- rameters.
25


Reviews:
[{"review_id": "GGWSIJ2Erh", "reviewer": "Reviewer_gqwS", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a novel approach to Personalized Federated Learning (PFL) based on an information-geometric projection framework for personalization in parametric Bayesian Federated Learning (BFL). The main idea appears to be original and (potentially) impactful; however, the empirical validation is performed solely on synthetic datasets, which makes the paper far less compelling then it could have been if it used real-world datasets that are available in the literature.", "claims_and_evidence_yn": "No", "claims_explanation": "As the evaluation is done solely on synthetic datasets, it is unclear to what extent the results are truly meaningful and/or applicable in real-world scenarios. See suggestion below on evaluating on real-world datasets", "audience_yn": "Yes", "audience_explanation": "In general, Federated Learning has many practical applicabilities. However, the current evaluation on synthetic datasets does NOT raise to the standards of a journal publication, especially when real-world datasets are available (see the paper's reference [Tan et al, 2023])", "requested_changes": "I will begin by quoting two consecutive paragraphs from \"Section V\" in [Tan et al, 2023]:\n\n>>> There are several FL benchmarking frameworks developed in recent years, including FLBench [75], Edge AIBench [76], OARF [77] and FedGraphNN [78]. LEAF [79] is one of the earliest and most popular benchmarking frameworks proposed for FL. At the time of writing, it provides six FL datasets covering a range of machine learning tasks including image classification, language modeling and sentiment analysis under both IID and non-IID settings. Examples datasets include the Extended MNIST [80] dataset split according to the writers of the character digits, the CelebA [81] dataset split according to the celebrity, and the Shakespeare [9] dataset split according to the characters in the play. A set of accuracy and communication metrics, along with implementation references for well-known approaches such as FedAvg, SGD and MOCHA are also provided. As LEAF extends existing public datasets from traditional machine learning settings, it does not fully reflect the data heterogeneity in FL scenarios. Although there are a few real-world federated datasets, such as a street image dataset for object detection [82] and a species dataset for image classification [83], they are often limited in size.\n\n>>> Despite the release of benchmark datasets for FL, they are not widely adopted in PFL research. The vast majority of PFL studies choose to simulate the non-IID setting by performing their own partitioning on a public benchmark dataset used in machine learning (e.g., MNIST [84], EMNIST [80], CIFAR-100 [85]), or creating a synthetic dataset [17], [64], [68]. \n\nIn this reviewer's opinion, this situation has become untenable: with the proliferation of PFL papers, it is simply unacceptable to continue evaluating the work on (semi-)synthetic datasets, especially when far more realistic ones are available [Tan et al, 2023]. The authors have a unique opportunity to perform a comprehensive evaluation on realistic/real-world data; if their results hold in these existing FL domains, they will set the bar against which all future approaches will be evaluated. If the current dataset are useful in making a particular point, they can be kept in the paper alongside the real-world ones. \n\nOTHER COMMENTS:\n- the authors should add 1-2 paragraphs explaining the relationships between the metrics used the paper and those from Table IV in [Tan et al, 2023]\n- please add an intuitive paragraph discussing in-depth why the proposed approach does not perform on Fashion-MNIST as well as on the other two datasets\n-", "broader_impact_concerns": "n/a", "additional_comments": "n/a", "extra_scores": {}, "extra_text": {}}, {"review_id": "ZKjVfYLmkl", "reviewer": "Reviewer_km55", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors' work contribute to Federated Learning (FL), in particular, to Private FL (PFL). The workflow in PFL consists of 3 iterative stages:\n1. A *learning phase* in which clients train privately their own local models with their own private data.\n2. An *aggregation phase* in which the server aggregates the local model parameters retrieved from each client into an updated global model. \n3. An *update phase* in which the server sends the updated global model parameters to each client.\n\nThe authors focus on the 3rd stage, allowing for the *personalization* of the updated global model. Roughly speaking, each client has the choice to replace their local model parameters by the updated global model parameters, keep their own local model parameters, or find a compromise between both. To achieve this, given a statistical divergence function, they propose to use a projection of the global model onto a ball centered at the local model; the choice of radius determines the trade-off between one model and the other. As this projection step could be possibly intractable, the authors show it can be recast as computing a suitable barycenter with respect to the same statistical divergence. The personalization step thus adds flexibility to the PFL workflow. To ensure that the implementation is tractable, they make the strong assumption that all models are multivariate normal distributions with diagonal covariance matrices. This allows them to use explicit expressions for the barycenters for the statistical divergences that they use.", "claims_and_evidence_yn": "No", "claims_explanation": "I find the work interesting and the evidence provided by their experiments sufficient. However, the proof of their main theoretical claim, namely, Theorem 1, seems to be lacking. The authors start from the projection problem, which is convex under Assumption 2 in the paper. They then form a Lagrangian to connect the problem of computing a projection to that of computing a barycenter. However, for this connection to be *an equivalence* the authors need to prove at least that strong duality holds. In the notation in Appendix A, one always has *weak duality*, that is,\n\n$$\\sup\\_{\\lambda \\geq 0}\\inf\\_{p} \\mathcal{L}\\_{\\lambda}(p) \\leq \\inf\\_{p} \\sup\\_{\\lambda \\geq 0} \\mathcal{L}_{\\lambda}(p).$$\n\nComputing the projection is equivalent to the problem in the right-hand side in the above inequality, whereas computing the barycenter for fixed $\\lambda$ corresponds to the problem\n\n$$\\inf\\_{p} \\mathcal{L}\\_{\\lambda}(p).$$\n\nHowever, it could be the case that *strong duality* does not hold, in which case\n\n$$\\sup\\_{\\lambda \\geq 0}\\inf\\_{p} \\mathcal{L}\\_{\\lambda}(p) < \\inf\\_{p}\\sup\\_{\\lambda \\geq 0} \\mathcal{L}\\_{\\lambda}(p).$$\n\nThis would imply that the best choice for $\\lambda$ does not match the minimal divergence obtained when computing the projection. In particular, one would *not* have an equivalence between computing a projection and computing a barycenter. This gap needs to be addressed by proving that, for the problems that the authors consider, strong duality actually holds. \n\nNote that even if strong duality holds, the optimal distributions for the left and right-hand sides may not be the same. For this stronger conclusion *total duality* must hold. \n\nTo be clear, my contention is not with the main contribution of the paper. In fact, I think that there is no issue in proposing the computation of a barycenter between the global and local model as a method for personalization since, as the authors also discuss, this enables a trade-off between local and global models. My main issue is with Theorem 1, for which the proof is, at this stage, incomplete.", "audience_yn": "Yes", "audience_explanation": "Private Federated Learning is a topic of interest for the community. Enabling personalization for each one of the clients is also a topic of interest, particularly as there is a trade-off in improving the performance of local models on local data at the expense of poor generalization, and possibly suboptimal global models.", "requested_changes": "1. In pg. 3, the second to last paragraph, change \"distributions ith respect\" by \"distributions with respect\".\n2. In Section 3.1, is the parametric distribution family $\\mathcal{Q}$ the same for each client and the server? I suppose this is the case, but I suggest stating this explicitly for an audience less familiar with PFL. \n3. In Section 3.2, for (6) it is stated that $h$ in the denominator is \"the Hessian approximation\". However, (6) shows a relation between scalars. It seems that either there is a typo, or an explanation of what $h$ is is missing.\n4. Appendix A is not referenced in the main text. It should be stated that the proof of Theorem 1 is provided in Appendix A. \n5. The proof of Theorem 1 must be corrected as indicated before. \n6. Appendix C is not referenced in the main text.\n7. Why do the authors use the parameter $\\lambda$ directly to characterize the trade-off? It seems that using $\\lambda/(1 + \\lambda)$ or $1/(1 + \\lambda)$ would be more interpretable, as these parameters take values on $[0, 1]$ instead of $[0, \\infty]$, and computing the barycenters uses these weights, as discussed in Remark 3.\n8. The experimental results show the effect of personalization, i.e., the value of $\\lambda$, on the local models. However, I did not find a similar discussion for the global model. One would expect that in the limit of overpersonalization, $\\lambda \\to \\infty$, the global model never gets truly updated beyond the first step. \n9. As far as I understand, the authors consider the same level of personalization for each client. What is the effect of letting each client select their own personalization factor?", "broader_impact_concerns": "I do not have broader impact concerns due to the nature of the contribution. However, the authors do provide a nice discussion about potential issues related to fairness in Apoendix D, particularly when overpersonalization may generate poor performance on underrepresented data.", "additional_comments": "No comments.", "extra_scores": {}, "extra_text": {}}, {"review_id": "7MYGf3jwMF", "reviewer": "Reviewer_huuH", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a novel personalization technique for Bayesian Federated Learning (BFL) utilizing information-geometric projection. The core concept involves projecting the global posterior onto a divergence-defined sphere encompassing a client’s local posterior, thereby enabling a customizable trade-off between global generalization and local specialization. Under mild convexity assumptions, the authors establish that this projection is equivalent to computing a two-point barycenter, resulting in a closed-form solution. This approach eliminates the need for additional fine-tuning, computation, or data access, making it “cost-free” personalization. The method is evaluated within a variational learning framework using the IVON optimizer and compared against several BFL baselines across three vision datasets. The results demonstrate favorable trade-offs in both global and local performance.\n\n**Strengths**\n- Frames personalization as an information-geometric projection and connects projection to barycenters via a clear theoretical result (Theorem 1).\n- As demonstrated through meticulous experimental validation, the method not only competes with but frequently surpasses the existing baselines.\n\n**Weaknesses**\n- Results rely heavily on Gaussian mean-field posteriors to achieve \"cost-free\" personalization.\n- The optimizer is highly sensitive to hyperparameters.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Overall, the core claims are well-supported,\n- The theoretical equivalence between projection and barycenter is clearly stated and relies solely on convexity (Assumption 2).\n- Experiments do support the claim that the method improves the global/local trade-off.\n- The entire experimental section is clearly justified, as it provides detailed information about the selection and construction of baselines.", "audience_yn": "Yes", "audience_explanation": "The paper sits at the intersection of:\n- Bayesian deep learning\n- Personalized federated learning\n- Information geometry\n\nThe method presents a novel conceptual perspective, offering practical benefits and ease of interpretation.  Furthermore, the experiments evaluate common benchmarks, suggesting this work provides a valuable approach to addressing data heterogeneity.", "requested_changes": "No adjustments are required.", "broader_impact_concerns": "There are no ethical concerns to raise regarding this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)