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

We thank the reviewer for their thorough review and feedback. We address their questions below.

1. Novelty is limited. Most of the framework is directly taken from previous works. Dirichlet-based model formulation is the same as [1] and [2], and the criteria are based on [3] and [4]. The only modification seems to be the stop-gradient operation based on the analysis on page 6~7.

We kindly disagree with this assessment. Our main innovation is suggesting a new approach to switching between local and global models based on uncertainty quantification. We believe that this methodological novelty is the main contribution of our paper. Everything else is already an implementation of a general idea, and we agree that we mostly use techniques that were known previously. However, the problem with the loss function identified by us was not known before despite this loss function has been used in a series of high impact papers. Thus, we think that it is another non-trivial contribution of our paper.

2. Cannot scale to more complex datasets like CIFAR-10. Significant drop in the in-distribution data in Table 2… An improvement from 10+ to 20+ cannot mitigate the performance drop from 70+ to 50+. This challenges the model ability in this Dirichlet-based framework, especially for even more complex datasets like CIFAR-100 and Tiny-ImageNet.

We want to emphasize that it is not the problem of the particular method. All the datasets dropped in quality on this dataset for a given specific data heterogeneity. Please note that our method still improved upon the others.

3. Additional computation. Learning a density model adds more work for the computer. Even for a simple model, the reason why the model doesn't perform well on hard datasets might be that this density model is not very strong. But making this special model stronger would require more computation.

We agree with this point. Learning an additional density model is indeed of extra cost and may be considered as the limitation of the model. However, we should emphasize that not only the density model is important, but also the feature extractor itself. If the resulting features are good, then we don’t need to train a stronger density model given the intuition that representations are well clustered. Also, sometimes simple density models work pretty well [1, 2].

Additionally, adding a density model is not just for the sake of improving the performance. It allows us to receive good uncertainty estimates and distinguish between different types of it. So the extra price is justified.

[1] Bertrand Charpentier, Oliver Borchert, Daniel Zügner, Simon Geisler, and Stephan Günnemann. Natural posterior network: Deep bayesian predictive uncertainty for exponential family distributions. In International Conference on Learning Representations, 2022.

[2] Vazhentsev, Artem, et al. "Efficient Out-of-Domain Detection for Sequence to Sequence Models." Findings of the Association for Computational Linguistics: ACL 2023. 2023.

Q1: Are there any explanations on why you can use entropy and average entropy as the criteria for Epistemic uncertainty and Aleatoric uncertainty respectively in Section 3.2?

The idea of using the entropy of the predicted Dirichlet distribution is rooted in the paper of [1, 2]. There is a connection between the sampling from the posterior of weights and the simplex, parameterized by Dirichlet. Specifically, Bayesian Model Averaging involves sampling from the posterior over weights, resulting in the posterior predictive distribution. This sampling from the posterior is an implicit sampling from the simplex, and we can parameterize the distribution over this simplex directly. The spread over the simplex is an indicator of the disagreement between models and thus can be used as an estimate of epistemic uncertainty. This spread can be measured, for example, as the entropy of the predictive Dirichlet. About aleatoric uncertainty - we follow ideas from [3]. Since aleatoric uncertainty is well-defined only for in-distribution samples, then every sample of the weights from posterior over weights should result in some ambiguous distribution over the simplex. That is why we find entropy first (which will be big for ambiguous samples), and then average it over all possible samples. Thus we result in the expected entropy as in equation (4).

[1] Malinin, A., & Gales, M. (2018). Predictive uncertainty estimation via prior networks. Advances in neural information processing systems, 31.

[2] Malinin, A., & Gales, M. (2019). Reverse kl-divergence training of prior networks: Improved uncertainty and adversarial robustness. Advances in Neural Information Processing Systems, 32.

[3] Kendall, A., & Gal, Y. (2017). What uncertainties do we need in bayesian deep learning for computer vision?. Advances in neural information processing systems, 30.

Thanks the authors for the rebuttal. Regarding the experimental results, I acknowledge that FedPN outperforms the baselines in smaller datasets, but I observe a more pronounced trade-off between InD and OOD performance in FedPN on CIFAR-10. Specifically, the improvement in OOD performance from 0 to 28% doesn't seem to fully compensate for the decline from 80% to 60% in comparison to the baselines. This raises my curiosity about FedPN's effectiveness on more complex datasets like CIFAR-100 and Tiny-ImageNet.

In terms of novelty, while I recognize some new analysis in your work, I find it to be somewhat incremental.

As such, I am inclined to maintain my original score.

We thank the reviewer for their thorough review and feedback. We address their questions below.

1. The methodological novelty is limited to applying existing approaches to a new problem.

We kindly disagree with this assessment. We are not considering any new problem but rather suggest a new approach to switching between local and global models based on uncertainty quantification. We believe that this methodological novelty is the main contribution of our paper. Everything else is already an implementation of a general idea, and we agree that we mostly use techniques that were known previously. However, the problem with the loss function identified by us was not known before despite this loss function has been used in a series of high impact papers. Thus, we think that it is another non-trivial contribution of our paper.

2. Ablation for the threshold selection.

This is an interesting and important point, and we are currently working on the corresponding experiment. As soon as it is ready, we will add it to the text.

• This also leads me to a concern about the conceptual scheme described in 2.1: a misestimated threshold may result in abstaining from predictions in a large portion of the data. Perhaps, the global model should not abstain at all?

This is true, that misestimated threshold may lead to the abstaining for a large portion of the data. However, the models (local and global) have different thresholds. One can set the threshold for the global moel high enough so that it rarely abstains. However, the applications that imply the possibility of abstention of any predcition (for example, when it possible to refer to human expert) may benefit from the abstention of the global model.

3. I am not sure about best practices in image classification these days, but the datasets seem simple.

Thank you for your comment. We wish to emphasize that the datasets we employed are standard in the federated learning (FL) community and widely recognized as adequate for evaluating federated algorithms. Seminal works in FL, such as FedProx [1], SCAFFOLD [2], and FedOpt [3], have also utilized these datasets. We believe our evaluation effectively demonstrates the viability of our proposed approach.

We acknowledge the reviewer's concerns regarding the complexity of the datasets. We will endeavor to include additional evaluations on a more complex dataset by the end of the discussion period or in the camera-ready version of our paper. However, we anticipate that the results will align with those already presented in our manuscript.

[1] Li, T., Sahu, A. K., Zaheer, M., Sanjabi, M., Talwalkar, A., & Smith, V. (2020). Federated optimization in heterogeneous networks. Proceedings of Machine learning and systems, 2, 429-450.

[2] Karimireddy, S. P., Kale, S., Mohri, M., Reddi, S., Stich, S., & Suresh, A. T. (2020, November). Scaffold: Stochastic controlled averaging for federated learning. In International conference on machine learning (pp. 5132-5143). PMLR.

[3] Reddi, S., Charles, Z., Zaheer, M., Garrett, Z., Rush, K., Konečný, J., ... & McMahan, H. B. (2020). Adaptive federated optimization. ICLR 2021

4. One issue I have with Table 2 is that the Mix score depends on the mixing proportions. I think it would be more illustrative to divide this table into InD and OOD tables and highlight with bold the best models in the respective tables.

Thank you for the suggestion! We have split the table into several ones and added them to the supplementary part of the paper.

5. Ablations are required to complement the fix proposed in 3.3. Otherwise, I remain unconvinced that the described issue “could be a potential issue in federated learning” or even exists.

We apologize for being imprecise by saying that this led to issues in federated learning. In fact, the issue is general, and it holds not only for the federated scenario but also for the centralized one. Thank you for pointing it out! The demonstration of the issue is already in Figure 2. However, to make it even more apparent, we will add the same plot for the loss function with the proposed fix. Even further, we slightly change the locations of the point clouds. The reason for that is the untrained flow by default initializes higher density in the origin. So we move the noisy blob to the point of (0, 1) and add new results.

Dear Reviewer,

Thank you for your valuable feedback!

Following your suggestions, we have conducted additional experiments, focusing on threshold selection ablation (detailed in Section B.4) and evaluating the effectiveness of the Stopgrad technique (described in Section B.5). We kindly invite you to review these additions, along with our general response to all reviewers. In relation to your query about recent publications, see for example these:

1. Yu, S., Hong, J., Wang, H., Wang, Z., & Zhou, J. (2022, September). Turning the curse of heterogeneity in federated learning into a blessing for out-of-distribution detection. In The Eleventh International Conference on Learning Representations

2. Kumar, S., Lakshminarayanan, A., Chang, K., Guretno, F., Mien, I. H., Kalpathy-Cramer, J., ... & Singh, P. (2022, September). Towards more efficient data valuation in healthcare federated learning using ensembling. In International Workshop on Distributed, Collaborative, and Federated Learning (pp. 119-129). Cham: Springer Nature Switzerland.

3. Li, B., Shi, Y., Kong, Q., Du, Q., & Lu, R. (2023). Incentive-Based Federated Learning for Digital Twin Driven Industrial Mobile Crowdsensing. IEEE Internet of Things Journal.

4. Zhai, R., Chen, X., Pei, L., & Ma, Z. (2023). A Federated Learning Framework against Data Poisoning Attacks on the Basis of the Genetic Algorithm. Electronics, 12(3), 560.

Thank you for the update.

Unfortunately, I am not sure that my concerns are addressed. I admit that I might not fully understand the results in B4, but it seems to me that we might as well use the global model if we care about OOD data, given that 1) using the local model in any capacity hurts OOD performance (i.e., there is always a trade-off between InD and OOD), and 2) selecting the threshold incorrectly may ruin the performance OOD while tuning it requires access to OOD data.

Thank you for providing B5. I suggest using a log scale for the y-axis.

Reading discussions with the other reviewers, it seems like there is still ongoing work on experiments with harder datasets. This is great and I wish the authors success in this endeavor.

We thank the reviewer for their thorough review and feedback. We address their questions below.

1. Comparison with Bayesian Baselines FedPA[1], FedEP[2], and FedPop[3]:

In the FedPA paper, the authors adopt a Bayesian approach to derive a global posterior from locally distributed data. This focus leads to an optimization process that results in the global model, lacking personalization. The adaptation of our switching procedure to their model is not straightforward. Similarly, the FedEP paper also focuses on the global posterior without offering personalized models, making it unclear how to compare our local/global models with their method.

As for the FedPop paper, while it adopts a Bayesian view and has personalization, they don’t have a notion of a global model. They indeed have a shared backbone (feature extractor) and many local personalized models. However, no proper global model. Moreover, they do not provide a method for separately evaluating epistemic and aleatoric uncertainties.

We wish to highlight that our paper's main contribution is the introduction of the switching framework. This framework is general, allowing for various implementations, whether they are based on Variational Inference (VI) or Markov Chain Monte Carlo (MCMC) methods. We chose to illustrate this framework with a Dirichlet-based model because of its relative ease of training, straightforward implementation, and intuitive way of measuring epistemic uncertainty.

2. Comparison is unfair for multiple reasons:

• Other models do not use a switching model

We appreciate the observation regarding the lack of switching strategies in other methods. Indeed, our approach is unique in this aspect. Implementing a similar switching strategy in other models is not a straightforward task.

This complexity arises from the complications involved in identifying a suitable metric for switching that aligns with the specific characteristics and objectives of each model. Our method was specially designed with a switching strategy that is integrated into its core framework, making the replication of this feature in other models non-trivial. This feature not only demonstrates the innovation of our approach but also underscores the challenges in adapting such strategies to different models. Moreover, not all the models used in the list of baselines we compared with are personalized. Hence, it is questionable what are those models between which we should switch.

• The proposed method is the only method geared toward UQ. Bayesian FL methods and/or the methods mentioned above should be included as those methods are (in)directly uncertainty quantification methods too.

Uncertainty quantification is an essential part of our approach that gives the full UQ capabilities for the resulting system of models. However, our main innovation is not in UQ itself but rather in its application to switching between local and global models. That is why our experiments are focused on classical performance metrics such as accuracy, and not on the evaluation of predictive uncertainty. Even if we have an estimate of uncertainty for some model, it is not clear how we should proceed with them for switching as some of the mentioned models do not have a global model at all. For example, in FedPop there is no global model, but only local ones (which are compositions of the global parameters and local ones sampled from some population prior).

• The paper claims that Dirichlet is an efficient method compared to MCMC and VI, but what about in terms of performance in federated settings?

Indeed, our approach utilizing the Dirichlet method can be seen as a lightweight form of VI, offering greater computational efficiency compared to standard VI techniques and MCMC. The latter, particularly, requires significant time both to converge to the target posterior and to obtain uncorrelated samples. Our paper's primary contribution is the introduction of a novel switching framework. We specifically chose the Dirichlet-based model within this framework for its computational lightness and its advantages in practical scenarios. While it's true that employing MCMC with many computational resources might yield superior results, our aim is not to demonstrate state-of-the-art performance. Rather, we focus on presenting a new concept in model selection that was previously unexplored, along with validating its effectiveness through practical application.