An Uncertainty-Aware Transformer for MRI Cardiac Semantic Segmentation via Mean Teachers

Abstract

Deep learning methods have shown promising performance in medical image semantic segmentation. The cost of high-quality annotations, however, is still high and hard to access as clinicians are pressed for time. In this paper, we propose to utilize the power of Vision Transformer (ViT) with a semi-supervised framework for medical image semantic segmentation. The framework consists of a student model and a teacher model, where the student model learns from image feature information and helps teacher model to update parameters. The consistency of the inference of unlabeled data between the student model and teacher model is studied, so the whole framework is set to minimize segmentation supervision loss and consistency semi-supervision loss. To improve the semi-supervised performance, an uncertainty estimation scheme is introduced to enable the student model to learn from only reliable inference data during consistency loss calculation. The approach of filtering inconclusive images via an uncertainty value and the weighted sum of two losses in the training process is further studied. In addition, ViT is selected and properly developed as a backbone for the semi-supervised framework under the concern of long-range dependencies modeling. Our proposed method is tested with a variety of evaluation methods on a public benchmarking MRI dataset. The results of the proposed method demonstrate competitive performance against other state-of-the-art semi-supervised algorithms as well as several segmentation backbones.

A Appendix

Table 4 gives detailed systematic IOU results under different assumptions of the ratio of labeled to total data, on the MRI Cardiac test set. It is pleasantly remarkable to see serviceable results being obtained with a proportion of labelled data as small as 1%, 2%, or 3% of the total. Given the small set of type-specific annotations that exist, they can now be put to good use by pairing them with large amounts of unlabeled data and making them available through our proposed method.

Table 5 and Table 6 reports the different approaches to modify the threshold $\tau$ of filtering certain or uncertain pixels with uncertainty estimation scheme, and the weight $\lambda$ of loss $\mathcal {L}_\mathrm{c}$ in each training iteration. We explore the fixed value, exponential ramp up [7], linear ramp up, cosine ramp down [9] and variants of them. Details of exponential ramp up, linear ramp up and cosine ramp down is illustrated in the following Eq. 8, 9, 10, respectively. Each experiment is conducted with different approaches under the other one either $\tau$ or $\lambda$ is fixed with exponential ramp up. The results illustrates different approaches of updating $\tau$, $\lambda$ in each training iteration cannot significantly improve the performance of proposed method, and all other experiments for $\tau$, $\lambda$ is with exponential ramp up.

$$\begin{aligned} \tau or \lambda = e^{-5\times (1-t_\mathrm{iteration}/t_\mathrm{maxiteration})^{2}} \end{aligned}$$

(8)

$$\begin{aligned} \tau or \lambda = t_\mathrm{iteration}/t_\mathrm{maxiteration)} \end{aligned}$$

(9)

$$\begin{aligned} \tau or \lambda = 0.5 \times (cosine(\pi \times t_\mathrm{iteration} / t_\mathrm{maxiteration})+1) \end{aligned}$$

(10)

Table 5. Ablation studies on the threshold setting of uncertainty in training process (the higher, the better)

Table 6. Ablation studies on the weight setting of consistency loss in training process (the higher, the better)

Figure 4 sketches randomly selected raw images with their corresponding uncertainty maps, and masks generated by proposed method at three different stages (from the beginning to the end) of the training process. In uncertainty maps, yellow represents the teacher ViT $f_\mathrm{t}$ is uncertain of prediction with the given pixels, and blue represents the teacher ViT $f_\mathrm{t}$ is certain of prediction with the given pixels. The uncertainty map is gradually moving from yellow to green in the training process as shown in Fig. 4. The threshold of certainty estimation is then applied with uncertainty map which results in masks, where the white represents that the prediction by teacher ViT $f_\mathrm{t}$ is certain enough to guide the student ViT $f_\mathrm{s}$ i.e. for calculation the consistency loss $\mathcal {L}_\mathrm{s}$, and the black represents that the pixels with uncertainty is temporally unavailable to be considered in consistency semi-supervision loss calculation. Please remind that both the background and ROI can be certain with the white simultaneously. Some typical example masks illustrates that model is only uncertain with the boundary of ROI as shown in Fig. 4, and finally the framework is very likely to be certain with the whole image with a proper threshold setting, that the uncertainty map is going to be blue, mask is going to be white in the end of training process.

Fig. 4.

Sample uncertainty maps, masks, and raw images during the training process (Color figure online)