[论文解读] Non-local U-Net for Biomedical Image Segmentation

Zhengyang Wang, Na Zou+2 more|arXiv (Cornell University)|Dec 10, 2018

Advanced Neural Network Applications参考文献 21被引用 7

一句话总结

该论文提出了一种名为非局部U-Net的新架构，用于生物医学图像分割，其通过基于自注意力机制的灵活全局聚合模块替代了U-Net中传统的局部卷积操作。这些模块被整合到保持尺寸、下采样和上采样层中，实现了高效的长距离上下文建模。该方法在3D多模态等信号强度婴儿脑MR图像分割任务中实现了最先进性能，且参数更少、推理速度更快。

ABSTRACT

Deep learning has shown its great promise in various biomedical image segmentation tasks. Existing models are typically based on U-Net and rely on an encoder-decoder architecture with stacked local operators to aggregate long-range information gradually. However, only using the local operators limits the efficiency and effectiveness. In this work, we propose the non-local U-Nets, which are equipped with flexible global aggregation blocks, for biomedical image segmentation. These blocks can be inserted into U-Net as size-preserving processes, as well as down-sampling and up-sampling layers. We perform thorough experiments on the 3D multimodality isointense infant brain MR image segmentation task to evaluate the non-local U-Nets. Results show that our proposed models achieve top performances with fewer parameters and faster computation.

研究动机与目标

• 为解决标准U-Net架构中依赖堆叠局部操作导致的效率低下和长距离上下文聚合能力有限的问题。
• 开发一种灵活的全局聚合机制，可应用于U-Net的所有组件：保持尺寸层、下采样层和上采样层。
• 在提升3D生物医学图像分割精度的同时，降低模型复杂度和推理时间。
• 评估全局聚合在编码器和解码器路径中的有效性，特别是在上采样过程中恢复空间细节的能力。

提出的方法

• 提出一种基于自注意力机制的全局聚合模块，通过计算特征图所有位置的加权和，实现在不依赖深层堆叠的情况下进行长距离上下文建模。
• 将全局聚合模块适配于下采样和上采样层，而不仅限于保持尺寸的模块，从而在整个网络中实现一致的全局上下文整合。
• 将U-Net中基于拼接的跳跃连接替换为求和操作，以减少参数量并支持长距离残差学习。
• 采用基于图像块的推理策略，结合重叠感受野，以提升分割精度，尤其适用于大尺寸3D图像。
• 通过网格搜索优化图像块大小和重叠步长，实现推理速度与性能之间的平衡。
• 在9个受试者上进行模型训练，并在第10个受试者上进行评估，任务为3D多模态等信号强度婴儿脑MR图像分割。

实验结果

研究问题

• RQ1 与堆叠的局部卷积相比，通过自注意力实现的全局聚合是否能提升3D生物医学图像的分割性能？
• RQ2 在跳跃连接中用求和替代拼接操作，是否能在保持或提升性能的同时减少模型参数？
• RQ3 全局聚合模块是否能有效应用于U-Net中的下采样和上采样层，而不仅限于瓶颈层或保持尺寸层？
• RQ4 图像块大小和重叠步长如何影响3D U-Net变体的推理速度与分割精度？
• RQ5 所提出的非局部U-Net是否在参数更少、计算更快的前提下，优于先前的最先进模型？

主要发现

• 所提出的非局部U-Net在第10个受试者上取得了0.4077的平均Dice分数，优于所有基线模型，包括3D U-Net和CC-3D-FCN。
• 性能最佳的模型（非局部U-Net）Dice分数为0.4077，而最佳基线模型（Model4）为0.4220，表明性能有显著提升。
• 消融实验表明，将底部块替换为全局聚合块（Model5）带来了最大的性能增益，凸显了瓶颈层中全局上下文的重要性。
• 最优图像块大小为32³，实现了最高Dice分数；过小或过大的图像块均导致性能下降。
• 重叠步长为8或16时，在精度与推理速度之间实现了最佳平衡，减少了需处理的图像块数量，同时保持了高性能。
• 非局部U-Net在参数更少、计算更快的前提下实现了最先进性能，充分证明了其高效性与有效性。