[论文解读] NexToU: Efficient Topology-Aware U-Net for Medical Image Segmentation

Pengcheng Shi, Xutao Guo+3 more|arXiv (Cornell University)|May 25, 2023

Medical Imaging and Analysis被引用 5

一句话总结

NexToU 提出了一种新颖的混合 U-Net 架构，通过整合视觉图神经网络模块（Pool GNN 和 Swin GNN），实现高效且具备拓扑感知能力的医学图像分割。通过利用全局与局部拓扑表征，并引入一种二元拓扑交互（BTI）模块，将计算成本从 $c(c-1)$ 降低至 $2(c-1)$，NexToU 在三个多样化数据集上实现了最先进性能，Dice 分数提升且 Hausdorff 距离降低。

ABSTRACT

Convolutional neural networks (CNN) and Transformer variants have emerged as the leading medical image segmentation backbones. Nonetheless, due to their limitations in either preserving global image context or efficiently processing irregular shapes in visual objects, these backbones struggle to effectively integrate information from diverse anatomical regions and reduce inter-individual variability, particularly for the vasculature. Motivated by the successful breakthroughs of graph neural networks (GNN) in capturing topological properties and non-Euclidean relationships across various fields, we propose NexToU, a novel hybrid architecture for medical image segmentation. NexToU comprises improved Pool GNN and Swin GNN modules from Vision GNN (ViG) for learning both global and local topological representations while minimizing computational costs. To address the containment and exclusion relationships among various anatomical structures, we reformulate the topological interaction (TI) module based on the nature of binary trees, rapidly encoding the topological constraints into NexToU. Extensive experiments conducted on three datasets (including distinct imaging dimensions, disease types, and imaging modalities) demonstrate that our method consistently outperforms other state-of-the-art (SOTA) architectures. All the code is publicly available at https://github.com/PengchengShi1220/NexToU.

研究动机与目标

• 解决现有 CNN 和 Transformer 主干网络在处理复杂、不规则解剖结构（尤其是管状血管）时，难以捕捉全局上下文与保持拓扑一致性的挑战。
• 克服当前模型在建模解剖结构之间包含与排除关系方面的局限，避免出现误分割与空间混淆。
• 将图神经网络（GNN）整合至 U-Net 框架中，以捕捉医学图像中非欧几里得的拓扑关系，同时保持计算效率。
• 开发一种轻量化、可微分的拓扑交互机制（BTI），利用二叉树结构编码层级关系，相比先前的 TI 模块显著降低 FLOPs。
• 在多个 2D 和 3D 医学影像数据集上实现最先进分割精度，参数更少且边界分割更清晰。

提出的方法

• 提出一种混合 U-Net 架构 NexToU，集成两种高效视觉图神经网络（EViG）模块：Pool GNN 用于全局拓扑表征，Swin GNN 用于局部窗口化特征学习。
• 利用 Pool GNN 识别图表示中的关键节点，以捕捉远距离解剖区域之间的长程依赖与结构上下文。
• 采用带移位窗口划分与反向操作的 Swin GNN，高效建模局部不规则结构（如血管）。
• 引入二元拓扑交互（BTI）模块，基于二叉树逻辑重构原始 TI 模块，以编码解剖结构之间的层级包含与排除关系。
• 将 TI 模块的计算复杂度从 $c imes (c-1)$ 降低至 $2 imes (c-1)$（$c$ 类分割），显著减少 FLOPs，同时保持拓扑感知能力。
• 将 BTI 模块集成至损失函数中，以训练过程中强制实现拓扑一致性，提升类别边界处的分割精度。

实验结果

研究问题

• RQ1 混合 CNN-GNN 架构能否有效学习医学图像分割中的全局与局部拓扑表征？
• RQ2 如何高效地将解剖结构之间的拓扑约束（如包含与排除关系）编码进深度学习模型？
• RQ3 所提出的 BTI 模块是否在降低计算成本的同时，相比标准拓扑交互模块提升了分割精度？
• RQ4 NexToU 是否能在多种成像模态与解剖结构上实现泛化，包括 2D 和 3D 数据集，涵盖不同疾病类型与分辨率？
• RQ5 引入拓扑感知学习在多大程度上改善了复杂结构（如脑血管）的边界分割并降低了个体间差异性？

主要发现

• 在 BTCV 数据集上，NexToU 实现了 87.84% 的平均 Dice 评分系数（DSC）与 6.33 mm 的 Hausdorff 距离，优于所有 SOTA 方法，包括 nnU-Net、Swin-Unet 和 PHTrans。
• 在 ICA 数据集上，NexToU 实现了 74.19% 的 DSC，为所有对比方法中的最高值，尤其在分割小而复杂的血管（如 AcomA 和 PcomA）方面表现显著提升。
• 在 RAVIR 数据集上，NexToU 实现了 78.21% 的 DSC，超过 nnU-Net（77.54%）与 U-Net++（75.35%），展现出在视网膜血管分割任务中的强大性能。
• 消融实验表明，加入 Pool GNN 和 Swin GNN 模块分别使 DSC 提升 1.48% 和 1.78%，而 BTI 进一部使 DSC 提升 0.61%，HD 降低 1.02 mm。
• 完整版 NexToU 模型仅使用 23.06M 参数即实现卓越性能，展现出高参数效率与在多样化解剖结构及成像模态下的强泛化能力。
• 可视化对比显示，NexToU 减少了误分类与空间混淆（如中脑动脉与脾脏区域），并优于 PHTrans 与 nnFormer，更好地保持了小血管的连通性。