Uncertainty-Based Extensible Codebook for Discrete Federated Learning in Heterogeneous Data Silos

Authors: Tianyi Zhang, Yu Cao, Dianbo Liu
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Federated learning, uncertainty, discretization

Abstract:
Federated learning (FL), aimed at leveraging vast distributed datasets, confronts a crucial challenge: the heterogeneity of data across different silos. While previous studies have explored discrete representations to enhance model generalization across minor distributional shifts, these approaches often struggle to adapt to new data silos with significantly divergent distributions. In response, we have identified that models derived from FL exhibit markedly increased uncertainty when applied to data silos with unfamiliar distributions. Consequently, we propose an innovative yet straightforward iterative framework, termed Uncertainty-Based Extensible-Codebook Federated Learning (UEFL). This framework dynamically maps latent features to trainable discrete vectors, assesses the uncertainty, and specifically extends the discretization dictionary or codebook for silos exhibiting high uncertainty. Our approach aims to simultaneously enhance accuracy and reduce uncertainty by explicitly addressing the diversity of data distributions, all while maintaining minimal computational overhead in environments characterized by heterogeneous data silos. Through experiments conducted on six datasets, our method has demonstrated its superiority, achieving significant improvements in accuracy (by 3%--22.1%) and uncertainty reduction (by 38.83%--96.24%), thereby outperforming contemporary state-of-the-art methods.

Decision: reject
Reviewer scores: 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1730535517218

thread_type:
official_review

<thread_object>
creator: reviewer_7ffp

<content>
summary:
This paper proposes UEFL, which utilizes an extensible codebook to address feature shifts in FL. Numerical results demonstrate the superior performance of the proposed method in both multi-domain learning and out-of-distribution generation tasks.

strengths:
1. Utilizing the codebook technique to address the feature shift in FL is innovative.

2. The authors conducted comprehensive experiments on multiple datasets, demonstrating the performance improvement of the proposed method in both traditional FL tasks and out-of-distribution generalization tasks.

weaknesses:
1. The algorithms presented in Section 3 do not align well with Algorithm 1. For instance, lines 226-228 are not thoroughly explained. It would be beneficial if the authors could provide additional details on the main pages.

2. The number of training epochs for the baseline models is set to be the same as that for UEFL. However, I am uncertain about the fairness of this approach, as (1) the optimal number of training epochs for UEFL may not be optimal for the baseline models; (2) the communication and computation costs associated with UEFL and the baseline models differ, which makes the comparison somewhat unfair.

questions:
It would be advantageous for the authors to include convergence curves relating to communication, computation, and the number of rounds. This would provide a clearer illustration of the effectiveness of UEFL.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730596957962

thread_type:
official_review

<thread_object>
creator: reviewer_3k1q

<content>
summary:
This paper presents a method for federated learning, where the goal is to address data heterogeneity across clients and enhance generalization ability of federated learning model across different data distributions. Specifically, the method leverages a codebook of representations to improve accuracy and reduce prediction uncertainty in federated settings. The authors validate the proposed framework in multi-domain learning, and domain generalization in leave-one-domain-out experiments.

The proposed framework is as follows: for each client model, first encode the input data to latent features, extract code words via k-means clustering of these latent features, and then discretize the latent features using the nearest code word; predictions are subsequently made based on the coded (discretized) latent vectors. The framework also introduces an uncertainty-based adjustment mechanism, where clients with uncertainty above a certain threshold use a larger codebook with additional code words. The server model is updated as the average over all clients.

strengths:
The authors address an important challenge in federated learning to tackle data heterogeneity and enhance model generalization across different distributions.

The presentation of the paper is easy to follow.

weaknesses:
I have the following questions regarding the work:

1. My first question is about the rationale behind the framework. Could the authors elaborate on how and why discretizing latent representations through a codebook helps achieve the goals of improved generalization and uncertainty reduction?

2. Regarding the server model, could the authors clarify how the codebook of the server model is updated?

3. Is there any convergence guarantee of Algorithm 1? Specifically, does the framework ensure that all clients will reach an uncertainty level below the threshold? Could the authors clarify the mechanism described in lines 322–323?

4. Could the authors clarify why reducing uncertainty in prediction is a valid objective? If the model achieves high certainty on incorrect predictions, could this exacerbate challenges in improving model performance? Additionally, if reducing uncertainty is a primary goal, beyond the proposed method of expanding client codebook, have the authors considered training client models to the point of neural collapse [1]?

5. Distribution shift is also addressed in other studies, such as [2,3]. It would strengthen the evaluation if the authors could expand their comparisons to include more related methods.



[1] Papyan, Vardan, X. Y. Han, and David L. Donoho. "Prevalence of neural collapse during the terminal phase of deep learning training." Proceedings of the National Academy of Sciences 117.40 (2020): 24652-24663.

[2] Nguyen, A. Tuan, Philip Torr, and Ser Nam Lim. "Fedsr: A simple and effective domain generalization method for federated learning." Advances in Neural Information Processing Systems 35 (2022): 38831-38843.

[3] Zhang, Hao, et al. "FedCR: Personalized federated learning based on across-client common representation with conditional mutual information regularization." International Conference on Machine Learning. PMLR, 2023.

questions:
It would be helpful if the authors could address my above questions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730711144128

thread_type:
official_review

<thread_object>
creator: reviewer_4szv

<content>
summary:
The paper addresses the challenge of data heterogeneity in Federated Learning (FL) by proposing a novel method called Uncertainty-Based Extensible-Codebook Federated Learning (UEFL). The method introduces a dynamic codebook that maps latent features to discrete vectors, allowing for the extension of the codebook based on the uncertainty of data distributions across different data silos (clients). By iteratively evaluating the uncertainty of the global model's predictions, UEFL adds new codewords to the codebook to improve performance on previously unseen or highly divergent data distributions. The proposed method is tested on six datasets, showing significant improvements in both accuracy (3% to 22.1%) and uncertainty reduction (38.83% to 96.24%) over state-of-the-art approaches.

strengths:
- The idea of using a dynamic, uncertainty-driven codebook to handle data heterogeneity in federated learning is innovative. By extending the codebook based on uncertainty, the method adapts to diverse and unseen data distributions more efficiently than traditional methods.

- The approach of starting with a small, shared codebook and gradually expanding it based on data-specific uncertainties is efficient. This design helps in minimizing computational overhead while ensuring that the model can adapt to new distributions without retraining the entire model.

weaknesses:
- In Table 1, the reported accuracy of FedAvg on CIFAR-100 is only 11%, which is surprisingly low. Typically, even with non-IID data, FedAvg achieves higher accuracy. This raises concerns about potential baseline suppression to exaggerate the performance of the proposed method. The authors need to clarify the experimental setup for FedAvg and ensure that the baseline methods are fairly represented. Without this clarification, the validity of the results is questionable.

- The experiments seem to rely on pretrained models. It would be essential to clarify whether the proposed UEFL method still performs well when trained from scratch. Relying on pretrained models can bias results, especially if the pretrained weights are already well-optimized for certain datasets. The authors should include experiments that train the models from scratch to validate the robustness and generalizability of UEFL. Additionally, the 11% accuracy for FedAvg on CIFAR-100 with a pretrained model is puzzling and further raises questions about the validity of the reported results.

- The paper primarily uses Monte Carlo Dropout for uncertainty estimation, which is a standard but limited approach. More advanced uncertainty quantification techniques (e.g., Deep Ensembles or Bayesian Neural Networks) could be explored or at least discussed as alternatives. The performance of UEFL might vary with different uncertainty estimation methods.

- The paper lacks a comprehensive analysis of the computational complexity of the proposed method. Although the authors claim that the codebook remains compact, it would be beneficial to quantify the computational and communication overhead in both the training and inference phases. A comparison of the complexity of UEFL with other FL methods in terms of memory usage and time complexity is missing.

- The paper mainly evaluates the method on six datasets, but does not provide a clear discussion of its scalability to larger numbers of clients. Given that FL is often applied to large-scale systems, it would be helpful to discuss how UEFL scales with increasing numbers of clients or more heterogeneous data distributions.

- Although the paper discusses the adaptive nature of the codebook, it would be useful to include an ablation study that explores the impact of different initial codebook sizes and how the codebook grows over iterations.

questions:
- The reported 11% accuracy for FedAvg on CIFAR-100 is unreasonably low, especially considering that the experiments used pretrained models. Could the authors clarify whether there were any specific preprocessing steps or hyperparameter choices that led to such low performance? Additionally, does UEFL still outperform FedAvg when both models are trained from scratch, without the benefit of pretraining?

- How does the growth of the codebook over multiple iterations affect the overall communication overhead in FL? Does the increased size of the codebook lead to significant additional communication costs during model aggregation, and how is this handled?

- How is the threshold for uncertainty (which triggers the expansion of the codebook) determined? Is this threshold set manually, or is it dynamically adjusted during training? Additionally, how sensitive is the method to this threshold?

- How does UEFL perform when the data distributions across silos are extremely divergent (for example, different modalities of data)? Does the codebook continue to expand indefinitely, or is there a mechanism to prevent overgrowth?
</content>
</thread_object>

</thread>


Metareview:
metareview:
Post-rebuttal, the majority of the reviewers still judged the contribution of the paper as fair. I appreciate the efforts the authors have put in their rebuttal, and the updates in the draft, very substantial and helpful in the appendices, but in the end the paper unfortunately does not pass the bar for acceptance. I can only encourage them further in updating the paper for resubmission, perhaps digging further in the theory that they so far cite from another paper.