FENDA-FL: Personalized Federated Learning on Heterogeneous Clinical Datasets

Thank you for taking the time to review our work and providing useful feedback on how our paper might be improved. We have significantly revised the paper to incorporate new experiments and discussion. Below we discuss each of the comments and outline how the revision aims to address them.

My major concern is a significant oversight in this work, which is the absence of discussions of an important line of related works - federated domain adaptation/generalization. This topic is closely related to the central topic of this paper. There are many advanced techniques for federated domain adaptation/generalization in recent years, for examples [1-2] and numerous other contributions in this domain. Yet, the authors seem to have omitted any discussion contrasting their proposed model with these works, nor have they incorporated them as baseline models for comparison. The decision to utilize the FENDA model, which appears potentially outdated in the domain adaptation field, raises concerns. Without comparing the proposed model against current SOTA methods in the field of federated domain adaptation, it is hard to assert that the domain adaptation strategy showcased here represents the state-of-the-art in FL.

[1] Yao, C. H., Gong, B., Qi, H., Cui, Y., Zhu, Y., & Yang, M. H. (2022). Federated multi-target domain adaptation. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (pp. 1424-1433). [2] Zhang, R., Xu, Q., Yao, J., Zhang, Y., Tian, Q., & Wang, Y. (2023). Federated domain generalization with generalization adjustment. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 3954-3963).

In the revision, we have included a thorough discussion of the intersection and differences between federated domain adaptation and personalized FL, which we consider here. The majority of this discussion occurs in Section 2 and Appendix J. As noted, personalized FL and federated domain adaptation are closely related fields. However, the settings that are considered differ in some significant ways. For example, [1] considers a server-hosted, labeled source dataset and client-based, unlabelled target datasets, which is quite distinct from the setting considered in our work. As such, [1] does not include benchmarks against standard FL approaches. Similarly, other works in federated domain adaptation, such as [3, 4, 5], also consider distinct settings with a mixture labeled and unlabeled datasets, distributed in various ways.

In this paper, we do not explicitly consider domain adaptation as the target task. Rather, we apply a technique from domain adaptation as a means of improving the robustness of personalized FL in the context of non-IID or heterogeneous datasets. However, the Generalization Adjustment approach proposed in [2] is very interesting. As an augmentation to existing FL approaches, it appears quite promising. The primary evaluation criterion in [2] is “leave-one-domain-out,” where only a subset of clients are used in training and generalization is measured exclusively on the held-out client domain, which is somewhat distinct from our setting. However, using it to improve both global and personal FL approaches is intriguing. We have included it as an important avenue for future work.

[1] Yao, C. H., Gong, B., Qi, H., Cui, Y., Zhu, Y., & Yang, M. H. (2022). Federated multi-target domain adaptation. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (pp. 1424-1433).

[2] Zhang, R., Xu, Q., Yao, J., Zhang, Y., Tian, Q., & Wang, Y. (2023). Federated domain generalization with generalization adjustment. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 3954-3963).

[3] X. Peng, Z. Huang, Y. Zhu, and K. Saenko. Federated adversarial domain adaptation. arXiv preprint arXiv:1911.02054, 2019.

[4] L. Song, C. Ma, G. Zhang, and Y. Zhang. Privacy-preserving unsupervised domain adaptation in federated setting. IEEE Access, 8:143233–143240, 2020. Doi: 10.1109/ACCESS.2020.3014264.

[5] Y. Shen, J. Du, H. Zhao, Z. Ji, C. Ma, and M. Gao. FedMM: A communication efficient solver for federated adversarial domain adaptation. In Proceedings of the 2023 International Conference on Autonomous Agents and Multiagent Systems, AAMAS ’23, pp. 1808–1816, Richland, SC, 2023. International Foundation for Autonomous Agents and Multiagent Systems. ISBN 9781450394321

Thanks to the authors for their response. However, I'll stand for my original score due to the limited related work discussion and baseline comparison.

Thank you for your attentive review of our work. We appreciate your time and the opportunity to improve our results along the lines of your suggestions. We have endeavored to revise the paper to thoroughly address the weaknesses identified in the review. Below we discuss each of the comments and outline how the revision aims to address them.

The main text of the submission exceeds 9 pages.

In the submission guidelines, authors were encouraged to include a “Reproducibility Statement” at the end of the main text but before the “References.” If we have read the instructions properly, the statement is not counted towards the page limit. Since it is the “Reproducibility” section that exceeds 9 pages, we believe that the submission is still compliant. However, if we are incorrect in such an interpretation, we would appreciate the correction and will ensure that the revision is, in fact, compliant.

This paper addresses the challenge of personalized federated learning using heterogeneous local datasets; however, the mechanism is not thoroughly explained. The proposed method draws inspiration from domain adaptation, where distribution shifts naturally occur. In this paper, local clients have datasets sampled from various distributions. The question arises: What types of distribution shifts (such as feature shifts or label shifts) can the proposed method effectively handle?

Please describe the mechanism of the proposed method in detail and elucidate its effectiveness in addressing different types of distribution shifts, particularly in the context of client dataset heterogeneity.

We agree that such a discussion is important and was not adequately present in our original draft. A formal description of the FENDA-FL approach is now included in Appendix J. Therein, we also discuss the admissibility of the three most common types of distribution changes: label shift, covariate shift, and concept drift. Because FENDA-FL includes a locally trained feature extractor and classification head, each of these shifts can theoretically be handled by the approach. For label and covariate shift, these modules may adapt feature representations and predictions to lean more or less heavily on global representations or modify prediction distributions to match local statistics. For concept drift, the relationship between labels and global features may be de-emphasized in favor of local representations. In extreme cases, where global information is not relevant, the local prediction head may learn to ignore global features altogether. We hope that this section more thoroughly, and clearly, describes the FENDA-FL design and its flexibility in non-IID data settings.

Thank you for your review of our work and offering your thoughts on where the paper could be improved. We have attempted to address each of the concerns raised and believe that the paper is markedly improved as a result. Below we discuss each of the comments and outline how the revision aims to address them.

It would be nicer to also see a table of the results, rather than a visual representation due to the number of settings. Minor comment: The color choices for the bar graphs seem to represent goodness, with red being worse and blue being better; however, the different colors are a little jarring and could cause confusion.

Thank you for your suggestion. Other reviewers also noted that the bar graphs did not present our results as clearly as we had hoped. Taking your suggestion, we have modified the presentation of the results to be in a tabular format. We hope that the representation of the results is improved. The bar graphs are still presented in Appendix K. However, if they should be removed, we are happy to do so.

The performance benefit compared to APFL shown in Fig. 2 is not very clear, also in Fig. 3, a similar observation can be made for FedAVG and APFL

With the change of results presentation, we hope that the readability and ease of analysis has been measurably improved. The experimental results are now presented in Tables 2-4. FENDA-FL outperforms APFL on the Fed-Heart-Disease task by a significant margin. While FENDA-FL is the second best method for Fed-IXI, the new architecture studies in Appendix F show that the gap between the two methods can be closed through thoughtful network design. For Fed-ISIC2019, each method under-performs global approaches. For both GEMINI Mortality and Delirium, FENDA-FL outperforms APFL in all metrics.

The idea of using domain adaptation seems like a natural pairing with FL. Can the authors include a discussion on the inherent differences / relationships between these 2 fields for better context? Google scholar reveals some papers such as https://arxiv.org/abs/1911.02054 and https://arxiv.org/abs/1912.06733 which may be relevant.

We agree that the pairing of techniques in domain adaptation and federated learning is quite natural. In the design of FENDA-FL, we have borrowed an approach from domain adaptation to improve personalized FL and aim to continue to do so in future work. Generally, domain adaptation (DA), including federated domain adaptation, considers the setting of training a model using labeled source datasets and unlabeled, or sparsely labeled, target datasets such that the model performs well in the target domain. As such, the space considers a distinct setting from that considered here but deals with similar data drift issues. For example, in [1], the server holds an unlabeled target dataset while labeled client-hosted datasets serve as the source.

In the revision, we have highlighted the differences and intersections between the disciplines in Section 2, including a discussion of the references provided and those suggested by Reviewer #4. Interestingly, [2] uses a slight variation of APFL and applies DP-SGD for privacy during client training and, as such, does not consider a true DA setting. We have included a reference to this work under settings not considered because the main difference is the use of DP-SGD in client-side training.

[1] X. Peng, Z. Huang, Y. Zhu, and K. Saenko. Federated adversarial domain adaptation. arXiv preprint arXiv:1911.02054, 2019.

[2] D. W. Peterson, P. Kanani, and V. J. Marathe. Private federated learning with domain adaptation. In the 33rd Conference on Neural Information Processing Systems (NeurIPS 2019), Proceedings of NeurIPS, Vancouver, Canada, 2019.

Is it possible to release a version of the code that works for open-source datasets?

This is very much possible and our intention! In releasing the experimental code associated with the paper, which works with the open-source datasets in FLamby, we will also be making the underlying library, which is not specific to the datasets studied, publicly available. The library is compatible with any open-source dataset that can be used with pytorch. Moreover, the library includes a number of examples to facilitate using the implemented methods for other experiments. Our hope is that this library will accelerate research in FL by making experimentation easier and the implementation of new techniques more straightforward.

Thank you for taking the time to carefully read and review our manuscript and for the thoughtful suggestions on how the work could be improved. We have conducted several new experiments along the lines suggested and worked to significantly improve the presentation of our results. Below we discuss each of the comments in order and outline how the revision aims to address them.

While adapting FENDA's concept to the FL scenario required some modifications, the application of FENDA in this methodology appears overly incremental. The approach of combining features to utilize both global and local information, as opposed to APFL's method of aggregating logits, is not particularly novel.

While the architectural differences between APFL and the FENDA-FL approach are not large, there are several significant benefits to the FENDA-FL design that we believe make the proposed modifications important. First and foremost, FENDA-FL outperforms APFL in many of the experiments and is quite competitive in the cases that it does not. Conceptually, the move from the aggregation of logits to the synthesis of feature representations is fairly straightforward. However, the change also implies that training can be done monolithically and that global and local module combinations can incorporate non-linear transformations rather than just weighted averaging.

As noted below, the approach also provides significant flexibility in the design of the local and global feature extraction modules and does not require additional hyper-parameters. In APFL, the local and global modules are exact copies, which limits any inductive bias one might inject into the module designs. As pointed out in another comment, we did not demonstrate the advantage of this flexibility as well as we could have. Appendix F now contains results for studies considering the benefit of adding inductive bias to the FENDA-FL model design.

Finally, the move to the FENDA-FL design is well founded on the success of FENDA. In our revision, we have endeavored to make these improvements more clear, while also recognizing that the change in design is straightforward.

The methodological contribution to the ICLR community unclear. The application of FENDA to federated learning is straightforward, "federated checkpointing" is conceptually a combination early stopping with federated learning.

Thank you for the opportunity to emphasize the contributions of this work. While the adaptation of FENDA is straightforward, we believe it is a useful conceptual bridge between personalized FL methods and domain adaptation. In future work, we hope to bring additional ideas from domain adaptation into the setting of personalized FL. With respect to checkpointing, the proposed characterization is certainly applicable. However, early stopping considers terminating training early, while in the proposed approach federated training continues, with other clients potentially benefiting from further training. Most existing works simply perform a fixed number of server rounds before evaluation.

While the approach is fairly simple, it helps mitigate the risk of overfitting, especially when client data imbalance is an issue. The combination of federated checkpointing and our proposed evaluation criterion aim to bring FL evaluation more in line with standard ML practices. As suggested in another comment, we have conducted experiments without federated checkpointing to demonstrate the advantage of its use. These results are included in Appendix E in the revision.

When considered as a whole, we believe that the experimental framework and results presented here constitute an important representation of effective and fair benchmarking in FL studies. We offer several improvements to evaluation, including checkpointing and various lenses through which FL methods should be evaluated. Finally, the experiments on the FLamby benchmark establish significant improvements for many methods on the selected tasks.

(Discussion Continued in a follow up comment)