Self-Supervised Pretraining
for 2D Medical Image Segmentation

Self-supervised learning methods learn representations by solving a pretext task where the supervisory signal emerges from the raw, unlabelled data. This paradigm enabled natural language processing (NLP) models with billions of parameters to learn from huge datasets and became the state-of-the-art approach for many problems in NPL[GPT, BERT]. Applying transfer learning to models that were pretrained on Imagenet[ImageNet] with supervised learning has been an effective and widely used method in computer vision. However, models pretrained with self-supervision are achieving state of the art for vision tasks as well, and the possibility of learning useful visual representations from uncurated and unlabelled data has been demonstrated[SEER].

However, constructing good pretext tasks for image data is not as intuitive as for text due to the high dimensionality and continuous nature of images. Early approaches used heuristic pretext tasks such as solving a jigsaw puzzle [JigsawSSL], predicting the relative position of image patches [PatchSSL], colourization [ColorizationSSL] and denoising [DenoisingSSL]. Recent methods ([CPCv2, SimCLR, MoCo, SwAV, BYOL, BarlowTwins] etc.) rely on the instance discrimination task, which treats augmented versions of every instance from a dataset as a separate class. A set of random augmentations are applied to each image, resulting in various views of it. The learning objective is then to assign similar latent vectors to different views of the same image, but different ones to views of different images. Methods that use such an instance discrimination task primarily differ in the set of augmentations and the approach to enforce this similarity of the latent vectors.

These methods aim to train a model that acts as an encoder, mapping images to an embedded vector space, also called latent vectors. These encoders could then be incorporated into downstream tasks as pretrained models.

The most common transfer learning approach in computer vision is fine-tuning a model that was pretrained for supervised image classification on ImageNet [ImageNet]. Another approach receiving notable attention recently is self-supervised pretraining either on ImageNet or on domain-specific data (i.e. on data that is similar or related to the problem tackled by transfer learning). Reed et al. [HierarchicalPretraining] investigate self-supervised pretraining pipelines including ImageNet pretraining (generalist pretraining), domain-specific (or specialist) pretraining and the sequential combination of the previous two, which they refer to as ”hierarchical pretraining”.

Medical image recognition problems are commonly solved using transfer learning from natural images, even though the distribution, frequency pattern, and important features of medical images are not the same as those of natural images [Morid2021]. Self-supervised pretraining on domain-specific data offers an alternative approach to this task, eliminating the mismatch in source and target dataset characteristics, and allowing a methodologically more precise pretraining. Moreover, domain-specific pretraining has the potential to reduce the number of labelled data, which has high costs in the medical domain.

Recently, a handful of papers have been published that adapt self-supervised learning methods to pretrain models on medical image analysis problems such as histopathology [ciga2021self], chest X-ray and dermatology classification [sowrirajan2021moco, BigSelfSupervisedModels]. These works show that MoCo [MoCo] and SimCLR [SimCLR] are effective pretraining methods on medical images without major modifications to the algorithms. They highlight the capability of these algorithms to outperform supervised ImageNet pretraining by 1-5% classification accuracy on various medical image recognition datasets.

Training pipelines consist of one or two pretraining stages followed by the downstream supervised segmentation training (see fig. 1). All reported metrics and learning curves are from the final segmentation training step. Between each training stage, we transfer the weights of the encoder network to the next stage, but every stage uses additional layers, which we initialize randomly. As explained in section 3.2, the applied self-supervised learning algorithm trains an encoder through two subnetworks (called the projector and predictor). The downstream segmentation network also has an additional module after the encoder: a decoder. All these subnetworks are randomly initialized at the beginning of each training phase using Kaiming-uniform initialization [KaimingHeInit], while the weights of the encoder are always loaded from the previous stage. In the first stage of every pipeline, the encoder is random-initialised as well.

Grill et al. [BYOL] proposed the Bootstrap Your Own Latent (BYOL) self-supervised learning method, which relies on the instance-discrimination task described in section 2.1 and a special asymmetric architecture to avoid representation collapse. BYOL trains an encoder ($f_{\theta }$ on fig. 2) to learn the representations of another fixed encoder ($f_{\xi }$ on fig. 2) for $x_1,x_2$ different views of the same input image. In other words, the $f_{\xi }$ fixed network produces some $h_2$ representation of the $x_2$ augmented version of the image, which the trained network should reproduce, even though its input is a different $x_1$ augmentation of the image. Empirical evidence shows, that the trained $f_{\theta }$ ”online” network can learn better representations than that of the fixed $f_{\xi }$ ”target” network. This allows for an iterative update procedure where the fixed network’s weights are also slowly updated, based on the improved weights of the online network. The latter one is trained by gradient-based learning, while the target network is updated using the exponential moving average of the online network’s weights. Gradients are not propagated through the target network, however, as training converges the weights of the two networks become similar and can be used for downstream tasks.

To achieve stable training, BYOL requires a $q_{\theta }$ predictor network in the online branch of the network. BYOL trains the online branch using the mean squared error between the $q_{\theta }\left(z_1\right)$ output of the predictor and the $z_2$ output of the $g_{\xi }$ target projector. See fig. 2 for an illustration of BYOL’s asymmetric architecture.

The “Automated Cardiac Diagnosis Challenge” dataset [ACDCdataset] is a publicly available dataset aiming to improve deep learning-based automated cardiac diagnosis. The dataset contains cardiac MRI records from 100 patients, distributed evenly in five groups, one healthy and four with different cardiac conditions. For each patient, the complete cardiac cycle²²2One cardiac cycle is the period between two heartbeats. is covered by the recordings, where each recording consists of 3D MRI scans, which include 6-18 high-resolution 2D slices. This gives approximately 25 thousand slices in total for the dataset. The resolution of the slices varies from $150\times 150$ to $500\times 500$ pixels.

Segmentation labels are provided for the end-diastolic (ED) and end-systolic (ES) frames of the cardiac cycle, which gives 1900 slices with segmentation masks. For the rest of the frames, accurate segmentation labels are not available, however, self-supervised learning algorithms can leverage this data as well. We refer to the slices of annotated frames as the labelled 2D segmentation dataset and include all 25k slices (labelled and unlabelled) in the unlabelled dataset. One can consider this a semi-supervised learning dataset, with labels only available for a subset of all slices.

Segmentation masks are provided with four classes, left and right ventricle (LV, RV), myocardium (MYO), and background (see fig. 7 as an example).

We selected a U-Net[UNet] architecture for the segmentation training due to the successful application of this architecture for medical image segmentation. The architecture and implementation of the encoder match with the ones we pretrain, which is necessary for transferring the pretrained weights to the segmentation model. Encoders pretrained on ImageNet expect three-channel, RGB images as input, however, MRI scans are monochromatic. Therefore, the first layer of these encoders needs to be adjusted for single-channel input, which we achieve by depth-wise summation of the weights of the first layer (after [PyTorchImageModels]).

We aim to test the hypothesis that efficient pretraining improves the downstream accuracy in low-annotated-data scenarios. To test this hypothesis, we run segmentation training on subsets of the labelled dataset, ranging from 1 sample to the full dataset. We also vary the pretraining strategy preceding the segmentation training. Each training runs for the same number of steps which is equivalent to 150 epochs on the full dataset and each run was repeated with 10 different randomly set initial parameter states (i.e. random seeds).

Figure (a)a shows learning curves for downstream segmentation training with different pretraining approaches. The figure shows that hierarchical pretraining methods converge faster than generalist or specialist pretraining on this semantic segmentation task. This indicates that representations learned during both stages retain information useful for the target task. However, two notable differences separate these hierarchical methods as well. First, pretraining involving only self-supervised steps leads to faster and more stable convergence compared to supervised ImageNet pretraining followed by domain-specific SSL pretraining. Second, the final segmentation Jaccard index (Intersection over Union, IoU score) for the latter approach is significantly lower than for any other method. Figure (a)a also reveals that self-supervised only hierarchical pretraining (BYOL ImageNet + BYOL ACDC) leads to downstream convergence in $\sim$200 steps, while other methods reach close to maximal IoU in 4-5 times more steps. The area under the learning curves (see fig. (b)b) highlight the stability of self-supervised only hierarchical pretraining even more clearly (BYOL ImageNet + BYOL ACDC on the figure). This figure also indicates that specialist self-supervised pretraining (BYOL ACDC) leads to similar results as mixing supervised and self-supervised pretraining (Supervised ImageNet + BYOL ACDC). Our findings show that among the investigated approaches, self-supervised only pretraining (i.e. BYOL ImageNet + BYOL ACDC pretraining) produces representations that achieve the fastest and most stable downstream convergence.

Figure 6 shows how downstream error scales with the available annotated data, in the case of different pretraining methods. Contradicting our hypothesis, the figure reveals that with low data quantities generalist pretraining approaches have lower test error compared to hierarchical pretraining.

Hestness et al. [PowerLawOfDL] investigates similar scaling behaviour as visible on fig. 6. They show the existence of a power-law region, where the reduction of test error follows a power-law scaling with increasing data quantities. Hestness et al. [PowerLawOfDL] also point out that by increasing the amount of training data above a certain threshold, the improvement of test error starts to decrease slower than a power-function of the data quantity. We can observe both scaling regions on fig. 6. Below 2% of the total available training data, that is  30 MRI slices, generalist pretraining approaches seem to follow power-law scaling (linear scaling on a log-log plot), while above 5% of available training data, the reduction of error is notably slower, indicating the transition to the irreducible error regime. The amount of training data that corresponds to this transition is important to machine-learning development and decision-making, because investing more in annotated data above this region has a much smaller return in error reduction, than with fewer data. In simple terms, labelling a given number of additional samples is increasingly less profitable above this transition region. In our experiments this transition region is around 30-75 randomly selected slices of MRI records, annotating more slices still improves downstream performance, however, the expected improvement after adding an annotated sample reduces.

We experiment with running the domain-specific step for a different number of epochs. As shown on fig. 8, we find that 3-4 epochs are sufficient to achieve the fast-converging, stable learning curves that we presented in section 5.1.