[论文解读] Continuous U-Net: Faster, Greater and Noiseless

Chun-Wun Cheng, Christina Runkel+4 more|arXiv (Cornell University)|Feb 1, 2023

Radiomics and Machine Learning in Medical ImagingMedicine被引用 3

一句话总结

本文提出连续U-Net，一种新颖的连续深度学习架构，用于医学图像分割，通过二阶常微分方程（ODEs）建模网络动态。通过用连续动态模块替代离散U-Net模块，该方法实现了更快的收敛速度、更强的抗噪能力，并通过伴随敏感度方法实现恒定的内存消耗，相较于现有U-Net变体在六个基准数据集上表现更优。

ABSTRACT

Image segmentation is a fundamental task in image analysis and clinical practice. The current state-of-the-art techniques are based on U-shape type encoder-decoder networks with skip connections, called U-Net. Despite the powerful performance reported by existing U-Net type networks, they suffer from several major limitations. Issues include the hard coding of the receptive field size, compromising the performance and computational cost, as well as the fact that they do not account for inherent noise in the data. They have problems associated with discrete layers, and do not offer any theoretical underpinning. In this work we introduce continuous U-Net, a novel family of networks for image segmentation. Firstly, continuous U-Net is a continuous deep neural network that introduces new dynamic blocks modelled by second order ordinary differential equations. Secondly, we provide theoretical guarantees for our network demonstrating faster convergence, higher robustness and less sensitivity to noise. Thirdly, we derive qualitative measures to tailor-made segmentation tasks. We demonstrate, through extensive numerical and visual results, that our model outperforms existing U-Net blocks for several medical image segmentation benchmarking datasets.

研究动机与目标

• 为解决离散U-Net架构的局限性，包括感受野固定、计算成本高以及对噪声敏感。
• 基于二阶ODE开发一种理论基础坚实的连续U-Net模块替代方案。
• 通过连续框架中的伴随敏感度方法，实现更快收敛和恒定内存消耗。
• 提升对数据噪声的鲁棒性，并增强在多样化医学影像数据集上的泛化能力。
• 为未来集成注意力机制或Transformer等高级组件的连续U-Net变体奠定基础。

提出的方法

• 网络结构为连续U-Net，其中编码器和解码器模块被由二阶ODE控制的动态模块所取代。
• 动态模块采用二阶神经ODE建模，实现更平滑、更灵活的特征学习轨迹。
• 应用伴随敏感度方法，通过反向ODE积分实现O(1)内存消耗，与模型深度无关。
• 理论分析证明，与标准U-Net模块相比，该方法收敛更快且对噪声更不敏感。
• 评估了不同ODE求解器（欧拉法、AB法、ABM法、RK4法），以确定分割任务的最优积分方法。
• 推导出定性度量，以针对特定分割任务定制网络，提升性能与稳定性。

实验结果

研究问题

• RQ1 基于二阶ODE的连续U-Net架构是否能在医学图像分割中超越标准离散U-Net模块？
• RQ2 通过连续ODE建模U-Net动态是否能实现更快收敛和更低内存消耗？
• RQ3 与标准U-Net变体相比，该方法在数据噪声下的表现如何？
• RQ4 连续框架能否提供关于鲁棒性和收敛性的理论保证？
• RQ5 连续U-Net是否能在无需额外机制（如注意力或Transformer）的情况下实现最先进性能？

主要发现

• 连续U-Net相比离散U-Net模块收敛显著更快，达到最优性能所需迭代次数更少。
• 通过伴随敏感度方法实现O(1)内存消耗，无论网络深度如何均保持计算高效。
• 在所有六个基准数据集中，连续U-Net均优于标准U-Net模块，甚至超越了集成注意力机制的模型。
• 在添加高斯噪声（σ = 0.2）条件下，连续U-Net的Dice分数仅比干净数据下的性能低0.1209，而带有Inception模块的U-Net在GlaS数据集上的Dice分数从0.7434降至0.1930。
• 四阶龙格-库塔（RK4）求解器在所有数据集和指标上均优于欧拉法、亚当斯-巴什福斯法和亚当斯-巴什福斯-莫顿法。
• 连续U-Net在所有数据集和指标上均表现出稳定且高水平的性能，而FO模块在STARE数据集上表现不一致，尤其明显。