[Paper Review] TransAttUnet: Multi-level Attention-guided U-Net with Transformer for Medical Image Segmentation

Bingzhi Chen, Yishu Liu+3 more|arXiv (Cornell University)|Jul 12, 2021

Advanced Neural Network ApplicationsComputer Science45 references84 citations

TL;DR

TransAttUnet combines Transformer-based self-attention and global spatial attention with multi-scale skip connections in a U-Net to improve medical image segmentation across multiple modalities, outperforming state-of-the-art baselines.

ABSTRACT

Accurate segmentation of organs or lesions from medical images is crucial for reliable diagnosis of diseases and organ morphometry. In recent years, convolutional encoder-decoder solutions have achieved substantial progress in the field of automatic medical image segmentation. Due to the inherent bias in the convolution operations, prior models mainly focus on local visual cues formed by the neighboring pixels, but fail to fully model the long-range contextual dependencies. In this paper, we propose a novel Transformer-based Attention Guided Network called TransAttUnet, in which the multi-level guided attention and multi-scale skip connection are designed to jointly enhance the performance of the semantical segmentation architecture. Inspired by Transformer, the self-aware attention (SAA) module with Transformer Self Attention (TSA) and Global Spatial Attention (GSA) is incorporated into TransAttUnet to effectively learn the non-local interactions among encoder features. Moreover, we also use additional multi-scale skip connections between decoder blocks to aggregate the upsampled features with different semantic scales. In this way, the representation ability of multi-scale context information is strengthened to generate discriminative features. Benefitting from these complementary components, the proposed TransAttUnet can effectively alleviate the loss of fine details caused by the stacking of convolution layers and the consecutive sampling operations, finally improving the segmentation quality of medical images. Extensive experiments on multiple medical image segmentation datasets from different imaging modalities demonstrate that the proposed method consistently outperforms the state-of-the-art baselines. Our code and pre-trained models are available at: https://github.com/YishuLiu/TransAttUnet.

Motivation & Objective

• Motivate improvements in medical image segmentation by addressing local bias of convolutional encoders.
• Propose a Transformer-based self-aware attention (SAA) module combining Transformer Self Attention and Global Spatial Attention.
• Introduce multi-scale skip connections to better fuse multi-scale decoder features.
• Demonstrate the method's effectiveness across diverse medical imaging modalities.
• Provide a fair comparison with strong baselines and ablation studies.

Proposed method

• Integrate a self-aware attention (SAA) module at the encoder-decoder bridge to fuse TSA and GSA with encoder features.
• Use Transformer Self Attention (TSA) with multi-head attention and learned positional encoding to model long-range dependencies.
• Apply Global Spatial Attention (GSA) to capture global context via position-aware channel interactions.
• Fuse TSA and GSA outputs with the encoder features through a learnable weighted combination (F_SAA).
• Implement multi-scale skip connections between decoder blocks to progressively aggregate features from varying semantic scales using residual or dense connections.
• Train with a joint Dice and BCE loss (L = alpha*L_BCE + beta*L_Dice) to balance pixel-wise accuracy and segmentation overlap.

Experimental results

Research questions

• RQ1 Can integrating Transformer-based self-attention and global spatial attention improve U-Net-based medical image segmentation?
• RQ2 Do multi-scale skip connections (residual/dense) better preserve fine details than traditional cascade connections?
• RQ3 How does TransAttUnet perform across diverse modalities and datasets compared to state-of-the-art baselines?
• RQ4 What is the impact of the SAA module and the multi-scale fusion on segmentation accuracy and boundary precision?

Key findings

• TransAttUnet variants outperform baseline U-Nets and several baselines across multiple datasets.
• The TransAttUnet_R (dense skip connections) achieves the highest Dice score on ISIC-2018 (90.74%).
• Transformer Self Attention (TSA) and Global Spatial Attention (GSA) jointly improve contextual modeling more than either alone.
• Multi-scale skip connections (residual or dense) provide better feature aggregation than one-off cascade connections, mitigating loss of fine details.
• Compared with MCTrans on ISIC-2018, TransAttUnet_R gives a Dice improvement (90.74% vs. 90.35%).