Source-linked AI summary
TransDeepLab: Convolution-Free Transformer-based DeepLab v3+ for Medical Image Segmentation
Reza Azad, Moein Heidari, Moein Shariatnia, Ehsan Khodapanah Aghdam, Sanaz Karimijafarbigloo, Ehsan Adeli, Dorit Merhof
TL;DR
Medical image segmentation requires both local detail and long-range contextual modeling, but CNNs and standard Transformers each have important limitations. TransDeepLab addresses this by extending DeepLabv3+ with hierarchical Swin-Transformer components and multi-scale window-based processing. Across varied medical segmentation tasks, it performs superior or comparably to contemporary Vision Transformer and CNN-based methods while reducing model complexity.
Problem
CNNs capture local features but have restricted receptive fields, whereas Transformers model global context but can be computationally expensive and weak on low-level pixel information.
Method
TransDeepLab models DeepLabv3+ with a pure Transformer architecture using hierarchical Swin-Transformer blocks and varying window sizes for multi-scale representation.
Results
TransDeepLab performs superior or comparably to contemporary Vision Transformer and CNN-based methods across various medical image segmentation tasks.
Takeaways & Limitations
The results support using a pure Swin-Transformer-based DeepLab architecture to model long-range dependencies in medical image segmentation.
Abstract
from arXiv · showhide
Convolutional neural networks (CNNs) have been the de facto standard in a diverse set of computer vision tasks for many years. Especially, deep neural networks based on seminal architectures such as U-shaped models with skip-connections or atrous convolution with pyramid pooling have been tailored to a wide range of medical image analysis tasks. The main advantage of such architectures is that they are prone to detaining versatile local features. However, as a general consensus, CNNs fail to capture long-range dependencies and spatial correlations due to the intrinsic property of confined receptive field size of convolution operations. Alternatively, Transformer, profiting from global information modelling that stems from the self-attention mechanism, has recently attained remarkable performance in natural language processing and computer vision. Nevertheless, previous studies prove that both local and global features are critical for a deep model in dense prediction, such as segmenting complicated structures with disparate shapes and configurations. To this end, this paper proposes TransDeepLab, a novel DeepLab-like pure Transformer for medical image segmentation. Specifically, we exploit hierarchical Swin-Transformer with shifted windows to extend the DeepLabv3 and model the Atrous Spatial Pyramid Pooling (ASPP) module. A thorough search of the relevant literature yielded that we are the first to model the seminal DeepLab model with a pure Transformer-based model. Extensive experiments on various medical image segmentation tasks verify that our approach performs superior or on par with most contemporary works on an amalgamation of Vision Transformer and CNN-based methods, along with a significant reduction of model complexity. The codes and trained models are publicly available at https://github.com/rezazad68/transdeeplab
1 Introduction
CNN-based medical segmentation models capture local features effectively but are limited by restricted receptive fields, while Transformers model global context yet can be computationally costly and weaker on low-level pixel information. Swin-Transformer addresses these trade-offs through hierarchical local self-attention with linear complexity.
- CNN-based segmentation: CNN segmentation architectures capture contextual semantic information and reconstruct precise maps through encoder-decoder paths with skip connections.U-Net-style models downsample features in the encoder and combine deep and shallow representations in the decoder.
- CNN-based segmentation: Restricted convolutional receptive fields limit CNNs’ ability to model long-range features and global information.
- Transformer-based segmentation: Transformers capture explicit global context and long-range relations but impose quadratic computational cost and weakly represent low-level pixel information.
- Swin-Transformer: Swin-Transformer uses hierarchical local self-attention in non-overlapping windows to achieve linear complexity relative to standard Vision Transformers.
2 Proposed method
TransDeepLab is a pure Transformer-based DeepLabv3+ architecture that replaces conventional convolutional components with hierarchical Swin-Transformer blocks. Its shifted-window encoder, multi-scale pyramid, cross-attention fusion, and decoder preserve both contextual and spatial information.
- Architecture: TransDeepLab extends DeepLabv3+ as a pure Transformer-based encoder-decoder for medical image segmentation.Both encoder and decoder are constructed from Swin-Transformer blocks.
- Encoder: Swin-Transformer blocks encode local semantic and long-range contextual representations from non-overlapping image patches.The encoder progressively reduces spatial dimensions while increasing feature dimensions to form hierarchical representations.
- Swin-Transformer blocks: Shifted-window attention alternates regular and shifted window-based self-attention across consecutive Swin-Transformer blocks.This design replaces global self-attention while retaining the block’s LayerNorm, MLP, and nonlinear components.
- Multi-scale representation: The Swin Spatial Pyramid Pooling block models DeepLab’s ASPP using varying window sizes to capture local and global information at multiple scales.Smaller windows target local information, whereas larger windows extract global information.
- Feature fusion: A cross-attention module fuses pyramid features, while scale attention selectively emphasizes informative feature maps and channel representations.The module concatenates features across pyramid levels and uses learned scaling coefficients to weight them.
- Decoder: The decoder upsamples deep features, concatenates them with low-level features, and uses cascaded Swin-Transformer blocks to recover full resolution.Combining shallow and deep features helps reduce spatial-detail loss caused by downsampling.
3 Experiments
The Synapse multi-organ CT experiment uses 30 abdominal CT scans containing 3,779 axial contrast-enhanced images. Data partitioning and quantitative reporting follow prior work.
- Synapse multi-organ CT dataset: The Synapse dataset contains 30 abdominal CT scans and 3,779 axial contrast-enhanced clinical images.Each scan contains 85–198 slices of 512×512 pixels.
- Experimental protocol: The experiments follow prior work for data partitioning and quantitative result reporting.
Skin Lesion Segmentation.
Skin lesion segmentation experiments use ISIC 2017, ISIC 2018, and PH2 datasets, with PH2 split according to a prior experimental setting. Models are trained with augmentation and regularized optimization on a single GPU.
- Datasets: Skin lesion segmentation uses the ISIC 2017, ISIC 2018, and PH2 datasets.The ISIC datasets provide dermoscopy images with corresponding ground-truth annotations.
- Training setup: TransDeepLab is trained for 200 epochs with SGD, batch size 24, softmax Dice and cross-entropy losses, and L2 regularization.
- Training setup: Rotation and flipping are used as data augmentation methods to diversify the training set.
3.3 Evaluation Results
TransDeepLab was evaluated quantitatively and visually on medical image segmentation tasks, achieving strong Synapse performance and improved skin-lesion boundary and fine-detail capture.
- 80.16% DSC and 21.25% HD were achieved on the Synapse multi-organ CT dataset, the best reported performance in Table 1.
- 84.08% dice was obtained for Kidney(L), 61.19% for Pancreas, and 78.40% for Stomach on Synapse.
- Synapse visualizations showed correctly detected and classified organ instances with only slight contour variations.
- On skin-lesion segmentation, TransDeepLab performed better than competitors on most evaluation metrics and captured boundaries and fine-grained details more accurately.
Results of Skin Lesion Segmentation
The supplied passage identifies overfitting as a concern for heavy deep networks trained on small medical-image datasets.
- Heavy deep networks trained on small medical image datasets are prone to overfitting.
Model Complexity
TransDeepLab is presented as a lightweight model, while its supplied comparison reports lower parameter count than several competing architectures.
- TransDeepLab has 21.14M parameters, compared with 27.17M for Swin-UNet, 54.70M for DeepLab, and 105M for Trans-UNet.
- The authors report that TransDeepLab combines lower model complexity with dominant evaluation-metric performance.
3.4 Ablation study
The ablation study examines replacing the Transformer encoder with a CNN backbone while keeping the decoder and SSPP module unchanged.
- A CNN encoder with the same decoder and SSPP module produced sub-optimal performance relative to the proposed design.
- The ablation results indicate that the Transformer module helps TransDeepLab perform segmentation to a certain degree.
CNN vs Transformer Encoder.
The cross-attention module fuses multi-scale Swin-Transformer representations by modeling interactions among feature levels and improves localization, particularly near boundaries.
- CNN vs Transformer Encoder.: Cross-attention captures interactions among feature levels according to token informativeness across scales.It is compared with concatenation followed by a fully connected layer for basic scale fusion.
- CNN vs Transformer Encoder.: Multi-scale Transformer representations support long-range contextual dependency learning and more precise localization in boundary areas.Visualizations show predictions closer to the ground truth and aligned with the real situation.
SSPP influence.
The SSPP module probes multiple feature scales to improve context representation, with its benefits depending on the number of levels and input resolution.
- SSPP influence.: SSPP improves context patterning by probing features at multiple scales to obtain multi-scale information.The ablation is reported on the Synapse dataset.
- SSPP influence.: Two-level SSPP mostly improves Dice score, while three-level SSPP provides notable Hausdorff-distance performance.The reported benefits differ across these evaluation measures.
- SSPP influence.: Increasing SSPP levels should follow higher-resolution input images to maintain efficiency and comply with the module's resolution requirements.
4 Conclusion
TransDeepLab is a pure Transformer-based DeepLabv3+ architecture using Swin-Transformer components for medical image segmentation. Across varied medical segmentation tasks, it builds long-range dependencies and outperforms other state-of-the-art Vision Transformers in the reported experiments.
- 4 Conclusion: TransDeepLab models the encoder-decoder DeepLabv3+ architecture with Swin-Transformer components as its fundamental building blocks.
- 4 Conclusion: Across various medical image segmentation tasks, TransDeepLab effectively builds long-range dependencies and outperforms other state-of-the-art Vision Transformers in the experiments.