[Paper Review] MaxViT-UNet: Multi-Axis Attention for Medical Image Segmentation

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

Brain Tumor Detection and Classification8 citations

TL;DR

This paper presents MaxViT-UNet, a hybrid CNN-Transformer encoder-decoder for medical image segmentation that uses a Hybrid Decoder with multi-axis self-attention to improve nuclei segmentation with modest memory and 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.

Motivation & Objective

• 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.

Proposed method

• 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.

Experimental results

Research questions

• 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?

Key findings

• 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.