UNet-2022: Exploring Dynamics in Non-isomorphic Architecture

论文地址:https://arxiv.org/pdf/2210.15566.pdf
代码地址:https://bit.ly/3ggyD5G

摘要     最近的医学图像分割模型大多是混合的，将自注意力和卷积层集成到非同构架构中。然而，这些方法的一个潜在缺点是它们未能直观地解释为什么这种混合组合方式是有益的，这使得后续工作难以在它们之上进行改进。为了解决这个问题，我们首先分析了自注意力和卷积的权重分配机制之间的差异。基于此分析，我们建议构建一个并行的非同构块，该块利用自注意力和卷积的优点，并具有简单的并行化。我们将生成的 U 形分割模型命名为 UNet-2022。在实验中，UNet-2022 在范围分割任务中明显优于同类产品，包括腹部多器官分割、自动心脏诊断、神经结构分割和皮肤病变分割，有时超过性能最佳的基线 4%。具体来说，UNet2022 大大超过了目前最受认可的分割模型 nnUNet。这些现象表明 UNet-2022 有可能成为医学图像分割的首选模型 引言

现存问题:

• （a）self-attention，不同的位置有不同的权重，而同一位置的所有通道共享相同的权重。分配的权重在通道维度上不是动态的，从而阻止了 self-attention 捕获不同通道之间的内部差异。
• （b）同一组卷积核权重在不同的空间位置之间共享，而动态权重则分配给不同的通道。
  

主要贡献:

• 我们提供了一个直观的解释，解释了为什么 self-attention 和卷积可以相互补充。核心区别在于权重分配机制的动态特性。自注意力解决了空间动态的重要性，但忽略了通道动态。相反，卷积将动态权重分配给不同的通道而不是空间位置。
• 我们提出了一种新的权重分配机制，将动态权重引入空间和通道维度。权重分配机制的实现非常简单，包括并行独立的自注意力和卷积模块。生成的非同构块为不同的空间位置和通道分配动态权重，使其能够捕获特征图中出现的复杂模式。

模型

模型简介:
UNet-2022 的编码器由一个卷积主干和三个阶段组成，每个阶段涉及三个并行的非同构 (PI) 块。对称地，解码器还包括三个阶段和一个反卷积主干。在每个下采样/上采样步骤中，我们增加/减少通道数并相应地减少/增加特征图的空间分辨率。跳过连接用于弥合低级细节和高级语义之间的差距。
图 2b 描述了并行非同构块的内部结构细节。我们使用深度卷积（DwConv）来减少参数的数量。自注意力和深度卷积并行化以分别探索空间和通道维度的动态，然后将其输出相加并传递到全连接 (FC) 层。在卷积干中，我们堆叠多个卷积层以提取高分辨率特征图。类似的操作也适用于反卷积茎，其中多个反卷积层被堆叠以产生最终的分割掩码。

Parallel non-isomorphic block:
如前所述，self-attention 和卷积强调不同维度的动态，使它们相互补充。受这一发现的启发，我们建议通过对自注意力和卷积的直接并行化将它们的优势整合到非同构块中。
其具体做法见图2（b）

Convolution stem
卷积干由许多堆叠的块（虚线框中的层）组成。每个块由两个卷积层组成，步幅分别为 2 和 1。每个卷积层后面跟着一个 GELU 和一个归一化层。每当特征图的空间分辨率降低一半时，我们相应地将其通道数加倍。其具体做法见图2（c）
为了适应高分辨率的图像，我们在卷积干中堆叠更多的块，以降低内存成本并提高计算效率。具体来说，块的数量，即 log2 P ，取决于我们手动设置的补丁大小 P。例如，对于中小输入分辨率，我们将 P 设置为 4，例如 224×224 和 320×320。当输入分辨率较高时，我们将 P 增加到 8，例如 512 × 512。
De-convolution stem
与卷积主干类似，反卷积主干的架构也因输入图像的分辨率而异。对于小（224 × 224）和中等（320 × 320）输入大小，反卷积茎仅由一个反卷积层组成，其内核大小为4，步幅为4。至于高分辨率输入（即,512×512），我们采用了与卷积茎相似的策略，即增加上采样块的数量（参见图 2d）。在实践中，反卷积茎涉及具有反卷积层的块，以实现对特征图进行上采样以产生最终掩模预测的目标。具体来说，反卷积块的数量为 (log2 P ) − 2，对于小 (P =4) 和中等 (P =4) 分辨率为 0，对于高分辨率输入 (P =8) 为 1 ， 分别。

损失函数:
crossentropy loss and dice loss

分别为1.2和0.8

实验 我们对各种医学图像分割任务进行了实验，包括腹部多器官分割（Synapse）、自动心脏诊断（ACDC）、神经结构分割（EM）、皮肤病变分割（ISIC-2016和PH2）。在 Synapse 和 ACDC 上，我们将 CT/MRI 体积数据转换为切片并将这些切片视为训练数据。实验结果表明，与之前基于 DCNN 和 Transformer 的分割模型相比，我们的 UNet-2022 实现了明显的性能提升。 **实验细节:** 发送到网络的补丁首先从完整图像中随机裁剪。然后，在训练过程中按给定顺序应用旋转、缩放、高斯噪声、高斯模糊、亮度和对比度调整、低分辨率模拟、伽马增强和镜像等增强。请注意，我们对每个图像应用积极的旋转增强，其中旋转角度将从 [0, 15] 中随机选择。 + 实验设备: 具有 11GB 内存的单个 NVIDIA 2080ti GPU

实验结果:

消融实验:
在本节中，我们首先通过将其与 Swin Transformer (ST) 块和 ConvNeXt (CNX) 块进行比较来研究提出的 PI 块的影响。然后，我们进一步研究了 PI 块中 self-attention 和 DwConv 层的影响。接下来，我们展示了基于 ImageNet 的预训练相对于从头开始训练的性能提升。最后，我们提出了一种新的后处理策略，我们发现它在经验上提供了稳定的任务性能改进。

首先，我们从 PI 块中移除自注意力层。结果，生成的构建块无法探索不同空间位置的动态，因为它们只包含卷积操作。这种失败可以通过表 6 第二行中可观察到的任务性能下降来验证。

引入预训练可将平均 DSC 提高 1.5%，并将 HD95 提高近 5mm。

结论 在本文中，我们从权重分配的角度分析了自注意力和卷积的优势。我们展示了自注意力和卷积将动态权重分配给特征空间中的不同维度。在此分析之上，我们建议通过简单地并行化自注意力和卷积操作来构建一个非同构块。由此产生的 UNet-2022 在一系列医学图像分割任务中实现了相当有竞争力的性能，有时比 nnUNet 高出近 4%。未来，我们将研究如何将自监督学习 [34–36,38,41] 适当地纳入 UNet-2022，因为我们发现预训练在提高分割性能方面起着至关重要的作用。

全文仅为个人理解, 如有错误欢迎指正!