Marginal loss and exclusion loss for partially supervised multi-organ segmentation

Many pioneering works have been done on multi-organ segmentation, using traditional machine learning methods or deep learning methods. In [30, 48, 41, 38, 36, 45], a multi-altas based strategy is used for segmentation, which registers an unseen test image with multiple training images and use the registration map to propagate the labels in the training images to generate final segmentation. However, its performance is limited by image registration quality. In [17, 7, 5], prior knowledge of statistical models is employed to achieve multi-organ segmentation. There are also some methods that directly use deep learning semantic segmentation networks for multi-organ segmentation [11, 43, 20, 22]. Besides, there are prior approaches that combine the above-mentioned different methods [6, 29] to achieve better multi-organ segmentation. However, all these methods rely on the availability of fully labelled images.

Very limited works have been done on medical image segmentation with partially-supervised learning. Zhou et al. [52] learns a segmentation model in the case of partial labeling by adding a prior-aware loss in the learning objective to match the distribution between the unlabeled and labeled datasets. However, it trains separate models for the fully labeled and partially labeled datasets, and hence involves extra memory and time consumption. Instead, our work trains a single multi-class network. Since only two loss terms are added, it needs nearly no additional training time and memory cost. Dmitriev et al.  [9] propose a unified, highly efficient segmentation framework for robust simultaneous learning of multi-class datasets with missing labels. But the network can only learn from datasets with single-class labels. Fang et al.  [10] hierarchically incorporate multi-scale features at various depths for image segmentation, further develop a unified segmentation strategy to train three separate datasets together, and finally achieve multi-organ segmentation by learning from the union of partially labeled and fully labeled datasets. Though this paper also uses a loss function that amounts to our marginal cross entropy, its main focus is on proposing the hierarchical network architecture. In contrast, we concentrate on studying the impact of the marginal loss including both marginal cross entropy and marginal Dice loss. Furthermore, it is worth mentioning that none of the above works considers the mutual exclusiveness, a well-known attribute between different organs. We propose a novel exclusion loss term, exploiting the fact that organs are mutually exclusive and adding the exclusiveness as prior knowledge on each image pixel.

A few existing methods have been developed on classification and object detection tasks using partially annotated data. Yu et al. [50] propose an empirical risk minimization framework to solve multi-label classification problem with missing labels; Wu et al.[46] train a classifier with multi-label learning with missing labels to improve object detection problem. Cour et al. [8] propose a convex learning formulation based on the minimization of a loss function appropriate for the partially labeled setting. Besides, as far as semi-supervised learning is concerned, a number of researches have been developed to solve  [15, 53, 47] classification problems or detection problems in the absence of annotations.

The loss function is generally proposed for a specific problem. A common idea for loss functions are based on classification tasks which optimize the intra-class difference and reduce the intra-class variation, for example contrastive loss  [13], triplet Loss  [35], center loss [44], large margin softmax loss [25], angular softmax [23] and cosine embeding loss [42]. The cross entropy loss [28] is the most representative loss function, which is commonly used in deep learning. There are also some loss functions designed to optimize the global performance for semantic segmentation, such as Dice loss  [28], Tversky loss  [33], combo loss [40], Lovasz-Softmax loss [1]. Besides, some losses are proposed specifically to improve a given loss function, for example, the focal loss [24] is developed based on cross-entropy loss [28] to better solve class imbalance problem. Here we focus on the cross-entropy loss and regular Dice loss that are most commonly used in multi-organ segmentation.

Suppose that, for a multi-class classification task with $N$ labels with its label index set as $\Omega_{N}=\{C_{1},C_{2},\ldots,C_{N}\}$, its data sample $x$ (i.e., an image pixel in image segmentation) belongs to one of $N$ classes, say class $C_{n}$, which is encoded as an $N$-dimensional one-hot vector ${\hat{y}}_{n}=[y_{1},y_{2},\ldots,y_{N}]$ with $y_{n}=1$ and all others $0$. A multi-class classifier consists of a set of response functions $\{a_{n}(x);~{}n\in\Omega_{N}\}$, which constitutes the outputs of the segmentation network $\Psi$. From these response functions, the posterior classification probabilities are computed by a softmax function,

To learn the classifier, the regular cross-entropy loss is often used, which is defined as follows:

Besides, the Dice score coefficient (DSC) is often used, which measures the overlap between the segmentation map and ground truth. The dice loss is defined as $1-DSC$:

For an image with incomplete segmentation label, it is possible that the pixels for some given classes are not ‘properly’ provided. To deal with such situations, we assume that there are a reduced number of $M<N$ classes in a partially-labeled dataset with its corresponding label index set as $\Omega^{\prime}_{M}=\{C^{\prime}_{1},C^{\prime}_{2},\ldots,C^{\prime}_{M}\}$. For each merged class label $m\in\Omega^{\prime}_{M}$, there is a corresponding subset ${\Phi_{m}}\subset\Omega_{N}$, which is comprised of all the label indexes in $\Omega_{N}$ that can be merged into the same class $m$. Because the labels are exclusive in multi-organ segmentation, we have $\Omega_{N}=\cup_{m\in\Omega^{\prime}_{M}}{\Phi_{m}}$.

Fig.2 illustrates the process of label merging, using an example of four organ classes $C_{i},i=1,2,3,4$. After the merging, there are two classes $C^{\prime}_{1}$ and $C^{\prime}_{2}$, with $C_{1}$ and $C_{2}$ are combined together to form a new merged label $C^{\prime}_{1}$ and $C_{3}$ and $C_{4}$ to form a new label $C^{\prime}_{2}$.

The classification probability for the merged class $m$ is a marginal probability

Also, the one-hot vector for a class $m\in\Omega^{\prime}_{M}$ is denoted as ${\hat{z}}_{m}=[z_{1},z_{2},\ldots,z_{M}]$, which is $M$-dimensional with $z_{m}=1$ and all others $0$.

Consequently, we define marginal cross-entropy loss and marginal Dice loss as follows:

We use marginal cross entropy as an example to perform the gradient calculation. Firstly, referring to Eqs. (1) and (4), the gradient of the output $m$ of a softmax node to the network node $a_{j}$ is:

where $\pi_{j}(\Phi_{m})$ is a boolean indicator function that tells if $j$ is in $\Phi_{m}$. $p$ and $q$ are the classification probabilities of regular and marginal softmax functions.The derivative gradient of $L_{mCE}$ to the network node $a_{j}$ is:

where ${\bar{m}}$ is the only class index that makes $\pi_{j}(\Phi_{m})=0$.

It happens in multi-organ segmentation tasks that some classes are mutually exclusive to each other. The exclusion loss is designed to add the exclusiveness as an additional prior knowledge on each image pixel. We define an exclusion subset for a class $n$ as $E_{n}$, which comprises all (or a part of) the label indexes that are mutually exclusive with class $n$. The exclusion label information is encoded as an N-dimensional vector ${\hat{e}}_{n}=[e_{1},e_{2},\ldots,e_{N}]$, which is obtained as:

Note that ${\hat{e}}_{n}$ is still an $N$-dimensional vector, but it is not an one-hot vector any more. Fig.3 shows the procedure of applying exclusion loss. Assuming that organ classes $C_{1}$, $C_{2}$ and $C_{3}$ are mutually exclusive, the labels of $C_{2}$ and $C_{3}$ form the exclusion subset $E_{C_{1}}$.

We expect that the intersection between the segmentation prediction $p_{n}$ from the network and $e_{n}$ is as small as possible. Following the Dice coefficient, the formula for the exclusion Dice loss is given as:

The exclusion cross-entropy loss is defined accordingly:

where $\epsilon$ is introduced to avoid the trap of $-\infty$. We set $\epsilon=1$.

We consider a partially-supervised multi-organ segmentation task that is common in practice (such as Fig. 1). For each partially annotated image, we restrict it with only one label. For clarity of description, we assume that $F$ denotes the fully-labeled segmentation dataset and $P_{i};i\in\{1,2,\ldots,C\}$ denotes a dataset of partially-annotated images that contain only a partial list of organ label(s). The datasets $P_{1:C}$ do not overlap in terms of their organ labels. For an image in $P_{i}$, there is a ‘merged’ background, which is the union of real background and missing organ labels. We jointly learn a single segmentation network $\Psi$ using $F\cup P_{1}\cup...\cup P_{C}$, assisted by the proposed loss functions.

For our experiments, we choose liver, spleen, pancreas, left kidney and right kidney as the segmentation targets and use the following benchmark datasets.

• 1.

  Dataset $F$. We use Multi-Atlas Labeling Beyond the Cranial Vault - Workshop and Challenge [21] as fully annotated base dataset $F$. It is composed of 30 CT images with segmentation labels of 13 organs, including liver, spleen, right kidney, left kidney, pancreas, and other organs (gallbladder, esophagus, stomach, aorta, inferior vena cava, portal vein and splenic vein, right adrenal gland, and left adrenal gland) we hereby ignore.

• 2.

  Dataset $P_{1}$. We refer to the task03 liver dataset from the Decathlon-10 [37] challenge as $P_{1}$. It is composed of 130 CT’s with annotations for liver and liver cancer region. We merge the cancer label into the liver label and obtain a binary-class (liver vs background) dataset.

• 3.

  Dataset $P_{2}$. We refer to the task09 spleen dataset from the Decathlon-10 challenge as $P_{2}$. It includes 41 CT’s with spleen segmentation label.

• 4.

  Dataset $P_{3}$. We refer to the task07 pancreas dataset from the Decathlon-10 challenge as $P_{3}$. It includes 281 CT’s with pancreas and its cancer segmentation label. The cancer label is merged into the pancreas label to obtain a binary-class (pancreas vs background) dataset.

• 5.

  Dataset $P_{4}$. We refer to KiTS [18] challenge dataset as $P_{4}$. Since the offered 210 CT segmentation makes no distinction between left and right kidneys, we manually divide it into left and right kidneys according to the connected component. Cancer label is merged into the according kidney label.

The spatial resolution of all these datasets are resampled to $(1.5\times 1.5\times 3)mm^{3}$. We split the datesets into training and testing. we randomly choose 6 samples from $F$, 26 samples from $P_{1}$ and 8 samples from $P_{2}$, 56 samples from $P_{3}$ and 42 samples from $P_{4}$ as testing. The others are used for training. Table 2 also provides a summary description of the datasets.

We set up the training of 10 deep segmentation networks for comparison as in Table 1.

• 1.

  $\Psi^{mc}_{F}$: a multiclass segmentation network based on $F$.

• 2.

  $\Psi^{b}_{P_{1}}$: a binary segmentation network for liver only based on $P_{1}$.

• 3.

  $\Psi^{b}_{P_{2}}$: a binary segmentation network for spleen only based on $P_{2}$.

• 4.

  $\Psi^{b}_{P_{3}}$: a binary segmentation network for pancreas only based on $P_{3}$.

• 5.

  $\Psi^{t}_{P_{4}}$: a ternary segmentation network for left kidney and right kidney only based on $P_{4}$.

• 6.

  $\Psi^{b}_{F+P_{1}}$: a binary segmentation network for liver only based on $F$ and $P_{1}$. Note that the spleen, pancreas, left kidney and right kidney labels in $F$ are merged into background.

• 7.

  $\Psi^{b}_{F+P_{2}}$: a binary segmentation network for spleen only based on $F$ and $P_{2}$. Note that the liver, pancreas, left kidney and right kidney labels in $F$ are merged into background.

• 8.

  $\Psi^{b}_{F+P_{3}}$: a binary segmentation network for pancreas only based on $F$ and $P_{3}$. Note that the liver, spleen, left kidney and right kidney labels in $F$ are merged into background.

• 9.

  $\Psi^{t}_{F+P_{4}}$: a ternary segmentation network for left kidney and right kidney only based on $F$ and $P_{4}$. Note that the liver, spleen, pancreas labels in $F$ are merged into background.

• 10.

  $\Psi^{mc}_{All}$: a multi-class segmentation network based on $F$, $P_{1}$, $P_{2}$, $P_{3}$ and $P_{4}$.

For training the above networks except $\Psi^{mc}_{All}$, we use the regular CE loss, regular Dice loss, and their combination. For training the network $\Psi^{mc}_{All}$, when involves partial labels we need to invoke the marginal CE loss, marginal Dice loss, and their combination. Further, for $\Psi^{mc}_{All}$ we experiment the use of exclusion Dice loss and exclusion CE loss.

Considering the impact of the varying axial resolutions of different data sets in the original CT image on the training process, we resample the 3D CT image to $(1.5\times 1.5\times 3)mm^{3}$ and then extract the patch with the shape $[190,190,48]$ as input to illustrate the merit of our loss functions. For comparison, we use the same parameter settings in all networks; therefore there is no inference time difference among them. During training, we use 250 batches per epoch and 2 patches per batch. In order to ensure the stability of model training, we set the proportion of patches that contain foreground in each batch to be at least 33%. The initial learning rate of the network is 1e-1. Whenever the loss reduction is less than 1e-3 in consecutive 10 epochs, the learning rate decays by 20%.

We train 3D nnUNet [19] for all segmentation networks. We choose the 3D nnUNet because it is known to be a state-of-the-art segmentation network. While there are other network architectures [10] that might achieve comparable performance, we expect similar empirical observations from our ablation studies even based on the other networks.

For the network $\Psi^{mc}_{All}$, we train it in two stages in order to prevent the instability caused by large loss value at the beginning of the training. In the first stage, we only use the fully annotated dataset $F$. The goal is to minimize the regular loss function using the Adam optimizer. The purpose of the first phase is to give the network an initial weight on multi-class segmentation in order to prevent the large loss value when applying the marginal loss functions. In the second stage, each epoch is trained jointly using the union of five datasets. In each epoch, we randomly select 500 patches from each training dataset with a batch size of 2. Depending on the source of the slice, we use either the regular loss, if from $F$, or the marginal loss and the exclusion loss, if from $P_{i}(i\in\{1,2,3,4\})$. In actual experiment, the first stage consists of 120 epochs and the second stage 80 epochs.

We use two standard metrics for gauging the performance of a segmentation method: Dice coefficient and Hausdorff distance (HD). A higher Dice coefficient or a lower HD means a better segmentation result. Table 3 shows the mean and standard deviation of Dice coefficients of the results obtained by the deep segmentation networks under different loss combinations and with different dataset usages, from which we make the following observations.

The effect of pooling together more data. The experimental results obtained by the models jointly trained from combinations of the datasets $F$ and $P_{i}(i\in\{1,2,3,4\})$ are generally better than those by the models trained from a single labeled dataset alone. As shown in Table 3 and Table 4, when comparing the performance of $\Psi^{b}_{F+P_{i}}$ vs $\Psi^{b}_{P_{i}}(i\in\{1,2,3,4\})$, the former generally outperforms the latter. For example, when using rCE+rDC as the loss, the mean Dice coefficient is boosted from .851 to .900 (the according HD is reduced by 37.5%). When comparing the performance of $\Psi^{mc}_{F+P_{i}}(i\in\{1,2,3,4\})$ vs $\Psi^{mc}_{F}$, again the former is better than the latter, the mean dice coefficient is increased from .874 to .900 (the according HD is reduced by 28.7%).

The importance of CE and Dice losses. When comparing the importance of CE and Dice losses, in general, it is inconclusive which one is better, depending on the setup. For example, the Dice loss works better on liver segmentation while the CE loss significantly outperforms the Dice loss on left kidney segmentation. Also fusing CE and Dice losses is in general beneficial in terms of our results as it usually brings a gain in segmentation performance. For example, when using $\Psi^{b}_{F}$, the average dice loss reaches .874 for rCE+rDC, while that for rCE and rDC is .846 and .826, respectively.