Multi-scale Hierarchical Transformer structure for 3D medical image segmentation

Transformers have demonstrated great potential in computer vision tasks which is benefited from its ability of long-range dependency modeling. However, its core material multi-head self-attention (MHSA) has high computational and spacial complexity. Besides, tokens embedded with 1D position sequence are unable to represent inductive bias of locality and neighbor contextual information, which are essential for high-resolution downstream vision tasks. Thus, transformer is hard to use in image segmentation tasks, especially in 3D medical images. In this paper, we propose a novel multi-scale hierarchical framework (MSHT) that efficiently incorporates the local modeling ability of CNN and long-range modeling ability of transformer for accurate 3D medical image segmentation. Unlike many prior Transformer-based solutions, the proposed MSHT first adopts parallel CNN and transformer as encoder block to extract the global and local feature representations. As the core component for our MSHT, a position embedding method for 3D medical image (3D IPE) is proposed to generate relational and absolute position encoding for tokens fed in transformer block. We conduct an extensive evaluation on both Lits2017 dataset and Kits2019 dataset. The results indicate that our MSHT leads to a substantial performance over other CNN-based methods on 3D image segmentation.

I. Introduction

Image segmentation is often the first step and still a longstanding challenge in medical image analysis. Since the U-Net [8] proposed, convolutional neural networks (CNNs) have become the predominant methods to solve this task. However, CNNs still suffer from the limited receptive field and fail to capture long-range dependency. Transformer entirely relies on self-attention mechanism to model long-range interaction and has achieved brilliant successes in Natural Languages Processing (NLP) field [10].Thus, lots of efforts are dedicated to explore efficient transformer-based structure in computation vison (CV) tasks. Some researches use CNNs to extract deep feature to feed transformer for further processing and regressing, such as DETR [2] and TransUNet [3]. Some researchers tried to apply pure transformer structure in vision task. ViT [5] is a thorough success, in which sequences of image patches can perform well with pure transformer structure on image classification tasks. However, when viewing image patches or feature maps as tokens, the sequence length is much larger than that in NLP applications, resulting in higher computational and space complexity owing to MHSA. Thus, this type structures allow low-resolution inputs, which is challenging to be directly adapted to pixel-level dense predictions. To address this problem, PVT [11] introduces a progressive shrinking pyramid structure to reduce the sequence length of token, which reduce computational complexity of large feature maps and achieve high resolution output. However, CV tasks often need the structure or spatial local information, which is ignored by the absolute positional encoding in transformer. To capture local position interaction, CMT [6] propose to extract local information by convolution operation and add the position information to the patches before feed in transformer module.