PaperlayerView exact revision on arXiv

Source-linked AI summary

AFTer-UNet: Axial Fusion Transformer UNet for Medical Image Segmentation

Xiangyi Yan, Hao Tang, Shanlin Sun, Haoyu Ma, Deying Kong, Xiaohui Xie

arXiv:2110.10403v1eess.IVcs.CVcs.LG

TL;DR

Medical image segmentation needs long-range 3D context, but single-slice transformers omit axial information and volumetric transformer processing is memory-intensive. AFTer-UNet combines 2D CNN feature extraction with axial fusion of intra-slice and inter-slice attention. Across three multi-organ datasets, it reports superior performance to current transformer-based models.

  • Problem

    Existing segmentation methods either omit axial-axis context by processing single slices or require memory-intensive 3D processing that limits long-range modeling.

  • Method

    AFTer-UNet uses a 2D CNN encoder-decoder with an axial fusion transformer that separately models and fuses within-slice and axial attention.

  • Results

    AFTer-UNet demonstrates superior performance over current transformer-based models across three multi-organ segmentation benchmarks.

  • Takeaways & Limitations

    The framework combines detailed CNN feature extraction with long-range intra-slice and inter-slice contextual modeling for 3D medical image segmentation.

Abstract

from arXiv · show

Recent advances in transformer-based models have drawn attention to exploring these techniques in medical image segmentation, especially in conjunction with the U-Net model (or its variants), which has shown great success in medical image segmentation, under both 2D and 3D settings. Current 2D based methods either directly replace convolutional layers with pure transformers or consider a transformer as an additional intermediate encoder between the encoder and decoder of U-Net. However, these approaches only consider the attention encoding within one single slice and do not utilize the axial-axis information naturally provided by a 3D volume. In the 3D setting, convolution on volumetric data and transformers both consume large GPU memory. One has to either downsample the image or use cropped local patches to reduce GPU memory usage, which limits its performance. In this paper, we propose Axial Fusion Transformer UNet (AFTer-UNet), which takes both advantages of convolutional layers' capability of extracting detailed features and transformers' strength on long sequence modeling. It considers both intra-slice and inter-slice long-range cues to guide the segmentation. Meanwhile, it has fewer parameters and takes less GPU memory to train than the previous transformer-based models. Extensive experiments on three multi-organ segmentation datasets demonstrate that our method outperforms current state-of-the-art methods.

1. Introduction

Medical image segmentation is clinically important but labor-intensive, while existing CNN and transformer approaches struggle to capture long-range 3D context efficiently. AFTer-UNet addresses this gap by fusing intra-slice and inter-slice information with an axial fusion mechanism.

  • 3D CT segmentation is time-consuming because volumes can contain hundreds of 2D slices, motivating robust automated tools.
  • CNN-based methods inadequately encode long-range interactions within slices and across neighboring slices, which are useful for anatomically variable organs.
  • 2D transformer methods model within-slice dependencies but omit axial-axis relationships naturally available in 3D medical images.
  • 3D transformer approaches face high memory demands, so local patches and restricted attention can lose information and limit long-range modeling.
  • AFTer-UNet fuses intra-slice and inter-slice information using an axial fusion transformer between a 2D CNN encoder and decoder.
  • The axial fusion mechanism separately computes axial and within-slice attention before fusing them, reducing 3D self-attention complexity.
  • Experiments on three multi-organ benchmarks report superior performance over current transformer-based models.

2. Related work

Medical segmentation evolved from atlas-based and CNN-based methods toward transformer-enhanced U-Net variants. Existing approaches either omit axial information by processing single slices or face computational costs when modeling 3D volumes directly.

  • TransUNet and Swin-Unet process single slices, leaving axial-axis information from 3D volumes unused.
  • Direct 3D self-attention is computationally expensive, motivating CoTr's deformable attention but introducing 3D-patch inputs and associated information loss concerns.

3. Method

AFTer-UNet retains a U-Net structure while inserting an axial fusion transformer to encode high-level context both within slices and across neighboring slices.

  • The architecture uses a 2D CNN encoder for fine-level features, an axial fusion transformer for high-level contextual information, and a 2D CNN decoder for pixel-level segmentation.

3.1. CNN encoder

The CNN encoder processes neighboring axial slices with U-Net-like convolutional blocks, then passes the final high-level feature map group to the axial fusion transformer under GPU-memory constraints.

  • Each input slice samples neighboring slices along the axial axis using a specified frequency and number of neighbors.
  • The CNN encoder follows U-Net with max-pooling blocks and two Conv2d-ReLU pairs augmented by instance normalization.
  • Encoder feature maps are indexed by feature level, with spatial dimensions reduced according to the level and channel count represented by Cb.
  • Only the final encoder feature map group is fed to the axial fusion transformer, combining higher-level semantics with a smaller memory footprint.

3.2. Axial fusion transformer encoder

The axial fusion transformer encodes high-level information within each slice and across neighboring slices, using feature maps produced by a CNN encoder. It separates axial and intra-slice attention to reduce computation while retaining both types of context.

  • Feature maps as input embeddings: CNN-extracted feature maps are directly fed into the axial fusion transformer, avoiding the patch partitioning and linear projection used in the original ViT setup.The smaller feature-map dimensions allow the model to use information from the whole image rather than a single local patch.
  • Feature maps as input embeddings: Learnable positional embeddings encode both locations within each feature map and positions among feature maps in the neighboring slice group.The resulting sequence is used as the transformer input.
  • Query-key-value matrices and self-attention: Each transformer block derives query, key, and value vectors from normalized representations and combines multiple attention heads through projection, MLP, normalization, and residual connections.The attention-head dimension is defined as Ch = CL/A.
  • Fusing axial information: Full 3D self-attention is memory-intensive, while single-slice attention misses neighboring-slice information that can affect segmentation quality.The paper reports that using fewer neighboring slices can produce poorer results.
  • Fusing axial information: Axial fusion applies inter-slice and intra-slice attention sequentially, combining contextual information from both axes before the MLP produces the block output.The fused encoding retains the feature-map group structure for subsequent decoding.
  • Fusing axial information: The axial strategy requires (HL · WL) + NA comparisons per vector instead of the (HL · WL) · NA comparisons required by the full self-attention model.This reduces the attention computation while preserving attention over spatial and axial dimensions.

3.3. CNN decoder

The CNN decoder mirrors the U-Net decoding path, progressively restoring resolution and refining the fused transformer features. Skip connections preserve low-level encoder details for segmentation.

  • CNN decoder: The decoder is mostly symmetric to the CNN encoder and uses Conv2d-ReLU blocks connected by upsampling layers with scale factor 2.Each block contains two Conv2d-ReLU pairs with instance normalization between convolution and ReLU.
  • CNN decoder: The decoder progressively upsamples the final fused transformer feature map to the input resolution and refines it with Conv2d-ReLU blocks.The fused representation from the last axial fusion transformer block is the decoder input.
  • CNN decoder: Skip connections between encoder and decoder retain low-level details for segmentation.This follows the U-Net design.
  • CNN decoder: For sampled neighboring slice groups, the network outputs segmentation map groups whose channels represent the organ classes.The final 3D prediction is formed after retaining the middle segmentation map and concatenating the resulting groups.
  • CNN decoder: The training loss is the sum of Dice loss and cross-entropy loss.

4. Experiments

Experiments evaluate AFTer-UNet on three multi-organ CT benchmarks, compare it with convolutional and transformer-based methods, and examine axial-fusion design choices, qualitative predictions, and memory use.

  • Setup: Experiments use BCV, Thorax-85, and SegTHOR, reporting average Sørensen–Dice coefficient (DSC) across abdominal or thoracic organs.BCV uses 18 training and 12 test scans; Thorax-85 uses 60 training and 25 test scans; SegTHOR uses 30 training and 10 validation scans.
  • Results: Long-range modeling benefits anatomically variable organs: 3D and transformer-based models achieve consistently higher DSC for esophagus and trachea, while results are comparable for large regular organs.AFTer-UNet combines 2D convolutions for fine details with axial fusion for inter-slice and intra-slice information.
  • Results on public datasets: AFTer-UNet outperforms U-Net by 4.34% and CoTr by 2.56% on BCV, while improving SegTHOR esophagus DSC by 2.91% over U-Net and 1.45% over CoTr.Across the public datasets, it further outperforms U-Net by 2.13% and CoTr by 0.69%.
  • Ablation study: Increasing the number of neighboring axial slices improves Dice scores, with more obvious gains for elongated organs such as esophagus and trachea.The model uses NA = 8 neighboring axial slices.
  • Ablation study: Increasing axial-fusion transformer layers improves average Dice scores, especially for esophagus and trachea, indicating stronger gains from compounding inter-slice cues.The evaluated layer settings are L = 1, 2, 4, and 6.
  • Ablation study: Lower sampling frequency improves performance because it provides denser inter-slice information and emphasizes nearer neighboring slices.The experiments compare Nf = 1, 2, and 4; the model uses Nf = 1.
  • Qualitative results: AFTer-UNet produces more consecutive and accurate esophagus predictions than previous methods in qualitative Thorax-85 results.The comparison includes ground truth, previous approaches, and AFTer-UNet, with orange rectangles highlighting model effectiveness.
  • Memory consumption and model parameters: Axial fusion yields dramatic computational savings, enabling AFTer-UNet training on a single RTX-2080Ti with 11GB memory.AFTer-UNet has 41.5M parameters, compared with 43.5M for TransUNet and 41.9M for CoTr.

5. Conclusion

AFTer-UNet is an end-to-end segmentation framework that fuses intra-slice and inter-slice context to guide segmentation, with effectiveness demonstrated on three datasets.

  • AFTer-UNet uses an axial fusion mechanism to combine intra-slice and inter-slice contextual information for final segmentation.
  • Experiments on three datasets demonstrate the model’s effectiveness compared with previous work.
Loading 2110.10403v1…