MedNeXt: Transformer-driven Scaling of ConvNets for Medical Image Segmentation

---

前言

论文：http://arxiv.org/abs/2303.09975
代码：暂无

缺乏大规模的带注释的医学分割数据集使得医学图像分割实现与自然图像分割相当的性能具有挑战性。卷积网络具有更高的归纳偏置（归纳偏置的解释可以看这个博文【机器学习】浅谈 归纳偏置 (Inductive Bias)）（我感觉意思是更适合处理图像），因此易训练为高性能。最近，ConvNeXt架构试图去通过模仿Transformer去现代化改造ConvNet。在这篇工作中，设计一个现代化和可扩展的卷积架构，针对数据稀缺的医疗环境的挑战进行定制。提出MedNeXt，这是一个基于Transformer的大卷积核分割网络，它引入了
(1) 用于医学图像分割的全 ConvNeXt 3D 编码器-解码器网络
(2) 残差 ConvNeXt 上采样和下采样块，以跨尺度保留语义丰富性
(3) 一种通过对小型内核网络进行上采样来迭代增加内核大小的新技术，以防止有限医疗数据的性能饱和
(4) 在 MedNeXt 的多个级别（深度、宽度、内核大小）进行复合缩放。
这导致了在CT和MRI模态和不同数据集大小的4个任务上的最先进的性能，代表了用于医学图像分割的现代化深度架构。

---

1. Introduction

Transformer被广泛应用，但是由于它有限的归纳偏置，所以被大型的注释数据集最大化性能收益所困扰。为了保持卷积固有的归纳偏置，同时利用Transformer的结构进行改进,最近提出的ConvNext重新建立卷积网络对自然图像处理的竞争性能。
ConvNeXt架构使用Transformer倒置的bottleneck，由深度层、扩展层、收缩曾组成，此外还有大型的深度内核复制长程表征学习，用庞大的数据集训练优于基于Transformer的网络。相反的，VGGNet堆叠小卷积核仍然是设计雨雪图像分割中卷积神经网络的主要技术。开箱即用的数据高效解决方案，如nnUNet ，使用标准UNet的变体，在广泛的任务中仍然有效。
ConvNeXt将Vision和Swin Transformer的远程空间学习能力与ConvNets固有的归纳偏置相结合。
倒置bottleneck设计允许我们扩展宽度(增加channels)不受kernel size的影响。
好好的利用可以得到以下的优势：
(1) 通过大的卷积核学习长程空间依赖
(2) 同时扩展多个网络级别
要实现这些需要有对抗大型网络对有限训练数据过度拟合的趋势。
近期有将大卷积核应用与医学图像分割,a large kernel 3D-Unet将卷积核分为深度核深度扩张的内核，以提高器官和脑肿瘤分割的性能 - 探索内核缩放,同时使用恒定数量的层和通道。3D-UX-Net使用了ConvNext,SwinUNETR的变压器被ConvNeXt块取代，以实现多个分段任务的高性能。
但是他只在标准卷积编码其中使用，限制了它的好处。
在这篇工作中我们最大限度地发挥ConvNeXt设计地潜力，同时独特地解决医学图像分割中有限数据集的挑战。

贡献:

1. purely of ConvNeXt blocks
2. Residual Inverted Bottlenecks
3. UpKern
4. Compound Scaling

2. Proposed Method

2.1 Fully ConvNeXt 3D Segmentation Architecture

我们通过采用ConvNeXt地总体设计作为类似3D-Unet构架中地构建的块来利用这些优势获得MedNeXt。我们也在上采样和下采样中使用这个ConvNeXt块，结果是成为第一个全使用ConvNeXt块的医学分割架构。
MedNeXt blocks有三层模仿Transformer块，被描述成一个C通道的输入。

1. Depthwise Convolution Layer:这一层包含一个Depthwise Convolution，卷积核的大小为k x k x k，接下来是归一化层(GN GroupNorm)，输出通道为C。深度卷积的性质允许该层中大的卷积核复制SwinTransformer的大注意力窗口。同时限制计算，将”heavy lifting“委托给膨胀层。（就是前面一个conv3x3,dilation,groups=channels,对空间信息进行融合，后面一个conv1x1对通道信息进行融合，而且它的参数量小，通道数增加，计算量不会很快增大）。

2. Expansion Layer：包含一个过完备的卷积层，输出通道数为CR，之后接一个GELU的激活函数，R的值越大允许网络在宽度方向上缩放（放大？），同时1x1x1的卷积核限制了计算量。该层有效地将宽度缩放与前一层中的感受野缩放分离。
   
   (a) MedNeXt的架构。该网络有4个编码器层、解码器层和1个bottleneck层。MedNeXt 模块也存在于上采样层和下采样层中。在每个解码层使用深度监督，在较低的分辨率下使用较低的损失权重（途中右边解码器的输出都被拿来算了loss）。所有残差都是累加的，而卷积则被填充以保持张量大小。
   
   (b) 升采样内核( UpKern )初始化一对具有相似配置的MedNeXt架构
   
   © MedNeXt在排行榜的性能。

3. Compression Layer（压缩层）：具有 1 × 1 × 1 个内核和 C 输出通道的卷积层，对特征图执行通道压缩。
   MedNeXt保留了卷积神经网络固有的归纳偏差，使得在稀疏医学数据集上的训练更加容易。我们的全 ConvNeXt 架构还支持在标准层和上采样层进行扩展（更多通道）和感受野（更大的内核）。除了深度缩放（更多层），我们探索了这3种正交类型的缩放设计了一个复合可扩展的MedNeXt用于有效的医学图像分割。

2.2 Resampling with Residual Inverted Bottlenecks

原始的ConvNeXt设计采用单独的下采样层，用标准的跨步卷积，同样的上采样层也可以用标准的跨步转置卷积。但是再重采样的时候它不能利用宽度或基于内核的ConvNeXt缩放。我们通过扩展Inverted Bottleneck到重采样块中来提升。这是通过在第一个深度层中分别插入跨步卷积或转置卷积来完成的，用于下采样和上采样 MedNeXt 块。MedNeXt充分利用Transformer - like Inverted bottlenecks的优势，以较低的空间分辨率在其所有组件中保留丰富的语义信息，这应该有利于密集医学图像分割任务。

2.3 UpKern:Large Kernel Convolutions without Saturation

大卷积核近似Transformers中的大注意力窗口，但仍然容易出现性能饱和。医学图像分割任务的数据量明显较少，性能饱和可能是大型核网络中的一个问题。我们借鉴了Swin Transformer V2的思想，用一个较小的注意力窗口训练另一个网络来初始化一个大的注意力窗口网络。
UpKern 允许我们通过对不兼容大小的卷积内核（表示为张量）进行三线性上采样，用兼容的预训练小内核网络初始化大型内核网络，从而迭代增加内核大小。这为MedNeXt提供了一种简单而有效的初始化技术，有助于大型内核网络克服医学图像分割常见的相对有限的数据场景中的性能饱和。

2.4 Compound Scaling of Depth, Width and Receptive Field

在多个级别（深度、宽度、感受野、分辨率等）上同时缩放提供的好处超出了在一个级别上缩放的好处。在 3D 网络中无限扩展内核大小的计算要求很快变得令人望而却步，并导致我们研究不同级别的同时扩展。与图 1a 一致，我们的缩放测试了块数 （B）、扩展比 （R） 和内核大小 （k） – 对应于深度、宽度和感受野大小。
我们进一步探索大内核大小，并针对每种配置试验k = {3， 5}，通过 MedNeXt 架构的复合扩展最大限度地提高性能。

3. Experimental Design

MedNeXt 在内核大小：{3， 5} 下的 5折交叉验证CV 结果优于 7 个基线——包括卷积、转换器和大型内核网络。

统一框架为所有网络提供了一个通用的测试平台，而不会在补丁大小、间距、增强、训练和评估等方面偏好一个网络。

这种数据集多样性显示了我们的方法在成像模式和训练集大小方面的有效性。

---

4. Result

1. 残差倒置bottleneck，特别是在上采样和下采样层中，在功能上使 MedNeXt（MedNeXt-B 重采样与标准重采样）能够用于医学图像分割。相反，缺少这些修改后的块会导致性能大大降低。这可能是由于在重采样时保留了特征图中的语义丰富性。
2. 训练大型内核网络以进行医学图像分割是一项艰巨的任务，从头开始训练的大型内核 MedNeXts 无法在 MedNeXt-B 中看到（UpKern vs From Scratch）。UpKern 在 BTCV 和 AMOS22 上提高了内核 5 × 5 × 5 的性能，而没有它的大型内核性能与小型内核的性能无法区分。
3. 大型内核的性能提升被认为是由于UpKern与更大的内核的结合，而不仅仅是更长的有效训练计划（Upkern vs Train 2×），因为经过训练的MedNeXt-B再次重新训练了内核3×3×3无法匹配其大型内核对应物。
   这突出表明，MedNeXt修改成功地将ConvNeXt架构转换为医学图像分割。我们进一步根据我们的基线（包括卷积、基于转换器和大型内核基线）在所有 4 个数据集上确定了 MedNeXt 架构的性能。我们在多个层面上讨论了MedNeXt的有效性。

总结

与自然图像分析相比，医学图像分割缺乏从扩展网络中受益的架构，这是由于固有的领域挑战，例如有限的训练数据。MedNeXt 以可扩展的变压器启发的全 ConvNeXt 3D 分割架构呈现，专为在有限的医学图像数据集上实现高性能而定制。我们展示了 MedNeXt 在 4 项具有挑战性的任务中针对 7 条强基线的先进性能。此外，与ConvNeXt类似，用于自然图像.我们提供复合可扩展的 MedNeXt 设计，作为标准卷积块的有效现代化，用于构建用于医学图像分割的深度网络。