AiA-UNet: Attention in Attention for Medical Image Segmentation

In medical image segmentation, U-Net has consistently played a vital role. Recently, the U-Net networks based on the Vision Transformer (ViT) architecture have become more and more popular. ViT exhibits superior capabilities in handling long-range dependencies and capturing global contextual information. However, it requires significant computational cost, and does not explore the optimal matching and the potential dependencies between different patches. To address the aforementioned issues, we propose a novel network framework, called AiA-UNet, for medical image segmentation. The AiA-UNet makes two main contributions. A convolutional self-attention mechanism is proposed to replace the self-attention module in ViT, effectively reducing computational complexity. Moreover, an Attention in Attention module (AiA) is applied within the ViT block. Experimental results on the Synapse multi-organ segmentation dataset demonstrate that AiA-UNet outperforms Trans-UNet by 5.40% and Swin-UNet by 3.75%. Code and models are available at https://github.com/xiaoqin1998/AiA-UNet.