TransAttUnet: Multi-level Attention-guided U-Net with Transformer for Medical Image Segmentation

Bingzhi Chen; Yishu Liu; Zheng Zhang; Guangming Lu; Adams Wai Kin Kong

Source-linked AI summary

TransAttUnet: Multi-level Attention-guided U-Net with Transformer for Medical Image Segmentation

Bingzhi Chen, Yishu Liu, Zheng Zhang, Guangming Lu, Adams Wai Kin Kong

arXiv:2107.05274v2Jul 12, 2021eess.IVcs.CV

TL;DR

Medical image segmentation needs accurate delineation, but convolutional architectures mainly capture local cues and inadequately model long-range dependencies. TransAttUnet combines multi-level guided attention with multi-scale skip connections in U-Net, and experiments report superior performance across five datasets, with limitations in memory, computation, and Transformer utilization.

Problem
Medical image segmentation is important for diagnosis and morphometric analysis, yet convolutional models have limited long-range contextual modeling.
Method
TransAttUnet combines a self-aware attention module using TSA and GSA with multi-scale skip connections inside a U-shaped encoder-decoder.
Results
Experiments on five medical image datasets report superior performance and generalizability compared with state-of-the-art baselines.
Takeaways & Limitations
The complementary attention and multi-scale connections are reported to improve semantic consistency and mitigate limitations of traditional U-shaped architectures.
Takeaways & Limitations
TransAttUnet relies on global self-attention, resulting in a large memory footprint and computation cost; Transformer use for biomedical images remains incomplete.

Abstract

from arXiv · show

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.

I. INTRODUCTION

Medical image segmentation supports diagnosis and morphometric analysis, but convolutional U-Net architectures have difficulty modeling long-range dependencies and preserving fine details. TransAttUnet addresses these limitations with guided attention and multi-scale decoder connections, and is reported to outperform state-of-the-art baselines across five datasets.

Motivation: Medical image segmentation delineates objects of interest from complex biomedical-image backgrounds for quantitative diagnosis and morphological analysis.Manual annotation can be accurate but is highly labor-intensive in clinical settings.
Limitations of Existing Architectures: Convolutional encoder-decoder architectures primarily extract local neighboring-pixel features and can limit long-range dependency modeling and low-level information transmission.Pooling, convolution, and sampling operations may compromise local information and segmentation quality.
Proposed Method: TransAttUnet integrates multi-level guided attention and multi-scale skip connections into a standard U-Net to improve medical image segmentation.The design jointly targets contextual modeling and feature aggregation across semantic scales.
Proposed Method: Transformer self attention and global spatial attention jointly model long-range interactions and global spatial relationships among encoder features.The combination is intended to support consistent feature representations and semantic embeddings.
Proposed Method: Residual or dense multi-scale decoder connections aggregate varying semantic scales and reduce fine-detail loss associated with large-scale upsampling.The connections use progressive feature transitions rather than relying only on one-step cascade connections.
Results: Experiments on five medical image datasets report superior performance and generalizability compared with state-of-the-art baselines.The datasets represent medical image segmentation challenges across different imaging modalities.

II. RELATED WORK

Prior U-Net research emphasizes attention-guided feature recalibration and multi-scale context fusion, while Transformer-based vision methods extend global context modeling. TransAttUnet belongs to this convergence of attention, multi-scale connections, and Transformer-based segmentation.

Variants of U-Net: U-Net variants are broadly organized around attention-guided approaches and multi-scale context approaches for medical image segmentation.These lines of work optimize feature selection or combine information across semantic scales.
Transformer for Medical Image Segmentation: TransAttUnet applies Transformer mechanisms within a U-Net framework to address the inability of conventional U-Net architectures to model long-range contextual interactions.Its design combines the related research directions rather than replacing the U-shaped architecture entirely.
Attention-Guided Approaches: Attention U-Net and Channel-UNet use attention mechanisms to recalibrate or highlight spatial features from encoder, decoder, and output representations.Their approaches generate or learn mappings that emphasize relevant feature responses.
Multi-Scale Context Approaches: UNet++, U2-Net, and MA-Unet aggregate features across semantic scales to retain detail and capture contextual information.Their redesigned or residual connections connect shallow and deep representations through multi-scale fusion.
Transformer for Various Vision Tasks: Transformer-based vision methods introduced self-attention into image recognition, object detection, and natural-image segmentation.Examples include ViT, DETR, and SETR.

2) Transformer for Medical Image Segmentation:

TransAttUnet uses a Transformer-based self-aware attention bridge within an encoder-decoder U-Net to capture global context and spatial relationships. Its Transformer self-attention pipeline projects encoder features into attention inputs, aggregates contextual values, and reshapes the result for decoding.

Overview of TransAttUnet: TransAttUnet retains a standard encoder-decoder U-shaped architecture while adding self-aware attention and multi-scale skip connections.The additions target long-range contextual information and feature aggregation across different semantic scales.
Overview of TransAttUnet: The self-aware attention module bridges the encoder and decoder and combines Transformer self attention with global spatial attention.These mechanisms capture wider contextual representations, long-range interactions, and global spatial relationships.
Transformer Self Attention: Transformer self attention uses multi-head attention with learned positional encoding to attend jointly to semantic information from global representation subspaces.The attention pipeline is depicted as part of the proposed architecture.
Transformer Self Attention: Encoder features F are linearly projected into query, key, and value matrices before scaled dot-product attention produces contextual feature representations.The attention-weighted values are reshaped into the final TSA output feature map Ftsa.

2) Global Spatial Attention:

Global Spatial Attention selectively aggregates global context and positional information into local encoder features, producing position-aware representations for TransAttUnet.

2) Global Spatial Attention:: GSA selectively aggregates global context and encodes broader positional information into local features to improve intra-class compactness and feature representations.Its stated role is to enrich local features with global spatial relationships.
2) Global Spatial Attention:: GSA reshapes encoder-derived feature maps into position representations and computes a softmax-normalized position attention map B.B captures interactions among all spatial positions, with n = h × w pixels.
2) Global Spatial Attention:: Each B_i,j measures the impact of position i on position j, enabling position-dependent aggregation across the image.The attention map is multiplied with W to produce aggregated features at each position.
2) Global Spatial Attention:: The resulting features are reshaped into the final GSA output F_gsa with channel and spatial dimensions c × h × w.This output is used as the global spatial attention representation.

3) Attention Embedding Fusion:

TransAttUnet fuses original encoder features with transformer and spatial attention embeddings, while multi-scale skip connections aggregate decoder features across semantic scales.

3) Attention Embedding Fusion:: The SAA fusion combines transformer self-attention features, global spatial attention features, and the original encoder features using weighted addition.The fusion is expressed as F_SAA = λ1F_tsa + λ2F_gsa + F_en.
3) Attention Embedding Fusion:: Unlike one-off cascade fusion, residual and dense connections progressively pass decoder information between blocks.Cascade concatenates all scales directly, whereas residual and dense schemes build connections step by step.
3) Attention Embedding Fusion:: The scale parameters λ1 and λ2 control the importance of self-attention and spatial-attention maps, respectively.Both parameters are initialized at zero and gradually increased during optimization.
3) Attention Embedding Fusion:: Multi-scale skip connections use upsampling, concatenation, and convolution to aggregate decoder features with varying semantic scales.These operations form residual or dense step-growth connections between decoder blocks.

2) Residual Connection:

The residual connection progressively transfers decoder features by combining each block’s upsampled input with its output before the next block.

2) Residual Connection:: Residual connections upsample each decoder block’s input features to the output resolution and concatenate them with the block’s output features.The concatenated representation becomes the input to the subsequent decoder block.
2) Residual Connection:: This residual scheme guides decoder upsampling through a sequence of transition operations across decoder stages.The proposed framework focuses on residual and dense multi-scale skip connection schemes.
2) Residual Connection:: Compared with one-off cascade connections, residual or dense step-growth connections gradually aggregate decoder features from different semantic scales.The stated goal is to generate more discriminative feature representations.
2) Residual Connection:: The paper states that these step-growth connections can mitigate fine-detail loss from over-upsampling and alleviate vanishing-gradient and overfitting problems.These are author-reported benefits of the residual or dense decoder connections.

IV. EXPERIMENTS

TransAttUnet is evaluated against state-of-the-art baselines on multiple segmentation datasets spanning skin lesions, lung fields, COVID-19 lesions, nuclei, and glands.

IV. EXPERIMENTS: The experiments compare TransAttUnet with state-of-the-art baselines and include ablation studies and visualization analysis of decoder stages.The evaluation is organized around benchmark comparisons, ablations, and decoder-stage visualizations.
IV. EXPERIMENTS: The study evaluates skin lesion segmentation on ISIC-2018 and lung field segmentation using JSRT, Montgomery, and NIH chest X-ray datasets.These experiments cover dermoscopy and frontal chest X-ray imaging modalities.
IV. EXPERIMENTS: Additional evaluations use Clean-CC-CCII CT images for COVID-19 pneumonia lesions, the Data Science Bowl dataset for nuclei, and GlaS for gland segmentation.The selected datasets include objects with varying sizes and textures, increasing segmentation difficulty and complexity.
IV. EXPERIMENTS: ISIC-2018 contains 2,596 annotated dermoscopy images, split into 2,076 training images and 520 testing images for lesion segmentation.The dataset supports automated melanoma-related segmentation applications.

5) GlaS:

The GlaS experiments evaluate TransAttUnet variants against attention-guided, multi-scale context, Transformer-based, and other segmentation baselines on a small dataset with complex objects.

Dataset: GlaS contains 165 colon histology images from 16 H&E-stained sections, with 85 images for training and 80 for testing.The dataset exhibits high inter-subject variability in stain distribution and tissue structure.
Method variants: TransAttUnet uses encoder-decoder guided attention and multi-scale skip connections to strengthen segmentation across semantic scales.The compared TransAttUnet C, D, and R variants differ in how decoder blocks are connected: cascade, dense, or residual operations.
Baselines: The evaluation compares TransAttUnet with attention-guided models, multi-scale context models, and Transformer-based baselines.The baselines include Attention U-Net, PraNet, Unet++, DoubleU-Net, MedT, Swin-Unet, and SegFormer, among others.
Evaluation: Mean Dice, IoU, accuracy, recall, and precision are used to quantify segmentation performance.These criteria are calculated pixel-wisely from true-positive, true-negative, false-positive, and false-negative values.

C. Experimental Results

TransAttUnet is evaluated against state-of-the-art baselines for skin-lesion and lung-field segmentation, combining quantitative comparisons with visual inspection of outputs and highlighted regions.

Skin lesion segmentation: On ISIC-2018, TransAttUnet achieves a DICE score of 90.74%, exceeding MCTrans at 90.35%.The comparison includes vanilla U-Net, attention-guided, multi-scale context, and Transformer-based approaches.
Skin lesion segmentation: On ISIC-2018, attention-guided and multi-scale context models outperform vanilla U-Net, while MCTrans reaches 90.35% DICE.Vanilla U-Net reports 67.40%, whereas Channel-UNet, PraNet, BCDU-Net, and DoubleU-Net report 84.82%, 87.46%, 85.10%, and 89.62%, respectively.
Lung-field segmentation: For lung-field segmentation, TransAttUnet obtains the highest DICE score of 98.88% on the combined JSRT, Montgomery, and NIH datasets.Its segmentation outputs are described as closer to the ground truths than those of the other baselines.
Lung-field segmentation: TransAttUnet improves over UNet from 96.17% to 98.88% DICE, a 2.71% improvement.The authors associate this improvement with encoder-decoder guided attention and multi-scale skip connections.
Visual comparison: Figure 4 compares quantitative results on ISIC 2018, the combined lung datasets, and Clean-CC-CCII, with boxes highlighting key lung regions.The visual comparison is intended to clarify differences between lung segmentation results and ground truths.

3) Evaluation on Pneumonia Lesion Segmentation:

TransAttUnet is tested for pneumonia-lesion segmentation on Clean-CC-CCII, where the authors report stronger quantitative performance than prior baselines and Transformer-based models.

Pneumonia lesion segmentation: Clean-CC-CCII is used to evaluate pneumonia-lesion segmentation from chest CT images with complex lesion textures and shapes.The experiment is intended to assess whether TransAttUnet is appropriate for this challenging segmentation task.
Quantitative results: TransAttUnet achieves the highest DICE score of 86.57% on Clean-CC-CCII, outperforming previous baselines.The authors report this result as evidence of the method’s segmentation performance on pneumonia lesions.
Variant comparison: TransAttUnet R outperforms TransAttUnet D with DICE scores of 86.57% versus 86.08%.The authors use this comparison to support residual decoder connections over dense decoder connections in this setting.
Baseline comparison: TransAttUnet outperforms Swin-Unet at 84.47% and SegFormer at 84.96% on the pneumonia-lesion task.The authors also report superiority over PraNet and note performance across lesions at different scales.

4) Evaluation on Nuclei Segmentation:

The experiments cover nuclei and gland segmentation, with TransAttUnet variants compared across datasets and visualizations to assess segmentation quality and feature representations.

Visual analysis: Figure 5 compares quantitative results for the Data Science Bowl and GLAS datasets, while Figure 6 visualizes decoder-stage feature maps from U-Net and TransAttUnet.The GLAS visualizations show better separation of glands from surrounding tissue for TransAttUnet.
Nuclei segmentation: On the 2018 Data Science Bowl dataset, TransAttUnet improves DICE from 75.76% for vanilla U-Net to 91.92%.The authors attribute the improvement to encoder-decoder guided attention and multi-scale skip connections.
Nuclei segmentation: TransAttUnet achieves 91.62% versus Attention U-Net at 90.93% on the Data Science Bowl comparison.The excerpt identifies TransAttUnet as superior to competitors from attention-guided and multi-scale context approaches.
Gland segmentation: On GLAS, TransAttUnet is evaluated for gland morphology quantification using comparisons with existing and Transformer-based segmentation baselines.The dataset contains multiple complex objects of interest and is described as small.
Gland segmentation: On GLAS, TransAttUnet outperforms MedT at 81.82%, Swin-Unet at 86.70%, and SegFormer at 87.36%.TransAttUnet D improves over the previous state-of-the-art KiU-Net by 5.86% in DICE and 8.85% in IoU.
Overall findings: Across the reported experiments, the authors state that TransAttUnet is effective and generalizable across varied medical image segmentation scenarios.The experiments span multiple datasets and imaging or object-segmentation settings.

D. Ablation Studies

Ablation experiments show that TransAttUnet’s attention and multi-scale skip-connection components contribute jointly to segmentation quality. Feature visualizations further indicate stronger contextual and semantic representations than vanilla U-Net, while the method remains computationally demanding.

Component ablations: Removing both TSA and GSA causes serious performance degradation, although the remaining multi-scale skip connections still outperform vanilla U-Net.This supports the effectiveness of decoder connections while showing the importance of the attention blocks.
Component ablations: Removing MSC lowers evaluation scores to 88.55%, 98.49%, 85.97%, 90.37%, and 87.83%, respectively.The reported drops support the contribution of multi-scale skip connections relative to the full model.
Component ablations: Removing TSA or GSA also degrades performance, indicating that the designed components complement and reinforce one another.The ablations compare TAU w/o TSA and TAU w/o GSA with the full model.
Feature visualization: Compared with vanilla U-Net, TransAttUnet benefits from long-range feature dependencies and global contextual information.Feature maps from decoder stages are visualized for both architectures to examine semantic discrimination.
Feature visualization: The multi-level non-local attention mechanisms provide more global semantic information to lower-stage decoders, generating more discriminative features.The comparison attributes the improved encoder representations to contextual guidance from the attention mechanisms.
Limitations: TransAttUnet relies heavily on global self-attention, resulting in a large memory footprint and high computation cost.The authors identify these costs as limitations and leave further improvements for future work.