Two-Stage Multi-task Self-Supervised Learning for Medical Image Segmentation

Medical image segmentation (MIS) aims at accurately delineating the target of interest (e.g., organs and lesions) in medical images, e.g., X-ray, CT, and MRI. So far, MIS has been significantly advanced by DL, demonstrated by state-of-the-art performance in a broad spectrum of tasks. Due to their over-parameterised nature, models in DL are data-hungry and typically require abundant annotated data to ensure good performance. Because the annotating process is expert-demanding and labour-intensive, MIS annotations are expensive and rare. As a result, DL learning with small annotated datasets has been in active research. So far, many endeavours have been focusing on improving model architecture [13, 14, 15] to more efficiently extract generalisable features. Other works have explored learning in certain settings with no or cheaper annotations. For instance, semi-supervised learning leverages abundant unlabelled imaging data to facilitate learning. Self-supervised learning [16] can train models solely with images in an unsupervised manner. Weakly supervised learning enables learning from weak but cheap annotations. Our work focuses on the fully supervised setting where only an annotated dataset is available, and there are no additional data or annotations.

Self-supervised learning (SSL) generally aims at feature representation learning from unlabelled data. In a common SSL paradigm, a model is pre-trained on large unlabeled data with carefully designed pretext tasks to serve as a general initialisation. It is expected that the knowledge learnt for solving the pretext tasks can be transferred to help solve various downstream tasks. Recent advances in SSL have achieved tremendous success in training DL-based models. So far, SSL has not only been widely demonstrated as an effective tool to combat data scarcity and boost performance but has also become a foundation technique for training large models [9]. The art of SSL is to design pretext tasks considering both the data domain and the downstream task property so that rich and useful knowledge can be extracted from the unlabeled upstream data to benefit downstream task performance. With abundant unlabelled data existing in the medical field, SSL for MIS has also been actively explored, most of which adapts the natural image counterparts to the medical image domain [8, 17]. The prosperity of SSL provides diverse unsupervised tasks that can serve as auxiliary tasks in our supervised setting.

Multi-task learning (MTL) is a machine learning paradigm that learns multiple related tasks simultaneously. MTL can enrich overall feature representation and boost task performance by sharing complementary knowledge obtained from solving different tasks or letting tasks act as regularisers of one another [18]. MTL techniques have been extensively explored by improving model architecture or optimisation process, aiming to better capture shared information, avoid conflicts among tasks, or balance the performance of different tasks. While MTL typically relies on a dataset with multiple annotations, model training with a single-task dataset can still take advantage of MTL by involving auxiliary tasks. This setting, sometimes known as auxiliary learning, is a special case of MTL, which focuses on only the performance of the target task. Existing auxiliary learning approaches commonly adopt the MTL formulation. Based on the joint training paradigm, different ways to adaptively adjust auxiliary task contributions have been explored. Most methods focus on task-level adaptiveness via selecting auxiliary tasks to learn from [4, 19] or reweighting the contribution of different tasks in loss [6, 20] or gradient aggregation [4]. To achieve this, some methods rely on heuristic rules to assess task relationships [4, 19], and others are driven by feedback from validation set [6, 20].

Knowledge distillation (KD) [7] is a popular model training technique in DL to enable knowledge transfer among models. KD was initially proposed for model compression [7, 21], where a large-sized pre-trained “teacher” model trains a small-sized “student” model by supervising the student with the teacher’s predictions. Through soft probabilistic labels, rich information can be transferred to the student model to facilitate its learning process [22, 23]. KD techniques have received active exploration due to the flexibility to transfer knowledge solely based on model outputs, posing no restriction on the architecture or the number of teacher models. So far, KD operations have been further developed to enable distilling knowledge with different representations, such as intermittent feature maps and inter-class relationships. The utility of KD has also been extended to transfer knowledge from multiple teacher models [24], a teacher trained together with the student [25], or even the student itself [26].

For conciseness, we use single-class segmentation on 2D single-channel images to derive the notations and formulations to be used in the following sections. Note that our method is naturally compatible with 3D, multi-channel, multi-modalities, and multi-class segmentation settings. We use $\bm{x}\in\mathbb{R}^{H\times W}$ to denote an image, where $H$ and $W$ represent its height and width, respectively. The binary segmentation mask annotation for each image is denoted as $\bm{y}\in{\{0,1\}}^{H\times W}$. A model for image segmentation can be represented as $f(\cdot;\bm{\theta})$, where $\bm{\theta}$ denotes its parameters.

In a typical training process, a model $f$ is trained by optimising its parameters $\bm{\theta}$ to solve the segmentation task represented by a training set $\mathbb{D}_{seg}^{tr}$, and the training process can be formulated as:

where $\bm{\theta}^{*}$ denotes the optimised parameters, and $\mathcal{L}_{seg}$ represents the loss function, which is often the Dice loss in the MIS domain. After training, the generalisation performance of $f(\cdot,\bm{\theta}^{*})$ is evaluated on a test set $\mathbb{D}_{seg}^{test}$ which does not overlap with $\mathbb{D}_{seg}^{tr}$.

Generally, we can denote the segmentation task as $\mathcal{T}_{seg}=\{\mathbb{D}_{seg}^{tr},\mathcal{L}_{seg}\}$. An auxiliary task $\mathcal{T}_{i}=\{\mathbb{D}_{i}^{tr},\mathcal{L}_{i}\}$ aims to train the model to enrich its knowledge representation in $\bm{\theta}$ by bringing in additional knowledge that cannot be learnt solely from $\mathcal{T}_{seg}$ and thus improve model’s generalisation performance. The dataset $\mathbb{D}_{i}^{tr}$ used by auxiliary task $\mathcal{T}_{i}$ is created based on $\mathbb{D}_{seg}^{tr}$, using either the images and labels in $\mathbb{D}_{seg}^{tr}$ or solely the images. Different tasks have different formulations and can be heterogeneous in terms of their input and output spaces. For example, in a segmentation task using whole images, $\bm{x}\in\mathbb{R}^{256\times 256}$ and $\bm{y}\in\{0,1\}^{256\times 256}$, wheras in a contrastive learning task using image patches, $\bm{x}\in\mathbb{R}^{128\times 128}$ and $\bm{y}\in\mathbb{R}^{1024}$. To differentiate the inputs and outputs among tasks, we use $\bm{x}_{seg}$ and $\bm{y}_{seg}$ to denote the image and label for segmentation and use $\bm{x}_{i}$ and $\bm{y}_{i}$ to denote the image and label for auxiliary task $i$. Our work involves multiple auxiliary tasks $\{\mathcal{T}_{i}\}_{i=1}^{N}$ to improve model training, where $N$ denotes the total number of auxiliary tasks. To enable knowledge transfer, every auxiliary task shares a certain part of the model with the segmentation task, parameterised by $\bm{\theta}_{sh}$, and every task, including the segmentation task, has its task-specific head. We use $\bm{\theta}_{seg}$ to denote the segmentation-specific parameters, and $\bm{\theta}_{i}$ to denote the parameters specific to auxiliary task $\mathcal{T}_{i}$.

Given a target MIS dataset and multiple auxiliary tasks, the proposed method aims to obtain a model with good performance on the MIS dataset by leveraging the auxiliary tasks. As shown in Figure 1, the method is composed of two training stages. First, the target segmentation task $\mathcal{T}_{seg}$ will be independently co-solved with each auxiliary task $\mathcal{T}_{i}$ in both joint training and pre-training modes, and the better model obtained by the two modes $f_{i}$ is selected via validation performance. As a result, multiple elite segmentation models $\{f_{i}\}_{i=1}^{N}$ are obtained, each containing useful and distinct knowledge transferred from an auxiliary task. Second, these models are treated as teacher models and are compressed into a single model $f_{s}$ via ensemble knowledge distillation [7]. In this way, the diverse knowledge obtained from multiple auxiliary tasks is aggregated to obtain an even more powerful student model, which is output for inference. The following sections describe the two training stages in detail.

This stage aims to obtain a segmentation model independently facilitated by each auxiliary task. In this way, the interference among auxiliary tasks can be avoided, and the most appropriate training mode for each auxiliary task can be independently applied to make best use of each auxiliary task. For auxiliary task $\mathcal{T}_{i}$, we train two models to co-solve $\mathcal{T}_{i}$ and segmentation $\mathcal{T}_{seg}$ in two modes: joint training and pre-training, respectively.

In the joint training mode, the model is trained to concurrently solve the segmentation and auxiliary tasks in a multi-task fashion. At each iteration step, we sample a batch of samples for the segmentation task and the auxiliary task, respectively, i.e., $\bm{x}_{seg},\bm{y}_{seg}\sim\mathbb{D}_{seg}^{tr}$ and $\bm{x}_{i},\bm{y}_{i}\sim\mathbb{D}_{i}^{tr}$. Then the inputs are fed to the model, and the training loss is calculated as

where $\lambda_{i}$ denotes a scalar coefficient to weight the loss contribution of task $i$. A larger $\lambda_{i}$ encourages the trained model to learn more from the auxiliary task and vice versa. All the model parameters $\bm{\theta}_{sh}$, $\bm{\theta}_{seg}$, and $\bm{\theta}_{i}$ are jointly updated to minimise $l^{joint}$.

The pre-training mode involves two training steps. First, the model is pre-trained to solve the auxiliary task via the task-specific loss and data:

to train $\bm{\theta}_{sh}$ and $\bm{\theta}_{i}$. Afterwards, the shared parameters $\bm{\theta}_{sh}$ are transferred to initialise the segmentation model in the fine-tuning step, where the model is trained to solve the segmentation task via

In the fine-tuning step, the pre-trained $\bm{\theta}_{sh}$ are continuously trained along with randomly initialised $\bm{\theta}_{seg}$.

After the two models are trained, we select the better one based on the validation performance evaluated by the validation set $\mathbb{D}_{seg}^{val}$. By selecting the most appropriate training mode for leveraging each auxiliary task, the chance of negative transfer is reduced. Also, because each auxiliary task takes effect independently, the conflicts between auxiliary tasks are avoided. Another advantage of independent training is that this step can be naturally distributed to different machines to accelerate training.

This stage aims to compress those models obtained by the previous stage, denoted as $\{f_{i}\}_{i=1}^{N}$ to a single model $f_{s}$ for inference. In this way, the diverse knowledge obtained from multiple auxiliary tasks is aggregated to obtain an even more powerful model. To achieve this, we employ an ensemble knowledge distillation [7]. The previously obtained models are treated as teachers to distil their knowledge into the student model $f_{s}$. For each input image, the average of all the teacher model’s probabilistic prediction maps, i.e., their ensemble prediction $\frac{1}{N}\sum_{i=1}^{N}f_{i}(\bm{x}_{seg})$, is used as the pseudo annotation to train the student model. Considering that the pseudo annotation inevitably carries erroneous predictions from the teacher models, only using $l^{KD}$ to supervise the student is prone to misleading. To address this issue, following common practice, we define the loss for student training as a linear combination of the loss w.r.t. the ground-truth annotation and the KD loss w.r.t. the pseudo annotation. Putting together, the total loss for training the student model is formulated as

where $\lambda^{KD}$ is a coefficient balancing the impact of ground-truth annotations and pseudo annotations. We set $\lambda^{KD}=1/(N+1)$ intuitively to treat every pseudo annotation from a teacher model as equally important as the ground-truth annotation.

PNE (SIIM-ACR Pneumothorax Segmentation Challenge) [12] is the largest public dataset for pneumothorax segmentation. The images take a subset of the ChestX-ray14 dataset [27] released by the National Institutes of Health (NIH), which consists of both positive and negative pneumothorax cases. In the positive cases, pixel-wise delineations of pneumothorax are provided. In this study, we take all 2669 samples with positive pneumothorax for experiments. In cases where multiple annotations are available, we merge the annotation masks by a union operation to include all the positive regions. All images are resized to $256\times 256$. We randomly split the dataset into 1601/267/801 samples for training/validation/testing.

Common setup. To simulate a scenario suffering data scarcity, we randomly select 200 samples for training, which is a common number of data accessible in the application. We train 2D U-Net [11] models with the backbone of ResNet-18 [28], which is commonly adopted for MIS tasks. To construct the U-Net architecture, we adopt group normalisation layers [29] and ReLU activation layers. For data augmentation, we use random translation, zooming, rotation, Gaussian noise, Gaussian blur, brightness jittering, contrast jittering, and gamma jittering following [1]. We use a batch size of 16 and the Dice loss function to train every segmentation model. We adopt the Rectified Adam (RAdam) optimiser [30] for model training, where the learning rate is initialised to 0.001 and decayed by the poly annealing scheduler with a power rate of 0.9 throughout the training process.

Auxiliary Tasks. We employ the following five representative auxiliary tasks, covering both segmentation annotation-based tasks (SDM-in and SDM-out) and unsupervised tasks (RKB, MoCo, and VICReg), to implement our method:

• •

  SDM-in and SDM-out: Surface distance map (SDM) is an alternative representation of segmentation mask, with a rigorous mapping between them through Euclidean distance transform. SDM prediction can help produce smooth segmentation boundaries and reduce spatially isolated errors [2]. In this study, we consider predicting two kinds of SDMs as auxiliary tasks. SDM-in is obtained by replacing the intensity of each pixel in the foreground with the distance to its closest background pixel and setting the background regions to 0. Conversely, SDM-out is obtained by replacing the intensity of each pixel in the background with the distance to its closest foreground pixel and setting the foreground regions to 0.

• •

  RKB (Rubik’s Cube) [8] is a predictive SSL task that trains a model to predict the correct order of a set of shuffled image patches cropped from the same image.

• •

  MoCo (Momentum Contrast-v3) [9] is a contrastive SSL task that trains a model to differentiate the image patches cropped from different images. The task is designed to contrast the latent feature representations of different image patches so that the features of patches cropped from the same image (positive pair) are pulled together, and the features of patches cropped from different images (negative pair) are pushed apart.

• •

  VICReg (Variance-Invariance-Covariance Regularization) [10] is another contrastive SSL task. It has a similar goal as MoCo but employs a different contrastive learning formulation via a variance loss term and an invariance loss term. In addition, it promotes feature diversity by a covariance loss term.

For each auxiliary task, we independently use it to conduct two experiments, each using the auxiliary task to facilitate training a segmentation model with either joint training or pre-training following the respective procedures described in Section III-C. The experimental results are reported in Table I. It can be observed that different auxiliary tasks work well with different training modes. SDM-in and SDM-out work well only via joint training while drastically degrading the segmentation performance with pre-training. By contrast, RKB, MoCo, and VICReg work better with pre-training while offering no or limited improvements when using joint training. These results suggest the appropriate training mode differs for each auxiliary task, and no single training mode works well for all the auxiliary tasks. Furthermore, to take full advantage of different auxiliary tasks, both joint training and pre-training modes need to be supported, which motivates the design of our proposed method.

This section compares the proposed method with several existing methods that leverage multiple auxiliary tasks to boost model performance. We compare the following methods, including intuitive methods for multi-task learning or pre-training, and existing methods proposed for leveraging multiple auxiliary tasks to boost model performance under scarce training data:

• •

  Joint Training is the most commonly used way to leverage auxiliary tasks. The segmentation task and the auxiliary tasks are used together to train the model via a weighted combination of their respective losses. Finally, the segmentation model is returned.

• •

  Multi-Task Pre-Training is a common way to perform model unsupervised pre-training with multi-task SSL, which is adapted to our supervised setting. It first pre-trains a backbone model on all the auxiliary tasks using Joint Training. Then, the pre-trained backbone is transferred to initialise the target segmentation model, which is fine-tuned on the segmentation dataset.

• •

  Ensemble is an intuitive way to aggregate the knowledge learnt from multiple auxiliary tasks. Similar to the first stage of the proposed method, it first uses each auxiliary task independently to facilitate the training of a segmentation model. Then, it directly uses the ensemble of all the obtained models for inference.

• •

  GCS (Gradient Cosine Similarity) [4] is an existing method that utilizes the similarity between the gradients from each auxiliary task and those from the primary task to determine whether to learn from each auxiliary task at every iteration step.

• •

  PCGrad (Projecting Conflicting Gradients) [31] is an MTL method adapted to our setting. At every iteration step, for each auxiliary task, the gradients that conflict with the segmentation task are removed.

• •

  OL-AUX (Online Learning for Auxiliary losses) [5] adaptively changes the loss weight based on the gradient inner product w.r.t the main task to decrease the long-term value of the main task loss.

• •

  AMAL (Adaptive mixing of Auxiliary Losses) [32] adaptively changes the loss weight based on gradient feedback from the validation set.

The results are reported in Table II. It can be seen that existing methods based on the joint training paradigm (all except for Multi-Task Pre-Training, Ensemble, and ours) can hardly integrate the advantages of multiple auxiliary tasks, evidenced by their equal or marginally better performance compared to the baselines of conventional training and joint training with a single auxiliary task. Among these methods, only PCGrad, which removes the gradients hampering solving the segmentation task can surpass all these baselines. This implies a severe interference among the learning heterogeneous tasks, which cannot be well tackled by certain designs trying to filter out negative transfer (GCS and PCGrad) or adaptively integrate the contributions of different auxiliary tasks (AMAL) within the joint training paradigm. Also, these compared methods cannot match the performance obtained by pre-training with a single auxiliary task of MoCo or RKB, showing the shortcoming of joint training to leverage these tasks. By contrast, those methods that determine the training mode for each auxiliary task (Ensemble and ours) significantly improve the performance over all the baselines and compared methods. This demonstrates the necessity to adaptively determine the training mode for different auxiliary tasks and the superiority of using each auxiliary task independently to avoid interference. In addition, our method outperforms Ensemble with a single model, indicating the effectiveness of using ensemble knowledge distillation to aggregate the knowledge from various auxiliary tasks.

We probe the number of training data that can be saved by our method. We train segmentation models using conventional training with different numbers of data, ranging from 200 to 1600 and report the corresponding segmentation performance in Table III. It can be seen that the performance of our method with only 200 training samples matches that obtained with 400 - 500 training samples with conventional training, indicating that our method can effectively boost the segmentation performance when the training data is limited.

As shown in Table IV, we start by using all 5 auxiliary tasks as a baseline and ablate each auxiliary task one by one to obtain the results with 4 auxiliary tasks. Then the best performer with 4 auxiliary tasks is set as the new baseline to further ablate an auxiliary task to conduct the experiments involving 3 auxiliary tasks. This process is repeated until only 2 auxiliary tasks are left. Finally, the results obtained using a single auxiliary task in our method are listed for reference. It can be observed from the results that the performance obtained using all 5 tasks is among the best performers, whereas using less than 5 auxiliary tasks cannot guarantee optimal performance. Besides, no single auxiliary task alone can obtain good performance. These results indicate the effectiveness of using all the auxiliary tasks to integrate their complementary advantages and ensure good performance. It is also shown that certain auxiliary tasks (RKB and MoCo in our case) play more important roles than the others, evidenced by those inferior results obtained without these two tasks.

We analyse the impact of the coefficient in the KD loss $\lambda^{KD}$ by setting it to certain values within $[0,1]$. $\lambda^{KD}=0$ corresponds to only using the segmentation loss for training, and $\lambda^{KD}=1$ corresponds to only using the distillation loss. Recall that our method is implemented by setting $\lambda^{KD}=1/(N+1)=0.83$. As shown in Table V, the best result is obtained at $\lambda^{KD}=0.7$, and our setting achieves the second best. The inferior performance obtained by a $\lambda^{KD}$ of 0 or 1 indicates the importance of incorporating the supervision from both the ground-truth annotations and pseudo annotations for good performance. The results also show that the KD method in Stage 2 is not sensitive to the choice of $\lambda^{KD}$ as values within $[0.6,0.9]$ yield similar results.