[论文解读] TransClaw U-Net: Claw U-Net with Transformers for Medical Image Segmentation

Yao‐Wen Chang, Menghan Hu+2 more|arXiv (Cornell University)|Jul 12, 2021

Advanced Neural Network Applications被引用 10

一句话总结

本文提出 TransClaw U-Net，一种混合深度学习架构，通过在编码器中结合卷积神经网络（CNNs）进行局部特征提取与视觉 Transformer（ViT）模块进行全局上下文建模，实现对医学图像中长距离空间依赖关系的建模。通过将基于 CNN 的特征学习与自注意力机制相结合，该模型在 Synapse 多器官 CT 分割数据集上实现了最先进性能，输入分辨率为 512×512 时平均 DSC 达到 80.39%。

ABSTRACT

In recent years, computer-aided diagnosis has become an increasingly popular topic. Methods based on convolutional neural networks have achieved good performance in medical image segmentation and classification. Due to the limitations of the convolution operation, the long-term spatial features are often not accurately obtained. Hence, we propose a TransClaw U-Net network structure, which combines the convolution operation with the transformer operation in the encoding part. The convolution part is applied for extracting the shallow spatial features to facilitate the recovery of the image resolution after upsampling. The transformer part is used to encode the patches, and the self-attention mechanism is used to obtain global information between sequences. The decoding part retains the bottom upsampling structure for better detail segmentation performance. The experimental results on Synapse Multi-organ Segmentation Datasets show that the performance of TransClaw U-Net is better than other network structures. The ablation experiments also prove the generalization performance of TransClaw U-Net.

研究动机与目标

• 解决 CNN 在医学图像分割中捕捉长程空间依赖关系能力有限的问题。
• 通过在编码器中引入基于自注意力机制的全局上下文建模，提升分割精度。
• 在解码器中保留原始 Claw U-Net 架构的细节特征恢复能力。
• 评估输入分辨率、补丁大小和跳跃连接数量对模型性能的影响。
• 验证所提出架构在腹部 CT 扫描中多个器官上的泛化性与鲁棒性。

提出的方法

• 编码器采用双流架构：一条分支使用标准卷积层提取浅层空间特征，另一条分支将特征图处理为补丁，以实现基于自注意力的全局表征学习。
• 视觉 Transformer 模块对编码特征图中的图像补丁之间应用多头自注意力，以学习长程依赖关系。
• 解码器保留 Claw U-Net 的自底向上上采样结构，以在分辨率恢复过程中保持细粒度空间细节。
• 在对应编码器与解码器层之间建立跳跃连接，以促进特征融合并提高定位精度。
• 采用交叉熵损失与 Dice 损失联合训练模型，通过超参数调优实现 Synapse 数据集上的最优性能。
• 消融研究系统评估了输入尺寸、补丁大小和跳跃连接数量对分割性能的影响。

Fig. 1 : The architecture of TransClaw U-Net

实验结果

研究问题

• RQ1 与标准 CNN 相比，将 Transformer 集成到 U-Net 编码器中，在长程特征学习方面有何改进？
• RQ2 在分割精度与计算成本之间取得平衡时，最优的输入图像分辨率是什么？
• RQ3 补丁大小如何影响视觉 Transformer 组件中的序列长度与模型性能？
• RQ4 改变跳跃连接数量对不同器官分割精度的影响如何？
• RQ5 所提出的 TransClaw U-Net 在腹部 CT 扫描中多样化的解剖结构上是否表现出良好的泛化能力？

主要发现

• 当输入分辨率为 512×512 时，TransClaw U-Net 在 Synapse 数据集上实现平均骰子相似系数（DSC）为 80.39%，优于 224×224 分辨率下的基线模型（DSC 为 78.09%）。
• 将输入分辨率从 224×224 提升至 512×512 显著提高了 DSC，表明更高分辨率的输入可增强特征表征能力并提升分割精度。
• 模型在 16×16 补丁大小下取得最高 DSC（78.09%），而更大的补丁大小（24×24）导致性能下降至 77.27%，尽管计算成本增加。
• 跳跃连接数量具有显著正向影响：完全移除跳跃连接会使主动脉的 DSC 下降 10.8%，胰腺的 DSC 下降 10.16%，证实其在特征融合中的关键作用。
• 消融研究证实，更大的输入尺寸、更小的补丁大小以及更多的跳跃连接可提升性能，但会增加内存占用与训练时间。
• 该模型在 Synapse 数据集中全部八个器官上均表现出强大的泛化能力，在最优设置下，肝脏（95.06%）、脾脏（91.16%）和胃（82.82%）的 DSC 均实现一致提升。

Fig. 2 : Comparison of segmentation results of different models on the Synapse multi-organ CT dataset.