[论文解读] Medical Image Segmentation Using Squeeze-and-Expansion Transformers

Shaohua Li, Xiuchao Sui+4 more|arXiv (Cornell University)|May 20, 2021

Advanced Neural Network Applications参考文献 43被引用 23

一句话总结

本文提出 Segtran，一种基于 Transformer 的医学图像分割框架，采用压缩与扩展 Transformer 机制，在保持高空间分辨率的同时实现大的有效感受野。通过整合可学习的正弦位置编码与双路径注意力机制，Segtran 在 2D 和 3D 医学影像任务中均优于 U-Net 及其变体，展现出最先进的精度与强大的跨域泛化能力，且仅需极少的超参数调优。

ABSTRACT

Medical image segmentation is important for computer-aided diagnosis. Good segmentation demands the model to see the big picture and fine details simultaneously, i.e., to learn image features that incorporate large context while keep high spatial resolutions. To approach this goal, the most widely used methods -- U-Net and variants, extract and fuse multi-scale features. However, the fused features still have small "effective receptive fields" with a focus on local image cues, limiting their performance. In this work, we propose Segtran, an alternative segmentation framework based on transformers, which have unlimited "effective receptive fields" even at high feature resolutions. The core of Segtran is a novel Squeeze-and-Expansion transformer: a squeezed attention block regularizes the self attention of transformers, and an expansion block learns diversified representations. Additionally, we propose a new positional encoding scheme for transformers, imposing a continuity inductive bias for images. Experiments were performed on 2D and 3D medical image segmentation tasks: optic disc/cup segmentation in fundus images (REFUGE'20 challenge), polyp segmentation in colonoscopy images, and brain tumor segmentation in MRI scans (BraTS'19 challenge). Compared with representative existing methods, Segtran consistently achieved the highest segmentation accuracy, and exhibited good cross-domain generalization capabilities. The source code of Segtran is released at https://github.com/askerlee/segtran.

研究动机与目标

• 解决 U-Net 及其变体因有效感受场较小而难以捕捉长距离上下文的局限性。
• 通过利用 Transformer 的无限有效感受场，同时保持高空间分辨率，提升医学图像分割性能。
• 设计一种专为医学图像分割量身定制的 Transformer 架构，克服自然语言预训练带来的局限性。
• 通过新颖的位置编码方案引入连续性归纳偏置，提升跨域泛化能力。
• 在多种医学影像任务中，通过极少的架构修改与超参数调优，实现最先进性能。

提出的方法

• 提出一种压缩与扩展 Transformer，其中压缩注意力模块通过正则化自注意力矩阵来减轻计算负担并稳定训练过程。
• 引入扩展模块，通过将特征投影到多个头并在注意力后融合，学习多样化表示。
• 采用可学习的正弦位置编码，强制实现空间连续性，相比标准可学习或固定正弦编码，更有利于图像数据的归纳偏置。
• 采用类似 U-Net 的编码器-解码器结构，在 Transformer 模块前后引入特征金字塔网络（FPNs），以维持并恢复空间分辨率。
• 采用 CNN 主干网络（如 ResNet-101 或 EfficientNet-B4）提取初始特征，随后将特征展平为序列以供 Transformer 处理。
• 基于像素坐标应用位置编码，以在自注意力计算过程中保留空间结构。

实验结果

研究问题

• RQ1 是否可通过增强注意力机制的 Transformer 架构，在上下文捕获能力受限的情况下，超越 U-Net 及其变体在医学图像分割任务中的表现？
• RQ2 与标准 Transformer 相比，压缩与扩展 Transformer 设计是否能提升医学影像中的特征表示与分割精度？
• RQ3 所提出的可学习正弦位置编码在多大程度上提升了归纳偏置与模型泛化能力？
• RQ4 当在成像特性不同的数据集（如 REFUGE20 与 RIM-One）上测试时，Segtran 在跨域泛化方面表现如何？
• RQ5 ImageNet 预训练对医学图像分割模型有何影响？其对不同架构的性能表现有何影响？

主要发现

• 在 REFUGE’20 挑战赛（视盘/视杯分割）中，Segtran 取得了最高的 Dice 评分，测试集平均 Dice 达到 0.938，位列前五名团队之中。
• 在 BraTS’19 挑战赛（3D 脑肿瘤分割）中，Segtran 超越 U-Net 与 DeepLabV3+，验证集平均 Dice 评分为 0.875。
• 当使用 EfficientNet-B4 作为主干网络时，Segtran 在所有任务的平均值上达到 Dice 评分 0.909，显著优于 U-Net（0.875）与 TransU-Net（0.901）。
• Segtran 展现出最佳的跨域泛化能力，从 REFUGE20 转移到 RIM-One 时，Dice 评分仅下降 14.2%，而 U-Net 与 DeepLabV3+ 分别下降 18.2% 与 19.3%。
• 在 ImageNet 上进行预训练使 Segtran 的平均 Dice 评分提升约 2.5%，表明在低数据场景下预训练具有显著优势。
• 尽管使用 ResNet-101 时参数量与 FLOPs 较高，但 Segtran 与 EfficientNet-B4 结合后实现了最优效率与性能，FLOPs 降低至 71.3G，参数量减少至 93.1M。