Keypoint-Augmented Self-Supervised Learning for Medical Image Segmentation with Limited Annotation

Summary

This paper tackles self-supervised learning for medical image segmentation. The main contribution is the keypoint identification in input images in self-supervised pre-training to consider long-range spatial dependencies. Each UNet encoder layer is followed by a keypoint-augmented fusion (KAF) layer, where keypoints are identified by SIFT on the input image and rescaled with respect to the resolution of each layer. Convolutional features at each keypoint are processed by a Vision Transformer to incorporate self-attention, provided to additional convolutional layers for extracting dense features from transformer outputs and finally concatenated with the input feature map. Resulting KAF layer features are used for contrastive training of UNet for both global representation learning similar to PCL [54], as well as local similarity learning via the SuperGlue graph neural network [42].

Strengths

Incorporating similarity learning from keypoints is interesting and novel in self-supervised medical image segmentation literature. Paper is well-written and the experiment setup is detailed. Cross-validation is used to compute error bounds in addition to performance metrics. Downstream segmentation performance with limited annotations outperforms several pre-training approaches across two benchmark datasets.

Weaknesses

While contributions are novel, literature review is missing a lot of works on self-supervised medical image segmentation via global/local similarity/contrastive learning: Taher et al. "CAiD: Context-Aware Instance Discrimination for Self-supervised Learning in Medical Imaging" 2022 Yan et al. "SAM: Self-supervised Learning of Pixel-wise Anatomical Embeddings in Radiological Images" 2021 Zheng et al. "MsVRL: Self-Supervised Multiscale Visual Representation Learning via Cross-Level Consistency for Medical Image Segmentation" 2023 Ouyang et al. "Self-supervised Learning for Few-shot Medical Image Segmentation" 2022 Fischer et al. "Self-supervised contrastive learning with random walks for medical image segmentation with limited annotations" 2023 Xie et al. "PGL: Prior-Guided Local Self-supervised Learning for 3D Medical Image Segmentation" 2020 Wang et al. "Self-supervised learning based transformer and convolution hybrid network for one-shot organ segmentation" 2022: also incorporates Vision Transformer for global feature learning. Experimental comparisons against Wang et al. 2022 would further strengthen the novelty.

Questions

How are the thresholds selected for identifying positive/negative keypoints?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

limitations and future work are discussed.

Summary

This paper focuses on improving medical image segmentation performance with limited labeled data. They (1) propose a modification to the standard UNet convolutional architecture by incorporating local features (at locations identified by SIFT) into the backbone, (2) show how their modification can be used to control self-supervised contrastive learning positive and negative pair selection, and (3) evaluate their proposed architectural change and self-supervised training strategy against relevant baselines. They find both the architectural change and self-supervised strategy improve segmentation performance.

Strengths

- Novel architectural change that impacts network performance both with random initializations and with self-supervised contrastive learning. - Improves performance against relevant and recent self-supervised baselines.

Weaknesses

- The GLCL baseline reported in Table 1 is much lower than the numbers reported in the GLCL original paper; the original GLCL paper actually reports performance better than the proposed method on the ACDC dataset. Please comment on why the reported GLCL performance in 14-18 points lower than the GLCL paper. - The proposed architectural change may increase compute requirements, but this is not discussed nor are any timing/computational load experiments reported. - There is limited dataset variety—both evaluation datasets are 3D cardiac anatomical segmentation tasks. Additional experiments, particularly including pathological segmentation, would strengthen the results and show the proposed method works for different segmentation targets. - Scaling the proposed solution to non-CNN architectures or 3d data may be difficult (as mentioned by the authors). - Minor: Table 1 caption needs to specify what metric you are reporting; Line 221 has a period where there should be a comma+”and.”

Questions

My questions and suggestions are covered in the Weaknesses section.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Limitations and impact adequately discussed.

Summary

SIFT keypoints are used in conjunction with a UNET to construct a transformer-like model for medical image segmentation. Global and local self supervised learning (SSL) losses are proposed. Correspondences identified using the SuperGlue method. Various experiments on two cardiac image datasets, including ablation studies, shows improves on other work in terms of segmentation dice score.

Strengths

The use of traditional keypoints as attention operator mechanism for UNet segmentation makes sense and is novel, to my knowledge. Improved segmentation results convince the reader that the approach has practical value.

Weaknesses

The method presented is based on 2D keypoint matching, issues with multiple 2D slices, wheras anatomy is 3D. It would be interesting to know if 3D keypoints would improve results. Table 1 should mention that the numbers shown are dice scores. Minor spelling Line 182 “based on the cloest L2 distance” -> “closest” Line 213 “patent ages ranging from 1 month to 21 year” -> “21 years” It would be interesting to know much the learning method here improves upon simple keypoint transfer segmentation as in [46]. Objective functions in equations (1) and (2) appear very similar to other recent keypoint matching work: Chauvin, Laurent, et al. "Efficient pairwise neuroimage analysis using the soft jaccard index and 3d keypoint sets." IEEE transactions on medical imaging 41.4 (2021): 836-845. The description of the datasets used does not mention the size and resolution of the images. Computational complexity of the proposed method is not mentioned. The paper lacks reproducibility. For instance, line 156 and line 181 allude to thresholds whose numerical values are not explicitly mentioned. Similarly, it is not clear what setup was used to train the Transformer block, including details such as learning rate, batch size, attention head size, and input embedding size. Additionally, it would be valuable to include statistical tests to demonstrate the superiority of the proposed method compared to existing approaches.

Questions

- It would be good to hear a discussion regarding 3D keypoint methods, and how much this method would improve upon simple keypoint transfer segmentation as in [46]. - In lines 208 and 209 it is stated that images are pre-processed using the methodology described in [54]. Could the preprocessing stage be briefly described in the paper? Are there aspects without which this approach would not work, e.g. registration of training data? - Could the authors explain why the Manhattan distance is preferred over other distance measures for keypoint filtering (as mentioned in line 181)? - Would it be possible for the authors to provide the values of the similarity index in Figure 2? This would allow us to compare those values with the ones presented in the previous Table 1. - Can the authors provide more insight into the justification for the following claim? Clearly there are various random transforms which would scramble image information, to which SIFT is not invariant. "(2) our proposed self-supervision helps maintain better local equivariance of the self-attention, i.e., with the same query point location, its interaction with other points remains identical no matter how the image is transformed."

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Several limitations exist: missing important details, computational complexity, description of data, transformer parameters (learning rate, batch size, attention head size, and input embedding size), statistical significance and comparison to based keypoint transfer segmentation.

Summary

The authors introduce an approach to improve the performance of the UNet model, specifically in the task of segmentation. They propose a keypoint augmented fusion layer to enhance the CNN-based encoder of UNet. This layer leverages self-attention to capture long-range spatial dependencies among keypoints extracted from multi-scale feature maps. In addition, the paper suggests the use of global and local contrastive losses to pretrain the model in a self-supervised manner. These losses help the model learn meaningful representations from unlabeled data, improving its performance. The proposed framework was evaluated on Cardiac CT and MRI datasets for segmentation tasks. The results demonstrate reasonable improvements over the state-of-the-art methods, particularly in scenarios with limited labeled data (few-shot settings).

Strengths

1) This work addresses limitations of traditional UNet models, which rely on convolutional neural networks (CNNs) and may struggle to capture long-range dependencies among spatial positions in an image. In this study, the authors propose a more effective approach to handle this issue. At each scale of the encoder, the authors employ Scale-Invariant Feature Transform (SIFT) to detect key points within the CNN feature maps. By applying SIFT, the authors aim to identify salient locations within the feature maps that can capture important spatial information. To facilitate correlation and interaction among these detected key points, a transformer model is employed. 2) Global contrastive loss is applied in an unique way on two types of representations - concatenated keypoint enhanced feature maps, as well as features from last layer of UNet. 3) instead of being used as a traditional image-wise similarity loss, here the local contrastive loss is proposed to exploit pixel-wise similarity (by measuring the correlation between keypoints within a spatial distance among adjacent slices)

Weaknesses

1) It is unclear why this work ignores Transformer-based UNet architectures (TransBTS [1], UNetr[2], Swin-Unet [3], etc) which employ self-attention to better capture the correlations. Though the authors propose to use it on keypoints, I find it difficult to comprehend the advantages here, except maybe computational efficiency. There are many works which do it on entire feature maps - why are those suboptimal? 2) Handcrafted features like SIFT/SURF to detect keypoints, though powerful, aren't the current SOTA. The performance can definitely be tested using finetuned object detection models (such as Faster RCNN [4]) In short - there should be ablation showing (1) advantage of using keypoints, (2) Different keypoint detection strategies, and most importantly (3) comparison with more recent Transformer-based UNet architectures. [1] TransBTS: Multimodal Brain Tumor Segmentation Using Transformer [2] UNETR: Transformers for 3D Medical Image Segmentation [3] Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation [4] Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks

Questions

1) Keypoint Robustness - The manner of application of global contrastive loss depends on the effectiveness of keypoints. Because the other global loss is used off-the-shelf from [54] as PCL loss. 2) As for local loss, the contribution appears to be more engineering (lines 175-200) i.e on how to use it among keypoints in adjacent slices. Here also, there is no ablation on the multiple thresholds selected in the form of spatial distance between keypoints or number of adjacent slices. 3) Standard non-cardiac benchmark datasets like BTCV, SegTHOR have not been used. To establish it as a generic method, the authors should definitely explore non-cardiac pan-organ datasets 4) I also fail to understand what customization the authors have proposed in their architecture to make it work in few-shot settings with limited annotations. 5) Can the authors please mention the unlabeled datasets used for SSL pretraining?

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

Not adequately addressed in the 'Discussion' section

Summary

The paper presents a self-supervised learning approach for medical image segmentation that uses matching keypoints to contrast local features that may be spatially distant from one another. The proposed model pretrains a segmentation network on unlabelled images, using three contrastive losses: a global loss on the keypoint-enhanced features, a global loss on pooled features in the last layer (as in PCL) and a local loss on features of matching keypoints (which are spatially close and have a high matching probability). This model exploits Keypoint-Augmented Fusion (KAF) layers which transforms the features of detected keypoints using self-attention transformer layers and then concatenates the transformed features to the original feature map (after constructing a sparse map that is diffused to a dense map via CNN layers). The SuperGlue network of Sarlin et al. is employed to compute the keypoints and their matching scores. The proposed approach is evaluated on two cardiac MRI segmentation datasets (ACDC and CHD) for which it shows superior performance compared to recent self-supervised approaches.

Strengths

* The idea of using keypoints in a local contrastive loss for self-supervised pre-training is novel (to my knowledge) and interesting. Although it relies on the accurate detection and matching of keypoints, it has the potential to overcome the limitations of current strategies, for instance, based on spatial proximity and transformation consistency (which can not be used across different images). * The method is sound and well presented. * Experiments are detailed and results show clear improvements over recent self-supervised approaches for segmentation in few shot settings. Likewise, ablation studies and visualization experiments are nicely done and showcase the proposed method. * The paper is quite well written.

Weaknesses

* While the proposed method is compared against recent self-supervised approaches, only one of these approaches (GLCL) tries to learn a local representation, which makes the comparison a bit unfair. Yet, several approaches have been proposed for this purpose, for instance, see references below. As I understand, these approaches can also contrast local features that are semantically related (not necessarily spatially related), for example, based on clustering as in Peng et al, pseudo-labels as in Zhong et al., or super-pixels as in Chaitanya et al. Wang, Zhaoqing, Qiang Li, Guoxin Zhang, Pengfei Wan, Wen Zheng, Nannan Wang, Mingming Gong, and Tongliang Liu. "Exploring set similarity for dense self-supervised representation learning." In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16590-16599. 2022. Peng, Jizong, Marco Pedersoli, and Christian Desrosiers. "Boosting Semi-supervised Image Segmentation with Global and Local Mutual Information Regularization." Machine Learning for Biomedical Imaging 1, no. MIDL 2020 special issue (2021): 1-29. Zhong, Yuanyi, Bodi Yuan, Hong Wu, Zhiqiang Yuan, Jian Peng, and Yu-Xiong Wang. "Pixel contrastive-consistent semi-supervised semantic segmentation." In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 7273-7282. 2021. Chaitanya, Krishna, Ertunc Erdil, Neerav Karani, and Ender Konukoglu. "Local contrastive loss with pseudo-label based self-training for semi-supervised medical image segmentation." Medical Image Analysis 87 (2023): 102792. * Some of the reported results are not consistent with the literature. Namely, Table 1 reports a Dice of 0.642 for GLCL-full with M=2 in ACDC whereas the original paper reports a Dice of 0.789, which is higher than results of the proposed method. * A potential weakness of the proposed method is that it requires to find matching local features across different images. Typically, descriptors like SIFT focus on points corresponding to sharp edges or blobs, which are present in cardiac structures. I suspect such points would be harder to find in structures with lower contrast such as prostate in MRI. Thus, it would be necessary to test the method on a broader range of organs and even other modalities like CT.

Questions

* What is the threshold used to match keypoints?How sensitive is the method to the number of matching keypoints? * Why are the results reported for GLCL-full lower than those of the original paper? * Why not test the method on other segmentation tasks? For example, the GLCL paper by Chaitanya et al. also test on prostate segmentation. * In the ablation study (Table 2), why not test the case where w1=w2=0? This would be useful since the main novelty of the method lies in L_local. * In eq (2), does the index N+1 correspond to a no-match class? Other comments: * p5: "based on the cloest L2 distance" --> closest

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

Some limitations are mentioned in the discussion.

Thanks for the detailed answers

I have read the authors' answers carefully and I am generally satisfied with them. Regarding the first point (Comparison with other local SSL methods), I think the paper by Chaitanya et al. also evaluates their method in a self-supervised pre-training setup. Moreover, while it is true that the method Peng et al. is tested in a semi-supervised setup, it was extended to self-supervised learning in a follow up work: Peng, Jizong, Ping Wang, Marco Pedersoli, and Christian Desrosiers. "Boundary-aware information maximization for self-supervised medical image segmentation." arXiv preprint arXiv:2202.02371 (2022).

Rebuttal response

I have read the other reviews and the authors' rebuttal. The authors have presented a thorough rebuttal responding to the raised questions/weaknesses and have presented new data showing useful extensions and ablations of their method. I think this is an interesting addition to the self-supervised body of work and am changing my score from 4 to 5.