[论文解读] SegStitch: Multidimensional Transformer for Robust and Efficient Medical Imaging Segmentation

Shengbo Tan, Zeyu Zhang+6 more|arXiv (Cornell University)|Aug 1, 2024

Brain Tumor Detection and ClassificationNeuroscience被引用 3

一句话总结

SegStitch 提出了一种新颖的3D医学图像分割框架，结合了基于冠状面补丁的自注意力机制与去噪常微分方程（ODE）模块，以增强长距离特征建模能力并提升模型鲁棒性。该方法在BTCV数据集上实现了11.48%的mDSC性能提升，在ACDC数据集上实现了6.71%的mDSC性能提升，同时相比UNetR将参数量减少了36.7%，FLOPs降低了10.7%。

ABSTRACT

Medical imaging segmentation plays a significant role in the automatic recognition and analysis of lesions. State-of-the-art methods, particularly those utilizing transformers, have been prominently adopted in 3D semantic segmentation due to their superior performance in scalability and generalizability. However, plain vision transformers encounter challenges due to their neglect of local features and their high computational complexity. To address these challenges, we introduce three key contributions: Firstly, we proposed SegStitch, an innovative architecture that integrates transformers with denoising ODE blocks. Instead of taking whole 3D volumes as inputs, we adapt axial patches and customize patch-wise queries to ensure semantic consistency. Additionally, we conducted extensive experiments on the BTCV and ACDC datasets, achieving improvements up to 11.48% and 6.71% respectively in mDSC, compared to state-of-the-art methods. Lastly, our proposed method demonstrates outstanding efficiency, reducing the number of parameters by 36.7% and the number of FLOPS by 10.7% compared to UNETR. This advancement holds promising potential for adapting our method to real-world clinical practice. The code will be available at https://github.com/goblin327/SegStitch

研究动机与目标

• 为解决标准视觉Transformer在3D医学图像分割中的局限性，特别是其高计算成本和局部特征建模能力差的问题。
• 通过混合注意力机制，提升3D器官与病灶分割中的长距离依赖学习与语义一致性。
• 通过集成神经ODE模块，实现连续的特征优化，从而增强模型的鲁棒性与泛化能力。
• 在不牺牲分割精度的前提下，降低模型复杂度与推理成本，实现临床部署。

提出的方法

• SegStitch采用基于冠状面补丁的查询机制，以保持3D体数据中空间与语义的一致性，避免全体积注意力计算，从而降低计算开销。
• 提出双路径注意力机制，结合粗粒度与细粒度自注意力，以更好地捕捉全局上下文与局部细节。
• 模型集成基于神经记忆ODE（nmODE）的去噪ODE模块，实现在推理过程中对特征图进行连续、动态的更新。
• ODE模块利用学习到的向量场（如sin²、sin、tanh）对特征进行优化，提升对输入扰动与数据噪声的鲁棒性。
• 通过跨补丁共享查询机制，增强补丁间的语义一致性，提升分割结果的一致性。
• 模型在3D医学影像数据集上端到端训练，采用交叉熵损失与Dice损失，并对ODE变体进行了消融实验。

实验结果

研究问题

• RQ1 基于冠状面补丁的注意力机制结合共享查询，是否能提升3D医学图像分割中的语义一致性并降低计算成本？
• RQ2 在存在噪声或扰动的医学影像中，集成ODE模块在提升模型鲁棒性与特征优化方面有何作用？
• RQ3 所提出的SegStitch框架在分割精度与效率方面，相较于SOTA Transformer模型（如UNetR）的性能提升程度如何？
• RQ4 哪种ODE公式（如sin²、sigmoid、tanh）在训练稳定性、泛化能力与分割性能之间实现了最佳平衡？

主要发现

• 与最先进方法相比，SegStitch在BTCV数据集上实现了11.48%的平均Dice相似系数（mDSC）提升。
• 在ACDC数据集上，SegStitch相比现有SOTA方法将mDSC提升了6.71%，在具有挑战性的心脏分割任务中表现出色。
• 与UNetR相比，该模型参数量减少了36.7%，FLOPs降低了10.7%，显示出显著的效率优势。
• 在ACDC数据集上，采用sin²激活函数的ODE模块取得了最佳验证损失（-0.8895）与mDSC（93.32%），优于sigmoid与自编码器消融实验。
• 定性结果表明，对如右心室与肾脏等细长器官的轮廓预测表现更优，证实了其在长距离依赖建模方面的增强能力。
• 注意力权重可视化结果表明，模型聚焦于显著的细长结构，表明其具备有效的全局特征学习能力。