Dirichlet-based Uncertainty Quantification for Personalized Federated Learning with Improved Posterior Networks

1 Introduction

The widespread adoption of deep neural networks in various applications requires reliable predictions, which can be achieved through rigorous uncertainty quantification. Although uncertainty quantification has been extensively studied in different domains under centralized settings [Lee et al., 2018; Lakshminarayanan et al., 2017; Gal and Ghahramani, 2016; Kotelevskii et al., 2022a], only a few works have considered this area within the context of federated learning (FL) [Linsner et al., 2022; Kotelevskii et al., 2022b]. Typically, in FL papers, algorithms result in using either a personalized local model or a global model. However, both these models could be useful in different cases by providing the tradeoff between the personalization of a local model and the higher reliability of the global one [Hanzely and Richt´arik, 2020; Liang et al., 2019].

In this paper, we introduce a new framework to choose whether to predict with a local or global model at a given point based on uncertainty quantification. The core idea is to apply the global model only if the local one has high epistemic (model) uncertainty [H¨ullermeier and Waegeman, 2021] about the prediction at a given point, i.e., the local model doesn't have enough information about the particular input point. In case the local model is confident (either in predicting a particular class or in the fact that it is observing an ambiguous object with high aleatoric (data) uncertainty [H¨ullermeier and Waegeman, 2021]), it should make the decision itself without involving the global one.

As a specific instance of our framework, we propose an approach inspired by Posterior Networks (PostNet; [Charpentier et al., 2020]) and its modification, Natural Posterior Networks (NatPN; [Charpentier et al., 2022]). This model is particularly useful for our purposes, as it enables the estimation of aleatoric and epistemic uncertainties without incurring additional inference costs. Thus, we can fully implement the switching between local and global models efficiently.

Related work. The literature presents various approaches to uncertainty quantification in FL settings. In [Linsner et al., 2022], the authors suggest that training an ensemble of K global models is the most effective for federated uncertainty quantification. Despite its effectiveness, this approach is K times more expensive compared to the classic FedAvg method. Another proposal comes from [Kotelevskii et al., 2022b], where the authors recommend using Markov Chain Monte Carlo (MCMC) to obtain samples from the posterior distribution. However, this method is computationally demanding due to its significant computational complexity.

Other works, such as [Chen and Chao, 2021; Kim and Hospedales, 2023], also present methods that could potentially estimate uncertainty in FL. However, these papers do not expressly address the opportunities and challenges of uncertainty quantification in their discussion. It's important to note that there are existing approaches of deferring classification to other models or experts in case of abstention, for example [Keswani et al., 2021]. Despite this, none of these approaches have been explored in a federated context, nor have they considered the corresponding constraints.

The central idea of PostNet and NatPN involves using Dirichlet prior and posterior distributions over the categorical predictive distributions [Malinin and Gales, 2018;

Malinin and Gales, 2019; Charpentier et al., 2020; Charpentier et al., 2022; Sensoy et al., 2018]. To parameterize the parameters of these Dirichlet distributions, the authors suggest introducing a density model over the deep representations of input objects. In NatPN, a Normalizing Flow [Papamakarios et al., 2021; Kobyzev et al., 2020] is employed to estimate the density of embeddings extracted by a trained feature extractor. This density is then used to calculate updates to the Dirichlet distribution.

Despite the success of the NatPN model [Charpentier et al., 2022], we identified certain issues with the loss function employed in NatPN, which become particularly critical when dealing with high aleatoric uncertainty regions. Intriguingly, the issue was not discovered in the authors' recent work [Charpentier et al., 2023], which investigated potential issues in the training procedure. Recent works [Bengs et al., 2022; Bengs et al., 2023] has unveiled other potential issues related to the training of Dirichlet models in general but have not offered solutions to address these challenges.

Contributions. Our paper stands out as one of the few studies that consider the usage of uncertainty quantification in FL scenarios. We aim to use uncertainty estimates to improve the prediction quality in personalized FL. Our contributions:

1. We present a new FL framework that is based on uncertainty quantification, which allows switching between using a personalized local or global model. To achieve that, we consider several types of uncertainties for local and global models and propose a procedure that ensures that the model is confident in its prediction if the prediction is made. Otherwise, our framework rejects the prediction if we are confident that there is a high uncertainty in the predicted class. To the best of our knowledge, we are the first who introduce the concept of switching between models based on uncertainty quantification in FL.
2. We introduce a specific realization of our framework FedPN, using Dirichlet-based NatPN model. For this particular model, we identify an issue in the loss function of NatPN (not known in literature before) that complicates disentanglement of aleatoric and epistemic uncertainties, and propose a solution to rectify it.
3. We conduct an extensive set of experiments that demonstrate the benefits of our approach for different input data scenarios. In particular, we show that the proposed model outperforms state-of-the-art personalized FL algorithms in the presence of out-of-distribution data while being on par with them in standard scenarios.

2 General Framework of Switching between Global and Personal Models

In this section, we introduce our framework, discussing the general idea and potential nuances. In the subsequent section, we will delve into a specific implementation.

2.1 Concept Overview

We are going to consider a FL setting with multiple clients, each having its personalized local model. We additionally assume that a non-personalized global model is also available.

We suppose that local and global models share the same architecture, but were trained differently. The global model was trained using a FL algorithm (e.g. FedAvg [McMahan et al., 2017]), hence using data (in a federated manner) from all the clients. Local models were trained by each of the clients locally, using only local data of a particular client. We assume both these models are available for each of the clients. The global model is typically expected to perform reasonably well on each client's data. Assuming that we already have trained global and local models, we consider a situation where clients have the option to use either the global model or their local models to make predictions for a new incoming object x.

The choice between local and global models for prediction depends on the multiple factors that contribute to their prediction quality. First of all, a shift in the distribution between the local data of a particular client and the global population may have a significant effect on the models' performance. Possible shifts include covariate shift, label shift, or different types of label noise. If the shift is significant, the global model might be very biased in the prediction for the particular client, while normally the local model is unbiased.

The second part of the picture is the size of the available data. Generally, the global model has more data and potentially, if no data shift is present, should outperform the local ones. However, the global model is usually trained with no direct access to the data stored by clients, which might degrade its performance. Eventually, the best-performing model will be the one that achieves a better bias-variance trade-off.

In this work, we propose a framework to choose between the pointwise usage of a local or global model for prediction. This decision is based on the uncertainty scores provided by the model. This approach aims to mitigate issues that arise when there is a distribution shift between a client's distribution and the global distribution, which can often hinder model performance. By using uncertainty to guide the model selection, we can enhance the model's ability to perform well even in the presence of such distribution shifts.

Both the local and global models can provide uncertainty estimates. These estimates include not only the total predictive uncertainty but also separate aleatoric and epistemic uncertainties. This feature allows for a more refined decisionmaking process, ensuring the most suitable model is used for prediction or choosing to abstain from making a prediction altogether in the face of high uncertainty. The resulting workflow is summarized in Figure 1.

It is worth noting, that our work is not the first which considers a joint usage of global and local models. However, we are the first to consider switching strategies between them that are based on uncertainty quantification. For example, one of the prior works [Hanzely and Richt´arik, 2020] employs the average distance of local models to their collective mean as a regularization technique, linking local models. Conversely, our approach relies on either purely local training or the global model, depending on model uncertainty for each sample. Additionally, the "averaged model" in their research is not akin to a global model in design; it does not guarantee universal performance across all data and is more similar to the MAML [Finn et al., 2017] approach.

2.2 Optimal Choice of Model

We start from the important fact that the local model is not exposed to the data shift between the general population and the particular client. Thus, if this model is sufficiently confident in the prediction, there is no need to involve the global model at all. However, it is important to distinguish between different types of uncertainty here. Usually, the total uncertainty of the model predicting at a particular data point can be split into two parts: aleatoric uncertainty and epistemic uncertainty [H¨ullermeier and Waegeman, 2021]. Aleatoric uncertainty is the one that reflects the inherent randomness in the data labels. The epistemic uncertainty is the one that reflects the lack of knowledge in the model choice because the model was trained on a data set of a limited size, or model misspecification. In what follows we will refer to epistemic as to "lack of knowledge" uncertainty, as the models we are using, neural networks, are known to be very flexible.

For our problem, it is extremely important to distinguish between aleatoric and epistemic uncertainties. Suppose the local model has low epistemic and high aleatoric uncertainty at some point. In that case, the model is confident that the predicted label is ambiguous, and the model should abstain from prediction. However, if the epistemic uncertainty is high, it means that the model doesn't have enough knowledge to make the prediction (not enough data), and the global model should make the decision. The global model, in its turn, may either proceed with the prediction if it is confident or abstain from prediction if there is high uncertainty associated with the prediction. Thus, in this context, for a fixed client and an unseen input, there are four interesting outcomes, see Table 1.

Note, that when local epistemic uncertainty is low, there is no need to refer to the global model. The assumption here is that the local model is better than the global one if it knows the input point well. Thus we consider only two options for the local model: low epistemic with low aleatoric and low epistemic with high aleatoric regardless of the confidence of Table 1: Possible scenarios for the input data point in the introduced setup. A particular input falls into one of the categories depending on the confidence in its prediction by local and global models. The disentanglement between aleatoric and epistemic uncertainties is crucial to optimally making the decision.

the global model. When the local epistemic uncertainty is high, it means that the client barely knows the input point, and thus we refer to the global model. For the global model, we look at the total uncertainty as we only care about the error of prediction which is determined by total uncertainty.

The particular implementation of the approach described above would depend on the choice of the machine learning model and the way to compute uncertainty estimates. The key feature required is the ability of the method to compute both aleatoric and epistemic uncertainties. In the next section, we propose the implementation of the method based on the posterior networks framework [Charpentier et al., 2020].

3 Dirichlet-Based Deep LearningModels

We have chosen to showcase Dirichlet-based models [Malinin and Gales, 2018; Charpentier et al., 2020; Charpentier et al., 2022] as a specific instance of our general framework. The intuition behind this decision lies in the fact that these models allow the distinction between various types of uncertainty and facilitate the computation of corresponding uncertainty estimates with minimal additional computational overhead. Furthermore, unlike ensemble methods [Beluch et al., 2018], there is no need to train multiple models. In comparison to approximate Bayesian techniques, such as MC Dropout [Gal and Ghahramani, 2016], Variational Inference [Graves, 2011] or MCMC [Izmailov et al., 2021], almost all expectations of interest can be derived in closed form and almost without computational overhead. This makes Dirichlet-based models an attractive and efficient option for implementing our proposed framework.

3.1 Basics of Dirichlet-BasedModels

In this section, we introduce the basics of Dirichlet-based models for classification tasks. To ease the introduction, let us start by considering a training dataset D = {xi, yi}Ni=1 , where N denotes the total number of data points in the dataset. We assume that labels yi belong to one of the K

classes. Typically, the Dirichlet-based approaches assume that the model consists of two hierarchical random variables, µ and θ. The posterior predictive distribution for a given unseen object x can be computed as follows:

where p(y | µ) is the distribution over class labels, given some probability vector (e.g., Categorical), p(µ | x, θ) is the distribution over a simplex (e.g., Dirichlet), and p(θ | D) is the posterior distribution over parameters of the model. However, for practical neural networks, the posterior distribution p(θ | D) does not have an analytical form and is computationally intractable. Following [Malinin and Gales, 2018], we suggest considering a "semi-Bayesian" scenario by looking at a point estimate of this distribution: p(θ | D) = δ(θ - ˆ θ), where ˆ θ is some estimate of the parameters (e.g., MAP estimate). Then the integral inside the brackets simplifies:

∫ p(µ | x, θ) p(θ | D)dθ = ∫ p(µ | x, θ) δ(θ - ˆ θ)dθ = p(µ | x, ˆ θ). In the series of works [Malinin and Gales, 2018; Malinin and Gales, 2019; Charpentier et al., 2020; Charpentier et al., 2022] the posterior distribution p(µ | x, ˆ θ) is chosen to be the Dirichlet distribution Dir(µ | αpost (x)) with the parameter vector αpost (x) = αpost (x | ˆ θ) that depends on the input point x. In these models, the prior over probability vectors µ takes the form of a Dirichlet distribution, representing the distribution over our beliefs about the probability of each class label. In other words, it is a distribution over distributions of class labels. This prior is parameterized by a parameter vector αprior , where each component αprior c corresponds to our belief in a specific class. PostNet [Charpentier et al., 2020] and NatPN [Charpentier et al., 2022] propose the idea that the posterior parameters αpost (x) can be computed in the form of pseudo-counts that are computed by a function:

where α(x) = α(x | ˆ θ) is a function of input object x that maps it to positive values.

Parameterization of α(x). In NatPN it is proposed to use the following parameterization:

In this parameterization, g(x) represents a feature extraction function that maps the input object x (usually highdimensional) to a lower-dimensional embedding. Subsequently, p(·) is a "density" function (parameterized by normalizing flow in the case of [Charpentier et al., 2022; Charpentier et al., 2020]), and f(·) is a function mapping the extracted features to a vector of class probabilities.

This parameterization offers several advantages. Firstly, since p(·) is expected to represent the density of training examples, it should be high for in-distribution data. Secondly, as the density is properly normalized, embeddings that lie far from the training ones will result in lower values of p(g(x)), thus leading to lower α(x). This means that for such input x, we will not add any evidence, and consequently, αpost (x) will be close to αprior .

3.2 UncertaintyMeasures for Dirichlet-Based Models

One of the advantages of using Dirichlet-based models is their ability to easily disentangle and quantitatively estimate aleatoric and epistemic uncertainties.

Epistemic uncertainty. We begin by discussing epistemic uncertainty, which can be estimated in multiple different ways [Malinin and Gales, 2021]. In this work, following the ideas from [Malinin and Gales, 2018], we quantify the epistemic uncertainty as the entropy of a posterior Dirichlet distribution, which can be analytically computed as follows

where αpost (x) = [αpost 1 (x), . . . , αpost K (x)], ˜ ψ is a digamma function and αpost 0 (x) = ∑Ki=1 αpost i (x).

Aleatoric uncertainty. Aleatoric uncertainty can be measured using the average entropy [Kendall and Gal, 2017], which can be computed as follows:

This metric captures the inherent noise present in the data, thus providing an estimate of the aleatoric uncertainty.

3.3 Loss Functions for Dirichlet-BasedModels

The loss function used in [Charpentier et al., 2022; Charpentier et al., 2020] is the expected Cross-Entropy (CE), also known as Uncertain Cross Entropy (UCE) [Biloˇs et al., 2019]. This loss function is a strictly proper scoring rule, which implies that a learner is incentivized to learn the true conditional p(y | x). For a given input x, the loss can be written as:

where αpost 0 (x) = ∑Ki=1 αpost i (x). Additionally, authors suggest penalizing too concentrated predictions, by adding the regularization term with some hyperparameter λ. The overall loss function looks as follows:

where H denotes the entropy of a distribution. This overall loss function is referred to by authors as Bayesian loss [Charpentier et al., 2022; Charpentier et al., 2020].

Issue with the loss function. Let us explore an asymptotic form of (5). For all x > 0, the following inequality holds:

Recalling the update rule (1) and using the specific parameterization of αprior c = 1 for all c, we conclude that all αpost c (x) > 1. Hence, we can approximate ˜ ψ(αpost c (x)) ≈ logαpost c (x).

To simplify the notation, we will use z = g(x) and denote by fy(z) predicted probability for the y-th (correct) class:

see the full derivation in Appendix, Section A.2.

We see from (7) that for an in-distribution case with high aleatoric uncertainty (all classes are confused and equally probable), the last term is canceled. Note, that the presence of entropy term (that is used in [Charpentier et al., 2022; Charpentier et al., 2020] resulting in the final loss of equation (6)), which incentivizes the learner to produce smooth prediction will only amplify the effect. This implies that no gradients concerning the parameters of the density model will be propagated. As a result, p(·), which is the density in the embedding space, may disregard regions in the embedding space that correspond to areas with a high concentration of ambiguous training examples. This violates the intuition of p(·) as a data density. Thus, uncertainty estimates based on p(·) cannot be used to measure epistemic uncertainty, as it ignores the regions with high aleatoric uncertainty.

In addition to the problem of confusing high-aleatoric and high-epistemic regions, we discovered that the proposed loss function defies intuition when confronted with "outliers." We define an "outlier" as an object with a predicted probability of the correct class is less than 1K . From (7), we see that the last term changes its sign precisely at the point where fy(z) = 1K . This implies that to minimize the loss function, we must decrease p(z) at these points, which seems counterintuitive as these z values correspond to objects from our training data.

Although the issue concerning equation (5) arises primarily in the asymptotic context, we can readily illustrate it through the loss function profiles (see Figure 2-left) for a range of fixed correct class prediction probabilities.

To further emphasize the problem, we examine the loss function landscape and provide a demonstrative example. Consider three two-dimensional Gaussian distributions, each with a standard deviation of 0.1 and centers at (-1, 0), (1, 0), and (0, 1) (see Figure 2-center). Left and right clusters are set to include objects of only one class, while the middle distribution contains uniformly distributed labels, representing a high aleatoric region. Each Gaussian contains an equal amount of data. We train the NatPN model using a centralized approach, employing both the loss function from equation (5) and the loss function with our fix; see equation (8). Subsequently, we evaluate the quality of the learned model by plotting the density in the center of the middle Gaussian. Ideally, we would expect the density to be equal for different numbers of classes K. However, we observe a different picture with a "density" estimate at the central cluster decreasing when one increases the number of clusters, see Figure 2-right. In this Figure, we plot the median of the estimated density for different loss functions. In the next section, we introduce a simple trick, which rectifies the behavior of the loss function, while without the correction density vanishes with K. We refer the reader to an additional toy example on image data in Appendix, Section B.6.

It is essential to emphasize that addressing these issues is critical for our framework, as we need to accurately differentiate between aleatoric and epistemic uncertainties to select the appropriate model for a given situation. In the following section, we propose a simple but efficient technique to rectify the aforementioned problem with the loss function, ensuring that our framework distinguishes between the different types of uncertainties and makes informed model choices.

Proposed solution. We propose to still use the parametric model to estimate density, but now our goal is to ensure that p(z) accurately represents the density of our training embeddings z = g(x). To achieve this, we propose maximizing the likelihood of the embeddings explicitly by incorporating a corresponding term into the loss function. Simultaneously, we aim to prevent any potential impact of the Bayesian loss on the density estimation parameters, maintaining their independence. Thus, we suggest the following loss function:

where λ, γ > 0 are hyperparameters, and StopGradp(g(x)) means that the gradient will be not propagated to the parameters of a density model, which parameterizes p(g(x)).

We should note that while the corrected loss function may look somewhat ad hoc, it is computationally simple to implement, which is very important for modern deep learning. Compared to the original loss function, our rectified loss does not introduce any additional computational overhead - in both cases (for the old loss and the new one) the gradient will be computed only once for each of the parameters.

4 FedPN: Specific Instance of the Framework

In the section, we consider specific instances of our framework. However, it is not limited to only one instance. Other possible options are discussed in Appendix, Section B.2.

4.1 Federated Setup

In this section, we show, how one can adapt NatPN for FL. Suppose we are given an array of datasets Di for 1 ≤ i ≤ b, where b represents the number of clients. Each Di = {xij ∈ Rd , yij }|Di| j=1 consists of object-label pairs. We construct our federated framework in such a way that all clients share the feature extractor g parameters ϕ and maintain personalized heads f i parameterized by θi. Furthermore, we retain a "global" head-model f, which is trained using the FedAvg [McMahan et al., 2017] method and which ultimately has parameters θ.

Following the NatPN approach, we employ normalizing flows to estimate the embedding density. As it is for head models, here we also learn two types of models - a local density model pi (using local data) parameterized by ψi and a global density model p (using data from other clients in a federated fashion) parameterized by ψ. It is important to note that both types of models are trained on the same domain since the feature extractor model is fixed for local and global models. We refer readers to Appendix, Section B.3. for the discussion of computational overhead.

4.2 Threshold Selection

Before discussing the results, it is essential to understand how we determine when to make predictions using a local model or a global one. This decision, resulting in a "switching" model, depends on a particular uncertainty score. This score can either be the logarithm of the density of embeddings, obtained using the density model (normalizing flows), or the entropy of predictive Dirichlet distribution (3). We found that both measures provide comparable behavior, and in the experiments for Table 2 we use the density of embeddings.

To apply this approach, we must establish a rule for how a client decides whether to use its local model for predictions on a previously unseen input object x or to delegate the prediction to the global model. One approach is to select some uncertainty values' threshold. This threshold can be chosen based on an additional calibration dataset. In our experiments, we split each client's validation dataset in a 40/60 ratio, using the smaller part for calibration.

The choice of the threshold is arguably the most subjective part of the approach. Ideally, we would desire to have access to explicit out-of-distribution data (either from another client "local OOD" or completely unrelated data "global OOD"). With this data, we could explicitly compute uncertainty scores for both types of data (in-distribution and outof-distribution) and select the threshold that maximizes accuracy. However, we believe it is unfair to assume that we have this data in the problem statement. Therefore, we suggest a procedure for choosing the threshold based solely on available local data.

To choose the threshold, we assume that for all clients, there might be a chance that some p% (typically 10%) of objects are considered outliers. We further compute the estimates of epistemic uncertainty (with either entropy or the logarithms of the density of embeddings) and select an appropriate threshold based on this assumption for each of the clients. For the high epistemic uncertainty points of the global model, similar thresholding can be performed to optimize its prediction quality. Additionally, we conducted an ablation study on threshold selection (see Appendix, Section B.5).

5 Experiments

For experiments, we employ seven diverse datasets: MNIST [LeCun et al., 1998b], FashionMNIST [Xiao et al., 2017], MedMNIST-A, MedMNIST-C, MedMNIST-S [Yang et al., 2021; Yang et al., 2023], CIFAR10 [Krizhevsky, 2009], and SVHN [Netzer et al., 2011]. The LeNet-5 [LeCun et al., 1998a] encoder architecture is applied to the first five datasets, while ResNet-18 [He et al., 2016] is used for CIFAR10 and SVHN. Building on the ideas from [Charpentier et al., 2020; Charpentier et al., 2022], we implemented the Radial Flow [Rezende and Mohamed, 2015] normalizing flow due to its lightweight nature and flexibility. In all our experiments, we focus on a heterogeneous data distribution across clients. We consider a scenario involving 20 clients, each possessing a random subset of 2 or 3 classes. The overall amount of data each client possesses is approximately equal. The FL process is conducted using FedAvg algorithm.

In the following sections, we present different experiments that highlight the strengths of our approach. We should note that we don't have a dedicated experiment to illustrate how the approach deals with local ambiguous data as the vision datasets we are considering have very few points of this type.

5.1 Assessing Performance of the SwitchingModel

In this section, we assess the performance of our method, where the prediction alternates between local and global models based on the uncertainty threshold.

For each dataset, we have three types of models. First, a global model is trained as a result of the federated procedure, following FedAvg procedure. Then, every client, after the federated procedure, retains the resulting encoder network while the classifier and flow are retrained from scratch using only local data. The third type of model, the "switching"

model, alternates between the first two models based on the uncertainty threshold set for each client.

Additionally, for each client, we consider the following three types of data: data from the same classes used during training (InD), data from all other classes (OOD), and data from all classes (Mix). For each of these datasets and data splits, we compute the average prediction accuracy (clientwise). The results of this experiment are presented in Table 2. Separate results on InD and OOD data are shown in Appendix, Section B.4.

In this experiment, we compare the performance (accuracy score) of different (personalized) FL algorithms: FedAvg [McMahan et al., 2017], FedRep [Collins et al., 2021], PerFedAvg [Fallah et al., 2020], FedBabu [Oh et al., 2022], FedPer [Arivazhagan et al., 2019], FedBN [Li et al., 2021], APFL [Deng et al., 2020]. Note, however, that our proposed method is the only which allows natural criteria for switching and selecting either the local or global models. For others, incorporation of the logic is not straightforward and might be a topic of further research.

Worth emphasising, that the state-of-the-art performance on the in-distribution data is not the ultimate goal of our paper, our method performs on par with other popular personalized FL algorithms. The remarkable thing about our approach is that it can be reliably used on any type of input data.

For our method, FedPN, we used the "switching" model. From Table 2, we observe that our model's performance for InD data is typically comparable to the competitors. For "Mix" data, our approach is the winner by a large margin for almost all datasets. Note, that this data split is the most realistic practical scenario - all the clients aim to collaboratively solve the same problem, given different local data.

However, due to the heterogeneous nature of the betweenclients data distribution (covariate shift), local models cannot learn the entire data manifold. Therefore, occasionally referring to global knowledge is beneficial while still preserving personalization when the local model is confident.

Note, that for CIFAR10 all the methods are not working well. For our method, it means that either the learned density model does not distinguish well between in- and out-ofdistribution data, or the threshold was not chosen accurately. In Table 3 we present results for our model on InD data when only local models were applied (so no "switching" is used, compared to Table 2). This might be useful in scenarios when we know that data comes from the distribution of the local data. In this case, there is no sense to switch to the global model. We see, that despite it was not the purpose of the work, our approach performs on par with other competitors, slightly outperforming them on some datasets. An additional ablation study on the performance of each model on different splits is in Appendix, Section B.7.

6 Conclusion

We proposed a personalized FL framework that leverages both a globally trained federated model and personalized local models to make final predictions. The selection between these models is based on the confidence in the prediction.

Empirical evaluation demonstrated that, under realistic scenarios with both InD and OOD data present, this approach outperforms both local and global models when used independently. While the model's capacity to handle OOD data is not perfect and depends on various factors, such as the quality of the global model and the selection of the threshold, our "switching" approach leads to improved performance.

Acknowledgements

The methodological part of research (Sections 3 and 4) was supported by the Russian Science Foundation grant 20-7110135.