PaperlayerView exact revision on arXiv

Source-linked AI summary

Class-Aware Adversarial Transformers for Medical Image Segmentation

Chenyu You, Ruihan Zhao, Fenglin Liu, Siyuan Dong, Sandeep Chinchali, Ufuk Topcu, Lawrence Staib, James S. Duncan

arXiv:2201.10737v5cs.CVcs.AIcs.LGeess.IV

TL;DR

Transformer segmentation is challenged by naive tokenization and single-scale representations, limiting capture of image features and multi-scale variation. CASTformer combines pyramid representations, a class-aware transformer, and adversarial training, and outperforms previous state-of-the-art models on three medical datasets.

  • Problem

    Transformer methods face suboptimal segmentation from patch tokenization and limited capture of rich anatomical features across scales.

    Source: §1 · pp. 1–2

  • Method

    CASTformer combines a pyramid structure, a class-aware transformer that progressively learns object regions, and adversarial training with a transformer-based discriminator.

    Source: §1 · pp. 2–3, §7 · p. 9

  • Results

    Across three medical datasets CASTformer outperforms previous state-of-the-art models; on Synapse multi-organ, its gains reach 5.88% absolute Dice improvement over prior methods.

    Source: §1 · p. 3, §7 · p. 9

  • Takeaways & Limitations

    The authors frame public computer-vision pretraining with a relatively small medical dataset as a practical setting requiring limited training data and annotations.

    Source: §1 · p. 2

Abstract

from arXiv · show

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.

1 Introduction

CASTformer addresses limitations in transformer-based medical image segmentation with multi-scale representations, adaptive class-aware sampling, and adversarial training. Its introduction also reports gains across three datasets and describes analyses of the model’s inner workings.

1 Introduction

  • Existing transformer segmentation methods are described as limited by single-scale, low-resolution features, fixed patch tokenization, and insufficient modeling of both global and local context.The introduction argues that these methods may miss object structure, fine spatial detail, and regions of interest.Source: §1 · p. 2
  • CASTformer combines a transformer pyramid for multi-scale context with class-aware sampling that selectively learns informative object regions.The introduction describes the sampled regions as a source of anatomical features for guiding segmentation.Source: §1 · p. 2
  • Its adversarial strategy pairs segmentation with a transformer discriminator for anatomical features and semantics, while pretrained checkpoints address the demand for large labeled datasets.For each input, the method combines the image with its predicted mask to create anatomical demonstrations.Source: §1 · pp. 2–3
  • The authors present a transformer-based GAN for 2D medical segmentation and analyze sampling behavior and training factors affecting performance.They report that progressively learning distinct contextual representations with the class-aware module yields more accurate, robust models for varied downstream medical imaging tasks.Source: §1 · p. 3
  • Up to 5.88% absolute Dice improvement over prior methods on Synapse multi-organ is reported within CASTformer’s gains across three datasets.The introduction also reports Dice scores of 82.55% and 73.82%.Source: §1 · p. 3

2 Related Work

Related work traces medical segmentation from CNN-based approaches to transformer methods, then places CASTformer alongside transformer-based adversarial learning. The section identifies object structure, region localization, and computational cost as relevant challenges.

CNN-based Segmentation Networks

  • CNNs were established medical segmentation methods, with prior work spanning 2D UNet, 2.5D FCN, and 3D CNN approaches.The cited examples include skip connections, residual blocks, and a 3D CRF constraint for brain lesion segmentation.Source: §2 · p. 4

Transformers in Medical Image Segmentation

  • Medical segmentation transformers combine CNN spatial representations with global dependencies or use pure transformers, but prior work is described as missing inherent object structures.CASTformer instead aims to estimate offsets toward regions of interest beyond naive tokenization.Source: §2 · p. 4

Transformer in Generative Adversarial Networks

  • Transformer-based GANs have been explored for image generation, but quadratic complexity can impose high computation costs that hinder real-world use.The section positions medical image segmentation as a further application for adversarial transformers.Source: §2 · p. 4

3 Method

CASTformer combines a multi-scale CNN–Transformer generator, iterative class-aware sampling, and adversarial training for medical image segmentation. Its encoder and decoder build multi-resolution features, while a transformer discriminator evaluates class-aware images.

3 Method

  • CASTformer’s four-stage generator combines an encoder, class-aware transformer, transformer encoder, and decoder.The stages share a patch embedding layer, class-aware layer, and Transformer encoder layers.Source: §3 · p. 4
  • A 40-layer convolutional stem creates multi-resolution feature maps that provide local contexts and support the feature pyramid.The method uses CNN-like features at multiple levels alongside transformer representations.Source: §3 · pp. 4–5
  • The class-aware transformer iteratively updates sampling locations to focus on useful object regions and propagate anatomical context across multi-scale representations.Sampling starts from a regularly spaced grid, with predicted offsets updating locations between steps.Source: §3 · p. 5
  • The transformer encoder models long-range context, while a lightweight All-MLP decoder fuses four resolution-specific feature maps to predict the segmentation mask.The decoder upsamples and concatenates the features before fusion.Source: §3 · pp. 5–6
  • A transformer discriminator classifies pixel-wise class-aware images formed from the input and segmentation mask, and the generator is trained with Dice, cross-entropy, and WGAN-GP losses.The discriminator starts from an ImageNet-pretrained R50+ViT-B/16 hybrid; the loss weights determine the relative importance of the three terms.Source: §3 · p. 6

4 Experimental Setup

Experiments use Synapse, LiTS, and MP-MRI datasets, with AdamW training and four segmentation metrics for evaluation.

4 Experimental Setup

  • Experiments cover the Synapse, LiTS, and MP-MRI benchmark datasets.Additional dataset details are in Appendix A.Source: §4 · p. 6
  • Models are trained for 300 epochs using AdamW, a batch size of 6, and a learning rate of 5e−4.The reported configuration uses 224×224 inputs, patch size 14, and a single 24GB NVIDIA GeForce RTX 3090 GPU.Source: §4 · pp. 6–7
  • Segmentation performance is evaluated with Dice, Jaccard, 95HD, and ASD.Table 1 reports quantitative segmentation results on the Synapse multi-organ CT dataset.Source: §4 · p. 7, Table 1 · p. 7

5 Results

On Synapse and LiTS, CASTformer reports substantial gains over prior best models. The LiTS visual comparison also highlights detailed anatomical information in its predictions.

  • The comparison evaluates CATformer and CASTformer against transformer-based segmentation methods across Synapse, LiTS, and MP-MRI.Source: §5 · p. 7
  • On Synapse, CASTformer achieved 82.55% Dice and 74.69% Jaccard, improving over TransUNet by +5.07% and +9.91%, respectively.Source: §5 · p. 7
  • CASTformer reached 73.82% Dice and 64.91% Jaccard on LiTS, improving over TransUNet by 5.88% and 4.66%, respectively.Source: §5 · p. 7
  • Figure 4 shows CASTformer producing superior LiTS visual results with detailed anatomical information, including tumor regions.Source: Figure 4 · p. 8

6 Analysis

The analysis examines transfer learning and model components, alongside additional architectural and training factors. Pretrained weights and the combination of class-aware and transformer-encoder modules are associated with improved Synapse results.

Analysis Scope

  • The analyses also cover iteration and sampling numbers, hyperparameters, GAN losses, and the model's inner workings.Source: §6 · p. 8

Transfer Learning

  • With both models pretrained, CASTformer gained +8.91% Dice and +12.01% Jaccard, while CATformer gained 7.33% Dice and 7.61% Jaccard.Source: §6 · p. 8
  • CASTformer reached 78.87% Dice with only the discriminator pretrained and 81.46% when only the generator was pretrained.Source: §6 · p. 8
  • Pretrained weights yielded higher accuracy than no pretraining, with particularly significant improvements on smaller datasets, and may support adaptation to smaller medical datasets.Source: §6 · pp. 8–9

Ablation of Model Components

  • Removing CAT or TEM individually improved baseline Dice by +2.61% and +3.87%, while using both modules outperformed either alone.Source: Table 3 · p. 8, §6 · p. 9

7 Conclusion and Discussion of Broader Impact

The conclusion presents CASTformer as a 2D medical image segmentation approach combining pyramid representations, class-aware learning, and adversarial training. The authors report strong results on three datasets and outline broader implications, open transparency questions, and future work.

7 Conclusion and Discussion of Broader Impact

  • CASTformer combines pyramid-based global and local multi-scale representations, class-aware learning of object regions, and adversarial training with a transformer-based discriminator.Source: §7 · p. 9
  • The authors propose CASTformer as a baseline for 2D medical image segmentation and call for more mechanistic explanations of neural-network behavior in clinical practice.Source: §7 · p. 9
  • The authors anticipate potential contributions to clinical segmentation systems and trustworthy clinical AI, while stating that the research poses no significant societal risks.Source: §7 · p. 9

Appendix to Class-Aware Adversarial Transformers for Medical Image Segmentation

A Datasets

The appendix describes Synapse, LiTS, and MP-MRI datasets, points to training-configuration details, and outlines CATformer's multilevel encoding and decoder pipeline.

A Datasets

  • Synapse uses 30 abdominal CT scans, LiTS uses 131 contrast-enhanced CT volumes, and the in-house MP-MRI dataset includes scans from 20 local patients with HCC.Source: A Datasets · p. 16

B More Implementation Details

C Model Architecture

  • CATformer applies independent class-aware attention to four ResNetV2 feature levels, processes each through CATformer-k blocks and transformer encoders, then feeds the outputs to a decoder for segmentation masks.Source: C Model Architecture · p. 16

D More Experiments: LiTS

This section points to the experimental results and records the training configuration and CATformer architecture settings.

D More Experiments: LiTS

E More Experiments: MP-MRI

On MP-MRI, CATformer and CASTformer are reported to outperform previous results, with CASTformer also exceeding CATformer. The section further examines sampling and loss choices and visualizes learned sampling and attention behavior.

F Effect of Iteration Number N

E More Experiments: MP-MRI

G Effect of Sampling Number n

H Hyperparameter Selection

I Importance of Loss Functions

J Visualization of Learned Sampling Location

  • Learned sampling points tend to attend to coherent anatomically similar regions, including liver, kidneys, and spleen, while visualizations also show some small-object cases.The authors hypothesize that anatomical variance makes small objects such as pancreas, aorta, and gallbladder harder to exploit.Source: J Visualization of Learned Sampling… · pp. 19–20, Figure 8 · p. 20

K Vision Transformer Visualization

  • Attention progresses from similar color-and-texture patches in early layers, to the query class or its boundary in middle layers, and to other classes in deeper layers.Figure 9 plots attention probabilities across layers and heads for a query patch.Source: K Vision Transformer Visualization · p. 20

L More Ablations on Decoder Modules

The decoder ablations favor the MLP-based decoder over FPN and report consistent gains from adding the class-aware transformer module. The section also compares segmentation losses.

L More Ablations on Decoder Modules

  • FPN uses a CNN-based decoder, whereas the proposed decoder is All-MLP; the paper motivates this choice as preserving context from local lower-layer and non-local higher-layer attention.The CAT module progressively learns anatomical regions correlated with semantic structures and adaptively adjusts sampling locations.Source: L More Ablations on Decoder Modules · pp. 20–22
  • The MLP-based decoder outperforms FPN on DSC, Jaccard, 95HD, and ASD, while adding CAT consistently improves segmentation performance on Synapse.The comparison uses a ResNet50 backbone under the same experimental setting.Source: Table 9 · p. 22, L More Ablations on Decoder Modules · p. 22

M More Ablations on Segmentation Losses

N More Ablations on Sampling Modules

The sampling-module comparisons and backbone ablations report advantages for the proposed components on Synapse. The section describes adaptive anatomical sampling and the effects of pretraining and adversarial training.

N More Ablations on Sampling Modules

  • The proposed sampling iteratively and adaptively focuses on discriminative regions of interest to refine anatomical tokens for segmentation.The authors contrast this with Deformable DETR's class-query attention locations and DAT's uniformly placed supporting points.Source: N More Ablations on Sampling Modules · p. 23
  • CATformer and CASTformer outperform DCN and Deformable DETR in the Synapse sampling-module comparison.The compared methods are DCN [100], Deformable DETR [101], DAT [102], and CATformer/CASTformer; no outcome versus DAT is stated.Source: Table 11 · p. 23, N More Ablations on Sampling Modules · p. 23

O More Ablations on Architecture Backbone

Related papers

Loading 2201.10737v5…