nnUNet原创团队全新力作！MedNeXt：医学图像分割新SOTA
Title：MedNeXt: Transformer-driven Scaling of ConvNets for Medical Image Segmentation
MedNeXt：用于医学图像分割的 ConvNets 的变压器驱动缩放
原论文链接：https://arxiv.org/abs/2303.09975
代码链接：暂未开源
导读：MedNeXt是nnUNet原创团队于2023年3月17日上传至arxiv上的新作品，该模型受ConNeXt启发，根据Transformer改进了现有的卷积网络，实现了医学图像分割领域的SOTA。ConvNeXt的任务是改进通用模型，因此根据Transformer改进了ResNet，而医学图像分割领域中的baseline是UNet，所以MedNeXt按照Transformer改进了UNet。MedNeXt在四个具有代表性的医学图像处理任务上实现了SOTA。除了①用Transformer改造UNet之外，MedNeXt还②改进了上采样和下采样块③提出了一个用小卷积核网络初始化大卷积核网络的方法④设计了可以对网络在深度、宽度和感受野大小上进行复合缩放的机制。主要干了这四件事，对应于第二章的四个小结。
摘要
对于基于 Transformer 的架构进行医学图像分割的兴趣当前已呈爆炸式增长。然而，医学数据集缺乏大规模注释，这使得在医学图像处理领域实现Transformer在自然图像处理领域的性能颇具挑战性。相比之下，卷积网络（ConvNet）具有更高的归纳偏差，也因此更易训练。近期，ConvNeXt架构试图通过模仿Transformer Block使标准的ConvNet“现代化”。在本文中，作者对此进行了改进，以设计一个现代化且可扩展的卷积架构，该架构专为数据稀缺的医学图像设计。作者提出了MedNeXt，一种受Transformer启发的大卷积核分割网络。提出了以下四点： 1) 用于医学图像分割的纯ConvNeXt 3D 编码器-解码器网络；2) 具有残差连接的ConvNeXt上采样和下采样block，以保持跨尺度的语义丰富性；3) 为防止在有限医疗数据下的性能饱和，提出了一种通过对小卷积核网络进行上采样来迭代增加卷积核尺寸的新技术；4）MedNeXt 可以实现在深度、宽度、卷积核大小上多个级别的复合缩放。这些设计使MedNeXt在不同规模的CT、多模态MRI数据集的4个任务上实现了最先进的性能，代表了用于医学图像分割的现代化深度架构。代码在此：[即将添加（to be added soon）]
1 Introduction
Transformer在医学图像领域已被广泛使用，包括纯Transformer结构或混合架构，其中一些已经实现了SOTA的性能。学习远程空间依赖性的能力是Transformer架构在视觉任务中的主要优势之一。然而，Transformer受到数据集规模的困扰，因为它们必须由大型标注数据集驱动自身的归纳偏差使性能优势最大化。虽然大型数据集（如ImageNet-1k、ImageNet-21k）在自然图像领域很常见，但在医学图像领域很少。为了获取Transformer架构优势的同时保留卷积固有的归纳偏差，最近一篇名为ConvNeXt的工作这样做了，并重新建立了卷积在自然图像处理领域的竞争力。 ConvNeXt 架构通过倒置bottle neck来模仿Transformers，由深度层、扩展层和收缩层组成，此外还用大卷积核来模仿Transformer对远距离关系的表示。该文利用大卷积核的ConvNeXt网络与大型数据集匹配，超越了当时最先进的基于Transformer的网络。相比之下，在医学图像分割中的主流卷积网络是堆叠小卷积核的VGGNet。使用UNet变体打造的开箱即用的高效解决方案，如nnUNet，在医学图像处理任务中仍是主流。
ConvNeXt 架构将ViT和Swin Transformers的远程空间表示学习能力与卷积固有的归纳偏差相结合。倒置瓶颈设计允许在不受卷积核大小影响的情况下扩展宽度（增加通道数）。这对医学图像分割的启发如下：1）通过大卷积核学习远距离语义关联性，2）可同时缩放多级别网络。要实现这一目标，需要一些技术来弱化大型网络在有限训练数据上的过拟合趋势。尽管如此，最近还是有学者尝试将大卷积核引入医学视觉领域。在3D LK convolution中，使用大核3D-UNet将核分解为深度核和深度扩张核，以提高器官和脑肿瘤分割的性能——探索核缩放，同时使用恒定数量的层和通道。 ConvNeXt 架构本身被用于 3D-UX-Net，该工作将SwinUNETR的Transformer部分由ConvNeXt块取代，在多个分割任务上实现了高性能。然而，3D-UX-Net仅在标准卷积编码器中部分使用这些块，没有完全发挥ConNeXt的优势。
在本文中，作者最大限度地发挥了ConvNeXt设计的潜能，同时专门解决了医学图像分割中数据有限的挑战。作者提出了第一个纯ConvNeXt 3D分割网络 MedNeXt，一个可扩展的编码器-解码器网络，并做出以下贡献：
利用一个纯ConvNeXt组成的架构，实现了ConvNeXt设计的网络宽度优势。 （第2.1节）
引入了Residual Inverted Bottlenecks来代替常规的上采样和下采样块，以在重采样时保持上下文的丰富性。修改后的残差连接改善了训练期间的梯度传播。（第2.2节）
提出了一种简单但有效的迭代增加内核大小的技术UpKern，通过使用预训练的上采样小卷积核网络进行初始化来防止大卷积核MedNeXts的性能饱和。（第2.3节）
提出应用多种网络参数的复合缩放的网络结构，允许宽度（通道）、感受野（卷积核大小）和深度（层数）同时缩放。（第2.4节）
2 提出的方法
2.1 纯ConvNeXt的3D分割架构
2.1 Fully ConvNeXt 3D Segmentation Architecture
ConvNeXt 块继承了 Transformers 的许多重要设计选择，旨在降低计算成本，同时扩大感受野以学习全局特征，实现了对标准ResNets的超越。本工作中采用ConvNeXt的设计来改进3D-UNet的宏观架构以得到最终的MedNeXt。此外，作者还将ConvNeXt块扩展到上采样层和下采样层。宏架构如图1a所示。
MedNeXt 块（黄色），输入通道为C，共有3层，和ConNeXt中的block几乎一样：
Depthwise Convolution Layer：该层包含一个内核大小为k×k×k的深度可分离卷积，然后是归一化，具有C个输出通道。使用GroupNorm代替LayerNorm来稳定小batch size时的性能（因为3D医学数据比较大，GPU一般只能塞得下小于等于2的batch）。使用深度可分离卷积的原因在ConvNeXt论文中有写，是为了模仿Transformer的架构并获相应的归纳偏置，而在Transformer中，特征提取在通道和空间上总是分开进行的，因此使用深度可分离卷积将通道和空间的提取分开进行。此外，还能降低参数量，降低后就有更大的空间去扩大宽度了。
Expansion Layer：与 Transformers类似，该层包含一个带有CR通道的过完备卷积层，其中R是扩展率，然后由GELU激活。R值变大则网络变宽，而1×1×1卷积核尺寸以限制计算量。重要的是，该层有效地将宽度缩放与上一层中的感受野缩放（k）缩放分离。
Compression Layer: 具有1×1×1卷积核和C 输出通道数的卷积层，对特征图执行逐通道的压缩。
MedNeXt是纯卷积的，保留了ConvNets的归纳偏差，可以更轻松地对稀疏医学数据集进行训练。完全ConvNeXt架构还可以在标准和上/下采样层上实现宽度（更多通道）和感受野（更大的卷积核）缩放。除了深度缩放（更多层），作者还探索了这 3 种正交类型的缩放，以设计可复合缩放的 MedNeXt 以实现有效的医学图像分割（第2.4节）。
2.2 使用残差倒置瓶颈进行重采样
2.2 Resampling with Residual Inverted Bottlenecks
ConvNeXt的重采样层是跨步卷积，然而这种设计无法实现缩放。因此，作者通过将倒置瓶颈扩展到MedNeXt中的重采样块来实现缩放。做法如图1（a）中的上采样块（绿色）和下采样块（蓝色）。相当于把MedNeXt块（黄色）中的DW卷积改成了跨步或转置，然后为了满足通道数的要求，残差连接得加一个1×1×1的卷积压缩通道。
2.3 UpKern：不饱和的大核卷积
2.3 UpKern: Large Kernel Convolutions without Saturation
大卷积核近似于Transformers中的大注意力窗口，但易出现性能饱和。尽管自然图像分类数据多，但ConvNeXt架构在7 × 7大小时就饱和了。而医学图像分割任务的数据少之又少，性能更容易饱和。为了提出一个解决方案，作者借鉴了Swin Transformer V2的灵感，用一个预训练的小窗口网络来初始化大窗口网络。具体来说，V2的作者提出将预训练的偏置矩阵进行空间插值到更大的尺寸，这样就不用重新训练了，而且性能也提高了。作者提出了一种类似的方法，但针对卷积核进行了定制，如图1（b） 所示，以克服性能饱和。 UpKern 通过对预训练小核网络进行三线性上采样来初始化大核网络，从而迭代地增加核大小。其他的大小相同的层（包括归一化层）都通过直接复制预训练层的权重来初始化。这为MedNeXt 带来了一种简单但有效的初始化技术，可帮助大卷积核网络克服医学图像分割常见的有限数据的性能饱和。
2.4 深度、宽度和感受野的复合缩放
2.4 Compound Scaling of Depth, Width and Receptive Field
复合缩放是指在多个级别（深度、宽度、感受野、分辨率等）上同时缩放，比只在单一级别上缩放的好处更多。在自然图像处理领域，大的baseline基本都提供了多种规模的配置，主要是提供了不同深度的模型，如ResNet101、ResNet50，SwinT、SwinL等等。提供了多个版本也就意味着模型更加灵活，可以根据任务的性质和数据量灵活缩放模型。而在医学图像处理领域，一直没有能够缩放的模型，之前都是一篇工作一个模型，无法调整。因此，作者希望设计一个可以缩放的模型，适应更多的场景。但医学图像数据少，盲目的增加深度往往容易过拟合。在3D网络中只缩放卷积核尺寸又会使计算量激增，这也让作者开始另辟蹊径，研究复合的缩放。作者测试了对block数 (B)、扩展率 (R) 和内核大小 (k)进行缩放，这些参数对应于深度、宽度和感受野大小。使用了MedNeXt的4种模型配置，如表 1（左）中详述。
基础的MedNeXt-S使用的C=32，R=2，B=2。然后对R和B进行了单独或同时的缩放，并对每种配置都测试了k=3或k=5的情况。
3 实验设计
3 Experimental Design
3.1 配置、实现和基线
3.1 Configurations, Implementation and Baselines
用PyTorch搭建，使用2种卷积核大小的MedNeXt 的4种配置进行实验。实验框架使用 nnUNet为主干（epochs = 1000，batches per epoch = 250），推理（50％ 补丁重叠），数据增强使用nnUNet中默认的。所有网络都用AdamW作为优化器进行训练。在训练和推理过程中，数据被重新采样到1.0 mm各向同性间距（原始间距的结果），分别对 3D和2D网络使用128×128×128和512×512的输入尺寸，batch size为 2 和14。学习率为 0.001，但KiTS19中的k=5除外，它使用0.0001来保证稳定性。对于Baselines，所有 Swin模型和3D-UX-Net使用0.0025的学习率，而ViT模型使用0.0001。使用Dice相似性系数 (Dice Similarity Coefficient，DSC) 和表面Dice相似性 (Surface Dice Similarity，SDC) 作为评估指标。所有模型使用 8:2分的5折交叉验证平均性能监督训练。
基线包括1个高性能的卷积网络nnUNet，4个卷积-变压器混合网络UNETR、SwinUNETR、TransBTS、TransUNet，1个纯Transformer网络，nnFormer以及1个非完全ConvNeXt 网络3D -UX-Net，一共7个用来对比。
3.2 数据集
3.2 Datasets
用了4个医学分割中的热门数据集，包括器官和肿瘤分割，来全面展示MedNeXt架构的优势：
Beyond-the-Cranial-Vault (BTCV) 腹部CT器官分割
AMOS22 腹部 CT器官分割
肾脏肿瘤分割挑战2019数据集 (KiTS19)
脑肿瘤分割挑战 2021 (BraTS21）
BTCV、AMOS22 和 KiTS19 数据集分别包含30、200和210个CT体，分别有13、15和2个类别，而BraTS21数据集包含 1251个MRI体和3个类别。这种多样性显示了本文的方法在成像模式和训练集大小方面的有效性。
4 结果和讨论
4 Results and Discussion
4.1 架构改进的性能消融
4.1 Performance ablation of architectural improvements
消融了AMOS22和BTCV数据集上的MedNeXt-B配置，以突出改进的功效，并证明普通ConvNeXt不能直接用，如在表 1（右）。
实验观察有三点：
Residual Inverted Bottlenecks，使MedNeXt能够用于医学图像分割。相反，如果没有Residual Inverted Bottlenecks会导致性能显着下降。这可能是由于在重采样时保留了特征图中丰富的语义。
为医学图像分割训练大核网络很有必要，重新训练的大型内核MedNeXts无法在MedNeXt-B中执行。而UpKern 在BTCV和AMOS22上将内核性能提高了。 说明UpKern是有效的。
大内核的性能提升被认为是由UpKern与更大内核的结合。说明UpKern确实同时提升了尺寸和性能，并不是简单的多训练几次就能实现的，因为他们试了一下多训练k=3的网络，并没有达到UpKern达到的性能。
这表明MedNeXt修改成功地将ConvNeXt架构转化为医学图像分割模型。
4.2 与基线的性能比较
4.2 Performance comparison to baselines
如表2所示，MedNeXt在性能上全面超越了现有SOTA们。并且在图1（c)中给出了MedNeXt在各大数据集官网上的排名，总之就是赢麻了！
5 结论
5 Conclusion
与自然图像分析相比，由于数据有限等固有的领域挑战，医学图像分割一直缺乏有效的可缩放网络的架构。在这项工作中，MedNeXt被呈现为可扩展的受Transformer启发的纯ConvNeXt 3D分割架构，专门为在有限的医学图像数据集上实现高性能而定制。针对7个强大的基线展示了MedNeXt在4个具有挑战的任务中的最先进性能。此外，类似于用于自然图像的ConvNeXt，本文提供了可复合缩放的MedNeXt 设计作为标准卷积的有效“现代化”，构建了专用与医学图像分割的深度网络。
相关论文链接
nnUNet：https://www.nature.com/articles/s41592-020-01008-z
ConvNeXt：https://openaccess.thecvf.com/content/CVPR2022/html/Liu_A_ConvNet_for_the_2020s_CVPR_2022_paper.html
除此之外，文中有些重要引用也增加了超链接。