nnFormer (Not-aNother transFORMER): 基于交叉Transformer结构的3D医疗影像分割网络

文章链接：https://arxiv.org/abs/2109.03201
开源代码和模型：https://github.com/282857341/nnFormer

▊ Highlight

1 相比较Swin-UNet，nnFormer在多器官分割任务上可以取得7个百分点的提升。

2 相较于传统的基于体素（voxel）计算self-attention的模式，nnFormer采用了一种基于局部三维图像块的计算方式，可以将计算复杂度降低90以上。

3 即使与目前最强劲的nnUNet相比，nnFormer仍然取得了小幅度的提升。

▊ 研究背景介绍

由于Transformer本身可以有效地捕捉和利用像素或体素之间的长期依赖（long-term dependencies），近期出现了非常多结合CNN和Transformer的针对医疗影像处理的模型和网络。其中大部分结果表明，在CNN中合适的位置嵌入类Transformer的结构，可以有效地提升网络的性能。

基于我们的观察，这些基于Transformer的医疗影像处理模型和网络通常可以分为两类:

i）仍然使用CNN作为主要的特征提取器，辅以类Transformer结构以捕捉特征中的全局信息，再将此信息嵌入到CNN中；

ii）直接使用纯Transformer结构进行处理。

不可否认的是第ii类的方法确实是一种更大胆的尝试，在计算机视觉领域也出现了诸如Swin Transformer一类的非常优秀的（纯）Transformer的网络模型。另外一方面，最近的一些研究和方法表明，第i类建立混合网络的方法或许可以更好地结合CNN和Transformer的特点，从而提供更高的模型性能。

但是正如我们在上面已经提到的，在医疗影像处理领域中，第i类方法通常是将CNN作为主力来使用。我们认为这种方式并不一定是最优的，如何更好地结合CNN和Transformer，这是我们提出nnFormer的主要背景和动机。

如上图所示，nnFormer是由厦门大学&香港大学提出的一种基于交叉式结构的Transformer模型，主要用来处理医学影像分析中的3D影像分割问题。目前文章、代码和模型均已放出。（链接在最上方）欢迎大家使用之后提出宝贵的意见和建议！

▊ 方法

先讲一下nnFormer主要用到的交叉式的主干结构（interleaved stem）和背后的一些思考。这里需要强调的是nnFormer使用的结合CNN和Transformer的方式是启发式的，是建立在我们从大量实验中总结得到的经验基础上的，因为并不具有很强的理论根据。

言归正传，如图1所示，nnFormer的主干结构包含了三部分的内容。第一部分是一个四层的卷积结构，主要用来将输入的影像转化为网络可以处理的特征。之所以使用四层的卷积来处理输入，有两个原因：

其一我们认为卷积网络可以更好的保留更加精确的位置信息；

其二卷积操作可以提供高分辨率的底层特征，这是后面应用Transformer blocks的基础。

具体来说，我们认为ViT里面那种patch-level的positional encoding的方式，其实是有点粗暴的，首先它没有显式地嵌入pixel-level的位置信息，其次位置信息并没有跟图像信息更好地融合就直接输入到transformer，这可能会带来一定的问题。

另外一方面，根据我们已有的经验，CNN在low-level的任务上还是非常强大的，直接忽视掉CNN这方面的优势或许并不可取。

经过卷积结构处理之后的特征会送入到Transformer blocks之中。跟Swin-UNet类似，nnFormer中的Transformer blocks也是基于Swin Transformer的（这里需要再次感谢Swin Transformer的作者们提供了非常好用的tools和models！）。

但是主要的不同点在于，nnFormer使用的是一种基于局部三维图像块的self-attention计算方式（文章中叫做 volume-based multi-head self-attention，也就是V-MSA）。

相较于传统的voxel和voxel之间计算self-attention的方式，V-MSA可以大大地降低计算的复杂度。经过我们的粗略估计，在Synapse和ACDC两个经典数据集上，V-MSA相较于MSA可以分别降低约98%和99.5%的计算复杂度。下图是V-MSA的计算复杂度，感兴趣的小伙伴可以具体算一下。

这里有个很有趣的现象，我们发现这些降低的复杂度主要集中在网络早期的计算过程中，伴随着特征空间维度的下降（H,W,D）以及通道输入（C）的增多，其实这种优势就不明显了。

不负责任的猜测一下，也许在high-level直接使用MSA效果会更好？其它的positional encoding和QKV-attention的计算过程，大家可以去读一下paper，应该很容易就能够理解。

除了卷积层和Transformer blocks之外，我们还是使用了Convolutional down-sampling（也就是stride>1的卷积）来进行下采样。进行下采样的原因有二：

其一，是多次下采样可以建立多尺度的特征金字塔结构（feature pyramid），这其实也是CNN能够在检测和分割领域取得很大进展的一大利器；

其二，下采样可以大大降低GPU显存的消耗，这同时也为nnFormer带来了一个很大的优势，单张英伟达2080GPU可以在不到两天的时间内就可以跑完1000个epoch（Synapse或者ACDC），所以小伙伴们再也不用担心缺卡的问题了，毕竟我们也是这么过来的 -_- 。

nnFormer的整体结构如下：

▊ 实验结果

最后简单说一下nnFormer的性能指标，在Synapse（多器官分割）和ACDC（心脏病诊断）上我们都提供了非常详细的性能对比结果。所有的结果都是多次实验结果的平均，所以即使微弱的提升也还是有一定说服力的-_-（大家轻喷）。

这里提一句，Synapse和ACDC是最近各种Transformer paper里面最常用的两个数据集，这也是我们为什么选择他们的原因。

如上图所示，在Synapse上，nnFormer可以大幅度的超过以往的Transformer结构。即使跟nnUNet相比，nnFormer也不遑多让，比如在胃（stomach）这一类别上，超越nnUNet约4个点，超过了SwinUNet约13个点。

再看ACDC：

这次我们的性能提升没有那么大了，相比较以往的Transformer结构大概有1.5-2个点提升。但是考虑到大家的结果都挺高了（>90），我们认为这种程度的提升还是可以接受的。

这里最后提一下我们经过ablation study得到的几个结论，与大家共勉：

1 预训练是真的有用！nnFormer使用的是自然图像上预训练的模型，如果使用医疗影像的预训练模型，效果应该还可以更好。另外一方面，nnUNet没有使用任何预训练就能够取得这么好的效果，也真是非常强大了。

未来我们可能会尝试更加轻量级的nnFormer或者引入自监督进一步提升性能！也欢迎感兴趣的小伙伴加入，我们可以一起做出nnFormer+或者nnFormer++。

2 最开始的卷积结构很有用。这也说明了目前基于局部的处理方式在图像处理里面还是有一席之地的。

3 Transformer blocks并不一定是越多越好。这个特点在医疗影像分割任务上尤其显著，因为分割的任务的数据量比较小，所以一个更加简单的网络结构或者加入一定程度的预训练是有必要的（老生常谈了这是）。

最后鸣谢一下Swin Transformer和nnUNet的作者们，其实往小里说，nnFormer不过是基于Swin Transformer和nnUNet的经验结合，technical上的novelty并不多。但是往大里说的话，nnFormer其实是一个很好的起点，可以启发更多的人投入到相关的topic中开发出更好的基于Transformer的医疗影像分析模型。也许大幅度超越nnUNet的那一天并不遥远。