Source-linked AI summary
Mixed Transformer U-Net For Medical Image Segmentation
Hongyi Wang, Shiao Xie, Lanfen Lin, Yutaro Iwamoto, Xian-Hua Han, Yen-Wei Chen, Ruofeng Tong
TL;DR
U-Net-based segmentation lacks explicit long-range modeling, while Transformer approaches face pre-training, computational, and inter-sample-correlation limitations. The proposed MTM combines LGG-SA and EA in MT-UNet, which achieves stronger performance than CNNs and other Vision Transformers on two datasets.
Problem
U-Net and its variants lack long-range correlation modeling, while Transformer-based approaches may require large-scale pre-training, have quadratic complexity, and ignore inter-sample correlations.
Method
MTM combines LGG-SA for efficient local- and global-context self-affinity modeling with EA for inter-sample correlations, forming the MT-UNet segmentation model.
Results
MT-UNet surpasses CNNs, Trans-Unet, and other Vision Transformers on two datasets, achieving 78.59% DSC on Synapse and 90.43% on ACDC.
Takeaways & Limitations
MT-UNet provides an efficient Vision Transformer for medical image segmentation with simultaneous inter- and intra-affinity learning and lower time complexity than the original Transformer encoder.
Abstract
from arXiv · showhide
Though U-Net has achieved tremendous success in medical image segmentation tasks, it lacks the ability to explicitly model long-range dependencies. Therefore, Vision Transformers have emerged as alternative segmentation structures recently, for their innate ability of capturing long-range correlations through Self-Attention (SA). However, Transformers usually rely on large-scale pre-training and have high computational complexity. Furthermore, SA can only model self-affinities within a single sample, ignoring the potential correlations of the overall dataset. To address these problems, we propose a novel Transformer module named Mixed Transformer Module (MTM) for simultaneous inter- and intra- affinities learning. MTM first calculates self-affinities efficiently through our well-designed Local-Global Gaussian-Weighted Self-Attention (LGG-SA). Then, it mines inter-connections between data samples through External Attention (EA). By using MTM, we construct a U-shaped model named Mixed Transformer U-Net (MT-UNet) for accurate medical image segmentation. We test our method on two different public datasets, and the experimental results show that the proposed method achieves better performance over other state-of-the-art methods. The code is available at: https://github.com/Dootmaan/MT-UNet.
1. INTRODUCTION
The paper motivates replacing purely convolutional or standard Transformer segmentation with MTM, which jointly learns intra-sample and inter-sample affinities while addressing efficiency concerns.
- U-Net variants struggle to model long-range correlations because convolution operations are intrinsically local.
- Transformers capture long-range dependencies through Self-Attention but typically require large-scale pre-training and incur quadratic computational complexity.
- Self-Attention models correlations within one sample while ignoring inter-sample correlations across the dataset.
- MTM combines efficient Local-Global Gaussian-Weighted Self-Attention with External Attention to learn intra- and inter-sample affinities simultaneously.
- The proposed MT-UNet is evaluated on two public datasets to verify its effectiveness for medical image segmentation.
2. METHODS
MT-UNet combines convolutional processing in high-resolution layers with MTMs in deeper layers. Its MTM uses local-global attention, Gaussian-weighted axial attention, and External Attention to model visual and dataset-level relationships efficiently.
- Mixed Transformer U-Net: MTMs replace deeper Transformer encoders, while upper layers retain convolution to preserve local high-resolution structure and reduce computational cost.
- Local-Global Gaussian-Weighted Self-Attention: LGG-SA applies local self-attention to fine-grained windows and global attention to coarse-grained context, emphasizing nearby regions with a learnable Gaussian matrix.
- Gaussian weighting: The learnable Gaussian coefficient also acts as a relative position bias, explicitly encoding relative relations that ordinary absolute positional embeddings cannot provide.
- Computational complexity: For fixed window size p, Local Self-Attention has O(n) complexity, Gaussian-weighted global attention has O(n√n), and overall LGG-SA has O(n√n).
- External Attention: External Attention uses shared memory units representing essential information from the entire dataset to exploit inter-sample correlations.
- Computational complexity: External Attention has O(n) complexity, so the overall MT-UNet time complexity remains O(n√n).
3. EXPERIMENTS
Experiments on Synapse and ACDC evaluate MT-UNet against alternative architectures and test the contributions of its components. MT-UNet achieves strong segmentation performance while its ablation results support the roles of local/global attention, Gaussian weighting, and External Attention.
- Datasets And Metrics: Synapse contains 30 contrast-enhanced abdominal CT cases, while ACDC contains 100 cardiac MRI exams with labels for cardiac structures.Synapse uses 18 training and 12 testing cases; ACDC uses 70 training, 10 validation, and 20 testing samples.
- Ablation Study: Removing either Local SA or Global SA reduces performance, showing that both attention granularities contribute to the model.The ablation study also reports that the Gaussian mask is necessary for focusing the network on local areas.
- Ablation Study: EA improves overall performance by 1.04% in DSC and 1.32mm in HD95.The reported gains follow the addition of External Attention in the ablation study.
- Experimental Results: MT-UNet surpasses other Vision Transformers in complex-shaped organ segmentation, including liver and left kidney cases.The authors attribute this to balanced perception of local and global context.
4. CONCLUSIONS
The paper concludes that MT-UNet is an efficient Vision Transformer for medical image segmentation. Its MTM learns inter- and intra-affinities through LGG-SA and EA, and the model outperforms other state-of-the-art Vision Transformers in the experiments.
- Conclusions: MT-UNet uses MTM to learn inter- and intra-affinities simultaneously through LGG-SA and EA.The model is characterized by lower time complexity and higher performance than other state-of-the-art Vision Transformers in the reported experiments.