[论文解读] `Project & Excite' Modules for Segmentation of Volumetric Medical Scans

Anne-Marie Rickmann, Abhijit Guha Roy+3 more|arXiv (Cornell University)|Jun 11, 2019

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

一句话总结

该论文提出了一种名为 'Project & Excite'（PE）模块的3D卷积神经网络（CNN）重校准模块，通过在三个主轴方向上对特征图进行投影后再进行激励，从而在保留空间信息方面优于传统的挤压-激励（SE）模块（后者采用全局平均池化）。该方法在具有挑战性的全脑和全身体积分割任务中，将3D U-Net的Dice分数提升了高达5%，同时仅使模型复杂度增加了2%。

ABSTRACT

Fully Convolutional Neural Networks (F-CNNs) achieve state-of-the-art performance for image segmentation in medical imaging. Recently, squeeze and excitation (SE) modules and variations thereof have been introduced to recalibrate feature maps channel- and spatial-wise, which can boost performance while only minimally increasing model complexity. So far, the development of SE has focused on 2D images. In this paper, we propose `Project & Excite' (PE) modules that base upon the ideas of SE and extend them to operating on 3D volumetric images. `Project & Excite' does not perform global average pooling, but squeezes feature maps along different slices of a tensor separately to retain more spatial information that is subsequently used in the excitation step. We demonstrate that PE modules can be easily integrated in 3D U-Net, boosting performance by 5% Dice points, while only increasing the model complexity by 2%. We evaluate the PE module on two challenging tasks, whole-brain segmentation of MRI scans and whole-body segmentation of CT scans. Code: https://github.com/ai-med/squeeze_and_excitation

研究动机与目标

• 解决3D全卷积网络（F-CNNs）在小解剖结构分割中因过拟合和训练数据有限而导致的性能瓶颈。
• 克服3D挤压-激励（SE）模块中因全局平均池化导致的空间上下文信息丢失问题，该问题对精细分割至关重要。
• 开发一种轻量化、高效的重校准模块，以增强特征图表示，同时显著降低模型复杂度的增加。
• 在3D U-Net架构中，通过体积医学扫描验证所提模块相较于标准3D SE模块和额外卷积层的优越性。

提出的方法

• 提出 'Project & Excite'（PE）模块作为SE模块在3D空间中的扩展，专为3D体积医学影像设计。
• 将3D SE中的全局平均池化替换为沿三个空间轴（H、W、D）的'投影'操作，以保留空间结构。
• 将沿各轴计算得到的三个投影向量（每轴一个）作为激励子网络的输入，以学习基于空间感知描述符的通道间依赖关系。
• 在3D U-Net的每个编码器和解码器模块后集成PE模块，实现在不显著增加FLOPs或参数量的前提下，为每一阶段提供全局感受野。
• 使用标准交叉熵损失进行训练，并以Dice系数作为评估指标，将PE与基线3D U-Net及3D通道SE（cSE）模块进行比较。
• 以即插即用的方式应用PE模块，确保与现有3D F-CNN架构（如3D U-Net）的兼容性。

实验结果

研究问题

• RQ1 通过轴向投影保留空间信息的3D重校准模块，是否能优于依赖全局平均池化的标准3D SE模块？
• RQ2 PE模块是否能提升3D体积扫描中对小解剖结构（如杏仁核或气管）的分割精度？
• RQ3 PE模块是否能在仅增加2%模型复杂度的前提下，实现与增加额外卷积层相当的性能提升？
• RQ4 PE模块在多种医学影像任务中是否均有效，例如全脑MRI和全身体积CT分割？
• RQ5 与基线模型和3D cSE模块相比，PE模块是否能减少对小结构的性能退化？

主要发现

• 与基线3D U-Net相比，PE模块在MALC数据集（全脑MRI）上将平均Dice分数提升了4%，在Visceral数据集（全身体积CT）上提升了3.5%。
• 对于小结构如下侧脑室、杏仁核和伏隔核，PE模块的Dice分数提升了3–36个百分点，而3D cSE模块则出现性能下降。
• 在气管和胸骨结构上，PE模块的Dice分数分别达到0.845和0.699，而3D cSE模块仅为0.491和0.427，表明其在小而复杂的结构上具有显著优势。
• PE模块仅使模型复杂度增加2%，却使Dice分数提升5%，优于增加两个额外卷积层（复杂度增加40%）的方案。
• 可视化结果证实，PE模型成功分割了如左侧下侧脑室和气管分叉等小结构，而基线3D U-Net和3D cSE模型均未能识别。
• PE模块在全脑和全身体积分割任务中均表现出一致的性能增益，表明其在3D医学图像分割中具有广泛适用性。