[论文解读] GCtx-UNet: Efficient Network for Medical Image Segmentation

Khaled Alrfou, Tian Zhao|arXiv (Cornell University)|Jun 9, 2024

Brain Tumor Detection and Classification被引用 6

一句话总结

GCtx-UNet 是一种轻量型的 UNet 风格架构，将 GC-ViT 的全局/局部注意力与 CNN 基于下采样/上采样相结合，在多种医学影像数据集上实现具有竞争力的分割准确性，同时保持较低的模型复杂度。

ABSTRACT

Medical image segmentation is crucial for disease diagnosis and monitoring. Though effective, the current segmentation networks such as UNet struggle with capturing long-range features. More accurate models such as TransUNet, Swin-UNet, and CS-UNet have higher computation complexity. To address this problem, we propose GCtx-UNet, a lightweight segmentation architecture that can capture global and local image features with accuracy better or comparable to the state-of-the-art approaches. GCtx-UNet uses vision transformer that leverages global context self-attention modules joined with local self-attention to model long and short range spatial dependencies. GCtx-UNet is evaluated on the Synapse multi-organ abdominal CT dataset, the ACDC cardiac MRI dataset, and several polyp segmentation datasets. In terms of Dice Similarity Coefficient (DSC) and Hausdorff Distance (HD) metrics, GCtx-UNet outperformed CNN-based and Transformer-based approaches, with notable gains in the segmentation of complex and small anatomical structures. Moreover, GCtx-UNet is much more efficient than the state-of-the-art approaches with smaller model size, lower computation workload, and faster training and inference speed, making it a practical choice for clinical applications.

研究动机与目标

• 在不高计算成本的前提下解决准确医学图像分割的需求。
• 在类 U-Net 架构中整合全局上下文建模与局部注意力。
• 利用 MedNet 预训练以提升同域上的表现，相较于 ImageNet 预训练。
• 在多个数据集（Synapse、ACDC、Polyp）上进行评估，以展示泛化性与效率。

提出的方法

• 使用将局部自注意力与全局上下文查询相结合的 GC-ViT 块，以建模长短程依赖。
• 引入下采样（Fused-MBConv）模块，以注入归纳偏置和通道间建模。
• 在 U 形结构中采用基于 GC-ViT 的编码器–瓶颈–解码器并具有跳跃连接。
• 对 MedNet（医学图像）进行 GC-ViT 预训练，并与 ImageNet 预训练进行比较。
• 使用 patchify 层将图像分割为重叠的补丁，随后在编码器中进行嵌入投影。

实验结果

研究问题

• RQ1 基于 GC-ViT 的模块是否能在参数数量更少的情况下达到与最先进的 CNN/Transformer 基线相当的分割性能？
• RQ2 在领域特定的医学数据（MedNet）上进行预训练，是否会比在自然图像上预训练（ImageNet）提升分割精度？
• RQ3 在多种医学影像任务（CT、MRI、息肉图像）上，GCtx-UNet 在 DSC 与 HD 方面的表现如何？
• RQ4 上采样/下采样设计与超参数对分割性能的影响是什么？
• RQ5 与同类方法相比，该架构在模型大小、FLOPs、训练时间和推理速度方面是否高效？

主要发现

• GCtx-UNet 在取得与最先进方法相竞争的结果的同时，拥有最小的模型规模（12.34M 参数）和在所评估方法中最低的 FLOPs。
• MedNet 预训练的 GCtx-UNet 在多个数据集上通常优于 ImageNet 预训练的变体。
• 在 Synapse 上，GCtx-UNet2（MedNet 预训练）达到 DSC 82.39% 与 HD 15.94 mm，是报道中较佳之一，并且计算量低于许多同行。
• 在 ACDC 上，GCtx-UNet2 的 DSC 分别为 91.23（RV）、89.88（心肌）、87.25（LV），超越了若干基于 Transformer 的和混合模型。
• 在息肉数据集中，GCtx-UNet2 对未见数据集（CVC-ColonDB、ETIS-LaribDB、CVC-300）具有强泛化能力，且 DSC 得分通常处于前列或接近前列。
• 消融研究确定了最佳损失组合（Dice 0.3，交叉熵 0.7）和学习率（0.0001），并显示转置卷积上采样结合 SE 模块可获得最佳性能。