[论文解读] UNet-2022: Exploring Dynamics in Non-isomorphic Architecture

Jiansen Guo, Hong-Yu Zhou+2 more|arXiv (Cornell University)|Oct 27, 2022

Radiomics and Machine Learning in Medical Imaging被引用 8

一句话总结

UNet-2022 提出了一种新颖的非同构 U-Net 架构，通过并行化自注意力与深度可分离卷积模块，联合利用动态空间与通道加权特征，实现了在多个医学图像分割任务中的最先进性能，包括腹部多器官与心脏分割，在 Dice 分数上相比 nnUNet 最高提升 4%。

ABSTRACT

Recent medical image segmentation models are mostly hybrid, which integrate self-attention and convolution layers into the non-isomorphic architecture. However, one potential drawback of these approaches is that they failed to provide an intuitive explanation of why this hybrid combination manner is beneficial, making it difficult for subsequent work to make improvements on top of them. To address this issue, we first analyze the differences between the weight allocation mechanisms of the self-attention and convolution. Based on this analysis, we propose to construct a parallel non-isomorphic block that takes the advantages of self-attention and convolution with simple parallelization. We name the resulting U-shape segmentation model as UNet-2022. In experiments, UNet-2022 obviously outperforms its counterparts in a range segmentation tasks, including abdominal multi-organ segmentation, automatic cardiac diagnosis, neural structures segmentation, and skin lesion segmentation, sometimes surpassing the best performing baseline by 4%. Specifically, UNet-2022 surpasses nnUNet, the most recognized segmentation model at present, by large margins. These phenomena indicate the potential of UNet-2022 to become the model of choice for medical image segmentation.

研究动机与目标

• 为自注意力与卷积在医学图像分割中互补优势提供直观解释。
• 解决现有混合架构中自注意力与卷积结合缺乏明确依据、可解释性不足的问题。
• 设计一种新型模块，实现在空间与通道维度上的动态权重分配。
• 通过简单而有效的架构创新，提升多样化医学影像任务中的分割性能。
• 验证基于 ImageNet 的预训练与推理策略调优在所提出的 UNet-2022 框架中的有效性。

提出的方法

• 提出一种并行非同构模块（PI 模块），通过自注意力与深度可分离卷积（DwConv）模块独立处理特征。
• 引入一种动态权重分配机制，为空间与通道维度分配自适应权重，结合两种模块的优势。
• 使用深度可分离卷积以减少参数量并缓解过拟合，同时保持性能。
• 采用类似 U-Net 的编码器-解码器结构与跳跃连接，将标准编码器替换为所提出的 PI 模块。
• 应用基于 ImageNet 的预训练以提升泛化能力，尤其在数据量较少的情况下。
• 通过滑动窗口推理（步长为 0.2× 裁剪尺寸）优化推理过程，以增强特征聚合。

实验结果

研究问题

• RQ1 为何结合自注意力与卷积的混合架构在医学图像分割中优于纯卷积或纯自注意力模型？
• RQ2 如何在统一框架中形式化比较并利用自注意力与卷积的动态权重分配机制？
• RQ3 在多样化医学影像任务中，并行非同构操作（自注意力与 DwConv）对分割性能有何影响？
• RQ4 与从零开始训练相比，基于 ImageNet 的预训练在多大程度上提升了所提出 UNet-2022 模型的性能？
• RQ5 在滑动窗口推理中，推理步长对最终分割质量有何影响？

主要发现

• 所提出的 PI 模块在 Synapse 数据集上达到 86.46 的平均 Dice 分数与 11.34 的 HD95，优于 Swin Transformer 模块（84.42 DSC，20.74 HD95）与 ConvNeXt 模块（84.96 DSC，16.60 HD95）。
• 从 PI 模块中移除自注意力组件后，平均 DSC 降至 85.20，HD95 上升至 14.96，证实了空间动态性的关键作用。
• 基于 ImageNet 的预训练使 UNet-2022 的平均 DSC 提升 1.5 个百分点，HD95 改善近 5mm，相比从零开始训练。
• 使用更小的推理步长（0.2× 裁剪尺寸）相比 0.5× 裁剪尺寸，HD95 改善 2.6mm，证明了更密集推理采样的优势。
• UNet-2022 在多个任务中（包括腹部多器官与皮肤病变分割）的 Dice 分数相比广泛采用的 nnUNet 模型最高提升 4%。
• 消融实验表明，自注意力与 DwConv 的并行集成比单独使用任一模块或顺序使用更有效。