Semi-Supervised Medical Image Segmentation with Equivariance Regularization

Abstract: The last decades of research in machine learning and deep learning have lead to enormous advancements in the field. One of the areas that stand to gain the most from this is the medical sector. However, the majority of deep learning models today rely on supervised learning and one considerable bottleneck in the sector's ability to adapt the technology is the need of large labeled datasets. Inspired by newly published semi-supervised methods for image classification, this work addresses the problem for the task of semantic segmentation (a task that is recurrent in medical imaging and cancer treatment) by introducing a semi-supervised method, named Equivariance Regularization (EquiReg). Using the EquiReg-method the model is trained to output equivariant predictions with respect to data augmentations using unlabeled data, in conjunction with standard supervised training using labeled data. Experiments with brain tumor MRI-scans from the BraTS 2019 dataset show that the EquiReg method can, when only a small percentage of data is labeled, boost performance by incorporating unlabeled data during training. Furthermore, the experiments show that the EquiReg-method also improves performance in the fully supervised case, by using the labeled data for both supervised and unsupervised training of the same model.