Improving Calibration and Out-of-Distribution Detection in Deep Models for Medical Image Segmentation

A. Training procedures for CNN-based medical image segmentation models

Large labeled datasets are an essential requirement for deep learning [8]. Since such datasets are difficult to come by in medical imaging, many strategies have been proposed to tackle this limitation. We briefly review some of these methods.

In transfer learning [9], the model is first trained on another domain/task and then fine-tuned for the target task. Transfer learning has been reported to improve the performance of CNN-based models for medical image segmentation [10], [11]. A limitation of transfer learning is that most public image datasets include 2D images, whereas medical images are 3D. Semi-supervised constitute a large and diverse body of techniques [4]. These methods utilize a mix of labeled and unlabeled data or data that have not been labeled in detail. These methods have been used in deep learning-based medical image segmentation [12]. One can use less accurate automatic methods to generate approximate segmentations on large corpora of images and use those to train a more accurate CNN-based model [13]. The potential benefits of semi-supervised methods strongly depend on the nature of the task and training data [14]. It may be easier to obtain rough segmentations on large datasets. Rather than treating such approximate segmentation labels as ground truth, one can use more intelligent methods [5]. Studies have reported successful applications of such methods [15].

An alternative to transfer learning is multi-task learning [9]. Multi-task learning aims at learning multiple tasks simultaneously. Studies have reported successful application of multi-task learning for medical image segmentation [16]. Training on mixed and heterogeneous data for medical image segmentation has been addressed in many prior works. These works appear under different titles such as domain adaptation, domain generalization, multi-site training, and joint training [17]. Many of these studies address domain shift, whereby a model trained on a source domain dataset fails to perform accurately and reliably on a target domain dataset. This is an important consideration in clinical settings [18], [19]. Another aim is to remove or reduce the need for labeled data in the target domain [20]-[22]. To deal with the changing image intensity and appearance between source and target domains, a wide range of solutions have been proposed [23]-[25]. Some studies incorporate domain-invariant prior information as well as unlabeled and weakly-labeled images [26]. Another example is the Synergistic Image and Feature Adaptation framework, which encourages image similarity and invariance of the learned features across domains [27]. Some studies resort to data simulation methods to improve model generalizability [28]. Federated learning methods are another class of related techniques, where the focus is on training using multi-site data while protecting institutional and patient privacy [29]. The above studies on multi-task learning have focused on segmentation accuracy. The effect of multi-task training on confidence calibration has not been explored in previous works.

B. Model calibration and uncertainty estimation

All machine learning models are bound to make wrong predictions. One would like the confidence of a model’s predictions to be proportional to the probability of making a correct prediction. Suppose that for test sample $x_i$ the model predicts the class ${\widehat{y}}_i$ with a probability ${\widehat{p}}_i$. In the ideal scenario with perfect confidence calibration, $\mathrm{P}(\widehat{y}=y\mid \widehat{p}=p)=p$ [30]. Standard deep learning models are poorly calibrated [6]. This is concerning for safety-critical applications such as medicine. Many methods have been proposed for improving the calibration of deep learning models. It has been shown that calibration can be improved by using a proper scoring rule as the loss function [6], [31], weight decay, avoiding batch normalization [6], training on adversarial examples [31], and Platt scaling [6]. For image classification, a more more accurate method was proposed by combining Platt scaling with histogram binning [32]. Another study proposed to train a separate model to map the un-calibrated output of a CNN to calibrated probabilities [33].

Another study showed that Mixup training, [34], can improve the model calibration in image classification [35]. Mixup was originally proposed as a data augmentation and regularization strategy [34]. In brief, it synthesizes additional training data using convex combinations of pairs of training data points and their labels. Thulasidasan et al. show that Mixup also improves the model’s confidence calibration in image classification [35].

For medical image segmentation, studies have proposed methods to estimate the prediction uncertainty [36] or to use the uncertainty for improving the segmentation accuracy [37]. However, little attention has been paid to methods for improving the calibration of CNN-based segmentation models. An example of the latter is [38], where authors show that the average prediction of a model ensemble is better calibrated than that of a single model. However, their method requires training as many as 50 models, which is inefficient.

C. Detecting out-of-distribution data and model failure

Another important problem is detection of OOD test data. A central assumption of every machine learning method is that the training and test data come from the same distribution. When a test sample comes from an entirely different distribution, the model should include a mechanism to detect the OOD sample. This is challenging with deep learning models because of their black-box nature and the complex mapping between their input and output. Advancements in network design have not made deep learning models more robust to OOD data [39]. Different methods have been proposed to increase the robustness of deep learning models to OOD data. For example, deep learning models trained with Mixup are less prone to making over-confident predictions on OOD test data in image classification [35]. One study suggests that histogram equalization and Adversarial Logit Pairing improve robustness to corrupted data [40]. However, the latter study noted that methods that work well on some datasets may fail on other datasets. More importantly, these studies focus on in-distribution data that have been perturbed, and do not address the true OOD data, on which the model is expected to fail.

Several studies have proposed methods for OOD detection in deep learning. One work proposed a simple method that used the statistical distribution of the softmax values for detecting OOD data [41]. It showed that the maximum softmax value was higher for in-distribution and correctly-classified data samples than for OOD data samples. More recent studies have improved upon this basic method [7], [42]. One study showed that higher OOD detection accuracies could be achieved by 1) temperature scaling and 2) input perturbation [42]. The proposed image perturbation pushes the image in the direction that increases its predicted class. Authors show that this way of perturbing the input images increases the separation between in-distribution and OOD images. Although this operation pushes both in-distribution and OOD images closer to their predicted class, the shift is larger for in-distribution data samples than it is for OOD data samples. For image classification, one study trained Gaussian discriminant models on the penultimate network layer and used the Mahalanobis Distance to detect OOD data samples [7]. The authors extended their method by carrying out the same computation on all network layers. A different solution, again for image classification, has been proposed in [43]. In addition to the prediction head, a confidence head is added to the end of the CNN. The model directly estimates the prediction confidence as a scalar value. During training, the model is allowed to obtain a “hint" about the true class by using a weighted average of the true class and the model-predicted class in computing the prediction loss. However, the model is penalized for low-confidence predictions. At test time, OOD data samples are detected by thresholding the confidence score.

Deep k-Nearest Neighbors (DkNN) detects OOD test samples based on nonconformity of their representations with the representations of the training set [44]. Given a test sample, $x_{\text{test}}$, DkNN finds its $k$ nearest neighbors in the training set in terms of representations in each of the network layers, $l\in 1\dots L$, and records the labels of those neighbors, ${\Omega }_l$. Nonconformity of $x_{\text{test}}$ with label $j$ is defined as $\alpha (x_{\text{test}},j)=\sum \begin{array}{c}\hfill \\ l\hfill \end{array}\mid i\in {\Omega }_l:i\ne j\mid$. DkNN also computes and stores nonconformity values for a calibration set, which is separate from the training set. Let us denote the set of nonconformity values of the calibration set with $A$. Then for a test sample $x_{\text{test}}$ and label $j$, DkNN computes $p(x_{\text{test}},j)=\mid \mid \{\alpha \in A:\alpha >\alpha (x\text{test},j)\}\mid ∕\mid A\mid$. The largest $p(x_{\text{test}},j)$ across all labels is referred to as prediction credibility. It quantifies the degree of support from the training set for the model prediction and is proposed as a measure of prediction confidence. We refer the reader to [44] for details.

Most methods proposed in prior works have been devised for 2D image classification settings and cannot be used for OOD detection in 3D medical image segmentation. This is because, in volumetric image segmentation, feature maps are extremely large and typically only tens of training images are available. Some studies have proposed to detect OOD data based on prediction uncertainty [45]. However, such methods are not accurate in semantic segmentation applications [46]. In fact, compared with image classification, OOD detection in image segmentation has received much less attention. A recent study found that methods proposed for OOD detection in image classification do not translate well to segmentation tasks [47]. For segmentation of street view images, one study proposed a dedicated neural network to detect OOD data samples [46]. Their method classified an image as in-distribution or OOD using a very large “background dataset" to represent the variety of scenes outside of the training data distribution. However, it is difficult to obtain or even define the background set in medical imaging. The authors of [48] have proposed a method to automatically generate OOD data using a generative adversarial network (GAN) model that is trained in parallel with the main image classification network. One study evaluated the performance of several of the state of the art OOD detection methods for classification of 2D medical images and found that no single method achieved consistently-satisfying results, especially on samples that were closer to the in-distribution data [49]. Another study used prediction uncertainty measures to identify OOD data in medical image segmentation [38]. However, they evaluated their method on data that were hard to segment, not on true OOD data. As shown below by our results, methods based on prediction uncertainty cannot accurately detect OOD data in medical image segmentation.

It is worth mentioning that a related topic to OOD detection is the topic of adversarial examples [50]. These are examples that are crafted to fool a model into making wrong predictions. Adversarial examples may be important in some applications, but they are beyond the scope of this paper, which focuses on natural OOD data samples. By “natural OOD" we mean OOD data that exist due to such factors as a change in subject age, scanning protocol, or similar factors, as opposed to artificially-crafted OOD data created by an adversary.

D. Contributions of this work

In this paper, we address the problems outlined above and make the following contributions.

• We show that multi-task learning can improve the confidence calibration of CNN models for medical image segmentation. Through extensive experiments on a diverse collection of medical image segmentation datasets, we show that the confidence calibration of deep learning-based medical image segmentation models improves with multi-task learning. Importantly, in general, multi-task learning does not negatively impact the segmentation accuracy, and on some datasets it may slightly improve accuracy as well, even compared with other well-known training strategies such as transfer learning.

• We propose a novel and accurate OOD detection method for CNN-based medical image segmentation. Our method is based on spectral analysis of CNN feature maps. We show that whereas methods based on prediction uncertainty or image classification are inaccurate for OOD detection in 3D medical image segmentation, our proposed method accurately detects OOD test data in different experimental settings.

A. Data

Table I summarizes the information about the datasets used in this work. The numbers of images in our datasets range from 15 to 400, which are typical of manually-labeled datasets in medical image segmentation, especially for in-house datasets because manual segmentation of complex 3D structures such as the brain cortical plate is time-consuming. Given the small size of the datasets, unless otherwise specified, we used a 3-fold cross-validation strategy. This way, all data are used for test, thereby increasing the power of the statistical significance tests. We always used a patient-wise data split. Computed Tomography (CT) images were normalized by mapping the Hounsfield Unit values in the range [−1000, 1000] to intensity range [0, 1]. Magnetic Resonance (MR) images were normalized by dividing each image by the standard deviation of voxel intensities. The first column in Table I shows the names that we use to refer to each dataset throughout this paper.

B. Network architecture and training details

The CNN architecture used in this work was based on the 3D U-Net [56], which we substantially modified by adding residual connections with short and long skip connections. The skip connections connect every fine feature map directly to all coarser feature maps in the encoder section of the network, similar to the DenseNet [57]. We set the number of features in the first stage of the encoder part of our network to 14, which was the largest allowed by our GPU memory. The model worked on 96³-voxel image blocks. During training, we sampled blocks from random locations in the training images, which acts as a form of data augmentation. Other data augmentation methods that we used during training included random flips and rotations (by integer multiples of π/2) in all directions as well as addition of random Gaussian noise to voxel intensity values. We also experimented with elastic deformation for data augmentation, [58], but we did not pursue that augmentation method because it negatively impacted segmentation of fine structures such as brain cortical plate. On a test image, a sliding window approach with a 24-voxel overlap between adjacent blocks was used to process the image. We used the negative of the Dice Similarity Coefficient (DSC) as the loss function and Adam [59] as the optimization method. We used an initial learning rate of 10⁻⁴, which was reduced by 0.90 after every 2000 training iterations if the loss did not decrease. If the loss did not decrease for two consecutive evaluations, we stopped the training. Since the focus of the study is on model calibration and OOD detection, we used the same settings mentioned above in all experiments.

C. Multi-task learning

In this work we advocate for training on heterogeneous data, which we refer to as “multi-task learning", although multi-task learning has also been used in other settings as we have explained in Section II-A. We train a single model on a mix of training datasets that can come from different imaging modalities with different organs of interest to be segmented. We do not use additional inputs to inform the model of the image modality or the organ that needs to be segmented. Furthermore, the network will still have a single segmentation head (i.e., single output layer). Given the input image from any modality and organ of interest, the network will output the predicted segmentation map of the organ of interest on the segmentation head. Given the GPU memory limit, we use a training batch size of one. For multi-task learning, we sample a block from one of the images in the training set and use that block and its corresponding ground truth segmentation map to update the model parameters. In other words, no changes are made to the network architecture and overall training procedures compared with training on a single dataset. The only point worth mentioning is the frequency of sampling from different training datasets. We sample from each dataset with a probability proportional to the inverse of the square root of dataset size, $1∕\sqrt{n}$. This way, if for example we train on two datasets with 10 and 100 images each, the probability of sampling an image from these two datasets will be 0.24 and 0.76, respectively ($\frac{\sqrt{10}}{\sqrt{10}+\sqrt{100}}=0.24$ and $\frac{\sqrt{100}}{\sqrt{10}+\sqrt{100}}=0.76$). We found that, especially on datasets with fewer images, using $1∕\sqrt{n}$ resulted in higher accuracy than using $1∕n$, which is equivalent to uniform sampling.

D. OOD detection using spectral analysis of feature maps

We propose a novel method for detecting OOD test data for CNN-based medical image segmentation models. As mentioned above, these models produce over-confident predictions even when a test sample is entirely outside the distribution of the training data. We show in Section IV that methods based on prediction uncertainty are unable to accurately detect OOD data. Moreover, because of the very large size of 3D medical images and their computed feature maps and small number of training images, methods based on analyzing the feature maps in their native space are ineffective. This is because these methods have been developed for scenarios where the size of the feature vector is on the order of hundreds or a few thousands and millions of training samples are available. We described several of these methods in Section II-C above. In Section IV, we show that such methods are not accurate for OOD detection in medical image segmentation applications. Instead, we propose computing the spectrum of the feature maps, which we define as the vector of singular values computed using singular value decomposition (SVD). Consider a test image $x_i$ and denote the feature map computed for this image at a certain stage (i.e., layer) of the network with $F_i\in {\mathbb{R}}^{w,h,d,n}$, where $w,h,d$ denote the dimensions of the feature map and $n$ is the number of features. We reshape $F_i$ as ${\mathbb{R}}^{whd,n}$ and compute the SVD of $F_i$ as $F_i=USV$, where $U$ and $V$ are orthonormal matrices and the diagonal matrix $S$ contains the singular values of $F_i$, which is referred to as its spectrum [60]. The vector of singular values, $s=\text{diag}(S)$, depends on the magnitude of the feature values, which in turn depend on the image voxel intensities. Moreover, the spectrum has a very large dynamic range. To eliminate these effects, we take the logarithm of the spectrum $s$ and then normalize it so that it has an ${\ell }_2$ norm of unity. We refer to the normalized logarithmic spectrum of the feature maps computed as explained above as “the spectral signature". We still denote this spectral signature with $s$ in the following.

Figure 2 shows example signatures, where a model trained on several datasets from Table I including CP-younger fetus and Liver-MRI-SPIR datasets but not including Pancreas and Hippocampus. The figure shows example spectral signatures of training images from these four datasets. Each dataset has a distinct spectral signature. Note that this model segments CP-younger fetus and Liver-MRI-SPIR accurately, but fails on Pancreas and Hippocampus, which have not been seen during training. Nonetheless, as we show in Section IV, methods based on prediction uncertainty cannot detect these as OOD.

We propose detecting ooD data based on the similarity of spectral signatures. We compute the spectral signatures of all training images and store them as $S_{\text{train}}$. Given a test image, $x_{\text{test}}^i$, we compute its spectral signature $s_{\text{test}}^i$. We define Out-Of-Distribution Measure (OODM) of $x_{\text{test}}^i$ as the Euclidean distance of its spectral signature, $s_{\text{test}}^i$, to its nearest neighbor in the training set:

We expect OODM to be smaller for test images coming from the distribution of the training set than for images from other distributions. We declare a test image $x_{\text{test}}^i$ to be OOD if $\text{OODM}(x_{\text{test}}^i)>\tau$. We determine the threshold $\tau$ using the training set. On the training set we compute the vector of ${\text{OODM}}_{\text{train}}$ using Eq. (1) using a leave-one-out strategy. We then compute $\tau$ as:

where we set $C=2.5$ for computing the detection accuracy. The value of $\tau$, and hence $C$, determine the trade-off between sensitivity and specificity. A larger $\tau$ reduces the false positive rate while also reducing the true positive rate. The area under the receiver-operating characteristic curve (AUC) is the standard measure that is used to characterize this trade-off [61], [62]. The receiver operating characteristic curve is the plot of the true positive rate versus the false positive rate at various settings of the detector threshold [61]. In our proposed model, we compute the AUC by changing our threshold, $\tau$.

One can apply the above method on any of the network’s feature maps. We found that using the deepest feature maps (i.e., feature maps closest to the output) led to more accurate OOD detection. This is in agreement with the fact that deeper layers provide more disentangled manifolds [63]. We applied the method on the last feature maps, which had 14 channels. Hence, the length of the spectral signatures in this work is 14. Figure 2 displays example histograms of OODM values for training data, in-distribution test data, and OOD test data, showing that OODM easily separates in-distribution from OOD data in this experiment.

We compare our proposed method with:

(a) A common method based on prediction uncertainty [45]. We train our model using dropout (with a rate of 10%) in each layer. At test time, we draw $N=10$ random dropout masks and compute the entropy of the mean of the segmentation probability maps, $H(\overline{p})=-\overline{p}\log (\overline{p})$ as the estimated voxel-wise prediction uncertainty map. We use the average of voxel-wise uncertainty on the predicted foreground as image-wise uncertainty. Similar to our proposed method, we compute a threshold similar to Eq. (2) on the training set. We refer to this method as UNC-Dropout.

(b) A method based on model ensembles [31]. Recently, this method was used in medical image segmentation in [38]. We refer to this method as UNC-Ensemble. (c) DkNN [44], explained in Section II-C We use 20% of images as the calibration set. From each image, we generate 100 data samples via augmentation. For a test image, we sample 10 blocks overlapping the predicted foreground and estimate credibility as the mean of $p(x_{\text{test}},j^{\prime })$, where $j^{\prime }$ is the most frequent label on the 10 blocks. We use credibility as a measure of confidence as suggested in [44]. This method can only be applied when the training set has more than one classes.

(d) Method of Outlier Exposure [64]. This method is based on training the model against a dataset of outliers. The proposed loss function includes an outlier exposure term that has the form ${\mathbb{E}}_{x^{\prime }\sim {𝒟}_{\text{out}}^{OE}}\left({ℒ}_{OE}(f(x^{\prime }),f(x),y\right)$, where ${𝒟}_{\text{out}}^{OE}$ is the outlier exposure dataset used during training. As suggested by [64], we use the cross-entropy, which is also the loss proposed by [48]. This method assumes the exact nature of the outlier data is unknown. Therefore, following the recommendations of [64], we used all other datasets listed in Table I. For example, if in an experiment we use “CP- younger fetus" as in-distribution and “CP- older fetus" and “CP- newborn" as OOD, we use all other 11 datasets in Table I as ${𝒟}_{\text{out}}^{OE}$. We refer to this method as Outlier Exposure.

(e) We compare with the method of [48], which is based on the inspection of softmax values as proposed by [41]. Specifically, OOD samples are detected as those for which the maximum softmax value is smaller than a threshold. However, authors of [48] propose two training strategies. First, they use an extra loss term to encourage the distribution of softmax values for OOD data samples to be close to a uniform distribution. They also use a GAN, whose task is to generate informative OOD training samples. The GAN model is trained in an alternative optimization framework in parallel with the main model. We refer to this method as Lee-2017.

(f) We compare with ODIN (Out-of-DIstribution detector for Neural networks) [42]. This method is based on the softmax values originally proposed by [41] that we have explained above. ODIN introduced two additional tricks: 1) temperature scaling, and 2) image perturbation. We have provided more detail on this method above in Section II-C.

(g) We also compare with the method proposed in [7], which followed the work of [42] and showed to be more accurate in image classification. As we briefly explained above, this method is based on the idea of a generative classifier. It assumes that the features can be modeled with class-conditional Gaussian distributions. Based on this assumption, the authors propose a confidence score based on the Mahalanobis Distance. Given the very large size of the feature maps, we use average pooling, as suggested by [7]. Since this method is based on the Mahalanobis Distance, we refer to it as Mah-Dist.

E. Evaluation metrics

We quantify segmentation accuracy using DSC, 95 percentile of the Hausdorff Distance (HD95) and Average Symmetric Surface Distance (ASSD). We assess model calibration by computing the Expected Calibration Error (ECE) and Maximum Calibration Error (MCE) [65]. ECE and MCE are computed by dividing the probability range, [0,1], into $K$ bins. Denote by $S_k$ the set of points (i.e., voxels in a semantic segmentation problem) whose predicted probabilities fall in the interval $[\frac{k-1}{K},\frac{k}{K}]$, for $k\in [1,K]$. Then ECE is defined as [6], [65]:

where $\text{acc}(S_k)$ is the average prediction accuracy for the points in $S_k$ and $\text{conf}(S_k)$ is the mean prediction probability for the points in $S_k$. MCE is defined as [6], [65]:

Values of calibration measures such as ECE can be dramatically influenced by the large percentage of background voxels. These voxels are usually correctly and confidently segmented by a relatively good model. Including these voxels in the computation of calibration measures artificially improves the values of these measures. Hence, we dilate the surface bound-ary of the ground-truth foreground (in 3D, in both directions in/out) in each image by 10 voxels and use the obtained mask for computing the calibration measures. A similar strategy was used in [38]. For OOD detection, we report accuracy, sensitivity, and specificity. Furthermore, we compute the area under the receiver-operating characteristic curve (AUC) by changing the value of $\tau$.

A. Feasibility and benefits of multi-task learning

As we mentioned above, we advocate for training a single model to segment different organs in different imaging modalities. To show that this is a viable approach, we trained a model on twelve datasets spanning various organs in MRI and cT images. We then trained twelve separate models, one on each of the datasets. We show a comparison of the test performance of these two training strategies in Table II, where the statistically significant differences have been marked using bold type. All statistical significance tests were performed using paired t-tests with a significance threshold of $p=0.01$.

The results show that in general multi-task learning improves model prediction calibration. The joint model was significantly better-calibrated than the dedicated models on 9 out of 12 datasets in terms of ECE and on 10 out of 12 datasets in terms of McE. only on the Hippocampus dataset the dedicated model was significantly better calibrated than the joint model. Hippocampus dataset included 260 images, compared with 15-32 images in eight of the other datasets used in this experiment. Therefore, this could be the influence of dataset size on the potential benefits of multi-task learning. However, there are three other large datasets in this experiment: CP-newborn ($n=400$), Pancreas ($n=281$), and KiTs ($n=300$). For Pancreas and KiTs, multi-task learning resulted in better-calibrated models in terms of all four calibration measures (ECE and MCE). For CP-newborn, too, multi-task learning resulted in a better-calibrated model in terms of ECE. We repeated the above experiment with a randomly-selected subset of $n=30$ images from this dataset. We observed that with this reduced dataset size, multi-task learning indeed resulted in statistically significantly better calibrated models. With $n=30$, a single model trained on the Hippocampus dataset achieved ECE and MCE of 0.16±0.04 and 0.32±0.08, respectively, whereas multi-task learning achieved ECE and MCE of 0.14±0.03 and 0.27±0.05, respectively, which were both significantly lower. Hence, multi-task learning in general improves the model prediction calibration, even for relatively large datasets. However, on some segmentation tasks (e.g., Hippocampus in this experiment) the positive effects of multi-task learning may disappear if large datasets are available. Figure 3 shows examples of estimated uncertainty maps. For this figure, we have intentionally selected test images on which the segmentation accuracy was relatively low. The figure shows that the model displays high segmentation uncertainty at the locations where segmentation error occurs, visually confirming that the model is well-calibrated.

In terms of segmentation accuracy, the results are somewhat mixed. Nonetheless, on average Table II shows that training a single model on a pool of heterogeneous datasets can achieve results that are as good as or even better than when dedicated models. Multi-task learning achieved significantly higher DSC on five datasets, significantly lower HD95 on seven datasets, and significantly lower AssD on four datasets. only on one of the datasets (Hippocampus) multitask learning was significantly worse. Given the small size of most of our datasets, the fact that a single model can learn to automatically recognize the context and accurately segment the organ of interest is interesting. Figure 4 shows example slices of several test images from different datasets and the computed segmentations. They show that the joint model accurately segments different organs in different imaging modalities. Please note that here our aim was to improve the model confidence calibration and not to improve the segmentation accuracy. As we showed above, the proposed strategy of training on heterogeneous data did improve the model confidence calibration on the overwhelming majority of the datasets. The fact that it also “on average" improved the segmentation accuracy is an incidental advantage.

We investigated the effect of mixup, [34], on the calibration and segmentation accuracy with the same twelve datasets used in Table II. As we mentioned above, recent works have shown that mixup can improve calibration of image classification models [35] and segmentation accuracy [66]. However, those studies used images from the same modality and organ. In our setting with a mix of twelve different datasets, mixup resulted in a consistent deterioration of model performance and calibration. on all twelve datasets, mixup resulted in lower DSC, often by large margins. For example on CP-younger fetus, Heart, and Hippocampus datasets, DSC dropped by approximately 0.09-0.13 compared with the results shown in Table II, which were statistically significant reductions. similarly, model calibration was worse when mixup was used. Therefore, mixup is very ineffective on datasets with heterogeneous modalities and organs of interest. Below, we present the results of using mixup on a more homogeneous set of data, i.e. CP segmentation in MRI.

Here we present an experiment to compare the multi-task learning approach with transfer learning using the three cortical plate datasets (see Table I). As shown in Figure 5, the shape and complexity of cortical plate evolves dramatically before and after birth. In addition, the sizes of the three datasets are highly unequal. Given the much smaller sizes of two of the datasets, transfer learning is the method that is recommended by some studies [4], [10]. Table III compares the results obtained with different transfer learning trials and the results obtained with multi-task learning. In each of the transfer learning trials, we first trained the model to convergence on one of the datasets and then fine-tuned it on another dataset. We then further fine-tuned the model on the third dataset. our fine-tuning strategy was “deep fine-tuning" [10]; we reduced the initial learning rate by half and fine-tuned all model layers. Shallow fine-tuning and keeping the initial learning rate produced inferior results. We also tried other fine-tuning curricula, i.e., orders of datasets used in fine-tuning, but did not achieve better results than those in Table III. For each dataset, we performed paired t-tests between the four results (i.e., the three transfer learning trials and the multi-task learning trial). statistically better results, at $p=0.01$, were marked with bold type.

Table III shows that the joint model had better-calibrated predictions than models trained on individual datasets as well as all three transfer learning trials on all three datasets in terms of ECE and MCE. In terms of accuracy, the multi-task learning approach was similar to or better than transfer learning. Although transfer learning improved the segmentation accuracy in some cases, the improvements were marginal. Multi-task learning, on the other hand, resulted in statistically significant improvements in segmentation accuracy. For the smallest dataset, i.e., CP-older fetus, multi-task learning significantly improved DSC, HD95, and ASSD. For the other two datasets, multi-task learning significantly improved ASSD. Figure 6 shows segmentation results on example test images from the three cortical plate segmentations for models trained on each of the three datasets, separately, and also for the joint model. The results show that a model trained on all three datasets can segment the test data from all three datasets with high accuracy. However, a model trained on each one of the datasets may perform very poorly on the test images from the other two datasets.

Table IV presents a comparison of multi-task learning with mixup on the same three CP datasets. Unlike with the twelve heterogeneous datasets in Table II, in this experiment mixup works and it does improve the segmentation accuracy of the model compared with standard training results presented in Table III. However, improvements in model calibration due to mixup are marginal, and ECE and MCE values achieved with multi-task learning are significantly better than those with mixup on all three datasets.

An additional appeal of a joint model that accurately segments all three datasets is that one would need to maintain only one set of model weights. When trained on one of the datasets, the model will perform poorly on the other datasets, one would need to maintain three separate models, and for a test image one would need to know which of the three datasets the image belongs to. Furthermore, the time required to train a single model on several datasets is generally less than the time needed to train separate models. For example, for CP segmentation, the training time for a model to segment cp-younger fetus, CP-older fetus, and CP-newborn is approximately 9 hours, whereas the training time for a model for each of these three datasets is approximately 5 hours, for a total of approximately 15 hours.

B. Detecting OOD test data

We present the results of OOD detection with our method and competing methods in four different experiments.

In the first experiment, we used a mixture of eight different datasets for training. These included CP-younger fetus, CP-older fetus, Prostate, Heart, Liver-CT, Liver-MRI-SPIR, Liver-MRI-DUAL-In and Liver-MRI-DUAL-Out datasets. We used test images from the same eight dataset as in-distribution data. As OOD data, we used Pancreas, Hippocampus, and Spleen datasets. Histograms of OODM have been shown in the lower part of Figure 2. Table V compares different methods for OOD detection. Our method perfectly detected the OOD images. UNC-Dropout failed, achieving an accuracy of 0.55. The other methods performed better than UNC-Dropout, but none of them achieves the level of accuracy of our method. Among the competing methods, UNC-Ensemble achieved better results, but it requires training tens of models. Following the recommendations of [38], we trained 50 models for this method.

In the second experiment, we trained a model on CP-newborn and applied the model on test images from the same dataset and the other two CP datasets. Figure 7 shows OODM histograms for this experiment. The OODM values for both CP-younger fetus and CP-older fetus fall outside of the distribution of CP-newborn. This model achieved DSC of 0.689 ± 0.095 and 0.781 ± 0.028 on CP-younger fetus and CP-older fetus datasets, respectively. These are very low compared with the results shown in Table III. Therefore, images from CP-younger fetus and CP-older fetus datasets are OOD. Our proposed method easily distinguished OOD data from in-distribution data. It is interesting to note that OODM values for CP-younger fetus are distributed farther away than those of CP-older fetus. This makes sense because CP-younger fetus is less similar to CP-newborn than CP-older fetus is. As shown in Table VI, our method accurately separated OOD data samples from in-distribution data samples. UNC-Dropout achieved an accuracy of 0.57, while UNC-Ensemble and Mah-Dist achieved 0.80. Outlier Exposure was the best of the competing methods, but still achieved accuracy, sensitivity, and specificity of 20% lower than the proposed method. DkNN cannot be applied in this experiment because the training set has only one class.

In another experiment, we trained our model on CP-younger fetus dataset and tested on the other two CP datasets and on four completely different datasets: Heart, Liver-CT, Hippocampus, and Pancreas. This model achieved a DSC of 0.788 ± 0.045 and 0.765 ± 0.061 on CP-older fetus and CP-newborn datasets, respectively, which are lower than the results in Table III. Also, as expected, on the other four datasets the model failed, achieving a mean DSC of 0.20-0.45. We present the results of this experiment separately for two CP datasets and the four non-CP datasets. Table VII shows the OOD detection accuracy results for different methods. For both CP and non-CP datasets our method achieved high detection accuracy and performed better than the other techniques. Among the competing methods, UNC-Ensemble and Outlier exposure achieved better results, but the accuracy for our method was much higher. Figure 8 displays the histograms of the OODM values for this experiment.

Finally, we report an experiment with the three liver MRI datasets (See Table I). The top section of Figure 9 shows a sample image from each of these datasets. This experiment demonstrates that OOD data are often not easy to distinguish visually. Our experiments show that a model trained on Liver-MRI-SPIR and Liver-MRI-DUAL-In accurately segments images from Liver-MRI-DUAL-Out (mean DSC= 0.89). Similarly, a model trained on Liver-MRI-SPIR and Liver-MRI-DUAL-Out achieved a mean DSC of 0.86 on Liver-MRI-DUAL-In. Even a model trained on Liver-MRI-DUAL-SPIR alone, accurately segmented Liver-MRI-DUAL-In and Liver-MRI-DUAL-Out. On the other hand, a model trained on Liver-MRI-DUAL-In and/or Liver-MRI-DUAL-Out failed on images from Liver-MRI-SPIR (mean DSC ≈ 0.40). These observations are not intuitive and are not at all easy to foretell by visually inspecting these images. Specifically, there are asymmetries that cannot be predicted by visual inspection. As an example, as mentioned above, Liver-MRI-DUAL-SPIR is OOD for a model trained on Liver-MRI-DUAL-In but Liver-MRI-DUAL-In is not OOD for a model trained on Liver-MRI-DUAL-SPIR. This example further highlights the importance and challenging nature of OOD detection in CNN-based medical image segmentation.

Figure 9(a) shows OODM histograms for an experiment where Liver-MRI-DUAL-In and Liver-MRI-DUAL-Out were used for training. The OODM values were computed on the test data from the same datasets and Liver-MRI-SPIR, which is OOD for this model. Our proposed method easily separates in-distribution from OOD data. Table VIII shows that UNC-Dropout and DkNN have low accuracies. UNC-Ensemble achieved better results, but still has an accuracy and AUC that are approximately 20% lower than our proposed method. Figure 9(b) shows the OODM for an experiment where Liver-MRI-SPIR and Liver-MRI-DUAL-In were used for training. The trained model works well on Liver-MRI-DUAL-Out dataset as well. As desired, the OODM values for most images from Liver-MRI-DUAL-Out fall below $\tau$, and hence correctly classified as in-distribution. Figure 9(c) shows the same for an experiment in which Liver-MRI-SPIR and Liver-MRI-DUAL-Out were used for training.

In terms of computation time, our method processes an image in approximately 5 seconds on a Linux machine with 32 GB of memory and an NVIDIA GeForce GTX 1080 GPU. With our implementation of UNC-Dropout, UNC-Ensemble, Mah-Dist, ODIN, and DkNN, they take approximately 4, 15, 5, 12, and 35 seconds, respectively. Another advantage of our OOD detection method is relative simplicity. The only hyperparameter in the OOD detection method itself is $C$ that determines the threshold, $\tau$. This value influences the trade-off between sensitivity and specificity. In all of our experiments we reported the value of AUC, which accounts for this trade-off. The accuracy/sensitivity/specificity values reported above were all obtained with $C=2.5$. The actual impact of the value of $C$ depends on the experiment. For example, for the experiment shown in Figure 8(a) using $C=2.0$ would result in sensitivity and specificity of 0.96 and 0.86, respectively, and using $C=3.0$ would result in sensitivity and specificity of 0.88 and 1.00, respectively. In Figure 8(b), on the other hand, sensitivity and specificity remain at 1.00 for all value in the range $C\in [2.0,3.0]$.

Applying the method on other feature maps resulted in lower accuracy than on the last layer. As an example, in the experiment with brain cortical plate datasets reported in Table V, applying the method on the 2nd deepest feature map resulted in accuracy, sensitivity, and specificity of 0.92, 0.89, and 0.92, respectively. Moreover, applying the method on the 3rd deepest feature map resulted in accuracy, sensitivity, and specificity of 0.90, 0.88, and 0.91, respectively. On the coarsest feature maps (i.e., the last encoder feature maps) the achieved accuracy, sensitivity, and specificity were 0.84, 0.83, and 0.86, respectively. Better accuracy achieved at deeper layers is likely due to the higher degree of disentanglement in deeper layers [63].

This study used a large number of dataset to evaluate the proposed methods. Nonetheless, the variability in medical image data is very high. It would be instructive to explore the potential of the proposed methods in multi-label and multi-class settings. In one experiment, we repeated the experiment reported in Table VII by considering a two-class segmentation for the Hippocampus data, where the Hippocampus dataset has two segmentation labels (anterior and posterior). Therefore, the output layer in this experiment had one extra channel for the Hippocampus dataset. For other datasets this channel is padded with zeros in the training data. With this setting, our OOD detection method achieved an accuracies of 0.94 and 1.00 on the CP and non-CP test data, which are similar to the results presented in Table VII. Another potentially important factor is image resolution. Increasing image resolution beyond the native (acquisition) resolution by upsampling did not significantly improve the segmentation accuracy or confidence calibration. even for fine structures such as the brain cortical plate. This may be due to the fact that after image upsampling the corresponding segmentations also need to be upsampled. For fine structures such as cortical plate, expert-provided labels are most accurate in the resolution used during manual annotation. Upsampling leads to inevitable errors in the training labels of fine structures that may contribute to model inaccuracy and poor calibration. Downsampling the images and labels substantially reduced the segmentation accuracy and calibration on fine structures such as cortical plate. significantly increasing/decreasing image resolution can also render an in-distribution test image OOD with respect to a model trained on images with very different resolutions. In terms of segmentation accuracy, another factor is the probability of sampling from different datasets. As we mentioned in section III-C, we sampled from each dataset with a probability proportional to the inverse of the square root of dataset size, $1∕\sqrt{n}$. This usually resulted in higher accuracy than uniform sampling (i.e., $1∕n$). As an example, using a uniform sampling for the experiment reported in Table II resulted in DSC, HD95, and ASSD of 0.87±0.04, 0.88±0.03, and 0.24 ± 0.08 for CP- younger fetus, and 0.87 ± 0.11, 8.9 ± 7.8, and 2.20 ± 2.28 for the Heart dataset. These are slightly worse than the results presented in Table II.