Source-linked AI summary
Multi-Compound Transformer for Accurate Biomedical Image Segmentation
Yuanfeng Ji, Ruimao Zhang, Huijie Wang, Zhen Li, Lingyun Wu, Shaoting Zhang, Ping Luo
TL;DR
Biomedical segmentation methods lack cross-scale pixel dependencies, semantic correspondence between categories, and consistency between feature representations and semantic embeddings. MCTrans addresses these gaps with multi-scale self-attention and learnable proxy-based cross-attention in a UNet-like network. It outperforms state-of-the-art methods across six biomedical segmentation datasets with a slight computation increase.
Problem
Existing biomedical segmentation methods do not adequately model cross-scale pixel dependencies, semantic category correspondence, or consistency between feature representations and semantic embeddings.
Method
MCTrans combines cross-scale Transformer-Self-Attention with Transformer-Cross-Attention using a learnable proxy embedding to model semantic relationships and enhance feature representations.
Results
MCTrans outperforms state-of-the-art methods across six challenging segmentation datasets with a slight computation increase.
Takeaways & Limitations
MCTrans provides a unified transformer-based framework for cross-scale context modeling and semantic relationship mining in biomedical image segmentation.
Abstract
from arXiv · showhide
The recent vision transformer(i.e.for image classification) learns non-local attentive interaction of different patch tokens. However, prior arts miss learning the cross-scale dependencies of different pixels, the semantic correspondence of different labels, and the consistency of the feature representations and semantic embeddings, which are critical for biomedical segmentation. In this paper, we tackle the above issues by proposing a unified transformer network, termed Multi-Compound Transformer (MCTrans), which incorporates rich feature learning and semantic structure mining into a unified framework. Specifically, MCTrans embeds the multi-scale convolutional features as a sequence of tokens and performs intra- and inter-scale self-attention, rather than single-scale attention in previous works. In addition, a learnable proxy embedding is also introduced to model semantic relationship and feature enhancement by using self-attention and cross-attention, respectively. MCTrans can be easily plugged into a UNet-like network and attains a significant improvement over the state-of-the-art methods in biomedical image segmentation in six standard benchmarks. For example, MCTrans outperforms UNet by 3.64%, 3.71%, 4.34%, 2.8%, 1.88%, 1.57% in Pannuke, CVC-Clinic, CVC-Colon, Etis, Kavirs, ISIC2018 dataset, respectively. Code is available at https://github.com/JiYuanFeng/MCTrans.
1 The University of Hong Kong
The paper lists the Shenzhen Research Institute of Big Data at The Chinese University of Hong Kong (Shenzhen) and SenseTime Research among its affiliations.
- The listed affiliations include the Shenzhen Research Institute of Big Data at The Chinese University of Hong Kong (Shenzhen) and SenseTime Research.
1 Introduction
The paper identifies limitations in existing biomedical segmentation context modeling and proposes MCTrans to address cross-scale dependencies and semantic relationships. MCTrans combines multi-scale contextual modeling with learnable semantic proxy embeddings in a UNet-like framework.
- Biomedical image segmentation supports computer-assisted diagnosis tasks including lesion detection and anatomical structure localization.
- Existing methods remain limited by inefficient non-local context modeling among arbitrary positions, hindering further accuracy improvements for complex views.
- Vision-transformer segmentation methods model global context at a single feature scale, overlooking cross-scale dependencies and consistency important for lesions with dramatic size changes.
- MCTrans introduces Transformer-Self-Attention for cross-scale pixel-level context modeling and Transformer-Cross-Attention with proxy embeddings for semantic category correspondence.
- MCTrans combines cross-scale contextual dependencies, semantic relationships, and a learnable proxy embedding within a UNet-like network evaluated on six challenging segmentation datasets.
- The overview processes CNN-derived multi-scale features through self-attention, proxy-based cross-attention, and progressive feature-map merging to produce segmentation results.
2 Related Work
Prior work uses attention and transformers to model contextual relationships in segmentation, but existing approaches do not fully exploit attention for biomedical image segmentation. The paper motivates a more specialized transformer design.
- Attention Mechanisms: Attention-based gates, channel attention, and non-local operations model target relevance, feature-channel dependencies, or pairwise pixel relationships.
- Transformer: TransUNet applies self-attention to the highest-level UNet features to collect long-range dependencies at a specific feature scale.
- Transformer: The paper argues that existing attention and transformer methods are not specifically designed for medical image segmentation, motivating a more tailored approach.
3 Multi-Compound Transformer Network
MCTrans inserts TSA and TCA between a UNet encoder and decoder to model multi-scale pixel context and category-level semantic relationships. Multi-scale tokens are processed through attention modules, then folded back into pyramid features for segmentation.
- Architecture: MCTrans places Transformer-Self-Attention and Transformer-Cross-Attention between the UNet encoder and decoder.TSA encodes contextual information across multiple features, while TCA uses learnable embeddings for semantic relationship modeling and feature enhancement.
- Tokenization: Multi-scale CNN features are unfolded into tokens, projected to a shared dimension, concatenated, and supplemented with positional embeddings.Features from levels i = 2,3,4 are combined before TSA processing.
- Decoder and supervision: The final encoded tokens are folded back into pyramid features and progressively merged with upsampling and convolution for segmentation prediction.An auxiliary classification loss also supervises the proxy embedding using unique elements from the ground-truth segmentation mask.
- Transformer-Self-Attention: TSA learns pixel-level contextual dependencies among multiple-scale features using stacked multi-head self-attention and feed-forward blocks.The module applies layer normalization before blocks and residual connections afterward.
- Transformer-Self-Attention: Deformable Self Attention replaces dense self-attention because the flattened multi-scale token sequence is too long for quadratic computation.DSA attends to a sparse, data-dependent subset of elements, reducing computation complexity.
- Transformer-Cross-Attention: TCA first models intra- and inter-class relationships among learnable proxy embeddings, then uses cross-attention to exchange information with enhanced feature tokens.This interaction is described as improving intra-class consistency and inter-class discriminability.
4 Experiments
MCTrans was evaluated across six biomedical segmentation datasets and through component, sensitivity, comparative, and computational-overhead experiments. Its modules improved segmentation accuracy, generalized across tasks, and delivered higher performance than UNet and other methods with modest overhead.
- Datasets and Settings: MCTrans was evaluated on six datasets spanning cell, polyp, and skin-lesion segmentation, with different modalities, dataset sizes, and foreground classes.The datasets were Pannuke, CVC-Clinic, CVC-ColonDB, ETIS-Larib, Kvasir, and ISIC2018.
- Analysis of the Network Components: 68.40% Dice on Pannuke exceeded the UNet baseline’s 64.92%; TSA alone reached 67.93%, while adding TCA raised the baseline and MCTrans to 67.16% and 68.40%.Removing auxiliary loss reduced MCTrans to 67.87%.
- Analysis of the Network Components: Increasing TSA and TCA capacity improved DSC only up to a peak, with TCA reaching its top at Nc = 4 before performance decreased.The sensitivity results indicate that larger transformer capacity is not uniformly better on a small dataset.
- Comparisons with State-of-the-art Methods: MCTrans outperformed competing methods on Pannuke and five lesion-segmentation datasets, demonstrating versatility across biomedical segmentation tasks.The comparisons used both VGG-Style and ResNet-34 feature extractors on Pannuke.
- Computational Overheads: 3.64% improvement over UNet was achieved with almost identical parameters and a slight computation increase; UNet++ used more computation while performing worse.The paper reports the comparison in terms of computational overhead and segmentation performance.
5 Conclusions
The paper proposes MCTrans, a transformer-based network for biomedical image segmentation that combines context modeling with semantic relationship mining. It addresses cross-scale dependencies and category relationships, and outperforms methods such as TransUnet on several public datasets.
- Conclusions: MCTrans is a transformer-based network for biomedical image segmentation that models rich context and semantic relationships through attention mechanisms.Its design targets cross-scale dependencies and semantic correspondence among categories.