[Paper Review] LeViT-UNet: Make Faster Encoders with Transformer for Medical Image Segmentation

Guoping Xu, Xingrong Wu+2 more|arXiv (Cornell University)|Jul 19, 2021

Advanced Neural Network Applications39 references54 citations

TL;DR

LeViT-UNet embeds a LeViT transformer-based encoder into a U-Net-like decoder to achieve fast, accurate 2D medical image segmentation with multi-scale feature fusion from both transformer and CNN blocks. It shows competitive accuracy and improved edge prediction on Synapse and strong generalization on ACDC.

ABSTRACT

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.

Motivation & Objective

• Motivate improving medical image segmentation by combining transformer-based global context with CNN local features.
• Propose a lightweight LeViT-based encoder integrated into a U-Net-style decoder.
• Develop a multi-scale feature fusion strategy to leverage both transformer and convolutional features.
• Evaluate on multiple medical segmentation benchmarks to assess accuracy and efficiency.

Proposed method

• Use LeViT as the encoder to extract global context with reduced FLOPs.
• Concatenate multi-scale features from both convolutional and transformer blocks in the encoder at the last stage.
• Keep a CNN-based decoder with cascaded upsampling and skip connections to recover resolution.
• Pretrain LeViT backbones on ImageNet-1k to initialize parameters.
• Compare three variants LeViT-UNet-128s, -192, -384 to study channel effects and performance.
• Conduct ablations on transformer presence, skip connections, and pretraining to understand their impact.

Experimental results

Research questions

• RQ1 Can a LeViT-based encoder improve segmentation accuracy while maintaining real-time-like efficiency in a U-Net framework?
• RQ2 Does fusing multi-scale transformer and CNN features enhance both global context and local detail for medical segmentation?
• RQ3 What is the impact of the number of transformer channels and skip connections on segmentation performance and boundary accuracy?
• RQ4 How does LeViT-UNet perform on standard medical datasets (Synapse, ACDC) compared with state-of-the-art CNN- and transformer-based methods?

Key findings

• LeViT-UNet-384 achieves 78.53% DSC and 16.84 mm HD on Synapse, outperforming several SOTA methods on boundary accuracy.
• On Synapse, LeViT-UNet variants reach competitive DSC across organs, with LeViT-UNet-384 providing the best HD among reported methods (16.84 mm).
• LeViT-UNet-384 attains 90.32 DSC and 93.76 DSC on ACDC RV and LV respectively, demonstrating strong cardiac segmentation performance.
• Increasing transformer channel count and incorporating transformer blocks consistently improves DSC and HD over non-transformer baselines.
• More skip connections generally improve performance, with notable gains for smaller organs like aorta and gallbladder.
• Pre-training helps larger transformer backbones (e.g., LeViT-UNet-384) but has mixed effects for smaller versions.