[论文解读] MaxViT-UNet: Multi-Axis Attention for Medical Image Segmentation

Abdul Rehman, Asifullah Khan|arXiv (Cornell University)|May 15, 2023

Brain Tumor Detection and Classification被引用 8

一句话总结

本论文提出 MaxViT-UNet，一种用于医学图像分割的混合 CNN-Transformer 编码器-解码器，采用带多轴自注意力的混合解码器以在 modest memory 和 computation 下提升细胞核分割性能。

ABSTRACT

Since their emergence, Convolutional Neural Networks (CNNs) have made significant strides in medical image analysis. However, the local nature of the convolution operator may pose a limitation for capturing global and long-range interactions in CNNs. Recently, Transformers have gained popularity in the computer vision community and also in medical image segmentation due to their ability to process global features effectively. The scalability issues of the self-attention mechanism and lack of the CNN-like inductive bias may have limited their adoption. Therefore, hybrid Vision transformers (CNN-Transformer), exploiting the advantages of both Convolution and Self-attention Mechanisms, have gained importance. In this work, we present MaxViT-UNet, a new Encoder-Decoder based UNet type hybrid vision transformer (CNN-Transformer) for medical image segmentation. The proposed Hybrid Decoder is designed to harness the power of both the convolution and self-attention mechanisms at each decoding stage with a nominal memory and computational burden. The inclusion of multi-axis self-attention, within each decoder stage, significantly enhances the discriminating capacity between the object and background regions, thereby helping in improving the segmentation efficiency. In the Hybrid Decoder, a new block is also proposed. The fusion process commences by integrating the upsampled lower-level decoder features, obtained through transpose convolution, with the skip-connection features derived from the hybrid encoder. Subsequently, the fused features undergo refinement through the utilization of a multi-axis attention mechanism. The proposed decoder block is repeated multiple times to segment the nuclei regions progressively. Experimental results on MoNuSeg18 and MoNuSAC20 datasets demonstrate the effectiveness of the proposed technique.

研究动机与目标

• Motivate hybrid CNN-Transformer architectures to capture both local and global context in medical image segmentation.
• Propose MaxViT-UNet, a UNet-like encoder-decoder with a Hybrid Decoder that fuses upsampled decoder features and skip connections using multi-axis attention.
• Reduce memory and computational burden while boosting discriminative power for separating nuclei from background.

提出的方法

• Introduce MaxViT-UNet, a UNet-type Hybrid CNN-Transformer architecture for medical image segmentation.
• Design a Hybrid Decoder that integrates upsampled lower-level decoder features (via transpose convolution) with skip-connection features from the hybrid encoder.
• Apply a multi-axis attention mechanism within each decoding stage for feature refinement.
• Repeat the proposed decoder block multiple times to progressively segment nuclei regions.
• Evaluate on MoNuSeg18 and MoNuSAC20 datasets to demonstrate effectiveness.

实验结果

研究问题

• RQ1 Can a CNN-Transformer hybrid architecture with multi-axis attention improve nuclei segmentation compared to standard UNet or pure Transformer approaches?
• RQ2 Does the proposed Hybrid Decoder effectively fuse encoder-decoder features with attention to enhance segmentation boundaries?
• RQ3 What is the impact of multi-axis attention on discriminating object vs. background in medical images?
• RQ4 How does the approach perform on MoNuSeg18 and MoNuSAC20 datasets in terms of segmentation quality and efficiency?

主要发现

• The MaxViT-UNet architecture demonstrates effectiveness on MoNuSeg18 and MoNuSAC20 datasets.
• Incorporating multi-axis attention within the decoder improves discrimination between nuclei and background regions.
• The Hybrid Decoder fuses upsampled decoder features with skip-connection encoder features before refinement, enabling progressive nucleus segmentation.
• The approach aims for nominal memory and computational burden while achieving improved segmentation performance.