[论文解读] Class-Aware Adversarial Transformers for Medical Image Segmentation

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

Radiomics and Machine Learning in Medical Imaging参考文献 24被引用 78

一句话总结

CASTformer 是一个基于 GAN 的二维医用图像分割 Transformer 框架，使用带类别感知 Transformer 模块的金字塔多尺度生成器以及基于 Transformer 的判别器以提升分割精度。它在 Synapse 和 LiTS 数据集上达到最先进的结果，在 Dice 和 Jaccard 分数上取得显著提升。

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.

研究动机与目标

• 通过解决现有基于 Transformer 的医用分割模型在多尺度、语义和上下文方面的局限性，推动分割性能的提升。
• 提出 CASTformer，具有金字塔结构的生成器、类别感知 Transformer 模块，以及基于 GAN 的训练，以增强全局和局部特征学习。
• 在多个医用影像基准上展示性能提升，并分析迁移学习与各组件的贡献。

提出的方法

• 引入一个 GAN 框架，包含基于 Transformer 的生成器（CATformer）以及从 Vision Transformers 预训练的判别器。
• 融入特征金字塔以学习用于分割的多尺度表征。
• 开发一个类别感知 Transformer（CAT）模块，迭代地对具辨识性的解剖区域进行采样。
• 使用 Transformer 编码器模块（TEM）以捕捉长距离上下文。
• 使用轻量级全 MLP 解码器实现高效的多尺度融合与掩码预测。
• 使用 WGAN-GP 目标函数与分割损失（Dice 和交叉熵）共同训练，以在真实感与准确性之间取得平衡。

实验结果

研究问题

• RQ1 金字塔结构的、基于 Transformer 的生成器是否能在多尺度医用图像分割上优于单尺度方法？
• RQ2 在 Transformer 中的类别感知采样策略是否能提升对解剖学上有意义区域的定位？
• RQ3 对抗性训练与基于 Transformer 的判别器是否能提升医用图像分割的保真度和语义一致性？
• RQ4 迁移学习和预训练的 CV 主干网络在有限医用数据集上的性能影响如何？
• RQ5 单独的 CASTformer 组件（CAT 模块、TEM 和 GAN 训练）对整体性能的贡献是什么？

主要发现

框架	平均值	主动脉	胆囊	左肾	右肾	肝脏	胰腺	脾	胃	Dice	Jaccard	95HD	ASD
CASTformer (ours)	82.55	74.69	22.73	5.81	89.05	67.48	86.05	82.17	95.61	67.49	91.00	81.55
CATformer (ours)	82.17	73.22	16.20	4.28	88.98	67.16	85.72	81.69	95.34	66.53	90.74	81.20

• CASTformer 在 Synapse 上取得了最先进的性能，Dice 82.55 和 Jaccard 74.69（CASTformer 行）。
• 在 LiTS 上，CASTformer 达到 Dice 73.82% 和 Jaccard 64.91%，在 Dice 和 Jaccard 上分别比 TransUNet 高出 5.88% 和 4.66%。
• CATformer（不含 GAN）也优于以往方法，在 Synapse 上的 Dice 为 82.17，Jaccard 为 73.22。
• 来自 CV 预训练主干的迁移学习显著提升性能，特别是在较小的数据集上。
• 阶梯消融显示类别感知 Transformer 与 TEM 都有显著贡献；移除任一者均会降低 Dice 增益。