Paper: Dirichlet-based Uncertainty Quantification for Personalized Federated Learning with Improved Posterior Networks (ICLR-2024)

Reviewer comment (Novelty):
Novelty is limited. Seems that most of the framework is directly taken from previous works, and this paper just combines and applies these in the federated learning setting as an application. For example, the 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.

Author response:
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.

Reviewer comment (Experiments):
Cannot scale to more complex datasets like CIFAR-10. From Table 2, it seems that there is a significant drop in the in-distribution data. 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.

Author response:
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.  

Dear Reviewer,

Thank you for the feedback.

In response to your comments, we have carried out additional experiments and included these in the revised manuscript. We also invite you to review our common response addressed to all reviewers for a comprehensive overview of the changes made.

Concerning the CIFAR-100 experiments, we have integrated them within our existing setup and observed that our model scales effectively to this dataset. However, to achieve optimal results, it has become apparent that an increase in the capacity of the flow model is necessary. We are working on this enhancement and are confident in presenting the finalized experiment in time for the Camera-Ready submission, should our paper be accepted.

We appreciate your feedback and hope that our revisions meet your expectations.

Reviewer comment (Experiments):
Additional computation. The authors mention that they have to use a density model to decide if they should use a global or local model. This extra step adds more work for the computer. Even if they use 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.

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

Reviewer comment (Novelty):
The methodological novelty is limited to applying existing approaches to a new problem.

Author response:
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.

Reviewer comment (Experiments):
The authors are right to note that the choice of the uncertainty threshold is the most subjective part of the approach. Ablations (where the threshold is varied) are required to inspect how important this choice is. This also leads 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?

Author response:
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.

Reviewer comment (Experiments):
I am not sure about best practices in image classification these days, but the datasets seem simple.

Author response:
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

Reviewer comment (Presentation):
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.

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

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Writing):
Some parts were unclear from the submission and I had to read cited sources to understand them. After eq. 2, properties of p(g(x)) are listed, but it isn’t explained why f(g(x)) is required (as far as I understood, this term is a class prior, but then I do not exactly understand why it is required along with $\alpha_{prior}$). Eq. 5 is hard to parse as cross-entropy, consider using $CE$ like Charpentier et al. (2020).

Author response:
Following [1], we consider the update in the parameters of the prior distribution as softmax prediction, multiplied by the corresponding density. Doing this, we approximate the amount of “pseudo-counts”. This is an important component of this method, as the density part p(g(x)) estimates the whole evidence budget, and the classification part f(g(x)) distributes it among the classes. 

Regarding the equation with the cross-entropy. Note that as a part of the loss function we use so-called Uncertain Cross-Entropy, which is the expectation of the standard Cross-Entropy with respect to the predicted Dirichlet distribution. We agree that the form of writing might seem unusual and we will add a note to make the connection clearer.


[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.

Reviewer comment (Theory):
Can there be scenarios where the scheme in 2.1 should be reversed? I.e., where the global model may have higher epistemic uncertainty.

Author response:
We think that such situations are possible but in a circumstances that are much more special than the ones considered in our paper. For example, it might be that global model is very complex and there is not enough data for the confident prediction at some part of the input space. However, local model for some agents might have significantly simpler structure that might be confidently learned with fewer data.

Reviewer comment (Experiments):
Could the general approach be extended to other Federated tasks such as Regression or RL?

Author response:
For regression the answer is yes. Since the method for which we quantify the uncertainty is based on the NatPN [1], which works for any likelihoods from the exponential family of distributions, including Gaussian, we can easily apply it to regression as well. For RL it looks entirely possible as well.

Reviewer comment (Theory):
From the bounds of $\psi(x)$ before eq. 7, wouldn’t it make more sense to approximate it as $log(x) - 1/x$? Would it change the conclusions about issues with the loss function?

Author response:
Thank you for the question! In fact, in our analysis, we assume that the model has sufficient evidence at the points, which lead to big alphas. Thus, the term 1/x can be omitted.
Regarding the validity of the result -- the experiment shown in Figure 2 computes the formula fairly. And we observe the described effect.

Reviewer comment (Experiments):
Authors claim to use Dirichlet methods for uncertainty quantification but do not compare with existing Bayesian federated methods such as FedPA[1], FedEP[2], and FedPop[3], which are also assumed to perform efficient Bayesian inference for (P)FL.

Author response:
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.

Reviewer comment (Experiments):
Comparison is unfair for multiple reasons: 1) Other models do not use a switching model. Technically speaking, other personalized FL models can benefit from "switching" too, but not exactly in the same way as this paper does. Also, the proposed method is the only method geared toward uncertainty quantification. Bayesian FL methods and/or the methods mentioned above should be included as those methods are (in)directly uncertainty quantification methods too. The paper claims that Dirichlet is an efficient method compared to MCMC and VI, but what about in terms of performance in federated settings?

Author response:
- 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.

Reviewer comment (Novelty):
The claim "In this paper, we proposed a personalized federated learning framework that leverages both a globally trained federated model and personalized local models to make final predictions." is misleading as there are previous papers utilizing both global and local models for prediction [4].

Author response:
Thank you for your comment and for referring to the work of Hanzely & Richtárik. We acknowledge this prior work and have cited it in our manuscript. We would like to clarify that our claim is not about being the first to consider both local and global models. Rather, our contribution lies in proposing a framework that utilizes these models differently, specifically by switching between purely local and global models based on uncertainty quantification.

The focus of Hanzely & Richtárik's work differs significantly from ours. They employ 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 individual 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 approach.

A key distinction between our work and the model by Hanzely & Richtárik is our use of the switching mixture of global and local models for inference, whereas they provide a transition between these models in their objective through a regularization penalty $\lambda$. In their model, there is no switching for inference: a single model is used to evaluate all local samples, with $\lambda$ serving as a hyperparameter. Here, $\lambda = 0$ corresponds to purely local models, and $\lambda = \infty$ to a global model. Lastly, their model is not Bayesian, unlike ours.

We hope this clarifies our framework's unique aspects and contributions in the context of existing literature.

Reviewer comment (Presentation):
The transition from Section 3 to federated settings is poorly developed. Section 3 has too much text in it and it is not clear what are the problems/complications of using Dirichlet methods in federated settings. Authors should condense Section 3 and move some parts to the Appendix as it is hard to understand what is the problem for federated learning.

Author response:
We would like to thank the reviewer for this important comment. We certainly agree that currently we do not have enough focus on the specific issues of FL and have some extra details that can be easily moved to the Appendix. We think that these issues are easy to resolve and we will make sure to improve the readability for the camera-ready version of the paper (if accepted).

Reviewer comment (Writing):
Federated loss was unclear to me. Are you using the loss in Eqn. (8)?

Author response:
Yes, we use the loss function from equation (8) for training our model.
Let’s break this loss into parts. The first one is the UCE (uncertain cross entropy) [1], which is the expectation of CE loss with respect to the predicted Dirichlet distribution, but without gradient update with respect to the parameters of a density model. This part pushes the model to make correct predictions and be confident. The second term penalizes predicted Dirichlet for overconfident predictions and is the same the one in [1]. The last term is the correction we suggested.

[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.