MaxViT-UNet: Multi-Axis Attention for Medical Image Segmentation

TL;DR

MaxViT-UNet tackles nuclei segmentation by uniting CNN-like feature extraction with Transformer-based global context via a MaxViT-based encoder and a novel Hybrid Decoder. The architecture employs a two-pronged MaxViT block (MBConv + multi-axis attention) to achieve local and global feature fusion with linear-like complexity, while a symmetric Decoder progressively refines segmentation with skip connections. Empirical results on MoNuSeg18 and MoNuSAC20 show significant improvements over CNN-based UNet and pure Transformer approaches, with ablation studies confirming the Hybrid Decoder’s contribution. The work yields a lightweight, plug-and-play framework that enhances segmentation accuracy and efficiency in medical imaging, with potential applicability to other 2D/3D modalities and ensemble strategies.

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.

Figures (6)

• Figure 1: Encoder-Decoder architecture of the Proposed MaxViT-UNet. The encoder generates hierarchical features at four scales. The proposed decoder first upscale the bottom-level features, merges them with skip-connection features, and applies the MaxViT-based hybrid feature processing blocks a couple of times to produce an output mask image for the "C" number of classes.
• Figure 2: Detailed architecture of the proposed Hybrid Decoder block. Features from the $i^{th}$ decoder stage are upscaled using ConvTranspose2D layer to match with size of the $(i+1)^{th}$ encoder stage coming from skip-connection. After the concatenation (concat) operation, the MaxViT-block is used a couple of times to merge the features efficiently.
• Figure 3: Data Pre-processing Pipeline visualized for MoNuSeg18 dataset. From left to right: Original Image (resized to $256 \times 256$), Random Affine (combination of Shift, Scale, and Rotate), Random Flip (either Horizontal or Vertical), PhotoMetric Distortion (changes the intensity of pixels), Padding (to ensure $256 \times 256$ image size), Final Augmented Input and Mask image are shown.
• Figure 4: Comparative plots of the proposed MaxViT-UNet with previous techniques on MoNuSeg18 and MoNuSAC20 challenge datasets. The left column displays the (a) Dice, (b) IoU, and (c) Training Loss on the MoNuSeg18 dataset, whereas the right column displays the (d) Dice, (e) IoU and (f) Training Loss on MoNuSAC20 dataset.
• Figure 5: Qualitative comparison of the proposed MaxViT-UNet with current methods on (a) MoNuSeg18 dataset; the colors white, red, and blue, respectively, indicate True-Positive, False-Positive, and False-Negative predictions. (b) The MoNuSAC20 dataset shows red, yellow, green, and blue representations of epithelial, lymphocyte, macrophage, and neutrophil, respectively.