Multi-task Attention-Based Semi-supervised Learning for Medical Image Segmentation

Abstract

We propose a novel semi-supervised image segmentation method that simultaneously optimizes a supervised segmentation and an unsupervised reconstruction objectives. The reconstruction objective uses an attention mechanism that separates the reconstruction of image areas corresponding to different classes. The proposed approach was evaluated on two applications: brain tumor and white matter hyperintensities segmentation. Our method, trained on unlabeled and a small number of labeled images, outperformed supervised CNNs trained with the same number of images and CNNs pre-trained on unlabeled data. In ablation experiments, we observed that the proposed attention mechanism substantially improves segmentation performance. We explore two multi-task training strategies: joint training and alternating training. Alternating training requires fewer hyperparameters and achieves a better, more stable performance than joint training. Finally, we analyze the features learned by different methods and find that the attention mechanism helps to learn more discriminative features in the deeper layers of encoders.

Recent Advances in Variational Autoencoders With Representation Learning for Biomedical Informatics: A Survey

“…VAEs have also succeeded in biological image analyses, and many studies show superior performance. The main research area based on the VAE use in medical imaging datasets includes: 1) Medical image data augmentation for downstream tasks include image classification [68,71,79,80], image segmentation [72][73][74][75][76]87], image restoration [85,86], and image reconstruction [72,[81][82][83].…”

“…Another approach for few shot learning is the used of semi-supervised learning in segmentation with VAEs: Sedai et al [73] proposed a semi-supervised VAEs-based method to segment optical cup, and utilized a small number of labeled data to accurately localize the anatomical structure. Chen et al [74] proposed a VAEs-based semi-supervised image segmentation method for Brain tumor and white matter hyperintensities segmentation. Qian et al [77] build a novel VAE for estimating object shape uncertainty in medical images.…”

Segmentation of Intracranial Hemorrhage Using Semi-Supervised Multi-Task Attention-Based U-Net

“…By forcing the network to learn features corresponding to a larger dataset, we effectively regularized our network against non-generalizing local minima that existed within the training process. In this paper, we solved this problem by using an autoencoder with a modified u-net initially proposed by Chen et al [34] on two separate, independent sources of data detailed in Section 5.1.…”

“…For the first two tasks, we used the u-net shown in Figure 1. For the third task, the encoder and decoder structure was identical to Figure 1, except that the residual connections were removed for the second (unsupervised) decoder and that the output included two feature maps rather than one [34]. Architecture of our modified u-net used in our models.…”

“…In this paper, we introduce a modified u-net and curriculum learning strategy, using a semi-supervised model based on an earlier work by Chen et al [34], to perform semantic segmentation on a small dataset of ICH obtained from the Al Hilla Teaching Hospital in Iraq [2]. We also used the RSNA Intracranial Hemorrhage Detection dataset as a collection of unlabeled images to test our semi-supervised learning strategy.…”