Class-Aware Adversarial Transformers for Medical Image Segmentation
Anonymous Author(s)
AffiliationAddress email
Abstract
Transformers have made remarkable progress towards modeling long-range dependencies within the medical image analysis domain.
However, current transformerbased models suffer from several disadvantages: (1) existing methods fail to capturethe important features of the images due to the naive tokenization scheme; (2) themodels suffer from information loss because they only consider single-scale featurerepresentations; and (3) the segmentation label maps generated by the modelsare not accurate enough without considering rich semantic contexts and anatomical textures.
In this work, we present CASTformer , a novel type of adversarialtransformers, for 2D medical image segmentation.
First, we take advantage of thepyramid structure to construct multi-scale representations and handle multi-scalevariations.
We then design a novel class-aware transformer module to better learnthe discriminative regions of objects with semantic structures.
Lastly, we utilize anadversarial training strategy that boosts segmentation accuracy and correspondinglyallows 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-basedapproaches on three benchmarks, obtaining 2.54%-5.88% absolute improvementsin Dice over previous models.
Further qualitative experiments provide a moredetailed picture of the model’s inner workings, shed light on the challenges inimproved transparency, and demonstrate that transfer learning can greatly improveperformance and reduce the size of medical image datasets in training, making
CASTformer a strong starting point for downstream medical image analysis tasks.
Codes and models will be made available to public.
Introduction
Accurate and consistent measurements of anatomical features and functional information on medicalimages can greatly assist radiologists in making accurate and reliable diagnoses, treatment planning,and post-treatment evaluation [1].
Convolutional neural networks (CNNs) have been the de-factostandard for medical image analysis tasks.
However, such methods fail in explicitly modelinglong-range dependencies due to the intrinsic locality and weight sharing of the receptive fieldsin convolution operations.
Such a deficiency in context modeling at multiple scales often yieldssub-optimal segmentation capability in capturing rich anatomical features of variable shapes andscales ( e.g. , tumor regions with different structures and sizes) [2, 3].
Moreover, using transformers than other, traditional CNN-based methods.
In parallel, transformers with powerful global relation
[Figure]
modeling abilities have become the standard starting point for training on a wide range of downstreammedical imaging analysis tasks, such as image segmentation [7, 10–13], image synthesis [14–16],and image enhancement [17–20].
Medical image semantic segmentation can be formulated as a typical dense prediction problem,which aims at performing pixel-level classification on the feature maps.
Recently, Chen et al.
[7]introduced TransUNet , which inherits the advantage of both UNet [21] and Transformers [4],to exploit high-resolution informative representations in the spatial dimension by CNNs and thepowerful global relation modeling by Transformers.
Although existing transformer-based approacheshave proved promising in the medical image segmentation task, there remain several formidablechallenges, because (1) the model outputs a single-scale and low-resolution feature representation; (2)prior work mainly adopts the standard tokenization scheme, hard splitting an image into a sequence ofimage patches of size 16 × 16 , which may fail to capture inherent object structures and the fine-grainedspatial details for the downstream dense prediction task; (3) compared to the standard convolution,the transformer architecture requires a grid structure, and thus lacks the capability to localize regionsthat contain objects of interest instead of the uninteresting background; and (4) existing methodsare usually deficient in ensuring the performance without capturing both global and local contextualrelations among pixels.
We argue that transformer-based segmentation models are not yet robustenough to replace CNN-based methods, and investigate several above-mentioned key challengestransformer-based segmentation models still face.
Inspired by recent success of vision transformer networks [3, 4, 22–28], we make a step towards amore practical scenario in which we only assume access to pre-trained models on public computervision datasets, and a relatively small medical dataset, which we can use the weights of the pre-trainedmodels to achieve higher accuracy in the medical image analysis tasks.
These settings are particularlyappealing as (1) such models can be easily adopted on typical medical datasets; (2) such a settingonly requires limited training data and annotations; and (3) transfer learning typically leads to betterperformance [29–32].
Inspired by such findings, we propose several novel strategies for expandingits learning abilities to our setting, considering both multi-scale anatomical feature representations ofinteresting objects and transfer learning in the medical imaging domain.
First, we aim to model multi-scale variations by learning feature maps of different resolutions.
Thus,we propose to incorporate the pyramid structure into the transformer framework for medical imagesegmentation, which enables the model to capture rich global spatial information and local multi-scalecontext information.
Additionally, we consider medical semantic segmentation as a sequence-tosequence prediction task.
The standard patch tokenization scheme in [4] is an art—splitting it intoseveral fixed-size patches and linearly embedding them into input tokens.
Even if significant progressis achieved, model performance is likely to be sub-optimal.
We address this issue by introducinga novel class-aware transformer module, drawing inspiration from a progressive sampling strategyin image classification [33], to adaptively and selectively learn interesting parts of objects.
This
[Figure]
I #
Offsets +
Bilinear Interpolation
Sampled Features
T r a n s f o r m e r s
H
Norm Norm
Multi-Head Self-Attention +
MLP +
Embedded Patches (L = 12) (b) Transformer Network (M = 4)
Second, we adopt the idea of Generative Adversarial Networks (GANs) to improve segmentationperformance and correspondingly enable a transformer-based discriminator to learn low-level anatomical features and high-level semantics.
The standard GANs are not guaranteed to prioritize the mostinformative demonstrations on interesting anatomical regions, and mixing irrelevant regions ( i.e. ,background) creates inferior contexts, which drastically underperform segmentation performance[2, 34, 3].
Additionally, it is well-known that they are notoriously difficult to train and prone to modelcollapse [35].
Training vision transformers is also tedious and requires large amounts of labeled data,which largely limits the training quality.
We use a more refined strategy, where, for each input, wecombine it with the predicted segmentation mask to create the image with anatomical demonstrations.
We also leverage the pre-trained checkpoints to compensate the need of large-dataset training, therebyproviding a good starting point with more discriminative visual demonstrations.
We refer to ourapproach as CASTformer , c
lass- a ware adver s arial t rans former s: a strong transformer-based methodfor 2D medical image segmentation.
Our contributions are summarized as follows:
• Novel Network Architecture: We make the first attempt to build a GAN using a transformer-
[Table]
based architecture for the 2D medical image segmentation task.
We incorporate the pyramid
[Table]
structure into the generator to learn rich global and local multi-scale spatial representa-
[Table]
tions, and also devise a novel class-aware transformer module by progressively learning
the interesting regions correlated with semantic structures of images.
To the best of our
knowledge, we are the first work to explore these techniques in the context of medical
imaging segmentation.
• Better Understanding Inner Workings: We conduct careful analyses to understand the
model’s inner workings, how the sampling strategy works, and how different training factors
lead to the final performance.
We highlight that it is more effective to progressively learn
distinct contextual representations with the class-aware transformer module, resulting in
more accurate and robust models that applied better to a variety of downstream medical
image analysis tasks (See Appendix J and K).
• Remarkable Performance Improvements: CASTformer contributes towards a dramatic
improvement across three datasets we evaluate on.
For instance, we achieve Dice scores
of 82 .
55% and 73 .
82% by obtaining gains up to 5 .
88% absolute improvement compared to
prior methods on the Synapse multi-organ dataset.
We illustrate the benefits of leveraging
pre-trained models from the computer vision domain, and provide suggestions for future
research that could be less susceptible to the confounding effects of training data from the
natural image domain.
Related Work
CNN-based Segmentation Networks Before the emergence of transformer-based methods, CNNswere the de facto methods in medical image segmentation tasks [36–50].
For example, Ronneberger et al.
[21] proposed a deep 2D UNet architecture, combining skip connections betweenopposing convolution and deconvolution layers to achieve promising performance on a diverse set ofmedical segmentation tasks.
Han et al.
[51] developed a 2.5D 24-layer Fully Convolutional Network(FCN) for liver segmentation tasks where the residual block was incorporated into the model.
Tofurther improve segmentation accuracy, Kamnitsas et al.
[52] proposed a dual pathway 11-layer 3D
CNN, and also employed a 3D fully connected conditional random field (CRF) [53] as an additionalpairwise constraint between neighboring pixels for the challenging task of brain lesion segmentation.
Transformers in Medical Image Segmentation
Recent studies [7, 10, 12, 13, 54–59] have focusedon developing transformer-based methods for medical image analysis tasks.
Recently, Chen et al.
[7]proposed TransUNet , which takes advantage of both UNet and Transformers , to exploit highresolution informative information in the spatial dimension by CNNs and the global dependenciesby Transformers.
Cao et al.
[10] explored how to use a pure transformer for medical image analysistasks.
However, the results do not lead to better performance.
These works mainly utilized hardsplitting some highly semantically correlated regions without capturing the inherent object structures.
In this work, beyond simply using the naive tokenization scheme in [4, 7], we aim at enabling thetransformer to capture global information flow to estimate offsets towards regions of interest.
Transformer in Generative Adversarial Networks Adversarial learning has proved to be a veryuseful and widely applicable technique for learning generative models of arbitrarily complex datadistributions in the medical domain.
As the discriminator D differentiates between real and fakesamples, the adversarial loss serves as the regularization constraint to enforce the generator G topredict more realistic samples.
Inspired by such recent success [60–65], Jiang et al.
[60] proposed tobuild a GAN pipeline with two pure transformer-based architectures in synthesizing high-resolutionimages.
Esser et al.
[25] first used a convolutional GAN model to learn a codebook of context-richvisual features, followed by transformer architecture to learn the compositional parts.
Hudson et al.
[3]proposed a bipartite self-attention on StyleGAN to propagate latent variables to the evolving visualfeatures.
Despite such success, it requires high computation costs due to the quadratic complexity,which fundamentally hinders its applicability to the real world.
Besides the image generation task,we seek to take a step forward in tackling the challenging task of 2D medical image segmentation.
Method
Our proposed approach is presented in Figure 1.
Given the input image x ∈ R H × W × 3 , similar to
TransUNet architecture [7], our proposed generator network G , termed CATformer , is comprisedof four key components: encoder (feature extractor) module, class-aware transformer module,transformer encoder module, and decoder module.
As shown in Figure 1, our generator has fourstages with four parallel subnetworks.
All stages share a similar architecture, which contains a patchembedding layer, class-aware layer, and L i Transformer encoder layers.
Encoder Module.
Our method adopts a CNN-Transformer hybrid model design instead of using apure transformer, which uses 40 convolutional layers, to generate multi-scale feature maps.
Such aconvolutional stem setting provides two advantages: (1) using convolutional stem helps transformersperform better in the downstream vision tasks [66–72]; (2) it provides high-resolution feature mapswith parallel medium- and low-resolution feature maps to help boost better representations.
In thisway, we can construct the feature pyramid for the Transformers, and utilize the multi-scale featuremaps for the downstream medical segmentation task.
With the aid of feature maps of differentresolutions, our model is capable of modeling multi-resolution spatially local contexts.
Hierarchical Feature Representation.
Inspired by recent success in object detection [73], wedeviate from TransUNet by generating a single-resolution feature map, and our focus is on extracting
CNN-like multi-level features F i , where i ∈ { 1 , 2 , 3 , 4 } , to achieve high segmentation accuracy byleveraging high-resolution features and low-resolution features.
More precisely, in the first stage, weutilize the encoder module to obtain the dense feature map F 1 ∈ R H 2 × W 2 × C 1 , where ( H 2 , W 2 , C 1 )is the spatial feature resolution and the number of feature channels.
In a similar way, we canformulate the following feature maps as follows: F 2 ∈ R H
× W 2 × ( C 1 · 4) , F 3 ∈ R H
× W 4 × ( C 1 · 8) , and
F 4 ∈ R H
× W 8 × ( C 1 · 12) .
Then, we divide F 1 into HW 16 2 patches with the patch size P of 16 × 16 × 3 ,and feed the flattened patches into a learnable linear transformation to obtain the patch embeddingsof size HW 16 2 × C 1 .
Class-Aware Transformer Module.
The class-aware transformer module (CAT) is designed toadaptively focus on useful regions of objects ( e.g. , the underlying anatomical features and structuralinformation).
Our CAT module is largely inspired by the recent success for image classification [33],but we deviate from theirs as follows: (1) we remove the vision transformer module in [33] to alleviatethe computation and memory usage; (2) we use 4 separate Transformer Encoder Modules (TEM),which will be introduced below; (3) we incorporate M CAT modules on multi-scale representationsto allow for contextual information of anatomical features to propagate into the representations.
Ourclass-aware transformer module is an iterative optimization process.
In particular, we apply theclass-aware transformer module to obtain the sequence of tokens I M, 1 ∈ R C × ( n × n ) , where ( n × n )and M are the number of samples on each feature map and the total iterative number, respectively.
Asshown in Figure 2, given the feature map F 1 , we iteratively update its sampling locations by addingthem with the estimated offset vectors of the last step, which can be formulated as follows:
[Equation]
where s t ∈ R 2 × ( n × n ) , and o t ∈ R 2 × ( n × n ) are the sampling location and the predicted offset vectorat t -th step.
Specifically, the s 1 is initialized at the regularly spaced sampling grid.
The i -th samplinglocation s i 1 is defined as follows:
[Equation]
where β yi =
⌊ i/n ⌋ , β xi
=
i − β yi ∗ n .
The step sizes in the y (row index) and x (column index)directions denote τ h = H/n and τ w = W/n , respectively.
⌊·⌋ is the floor operation.
We can definethe initial token on the input feature map in the following form: I ′ t
=
F i ( s t ) , where t ∈ { 1 , . . . , M } ,and I ′ t ∈ R C × ( n × n ) denotes the initial sampled tokens at t -th step.
We set the sampling function asthe bilinear interpolation, since it is differentiable with respect to both the sampling locations s t andthe input feature map F i .
We do an element-wise addition of the current positional embedding of thesampling locations, the initial sampled tokens, and the estimated tokens of the last step, and then wecan obtain the output tokens at each step:
[Equation]
where W t ∈ R C × 2 is the learnable matrix that embeds s t to positional embedding S t ∈ R C × ( n × n ) ,and ⊕ is the element-wise addition.
Transformer ( · ) is the transformer encoder layer, as we will showin the following paragraphs.
We can compute the estimated sampling location offsets as:
[Equation]
where θ t ( · ) ∈ R 2 × ( n × n ) is the learnable linear mapping for the estimated sampling offset vectors.
It is worth noting that these operations are all differentiable, thus the model can be learned in anend-to-end fashion.
Transformer Encoder Module.
The transformer encoder module (TEM) is designed to modellong-range contextual information by aggregating global contextual information from the completesequences of input image patches embedding.
In implementations, the transformer encoder module follows the architecture in ViT [4], which is composed of Multi-head Self-Attention (MSA), and
MLP blocks, which can be formulated as:
[Equation]
where i = 1 . .
M , and LN( · ) is the layer normalization.
H ∈ R ( P 2 · C ) × D and H pos ∈ R N ×
Ddenote the patch embedding projection and the position embedding.
Decoder Module.
The decoder is designed to generate the segmentation mask based on four outputfeature maps of different resolutions.
In implementations, rather than designing a hand-crafteddecoder module that requires high computational demand, we incorporate a lightweight All-MLPdecoder [74], and such a simple design allows us to yield a powerful representation much moreefficiently.
The decoder includes the following criteria: 1) the channel dimension of multi-scalefeatures is unified through the MLP layers; 2) we up-sample the features to 1/4th and concatenatethem together; 3) we utilize the MLP layer to fuse the concatenated features, and then predict themulti-class segmentation mask y ′ from the fused features.
Discriminator Network.
We use the R50+ViT-B/16 hybrid model pre-trained on ImageNet from
ViT [4] as a starting point for our discriminator design, in this case using the pre-trained strategies tolearn effectively on the limited size target task data.
Then, we simply apply a two-layer multi-layeredperceptron (MLP) to make a prediction about the identity of the class-aware image.
Followingprevious work [2], we first utilize the ground truth image x and the predicted segmentation masky ′ to obtain the class-aware image ˜
x ( i.e. , pixel-wise multiplication of x and y ′ ).
It is importantto note that this construction re-uses the pre-trained weights and does not introduce any additionalparameters.
D seeks to classify between real and fake samples [75].
G and D compete with eachother through attempting to reach an equilibrium point of the minimax game.
Using this structureenables the discriminator to model long-range dependencies, making it better assess the medicalimage fidelity.
This also essentially endows the model with a more holistic understanding of theanatomical visual modality (categorical features).
Training Objective.
As to the loss function and training configurations, we adopt the settings used in
Wasserstein GAN (WGAN) [76], and use WGAN-GP loss [77].
We jointly use the segmentation loss[7, 13] and WGAN-GP loss to train G .
Concretely, the segmentation loss includes the dice loss andcross-entropy loss.
Hence, the training process of CASTformer can be formulated as:
[Equation]
where λ 1 , λ 2 , λ
3 determine the importance of each loss term.
See Appendix H for an ablation study.
Experimental Setup
Datasets.
We experiment on multiple challenging benchmark datasets: Synapse 1 , LiTS, and MP-MRI.
More details can be found in Appendix A.
Implementation Details.
We utilize the AdamW optimizer [78] in all our experiments.
For trainingour generator and discriminator, we use a learning rate of 5 e − 4 with a batch size of 6, and traineach model for 300 epochs for all datasets.
We set the sampling number n on each feature map andthe total iterative number M as 16 and 4, respectively.
See Appendix B, C, H for details on thetraining configuration, model architecture and hyperparameters.
We also adopt the input resolutionand patch size P as 224 × 224 and 14, respectively.
We set λ 1 = 0 .
5 , λ 2 = 0 .
5 , and λ 3 = 0 .
1 in thisexperiments.
In the testing stage, we adopt four metrics to evaluate the segmentation performance:
Dice coefficient (Dice), Jaccard Index (Jaccard), 95% Hausdorff Distance (95HD), and Average
Symmetric Surface Distance (ASD).
All our experiments are implemented in Pytorch 1.7.0.
We trainall models on a single NVIDIA GeForce RTX 3090 GPU with 24GB of memory.
[Figure]
[Table]
Results
We compare our approaches ( i.e. , CATformer and CASTformer ) with previous state-of-the-arttransformer-based segmentation methods, including UNet [21], AttnUNet [79], ResNet50 + UNet[21], ResNet50 + AttnUNet [79], SETR [80], CoTr w/o CNN-encoder [13], CoTr [13], TransUNet[7], SwinUnet [10] on Synapse, LiTS, and MP-MRI datasets.
More results are in Appendix D and E.
Experiments: Synapse Multi-organ.
The quantitative results on the Synapse dataset are shown in
Table 1.
The results are visualized in Figure 3.
It can be observed that our CATformer outperforms theprevious best model by a large margin and achieves a 4 .
69% − 8 .
44% absolute improvement in Diceand Jaccard, respectively.
Our CASTformer achieves the best performance of 82 .
55% and 74 .
69% ,dramatically improving the previous state-of-the-art model ( TransUNet ) by +5 .
and +9 .
91% ,in terms of both Dice and Jaccard scores.
This shows that the anatomical visual information is usefulfor the model to gain finer control in localizing local semantic regions.
As also shown in Table 1, our
CASTformer achieves absolute Dice improvements of +2 .
77% , +2 .
51% , +1 .
35% , +4 .
95% on largeorgans ( i.e. , left kidney, right kidney, liver, stomach) respectively.
Such improvements demonstratethe effectiveness of learning the evolving anatomical features of the image, as well as accuratelyidentifying the boundary information of large organs.
We also observed similar trends that, comparedto the previous state-of-the-art results, our CASTformer obtains 89 .
48% , 67 . 49% in terms of
Dice on small organs ( i.e. , aorta, gallbladder, pancreas) respectively, which yields big improvementsof +1 .
82% , +0 .
95% , +10 .
This clearly demonstrates the superiority of our models, allowingfor a spatially finer control over the segmentation process.
Experiments: LiTS.
To further evaluate the effectiveness of our proposed approaches, we compareour models on the LiTS dataset.
Experimental results on the LiTS CT dataset are summarized
[Figure]
in Appendix Table 6.
As is shown, we observe that our CATformer yields a 72 .
39% Dice score,outperforming all other methods.
Moreover, our CASTformer significantly outperforms all previousapproaches, including the previous best TransUNet , and establishes a new state-of-the-art of 73 . 82%and 64 .
91% in terms of Dice and Jaccard, which are 5 .
88% and 4 .
66% absolute improvements betterthan TransUNet .
For example, our CASTformer achieves the best performance of 95 . 88% Diceon the liver region by 2 .
48% , while it dramatically increases the result from 42 .
49% to 51 .
76%on the tumor region, demonstrating that our model achieves competitive performance on liver andtumor segmentation.
As shown in Figure 4, our method is capable of predicting high-quality objectsegmentation, considering the fact that the improvement in such a setting is challenging.
Thisdemonstrates: (1) the necessity of adaptively focusing on the region of interests; and (2) the efficacyof semantically correlated information.
Compared to previously high-performing models, our twoapproaches achieve significant improvements on all datasets, demonstrating their effectiveness.
Analysis
In this section, we conduct thorough analyses of our CASTformer in the following aspects: transferlearning, model components, effects of iteration number N (Appendix F), sampling number n (Appendix G), hyperparameter selection (Appendix H), different GAN-based loss functions (Appendix
I), and gain a better understanding of the model’s inner workings (Appendix J and K).
Transfer Learning.
We consider whether we can leverage the pre-trained model commonly usedin computer vision literature [4], to provide more evidence for the beneficial impact of our networkperformance.
We use CATformer ( G ) as the baseline and evaluate all the settings on the Synapsemulti-organ dataset.
To put our results in perspective, we compare with six ways of using pre-trained
R50+ViT-B/16 hybrid model from ViT [4] for transfer learning, namely (1) CATformer :
w/o pretrained and w/ pre-trained; (2) CASTformer : both w/o pre-trained, only w/ pre-trained D , only w/pre-trained G , and both w/ pre-trained G and D .
multi-organ dataset.
[Table] [Table]
The results are in Table 2.
Overall, we observe that using “w/ pre-trained” leads to higher accuracythan “w/o pre-trained”, with significant improvements for the smaller sizes of datasets, suggestingthat using “w/ pre-trained” provides us a good set of initial parameters for the downstream tasks.
Withusing pre-trained weights, CATformer outperforms the setting without using pre-trained weights bya large margin and achieves 7 .
33% and 7 .
61% absolute improvements in terms of Dice and Jaccard.
CASTformer (“both w/ pre-trained”) also yields big improvements ( +8 . 91% and +12 .
01% ) in Diceand Jaccard.
This suggests that CASTformer is better at both initializing from the pre-trained modelsand better at gathering the anatomical information in a more adaptive and selective manner.
As shownin Table 2, surprisingly, there is a significant discrepancy between only using “w/ pre-trained D ” and“w/ pre-trained G ”: for example, CASTformer achieves 78 .
87% in Dice with only w/ pre-trained D ,while CASTformer achieves 81 .
46% if only G uses the pre-trained weights.
This demonstrates thatonly using pre-trained weight in D might be the culprit for the exploitation of anatomical information.
Our results suggest that (1) utilizing pre-trained models in the computer vision domain can help themodel quickly adapt to new downstream medical segmentation tasks without re-building billionsof anatomical representations; (2) we find that leveraging pre-trained weights can further boost theperformance because it can mitigate the discrepancy between training and inference; and (3) it alsocreates a possibility to adapt our model to the typical medical dataset with the smaller size.
Ablation of Model Components.
Our key observation is that it is crucial to build high-qualityanatomical representations through each model component.
To show the strengths of our approach,we examine the following variants and inspect each key component on the Synapse multi-organsegmentation dataset: (1) Baseline : we remove the class-aware transformer module and the transformer encoder module in our CATformer as the baseline, similar to TransUNet defined in [7];(2) CATformer w/o CAT : we only remove the class-aware transformer module in our CATformer ;(3) CATformer w/o TEM : we only remove the transformer encoder module in our CATformer ; (4)
CATformer : this is our G model; and (5) CASTformer : this is our final model described in Section3.
Table 3 summarizes the results of all the variants.
As is shown, we observe that compared tothe baseline model, both CATformer w/o CAT and CATformer w/o TEM are able to develop abetter holistic understanding of global shapes/structures and fine anatomical details, thus leadingto large performance improvements ( +2 .
61% and +3 .
in terms of Dice.
Our results show theclass-aware transformer module is useful in improving the segmentation performance, suggestingthat the discriminative regions of medical images are particularly effective.
Finally, one thing worthnoticing is that incorporating both the class-aware transformer module and the transformer encodermodule performs better than only using a single module, highlighting the importance of two modulesin our CATformer .
Conclusion and Discussion of Broader Impact
In this work, we have introduced CASTformer , a simple yet effective type of adversarial transformers,for 2D medical image segmentation.
The key insight is to integrate the multi-scale pyramid structureto capture rich global spatial information and local multi-scale context information.
Furthermore,
CASTformer also benefits from our proposed class-aware transformer module to progressively andselectively learn interesting parts of the objects.
Lastly, the generator-discriminator design is used toboost segmentation performance and correspondingly enable the transformer-based discriminatorto capture low-level anatomical features and high-level semantics.
Comprehensive experimentsdemonstrate that our CASTformer outperforms the previous state-of-the-art on three popular medicaldatasets considerably.
We conduct extensive analyses to study the robustness of our approach, andform a more detailed understanding of desirable properties in the medical domain ( i.e. , transparencyand data efficiency).
Overall, we hope that this model can serve as a solid baseline for 2D medical image segmentationand motivate further research in medical image analysis tasks.
It also provides a new perspective ontransfer learning in medical domain, and initially shed novel insights towards understanding neuralnetwork behavior.
As such pattern is hard to quantify, we expect more mechanistic explanations for clinical practise.
We also plan to optimize the transformer-based architectures for the downstreammedical image analysis tasks both in terms of data and model parameters.
Broader Impact.
We acknowledge that our research will not pose significant risks of societal harmto society.
Our work is scientific nature and will have the potential to positively contribute to anumber of real-world clinical applications that establish high-quality and end-to-end medical imagesegmentation systems.
We expect our approach to contribute to the grand goal of building moresecured and trustworthy clinical AI.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section 7, and Broader
Impact on Page 9 .
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See
Section 7, and Broader Impact on Page 9 .
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] See Sections
and supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Sections 4 and supplemental material.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See Sections 4 and supplemental material.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Sections 4 and supplementalmaterial.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?
[Yes] See Sections 4 and supplemental material.
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[Yes] See Sections 4 and supplemental material.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional ReviewBoard (IRB) approvals, if applicable?
[N/A](c)
Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]
Appendix toClass-Aware Adversarial Transformersfor Medical Image Segmentation
Sec.
A provides additional details on the training datasets.
Sec.
B provides additional details on the implementation.
Sec.
C provides additional model details.
Sec.
D provides more quantitative results on the LiTS dataset.
Sec.
E provides more experimental results on the MP-MRI dataset.
Sec.
F provides additional analysis on Iteration Number N .
Sec.
G provides additional analysis on sampling number n .
Sec.
H provides additional details on the hyperparameter selection.
Sec.
I provides additional GAN-based loss details.
Sec.
J takes a deeper look and understand our proposed class-aware transformer module.
Sec.
K forms a better understand the attention mechanism in our proposed CASTformer .
A Datasets
Synapse : Synapse multi-organ segmentation dataset includes 30 abdominal CT scans with 3779 axialcontrast-enhanced abdominal clinical CT images.
Each CT volume consists of 85 ∼ 198 slices of512 × 512 pixels, with a voxel spatial resolution of ([0 .
54 ∼ 0 .
54] × [0 .
98 ∼ 0 .
98] × [2 . 5 ∼ 5 . 0])
mm 3 .
The dataset is randomly divided into 18 volumes for training (2212 axial slices), and 12 for validation.
For each case, 8 anatomical structures are aorta, gallbladder, spleen, left kidney, right kidney, liver,pancreas, spleen, stomach.
LiTS : MICCAI 2017 Liver Tumor Segmentation Challenge (LiTS) includes 131 contrast-enhanced3D abdominal CT volumes for training and testing.
The dataset is assembled by different scannersand protocols from seven hospitals and research institutions.
The image resolution ranges from0.56mm to 1.0mm in axial and 0.45mm to 6.0mm in z direction.
The dataset is randomly dividedinto 100 volumes for training, and 31 for testing.
MP-MRI : Multi-phasic MRI dataset is an in-house dataset including multi-phasic MRI scans of 20local patients with HCC, each of which consisted of T1 weighted DCE-MRI images at three-timepoints (pre-contrast, arterial phase, and venous phases).
Three images are mutually registered to thearterial phase images, with an isotropic voxel size of 1.00 mm.
The dataset is randomly divided into48 volumes for training, and 12 for testing.
B More Implementation Details
The training configuration and hyperparameter settings are summarized in Table 4.
C Model Architecture
We present the detailed architecture of CATformer ’s encoding pipeline in Section 5.
We use input/output names to indicate the direction of the data stream.
CATformer applies independent class-awareattention on 4 levels of features extracted by the ResNetV2 model.
Each feature level L- k is processedby CATformer - k , consisting of 4 blocks of class-aware transformer modules, followed by 12 layersof transformer encoder modules.
Outputs from all four feature levels are fed into the decoder pipelineto generate the segmentation masks.
[Table] [Table]
D More Experiments: LiTS
results are summarized in Table 6.
[Table]
E More Experiments: MP-MRI
Experimental results are summarized in Table 7.
Overall, CATformer and CASTformer outperform the previous results in terms of Dice and Jaccard.
Compared to SETR, our CATformer and
[Table]
CASTformer perform 1 .
78% and 2 .
54% higher in Dice, respectively.
We also find CASTformerperforms better than CATformer , which suggests that using discriminator can make the model betterassess the medical image fidelity.
Figure 5 shows qualitative results, where our CATformer and
CASTformer provide better anatomical details than all other methods.
This clearly demonstrates thesuperiority of our models.
All these experiments are conducted using the same hyperparameters inour CASTformer .
[Figure]
F Effect of Iteration Number N
We explore the effect of different iteration number N in Figure 6 (a).
Note that in the case of N = 1 ,the sampling locations will not be updated.
We find that more iterations of sampling clearly improvenetwork performance in Dice and Jaccard.
However, we observe that the network performancedoes not further increase from N = 4 to N = 6 .
In our study, we use N = 4 for the class-awaretransformer module.
G Effect of Sampling Number n
We further evaluate the effect of sampling number n of the class-aware transformer module in Figure6 (b).
Empirically, we observe that results are generally well correlated when we gradually increasethe size of n .
As is shown, the network performance is optimal when n = 16 .
[Figure]
H Hyperparameter Selection
[Figure]
We carry out grid-search of λ 1 , λ 2 , λ 3 ∈ { 0 . 0 , 0 . 1 , 0 . 2 , 0 . 5 , 1 . 0 } .
As shown in Figure 7, with acarefully tuned hyperparameters λ 1 = 0 .
5 , λ 2 = 0 .
5 , and λ 3 = 0 .
1 , such setting performs generallybetter than others.
[Figure]
I Importance of Loss Functions
[Figure]
One main argument for the discriminator is that modeling long-range dependencies and acquiringa more holistic understanding of the anatomical visual information can contribute to the improvedcapability of the generator.
Besides the WGAN-GP loss [77], the minimax (MM) GAN loss [75],the Non-Saturating (NS) GAN loss [86], and Least Squares (LS) GAN Loss [87] are also commonlyused as adversarial training.
We test these alternatives and find that, in most cases, using WGAN-GPloss achieves comparable or higher performance than other loss functions.
In addition, models trainedusing MM-GAN loss perform comparably to those trained using LS-GAN loss.
In particular, ourapproach outperforms the second-best LS-GAN loss [87] by 1 . 10 and 2 .
49 points in Dice and Jaccardscores on the Synapse multi-organ dataset.
It demonstrates the effectiveness of the WGAN-GP lossin our CASTformer .
[Figure]
[Table]
J Visualization of Learned Sampling Location
To gain more insight into the evolving sampling locations learned by our proposed class-awaretransformer module, we visualize the predicted offsets in Figure 8.
We can see that particular
[Figure]
sampling points around objects tend to attend to coherent segmented regions in terms of anatomicalsimilarity and proximity.
As is shown, we show the classes with the highly semantically correlatedregions, indicating that the model coherently attends to anatomical concepts such as liver, right/leftkidney, and spleen.
These visualizations also illustrate how it behaves adaptively and distinctivelyto focus on the content with highly semantically correlated discriminative regions ( i.e. , differentorgans).
These findings can thereby suggest that our design can aid the CATformer to exercise finercontrol emphasizing anatomical features with the intrinsic structure at the object granularity.
As isindicated (Figure 8 last column), we also find evidence that our model is prone to capture some smallobject cases ( e.g. , pancreas, aorta, gallbladder).
We hypothesize that it is because they contain moreanatomical variances, which makes the model more difficult to exploit.
K Vision Transformer Visualization
In this section, we visualize the first 12 class-aware transformer layers on sequences of 28 × 28 featurepatches in the encoder pipeline.
In Figure 9, we plot the attention probabilities from a single patchover different layers and heads.
Each row corresponds to one CAT layer; each column corresponds toan attention head.
As we go deeper into the network, we are able to observe three kinds of attentionbehaviors as further discussed below.
Attend to similar features: In the first group of layers (layer 1 through 4), the attention probabilityis spread across a relatively large group of patches.
Notably, these patches correspond to areas in theimage with similar color and texture to the query patch.
These more primitive attention distributionsindicate that the class-awareness property has not yet been established.
Attend to the same class and its boundary: In the middle layers of the transformer model, mostnoticeable in the 5 th and 6 th layers, the attention probabilities start to concentrate on areas that sharethe same class label as the query patch (layer 5-2).
In some other instances, the model attends to theboundary of the current class (layer 5-3, 5-6).
I m a g e
L a y e r L a y e r
L a y e r L a y e r
L a y e r L a y e r
L a y e r L a y e r
L a y e r L a y e r
L a y e r L a y e r
Attend to other classes:
In the deeper layers of the model, the attention probability mainly concentrates on other classes.
This clearly demonstrates persuasive evidence that the model establishes classawareness, which is helpful in the downstream medical segmentation tasks.
L More Ablations on Decoder Modules
In this section, we explore another state-of-the-art backbone proposed by Lin et al.
[73], termed
Feature Pyramid Network ( FPN ).
FPN utilizes a top-down pyramid with lateral connections to construct the semantically strong multi-scale feature pyramid from a single-scale input.
The major differencesbetween FPN and our work are as follows:
• The former utilizes a CNN-based decoder ( FPN [73]), and ours uses an All-MLP-based
decoder.
In particular, our motivation comes from the observation that the attention oflower layers tends to be local, and those of the higher layers are highly non-local [74].
As the decoder design plays an important role in determining the semantic level of the
latent representations [9] and Transformers have the larger receptive fields compared to
CNNs, how to use large receptive fields to include context information is the key issue
Prior work [
] suggests that the use of MLP-based decoder design
can
be a very effective tool in learning additional contextual information to build powerful
representations.
The key idea is to essentially take benefits of the Transformer-induced
features by leveraging the local attention at the lower layers and highly non-local (global)
attention at the higher layers to formulate the powerful representations [74].
To this end, we
utilize an MLP-based decoder instead of a CNN-based decoder to preserve more contextual
information, specifically for medical imaging data, including more anatomical variances.
• We devise the class-aware transformer module to progressively learn interesting
anatomical
regions correlated with semantic structures of images, so as to guide the segmentation of
objects or entities.
We study the model’s qualitative behavior through learnable samplinglocations inside the class-aware module in Figure 8.
As indicated, sampling locations areadaptively adjusted according to the interesting regions.
The table below shows the comparision results of using an FPN decoder, MLP-based decoder, andthe class-aware transformer (CAT) module, all of which include the backbone feature extractor( ResNet50 ), on the Synapse multi-organ CT dataset.
All the experiments are conducted underthe same experimental setting in Section 4.
As we can see, adopting the MLP-based decoder canoutperform the state-of-the-art FPN decoder in terms of DSC, Jaccard, 95HD, and ASD, respectively.
Similarly, incorporating the CAT module can also consistently improve the segmentation performanceby a large margin on the Synapse multi-organ CT dataset.
The results prove the robustness ofour MLP-based decoder and the effectiveness of our proposed CAT module for medical imagesegmentation.
[Table]
M More Ablations on Segmentation Losses
To deal with the imbalanced medical image segmentation, Lin et al.
[72] proposed Focal loss interms of the standard cross entropy to address the extreme foreground-background class imbalanceby focusing on the hard pixel examples.
The table below shows the results of the loss function.
Wefollow γ = 2 in the original paper.
As we can see, the setting using Focal loss and the other ( i
[Table]
e. ,
Dice + Cross-Entropy) achieve similar performances.
[Table]
N More Ablations on Sampling Modules
In this section, we investigate the effect of recent state-of-the-art sampling modules [88–90].
However,the motivation and the sampling strategy are different from these works [88–90].
Our motivationcomes from the accurate and reliable clinical diagnosis that rely on the meaningful radiomic featuresfrom the correct “region of interest” instead of other irrelevant parts [91–94].
The process ofextracting different radiomic features from medical images is done in a progressive and adaptivemanner [92, 93].
DCN [88] proposed to learn 2D spatial offsets to enable the CNN-based model to generalize thecapability of regular convolutions.
Because CNNs only have limited receptive fields compared to
Transformers, DCN focuses on local information around a certain point of interest.
In contrast, our
CATformer / CASTformer take benefits of the Transformer-induced features by leveraging the localattention at the lower layers and highly non-local (global) attention at the higher layers to formulatethe powerful representations.
Deformable DETR [89] incorporated the deformation attention to focus on a sparse set of keys ( i.e. ,global keys are not shared among visual tokens).
This is particularly useful for its original experimentsetup on object detection.
Since there are only a handful of query features corresponding to potentialobject classes, deformable DETR learns different attention locations for each class.
In contrast, ourapproach aims at refining the anatomical tokens for medical image segmentation.
To this end, weproposed to iteratively and adaptively focus on the most discriminative region of interests.
Thisessentially allows us to obtain effective anatomical features from spatial attended regions within themedical images, so as to guide the segmentation of objects or entities
DAT [90] introduced deformable attention to make use of global information ( i.e. , global keys areshared among visual tokens) by placing a set of the supporting points uniformly on the feature maps.
In contrast, our approach introduces an iterative and progressive sampling strategy to capture themost discriminative region and avoid over-partition anatomical features.
The table below shows the comparison results between DCN [88], Deformable DETR [89], DAT [90],and ours ( CATformer / CASTformer ) on the Synapse multi-organ CT dataset.
As we can see, ourapproach ( i.e. , CATformer / CASTformer ) can outperform existing state-of-the-art models, i.e. , DCN[88], and Deformable DETR .
[Table]
O More Ablations on Architecture Backbone
In this section, we conduct the ablation study on the Synapse multi-organ CT dataset to compareour approach with the recent state-of-the-art architecture ( SwinUnet ) [10].
The table below showsthe results of our proposed architecture ( e.g. , Swin-class-aware transformer (Swin-CAT) module,multi-scale feature extraction module) are superior compared to the other state-of-the-art method onthe Synapse multi-organ CT dataset.
All the experiments are conducted under the same experimentalsetting in Section 4.
For brevity, we refer our CATformer and CASTformer using SwinUnet asthe backbone to Swin-CATformer and Swin-CASTformer .
As we can see, using SwinUnet asthe backbone, the following observations can be drawn: (1) “w/ pre-trained” consistently achievessignificant performance gains compared to the “w/o pre-trained”, which demonstrates the effectivenessof the pre-training strategy; (2) we can find that incorporating the adversarial training can boostthe segmentation performance, which suggests the effectiveness of the adversarial training strategy; and (3) our Swin-CASTformer with different modules can also achieves consistently improvedperformance.
The results prove the superiority of our proposed method on the medical imagesegmentation task.
[Table] [Table]