mednext

2023年3月发表在arxiv　　

　　现阶段，人们对于transformer的兴趣激增，然而，由于缺乏大规模的注释医疗数据集，使得实现与自然图像相同的性能具有挑战性，相比之下，卷积网络具有更高的归纳偏差，因此，很容易被训练到高性能。现阶段，ConvNext试图通过镜像transformer来使得ConvNet达到现在的需要，在这项工作中，我们对此进行了改进，设计了一个现代化的和可伸缩的卷积体系结构来应对数据缺乏的挑战，我们称之为MedNeXt，一个受transformer启发的大型卷积核分割网络，他是（1）一个完全的ConvNeXt 3D编码器解码的神经网络 （2）残差ConvNeXt的上、下采样块来保存语义丰富跨尺度（s to preserve semantic richness across scales）没太搞懂什么意思（3）通过上采样小型卷积神经网络迭代增加卷积核尺寸的技术，来防止在受限的医学图像数据上性能达到饱和 （4）在MedNeXT中的多层次（深度、宽度、内核大小）的复合缩放，这就引导了在CT和MRI模式上的4个任务和不同的数据集尺寸上的优越性能

　　Transformer作为一种混合结构或者是单一技术的组件，已经被广泛应用到了医学图像分割，也因此诞生了领先的性能。学习long-rang dependencies是transformer在视觉任务中的主要优势之一.对于long-rang dependencies 的理解是让两个长距离或者短距离的像素点产生关系。就像下图中的点1像素点的落在耕地上，而其它表示耕地的像素点还有很多，利用attention等机制来使得表示耕地的这些像素点产生联系即long-range dependencies。同样也可以使得表示建筑物的像素点之间产生联系。这样有利于利用像素点之间的关系来进行特征提取。（图片来源STANet通过BAM提取图片得到）

 

接着上面讲，由于其有限的归纳偏差，transformer受到需要大型注释数据集来最大化性能效益的困扰。但是，缺乏高质量的注释图片在医学图像领域是非常常见的，为了在利用Transformer的同时保持固有的归纳偏差。ConvNeXt被引入重建对自然图像的卷积神经网络优秀的性能，ConvNeXt架构使用了一个倒置的瓶颈镜像的Transformer，由一个深度层、一个扩展层和一个收缩层组成,除了大型的深度卷积核来复制远程表示学习以外。作者将两个带着大量数据集的大型卷积ConvNeXt结合以超越以前最先进的基于Transformer的网络。相比之下，VGGNet [28]方法仍然是医学图像分割中设计卷积神经网络的主要技术，开箱即用的数据高效解决方案，如nnUNet [13]，使用标准UNet [5]的变体，在广泛的任务中仍然有效。

　　ConvNeXt架构结合了Vision [7]和Swin Transformer的远程空间表示学习能力，还具有卷积神经网络固有的归纳偏差。此外，倒置的瓶颈设计允许我们缩放宽度（增加通道），同时不受卷积核大小的影响，在医学图像分割中对此结构的有效使用将从一下几点受益：（1）凭借大型卷积核学习long-rang dependencies （2）不那么直观的同时扩展多个网络等级。要实现这一点，需要技术来对抗大型网络对有限训练数据过度拟合的趋势。尽管如此，最近还是有人尝试将大型卷积核技术引入医学视觉领域。在[18]中，利用大型卷积核3D-UNet [5]，将核分解为深度的和深度扩张的核，以提高器官和脑肿瘤分割的性能，探索内核缩放（还是内核尺寸，我也不太好翻译），同时使用恒定数量的层和通道。ConvNeXt架构本身被用于3D-UX-Net [17]，其中SwinUNETR [8]的Transformer被ConvNeXt块取代，以在多个分割任务上实现高性能，然而，3D-UX-Net只在标准卷积编码器中部分使用这些块，这限制了它们可能获得的收益。

 

 

 

在这项工作中，我们最大限度地发挥了ConvNeXt设计的潜力，同时独特地解决了医学图像分割中有限的数据集的挑战。.我们提出了第一个完全的ConvNeXt三维分割网络-----MedNeXt，这是一个可伸缩的编码器-解码器网络，并做出以下贡献：（1）我们利用了一个完全由ConvNeXt块构成的网络发挥出了ConvNeXt设计时的全网络范围的优势 （2）我们引入残差倒置瓶颈来代替常规的上采样块和下采样块，来保护再采样的上下文丰富度来优化密集的分割任务，改进的残差连接特别改善了训练过程中的梯度流。（3）我们介绍了一种简单而有效的迭代增加核大小的技术------UpKern，通过对训练好的上采样小内核网络进行初始化，以防止大型内核MedNeXts上的性能饱和。（4）我们提出应用多个网络参数的复合缩放(Compound Scaling)来满足我们的网络设计，满足宽度（通道）、接受域（内核大小）和深度（层数）缩放的正交性

　　MedNeXT实现了领先的性能，超越了基于Transformer、卷积和大型卷积核（kernel）网络，我们展示了在4中不同模态（CT、MRI）和尺寸上的工作上的表现优势，包括器官和肿瘤的分割，我们建议MedNeXt作为一种强大的和现代化的替代方案，用于构建用于医学图像分割的深度网络

2、提出的方法

2.1一个全ConvNeXt 3D分割网络

在之前的工作中，ConvNeXt从vision Transformer和Swin Transformer中提取了优点，形成了一个卷积结构，所以，ConvNeXt块继承了很多重要的设计选择，限制计算开销，同时增加接受域宽度来学习全局特征，这证明了性能要优于标准resnet，.在这项工作中，我们利用这些优势，采用ConvNeXt的一般设计作为类似3d-unet的[5]宏架构的构建块来获得MedNeXt，同样的，我们还将这些块扩展到向上和下采样层，从而产生了第一个用于医学图像分割的完全ConvNeXt架构。MedNeXt块体系结构如图1a所示。具有3层镜像变压器块，并且对c通道输入的描述如下：

 

　　①深度卷积层：这一层包含一个卷积大小为k×k×k的深度卷积，然后用C个输出通道进行归一化，我们使用通道级的组规范（GroupNorm）[32]来保证small batch的稳定性,而不是原来的层规范（layerNorm),深度自然卷积允许在这一层大的卷积核来复制SWin-transformer的大型注意力窗口，同时限制计算，将重活移交给拓展层

　　②拓展层：和Transformer的类似设计相对应，该层包含一个具有CR个输出通道的超完整卷积层，其中R为膨胀比，然后是GELU [12]激活。大的R值允许网络按宽度缩放，而1×1×1内核限制计算。需要注意的是，这一层有效地与前一层的接受域（内核大小）缩放解耦

　　③压缩层：1×1×1内核和C输出通道的卷积层执行特征图的通道级压缩。

　　MedNeXt是卷积的，并保留了卷积vnets固有的归纳偏差，这种偏差允许在稀疏的医疗数据集上进行更容易的训练。我们的完全ConvNeXt架构还允许在标准采样层和上/下采样层上进行宽度（更多的通道）和接受域（更大的内核）的缩放。除了深度缩放（更多的层），我们探索了这三种正交的缩放类型，以设计一个复合可伸缩的MedNeXt，用于有效的医学图像分割

2.2残差倒置瓶颈

　　最初的ConvNeXt架构利用由标准步长卷积组成的独立下采样层。与之相反的便是应用转置卷积来进行上采样操作。然而，这种朴素的设计并未能充分该架构的优势。因此，本文通过将倒置的瓶颈层扩展到MedNeXt中的重采样块来改进这一点。

具体实现上，可以通过在第一个DW层中分别插入步长卷积或转置卷积来完成，以实现可以完成上、下采样的MedNeXt块，如上图绿色和蓝色部分所示。此外，为了使梯度流更容易，作者添加了具有 1 × 1 × 1 卷积或步长为 2 的转置卷积的残差连接。如此一来。便可以充分利用类似 Transformer 的倒置瓶颈的优势，以更低的计算代价保留更丰富的语义信息和所需的空间分辨率，这非常有利于密集预测型的医学图像分割任务。

 

 

2.3无饱和的大核卷积 UpKern

　　大家都知道，提高卷积核的大小意味着增大网络的感受野从而有效提升网络的性能。然而，需要注意的是，这仅仅是理论感受野，而非实际感受野。So sad~~~

因此，最近有许多工作都在探索大卷积核的魔力，据笔者有限的知识储备，目前看到过最高的极限是扩展到61 x 61，有兴趣的读者可以自行去翻阅『CVHub』。讨论回ConvNeXt本身，其卷积核的极限只到7 x 7，根据原著所述再往上增大就“饱和”了。所以，针对医学图像分割这类本身数据就很稀缺的任务来说，如何才能有效的应用和发挥该架构的优势呢？下面看看作者是如何做的。

为了解决这个问题，作者们首先借鉴了Swin Transformer V2的灵感，其中一个大的注意力窗口网络是用另一个较小的注意力窗口训练的网络进行初始化的。此外，作者提议将现有的偏差矩阵空间插值到更大的尺寸作为预训练步骤，而不是从头开始训练，后续的实验也充分验证了此方法的有效性。

 

 

如上图所示，作者对针对卷积核进行了相应的“定制化”，以克服性能饱和问题。其中，UpKern允许我们通过对大小不兼容的卷积核（表示为张量）进行三维线性上采样来初始化具有兼容的预训练小卷积核网络的大卷积核网络，从而迭代地增加卷积核大小。所有其他具有相同张量大小的层（包括归一化层）都通过复制未更改的预训练权重来初始化。

综上所述，以上操作为 MedNeXt 带来了一种简单但有效的初始化技术，可帮助大卷积核网络克服医学图像分割常见的相对有限的数据场景中的性能饱和。

2.4Compound Scaling of Depth, Width and Receptive Field 

下表展示的是作者针对不同的维度采用复合缩放策略来最大化性能的表现：可以看到，相较于常规的上采样和下采样模块，本文方法能够更好的适应不同的任务。

实验：

本文在AMOS22和BTCV数据集上进行了相关的实验，以展示所提方法的有效性，同时也证明了直接应用普通的ConvNeXt架构连现有的分割基线(如nnUNet)都打不过。

从上表可以看出，MedNeXt针对现有的四个主流的医学图像分割数据集均取得了SOTA的性能，而无需额外的训练数据。尽管存在任务异质性（脑和肾肿瘤、器官）、模态（CT、MRI）和训练集大小（BTCV：18 个样本 vs BraTS21：1000 个样本），但 MedNeXt-L 始终优于当前最先进的算法，如nnUNet。此外，借助 UpKern 和 5 × 5 × 5 的卷积核，MedNeXt 利用完全复合缩放进一步改进其自身网络，在器官分割（BTCV、AMOS）和肿瘤分割（KiTS19、BraTS21）方面进行全面改进。

此外，在官方的Leaderboard上，MedNeXt 在 BTCV 任务上轻松击败nnUNet。值得注意的是，这是目前为止仅受监督训练且没有额外训练数据的领先方法之一(DSC：88.76，HD95：15.34)。同样地，对于AMOS22数据集，MedNeXt不仅超过了nnUNet，而且一直占据榜首(DSC：91.77，NSD：84.00)！最后，MedNeXt在另外两个数据集，即KITS19和BraTS21均获得了不错的表现，这一切得益于其优秀的架构设计。

总结：

与自然图像任务相比，由于有限的训练数据等固有的领域挑战，医学图像分割缺乏受益于缩放网络的架构(如ConvNeXt)。本文提出了一种具备高度可扩展性的类ConvNeXt的 3D 分割架构，其在有限的医学图像数据集上优于其它 7 个顶流方法，当中就包含了非常强的nnUNet。MedNeXt设计作为标准卷积块的有效替代，完全可作为医学图像分割领域的新网络架构标杆之作！