Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

GASA-UNet: Global Axial Self-Attention U-Net for 3D Medical Image Segmentation

Chengkun Sun, Russell Stevens Terry, Jiang Bian, Jie Xu

TL;DR

GASA-UNet is introduced, a refined U-Net-like model featuring a novel Global Axial Self-Attention (GASA) block that has demonstrated promising improvements in segmentation performance, particularly for smaller anatomical structures, as evidenced by enhanced Dice scores and Normalized Surface Dice (NSD) on three benchmark datasets.

Abstract

Accurate segmentation of multiple organs and the differentiation of pathological tissues in medical imaging are crucial but challenging, especially for nuanced classifications and ambiguous organ boundaries. To tackle these challenges, we introduce GASA-UNet, a refined U-Net-like model featuring a novel Global Axial Self-Attention (GASA) block. This block processes image data as a 3D entity, with each 2D plane representing a different anatomical cross-section. Voxel features are defined within this spatial context, and a Multi-Head Self-Attention (MHSA) mechanism is utilized on extracted 1D patches to facilitate connections across these planes. Positional embeddings (PE) are incorporated into our attention framework, enriching voxel features with spatial context and enhancing tissue classification and organ edge delineation. Our model has demonstrated promising improvements in segmentation performance, particularly for smaller anatomical structures, as evidenced by enhanced Dice scores and Normalized Surface Dice (NSD) on three benchmark datasets, i.e., BTCV, AMOS, and KiTS23.

GASA-UNet: Global Axial Self-Attention U-Net for 3D Medical Image Segmentation

TL;DR

GASA-UNet is introduced, a refined U-Net-like model featuring a novel Global Axial Self-Attention (GASA) block that has demonstrated promising improvements in segmentation performance, particularly for smaller anatomical structures, as evidenced by enhanced Dice scores and Normalized Surface Dice (NSD) on three benchmark datasets.

Abstract

Accurate segmentation of multiple organs and the differentiation of pathological tissues in medical imaging are crucial but challenging, especially for nuanced classifications and ambiguous organ boundaries. To tackle these challenges, we introduce GASA-UNet, a refined U-Net-like model featuring a novel Global Axial Self-Attention (GASA) block. This block processes image data as a 3D entity, with each 2D plane representing a different anatomical cross-section. Voxel features are defined within this spatial context, and a Multi-Head Self-Attention (MHSA) mechanism is utilized on extracted 1D patches to facilitate connections across these planes. Positional embeddings (PE) are incorporated into our attention framework, enriching voxel features with spatial context and enhancing tissue classification and organ edge delineation. Our model has demonstrated promising improvements in segmentation performance, particularly for smaller anatomical structures, as evidenced by enhanced Dice scores and Normalized Surface Dice (NSD) on three benchmark datasets, i.e., BTCV, AMOS, and KiTS23.
Paper Structure (16 sections, 3 equations, 4 figures, 4 tables)

This paper contains 16 sections, 3 equations, 4 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Multi-Head Self-Attention
  4. Medical Vision Transformer
  5. Axial Attention
  6. GASA-UNet
  7. GASA Block
  8. GASA-UNet
  9. Experiments
  10. Datasets
  11. Experimental Setup
  12. Segmentation Results
  13. Ablation Experiments
  14. Computational Load
  15. Discussion
  16. ...and 1 more sections

Figures (4)

  • Figure 1: Comparison of Axial Self-Attention blocks: (a) Traditional Axial Self-Attention block, and (b) Proposed GASA block.
  • Figure 2: Schematic illustration of the GASA block. The input features undergo three $2\times2$ convolutions along the W, H, and D axes. The resulting patches are concatenated and processed by the MHSA module with adjustable heads and dimensions. Each output attention value is expanded into a specific slice and then into the entire 3D feature space. Following channel concatenation, these axial attention values are combined to form the GASA feature, which is further enhanced with absolute positional embedding. The input feature is integrated with the GASA feature at the channel level. Notably, the number of GASA feature channels increases from $n$ to $3n$ through a manual design process.
  • Figure 3: Visual comparison of segmentation results by different axial attentions on representative samples from BTCV, AMOS, and KiTS23 datasets. The +MAT results on AMOS were omitted due to its inability to process asymmetric 3D features.
  • Figure 4: Overview of the three PE strategies and visual comparison on selected samples using three PE strategies.