ROAM: Random Layer Mixup for Semi-Supervised Learning in Medical Imaging

Abstract

Medical image segmentation is one of the major challenges addressed by machine learning methods. Yet, deep learning methods profoundly depend on a large amount of annotated data, which is time-consuming and costly. Though, semi-supervised learning methods approach this problem by leveraging an abundant amount of unlabeled data along with a small amount of labeled data in the training process. Recently, MixUp regularizer has been successfully introduced to semi-supervised learning methods showing superior performance. MixUp augments the model with new data points through linear interpolation of the data at the input space. We argue that this option is limited. Instead, we propose ROAM, a \textit{random layer mixup}, which encourages the network to be less confident for interpolated data points at randomly selected space. ROAM generates more data points that have never seen before, and hence it avoids over-fitting and enhances the generalization ability. We conduct extensive experiments to validate our method on publicly available datasets on whole-brain image segmentation (MALC) and Lung and Infection segmentation in COVID-19. ROAM achieves state-of-the-art (SOTA) results in fully supervised (89.5\%) and semi-supervised (87.0\%) settings with a relative improvement of up to 2.40\% and 16.50\%, respectively for the whole-brain segmentation. Similarly, ROAM achieves superior performance for COVID-19 segmentation beating SOTA SSL settings.