Source-linked AI summary
Medical Image Segmentation Using Squeeze-and-Expansion Transformers
Shaohua Li, Xiuchao Sui, Xiangde Luo, Xinxing Xu, Yong Liu, Rick Goh
TL;DR
Medical image segmentation must capture large context while retaining high spatial resolution, a balance that conventional CNN-based approaches struggle to achieve. Segtran uses transformer-based pairwise contextualization with Squeeze-and-Expansion attention and learnable sinusoidal positional encoding. Across two 2D and one 3D tasks, it consistently outperformed existing methods and generalized well to new domains.
Problem
Medical image segmentation requires large contextual coverage and high spatial resolution simultaneously, while conventional multi-scale CNN features can retain small effective receptive fields focused on local cues.
Method
Segtran combines transformer-based feature contextualization with Squeeze-and-Expansion attention, learnable sinusoidal positional encoding, and feature pyramids for upsampling.
Results
Across two 2D and one 3D medical image segmentation tasks, Segtran consistently outperformed U-Net variants and DeepLabV3+.
Takeaways & Limitations
Segtran provides a transformer-based segmentation framework that captures global context while preserving fine details and generalizes well to new domains.
Takeaways & Limitations
Using ResNet backbones gives Segtran significantly higher model size and FLOPs, making this choice unsuitable for efficiency-sensitive scenarios.
Abstract
from arXiv · showhide
Medical image segmentation is important for computer-aided diagnosis. Good segmentation demands the model to see the big picture and fine details simultaneously, i.e., to learn image features that incorporate large context while keep high spatial resolutions. To approach this goal, the most widely used methods -- U-Net and variants, extract and fuse multi-scale features. However, the fused features still have small "effective receptive fields" with a focus on local image cues, limiting their performance. In this work, we propose Segtran, an alternative segmentation framework based on transformers, which have unlimited "effective receptive fields" even at high feature resolutions. The core of Segtran is a novel Squeeze-and-Expansion transformer: a squeezed attention block regularizes the self attention of transformers, and an expansion block learns diversified representations. Additionally, we propose a new positional encoding scheme for transformers, imposing a continuity inductive bias for images. Experiments were performed on 2D and 3D medical image segmentation tasks: optic disc/cup segmentation in fundus images (REFUGE'20 challenge), polyp segmentation in colonoscopy images, and brain tumor segmentation in MRI scans (BraTS'19 challenge). Compared with representative existing methods, Segtran consistently achieved the highest segmentation accuracy, and exhibited good cross-domain generalization capabilities. The source code of Segtran is released at https://github.com/askerlee/segtran.
1 Introduction
Medical image segmentation requires both broad context and fine spatial detail, but conventional CNN-based designs face a receptive-field and resolution trade-off. Segtran addresses this with transformer-based contextualization and reports stronger performance across multiple medical segmentation tasks.
- Medical image segmentation supports computer-aided diagnosis by delineating anatomical structures and regions of interest for quantification and pathology localization.
- CNN-based segmentation faces a trade-off: larger context generally comes with lower feature resolution, while fine-grained masks require high spatial resolution.
- U-Net and its variants combine encoder-decoder features across multiple scales, but deeper convolutions can still limit effective receptive fields.
- Segtran uses transformer layers to contextualize flattened local features, with positional encodings and feature pyramids supporting spatially detailed segmentation.
- The effective receptive-field visualization shows Segtran’s non-negligible gradients dispersed across the image, unlike the concentrated gradients of U-Net and DeepLabV3+.The visualization uses gradients back-propagated from the image center on 576 × 576 inputs.
- Segtran was evaluated on optic disc/cup, colon polyp, and brain tumor segmentation, consistently outperforming U-Net variants and DeepLabV3+.
2 Related Work
Related work includes transformer-based detection and biomedical segmentation approaches, CNNs augmented with positional encodings, and strategies for enlarging U-Net receptive fields. These approaches differ in applicability, empirical consistency, and effects on spatial precision or model complexity.
- DETR-inspired approaches use transformers for contextualized representations, while panoptic DETR’s two-stage design is described as unsuitable for medical image segmentation.
- CNNs extended with positional-encoding channels produced mixed segmentation results, whereas Segtran’s ablation study found positional encodings helpful to some degree.
- Increasing U-Net receptive fields through extra downsampling can raise parameter counts and overfitting risk, while larger convolution strides sacrifice feature-map spatial precision.
3 Squeeze-and-Expansion Transformer
The Squeeze-and-Expansion Transformer adapts self-attention to images by reducing attention complexity, increasing representational diversity, modeling symmetric image-unit relationships, and encoding spatial continuity.
- Motivation: Traditional transformers use large N × N attention matrices that can be vulnerable to noise and overfitting, motivating lower-rank attention.Image relationships are often symmetric, while standard key and query projections independently model asymmetric relationships.
- Squeezed Attention Block: The Squeezed Attention Block uses M learned inducing points to replace full attention with two compact N × M interactions.Input units attend to codebook embeddings, which are projected and then attend back to the inputs; usually M ≪ N.
- Expanded Attention Block: The Expanded Attention Block combines Nm complete feature representations from separate transformer modes through dynamic mode attention.Unlike multi-head attention, EAB modes share the same feature space, preserving representation power when individual modes are removed.
- Squeeze-and-Expansion Transformer: The proposed architecture combines SAB with EAB, using tied query and key projections to model symmetric image relationships.The paper identifies this combination as its novel transformer architecture and describes EAB as a mixture-of-experts design.
- Learnable Sinusoidal Positional Encoding: Learnable Sinusoidal Positional Encoding combines sine and cosine functions of learnable linear coordinate transformations to impose adaptable spatial continuity.Coordinates are normalized for consistency across image sizes, and the encoding extends to 3D coordinates for volumetric inputs.
4 Segtran Architecture
Segtran combines CNN features, feature pyramids, positional encodings, Squeeze-and-Expansion transformer layers, and a segmentation head to contextualize image features while retaining spatial resolution.
- Segtran uses a CNN backbone, input/output feature pyramids, learnable sinusoidal positional encoding, Squeeze-and-Expansion transformer layers, and a segmentation head.
- Visual features are added to positional encodings, then flattened into a one-dimensional sequence of image units before transformer processing.
- Each transformer layer computes pairwise interactions and outputs contextualized features with the same number of spatial units.
- Feature Pyramids: The input FPN upsamples backbone features from 1/16 to 1/8 of the original image resolution before transformer processing.A 1 × 1 convolution aligns channels, and bilinear interpolation performs the ×2 upsampling.
- Feature Pyramids: The output FPN upsamples transformer features by ×4 and combines them with earlier backbone features, producing maps at half the original image resolution.Its two steps use f 12 = upsample×2(f 2) + conv12(f 1) and g1234 = upsample×4(g34) + conv24(f 12).
- Segmentation Outputs: Figure 5 illustrates three-class optic disc/cup masks and two-class polyp masks using the stated grayscale class encodings.
- Feature Pyramids: The adopted bottom-up FPN empirically performs better than the original top-down FPN because richer semantics in top layers are better preserved.
- Segmentation Head: The segmentation head is a 1 × 1 convolution that outputs confidence scores for each class in the mask.
5 Experiments
Experiments evaluated Segtran across two 2D and one 3D medical segmentation task, including ablations, baseline comparisons, generalization, and computational efficiency. Segtran generally achieved strong performance across tasks, while backbone choice affected efficiency.
- Tasks: Three tasks were evaluated: optic disc/cup segmentation, polyp segmentation, and brain tumor segmentation in MRI scans.
- Tasks: REFUGE20 used 1,200 training and 400 validation fundus images, with additional Drishti-GS and RIM-ONE data used during training.
- Ablation Studies: The ablations compared Squeeze-and-Expansion with multi-head transformers, positional-encoding schemes, and different transformer depths.
- Ablation Studies: Squeeze-and-Expansion transformers outperformed traditional multi-head transformers, with both squeeze and expansion blocks contributing to improved performance.
- Ablation Studies: Removing positional encoding reduced Segtran performance by only 1~2%, with CNN features offered as a possible source of positional information.
- Ablation Studies: Performance improved as transformer depth increased from 1 to 3 layers, but adding another layer caused a drop attributed to possible overfitting.
- Baseline Comparisons: Across all tasks, transformer-based SETR, TransU-Net, and Segtran achieved the best performance; Segtran gained greater advantages with EfficientNet-B4.
- Baseline Comparisons: Segtran with EfficientNet-B4 ranked among the top five teams in both REFUGE20 semifinal and final leaderboards.
6 Conclusions
Segtran is a transformer-based medical image segmentation framework that contextualizes features using transformers’ unlimited receptive fields. It captures global context and fine details, consistently outperforming existing methods across two 2D and one 3D task while generalizing well to new domains.
- Segtran is a transformer-based medical image segmentation framework that leverages transformers’ unlimited receptive fields to contextualize features.
- Segtran captures both the global picture and fine details, supporting strong segmentation performance.
- Across two 2D and one 3D medical image segmentation task, Segtran consistently outperformed existing methods and generalized well to new domains.