Source-linked AI summary
A Robust Volumetric Transformer for Accurate 3D Tumor Segmentation
Himashi Peiris, Munawar Hayat, Zhaolin Chen, Gary Egan, Mehrtash Harandi
TL;DR
Volumetric segmentation must preserve spatial information across all axes while modeling both local and global cues. VT-UNet addresses this with a purely transformer-based 3D encoder-decoder using windowed attention and encoder-guided decoder attention. The model is reported to achieve consistent improvements over existing methods, with computational efficiency and robustness to artefacts, while the complexity analysis uses simplifying assumptions.
Problem
The paper addresses the challenge of segmenting 3D medical volumes while preserving inter-slice dependencies and spatial information across the volume.
Method
VT-UNet uses a purely transformer-based volumetric encoder-decoder with self-attention, 3D patch merging, and decoder cross-attention from encoder keys and values.
Results
The model achieves consistent improvements over existing state-of-the-art volumetric segmentation methods and learns representations robust against artefacts.
Takeaways & Limitations
The proposed network is presented as a computationally efficient approach for volumetric medical image segmentation that can handle large 3D volumes.
Takeaways & Limitations
The computational-complexity analysis assumes equal key, value, and token-channel dimensions, omits softmax FLOPs, and abstracts multi-head attention.
Abstract
from arXiv · showhide
We propose a Transformer architecture for volumetric segmentation, a challenging task that requires keeping a complex balance in encoding local and global spatial cues, and preserving information along all axes of the volume. Encoder of the proposed design benefits from self-attention mechanism to simultaneously encode local and global cues, while the decoder employs a parallel self and cross attention formulation to capture fine details for boundary refinement. Empirically, we show that the proposed design choices result in a computationally efficient model, with competitive and promising results on the Medical Segmentation Decathlon (MSD) brain tumor segmentation (BraTS) Task. We further show that the representations learned by our model are robust against data corruptions. \href{https://github.com/himashi92/VT-UNet}{Our code implementation is publicly available}.
1 Introduction
VT-UNet addresses the challenge of segmenting intact 3D medical volumes without losing inter-slice information. It combines transformer-based volumetric modeling with computational efficiency, strong segmentation performance, and robustness to artefacts.
- Research motivation: Prior 3D transformer segmentation methods often process volumes as 2D slices, potentially losing critical inter-slice dependencies.The paper identifies preserving volumetric information across all three views as a central challenge.
- Research gap: Purely transformer-based architectures that retain intact volumetric input were previously unexplored, motivating VT-UNet.Hybrid convolution-transformer approaches preserve 3D data, but the paper targets a convolution-free alternative.
- Contributions: VT-UNet is designed to limit model parameters and FLOPs while maintaining competitive volumetric segmentation performance.The introduction links these efficiency goals to the proposed architectural choices and Fig. 1.
- Contributions: The model achieves state-of-the-art volumetric segmentation results and enhanced robustness to data artefacts in the reported evaluations.The introduction presents both segmentation quality and corruption robustness as primary empirical outcomes.
2 Methodology
VT-UNet uses volumetric transformer blocks, hierarchical token processing, and a decoder combining self-attention with encoder-guided cross-attention. Windowing and patch merging reduce the computational burden while Fourier positional encoding supplements spatial information.
- Architecture: VT-UNet is a purely transformer-based encoder-decoder architecture without convolution modules in its main design.The architecture includes volumetric shifted windows, an encoder-decoder structure, and a fusion module.
- Windowed attention: The encoder alternates regular and shifted volumetric windows, shifting windows by (2, 2, 2) tokens in the illustrated configuration.For an 8×8×8 volume with 4×4×4 windows, the example changes from 8 regular windows to 27 shifted windows.
- VT Encoder: The encoder represents 3D inputs as tokens using patch partitioning, linear embedding, 3D patch merging, and successive volumetric transformer blocks.The input is a D×H×W×C volume and the output is a D×H×W×K voxel-label volume.
- VT Encoder: Patch merging concatenates each non-overlapping 2×2 token group and linearly projects it to doubled channel dimensionality, creating feature hierarchies.These hierarchies support finer dense-prediction details and reduce FLOPs by a factor of 16 after each encoder block.
- VT Decoder: Each decoder block combines self-attention on decoder tokens with cross-attention using encoder keys and values from the corresponding stage.The cross-attention branch uses decoder queries and encoder-derived keys and values, with regular and shifted windowing.
- Fusion and positional encoding: The fusion module linearly combines self-attention and cross-attention outputs, with α = 0.55 controlling their contribution and Fourier positional encoding added to the self-attention tokens.The fused tokens are passed to the next decoder block.
- Computational complexity: Windowed self-attention lowers computational load because attention is applied within groups of tokens rather than across the full volumetric sequence.The paper notes that volumetric inputs have far more tokens than windows, making windowing especially useful.
3 Related Work
Transformer-based medical segmentation has progressed from hybrid 2D-slice approaches toward purely transformer-based designs, but preserving full 3D volumetric structure remained unexplored.
- TransUNet replaces a CNN bottleneck with a ViT while processing 2D slices of 3D volumes.
- Swin-UNet is a purely transformer-based medical segmentation network that reports better segmentation results than TransUNet.
- TransBTS combines a CNN encoder-decoder with a Transformer bottleneck for 3D medical image segmentation.
- nnFormer uses 3D Swin Transformer blocks in its encoder and decoder with interleaved convolutional stems.
- The computational-complexity derivation assumes C_k = C_v = C, omits softmax FLOPs, and notes that multi-head attention lowers computation across smaller head dimensions.
4 Experiments
Experiments evaluate VT-UNet on MSD BraTS using segmentation, qualitative, robustness, and ablation analyses, with comparisons against transformer and CNN baselines.
- Implementation Details: The study uses 484 MSD BraTS MRI scans split into 80% training, 15% validation, and 5% testing sets.Training uses AdamW with a 1e−4 learning rate for 1000 epochs and batch size 4 on one Nvidia A40 GPU.
- Segmentation Results: Table 1 evaluates Enhancing Tumor, Tumor Core, and Whole Tumor using Dice Sørensen coefficient and Hausdorff Distance.
- Segmentation Results: VT-UNet achieves the best overall performance in DSC and HD among the compared transformer-based and CNN-based methods.
- Qualitative Segmentation Results: Table 4 presents qualitative results for tumor-component colors, volumetric tumor predictions, and segmentation boundaries.
- Robustness Analysis: Robustness analysis synthetically adds motion, ghosting, and spike artefacts to MRI images at inference time and compares VT-UNet with nnFormer.
- Ablation Study: The ablation results indicate that Parallel CA, SA, and FPE components each contribute to model performance.
5 Conclusion
The paper concludes that VT-UNet is computationally efficient for large 3D volumes and learns representations robust to artefacts, with consistent improvements in volumetric segmentation.
- VT-UNet is presented as a computationally efficient volumetric transformer for large-sized 3D medical image volumes.
- The learned representations are reported to be robust against artefacts.
- Results show consistent improvements over existing state-of-the-art methods in volumetric segmentation.
6 Supplementary Material
The supplementary material explains self-attention, volumetric token processing, decoder positional encoding, and the joint loss used to train VT-UNet.
- Preliminaries: Self-attention generates new tokens as weighted combinations of input tokens to better represent the signal.
- Preliminaries: The attention formulation learns value, key, and query projections from input tokens and uses softmax-based similarities to define attention weights.
- Breaking Symmetry Cont.: Three-dimensional Fourier positional encoding supplements decoder tokens to preserve information across continuous slices and improve feature conditioning.
- Loss Function: VT-UNet jointly minimizes Dice Loss and voxel-wise Cross Entropy loss during training.