Fast and Efficient Brain Extraction with Recursive MLP Based 3D UNet

Abstract

Extracting brain from other non-brain tissues is an essential step in neuroimage analyses such as brain volume estimation. The transformers and 3D UNet based methods achieve strong performance using attention and 3D convolutions. They normally have complex architecture and are thus computationally slow. Consequently, they can hardly be deployed in computational resource-constrained environments. To achieve rapid segmentation, the most recent work UNeXt reduces convolution filters and also presents the Multilayer Perception (MLP) blocks that exploit simpler and linear MLP operations. To further boost performance, it shifts the feature channels in MLP block so as to focus on learning local dependencies. However, it performs segmentation on 2D medical images rather than 3D volumes. In this paper, we propose a recursive MLP based 3D UNet to efficiently extract brain from 3D head volume. Our network involves 3D convolution blocks and MLP blocks to capture both long range information and local dependencies. Meanwhile, we also leverage the simplicity of MLPs to enhance computational efficiency. Unlike UNeXt extracting one locality, we apply several shifts to capture multiple localities representing different local dependencies and then introduce a recursive design to aggregate them. To save computational cost, the shifts do not introduce any parameters and the parameters are also shared across recursions. Extensive experiments on two public datasets demonstrate the superiority of our approach against other state-of-the-art methods with respect to both accuracy and CPU inference time.

G. Shangguan and H. Xiong—Equal contribution.