Paper: Accelerating Non-IID Federated Learning via Heterogeneity-Guided Client Sampling (ICLR-2024)

Reviewer comment (Presentation):
The organization of Section 3 could be improved. It currently includes background ideas, existing approaches, and the proposed method, making it difficult to identify the novel aspects of this work. A possible solution could be to split it into two sections: "Preliminaries" and "Proposed Method".

Author response:
**3.1A**: Thanks for the suggestions! We are committed to following up and improving the paper’s presentation and organization. We have already taken steps to clarify the novel aspects of the work – please see our response to reviewer V18X30 (parts of which we incorporated in the revised paper).

Reviewer comment (Experiments):
The proposed method has only been proven effective on image classification tasks (FMNIST, CIFAR-10, and CIFAR-100). It remains unclear whether the method can be extended to non-classification tasks (e.g., regression) or non-image tasks (e.g., NLP).

Author response:
**3.2A**: As stated in Section 3.2.1, our method is applicable to any multi-class classification tasks. Moreover, the method is model-agnostic since we make no assumptions on the model architecture. While our evaluation has focused on the widely used image datasets (FMNIST and CIFAR-10/100), the method could in principle be applied to e.g. transformer-based models used for text classification.

Following up on the reviewer’s comment that the experimental results on NLP tasks would further strengthen the significance of our work, we conducted additional experiments involving text classification on the THUC news dataset (https://github.com/thunlp/THUCTC) in Chinese language with 10 labels. There are 18000 samples in the training set, and 10000 samples in each the validation and test set.  For consistency with the main paper, we allocated data to 50 clients by emulating heterogeneous data distributions scenarios with parameter of $\alpha$ set to: (1) {0.001, 0.01, 0.1, 0.2,1}; (2) {0.001, 0.002, 0.01, 0.1,0.5}; and (3) {0.001, 0.002, 0.005, 0.01, 0.1}. We trained **TextRNNs** [1] with BiLSTM architecture (very different from the CNNs used in our prior experiments) as the classifiers with Adam optimizers. Further details of the experimental settings could be provided in the appendix of the revised submission. The results, reported in Table 3 below, show that our method outperforms baselines in all the settings, demonstrating efficacy of our proposed algorithm in an NLP task. For completeness, we also provide Table 4 showing the number of global rounds needed for these experiments to achieve 80% test accuracy.

Table 3: Test accuracy of TextRNN trained for a text classification task on THUC news dataset after 100 rounds of training.

|  | setting (1) | setting (2) | setting (3) |
|---|:---:|:---:|:---:|
| Random | 78.9 | 74.9 | 72.7 |
| pow-d | 80.0 | 75.4 | 66.5 |
| CS | 80.6 | 82.8 | 79.4 |
| DivFL | 73.0 | 68.9 | 72.1 |
| FedCor ($\beta = 0.9$)  | 81.2 | 81.3 | 76.4 |
| HiCS-FL  | **83.2** | **83.9** | **79.7** |

Table 4: The number of global rounds needed to achieve **80%** test accuracy. The settings are the same as in the experiments reported in Table 3. 

|  | setting (1) | setting (2) | setting (3) |
|---|:---:|:---:|:---:|
| Random | 56 | 83 | 179 |
| pow-d | 67 | 109 | 159 |
| CS | 48 | 74 | 107 |
| DivFL | 189 | 289 | 355 |
| FedCor  | 64 | 100 | 131 |
| HiCS-FL  | **25** | **27** | **91** |


Reference:

[1] Liu P, Qiu X, Huang X. Recurrent neural network for text classification with multi-task learning[J]. arXiv preprint arXiv:1605.05101, 2016.

Reviewer comment (Experiments):
The significance of the proposed method over existing client sampling-based approaches such as pow-d, CS, and DivFL is clear. However, the necessity for client sampling is not. There are numerous methods addressing non-iidness in FL, not just FedProx, but also Scaffold [Karimireddy et al. 2020], Scaffnew [Mishchenko et al., 2022], FedNova [Wang et al., 2020], FedOpt [Reddi et al., 2021]. The contribution of the proposed method would be clearer if it were compared against these non-sampling-based methods quantitatively in terms of training speeds and communication/computation overheads.

Author response:
**3.3A**: Scaffold, FedNova and FedOpt address the data heterogeneity problem in FL using regularization or adaptive optimizers, which is indeed a different strategy than client selection utilized by the methods [1]-[4]. We view the two sets of techniques as being complementary rather than competing with each other. It would certainly be very interesting to further explore integration of these two sets of approaches but a comprehensive study on that topic would be rather involved and goes beyond the scope of the current paper.

**3.4A**: As stated in the paper (please see the second paragraph in Section 4.2 and Table 1), the communication cost of our method is identical to the communication cost of FedAvg – our method does not require any additional information and thus incurs no communication overhead. Regarding the computation overhead: The only additional computation is to estimate data heterogeneity and perform clustering based on the bias updates. These operations are negligible compared to the computational efforts required to update the model. It is helpful that the reviewer raised this question, we will make sure to highlight these points in the revised paper; currently, most of the discussion on this topic is in Appendix A.13.


References:

[1] Chen W, Horvath S, Richtarik P. Optimal client sampling for federated learning[J]. arXiv preprint arXiv:2010.13723, 2020.

[2] Cho Y J, Wang J, Joshi G. Client selection in federated learning: Convergence analysis and power-of-choice selection strategies[J]. arXiv preprint arXiv:2010.01243, 2020.

[3] Fraboni Y, Vidal R, Kameni L, et al. Clustered sampling: Low-variance and improved representativity for clients selection in federated learning[C]//International Conference on Machine Learning. PMLR, 2021

[4] Balakrishnan R, Li T, Zhou T, et al. Diverse client selection for federated learning via submodular maximization[C]//International Conference on Learning Representations. 2022.

Reviewer comment (Experiments):
The comparison methods employed in the study are both limited in scope and outdated. It is strongly recommended that the authors thoroughly explore recent and state-of-the-art methodologies to offer a more comprehensive and up-to-date evaluation of their proposed approach.

Author response:
**2.3Q**: *“The comparison methods employed in the study are both limited in scope and outdated. It is strongly recommended that the authors thoroughly explore recent and state-of-the-art methodologies to offer a more comprehensive and up-to-date evaluation of their proposed approach.”*

**2.3A**: The baselines in our experiments include pow-d (2020 AISTAT) [5], clustered sampling (2021 ICML) [6], and DivFL (2022 ICLR) [7], state-of-the-art methods for client selection in FL systems. Moreover, these three methods are explainable and come with theoretical analysis/guarantees, as the method in our paper does. Following the suggestion of reviewer yWNZ, we added one more baseline to the experiments, FedCor (2022 CVPR) – please see our response to reviewer yWNZ. There indeed exist a number of other client selection methods, e.g., those that rely on reinforcement learning techniques. However, those methods tend to be empirical in nature, and challenging to analyze/interpret. That being said, we would certainly welcome suggestions from the reviewer regarding additional baselines, especially if said baselines are analytically well-grounded.

Reference:

[1] Li T, Sahu A K, Zaheer M, et al. Federated optimization in heterogeneous networks[J]. Proceedings of Machine learning and systems, 2020

[2] Acar D A E, Zhao Y, Navarro R M, et al. Federated learning based on dynamic regularization[J]. arXiv preprint arXiv:2111.04263, 2021.

[3] Li Q, He B, Song D. Model-contrastive federated learning[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021

[4] Chen W, Horvath S, Richtarik P. Optimal client sampling for federated learning[J]. arXiv preprint arXiv:2010.13723, 2020.

[5] Cho Y J, Wang J, Joshi G. Client selection in federated learning: Convergence analysis and power-of-choice selection strategies[J]. arXiv preprint arXiv:2010.01243, 2020.

[6] Fraboni Y, Vidal R, Kameni L, et al. Clustered sampling: Low-variance and improved representativity for clients selection in federated learning[C]//International Conference on Machine Learning. PMLR, 2021

[7] Balakrishnan R, Li T, Zhou T, et al. Diverse client selection for federated learning via submodular maximization[C]//International Conference on Learning Representations. 2022.

Reviewer comment (Experiments):
This paper didn't provide enough experiments to support this method. Please add experiments for more complicated networks, current test accuracy is too low to verify the efficiency of the algorithm.

Author response:
**1.2Q**: *“This paper didn't provide enough experiments to support this method. Please add experiments for more complicated networks, current test accuracy is too low to verify the efficiency of the algorithm. Please provide experiments to benchmark with STOA methods like FedCor with optimizers like Adam. ”*

**1.2A**: We thank the reviewer for the suggestions regarding incorporating among baselines other model architectures and the SOTA methods such as FedCor [2]. Following up on this request, we conducted additional experiments and are reporting the results below. However, please note that the performance of FedCor, as reported in [2], is comparable to the performance of HiCS-FL, as reported in our manuscript, even though our manuscript explores heterogeneous data settings which are considerably more challenging. Indeed, data heterogeneity, rather than the capacity of the networks, is the main reason for the relatively low accuracy (i.e., the low accuracy in data heterogeneous settings is caused by the client-drift problem, as previously argued in e.g. [1]). As illustrated in Appendix A.10, clients’ local datasets may be small; in extreme cases, there could be only one or two classes present in those datasets. In such settings, local training of deep models (e.g., ResNet) will often lead to overfitting, ultimately causing relatively low test accuracy.

The new experiments: For a fair comparison of our method to FedCor we ran experiments with ResNet20 and CNNs on CIFAR-10/100, where the models are trained using the Adam optimizer; the FedCor’s code is obtained from the official repository (https://github.com/Yoruko-Tang/FedCor/tree/fedcor_new). The results are reported in Table 1 and Table 2 below. As can be seen there, using a deeper network does not improve the test accuracy in the settings with extremely heterogeneous client data (i.e., overfitting still happens). Nevertheless, since our method (HiCS-FL) selects clients with more balanced data, it succeeds to a degree to mitigate overfitting and generally outperforms the competing methods. FedCor achieves competitive performance only in setting (3), which involves CNNs. 

Table 1: Test accuracy of ResNet20 trained for a classification task on CIFAR10. The FL system has 50 clients, the number of global rounds is set to 500, two local epochs. As in the main papr, we consider 3 heterogeneity settings parameterized via the values of the Dirichlet distribution parameter alpha: (1) {0.001, 0.01, 0.1, 0.5, 1}; (2) {0.001, 0.002, 0.005, 0.01, 0.5]}; (3) {0.001,0.002,0.005,0.01,0.1}. Note that beta in FedCor is the annealing coefficient controlling the sampling strategy [2]; we set it to the recommended value.

| | setting (1) | setting (2) | setting (3) |
|---|:---:|:---:|:---:|
| Random | 48.7 | 33.3 | 22.5 |
| pow-d | 41.9 | 40.9 | 30.3 |
| CS | 60.2 | 31.8 | 28.5 |
| DivFL | 39.6 | 26.3 | 22.3 |
| FedCor ($\beta = 0.9$) | 53.7 | 43.1 | 25.5 |
| HiCS-FL  | **61.7** | **50.2** | **35.3** |


Table 2: Test accuracy of CNNs trained for a classification task on CIFAR10. The other settings are the same as in Table 1.

| | setting (1) | setting (2) | setting (3) |
|---|:---:|:---:|:---:|
| Random | 55.6 | 36.1 | 35.1 |
| pow-d | 52.4 | 34.8 | 33.6 |
| CS | 61.4 | 46.3 | 43.9 |
| DivFL | 43.6 | 35.6 | 27.1 |
| FedCor ($\beta = 0.9$)  | **64.3** | 52.6 | 42.1 |
| HiCS-FL  | 63.9 | **61.7** | **50.5** |

Reviewer comment (Experiments):
This paper didn't evaluate cases with client inaccessibility issues, though which were mentioned in the abstract. Please provide experiments when a given fraction of clients are unavailable and show the performance of your methods with other STOA methods.

Author response:
**1.3Q**: *“This paper didn't evaluate cases with client inaccessibility issues, though which were mentioned in the abstract. Please provide experiments when a given fraction of clients are unavailable and show the performance of your methods with other STOA methods”*

**1.3A**: The reviewer is likely referring to the following sentence in Abstract: “Particularly challenging are the settings where due to resource constraints only a small fraction of clients can participate in any given round of FL.” What we mean is that in order to satisfy communication constraints, the server selects only a subset of clients out of the entire cohort of clients – more precisely, the assumption is that the number of available clients exceeds the limit imposed by bandwidth constraints, and that the server selects a subset of the available clients. This point is further clarified in Introduction (the 1st sentence of the 3rd paragraph). Similar settings are considered by the authors of the baselines [2-5] to which we compare our proposed method. While the scenario where, unbeknownst to the server, a selected client is unavailable has not been within the scope of our work thus far, the reviewer makes a valid suggestion – it would indeed be interesting to consider it as part of future work.

Reviewer comment (Writing):
Writing issues: some part of this paper is a little confusing and there are some typos in certain figures. For example, in Figure 6, the y-axis label is too long, explain it in the caption of the figure.

Author response:
**1.4Q**: *“some parts of this paper are a little confusing and there are some typos in certain figures. For example, in Figure 6, the y-axis label is too long, explain it in the caption of the figure”*

**1.4A**: Good catch, thanks for pointing this out! We are committed to improving the readability and fixing such errors in the final version of the manuscript, should it be accepted.

Reviewer comment (Experiments):
My greatest concern is that this method may not be extended to cases other than SGD as mentioned in this paper. Even though there are proofs in the appendix, no experiments are provided to demonstrate that. The appendix mentioned your method could be extended to the non-SGD optimizers. Please provide experiments to benchmark with STOA methods like FedCor with optimizers like Adam. The reason I recommended FedCor is that it doesn't need a loss from all clients and shows good performances in terms of convergence speed and final accuracy.

Author response:
**1.1Q**. *“My greatest concern is that this method may not be extended to cases other than SGD as mentioned in this paper.”*

**1.1A**: Actually, the experiments on CIFAR-10/100 data presented in the main paper are conducted using Adam optimizer, as stated in Section 4.1 (please see the first sentence of the second and third paragraph there). Moreover, Appendix A.9 “Optimization Algorithms Beyond SGD” presents an analysis of the proposed scheme for (a) SGD with momentum and (b) Adam optimizer. This analysis is omitted from the main paper only due to space limitations. Our apologies if our statements in the main paper were insufficiently clear, we will certainly clarify in the revised manuscript that the method applies to optimizers beyond SGD.