Source-linked AI summary
TransNorm: Transformer Provides a Strong Spatial Normalization Mechanism for a Deep Segmentation Model
Reza Azad, Mohammad T. AL-Antary, Moein Heidari, Dorit Merhof
TL;DR
Medical image segmentation needs both global contextual modeling and precise local detail, while CNNs and pure Transformers each have complementary weaknesses. TransNorm combines CNN and Transformer features and uses Transformer-derived Spatial Normalization to recalibrate skip connections. Across three segmentation tasks, it is reported as effective, including 78.40% DSC and 30.25% HD on Synapse and an 8% F1 improvement over U-Net on ISIC 2017.
Problem
CNNs have limited long-range contextual modeling, while pure Transformer segmentation can lack local information and localization capacity.
Method
TransNorm combines CNN and Transformer features and applies a Transformer-derived Spatial Normalization attention mechanism to recalibrate decoder skip connections.
Results
TransNorm is reported to outperform competing approaches across several datasets, achieving 78.40% DSC and 30.25% HD on Synapse and improving F1 by 8% over U-Net on ISIC 2017.
Takeaways & Limitations
Combining CNN semantic features, Transformer long-range context, and adaptive skip-connection attention supports smooth and precise segmentation results.
Abstract
from arXiv · showhide
In the past few years, convolutional neural networks (CNNs), particularly U-Net, have been the prevailing technique in the medical image processing era. Specifically, the seminal U-Net, as well as its alternatives, have successfully managed to address a wide variety of medical image segmentation tasks. However, these architectures are intrinsically imperfect as they fail to exhibit long-range interactions and spatial dependencies leading to a severe performance drop in the segmentation of medical images with variable shapes and structures. Transformers, preliminary proposed for sequence-to-sequence prediction, have arisen as surrogate architectures to precisely model global information assisted by the self-attention mechanism. Despite being feasibly designed, utilizing a pure Transformer for image segmentation purposes can result in limited localization capacity stemming from inadequate low-level features. Thus, a line of research strives to design robust variants of Transformer-based U-Net. In this paper, we propose Trans-Norm, a novel deep segmentation framework which concomitantly consolidates a Transformer module into both encoder and skip-connections of the standard U-Net. We argue that the expedient design of skip-connections can be crucial for accurate segmentation as it can assist in feature fusion between the expanding and contracting paths. In this respect, we derive a Spatial Normalization mechanism from the Transformer module to adaptively recalibrate the skip connection path. Extensive experiments across three typical tasks for medical image segmentation demonstrate the effectiveness of TransNorm. The codes and trained models are publicly available at https://github.com/rezazad68/transnorm.
I. INTRODUCTION
Medical image segmentation supports diagnosis and treatment but remains challenging because CNN-based methods have limited receptive fields and lose deep-layer details. U-Net improves detail recovery through skip connections, while Transformers offer stronger global-context modeling, motivating their combination.
- Medical image segmentation enables computer-aided diagnosis, treatment planning, and surgical pre-assessment, but manual annotation is labor-intensive and error-prone.
- U-Net uses symmetric encoder-decoder paths with skip connections to combine high-resolution encoder features with decoder features for more accurate representations.
- CNNs improve context modeling only within a confined receptive field, limiting their ability to capture long-range semantic dependencies in varied medical images.
- Transformers provide a complementary direction because self-attention can model global context and long-term features beyond convolutional receptive fields.
- The paper motivates attention-based skip-connection design to strengthen important regions during segmentation and improve feature fusion between network paths.
B. TRANSFORMERS
Transformer-based segmentation methods address CNNs’ weak global-context encoding but can underrepresent local information. TransNorm combines CNN and Transformer features and adds a two-level attention mechanism that recalibrates skip-connection features.
- TransUNet first combines convolutional feature extraction with a ViT encoder, then upsamples encoded features through cascaded decoder layers for segmentation.
- CNN methods often produce poor global-context encoding, whereas Transformer methods capture long-range connectivity but may lack local-information learning and generalization.
- C. COMBINING CNNS WITH TRANSFORMER ARCHITECTURE: The method concatenates Transformer and CNN features, then applies a two-level attention mechanism to adaptively recalibrate feature combinations on decoder skip connections.
- A. ENCODER: TransNorm uses a CNN encoder for semantic features and a Transformer branch at the network bottleneck to learn long-range contextual dependencies.
- A. ENCODER: The encoder maps the input image to lower-resolution feature maps through a parameterized encoding network before subsequent Transformer processing.
B. LONG-RANGE CONTEXTUAL REPRESENTATION
TransNorm reshapes encoder features into patch sequences, adds positional information, and processes them with Transformer stages to capture long-range context. The resulting Transformer features are concatenated with CNN encoder features for complementary semantic and contextual representations.
- B. LONG-RANGE CONTEXTUAL REPRESENTATION: The method reshapes encoder features into flattened nonoverlapping patches and maps them into an embedding sequence for Transformer processing.Each patch has size P × P, and the sequence contains N patches.
- B. LONG-RANGE CONTEXTUAL REPRESENTATION: Learned positional embeddings are added to patch embeddings because self-attention is permutation-invariant and otherwise does not retain patch locations.
- B. LONG-RANGE CONTEXTUAL REPRESENTATION: The Transformer encoder uses K stages composed of Multihead Self-Attention and Multi-Layer Perceptron blocks.
- B. LONG-RANGE CONTEXTUAL REPRESENTATION: Transformer features are concatenated with CNN encoder features to form complementary semantic and long-range contextual representations.These fused features support the proposed Spatial Normalization mechanism in the decoder skip connections.
C. ATTENTION GATE
The attention gate recalibrates skip-connection features in two levels: channel attention first, followed by Transformer-derived spatial normalization. This design directs decoding toward informative regions while suppressing less relevant content.
- C. ATTENTION GATE: The two-level gate first computes channel attention to recalibrate object-recognition features, then applies a Transformer-derived spatial map.The channel map is 1D, whereas the spatial map is 2D over the feature-map locations.
- C. ATTENTION GATE: Channel normalization uses global average pooling and a fully connected layer to adaptively scale each channel.
- C. ATTENTION GATE: The spatial attention map Ws focuses on informative regions so the segmentation model can emphasize or suppress locations during mask generation.
- C. ATTENTION GATE: Figure 3 depicts the proposed gate as normalizing both channel and spatial information on each decoding path.
- C. ATTENTION GATE: Generated attention maps are incorporated into skip connections to fuse high-level and low-level feature maps during decoding.
D. DECODER
The decoder is a pure CNN that progressively upsamples feature maps to the original image resolution for pixel-level segmentation. It uses Spatial Normalizer outputs to introduce high-level semantic information into the upsampling path.
- D. DECODER: The decoder progressively upsamples feature maps to the original image space and produces pixel-level segmentation through transpose convolution.
- D. DECODER: Cascaded upsampling recovers resolution from the preceding layer as the decoder reconstructs the segmentation output.
- D. DECODER: Spatial Normalizer feature maps are used in the upsampling layer to promote global information with high-level semantic content.
- D. DECODER: The experiments cover three medical image segmentation tasks and include quantitative, visualization, and ablation analyses.
1) Synapse multi-organ segmentation
This section describes the datasets and evaluation metrics used for multi-organ, skin-lesion, and multiple-myeloma segmentation. Synapse is evaluated with DSC and HD, while skin-lesion and myeloma tasks use classification and overlap metrics.
- 1) Synapse multi-organ segmentation: The BTCV dataset contains 30 abdominal CT scans, 3779 axial images, and annotations for 13 organs.
- 1) Synapse multi-organ segmentation: ISIC 2017 provides 2000 dermoscopy images at 576 × 767 pixels for lesion segmentation and related clinical-image tasks.
- 1) Synapse multi-organ segmentation: ISIC 2018 contains more than 10,000 dermoscopic images spanning seven disease types, resized to 256 × 256 pixels.
- 1) Synapse multi-organ segmentation: The PH2 dataset contains 200 dermatoscopic images of melanocytic skin lesions, while the SegPC challenge provides 290 annotated multiple-myeloma training images.
- 1) Synapse multi-organ segmentation: Synapse is evaluated using average DSC and HD, whereas skin-lesion and multiple-myeloma segmentation use accuracy, sensitivity, specificity, F1-Score, and mIoU.
C. IMPLEMENTATION DETAILS
The implementation used PyTorch on a single RTX 3090 system, with fixed training settings for both baseline and proposed models.
- Both models were trained for 100 epochs with batch size 16, initial learning rate 1e−3, and decay rate 1e−4.
- Training used normally distributed initialization, no data augmentation, and batch normalization in each block to facilitate convergence.
D. COMPARISON RESULTS
TransNorm was evaluated against comparable methods across multi-organ, skin-lesion, and myeloma segmentation tasks. The reported results generally favor TransNorm, including best performance on Synapse and SegPC.
- 1) Synapse multi-organ segmentation: 78.40% DSC and 30.25% HD were achieved on Synapse, improving over recent TransU-Net variants by 0.98% DSC and 1.44 mm HD.Table 1 reports the comparison against state-of-the-art approaches.
- Visual results show masks close to ground truth on Synapse, smooth lesion boundaries on ISIC 2017, and precise separation under overlapping backgrounds.The cited figures provide qualitative examples for Synapse and ISIC 2017.
- 2) Skin Lesion Segmentation: TransNorm improved the U-Net baseline’s F1 score by 8% on ISIC 2017 and outperformed state-of-the-art approaches in almost all metrics.It also performed better than TransU-Net on this dataset.
- 2) Skin Lesion Segmentation: TransNorm outperformed both CNN- and Transformer-based approaches on ISIC 2018, including cases with highly overlapping lesion backgrounds.Table 3 provides the quantitative comparison, while Figure 6 illustrates precise segmentation under background overlap.
- 2) Skin Lesion Segmentation: On PH2, TransNorm performed better than state-of-the-art models except on sensitivity.The comparison is reported in Table 4.
- 3) Multiple Mylomia Segmentation: TransNorm obtained the best overall performance against SegPC 2021 challenge winners while using fewer parameters than those competing approaches.The cited comparison contrasts TransNorm with teams using combinations of instance-segmentation, Mask-RCNN, and attention-DeepLab models.
E. ABLATION STUDY
The ablation study examined attention, skip connections, input resolution, and model scaling on ISIC 2018. Its design varied these factors to assess their influence on overall performance.
- A comprehensive ISIC 2018 ablation study investigated attention, skip connections, input resolution, and model scaling.
1) Attention Mechanism
The attention ablation study evaluated the proposed two-level attention mechanism and the effect of varying skip-connection counts. The reported rationale links channel and spatial information with global relationship modeling and low-level spatial detail.
- 1) Attention Mechanism: The ablation results across five metrics indicate the significance of the proposed two-level attention mechanism.
- 1) Attention Mechanism: Integrating channel and spatial information on decoding paths leverages Transformer modeling of global relationships while exploiting complementary feature information.
- 1) Attention Mechanism: Skip connections supply low-level spatial information to the decoding path, and their number was varied experimentally to assess its effect on performance.The corresponding experiments are reported in Table 7.
3) Input Resolution Influence
The experiments examine how input resolution and model scaling affect TransNorm on ISIC 2018. Higher-resolution inputs provide finer object detail, while increasing model parameters yields limited inference-time benefit.
- Input Resolution Influence: Higher-resolution inputs provide finer details of the object of interest and can influence segmentation performance on ISIC 2018.The experiments compare the standard 256×256 resolution with 512×512 inputs.
- Experimental Scope: The reported resolution and scaling experiments are conducted on the ISIC 2018 dataset.
- Model Scaling Influence: Increasing the number of parameters improves training-set performance but provides little improvement at inference time because of overfitting.The authors therefore report that model scaling is not computationally efficient for this design.