Distance-Aware Non-IID Federated Learning for Generalization and Personalization in Medical Imaging Segmentation

Authors: julia  alekseenko
nicolas  padoy
Venue: MIDL_io_2024_Conference
Published: 2024-01-31 21:43:54

Abstract: Federated learning (FL) in healthcare suffers from non-identically distributed (non-IID) data, impacting model convergence and performance. While existing solutions for the non-IID problem often do not quantify the degree of non-IID nature between clients in the federation, assessing it can improve training experiences and outcomes, particularly in real-world scenarios with unfamiliar datasets. The paper presents a practical non-IID assessment methodology for a medical segmentation problem, highlighting its significance in medical FL. We propose a simple yet effective solution that utilizes distance measurements in the embedding space of medical images and statistical measurements calculated over their metadata. Our method, designed for medical imaging and integrated into federated averaging, improves model generalization by downgrading the contribution from the most distant client, treating it as an outlier. Additionally, it enhances model personalization by introducing distance-based clustering of clients. To the best of our knowledge, this method is the first to use distance-based techniques for providing a practical solution to the non-IID problem within the medical imaging FL domain. Furthermore, we validate our approach on three public FL imaging radiology datasets (FeTS, Prostate, and Fed-KITS2019) to demonstrate its effectiveness across various radiology imaging scenarios.

Main paper: 
 introduction: 
 Federated learning (FL) in healthcare aims to achieve data collaboration while preserving privacy. It enables multiple institutions or healthcare entities (clients) to jointly train or evaluate artificial intelligence (AI) models without sharing raw, sensitive patient data. This collaborative and privacy-preserving approach improves predictive models, personalized treatments, and disease detection, leveraging diverse datasets from various institutions and democratizing the power of distributed clients. FL promotes inclusive model training, incorporating diverse populations for robustness and generalizability. Studies like (Sheller et al., 2020; Dou et al., 2021 ) highlight FL's efficacy in medical applications, showcasing the power of algorithms like federated averaging (FedAvg) (McMahan et al., 2017) .While traditional FL algorithms, like FedAvg, assume uniform data distribution across clients, in contrast, real-world applications often face non-independently and identically distributed (non-IID) data challenges, where data across clients lacks uniformity. Factors like disease manifestation, patient populations or image acquiring protocols contribute to this heterogeneity, impacting model convergence and performance (McMahan et al., 2017) . Recent approaches on improving the generalization of the global model, such as FedProx by Li et al. (Li et al., 2020) , regulate local updates to improve model generalization, solutions like FedBN (Li et al., 2021b) and FedCross (Xu et al., 2022) address non-IID scenarios by optimizing feature spaces or sequentially training the global model across clients. Another essential strategy in FL is personalization, which involves training a specific model for each client while leveraging insights from others. Recent advancements in personalized FL include training one model per participating institution through adaptations of metalearning (Fallah et al., 2020; Acar et al., 2021) , multi-task learning (Marfoq et al., 2021) , utilizing partial model sharing (Pillutla et al., 2022) , local fine-tuning (Li et al., 2021a; Yu et al., 2020) , and clustering solutions (Ghosh et al., 2022; Manthe et al., 2023) .However, the majority of proposed works aim to accept client data distributions as non-IID without measuring the heterogeneity of the federation and integrating this information into the pipeline. Assessing non-IID characteristics can provide crucial insights into training challenges and generalizability, for instance, (Zhao et al., 2018) observed decreased accuracy in federated models with higher Earth Mover's Distance (EMD) in non-IID image datasets. Yet, their study only evaluated classification problems with basic datasets like MNIST and CIFAR-10, limiting its applicability to healthcare. In the medical domain, (Luo et al., 2023) proposed analyzing data distributions related to site, tumor type, tumor size, dataset size, and tumor intensity. They demonstrated a significant negative correlation between the Dice score ratio and data distribution distances, particularly with EMD, in medical image segmentation. However, practical solutions for improving the federation based on these findings have not been proposed.Our main contribution lies in integrating EMD distance-related insights into federated averaging to offer optimal training for both generalization and personalization strategies, surpassing the performance of traditional FL methods. Assuming that this approach may not capture all non-IID characteristics, we explore non-IID measurements within the embedding space of data (i.e., medical images). We utilize a publicly-available pre-trained model to extract rich and meaningful embeddings, and then calculate the Euclidean distance (EUC) based on them. To the best of our knowledge, this is the first work in the medical imaging FL to propose this concept. 
method: 
 Our methodology, presented in Figure 1 , aims to achieve two important yet opposing goals in FL: a) model generalization, focusing on improving the accuracy and generalization of the model, and b) model personalization, tailoring the model for the highest accuracy at the client level. Our proposed methodology can be subdivided into two (2) steps. The first step involves measuring the degree of data heterogeneity (non-IID) among clients. To achieve this, our methodology runs two approaches in parallel. The first approach relies on statistical distances (i.e., EMD) based on metadata from medical images (Subsection 2.1), while the second one calculates EUC in the embedding space of these images (Subsection 2.2). For this, a publicly pre-trained model is deployed from the server to each client for extracting embeddings. Subsequently, we extend the federated averaging algorithm by incorporating the down-weighting of the most distant client to enhance generalization. We also augment the personalization strategy by clustering closely-related clients (Subsection 2.3). While we hypothesize that the federation is trusted, we recognize the potential for privacy enhancement. However, the investigation into the compatibility of our methodology with privacy-preserving techniques (Jin et al., 2023) , (Wei et al., 2020) is beyond the scope of this study.To assess the non-IID nature of the federation, characterized by disparities between each client data, we calculate the EMD distance on metadata available in training medical images from each client. This metadata includes maximum intensity values and label volume values related to specific use cases. For instance, in the Federation of Tumor Segmentation (FeTS), where three labels (WT: whole tumor, TC: tumor core, ET: expanding tumor) are segmented, we extract volume values corresponding to these labels. The choice of this metadata is based on the assumption that it can be consistently extracted for all experiments, making it inherently available for all medical radiological images.Subsequently, the metadata vectors V M i for i = 1, 2, . . . , N obtained from each client (where N is the total number of federated clients), representing the extracted metadata, are transmitted to the server. Here, the EMD is computed for each client in relation to other clients within the federation. The choice of this distance is intentional as it has demonstrated its robustness to capture dissimilarity between probability distributions and thus provide valuable insights into the distributional disparities among clients (Luo et al., 2023) . Other metrics may be considered in alternative domains.In contrast to computing non-IID solely based on metadata, which may not capture the intricate characteristics within the data, we investigate computation in the embedding space of data for its capacity to capture deeper and richer data representations. Given the domain (i.e., medical imaging), a large publicly-available pre-trained MedicalNet model (Chen et al., 2019) is utilized for embedding extraction in our proposed methodology. Training on the diverse 3DSeg-8 dataset, covering a wide range of modalities, organs, and pathologies, has resulted in the development of a set of heterogeneous 3D neural network models. We adopt the pre-trained 3D-ResNet18 network to extract embedding features from the local data on each client. Specifically, image features are extracted from the network' bottleneck layer, known for providing concise and essential representations of input data. We compute the angles between Principal Component Analysis (PCA) components (where PCA = 2) to validate their alignment (more details in Appendix E). These components compress vectors before transmission to the server, thereby enhancing privacy as reconstruction accuracy drops by decreasing the number of components (Reddy and Jaya, 2021).Subsequently, all client embedding vectors, denoted as V Ei for i = 1, 2, . . . , N , each with a dimension of client samples×(512, 2) and flattened to client samples×(1024), are transmitted to the server. These vectors are then utilized for the computation of EUC according to Equation (1):EQUATION)where V Ei and V Ej represent feature vectors from any two clients.For generalization, we propose downgrading the most distant client (outlier) in the federation. To achieve this, we utilize the distances between clients computed during the non-IID assessment step and incorporate this information into weights (ω) assigned to each client. This approach is implemented using the FedAvg algorithm for the proof of concept, while noting that any FL averaging method could be employed. The proposed update of the global model in FedAvg w is defined in Equation (2):EQUATIONwhere w t+1 is the updated global model, w i,t represents the weight for client i at iteration t, f i (w t ) is the local objective function for client i, and N is the total number of clients.To identify a client for down-weighting, let D EMD and D EUC represent the matrices of EMD and EUC distances between clients, respectively. By summing along each column, we identify the client with the highest total (distant client) according to Equation (3), where D represents either D EUC or D EMD :i max = arg max i   N j=1 D ij   .(3)Following that, we degrade its contribution by applying the arbitrary weight values (ω) of 0.1, 0.3 and 0.5 to demonstrate the trend in improving the performance. This procedure enables us to assess the impact of reducing the influence of a single distant client (i max ) on the entire federation, thereby providing a clear illustration of the correlation between computed distances and the performance of the global model. The base weights (ω) for the non-downgraded clients are 1.Then, in our personalization strategy, we use the same distance matrices D EMD and D EUC to build clusters of closely-related clients. The clustering algorithm minimizes total distances within each cluster (i.e., C 1 and C 2 ). The process is outlined in Algorithm (1). This clustering approach promotes effective collaboration and information exchange among clients with closer data distributions within each cluster. We limited the evaluation to only two clusters to show the benefit of the clustering approach and its impact on improving performance per client, while maintaining a reasonable number of experiments.max i max ← arg max i j D ij ; C 2 ← {i max } ; // Initialize cluster C 2 with i max C 1 ← {i | i ̸ = i max } ; // Initialize 
datasets: 
 We used three publicly available FL datasets for our study. FeTS 2021 (Pati et al., 2021) consists of glioblastoma multi-modality MRIs from multiple sites, with WT, TC, and ET segmentations. For our experiment, we selected four (4) clients, ensuring a balanced distribution of samples (Hospital 6 : 34 samples, Hospital 13 : 35 samples, Hospital 20 : 33 samples, and Hospital 21 : 35 samples). The multi-site prostate MRI segmentation dataset (Liu et al., 2020b) , (Liu et al., 2020a) features T2-weighted MRIs with prostate segmentation masks. We used four (4) balanced clients for our experiments: Client 1 : 39 samples, Client 2 : 32 samples, Client 3 : 40 samples, Client 4 : 39 samples. The Fed-KITS2019 dataset (Terrail et al., 2022) focuses on kidney and tumor segmentation in CT scans. We created a 5-client federated version, excluding one site (Client 6 : 30 samples) for a balanced distribution of samples: Client 1 : 12 samples, Client 2 : 14 samples, Client 3 : 12 samples, Client 4 : 12 samples, Client 5 : 16 samples. Our focus on balanced and small federations promotes equal and rapid contributions from each client, facilitating equitable evaluation of the FL model across datasets. The data in each group were divided into training (80%) and validation (20%) sets, as originally proposed.Additionally, we redistributed labels among clients to create non-IID federations, diversifying client distributions. We followed the methodology proposed in (Luo et al., 2023) . For non-IID federations, we aimed to maintain consistency in assigning training and validation samples across each set. However, if they were not available on the same client, adjustments were made, potentially resulting in differences in the selected samples, while preserving set sizes. Please refer to Appendix A for more details regarding the data, and Appendix D for information on building non-IID federations.We used a 3D U-Net network (Ronneberger et al., 2015) along with the SGD optimizer with a learning rate of 0.01 and momentum of 0.9, employing DiceLoss for training, following a standard protocol for medical segmentation. Training was conducted using FedAvg (McMahan et al., 2017) , along with its weighted variant (Equation 2), across 25 global epochs as a balance between model convergence and optimized experiment time. We assessed performance using the Dice Metric, considering inter-client standard deviation for variation. We compared our methods against two representative FL algorithms: FedProx (Li et al., 2020) and DITTO (Li et al., 2021a) . Please refer to Appendix B for further details. 
generalization optimization: 
 To further reduce computation of D EM D for FeTS clients with many modalities and labels, we considered only the metadata where EMD values are maximum (EMD max ) among others as having the greatest negative impact. This applies to both maximum intensity (EMD maxI ) and label volume values (EMD maxL ). Please refer to Appendix C for more details. Table 1 presents the correlation between EMD and the performance of the FedAvg global model, indicating that as EMD values increase, performance decreases. We refine the federation with the lowest performance A,B,C to enhance generalization and personalization strategies. Analyzing the distance matrices provides valuable insights into the relationships between different clients (Table 2 ). In FeTS A , both EMD and EUC highlight client 1 as the most distant within the federation while for Prostate B , client 4 emerges as the most distant. In Fed-KITS2019 C , while both metrics suggest significant differences between clients, the most distant client differs. This discrepancy is attributed to the nature of the data each distance operates on. EMD analyzes metadata distribution, while EUC operates in the embedding space, potentially capturing different data features. As a result, certain clients may appear more distant in one assessment compared to the other. Table 3 compares results across various learning approaches, including our down-weighting strategy (FedAvg w ) for distant clients. In FeTS A , FedAvg w leads to a relative increase in performance compared to the default FedAvg approach and FedProx. For Prostate B , it significantly enhances performance (+5% vs FedAvg, +16.6% vs FedProx). Similarly, in Fed-KITS2019 C , it improves Dice performance regardless of distance metric (EMD or EUC).  
conclusion and discussion: 
 Our study underscores the significance of assessing client data heterogeneity (non-IID) in medical imaging FL to optimize both generalization and personalization goals. We propose a down-weighting strategy to enhance global model performance across datasets by reducing the impact of a distant client. Additionally, we advocate for distance-based clustering of clients as a personalization solution to enhance medical imaging segmentation accuracy across diverse datasets.While promising, our study is limited to medical imaging, particularly volumetric radiographic datasets, and prioritizes balanced scenarios and small federations for faster computation and proof-of-concept purposes. Future research should explore unbalanced scenarios, larger federations, alternative architectures for embedding extraction to broaden the applicability of our proposed strategies. The formula to compute the angle between two PCA components v and w is given by: Angle(v, w) = arccos v • w ∥v∥ • ∥w∥ where • represents the dot product, ∥v∥ and ∥w∥ represent the magnitudes of vectors v and w respectively, and arccos is the inverse cosine function. We iterate over pairs of clients and computes the angle between the corresponding components for a random sample on each client. 
 
 Appendix:  


Decision: The reviewers emphasized the novelty of the method and the quality of the experimental evaluation. The authors engaged in the rebuttal process.

Reviews: review: paper_summary: paper_summary: This paper addresses non-IID data challenges in healthcare federated learning (FL) by proposing a novel assessment methodology for medical segmentation, leveraging distance measurements and metadata analysis to enhance FL performance. It introduces a novel approach in medical imaging FL that minimizes outlier effects and employs distance-based clustering for model personalization. Validated on three radiology imaging datasets, the proposed method significantly improves model generalization and personalization across diverse scenarios.
main_review: main_review: I am satisfied with the response regarding concerns about hyper-parameters and have increased my score to a borderline accept. It would be beneficial to emphasize the key hyper-parameters in the experimental section. Additionally, using 'i' to represent individual clients seems unnatural, as it is commonly used to denote an index.
strength_weakness: strength_weakness: - The experiments fail to highlight the key hyperparameters influencing the proposed method's performance and omit some ablation studies on these significant hyperparameters.
- Algorithm 1 appears to be not self-contained and need more captions; the variable "i" is unclear, as it is also used as an index for D_{i,j}, creating confusion.
questions: questions: N/A
paper_summary: paper_summary: This work develops novel federated learning techniques to mitigate the loss in performance due to the non-IID nature of realistic decentralized datasets. The authors propose two distances to quantify the degree of dissimilarity of the data distribution of two clients in a federated fashion. They use these distances both in a rescaling fashion in standard weighted Federated Averaging (FedAvg), and as a distance matrix in standard clustered federated learning. The method is validated on three different medical image segmentation datasets.
main_review: main_review: I appreciated the discussion period with the authors. The clarification of the results as well as the addition of different Appendices greatly improved the readability of the paper and its value. I think that both the two proposed Non-IID metrics and the choice of datasets to work on are interesting and the community would benefit from it being published.

However, some of my concerns remain. The choice of clustering algorithm is questionable, although researchers applying this method could use something else. For one of the Non-IID metrics (EUC distance), clients compute local PCAs before communicating their embeddings and computing euclidean distances between these server-side. Appendix E shows that the components of each local PCA can be quite different from one client to another, confirming my concern. The angles between principal components of different clients are relatively small (at least for the first two components) for datasets FeTS and Prostate, but quite high for the last dataset. I cannot say if one can fully trust the current results for the EUC distance matrices for this last dataset, although in practice this problem could be handled by Federated PCAs or local PCA components communications, making the method probably viable.

Thus I keep my rating to borderline. I could consider it accepted depending on MIDL standards, if the authors can add in the article one or two sentences to explain why Appendix E is added, and what can be done to counter this problem.
strength_weakness: strength_weakness: **Euclidean Embedding Distance computation** The description of the computation of this distance is unclear to me. In particular page 4 on client embedding vectors
> "Principal Component Analysis (PCA) with 32 components is used to compress them before transmission to the server, enhancing privacy",

followed by 

> "Subsequently, all client embedding vectors, denoted as $V_{Ei}$ for i = 1, 2, . . . , N , each with a dimension of (512, 32) and flattened to (16384), are transmitted to the server." 
    
I am unsure where does the dimension (512,32) for the vector of one client come from. Moreover, how is the PCA actually computed? I would assume either on server-side after transmission, or in a federated fashion? It seems to me that the transmitted vectors $V_{Ei}$ and $V_{Ej}$ of two clients which independently computed a PCA do not lie in the same sub-space, I don't think an euclidean distance between them would be usable.


**Clustering algorithm** The choice of clustering algorithm seems quite arbitrary and task specific. It always gives a cluster $C_1$ of 2 clients. I understand that small federations is the use case of interest, but more established clustering algorithms such as spectral or agglomerative clustering would potentially make the method more broadly applicable. Could you please elaborate on the choice of the clustering algorithm?
    
**Experiments description** The experiments of part 4.1 are only partially described. I could not find information on the actual non-IID splits apart from a citation, even in Appendices. It makes reading Table 1 difficult and further results quite misleading. It would greatly improve the readability of the paper to add more information on the different federations, potentially in Appendix, as well as writing the federation used for each dataset in Tables 2, 3 and 4 in the captions as I am still unsure about this information.
    
**Comparison to other methods** Only few state-of-the-art methods are compared in the paper. I don't think it is a problem for part 4.1 as its goal as I understand is mainly to validate the usage of the introduced non-IID distances. However, the clustering federated learning results would benefit to be compared with state of the art clustered methods, with the closest to my knowledge being CFL : Avishek Ghosh et al., “An Efﬁcient Framework for Clustered Federated Learning,”. The only difference with the proposed method lies in the distance matrix as they use the cosine similarity between local updates of two clients. If possible, comparing the clustering results of CFL and the proposed method could strengthen the results.
questions: questions: See the weaknesses above
- Clarification of the computation of the EUD,
- Clarification of the experiments,
- Some comments on the clustering algorithm,
- Ideally a comparison with CFL.
paper_summary: paper_summary: This work presents a practical non-IID assessment methodology for a medical segmentation problem, highlighting its significance in medical FL.
main_review: main_review: I would like to see the author add more comparisons with other related works that I have mentioned. The novelty and contribution also need to be clarified, because the current version lacks strong motivation.
strength_weakness: strength_weakness: The first concern is regarding the novelty. The non-iid is a common issus in federated learning, as well as MRI segmentation. And the authir use a very commonly used technique to solve it, i.e., personalized FL. For example, Specificity-Preserving Federated Learning for MR Image Reconstruction TMI 2022. It is hard to see the novelty of this work. Thus, I would like to see the motivation for the author to use such a method or the new contribution of this work, not only the downstream task, i.e., segmentation.

The second concern is regarding the related works, I would like to see a more complete review on the medical FL. Actually, benefited from the large scale pretrained models, many medical FL also build on them to significantly improve the medical FL performance. For example, Learning Federated Visual Prompt in Null Space for MRI Reconstruction, CVPR 2023.
questions: questions: see the Weaknesses

Comments: title: Insufficient justification of client local PCAs
no need to cite if not relevant
Reviewer Recommendation for Citation
comment: Dear Reviewer,

Thank you for your active participation in the discussion and review of our manuscript. We appreciate your suggestion to compute angles between the principal directions of each client.

In response to your concern, we have taken steps to compute the angles between components and have incorporated the findings into our manuscript, as detailed in Appendix E. We adjusted the step of the PCA computation by choosing the components with the smallest angles between them (in particular, 2 components). Additionally, we recomputed the distances in Table 2. We obtained the same findings in terms of both a distant client and clustering assessment.

We believe that these revisions enhance the clarity of the computation of PCA on the client side. However, we acknowledge that Federated PCA could be considered for further studies, as proposed.

Thank you once again for your feedback!
Dear Authors,

I first thank you for your rebuttal. Most of my remarks were relatively minor apart from my concern about client local PCAs. I agree with what is said for these minor comments, and greatly appreciated the modifications of the article, especially Appendix D. 

However, I don't think that your argument on the validity of computation of EUC is enough. Clients do perform independent PCAs in your method, there is a possibility that the embedding vectors transmitted to the server are not comparable using a simple coordinate-based euclidean distance. I understand that the same model and pre-processing were used to obtain the pre-PCA embeddings. But at the same time, this model is used to capture differences between the probably heterogeneous data of different clients. I would not be surprised if their computed PCAs were at least slightly different. This "slightly" is the important part to me, we just do not want the differences in PCAs to affect to much the computed distances. I think this could be solved by either a client computing its PCA and communicating it to the others, using Federated PCAs (https://arxiv.org/pdf/1907.08059.pdf), or something else anyway.

Thus, I do not question the method in itself which I find very interesting, just the current definition of the PCAs and the presented values in the EUC distance matrices of the paper *which, without further verification, could be random*. For the sake of validation of the paper experiments, I would please ask you to compute angles between the principal directions of each client (angles between first directions of client 1, ..., 5, angles between second directions of 1, ..., 5, etc.) or any other form of distance between projection subspaces for at least one or two datasets, to verify your argument for the experiments of the paper. A small discussion of this aspect could be added in the Method part, proposing either local PCAs or another form for future applications of your method.

All the best,
-.
Dear Reviewer, 

Thank you for your thorough review of our manuscript, and for providing detailed feedback on both the strengths and weaknesses of our work. We appreciate the time and effort you have invested in evaluating our research. We have carefully considered each of your comments and would like to address them below: 

Euclidean Embedding Distance computation: To clarify, the dimension (512,32) for the vector of one client is derived from the embedding vectors obtained from the client's local data, specifically from one MRI or CT scan (producing (512, 32)). We applied PCA=32 at the end of the bottleneck layer of the pre-trained network to compress these vectors. In our manuscript revision, we will provide a more detailed explanation of this computation. Regarding the question of whether computing PCA in this separate manner could affect the value of such spaces for Euclidean Distance computations, we believe that two key factors support its validity: 1) the consistent use of the same preprocessing data pipeline, and 2) the utilization of the same pretrained feature extractor (MedicalNet). These two points lead us to assume that the final features reside in the same subspace. 

Clustering algorithm: We acknowledge your concern regarding the choice of clustering algorithm and its potential impact on the method's applicability. While the current clustering algorithm may seem arbitrary and task-specific, we intentionally opted for simplicity and suitability for small federations, aligning with our specific use case. Our goal was to demonstrate that even with a straightforward clustering method, we can effectively enhance performance using our proposed methodology. The value of our approach lies not in the sophistication of the clustering algorithm, but rather in how we implement its clustering and what distance we rely on. Furthermore, it is important to note that our methodology is not limited to a particular clustering approach. Any clustering method aimed at grouping closely-related clients could be integrated with our proposed technique. 

Experiments description: We apologize for the lack of detailed description of the experiments in Section 4.1. In our Datasets Section (3.1), we briefly outlined the process of obtaining the federation by redistributing labels among clients to create non-IID federations, diversifying client distributions. This methodology was based on the approach proposed in Luo G, Liu T, Lu J, Chen X, Yu L, Wu J, Chen DZ, Cai W. Influence of Data Distribution on Federated Learning Performance in Tumor Segmentation. Radiol Artif Intell. 2023 Apr, specifically focusing on label size as a criterion. To ensure consistency in assigning training and validation samples across each set within non-IID federations, we aimed to maintain uniformity. However, adjustments were made if the required samples were not available on the same client, potentially resulting in differences in the selected samples while preserving set sizes. Though, we will provide additional information on the non-IID federations used in the Appendix to clarify the experimental setup.  

The other minor improvements proposed by you will be integrated for better readability. 

Thank you once again for your valuable feedback!
Thank you for bringing this up.

There is no need to cite if you do not think that these papers are relevant, you can address the reviewers comments by including citations to other more relevant work in your opinion. This has been noted and will be kept in mind during decisions.

Best regards and thanks for engaging with the review process!
Dear Reviewer, 

Thank you for your thorough evaluation of our work and for providing detailed feedback on the novelty and contribution of our methodology. We have carefully considered your comments and would like to address each concern raised. 

Regarding the novelty and contribution of our methodology, we understand your perspective on the familiarity of non-IID issues and personalized federated learning techniques within the context of medical federated learning (FL). However, we believe our work offers several significant contributions to the field. 

Firstly, our methodology introduces a practical non-IID assessment specifically tailored for medical segmentation problems. Our work is related to the area of personalization, but we are broadly examining the question of non-IID in medical FL. In particular, we propose a methodology that not only effectively measures non-IID but also serves as a tool for improving performance within federation. We find this practically useful because working with a new dataset requires initial steps to measure the level of non-IID in order to choose the optimal strategy. The decision to consider the segmentation use case is based on its interest from the community and the number of available federated learning benchmarks (FeTS, FLamby, FL-PoST, and others), which provide the basis for fair comparison. Additionally, we observe that there has not been much research done on medical image segmentation and non-IID, as also stated by Luo G, Liu T, Lu J, Chen X, Yu L, Wu J, Chen DZ, Cai W. Influence of Data Distribution on Federated Learning Performance in Tumor Segmentation. Radiol Artif Intell. 2023 Apr. 

We appreciate your reference to "Specificity-Preserving Federated Learning for MR Image Reconstruction TMI 2022" as a relevant work in the field of personalized strategies. Indeed, the approach proposed in that work aligns with our literature review. The authors propose to divide the MR reconstruction model into two parts: a globally shared encoder to obtain a generalized representation at the global level, and a client-specific decoder to preserve the domain-specific properties of each client. In the reference (Krishna Pillutla, Kshitiz Malik, Abdel-Rahman Mohamed, Mike Rabbat, Maziar Sanjabi, and Lin Xiao. Federated learning with partial model personalization) the authors consider this way of partitioning deep learning models where some layers (like decoder part can be personalized), which is exactly proposed in Specificity-Preserving Federated Learning for MR Image Reconstruction TMI 2022. To sum up, while the area of application differs, with our focus on medical segmentation rather than MRI reconstruction, the fundamental concept of “partitioning” deep learning models to enable personalization remains consistent.  

We believe that the work you highlighted “Learning Federated Visual Prompt in Null Space for MRI Reconstruction”, as well as, for example, "FedFMS: Exploring Federated Foundation Models for Medical Image Segmentation," align with our research direction proposed in the manuscript. While the use of large-scale pretrained models is indeed a promising avenue for improving FL performance, we believe that it could be also referred to the pretraining methods (the authors again use the ideas of pretraining on massive data and later fine-tuning). Our focus is on the importance of measuring non-IID data and strategically splitting clients into groups based on distance measurements or using degradation of a distant client. In fact, our methodology could potentially enhance the effectiveness of pretrained models by ensuring that clients with similar data distributions are grouped together for more efficiency. 

Thank you for bringing this to our attention as possible further direction of our research!
Dear Reviewer, 

Thank you for the detailed review of our submission. We appreciate the time and effort you have put into evaluating our work. We would like to address the concerns raised regarding hyperparameters and algorithm clarity. 

Hyperparameters: We acknowledge your feedback on hyperparameters. However, it is important to note that our proposed method intentionally aims for simplicity and general applicability. The algorithm primarily leverages classical mathematical tools and existing metrics, such as Earth Mover’s Distance and Euclidean Distance, rather than introducing additional hyperparameters. The only new hyperparameter introduced in our approach is the weight parameter mentioned, particularly in Table 3, which adjusts the contribution of a distant client to the overall model generalization process. We provided the values for this parameter in the article to ensure transparency and reproducibility. Furthermore, the algorithm used for clustering, as detailed in the manuscript, does not depend on hyperparameters. It employs distance-based clustering techniques that are well-established and do not require tuning specific parameters for effective implementation. In summary, while we appreciate the importance of hyperparameters in machine learning algorithms, our approach intentionally minimizes their presence to maintain simplicity and applicability across different scenarios. 

Algorithm Clarity: To address the weakness regarding the clarity of Algorithm 1 and the confusion surrounding the variable "i," we would like to provide additional clarification. In the manuscript, when presenting the generalization and personalization strategies, we clarify that the variable "i" represents individual clients within the federation, while "i_max" denotes the client identified as the most distant from the others. Additionally, when defining the distance matrix Di,j , we use the same variable "i" to denote clients within the federation, emphasizing that this matrix captures the relationships between clients in the federation. To further enhance clarity, we will include comments in Algorithm 1, including one stating that "i" represents a client in the federation. 

Thank you once again for your valuable feedback, and we look forward to addressing any further questions or concerns you may have.
Dear Program and Area Chairs,

We have carefully considered all the comments and suggestions provided by the reviewer. However, I would like to bring to your attention a specific request made by Reviewer aSJy.

Reviewer aSJy has expressed satisfaction with our manuscript overall but has recommended the inclusion of citations to two papers, namely:

1. Feng, Chun-Mei, et al. "Specificity-preserving federated learning for MR image reconstruction." IEEE Transactions on Medical Imaging (2022).
2. Feng, Chun-Mei, et al. "Learning federated visual prompt in null space for MRI reconstruction." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

We are more than willing to include any citations that directly contribute to the manuscript's scientific merit. However, both of these papers focus on MRI reconstruction, whereas ours explores medical image segmentation. These papers are not directly related to our paper's scope, methodology, or findings. As it is evident that these papers are from the same author, we believe Reviewer aSJy has asked us to cite these papers to improve their citations. It is necessary to bring to your attention that such a practice should be discouraged. Furthermore, incorporating these unrelated citations might detract from the coherence of our paper, potentially confusing readers regarding the focus and contribution of our study.

Therefore, we kindly request your guidance on proceeding with this recommendation while maintaining the submission's integrity.

Thank you!