Title: Integrating Deep Metric Learning with Coreset for Active Learning in 3D Segmentation
Venue: NeurIPS 2024
Accepted: yes
Keywords: Active Learning

Medical Imaging

Segmentation

Deep Metric Learning
Abstract: Deep learning has seen remarkable advancements in machine learning, yet it often demands extensive annotated data. Tasks like 3D semantic segmentation impose a substantial annotation burden, especially in domains like medicine, where expert annotations drive up the cost. Active learning (AL) holds great potential to alleviate this annotation burden in 3D medical segmentation. The majority of existing AL methods, however, are not tailored to the medical domain. While weakly-supervised methods have been explored to reduce annotation burden, the fusion of AL with weak supervision remains unexplored, despite its potential to significantly reduce annotation costs. Additionally, there is little focus on slice-based AL for 3D segmentation, which can also significantly reduce costs in comparison to conventional volume-based AL. This paper introduces a novel metric learning method for Coreset to perform slice-based active learning in 3D medical segmentation. By merging contrastive learning with inherent data groupings in medical imaging, we learn a metric that emphasizes the relevant differences in samples for training 3D medical segmentation models. We perform comprehensive evaluations using both weak and full annotations across four datasets (medical and non-medical). Our findings demonstrate that our approach surpasses existing active learning techniques on both weak and full annotations and obtains superior performance with low-annotation budgets which is crucial in medical imaging. Source code for this project is available in the supplementary materials and on GitHub: https://github.com/arvindmvepa/al-seg.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper is concerned with the task of 3D segmentation in medical images. Motivated by the high cost of annotations by radiologists, the authors propose a novel active learning method based on the coresets. Here, instead of using the euclidean metric for defining the similarity between different samples, the authors propose to learn a metric via a contrastive loss. The authors propose a group contrastive loss based on the NT-Xent loss which leverages intrinsic similarities in the data (different patients, different volumes of the same patient, adjacent slices) and evaluate their method both on weakly-supervised (with scribble annotations) and fully-supervised tasks and using three datasets from the medial domain and one video segmentation dataset.\n\nSTRENGTHS:\nThe paper is well written and methods to improve medical image segmentation with limited annotations are of high interest to the community. The related work is comprehensive and the authors motivate their work well. Overall methods are sound and the evaluation is enough to support the claims.\n\nWEAKNESSES:\n**3D vs. 2D segmentation models**\n- To the best of my knowledge 3D models (i.e., *not* slice-based models) generally outperform 2D models in most medical image segmentation tasks [a,b]. While I understand the authors' motivation for using 2D (i.e., slice-based) methods, I find their statements (e.g., L. 36-38, 82-83)\n    > There were also inefficiencies in previous studies, such as choosing whole volumes instead of individual slices in 3D segmentation, which increased costs [37]. (L. 82-83)\n    \n    surrounding this choice somewhat misleading. I would appreciate a more honest discussion where the authors phrase their focus on slice-based methods less as an advantage and more as a limitation of the present work. They could also briefly discuss potential extensions of their work to 3D segmentation. I suppose that generally, this is possible when the adjacent slice group contrastive loss is removed (potentially one could still consider adjacent 3D patches or something similar). The ablation studies also suggest that the adjacent slice group loss doesn't really help much and the authors did not include it in their final loss formulation (L. 582).\n\n- The restriction to slice-based methods and in particular the slice group contrastive loss imply that the proposed methods would assume the annotations acquired during the active learning to also be slice-based. However, for medical image segmentation it might be much easier to acquire annotations for $N$ slices within the same volume than it is to acquire annotations for $N$ slices across different volumes because readers would usually use some interactive segmentation tool and segmentations are similar for adjacent slices. \n\n**Ablation studies**\n\nIn Tab. 3 the authors present an ablation study where they study the effect of different group losses on the overall performance. I find this study very interesting, however:\n- It seems that ablation studies where only conducted on a single dataset. If so, on which dataset? \n- How comparative are the results, given that including more contrastive losses lead to larger epochs and more update steps by design of the batch sampler? Can the authors clarify?\n- How did the authors choose the weighting of each loss term here?\n\n**Typos and minor comments:**\n\n- \"NT_Xent loss\" -> \"NT-Xent loss\"\n- L. 246: \"Training speed for for the [...]\"\n- Fig. 2 could be improved by just adding a column ground truth segmentation and marking the segmentation by the model in a single color. The blue and red colors can be somewhat confusing. Also, which dataset is this from? I suspect the ACDC dataset? Then, shouldn't there be three classes (RV, LV, Myocardium)?\n\n---\n[a] Avesta A, Hossain S, Lin M, Aboian M, Krumholz HM, Aneja S. Comparing 3D, 2.5D, and 2D Approaches to Brain Image Auto-Segmentation. Bioengineering. 2023; 10(2):181. https://doi.org/10.3390/bioengineering10020181\n\n[b] Wasserthal, J., Breit, H. C., Meyer, M. T., Pradella, M., Hinck, D., Sauter, A. W., Heye, T., Boll, D. T., Cyriac, J., Yang, S., Bach, M., & Segeroth, M. (2023). TotalSegmentator: Robust Segmentation of 104 Anatomic Structures in CT Images. Radiology. Artificial intelligence, 5(5), e230024. https://doi.org/10.1148/ryai.230024\n\nQUESTIONS:\n- In  L. 582 the authors write that they did not include the slice group loss in the optimal formulation of their loss but in L. 217-220 it seems they included this term in their experiments for all datasets (also if the slice group loss was not included in training of CHAOS, MS-CMR, DAVIS) this would mean that the authors only used the volume group contrastive loss for all these experiments? Can the authors please clarify this in the manuscript?\n- How were the individual group losses weighted?\n- In L. 224 the authors mention that segmentation requires more training data for natural images. Could you please provide more details on this or a reference?\n- Can you better explain the different groups you have, in particular the patient groups? I suspect that these are different pathology subgroups + healthy group (e.g. in the ACDC dataset we would have 5 patient groups)\n- On which dataset is the ablation study in Tab. 3 performed?\n- Were reference methods trained for the same number of update steps and on the same amount of training data as the proposed method?\n- For the random baseline, did the authors choose samples randomly or weigh samples on a per-volume/per-patient level? Otherwise patients for which more volumes/slices are acquired would overrepresent in the coreset, right?\n\nLIMITATIONS:\nThe authors discuss limitations of their method in Sec. 5. Maybe the authors can also discuss the possibility of including other groups (e.g. pathologies, patient age/sex/..., ...), in the future? I would also appreciate a more honest discussion of the limitations and implications of slice-based segmentation (see **3D vs. 2D segmentation models** above)", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a novel active learning approach based on group-based contrastive learning, aiming to optimize data selection and model efficiency in medical image segmentation tasks. By combining NT_Xent loss with various group contrastive losses, it addresses the Coreset problem in active learning, particularly in handling the diversity and complexity of medical image data. The paper provides detailed descriptions of the method's construction and training process, validated through experiments on multiple classical and novel datasets to demonstrate the model's superiority and robustness under different annotation budgets and data complexities. The research findings hold significant academic and practical significance for enhancing workflows in medical image analysis and have profound implications for advancing machine learning applications in medicine.\n\nSTRENGTHS:\nFirstly, by introducing a group-based contrastive learning approach, the paper significantly enhances model performance in medical segmentation tasks. Secondly, it innovatively combines NT_Xent loss with various group contrastive losses, effectively addressing the Coreset problem in active learning and tackling challenges posed by data diversity and complexity in medical image segmentation. The methodology and experimental design are of high quality, providing clear descriptions of contrastive learning loss applications and their integration within the active learning framework. Experimental results demonstrate the model's strong performance under varying annotation budgets and data complexities, highlighting its relevance for medical image analysis and its potential value to the NeurIPS community.\n\nWEAKNESSES:\n1. The description is confusing. For instance, the first chapter mentions two main contributions, but the abstract and methods sections also discuss the loss function, which is not explained in the contributions of the first chapter.\n\n2. The paper's description mainly focuses on solving the task of medical image annotation, yet it also includes experiments on natural data. I believe the use of video data does not align well with the primarily medical theme described in the paper.\n\n3. Even if the method is shown to be extendable to natural data, using only one natural dataset is insufficient. To enhance its applicability and impact across broader domains within the NeurIPS community, I suggest using more varied types of data (e.g., natural images, video sequences).\n\n4. Writing details. There are inconsistencies in figure captions, with some having periods and others not. The numbers in Table 1 exceed the table lines. Some tables comparing with other methods lack proper citations. The colors in the figures are too varied; please use black, white, and gray tones to highlight professionalism.\n\n5. The explanation of the annotated data volume for weak supervision and full supervision is unclear.\n\n6. Although the text mentions experiments conducted under low annotation budgets (2%-5%) and fully annotated scenarios, comparing the performance of different methods, it does not clearly explain why a 2%-5% annotation data volume is used in the fully annotated scenarios.\n\n7. In practical experimental design, the reason for using a 2%-5% annotation data volume under full supervision is usually to compare the performance of weak supervision and full supervision methods on the same data volume basis. However, weak supervision and full supervision are different concepts, and using a partial annotation data volume for both does not align with the concept of full supervision.\n\nQUESTIONS:\n1. The main point of the authors' discussion is the time-consuming and labor-intensive nature of medical annotation. However, they added a dataset in a natural scenario, which seems somewhat off-topic. Is there a reason for this?\n\n2. Why are fully supervised and semi-supervised methods compared using different annotation increments? Fully supervised data is entirely annotated. How does this relate to the primary focus of the paper on semi-supervised scenarios?\n\n3. The paper introduces different loss functions, but what are the weight values for these different loss functions under different experimental settings? The paper does not clarify this point.\n\n4. The ACDC dataset includes both weakly supervised and fully supervised experiments. Why does the MS-CMR dataset only have weakly supervised experiments, while the CHAOS and DAVIS datasets have fully supervised experiments?\n\nLIMITATIONS:\n1. Annotation costs and efficiency: Although the paper discusses the potential of active learning to reduce annotation costs, it would be beneficial to explicitly discuss the trade-offs between annotation quality and the computational expense of implementing the proposed methodology. This discussion could provide insights into practical considerations for real-world deployment.\n\n2. Section 4.5 ablation study focuses solely on the combination of loss functions. Consider other ablation experiments, such as the scribble-weighted distance.\n\n3. Due to limited computational resources, the researchers selected a single architecture for each type of annotation and limited the number of experimental runs and model training. This may result in an incomplete evaluation of other architectures and models' performance.\n\n4. Video segmentation models typically require extensive pre-training to achieve good metrics. However, due to resource constraints, the researchers did not conduct such experiments, potentially leading to suboptimal performance in video segmentation tasks.\n\n5. The paper does not adequately address the issue of training set bias, which may result in unfair outcomes for minority groups. Future work should consider this issue more comprehensively and focus on the method's applicability across different domains.\n\n6. The proposed method mainly targets 3D segmentation tasks in the medical field. Further validation and adjustments may be necessary for applications in other fields.\n\n7. Was the selection of annotated data random? Did it consider data distribution bias? This could result in unfair outcomes for minority groups.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a new metric learning approach to improve the Coreset method (an optimization framework to seek the most diverse samples for model learning in active learning) with applications for slice-based active learning in 3D medical segmentation. The main idea is to utilize data groups extracted from 2D slides in 3D volumes and build group-based contrastive learning to enhance feature representations. The authors conduct experiments on four datasets and compare them with several baselines in active learning. The experiment shows good overall performance.\n\nSTRENGTHS:\nOne of the main strengths of this paper is that the authors compared diverse baselines and benchmarks with some of the latest methods, for e.g., [19,37]. Also, their method achieves good performance with significant improvements on four datasets. The author also provides ablation studies to analyze the contributions of key components in Table 3.\n\nWEAKNESSES:\nThere are two main critical concerns of Reviewers:\n\n(a) The way the author improves the Coreset with deep metric learning using group-based contrastive learning is limited in its novelty. In particular, group-based contrastive learning has already been introduced in prior work [65], and several advanced versions have been introduced. Although the authors modify with two sums in Equation (4), this is a minor point.  In short, given the claim that deep metric learning is the core contribution, authors should further tailor contrastive learning for 3D medical tasks, for e.g., making it robust under domain shift or considering other factors like hierarchical relations [1,2]. Also, fusing different metric learning using contrastive or generative learning is also a promising direction. \n\n(b) Although experiments are promising, they are tested only with a single architecture, which is not enough to conclude whether the method is generalized to different architectures. As can be seen in Table 2 (DICE scores for DAVIS), model accuracy is still under an acceptable threshold; this suggests that further evaluation of current foundation models [3] (using SAM) or advanced pre-trained medical models [4] (has versions for both ResNet and ViT) is necessary to boost results further. In summary, it's important to provide further experiments with another architecture (preferred large-pre-trained models) to validate the robustness of the proposed active learning. Besides 3D volumes, considering one or two test cases with video is also interesting, where the group of similar 2D slices becomes a group of similar frames. \n\n[1] Taleb, Aiham, et al. \"3d self-supervised methods for medical imaging\", NeurIPS 2020\n\n[2] Zheng, Hao, et al. \"Hierarchical self-supervised learning for medical image segmentation based on multi-domain data aggregation\" MICCAI 2021\n\n[3] Ma, Jun, et al. \"Segment anything in medical images.\" Nature Communications 2024\n\n[4] MH Nguyen, Duy, et al. \"Lvm-med: Learning large-scale self-supervised vision models for medical imaging via second-order graph matching\", NeurIPS 2023\n\n[65] Sohn, Kihyuk. \"Improved deep metric learning with multi-class n-pair loss objective.\" NeurIPS 2016\n\nQUESTIONS:\nMy questions include:\n\n(a) can the author explain the 'group' definition in Section 3.1? It is confusing whether this 'group' refers to a set of similar 2D slides inside a 3D volume of one patient or it refers to a set of 2D slides of different patients. Reviews did not understand what the author wrote in lines 123-128.\n\n(b) how the algorithm 1 be extended if each patient has different data modalities like CT & MRI?\n\nLIMITATIONS:\nAs mentioned in the weakness section, this paper needs to make a major improvement by either adding more experiments (with another architecture) or further improving the novelty in the way it formulates deep metric learning with contrastive ideas. The writing also need to be improved, especially when explaining the 'group' concept to avoid confusing readers.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents an active learning framework based on deep metric learning for 3D medical image segmentation. The proposed approach extends the CoreSet algorithm for active learning by computing distances on features learned via a contrastive loss. This contrastive loss exploits low-cost labels to group samples (2D slices) based on their spatial position and corresponding subject. Additionally, a weakly-supervised version of the method, based on scribbles, is presented. In this version, scribbles are used in an auxiliary distance term measuring the intensity difference for labelled pixels. Experiments are conducted on three medical imaging dataset and one video segmentation dataset to evaluate the method's performance. Results shows the method to outperform recent active learning approaches and baselines in the majority of test cases.\n\nSTRENGTHS:\n* The paper introduces a simple yet efficient way to incorporate contrastive learning in the CoreSet algorithm for active learning in segmentation.\n\n* The extension of active learning to a weakly-supervised setting is also interesting and could have useful applications.\n\n* Experiments are comprehensive. The proposed method is tested on four datasets and compared against several related approaches. Results show clear improvements in most cases.\n\n* The paper is well written, and the method can be understood easily.\n\nWEAKNESSES:\n* Despite its novelty, the board idea of the proposed method is similar to previous work using contrastive learning for core set selection, for instance:\n\nJu J, Jung H, Oh Y, Kim J. Extending contrastive learning to unsupervised coreset selection. IEEE Access. 2022 Jan 13;10:7704-15.\n\nJin Q, Yuan M, Qiao Q, Song Z. One-shot active learning for image segmentation via contrastive learning and diversity-based sampling. Knowledge-Based Systems. 2022 Apr 6;241:108278.\n\nWhile these previous approach do not use group labels as the proposed method, a similar strategy was employed for self-supervised segmentation of medical images:\n\nChaitanya K, Erdil E, Karani N, Konukoglu E. Contrastive learning of global and local features for medical image segmentation with limited annotations. Advances in neural information processing systems. 2020;33:12546-58.\n\n* The comparison against other approaches might be unfair since the contrastive (pre) training, which is used to select the core set, also boosts the performance of the segmentation model (regardless of which labelled samples are used for training). This can be seen in the work of Chaitanya et al (see reference above), where a segmentation trained with the same contrastive loss achieves a performance near full supervision with very few samples.\n\n* The scribbled weighted distance presented in Section 3.2 lacks motivation. It is unclear to me how the pixel-wise difference in intensity gives useful information for non-registered images (potentially of different subjects).\n\nQUESTIONS:\n* Can authors clarify the contribution of their work with respect to previous work on active learning based on a contrastive technique?\n\n * What is the meaning/definition of h_{S0} in Eq (3) ?\n\n* Do you assume that pixels x^i_1 and x^i_2 correspond to the same anatomical structure in the scribble-weighted distance (i.e., images are registered)? Why compute the distance on pixel intensities?\n\n* Why is the method worse than random selection on the CHAOS dataset (2% fully-supervised) ?\n\n* Can you compare your method against random/maxEntropy selection on a segmentation model pre-trained using the SAME contrastive loss as your model?\n\nLIMITATIONS:\n* The last section of the paper mentions relevant limitations of the work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}