[Paper Review] Class-Aware Adversarial Transformers for Medical Image Segmentation

Chenyu You, Ruihan Zhao+6 more|PubMed|Jan 26, 2022

Radiomics and Machine Learning in Medical ImagingMedicine24 references78 citations

TL;DR

CASTformer is a GAN-based transformer framework for 2D medical image segmentation that uses a pyramid multi-scale generator with a class-aware transformer module and a transformer-based discriminator to boost segmentation accuracy. It achieves state-of-the-art results on Synapse and LiTS datasets with notable improvements in Dice and Jaccard scores.

ABSTRACT

Transformers have made remarkable progress towards modeling long-range dependencies within the medical image analysis domain. However, current transformer-based models suffer from several disadvantages: (1) existing methods fail to capture the important features of the images due to the naive tokenization scheme; (2) the models suffer from information loss because they only consider single-scale feature representations; and (3) the segmentation label maps generated by the models are not accurate enough without considering rich semantic contexts and anatomical textures. In this work, we present CASTformer, a novel type of adversarial transformers, for 2D medical image segmentation. First, we take advantage of the pyramid structure to construct multi-scale representations and handle multi-scale variations. We then design a novel class-aware transformer module to better learn the discriminative regions of objects with semantic structures. Lastly, we utilize an adversarial training strategy that boosts segmentation accuracy and correspondingly allows a transformer-based discriminator to capture high-level semantically correlated contents and low-level anatomical features. Our experiments demonstrate that CASTformer dramatically outperforms previous state-of-the-art transformer-based approaches on three benchmarks, obtaining 2.54%-5.88% absolute improvements in Dice over previous models. Further qualitative experiments provide a more detailed picture of the model's inner workings, shed light on the challenges in improved transparency, and demonstrate that transfer learning can greatly improve performance and reduce the size of medical image datasets in training, making CASTformer a strong starting point for downstream medical image analysis tasks.

Motivation & Objective

• Motivate improved segmentation by addressing multi-scale, semantic, and contextual limitations of existing transformer-based medical segmentation models.
• Propose CASTformer with a pyramid-structured generator, class-aware transformer module, and GAN-based training to enhance both global and local feature learning.
• Demonstrate performance gains on multiple medical imaging benchmarks and analyze transfer learning and component contributions.

Proposed method

• Introduce a GAN framework with a transformer-based generator (CATformer) and a discriminator pre-trained from Vision Transformers.
• Incorporate a feature pyramid to learn multi-scale representations for segmentation.
• Develop a class-aware transformer (CAT) module that iteratively samples discriminative anatomical regions.
• Employ a transformer encoder module (TEM) to capture long-range context.
• Use a lightweight all-MLP decoder for efficient multi-scale fusion and mask prediction.
• Train with a WGAN-GP objective plus segmentation losses (Dice and cross-entropy) to balance realism and accuracy.

Experimental results

Research questions

• RQ1 Can a pyramid-structured, transformer-based generator improve multi-scale medical image segmentation over single-scale approaches?
• RQ2 Does a class-aware sampling strategy within a transformer enhance localization of anatomically meaningful regions?
• RQ3 Do adversarial training and a transformer-based discriminator improve segmentation fidelity and semantic consistency on medical images?
• RQ4 How do transfer learning and pre-trained CV backbones impact performance on limited medical datasets?
• RQ5 What is the contribution of individual CASTformer components (CAT module, TEM, and GAN training) to overall performance?

Key findings

• CASTformer achieves state-of-the-art performance on Synapse with Dice 82.55 and Jaccard 74.69 (CASTformer row).
• On LiTS, CASTformer attains Dice 73.82% and Jaccard 64.91%, outperforming TransUNet by 5.88% and 4.66% in Dice and Jaccard, respectively.
• CATformer (without GAN) also outperforms prior methods, with Dice 82.17 and Jaccard 73.22 on Synapse.
• Transfer learning from CV pre-trained backbones significantly boosts performance, especially on smaller datasets.
• Ablation shows both the class-aware transformer and TEM contribute meaningfully; removing either reduces Dice gains.