PaperlayerView exact revision on arXiv

Source-linked AI summary

SegFormer3D: an Efficient Transformer for 3D Medical Image Segmentation

Shehan Perera, Pouyan Navard, Alper Yilmaz

arXiv:2404.10156v2cs.CV

TL;DR

High-performing Transformer approaches to 3D segmentation can be computationally demanding, while limited medical datasets pose generalization challenges. SegFormer3D uses a hierarchical multiscale Transformer and an all-MLP decoder. The authors report competitive results across three datasets alongside substantial reductions in model size and complexity.

  • Problem

    High-performing Transformer architectures for 3D medical segmentation are computationally complex, and limited datasets can challenge generalization and convergence.

    Source: §1 · p. 2, §4.4 · p. 7

  • Method

    SegFormer3D combines hierarchical multiscale feature encoding and compressed self-attention with an all-MLP decoder for volumetric segmentation.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    Across BraTS, Synapse, and ACDC, SegFormer3D is reported competitive; its 4.5 million parameters and 17 GFLOPS accompany claimed 34× parameter-count and 13× model-complexity reductions vs SOTA.

    Source: §1 · p. 2, §4.1 · p. 6, §4.2 · p. 6

  • Takeaways & Limitations

    The authors argue that lightweight ViTs are a valuable research direction for medical imaging, reporting competitive results without pretraining.

    Source: §1 · p. 2

Abstract

from arXiv · show

The adoption of Vision Transformers (ViTs) based architectures represents a significant advancement in 3D Medical Image (MI) segmentation, surpassing traditional Convolutional Neural Network (CNN) models by enhancing global contextual understanding. While this paradigm shift has significantly enhanced 3D segmentation performance, state-of-the-art architectures require extremely large and complex architectures with large scale computing resources for training and deployment. Furthermore, in the context of limited datasets, often encountered in medical imaging, larger models can present hurdles in both model generalization and convergence. In response to these challenges and to demonstrate that lightweight models are a valuable area of research in 3D medical imaging, we present SegFormer3D, a hierarchical Transformer that calculates attention across multiscale volumetric features. Additionally, SegFormer3D avoids complex decoders and uses an all-MLP decoder to aggregate local and global attention features to produce highly accurate segmentation masks. The proposed memory efficient Transformer preserves the performance characteristics of a significantly larger model in a compact design. SegFormer3D democratizes deep learning for 3D medical image segmentation by offering a model with 33x less parameters and a 13x reduction in GFLOPS compared to the current state-of-the-art (SOTA). We benchmark SegFormer3D against the current SOTA models on three widely used datasets Synapse, BRaTs, and ACDC, achieving competitive results. Code: https://github.com/OSUPCVLab/SegFormer3D.git

1. Introduction

The introduction frames 3D volumetric segmentation as an important medical-imaging task and identifies efficiency and limited-data challenges for Transformer models. It presents SegFormer3D as a lightweight hierarchical Transformer with multiscale features and an all-MLP decoder.

  • 3D volumetric segmentation is described as a fundamental medical-image-analysis task, with tumor and multi-organ localization among its diagnostic and treatment applications.Source: §1 · p. 1
  • ViTs capture global relationships, while prior medical segmentation work combined Transformer encoding with CNN feature refinement in decoder stages.The introduction contrasts this global attention with CNNs’ local inductive bias.Source: §1 · p. 1
  • Long 3D input sequences increase Transformer computational cost, and medical imaging often lacks large pretraining datasets while small datasets challenge generalization.The introduction also notes that repetitive medical-image structures may be compressible.Source: §1 · p. 2
  • SegFormer3D encodes input volumes at multiple scales and compresses the embedded sequence to a fixed ratio with efficient self-attention.Its design follows a hierarchical Transformer and aims to capture coarse-to-fine features while reducing model complexity.Source: §1 · p. 2
  • Overlapping patch embeddings preserve voxel continuity, while positional-free encoding addresses resolution mismatch and an all-MLP decoder aggregates features for segmentation.The introduction reports evaluation on Synapse, ACDC, and BraTS.Source: §1 · p. 2
  • 4.5 million parameters and 17 GFLOPS accompany reported 34× parameter-count and 13× model-complexity reductions vs SOTA, with competitive results without pretraining.The authors present lightweight ViTs as a valuable research direction for medical imaging.Source: §1 · p. 2

2. Related Work

Related work covers convolutional architectures for medical-image segmentation and Transformer–convolution hybrids that combine global context with local feature processing. The cited hybrid models are described as substantially complex in model size and computation.

  • CNN segmentation work includes 3D U-Net, V-Net, and nnU-Net, alongside hierarchical feature extraction and deformable convolutions for irregular shapes.The passage also describes downsampling volumes and extending 2D convolutions to 3D.Source: §2 · p. 3
  • Transformer–convolution hybrids include U-shaped CNN decoders, Transformer bottlenecks, and large-kernel or dual-encoding designs for spatial dependencies.Examples named include TransUNet, UNETR, LoGoNet, nnFormer, and TransBTS.Source: §2 · p. 3
  • The cited Transformer–convolution methods are characterized as having considerable model and computational complexity.Source: §2 · p. 3

3. Method

SegFormer3D combines a hierarchical volumetric Transformer with efficient attention and a streamlined linear-layer decoder. The paper reports lower model size and complexity versus SOTA while retaining competitive BRaTS performance.

3. Method

  • 4.5 million parameters and 17 GFLOPS accompany reported reductions of 34× in parameter count and 13× in complexity versus SOTA.Table 1 compares SegFormer3D with SOTA models by size (M) and complexity.Source: §3 · p. 3, Table 1 · p. 3
  • Overlapped patch merging retains neighborhood transitions, while mix-ffn learns positional cues without fixed encodings for resizing volumetric inputs.Source: §3 · pp. 3–4
  • For long 3D sequences, efficient attention lowers complexity from O(n^2) to O(n^2/r), with r set to 4×, 2×, 1×, and 1× across four encoder stages.A 128³ MRI volume produces a sequence length of 32,768, and the original attention operation has O(n^2) complexity.Source: §3 · p. 3
  • BRaTS Table 2 ranks average performance across classes and reports category-wide outperformance; Fig. 2 compares qualitative segmentations across solutions.Source: Table 2 · p. 4, Figure 2 · p. 4
  • Instead of successive 3D convolutions, the decoder uses linear layers; standardized stage features are upsampled, concatenated, fused, and projected into masks.Source: §3 · p. 4

4. Experimental Results

SegFormer3D was evaluated without external pretraining on BraTS, Synapse, and ACDC using the same datasets and evaluation methods as comparison architectures. It was competitive across these benchmarks while using a lightweight design.

4. Experimental Results

  • All three datasets were evaluated without external pretraining, using the same datasets and evaluation methods as the comparison architectures.Experiments used a single Nvidia RTX 3090 GPU; models shared a learning-rate policy, AdamW optimization, and 1000-epoch training.Source: §4 · pp. 4–5

4.1. Results on Brain Tumor Segmentation (BraTs)

  • On BraTS, SegFormer3D was a competitive lightweight architecture against larger CNN and Transformer models, with 4.5 million parameters and 17.5 GFLOP complexity.The reported measures cover whole tumor, enhancing tumor, and tumor core; the authors also describe its qualitative results as superior to established architectures.Source: §4.1 · pp. 5–6

4. Experimental Results

  • SegFormer3D was competitive on BraTS, ranked second only to nnFormer on Synapse, and was within 1% of ACDC SOTA.The Synapse comparison ranks average performance across classes; the ACDC comparison is against models averaging 34× more parameters and 13× more computational complexity.Source: Table 3 · p. 5, §4.2 · p. 6, §4.1 · p. 6

4.2. Results on Multi-Organ CT Segmentation (Synapse)

  • On Synapse, SegFormer3D ranked second only to nnFormer and produced competitive results with 4.5M parameters.The results cover a complex multi-class organ-segmentation task; the table caption describes rankings by average performance across classes.Source: Table 3 · p. 5, §4.2 · p. 6

4.3. Results on Automated cardiac diagnosis (ACDC)

  • On ACDC, SegFormer3D was within 1% of SOTA performance, without large-scale pretraining on the small dataset.The comparison follows the training and inference pipelines specified in [32] and reports Dice accuracy.Source: Table 4 · p. 6, §4.3 · p. 7

4.4. Conclusion

  • The conclusion argues that lightweight, efficient architectures can broaden access to 3D medical segmentation with minimal computational resources and no additional pretraining.It presents developing high-performance lightweight models as a research direction with practical relevance to medical imaging.Source: §4.4 · p. 7
Loading 2404.10156v2…