Dense Self-Supervised Learning for Medical Image Segmentation

Authors: maxime  seince
loïc a le folgoc
luiz augusto facury de souza
elsa  angelini
l  seince
e  facury de souza
  angelini
Venue: MIDL_io_2024_Conference
Published: 2024-01-26 15:13:49

Abstract: Deep learning has revolutionized medical image segmentation, but it relies heavily on high-quality annotations. The time, cost and expertise required to label images at the pixel-level for each new task has slowed down widespread adoption of the paradigm. We propose Pix2Rep, a self-supervised learning (SSL) approach for few-shot segmentation, that reduces the manual annotation burden by learning powerful pixel-level representations directly from unlabeled images. Pix2Rep is a novel pixel-level loss and pre-training paradigm for contrastive SSL on whole images. It is applied to generic encoder-decoder deep learning backbones (e.g., U-Net). Whereas most SSL methods enforce invariance of the learned image-level representations under intensity and spatial image augmentations, Pix2Rep enforces equivariance of the pixel-level representations. We demonstrate the framework on a task of cardiac MRI segmentation. Results show improved performance compared to existing semi- and self-supervised approaches; and a 5-fold reduction in the annotation burden for equivalent performance versus a fully supervised U-Net baseline. This includes a 30% (resp. 31%) DICE improvement for one-shot segmentation under linear-probing (resp. fine-tuning). Finally, we also integrate the novel Pix2Rep concept with the Barlow Twins non-contrastive SSL, which leads to even better segmentation performance.

Main paper: 
 introduction: 
 Medical image segmentation has seen tremendous progress with the advent of deep learning (Ronneberger et al., 2015; Milletari et al., 2016; Kamnitsas et al., 2017) . The drawback of this paradigm is its reliance on large quantities of data, annotated at the pixel-level, to train strong segmentation models. These pixel-level annotations are costly to obtain, and take precious time from medical experts.To circumvent this burden, techniques have emerged in recent years that better exploit more widely available unlabeled data. Semi-supervised approaches e.g., pseudo-labels (Lee, 2013; Bai et al., 2017; Tran et al., 2022) and mean teacher (Tarvainen and Valpola, 2017; Yu et al., 2019) , balance a supervised segmentation loss on a small labeled dataset with a consistency loss on the larger unlabeled dataset, yielding improved segmentation. Other semi-supervised approaches include Bayesian deep learning e.g., Dalca et al. (2018) introduce anatomical priors that can be learnt using unlabeled or unpaired data.Self-Supervised Learning (SSL) follows an alternative route whereby deep representations are directly learned from unlabeled data. Early methods trained these representations by solving pretext tasks e.g., relative position prediction (Doersch et al., 2015) , image recolorization (Zhang et al., 2016) , jigsaw puzzles (Noroozi and Favaro, 2016) or rotation prediction (Gidaris et al., 2018) . These methods are designed for image classification as a primary downstream task, thus an image is encoded to an image-level vector representation. Many recent methods for image-level representation learning coexist in the state-of-the-art, based on contrastive learning (Chen et al., 2020; He et al., 2020) , on redundancy-reduction (Zbontar et al., 2021) , on self-distillation (Grill et al., 2020; Caron et al., 2021) , on Masked Image Modeling (He et al., 2022 ) and many more.We propose instead a framework for pixel-level (dense) representation learning, dubbed Pix2Rep, which can be used to pretrain encoder-decoder architectures, such as U-Nets. Whereas most aforementioned methods rely on invariance under certain intensity-based augmentations (brightness & contrast, Gaussian noise, etc.) and geometric augmentations (crops), Pix2Rep is based on equi variance under geometric transformations. For the task of cardiac MRI segmentation, we propose rotations and intensity reversals as additional augmentations that further improve results. Finally, we investigate the performance of Pix2Rep in various data regimes (one-shot, few-shot segmentation, or large annotated data), both under linear probing and fine-tuning. 
related work: 
 Comparatively, fewer pixel-level representation learning methods have been proposed so far. Kalapos and Gyires-Tóth (2023); Punn and Agarwal (2022) propose to pretrain a U-Net encoder (a.k.a. its downsampling branch) using image-level SSL (BYOL/Barlow Twins). The U-Net decoder however is randomly initialized before fine-tuning on the downstream segmentation task. Tang et al. (2022) pretrain a Swin UNETR encoder using a combination of image-level contrastive learning, pretext task and masked image modeling. Zeng et al. (2021) exploit the positional information of slices within stacks for contrastive pretraining of a U-Net encoder. Chaitanya et al. (2020) manage pretraining of the first decoder layers by introducing a local contrastive loss that relies on rough alignment of subject volumes. For contrastive pretraining of the whole decoder (Xie et al., 2021) , positive pairs of pixels need to be defined. Zhong et al. (2021) constrain the two augmented views to differ only up to intensity transformations, so as to form positive pairs from pixels at identical locations in the two views. Hu et al. (2021) ; Zhao et al. (2021) regard as positive all pixels sharing the same label, at the cost of pretraining only on the (potentially smaller amount of) labeled data. Wang et al. (2021) ; Bardes et al. (2022) define pixels with highly similar features as positives. The closest related work to our proposed approach are those of O. Pinheiro et al. (2020) ; Yan et al. (2022) ; Goncharov et al. (2023) , which define positive pairs to be pixels that describe the same physical location in the scene on different augmented patches, that differ up to intensity-based and geometric augmentations. Our framework targets equivariance rather than equivalence to random spatial transforms and works at whole image level for augmentations. O. Pinheiro et al. (2020) experiment on natural -not medical-images. Goncharov et al. (2023) ; Yan et al. (2022) focus less on few-shot segmentation. x an unlabeled training image; ϕ ∼ T s a random spatial transformation; t, t ′ ∼ T i two random intensity transformations; g a projection head. We train pixel representation maps output by f to be equivariant under ϕ and invariant to t, t ′ by maximizing agreement between the outputs of the two branches, via a pixel-level contrastive loss. 
methods: 
 We consider an arbitrary (trainable) neural network backbone f : R H×W ×C → R H×W ×D that maps input images to pixel representation maps. f can be any encoder-decoder network and we opt for simplicity for a U-Net (Ronneberger et al., 2015) , without its final segmentation head (≡ typically 1 × 1 conv + softmax)1 . f can be interpreted as embedding C-dimensional input pixels to D-dimensional vector representations, accounting for local and global context. At the core of the proposed self-supervised pretraining of f is a contrastive loss which forces the pixel representations derived from f to be invariant under the action of intensity augmentations and equivariant under the action of spatial transformations. The method is summarized in Fig. 1 .Pix2Rep pretraining relies exclusively on an unlabeled dataset D ≜ {x ∈ R H×W ×C }. Apart from the encoder-decoder f (•) which extracts pixel-level representation maps, the framework consists of three major components.(1) For any input image x, a stochastic data augmentation module generates two random intensity-based transformations t, t ′ ∼ T i (incl. brightness & contrast augmentation, Gaussian noise, bias field, and intensity reversal), resulting in two views v ≜ t(x), v ′ ≜ t ′ (x). In addition, the stochastic data augmentation module generates a single random spatial transformation ϕ ∼ T s (incl. zooms, flips, rotations), which is applied asymmetrically to v, v ′ .(2) A small projection head g : R H×W ×D → R H×W ×d transfers pixel representation maps to the space where the contrastive loss is applied. The projection head g(•) ≜ W (2) * σ(W (1) * •) consists in a MLP with one hidden layer (where σ is a ReLu non-linearity), implemented using 1 × 1 convolutions. We have experimented with deeper projection heads (up to 3 hidden layers) without noticing significant benefits. Then:• For the view v ′ : we first transform it to the new spatial viewpoint by applying ϕ,yielding 2 ϕ • v ′ ≜ v ′ • ϕ -1 ∈ R H×W ×C , then pass the spatially-transformed image through g • f , yielding z ′ ≜ (g • f )(ϕ • v ′ ).• For the view v: we first apply g•f , yielding the pixel representation map z = (g•f )(v), before transporting z to the new viewpoint:ϕ • z = ϕ • (g • f )(v).Note that ϕ is applied exactly once along the two computational branches, hence ϕ• z = ϕ • (g • f )(t(x)) and z ′ = (g • f )(ϕ • t ′ (x))) share the same viewpoint.(3) Thirdly, a contrastive loss L is defined for a pixel-level contrastive prediction task. Given any pixel coordinate s ∈ R 2 , features (ϕ • z)(s) ∈ R d and z ′ (s) ∈ R d correspond to the same anatomical point in the two views, thus they form natural positive pairs. Natural negative pairs are obtained from samples from the same view at all other pixel coordinates, or from other images in the minibatch at any pixel coordinate. However, creating one positive pair per pixel coordinate would yield hundreds of thousands of positive pairs and potentially billions of negative pairs. Instead, we sample M random pixel coordinates {s (Chen et al., 2020) , we use the InfoNCE loss l i,j of Eq. (1) for each positive pair (u i , u j ):EQUATIONwhere τ denotes a temperature parameter, and sim(u, u ′ ) ≜ u T u ′ /∥u∥ 2 ∥u ′ ∥ 2 denotes the cosine similarity. The total loss is aggregated by summing over all l i,j , including symmetrizing the roles of i, j. Assuming without loss of generality that positive pairs have consecutive indices i = 2k -1 and j = 2k in the list of sampled pixels, this yields Eq. (2):EQUATION)Figure 2 : Pixel embedding similarity maps. Large images: query images in which we select a query pixel (highlighted in red). For each query, we display two test images, with the pixel closest (in embedding space) to the query pixel highlighted in red.Similarity maps (cosine similarity between pixel embeddings) are also shown.Pix2Rep-v2. The computational complexity of contrasting positive and negative pairs of pixels limits us to sample M coordinates for the InfoNCE contrastive loss. Alternatively, we propose to replace this contrastive loss with a loss based on Barlow Twins (Zbontar et al., 2021) : we call this variant Pix2Rep-v2. It minimizes the Barlow Twins loss defined from the cross-correlation matrix C ∈ R d×d between twin pixel embeddings (ϕ • z)(s) and z ′ (s) ∈ R d , aggregated over all pixel coordinates s and the whole minibatch. Although it aggregates information from the whole pixel representation maps (rather than samples), Pix2Rep-v2 has a reduced memory footprint compared to Pix2Rep's contrastive loss. Downstream segmentation. For a given segmentation task, we initialize the encoderdecoder f with the pretrained weights (discarding the projection head g), and add a taskspecific, learnable segmentation head (1 × 1 conv + softmax), projecting pixel representations to class probabilities. We keep f frozen, and only train the segmentation head (called linear probing), or allow f to be fine-tuned from the supervised data (called fine-tuning). 
experiments: 
 We demonstrate the self-supervised pretraining on a downstream task of cardiac MRI segmentation.Data. The ACDC dataset (Bernard et al., 2018) consists of 3D short-axis cardiac cine MR images of 150 subjects, including expert annotations at End-Systole and End-Diastole for the left ventricle, right ventricle and myocardium. It is split into a training-validation set (100 images) and a test set (50 images). Slices are intensity-normalized using min-max normalization (using the 1st and 99th percentiles), cropped and resized to 128 × 128. We quantify the contribution of the Pix2Rep pretraining on the performance on a downstream segmentation task. The 3D Dice similarity coefficient is used as the evaluation metric. We report the average Dice over the test set over the segmented structures. For each evaluated method, each reported score is an average over five runs (training+test).Comparison to other methods. A natural baseline is obtained by keeping the same backbone U-Net, initialized with random (rather than pretrained) weights, then trained using varying amounts of labeled data X tr (|X tr | ∈ {1, 2, 5, 10, 20, 50, 100}) . We refer to it as the fully-supervised baseline.In addition, we compare with several methods in the state-of-the-art. We implement a mean teacher (Tarvainen and Valpola, 2017) semi-supervised model (using the same U-Net backbone), where the supervised loss applies to the available labeled slices, and the unsupervised consistency loss uses all slices from the 100 raw training volumes. As for selfsupervised models, we compare to (1) encoder-only pretraining of the same U-Net backbone, using image-level SimCLR, representative of Kalapos and Gyires-Tóth (2023) ; Punn and Agarwal (2022); (2) encoder-decoder (Pix2Rep) pretraining without rotations (crops only) and without intensity reversals as augmentations; and (3) the method of Chaitanya et al. (2020) . Unless stating otherwise, all self-supervised approaches use fine-tuning in the second stage rather than linear probing, as this yields higher performance. For the proposed approach, we report numbers both for linear-probing and fine-tuning.Backbone and Implementation details. We use a standard U-Net architecture consisting of five DoubleConv encoding blocks ({Conv,BatchNorm,ReLu,Conv,BatchNorm,ReLu}) with MaxPooling downsampling and five DoubleConv decoding blocks with Trans-Conv upsampling, starting with 128 feature maps and doubling after each block up to 2048 at the bottleneck. We experimented with changing the number of feature maps in the last layer (cf. ablation study), starting at 64, noting a monotonic improvement in performance up to the maximum tested value of 1024. We report the performance for n f t := 1024. For the fully-supervised baseline instead, we noticed decreased performance above and below n f t := 128 and report performance for this setting.The framework is implemented in PyTorch (https://github.com/pix2rep/). The spatial transformation ϕ is applied in auto-differentiable manner via affine grid, grid sample. Intensity reversal augmentations correspond to x := 1 -x. We pretrain using the proposed approach for 200 epochs with the Adam optimizer. The batch size is set to 8, the number of sampled pixels to 1000. The learning rate for the backbone is set to 5 • 10 -4 . For fine-tuning (100 epochs), the learning rate of the backbone is set to 5 • 10 -5 and of the classification head to 10 -2 . For linear probing (100 epochs), the learning rate is set to 10 -2 .Results. To gain qualitative insight into the learnt pixel representations, Fig. 2 plots the cosine similarity between a query pixel embedding and other pixel embeddings in two other representative images. The semantics of the anatomical structures are captured to some extent, as pixels anatomically similar to the query pixel have the highest similarity with it. Table 1 summarizes the main results. The proposed dense contrastive pretraining yields higher Dice than the fully-supervised baseline for all data regimes, as also shown in Fig. 3 .The gap is largest for few-shot segmentation; there is a 29.8% (resp. 31.2%) improvement in Dice for one-shot segmentation in linear probing (resp. fine-tuning). Of note is the limited drop in performance with linear probing even in the large data regime, highlighting the quality of the unsupervised nonlinear pixel representations. With fine-tuning, there is roughly a 5× reduction in the annotation burden for equivalent performance vs. the fully supervised baseline. Pretraining also reduces the variability in performance resulting from the choice of training subject, as shown by lower standard deviations in Fig. 3 .Ablation study. Table 2 shows the incremental contributions of various modifications on the proposed framework, starting from a suboptimal setup, in the few-shot segmentation setting (for |X tr | = 5). The inclusion of image intensity reversals (i.e., x := 1 -x) and rotations as augmentations has a significant impact on performance.Avg. Dice Defaults w/ 10 epochs pretraining 39.7 ± 1.0 10 epochs → 100 epochs 51.2 ± 2.7 64 feature maps → 1024 feature maps 73.4 ± 0.9 Add 'intensity reversal' augmentation 82.5 ± 3.0 Add rotation augmentation 84.5 ± 1.0 Linear-probing → Fine-tuning 89.2 ± 1.3Table 2 : Ablation study (|X tr | = 5). Self-supervised pretraining benefits from the increased number of feature maps as well as from the additional augmentations.Pix2Rep-v2.  
discussion & conclusion: 
 We have introduced Pix2Rep, a novel framework for pixel-level (dense) self-supervised representation learning, that allows to pretrain encoder-decoder architectures such as U-Nets directly from unlabeled images. We have shown performance gains on a downstream cardiac MRI segmentation task. Especially in the few-shot segmentation regime for the most challenging structure, we got 83% Dice with 5 training subjects vs. 70% for the fully-supervised baseline, and closest SOTA method 3% below. As future work, in addition to comparing to Yan et al. (2022) , we plan to evaluate the framework on various segmentation tasks and assess the generalizability of learned representations for novel tasks (foundation model).Yuanyi Zhong, 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 (ICCV), pages 7273-7282, October 2021.  
 
 Appendix: For downstream segmentation tasks, the segmentation head R H×W ×D → R H×W ×K , with K the number of classes, is plugged back at the end of the pretrained backboneThe action ϕ • v ′ of a spatial transformation ϕ on an image v ′ is v ′ • ϕ -1 . This is well-known in the image registration literature. Following standard practice, we randomly sample ϕ -1 directly (rather than ϕ) to circumvent the numerical inversion. 


Decision: This paper extends image-level contrastive self-supervised learning (SSL) such as SimCLR to pixel-level representations. While the reviewers raised their scores somewhat after the rebuttal, the final ratings were weak accept, borderline accept and borderline reject. The reviewer who originally gave a strong reject and raised to borderline reject after the rebuttal still questioned the claimed novelty of the work. 

The main claim of the paper previous to the rebuttal appeared to be that contrastive SSL for pixel level representations was not studied before. However, as was pointed out in the reviewer comments this statement is too strong as there has been previous methods which use contrastive learning at the pixel/voxel level for SSL. The authors pointed out differences of their work and clarified in response that "while rotations have been used for creating positive pairs in the context of image-level contrastive SSL, they have not been used to our knowledge for creating positive pairs in the context of pixel-level contrastive SSL." However, in another response the authors stated "Our paper proposes a novel method, extending image-level contrastive SSL such as SimCLR to a pixel-level loss function." which seems inconsistent with the previous statement. Other references in the paper such as vox2vec are also clearly pixel-level contrastive formulations  even if they don't use equivariance under rotations. 

Another issue I see is that the paper does not present their contribution in the context of the most related methods. The authors stated in the rebuttal "Since we introduce a novel framework, we chose to compare our pixel-level loss to image-level loss functions on SOTA contrastive SSL architectures." While these might be the most related conceptually, I think the method should also be compared against existing dense contrastive SSL methods such as works that were already cited in the original submission, e.g. vox2vec. 

Overall, I think this paper is borderline due to the weaknesses discussed above; however, I think the method presented has some degree of novelty which slightly outweighs the criticisms.

Reviews: review: paper_summary: paper_summary: This study presents a few shot learning, and SSL based segmentation method, Pix2Rep, to reduce the annotation burden for medical image segmentation. Results, obtained from cardiac segmentation experiments, show that the proposed method is beter than other self-supervised and semi-supervised methods, and similar performance to fully supervised method with 5-times reduction in efforts.
main_review: main_review: I have decided to change my decision for three steps to borderline accept.
Authors clarified some misunderstandings, this is good.
Innovation, I have still some problems as those concepts are already studied, and more heavy and structured data augmentation already gives rise into dice scores, and last but not least the 3D should appreciate more information while 2D is preference, not totally clear.
strength_weakness: strength_weakness: -SSL and a few shot learning based segmentation are heavily explored in medical domains, not all studies are recognized, and comparisons are weak in results section. 
-intensity reversal and rotations are proposed already for data augmentation. Not new. Therefore, it is not entirely clear what authors are presenting as novelty.
questions: questions: - any chance to explain how the method operates on 3D?
- why not use individual object dice to give an idea where it fails and performs better?
- not mentioned if mri is preprocessed or not.
- how about using more adaptive data augmentation and try comparison in that way, instead of comparing with limited augmentation to none with other baselines?
- what is new in this paper, truly? all the aforementioned methods and strategies were all presented before, and amount and the true idea behind the geometric transformations and its additive value to data augmentation during SSL is not exclusively explored. This is an application paper and hard to judge then how to position the paper theoretically, as there is only single segmentation approach (MRI cardiac).
paper_summary: paper_summary: Paper demonstrates a framework for pixel-level contrastive representation learning and demonstrates usefulness of learned representations to downstream segmentation task of medical images in various data-regimes. Intensity reversal augmentations and rotations were proposed to farther improve results.
main_review: main_review: Based on the clarification of the authors during rebuttal and better understanding of the contribution, I have increased my score to borderline (reject). I still believe the novelty is rather small (adapted from computer vision), so I am not voting directly for acceptance, but I also would think the paper has some merit to be discussed at MIDL.
strength_weakness: strength_weakness: The proposed method lacks methodological novelty. As authors note in related works dense contrastive learning by itself is not a new method for pre-training on medical images datasets. Application of different types of augmentations is a core component in contrastive learning and also investigated a lot. In contrast to claimed in the paper contrastive learning under rotation augmentations and also elastic deformations demonstrated in following work: ‘SAM: Self-supervised Learning of Pixel-wise Anatomical Embeddings in Radiological Images’ by Ke Yan et. al.  (https://arxiv.org/pdf/2012.02383.pdf)
questions: questions: The authors should address the novelty concerns regarding the proposed method. It’s not clear what kind of intensity augmentations (if at all) are used as ‘limited augmentations’ for the baseline method and what their strength.
paper_summary: paper_summary: Pix2Rep is a self-supervised learning approach for few-shot medical image segmentation, significantly reducing the need for extensive annotations. It works with any deep learning architecture and focuses on learning pixel-level representations through equivariance, showcasing superior performance in cardiac segmentation tasks with substantially less annotation effort.
main_review: main_review: The final rating reflects a comprehensive evaluation of the author's response to feedback throughout the revision process. Each comment was meticulously considered, leading to significant improvements in the clarity, depth, and accuracy of the work. The author's dedication to refining their manuscript based on the constructive critiques provided is commendable. This thorough engagement not only enhanced the overall quality of the submission but also demonstrated the author's commitment to excellence and receptiveness to feedback. The enhancements made as a result of addressing the comments have substantially contributed to the work's coherence, making it a valuable contribution to its field. Therefore, the final rating is justified by the author's successful efforts to integrate all provided feedback, ultimately elevating the work's academic and practical implications.
strength_weakness: strength_weakness: While Pix2Rep presents a compelling solution to the annotation challenge, its effectiveness across a broader range of medical imaging modalities and segmentation tasks outside of cardiac imaging remains to be fully assessed. The reliance on self-supervised learning may introduce complexities in tuning and optimizing the model for specific tasks, potentially requiring domain-specific adjustments. Moreover, the impact of the approach on computational resources, given the need for processing large volumes of unlabeled data, and its scalability to very large datasets are concerns that need further exploration.
questions: questions: No comments

Comments: title: Please read and respond to author comments
Authors' response to Reviewer t8ZT
Authors' response to Reviewer nbJh (2/2)
Authors' response to Reviewer nbJh (1/2)
Authors' response to reviewer Udt4
comment: No further comments
We thank all three reviewers for taking the time to respond, and for their positive feedbacks on our rebuttal, especially nbJH who has agreed to improve his/her score (quoting: *"I hence improved my final score"*).
thank you for the comments
Thank you for clarifications and incorporation of suggestions. Though I still doubt the claimed novelty of the method. Learning self-supervised pixel-level representations through equivariance was used for natural images (while using real flow transformations rather than synthetic ones). See this paper for example: 'Self-Supervised Representation Learning from Flow Equivariance' - https://openaccess.thecvf.com/content/ICCV2021/papers/Xiong_Self-Supervised_Representation_Learning_From_Flow_Equivariance_ICCV_2021_paper.pdf. In this case the novelty of the method is reduced to application of equivariance to self-supervised learning for medical images while using synthetically generated transformations. Nevertheless it can be interesting for medical imaging community to see such a boost in results while using this method. I hence improved my final score
Thanks for addressing those concerns. I would accept it now.
Thanks for addressing those concerns. I would accept it now.
Thanks for addressing those concerns. I would accept it now.
Thanks for addressing those concerns. I would accept it now.
Thanks for addressing those concerns. I would accept it now.
Thanks for addressing those concerns. I would accept it now.
Dear reviewers. The authors have posted responses to your reviews. Please take the time to read and respond before March 27.
**Positive comment overall:** “Its demonstrated ability to significantly reduce
annotation requirements while maintaining, or even improving, segmentation performance is a notable achievement.” + “The method's flexibility, applicability to existing deep learning models, and substantial improvements in efficiency and effectiveness provide a strong basis for a positive preliminary assessment”

**R3.1** “further validation across diverse imaging modalities and segmentation challenges will be crucial”

- This point is the focus of our current on-going investigations.  

&nbsp;

**Questions To Address In The Rebuttal:** 

None
**Questions To Address In The Rebuttal:**

**R2.5** “The authors should address the novelty concerns regarding the proposed method.  It’s not clear what kind of intensity augmentations (if at all) are used as ‘limited augmentations’ for the baseline method and what their strength.”

- We have rewritten the abstract (see above or in the revised paper) to improve the clarity of our paper on its novelty component.

- In Table 1, the row “Enc.-dec. Pretraining (limited augm.)” refers to our proposed approach, without rotations and intensity reversal (only intensity transforms, translations and crops). We agree that “limited augm.” is a misnomer, which we have changed to “Proposed (w/o rot. & int. reversal)”.
**Comments:**

**R2.1** “The paper does not bring methodological novelty but rather reimplement existing methods.”

- (cf. R1.5) We were not clear enough on the aim of our paper, which is to introduce a novel formulation of SimCLR and Barlow Twins at pixel level.

**R2.2** “As authors note in related works dense contrastive learning by itself is not a new method for pre-training” + “Application of different types of augmentations is a core component in contrastive learning and also investigated a lot”+ “ [..] demonstrated in following work: “‘SAM: Self-supervised Learning of Pixel-wise Anatomical Embeddings in Radiological Images’ by Ke Yan et. al.”

- While we agree that dense SSL is not new, we propose a new formulation for this task, using a principle of equivariance under spatial transforms which was not proposed before.

- We now cite the paper by Ke Yan et al. in the revised version of our paper, along with other similar methods (VADeR and vox2vec) we already cited in the Related Work section. We have overseen Ke Yan et al. paper and are very grateful to the reviewer for pointing this out.

- The method of Ke Yan et al differs from our work in several aspects:

  1. Ke Yan et al. generate two views per pixel via random "data" augmentation (at image level per Fig 1 "We first do random data augmentation to an unlabeled image") and then local random crops around a given pixel where they select positive pairs only in the overlap between the two crops. In contrast, we rely on a principle of equivariance (geometric transform performed at different stages of the encoding-decoding) and work with the whole image encoding instead of crops. Concretely, we apply a single spatial transform asymmetrically (before or after the encoder-decoder backbone) to two copies of the whole image (that also differ up to intensity transforms). This results in two views that have complete spatial overlap, from which we can easily sample positive pairs at any pixel locations. 
  2. On the other hand, in Table VI Ke Yan et al. report some tests on different “data augmentations” including rotations and deformations. This table reports great benefits from "Def.+Rot." for the hand Xray landmark detection task.
  3. Using crops raised the following issue in Ke Yan et al. paper: "SAM is robust to the overlap ratio of the two augmented patches in training, but the overlap upper bound should not be too small." . We don't face such issues in constraining spatial transforms as we use whole images that remain aligned after similar geometric transforms performed in different orders. 
  4. From our reading of their official github code, Ke Yan. et al ended up using only random crops to generate paired views, without rotations in the released version.
  5. Ke Yan. et al. did not report the importance of rotation augmentation in their segmentation use-case on liver.
  6. Ke Yan et al. report in Table VIII few-shot learning for segmentation with a loss of ~7% in Dice score from N=105 to N=5 in the labeled training set. We report a loss of only 4% in Table 1 for comparable training set sizes and Dice values (above 90% in low-regime).
  7. Other differences include: backbone architecture, negative pairs sampling strategy, use of pretraining on ImageNet in Ke Yan et al.
  8. Ke Yan et al. have tested more advanced strategies for negative pairs sampling that are very relevant and that we will compare to in future work.

**R2.3** “The generalisability of the proposed method is not addressed” + “ACDC dataset that is also used for pre-training”

- Since we are introducing a new method in a conference paper, we focused on a single use-case for cardiac MRI segmentation, using an open-access cohort for reproducibility.

- Indeed we performed pre-training and fine-tuning on the same cohort which is a limitation on the scope of this preliminary work. We made sure to stratify training and test sets on disease status following the ACDC partition.

- This initial publication using an open-access image cohort as our use-case will guarantee research reproducibility once we share our github repository upon acceptance of the paper.

**R2.4** “It can be interesting to visualise feature maps of pertained model to see how good is clustering for different classes.”

- We thank the reviewer for this very pertinent suggestion. We provide in the updated paper version (in the Supplement C section), some visual illustrations of our pixel-level representations learned with our SSL pretrained model (before fine-tuning for the segmentation task). The 2D t-SNE embedding of pixel representations, and a visual mapping in some predefined color space confirm that our proposed contrastive SSL method already realizes a coarse segmentation of the different cardiac structures, without any supervision.
**Comments:**

**R1.1** “This study presents a few shot learning, and SSL based segmentation method, Pix2Rep” + ”I believe that the paper is an application paper”

- Our paper proposes a novel method, extending image-level contrastive SSL such as SimCLR to a pixel-level loss function.

- This pixel-level loss function is derived from a principle of equivariance under spatial transforms (Fig. 1), which to our knowledge has not been proposed before.

**R1.2** “comparisons are weak in results section”

- Since we introduce a novel framework, we chose to compare our pixel-level loss to image-level loss functions on SOTA contrastive SSL architectures.

- We also compared to the pixel-level SSL method of Chaitanya et al (NeuRIPS 2020), as well as to a semi-supervised approach (Mean Teacher). 

- Also, we provided an ablation study to compare results with and without key components of our encoding method, such as rotation equivariance (Table 1 “Enc.-dec. Pretraining (limited augm.)” where "limited augm" means without rotations and intensity reversal) as well as the influence of training hyperparameters (Table 2)

- The comparisons are conducted using the same backbone, to focus on assessing only the impact of our contributions.

- We focused on a single segmentation use-case for a first publication of our method in a conference paper.

**R1.3** “intensity reversal and rotations are proposed already for data augmentation. Not new.” + “It seems to me that heavy and more engineered data augmentation will increase the results in any way”

- Indeed, we aimed to compare our pixel-level loss to image-level loss for contrastive SSL using standard image augmentation transforms.

- But, while rotations have been used for creating positive pairs in the context of _image-level_ contrastive SSL, they have not been used to our knowledge for creating positive pairs in the context of _pixel-level_ contrastive SSL.

&nbsp;

**Questions To Address In The Rebuttal:**

**R1.4** “any chance to explain how the method operates on 3D?”

- The steps presented in the Method, and summarized in Fig. 1, are valid in 3D as well as in 2D

- Our choice of 2D versus 3D is motivated by the increased computational cost at pixel-level with our method. We agree that taking our method to 3D might be beneficial but it will raise computational challenges since we work with whole images rather than patches.

**R1.5** “why not use individual object dice to give an idea where it fails and performs better?”

- We have added in the revised paper, in the Supplement B section, all the individual Dice scores per anatomical structures for all experiments. For instance, for |X_tr|=5, we report Dice values for (LV, RV, MYO, BG) as: for the fully-supervised baseline (82.7%,69.8%,69.2%,95.3%) and for our method Pix2Rep+Fine-tuning (93.9%,83.0%,86.3%,97.6%). The fact that larger errors are observed on the RV and MYO, compared to the LV, is consistent with the literature and is observed across all the |X_tr| values and tested methods.

**R1.6** “not mentioned if mri is preprocessed or not.”

- As mentioned in the text, axial slices were intensity-normalized using min-max normalization (using the 1st and 99th percentiles), cropped around the heart region and resized to 128 × 128 pixels, which is standard for this dataset.

**R1.7** “how about using more adaptive data augmentation and try comparison in that way, instead of comparing with limited augmentation to none with other baselines?”

- In Table 1, the row “Enc.-dec. Pretraining (limited augm.)” refers to our proposed approach, without rotations and intensity reversal (but using intensity transforms, translations and crops). Using translations and crops only is the standard in pixel-level SSL literature, so we used it as a point of comparison. “Limited augm.” is a misnomer, we have changed it to “Proposed (w/o rot. & int. reversal)”.

- Introducing novel adaptive image augmentation strategies is a very relevant field of investigation for an extension of our work.

**R1.8** “what is new in this paper, truly? all the aforementioned methods and strategies were all presented before, and amount and the true idea behind the geometric transformations and its additive value to data augmentation during SSL is not exclusively explored. This is an application paper and hard to judge then how to position the paper theoretically, as there is only single segmentation approach (MRI cardiac).”

- We clearly failed in our writing to emphasize the novelty of our method. We have now reformulated the abstract with color-coded changes to clarify this point.

- We focused on a single segmentation use-case for a first publication of our method in a conference paper.
**All reviewers entered** “**confidence**: 4: The reviewer is confident but not absolutely certain that the evaluation is correct”

We have posted in individual rebuttals our responses to some negative comments and further details on requested questions to address. We understand from the first two reviews that we were not clear enough on the novelty of our proposed method and we now propose the following revision of the abstract (below).

We have posted a revised version of our paper (new text in magenta) which contains the following updates: (1) updated abstract, (2) a reference to the paper ‘SAM: Self-supervised Learning of Pixel-wise Anatomical Embeddings in Radiological Images’ by Ke Yan et. al. mentioned by **reviewer nbJh**, (3) updated conclusion (4) Appendix B with Dice results per organ, (5) Appendix C with a visualization of pixel-level embeddings with a t-SNE plot and with a colored parametric map. 

We have added the following sentence in the conclusion:
>Especially in the few-shot segmentation regime for the most challenging structure, we got 83% Dice with 5 training subjects vs. 70% for the fully-supervised baseline, and closest SOTA method 3% below. As future work, in addition to comparing to Yan et al. (2022) [...]   

&nbsp;
   
***

**Revised abstract (changes in bold)**:

Deep learning has revolutionized medical image segmentation, but it relies heavily on high-quality annotations. The time, cost and expertise required to label images at the pixel-level for each new task has slowed down widespread adoption of the paradigm. We propose Pix2Rep, a self-supervised learning (SSL) approach for few-shot segmentation, that reduces the manual annotation burden by learning powerful pixel-level representations directly from unlabeled **images**. Pix2Rep **is a novel pixel-level loss and pre-training paradigm for contrastive SSL on whole images. It is applied to generic encoder-decoder** deep learning backbones (e.g., U-Net). Whereas most SSL methods enforce invariance of the learned image-level representations under intensity and spatial image augmentations, Pix2Rep enforces equivariance of the pixel-level representations. We demonstrate the framework on a task of cardiac **MRI** segmentation. Results show improved performance compared to existing semi- and self-supervised approaches; and a 5-fold reduction in the annotation burden for equivalent performance versus a fully supervised U-Net baseline. This includes a 30% (resp. 31%) DICE improvement for one-shot segmentation under linear-probing (resp. fine-tuning). **Finally, we also integrate the novel Pix2Rep concept with the Barlow Twins non-contrastive SSL, which leads to even better segmentation performance.**