[论文解读] D-Former: A U-shaped Dilated Transformer for 3D Medical Image Segmentation

Yixuan Wu, Kuanlun Liao+5 more|arXiv (Cornell University)|Jan 3, 2022

AI in cancer detection被引用 9

一句话总结

D-Former 提出了一种用于 3D 医疗图像分割的 U 形扩张 Transformer，通过使用扩张多头自注意力机制在不增加 token 数量的前提下捕捉跨 patch 的长距离依赖关系，从而降低计算成本。该方法在 Synapse 和 ACDC 数据集上实现了最先进性能，且 FLOPs 和参数量低于其他竞争的 CNN 和 Transformer 模型，通过一种新颖的动态位置编码机制提升了空间建模能力。

ABSTRACT

Computer-aided medical image segmentation has been applied widely in diagnosis and treatment to obtain clinically useful information of shapes and volumes of target organs and tissues. In the past several years, convolutional neural network (CNN) based methods (e.g., U-Net) have dominated this area, but still suffered from inadequate long-range information capturing. Hence, recent work presented computer vision Transformer variants for medical image segmentation tasks and obtained promising performances. Such Transformers model long-range dependency by computing pair-wise patch relations. However, they incur prohibitive computational costs, especially on 3D medical images (e.g., CT and MRI). In this paper, we propose a new method called Dilated Transformer, which conducts self-attention for pair-wise patch relations captured alternately in local and global scopes. Inspired by dilated convolution kernels, we conduct the global self-attention in a dilated manner, enlarging receptive fields without increasing the patches involved and thus reducing computational costs. Based on this design of Dilated Transformer, we construct a U-shaped encoder-decoder hierarchical architecture called D-Former for 3D medical image segmentation. Experiments on the Synapse and ACDC datasets show that our D-Former model, trained from scratch, outperforms various competitive CNN-based or Transformer-based segmentation models at a low computational cost without time-consuming per-training process.

研究动机与目标

• 解决 3D 医疗图像分割 Transformer 中全局自注意力带来的高计算成本问题。
• 在保持空间分辨率并减少参数量的同时，改进 3D 医疗图像中的长距离特征建模能力。
• 设计一种高效、可扩展的 Transformer 架构，适用于计算资源有限的 3D 医疗成像任务。
• 集成一种能够自适应序列依赖位置关系的动态位置编码，以提升空间理解能力。

提出的方法

• 提出一种扩张自注意力机制，以交替方式计算局部和全局自注意力，通过扩张扩展感受野，同时不增加 patch 数量。
• 引入一种 U 形编码器-解码器架构 D-Former，结合局部感受野模块（LSM）和全局感受野模块（GSM），通过交替的扩张注意力机制实现。
• 在首个局部自注意力模块之前引入动态位置编码（DPE）模块，以灵活学习位置感知表征。
• 采用分层设计，每个阶段包含交替的 LSM 和 GSM 模块，以逐步捕捉多尺度上下文特征。
• 应用可学习的位置嵌入，自适应于输入相关的空间关系，从而改善 3D 空间结构的建模能力。
• 从零开始训练模型，无需预训练，即可在低计算成本下实现高性能。

实验结果

研究问题

• RQ1 扩张自注意力机制是否能在降低 3D 医疗图像分割计算成本的同时，保持长距离依赖关系的建模能力？
• RQ2 在 3D 医疗成像任务中，所提出的 D-Former 架构与标准 Transformer 和 CNN 相比，在分割精度和 FLOPs 方面表现如何？
• RQ3 动态位置编码的最佳放置位置和设计方式是什么，以提升 Transformer 中 3D 空间表征能力？
• RQ4 用扩张全局感受野模块替代全局自注意力是否能在不增加 token 数量的前提下提升性能？
• RQ5 当在 Synapse 和 ACDC 等 3D 医疗图像数据集上从零开始训练时，D-Former 模型的有效性如何？

主要发现

• D-Former 在 Synapse 数据集上实现了 88.83% 的平均 Dice 系数（DSC），优于包括 Swin-Unet 和 nnFormer 在内的最先进模型。
• 与标准 Transformer 相比，该模型在保持优越分割精度的同时，显著降低了 FLOPs 和参数量。
• 所提出的动态位置编码（DPE）相比绝对位置编码提升了 4.05% 的平均 DSC，相比相对位置编码提升了 2.48%。
• 将 DPE 模块放置在首个局部自注意力模块之前可获得最佳性能，证实其在空间建模中的关键作用。
• 采用扩张注意力的全局感受野模块（GSM）相比仅使用局部感受野模块的模型，精度提升了 0.66%，证明其在长距离建模中的有效性。
• 消融实验表明，扩张注意力机制和 DPE 均对实现低计算成本下的最先进性能至关重要。