Improving Identically Distributed and Out-of-Distribution Medical Image Classification with Segmentation-Guided Attention in Small Dataset Scenarios

Authors: mariia  rizhko
lauren  erdman
mandy  rickard
armando j lorenzo
kunj  sheth
daniel  alvarez
megan a bonnett
christopher s cooper
gregory e tasian
anna  goldenberg
Venue: MIDL_io_2024_Conference
Published: 2024-02-01 02:36:06

Abstract: We propose a new approach for training medical image classification models using segmentation masks, particularly effective in small dataset scenarios. By guiding the model’s attention with segmentation masks toward relevant features, we significantly improve accuracy for diagnosing Hydronephrosis. Evaluation of our model on identically distributed data showed either the same or better performance with improvement up to 0.28 in AUROC and up to 0.33 in AUPRC. Our method showed better generalization ability than baselines, improving from 0.02 to 0.75 in AUROC and from 0.09 to 0.47 in AUPRC for four different out-of-distribution datasets. The results show that models trained on smaller datasets using our approach can achieve comparable results to those trained on datasets 25 times larger.

Main paper: 
 introduction: 
 Can a model trained to predict medical diagnoses for young children maintain the same level of accuracy for older children? If a model is trained with data from one hospital, will it perform just as efficiently with data from another? How precise will a model be when analyzing images produced by different machines? Machine Learning (ML) models struggle with these scenarios since they are trained on identically distributed (i.i.d.) data. However, their accuracy can vary significantly when these models are tested on out-of-distribution (OOD) data. The problem is known as domain shift between i.i.d source and OOD target domains. It occurs when models are not trained to deal with the domain shift in mind. This issue is significant for the medical field, where labeled data is limited, and training different models for each scenario is impractical.Domain shift is a challenge that extends beyond healthcare. The task of addressing domain shift is known as Domain Generalization (DG), a problem that exists in almost every application of ML (Zhou et al., 2022) . For example, in the semantic segmentation task in autonomous driving, a model trained on urban data may fail in rural settings (Hoffman et al., 2018; Ros et al., 2016) , potentially leading to accidents. In personal identification systems, the model trained on well-illuminated images may not recognize a person in dim lighting (Sun et al., 2019; Li et al., 2020) , potentially preventing access to their home if the lights are broken. Even with seemingly simple tasks like recognizing handwritten digits, ML models can underperform due to minor variations like ink color (Ganin and Lempitsky, 2015) . Similarly, in the medical domain, a model trained on images collected with one protocol might be ineffective for images collected through another (Liu et al., 2020) . These examples underline the significance of the problem across different domains.An excellent survey (Zhou et al., 2022) categorizes various DG methodologies. The Domain Alignment (Li et al., 2018b,d) methods focus on learning a mapping function between the domain and target datasets. Meta-learning (Li et al., 2018a; Balaji et al., 2018) approaches divide data into meta-train and meta-test sets, where a model is trained on the meta-train set and evaluated on the meta-test set. The methods separate domain-specific and domain-agnostic features within datasets in the Learning Disentangled Representations (Li et al., 2017; Ilse et al., 2020) category. While Domain Alignment, Meta-learning, and Learning Disentangled Representations offer promising approaches, they require a labeled target dataset during training on a domain dataset. The target dataset is usually unavailable during training in the medical domain, so other approaches should be used.The DG survey (Zhou et al., 2022 ) also covers the methodologies that do not require a target dataset during the training. Data augmentations (Volpi et al., 2018; Volpi and Murino, 2019; Xu et al., 2020) simulate a domain shift by changing images. Ensemble learning (Xu et al., 2014; Cha et al., 2021) trains the same model with a different random seed for weight initialization or data split. Self-supervised learning (Carlucci et al., 2019; Bucci et al., 2021) lets a model learn generic features of your data first and then fine-tune the model for a downstream task. Regularization Strategies (Wang et al., 2019; Huang et al., 2020) learn generalized features by focusing on global structure instead of local patterns or by masking out over dominant features. All of these approaches are generally considered to make more robust models. However, when trained on small datasets, which is usually the case for the medical domain, their performance might suffer on i.i.d and OOD data. before after before afterFigure 1 : Illustration of our approach. Before: focus on both organ and background noise.After: targeted focus on the critical organ.To address this limitation we utilize Gradient-weighted Class Activation Mapping (Grad-CAM) (Selvaraju et al., 2017) to create attention maps. GradCAM is a visual explanation method in computer vision that highlights the regions in an input image that influence the model's outcome the most. Prior work showed that this attention mechanism could be learned, resulting in a better performance for image segmentation (Li et al., 2018c) and classification tasks (Fukui et al., 2019) . This idea has been adapted for the medical domain: it showed improved accuracy for thyroid nodules (Lu et al., 2022) , for chest X-ray abnormality localization and diagnosis (Ouyang et al., 2020b) , for diagnosis of COVID-19 (Ouyang et al., 2020a) and dementia (Lian et al., 2020 (Lian et al., , 2019)) . However, these prior studies used large datasets, ranging from 2,000 MRI scans to as many as 1.2 million images from ImageNet, and they focused only on i.i.d. data. In contrast, we apply the idea to DG tasks and demonstrate its effectiveness on small datasets with less than 100 images.Our study addresses the typical scenario in the medical field where models are trained on small datasets. Typically, these models learn specific "useful" noise patterns, leading to high performance on similar (i.i.d) test datasets. However, their accuracy declines when applied to new images without these noise patterns. Beyond basic classification tasks, models should also be trained to disregard features known a priori to be irrelevant. For instance, in kidney ultrasound classification (see Figure 1 ), the model should focus only on the kidney, ignoring background noise. Often, in medical imaging, additional information like segmentation masks is available. We adapt gradient-based techniques to utilize segmentation masks for medical imaging with small dataset scenarios. This adaptation allows us to effectively train models on small datasets and improve their performance when tested on i.i.d and OOD data without having target datasets during training. 
method: 
 Let X be the input image space and Y the label space. A domain is defined as a jointdistribution D = (X , Y), which contains image-label pairs {(x (n) , y (n) )} N n=1, where N is the number of samples. Our goal is to learn a classification model F θ : X → Y using the source domain D for generalization across unseen target domains {D 1 tg , D 2 tg , . . . , D K tg } set of K target domains. In a source domain for input images X , we have corresponding segmentation masks M that will be utilized for our method. Note that there are no requirements for masks in target domains. The core idea of our method is to force the model to learn two things simultaneously: the attention mechanism learning task and the classification task itself.Attention Map Calculation. Given an input image, a classification model processes it up to a target layer. Let A k represent the activation of the k-th feature map at this layer. The gradient of the score for class c, denoted y c , with respect to the activations A k of the feature map is computed. This gradient is represented as ∂y c ∂A k . To obtain the neuron importance weights α c k we apply Global Average Pooling (GAP) to these gradients. This is given by:α c k = GAP ∂y c ∂A k (1)The Class Activation Map (CAM) for class c, denoted as L c Grad-CAM , is a weighted sum of the feature maps, weighted by α c k , and passed through a ReLU function:L c Grad-CAM = ReLU k α c k A k (2)The final Attention Map A is achieved by resizing L c Grad-CAM to the dimensions of the input image.Attention Loss. It is a custom loss function, denoted as L Attention , that incorporates the difference between the Grad-CAM Attention Map A and a given ground truth attention mask M by calculating the mean squared error (MSE) between A and M:L Attention = 1 N N i=1 (A i -M i ) 2 (3)where N is the total number of pixels in the image, and i indexes these pixels. This loss function measures the alignment between the regions highlighted by the Grad-CAM and those indicated by the attention mask. The objective of the training is to minimize this loss, thereby encouraging the model to focus more on areas marked as important by the mask.Overall Loss. The Binary Cross-Entropy Loss L BCE , given the predicted outputs y pred and the true labels y true , is defined as:EQUATION)where N is the number of samples and i indexes these samples. The overall loss is a weighted combination of the Binary Cross-Entropy Loss L BCE and the Attention Loss L Attention .L = αL BCE + βL Attention (5)where α and β are weighting coefficients that balance the two components of the loss. In all our experiments α = β = 1.By combining these two losses, the model not only focuses on minimizing the prediction error but also emphasizes the alignment of the attention maps with the important regions as marked by the attention masks. 
experiments: 
 Hydronephrosis (HN) is a medical condition characterized by the swelling of one or both kidneys due to a urine buildup. It can affect people of any age and is spotted in up to 5% of babies during routine pregnancy ultrasound scans. However, surgical intervention is only required in 20% of these cases (Dos Santos et al., 2015) . To determine which cases need intervention, patients receive repeated invasive scans and ultrasounds to monitor whether the HN is causing functional damage or resolving without the need for surgery. Previously, deep learning models have used postnatal ultrasound images to predict surgical intervention in HN from the first ultrasound (Erdman et al., 2020) , further investigated predicting HN grades (Smail et al., 2020) and risk scores (Tabrizi et al., 2021) . While prior models (Erdman et al., 2020; Smail et al., 2020; Tabrizi et al., 2021) worked well for i.i.d data, they showed lower performance on smaller and OOD datasets.Datasets. We use five datasets from four pediatric hospitals in North America containing ultrasounds of kidneys. The variation of the data comes not only from its collection across various hospitals but also from differences in patient demographics and imaging equipment, attributing to its OOD characteristics. For example, the average patient age in the Hospital for Sick Children (SickKids) is 53 weeks, while in the Children's Hospital of Philadelphia (CHOP), it is 313 weeks. These variances provide a robust setting for evaluating models for DG.The source domain dataset D from SickKids has 2542 ultrasounds. 20% of D is held out to create i.i.d. test dataset D test . The rest of the 2048 images are used for training and validation. Only 83 out of 2048 images have corresponding kidney segmentation masks. We will utilize these 83 images for training baseline models. The same 83 images with their segmentation masks will also be used to train our model. To assess the robustness of our model, we will further use the complete set of 2048 images, which is bigger by 25 times, to train additional baseline models. We split D into train D train and validation D val sets, ensuring each patient's images appear in only one set. D train has only 66 images with kidney segmentation masks, creating D seg train with 51 non-surgical and 15 surgical cases. Similarly, the validation set D seg val is a subset of D val and has 10 non-surgical and 7 surgical cases. We evaluate models on four distinct target domain datasets. The first, T SickKids , includes data from 202 patients at SickKids, having 711 images, of which 75 are positive. Despite being collected at the same hospital as the training dataset, patient demographics and imaging equipment variations make this dataset OOD. The second dataset, T Stanf ord , is from the Stanford Children's Hospital (Stanford) and includes data from 103 patients, with 551 images (27 positive). The third, T U Iowa , is from the University of Iowa Children's Hospital (UIowa) with 91 patients and 97 images (56 positives). Lastly, T CHOP comes from CHOP with 89 patients and 89 images, 60 of which are positive. The datasets summary is shown in Table 1 .Baselines. ResNet-18, ResNet-50 (He et al., 2016) , ViT-Tiny, and ViT-Base (Dosovitskiy et al., 2020) were trained on D seg train and validated on D seg val , utilizing Binary Cross Entropy Loss L BCE only. We tested two weights initialization methods: Kaiming uniform initialization (random) (He et al., 2015) and using weights pre-trained on ImageNet (Deng et al., 2009) . Hyperparameters were tuned via Bayesian optimization (Snoek et al., 2012) to minimize the loss on D seg val , more details in Appendix A.Consistent image transformations For further analysis, we also trained additional baseline models on the larger datasets D train and D val ; all experimental setups were the same.Our model. We trained the ResNet-18 model on D seg train and validated it on D seg val , utilizing our proposed loss function as described in Equation 5. The model's weights were initialized using pre-trained ImageNet weights. All other experimental setups, including hyperparameters search, image transformations, and training duration, were consistent with those used in the baseline models.3.1. Baselines vs. Our Model trained on the Small Datasets.The results in Table 2 and Table 3 show Area Under the Receiver Operating Characteristic (AUROC) and Area Under the Precision-Recall Curve (AUPRC) of the baseline models and our model, all trained on D seg train and validated on D seg val .I.i.d. Comparison.Table 2 presents a comparative analysis of the models performance on i.i.d. test dataset D test . Note, D test is a held-out test dataset from the whole dataset D and has 494 images, while the models are trained on the small subsets D seg train and D seg val with 66 and 17 images respectively. This comparison reflects each model's ability to generalize to new data with a similar distribution to the training set. Interestingly, only three models, including our own, were able to effectively generalize to D test , showing 0.81-0.83 AUROC and 0.48 AUPRC.Table 3 presents the performance of the models across four different OOD datasets T SickKids , T Stanf ord , T U Iowa , and T CHOP . Notably, our model consistently outperformed all baselines across all OOD datasets. To further analyze our model, we trained additional baselines on the whole train dataset D train and validation dataset D val with a total of 2078 images. We compared the baselines to our model trained on D seg train and validated on D seg val with a total of 83 images. I.i.d. Comparison. Table 4 shows the overall performance of baselines and our model on the held-out i.i.d. test dataset D test . All models, including our trained only on 4% of the data, have comparable AUROC (0.82 -0.87) and AUPRC (0.47 -0.54), which means all models generalize well to unseen images from the same i.i.d. distribution.Table 5 shows models' performance on OOD datasets T SickKids , T Stanf ord , T CHOP , and T U Iowa . Out of 9 models that perform well on i.i.d. data, only three models, including ours, transfer well to all OOD datasets. It demonstrates the effectiveness of using our approach, considering that our model trained on 25 times less data could generalize well to OOD data.  
conclusion: 
 This paper presented a new method for improving medical image classification models using segmentation masks, especially effective in small dataset scenarios (less than 100 images). By utilizing a specialized loss function, our model demonstrated remarkable performance on both i.i.d. and OOD datasets despite limited training data. It matched or exceeded the performance of other models trained on similar-sized datasets in i.i.d. scenarios and consistently outperformed all baselines in OOD settings. Notably, our model, trained on just 4% of the data, showed the same or even better performance as baselines trained on significantly larger datasets in i.i.d. and OOD settings. The implications of these results are promising. Creating segmentation masks, which our method relies on, could be more feasible than gathering extensive data on rare diseases. Additionally, our model's ability to transfer across different hospitals could reduce the need for unique models for each medical setting. 
 
 Appendix:  


Decision: This paper introduces a loss function between the attention of a model computed by GradCAM and segmentation masks to improve image classification performance in small dataset scenarios. The idea was in general found interesting by the reviewer, but there were some concerns about novelty, comparisons to other SOTA approaches, difficulty in generation of segmentation masks. Novelty concerns were resolved in the rebuttal phase though concerns around limited comparisons and difficulty in obtaining segmentation masks remained to some degree. Overall, I agree with the reviewers that this is a borderline paper with positives slightly outweighing the negatives.

Reviews: review: paper_summary: paper_summary: This is a research paper that discusses the challenges of training machine learning models to analyze medical images. The paper proposes a new approach that uses attention maps to help models focus on the relevant parts of an image, such as a specific organ, while ignoring irrelevant background noise. The paper also introduces a new metric called Attention Score to measure how well the model is paying attention to the relevant areas. The paper highlights the importance of addressing the problem of domain shift, where models trained on one dataset may not perform well on another dataset. The proposed approach is shown to be effective on small datasets with less than 100 images.
main_review: main_review: This is a research paper that discusses the challenges of training machine learning models to analyze medical images. The paper proposes a new approach that uses attention maps to help models focus on the relevant parts of an image, such as a specific organ, while ignoring irrelevant background noise. The paper also introduces a new metric called Attention Score to measure how well the model is paying attention to the relevant areas. The paper highlights the importance of addressing the problem of domain shift, where models trained on one dataset may not perform well on another dataset. The proposed approach is shown to be effective on small datasets with less than 100 images.

From the following comments of the authors, now I understand better, resulting in change of my intial evaluation on this paper one step higher:

The segmentation masks are manually created by medical professionals our innovation
Its innovation lies in integrating these maps with a novel loss function
strength_weakness: strength_weakness: This paper lacks novelty. It adopts Gradientweighted Class Activation Mapping (Grad-CAM) (Selvaraju et al., 2017) to compute an attention map. Also, their training set is not randomly selected, but the ones with segmentation masks. Thus, the fair comparison of their method to other methods or models is questionable.
questions: questions: If the segmentation guided approach is applied to the other models in the experiments, this paper might be considered as a more acceptable paper.
paper_summary: paper_summary: The paper presents a method for improving medical image classification using segmentation masks and is particularly beneficial in scenarios with small datasets. This approach significantly improves model accuracy in diagnosing hydronephrosis by guiding the model's focus toward relevant features, showing substantial performance enhancements both on in-distribution (i.i.d) and out-of-distribution (OOD) datasets. The results demonstrate that models trained on smaller datasets using this method can match or outperform those trained on much larger datasets, emphasising the potential of segmentation-guided attention mechanisms in medical image analysis.
main_review: main_review: The authors have addressed my comments on the rebuttal process and I would like to thank them for their efforts. The paper introduces a method for training image classification models by guiding the model’s attention with segmentation masks. The approach is particularly efficient in small dataset scenarios. While the paper would benefit from additional experimental validations - especially to support the claims of DG-, the presented material is satisfactory. The reasons to accept this paper slightly outweigh the reasons to reject it.
strength_weakness: strength_weakness: * One of the weaknesses of the method is its dependency on segmentation masks which limits the practicality of the approach. In reality these masks may not be readily available or may be time-consuming and costly to produce.
* The paper lacks comparison with other state-of-the-art approaches. Experiments of the actual method are exhaustive and complete (even the hyperparameter selection is provided) but benchmarks with other methods are lacking making it hard to evaluate the actual improvement offered by the new approach.
* The dataset used is private, limiting reproducibility. Also, the data is not diverse in terms of demographics which might weaken the authors’ argument that the method works well on OOD datasets. No details were provided in terms of settings or imaging equipment used when collecting data.
questions: questions: Providing a fair and thorough comparison with at least one other DG method.
paper_summary: paper_summary: This paper presents an approach for training medical image classification models using segmentation masks and small dataset. The evaluation focused on improving the accuracy for diagnosing hydronephrosis. The major outcomes of this work is that the results show that models trained on small dataset can achieve comparable results to models based transfer learning trained on datasets 25 times larger. 
This paper investigates a significant issue arising when there is a limited dataset available for constructing ML/DL models.
main_review: main_review: I would like to thank the authors for their efforts during the rebuttal process. They have addressed some of concerns during the rebuttal process. However, the dependency on segmentation masks limitation is not fully addressed. Following the stated argument that creating segmentation masks is relatively straightforward, I would recommend adding the automatic segmentation to the proposed approach and test the method with manual/ automatic segmentations. In addition, the attention loss relies on the attention map from Grad-CAM. Different visualization methods produce different activation maps; how this impact on the proposed method is not discussed in the paper.
strength_weakness: strength_weakness: •	The novelty in this work is limited.

•	The literature review section lacks depth.

•	The idea of training classification models to ignore features of the ROI surrounding areas in medical applications is not convincing. Substantiating such a claim requires input from subject specialists, such as medical experts. It might be valid for some medical applications. However, in several applications, e.g. thyroid nodule classification in US, the identification of the nodule characteristics such as echogenicity requires information from the surrounding area of the nodule. 

•	This work assumes additional information like segmentation masks is available. This is not the case in several medical images. In addition the identifications of the ROI boundary in US imaging is highly subjective. 

•	The rational of selecting Grad-CAM is not clear. The are several enhancements to the Grad-CAM method, including Ablation-CAM [1] and Clustered-CAM [2]. These methods have demonstrated to be more accurate in US image classification task [4]. 

•	Several details about the datasets (for example the distributions of the image sizes, imaging equipment or machine maker, ..) are not presented in the paper. 

•	51 non-surgical and 15 surgical cases were used. There is no discussion about the impact of class imbalance on the proposed model.

•	The method of producing kidney segmentation masks is not described. It is not clear if the segmentation mask includes the whole boundary of the organ or just the inner boundary.

•	The rational of selecting the baseline networks is not clear. 

•	The paper does not present any future work.

[1]. "Ablation-cam: Visual explanations for deep convolutional network via gradient-free localization." proceedings of the IEEE/CVF winter conference on applications of computer vision. 2020.

[2]. "Clustered-CAM: Visual Explanations for Deep Convolutional Networks for Thyroid Nodule Ultrasound Image Classification." In Medical Imaging with Deep Learning. 2021. 

[3]. Classification of breast lesions in ultrasound images using deep convolutional neural networks: transfer learning versus automatic architecture design. Medical & Biological Engineering & Computing. 2024.
questions: questions: •	Provide the key contributions in the introduction section.
•	Comparison with other visualization methods (such as Ablation-cam and/or Clustered-CAM) will improve the quality of the paper.
•	I would like to see a discussion about the impact of class imbalance on the proposed model.
•	I would really like to see more discussion, medical views, and definition of ignoring what so called "background noise".
•	The paper could benefit from a more in-depth literature review.
•	I would like to see more about the future directions of this research.
paper_summary: paper_summary: In this paper, the authors propose to employ attention loss based on Grad-CAM to penalize the model for focusing on undesired image regions. Those regions are defined by segmentation masks. The authors note that this has already been done to improve model performance, but has not been shown on small datasets and OOD samples. They compare this approach with other models trained using standard binary cross entropy as baselines.
main_review: main_review: The paper show an interesting validation study that is useful for the community. It applies not only to ultrasound images but could extend to many other medical image analysis tasks. However, the paper is a bit hard to follow especially due to confusing data set naming. Also, the line between prior work and own contribution could be clearer. Also, one quite obvious experiment is missing
strength_weakness: strength_weakness: * It could be more clear, which findings (in terms of model performance improvement) have already been made by others and which findings are contributed here. Some important prior work is given in the methods section instead of the introduction
* The paper focuses on both sample efficiency and domain generalization but the introduction only focuses on the latter, while the title only names the former
* The Grad-CAM visualizations in Figure 2 only show the attention of the proposed approach. Without a comparison to the baseline models, this is of limited use. Figure 1 shows such a comparison. It is not clear why those figures are separate.
* The naming of the many different data sets is confusing
* Using multiple cohorts / domains for training is a straightforward way to improve domain generalization. This should have been compared here. How much would that alone improve performance? And how much could attention loss then add to that performance
questions: questions: see weaknesses

Comments: title: increase my evaluation by one step higher
Please read and respond to author comments
comment: From the following comments of the authors, now I understand better, resulting in change of my intial evaluation on this paper one step higher:

* The segmentation masks are manually created by medical professionals our innovation
* Its innovation lies in integrating these maps with a novel loss function
Thank you for your revision. I understand the rationale behind your answer to my last point regarding using multiple domains for training. However, you are proposing a method that improves domain generalization. One other way of improving domain generalization is using multiple domains during training. Of course, often times this data is not available. But you do have such data and could easily perform a comparison of these 2 approaches and assess how far your approach would lead to domain generalization compared to collecting multiple domains to train on.
Dear reviewers. The authors have posted responses to your reviews. Please take the time to read and respond before March 27.
We would like to thank everyone for their time and feedback! We are encouraged by the positive remarks on the significance of our work in studying identically distributed and out-of-distribution medical image classification in small dataset scenarios.

We have responded to each reviewer individually, and following the feedback, we have made the following improvements:

1. Updated the introduction section and changed the title to better reflect the scope of our work
2. Added Appendix C with information about the datasets to showcase various data distribution
3. Included Appendix D, in which we discuss future work  

We note the favorable comments from the reviewers, indicating that our approach is novel [T72m, 4V4q], the evaluation is comprehensive [T72m], and the methodology is rigorously explained [4V4q]. Thank you also for highlighting that attention Score is a valuable contribution to the field [T72m], our method could be useful for other ML/DL applications [HeNo], and the problem we solve is very important in medical image analysis [xY5m].

We thank everyone for their feedback and engagement and are ready to answer any further questions.
Thank you for your constructive feedback! Below, we address the key points raised in your review:

**Focus on Surrounding Areas:** We acknowledge your concern regarding the relevance of surrounding areas in various medical applications. In our experiments, we focus on Hydronephrosis which is the swelling of kidneys, and all the signs of Hydronephrosis are located inside the kidney. When transferring this approach to other medical conditions one should take a medical opinion on what is considered important in the ultrasound.

**Dependency on Segmentation Masks:** This is true that we need extra segmentation masks to train the model. However, in specific medical conditions such as Hydronephrosis, creating segmentation masks for kidneys is relatively straightforward and cost-effective compared to the significant effort and resources needed to collect larger datasets. Our approach leverages these masks to train models that can be easily adapted and reused across different hospitals, offering a practical solution in scenarios where data availability is a major constraint. As we show in Table 4 and Table 5, to achieve the comparable results you will need to collect 25 times less data.

**Grad-CAM and Its Modifications:** We agree that many methods could always be considered. However, we chose to focus on the baseline implementations of Grad-CAM to get an unbiased view of its performance in this context. Our method is not mutually exclusive from other ones. Anyone reading this paper can pick and choose whichever specific Grad-CAM method they want to use and adapt it to our proposed loss function. All you need to do is to update Equation 2 depending on a specific implementation of Grad-CAM. Everything else will be the same.

**Diversity of Data:** In response to your feedback, Appendix C now includes comprehensive details on patient demographics, imaging equipment variations, and dataset distribution, reinforcing the robustness of our model against data shifts.

**Class Imbalance:** We appreciate your pointing out the significance of class imbalance. For all experiments, we provide both AUROC and AUPRC to provide a balanced view of our model's performance, accounting for class distribution. 

**Literature Review and Future Work:** We have updated the literature review and introduced Appendix D with Future Work.
Thank you for your insightful comments and the opportunity to enhance our manuscript. Below we address the key points raised in your review:

**Dependency on Segmentation Masks:** We appreciate your concern regarding the practicality of requiring segmentation masks. However, in specific medical conditions such as Hydronephrosis, creating segmentation masks for kidneys is relatively straightforward and cost-effective compared to the significant effort and resources needed to collect larger datasets. Our approach leverages these masks to train models that can be easily adapted and reused across different hospitals, offering a practical solution in scenarios where data availability is a major constraint. As we show in Table 4 and Table 5, to achieve comparable results you will need to collect 25 times less data.

**Diversity of Data:** The limitation regarding data diversity is acknowledged in our manuscript. While our current datasets are primarily from Canada and the USA, we demonstrated our model's effectiveness (Table 3 and Table 5) across different hospitals and imaging equipment within these regions. We've added additional information about the datasets in Appendix C, highlighting variations in patient demographics and imaging equipment to confirm our claims about Out-of-Distribution performance. We’ve also added Appendix D about future work in which we will focus on extending this validation to datasets from other countries.

**Future Work:** Our initial study was primarily designed to explore and validate the effectiveness of our segmentation-guided attention mechanism in scenarios characterized by small datasets and domain shifts — a niche that, to our understanding, had not been thoroughly investigated previously. Given this focus, and the complexities involved in directly comparing our method with a broad range of state-of-the-art domain generalization methods — each with its unique assumptions, requirements, and implementation details — we chose to concentrate our efforts on demonstrating the significant contributions our approach offers within its intended context. Nevertheless, your point is very important and we plan to focus on answering it in future work.

**Typos:** Thank you for the careful read-through, we have fixed the typos.
We greatly appreciate your insights and the opportunity to clarify aspects of our work. Below, we address each of your concerns:

**Grad-CAM and Novelty:** We want to flag a potential misunderstanding here. We do use Grad-CAM to compute the attention maps but our innovation lies in integrating these maps with a novel loss function. This function penalizes the model for focusing on irrelevant regions, improving accuracy in diagnosing Hydronephrosis, especially in small datasets — a key novelty of our approach.

**Training Dataset Selection:** Our methodology includes a fair comparison through two baseline settings: 1) Models trained on images with segmentation masks (same as our model), and 2) Models trained on a dataset 25 times larger. This dual approach ensures a balanced evaluation of our model's effectiveness cause in the first case all models are trained on the same amount of data and the same images and in the second case only baseline models are trained on a bigger dataset.

**Segmentation Masks Origin:** We want to highlight another potential misunderstanding. The segmentation masks are manually created by medical professionals, underscoring the practical and real-world applicability of our method in leveraging existing clinical annotations for model improvement.

**Ablation Study:** The comparison between ResNet-18 models using Binary Cross Entropy loss versus our Attention loss (Tables 2 and 3) serves as an ablation study. This demonstrates significant improvements with our approach, both in identically distributed data and, notably, in OOD scenarios—highlighting the effectiveness of our loss function.

**Future Work:** Recognizing the value of applying our segmentation-guided approach to other models, we aim to explore this in future research. Limited by computational resources, we couldn't include these experiments within the rebuttal period but have noted this in Appendix D in our revised manuscript.

We believe these clarifications and our manuscript's revisions directly address your concerns, further underscoring our contributions to advancing medical image classification.
Thank you for your time and feedback! We have revised our manuscript to address your concerns as follows:

**Prior Work:** Updated the introduction section so it contains the important prior work instead of the methods section.

**Title Revision:** Thank you for the suggestion! We have updated the title to “Improving Identically Distributed and Out-of-Distribution Medical Image Classification with Segmentation-Guided Attention in Small Dataset Scenarios“ so it more accurately represents the scope of our work.

**Grad-CAM Figures and Datasets Namings:** Updated the figures and standardized dataset names for clarity.

**Why don’t we use multiple domains for training?** We assume that only the source dataset is available during the training. So there will be no need to collect data for each hospital separately but rather just apply the already trained model from another hospital.

We appreciate your insights, which have strengthened our manuscript.