Source-linked AI summary
LeViT-UNet: Make Faster Encoders with Transformer for Medical Image Segmentation
Guoping Xu, Xingrong Wu, Xuan Zhang, Xinwei He
TL;DR
Medical image segmentation needs better global-context modeling than local CNN operations provide, while transformer-based approaches remain computationally demanding. LeViT-UNet uses a LeViT transformer-convolution encoder within U-Net and reuses multi-scale features through skip connections. It reports competitive or superior performance on Synapse and ACDC while trading off accuracy and efficiency.
Problem
CNN-based segmentation is limited in capturing global context and long-range relations, while transformer-based segmentation requires further efficiency improvements because of high computational complexity.
Method
LeViT-UNet integrates a LeViT hybrid transformer-convolution encoder with a U-Net decoder and passes multi-scale encoder features through skip connections.
Results
LeViT-UNet achieves competitive performance against fast CNN-based models, and LeViT-UNet-384 surpasses previously reported transformer-based methods including TransUNet and Swin-UNet.
Takeaways & Limitations
The study demonstrates a transformer-based segmentation architecture that balances accuracy and efficiency across Synapse and ACDC evaluations.
Takeaways & Limitations
The input resolution is down-scaled to 1/16 before transformer blocks, which may affect segmentation performance, and further efficiency optimization remains needed against fast CNN methods.
Abstract
from arXiv · showhide
Medical image segmentation plays an essential role in developing computer-assisted diagnosis and therapy systems, yet still faces many challenges. In the past few years, the popular encoder-decoder architectures based on CNNs (e.g., U-Net) have been successfully applied in the task of medical image segmentation. However, due to the locality of convolution operations, they demonstrate limitations in learning global context and long-range spatial relations. Recently, several researchers try to introduce transformers to both the encoder and decoder components with promising results, but the efficiency requires further improvement due to the high computational complexity of transformers. In this paper, we propose LeViT-UNet, which integrates a LeViT Transformer module into the U-Net architecture, for fast and accurate medical image segmentation. Specifically, we use LeViT as the encoder of the LeViT-UNet, which better trades off the accuracy and efficiency of the Transformer block. Moreover, multi-scale feature maps from transformer blocks and convolutional blocks of LeViT are passed into the decoder via skip-connection, which can effectively reuse the spatial information of the feature maps. Our experiments indicate that the proposed LeViT-UNet achieves better performance comparing to various competing methods on several challenging medical image segmentation benchmarks including Synapse and ACDC. Code and models will be publicly available at https://github.com/apple1986/LeViT_UNet.
1 Introduction
CNN-based segmentation methods struggle to capture global context and long-range relations because convolutions are local. LeViT-UNet addresses this with a faster transformer encoder and multi-scale feature reuse, achieving competitive or superior performance across comparisons.
- CNN-based approaches have insufficient capability to capture explicit global context and long-range relations because convolution operations are intrinsically local.
- These limitations are especially relevant for objects with large inter-patient variation in shape, scale, and texture.
- LeViT-UNet uses LeViT as a fast encoder and aims to improve the speed and accuracy of transformer-based medical image segmentation.
- LeViT-UNet achieves competitive performance against fast CNN-based models, while LeViT-UNet-384 surpasses TransUNet and Swin-UNet in the reported comparison.
- The architecture passes multi-scale feature maps from transformer and convolutional blocks to the decoder through skip connections.
2 Related Works
Related work progresses from conventional CNN-based U-Net variants and CNN attention mechanisms toward transformer architectures for vision and medical segmentation.
- CNN-based methods: CNN-based methods, including U-Net and its variants, have been extensively studied for medical image segmentation.
- Self-attention mechanisms to complement CNNs: Self-attention modules have been integrated into CNNs to model channel-wise, spatial, or semantic interdependencies.
- Transformers: Vision transformers and transformer-based segmentation models such as TransUNet and Swin-UNet apply self-attention or hybrid transformer-CNN designs to segmentation.
3 Method
LeViT-UNet replaces the conventional U-Net encoder with a LeViT hybrid of convolutional and transformer blocks, while retaining a convolutional decoder and skip-connected multi-scale features.
- Overall architecture: LeViT-UNet applies a LeViT module in the encoder to extract long-range structural information while retaining the U-Net decoder design.
- LeViT as Encoder: The LeViT encoder combines convolutional blocks for resolution reduction with transformer blocks designed to reduce transformer floating-point operations.
- LeViT as Encoder: The encoder uses four 3x3 convolutions with stride 2 before transformer processing.
- LeViT as Encoder: Transformer blocks contain MLP and multi-head attention modules with batch normalization and residual connections.
- LeViT as Encoder: The attention computation uses query, key, and value matrices, while attention bias supplies positional information within each block.
- Decoder: Decoder features are combined through skip connections and cascaded CNN upsampling blocks to recover the full H x W resolution.
4.1 Dataset
The evaluation uses Synapse abdominal CT and ACDC cardiac MR datasets, with separate training and validation or testing splits.
- Synapse: Synapse contains 30 abdominal CT scans and 3779 axial images, using 18 cases for training and 12 for validation.
- Synapse: Synapse performance is evaluated using average Dice Similarity Coefficient and Hausdorff Distance across eight abdominal organs.
- ACDC: ACDC contains cine-MR data from 150 patients, with 100 annotated volumes split into 80 training samples and 20 testing samples.
4.2 Implementation details
Experiments used standardized software, augmentation, hardware, image resolution, optimization, and loss settings.
- Training used random flipping and rotations to increase data diversity.
- All models were trained on 224x224 images using an Nvidia 3090 GPU with 24GB memory.
- The experiments used Python 3.8, PyTorch 1.8.0, and Ubuntu 18.04.1 LTS.
- Training used Adam with learning rate 1e-5 and weight decay 1e-4.
- The objective combined cross entropy and Dice loss.
4.3 Experiment results on Synapse dataset
On Synapse, LeViT-UNet was evaluated against state-of-the-art and efficient segmentation methods using accuracy, boundary quality, speed, and ablations. Results favored LeViT-UNet-384 overall, while ablations examined transformer blocks, skip connections, and pretraining.
- Compare state-of-the-art methods: LeViT-UNet variants were evaluated against state-of-the-art methods using average DSC and HD for generalization and robustness.The variants were LeViT-UNet-128s, LeViT-UNet-192, and LeViT-UNet-384.
- Compare state-of-the-art methods: LeViT-UNet-384 achieved the best average HD at 16.84 mm on the Synapse multi-organ CT dataset.This was reported as an improvement of about 14.8 mm and 4.7 mm over recently state-of-the-art methods.
- Compare state-of-the-art methods: LeViT-UNet outputs were relatively smoother, while TransUNet, UNet, and DeepLabv3+ more often under-segmented or over-segmented organs.Examples included stomach under-segmentation by TransUNet and DeepLabv3+, and over-segmentation by UNet.
- Compare with fast segmentation methods: LeViT-UNet-384 achieved 78.53% mDSC and 16.84mm mHD, reported as best among the methods in Table 2.The comparison also considered parameter count and inference speed against fast segmentation methods.
- Compare with fast segmentation methods: LeViT-UNet had fewer parameters than TransUNet but still needed improvement relative to CGNet, ContextNet, and ENet.ENet and FPENet were slightly faster than LeViT-UNet-128s, while their HD still required improvement.
- Ablation study: Adding transformer blocks improved DSC and HD, with DSC gains of 1.25%, 0.25%, and 4.84% across the compared settings.Increasing input channels from 128 to 384 also reduced HD substantially in the transformer-block configurations.
- Ablation study: Adding more skip connections improved performance, with larger gains for smaller organs such as the aorta, gallbladder, and kidneys.The compared settings inserted skip connections at progressively smaller resolution scales.
- Ablation study: Pretraining had little influence on LeViT-UNet-128s and LeViT-UNet-192, but helped LeViT-UNet-384, which has about 52.1 million parameters.The smaller variants had 15.9 million and 19.9 million parameters, respectively.
4.4 Experiment results on ACDC dataset
LeViT-UNet generalizes to automated cardiac segmentation on the ACDC MR dataset, achieving better DSC results and comparable performance to Swin-UNet and TransUNet.
- LeViT-UNet achieves better ACDC cardiac-segmentation results in terms of DSC.
- 90.08% and 90.32% DSC are achieved by LeViT-UNet-192 and LeViT-UNet-384, respectively.
- LeViT-UNet achieves comparable DSC to Swin-UNet and TransUNet on ACDC.
4.5 Discussion
The discussion identifies the current LeViT-UNet design and outlines future work on multi-scale feature fusion, efficiency-accuracy balance, and 3D segmentation.
- LeViT serves as the encoder in UNet, with feature maps from three Transformer blocks concatenated after upsampling.
- Future work will explore fusing multi-scale global feature maps from Transformer blocks.
- Down-scaling the input image to 1/16 before Transformer blocks reduces computation complexity but may affect segmentation performance.
- The authors plan to design more efficient Transformer-based architectures and explore LeViT-UNet for 3D medical image segmentation.
5 Conclusion
The paper presents LeViT-UNet as a Transformer-CNN U-Net architecture for combining global context with low-level feature extraction. Experiments report superior performance, generalization ability, and an accuracy-efficiency trade-off relative to current methods.
- LeViT-UNet integrates LeViT into a U-Net-like architecture for general medical image segmentation.
- The architecture uses Transformers to build global context while retaining CNNs for low-level feature extraction.
- Experiments demonstrate superior performance and good generalization ability compared with current state-of-the-art methods.
- LeViT-UNet provides a trade-off between segmentation accuracy and efficiency.