Distill DSM: Computationally efficient method for segmentation of medical imaging volumes

Harsh Maheshwari, Vidit Goel, Ramanathan Sethuraman, Debdoot Sheet

Keywords: Volumetric segmentation, parameter efficient 3D CNN, deep learning, distillation

TL;DR: We propose a novel module Distill Depth shift module (Distill DSM) for efficiently using the information along the depth dimension with negligible increase in the parameters compared to 2D CNN.

Abstract: Accurate segmentation of volumetric scans like MRI and CT scans is highly demanded for surgery planning in clinical practice, quantitative analysis, and identification of disease. However, accurate segmentation is challenging because of the irregular shape of given organ and large variation in appearances across the slices. In such problems, 3D features are desired in nature which can be extracted using 3D convolutional neural network (CNN). However, 3D CNN is compute and memory intensive to implement due to large number of parameters and can easily over fit, especially in medical imaging where training data is limited. In order to address these problems, we propose a distillation-based depth shift module (Distill DSM). It is designed to enable 2D convolutions to make use of information from neighbouring frames more efficiently. Specifically, in each layer of the network, Distill DSM learns to extract information from a part of the channels and shares it with neighbouring slices, thus facilitating information exchange among neighbouring slices. This approach can be incorporated with any 2D CNN model to enable it to use information across the slices with introducing very few extra learn-able parameters. We have evaluated our model on BRATS 2020, heart, hippocampus, pancreas and prostate dataset. Our model achieves better performance than 3D CNN for heart and prostate datasets and comparable performance on BRATS 2020, pancreas and hippocampus dataset with simply 28\% of parameters compared to 3D CNN model.

Registration: I acknowledge that publication of this at MIDL and in the proceedings requires at least one of the authors to register and present the work during the conference.

Authorship: I confirm that I am the author of this work and that it has not been submitted to another publication before.

Paper Type: validation/application paper

Primary Subject Area: Segmentation

Secondary Subject Area: Application: Radiology