[MLP]UNeXt: MLP-based Rapid Medical Image Segmentation Network

from JHU

Abstract

UNet及其Transformer版本-TransUNet在医学图像分割任务中占据主要地位，但是这些算法由于参数量太大、计算复杂、处理缓慢无法用于快速图像分割。因此本文提出基于卷积MLP的图像分割网络UNeXT,包含一个早期的卷积阶段和后续的MLP阶段，其中卷积用于提取特征，MLP用于建模特征表述。
为了进一步提升性能，本文建议输入MLP时改变通道，从而让网络聚焦学习局部特征，而在隐空间中使用标记化的MLP减少了网络参数量和计算量，也能获得更好的特征表示帮助分割。UNeXT还包含编解码之间的跳跃连接。在多个医学图像分割数据集上的实验结果表示，UNeXT的参数量压缩了72x，计算复杂度降低了68x,推理速度提升了10倍，同时获得了优于SOTA的分割结果。

Section I Introduction

医学图像分析在疾病诊断等方面至关重要，其中一项主要任务是医学图像分割任务。UNet是一种高效的encoder-decoder卷积神经网络，目前已经成为医学图像分割的主流框架，并且衍生出了一系列变体，如UNet++,UNet+++，3D UNet，V-Net，Y-Net，KiUNet等。近期也有需要基于Transformer的模型被提出，得益于可以良好的提取全局依赖，进一步提升了分割性能。
TransUNet将ViT结构改进到UNet中进行了2D图像的分割，还有MedT,TransBTS,UNETR等基于Transformer的网络结构用于分割，但是他们大多聚焦于提升分割性能，很少有关注到计算复杂度、推理时间或者参数量的优化，而这些对实际应用场景也至关重要，因为大多数模型都是在实验室环境下测试的，可以借助高算力平台（如GPU）。
但是需要将这些模型迁移到临床场景，也被称之为护理点成像，用于帮助临床医生诊断和改善患者的护理。一方面计算机辅助诊断可以缩短诊断和等待时间，提升患者满意度，比如借助即时超声生成用于检查肺部胸膜不规则，借助摄像头成像诊断皮肤病，借助MRI图像进行疾病的诊断以上都属于医学图像分割、配准等任务。基于深度学习的解决方案如UNet和TrasUNet具有较大的计算复杂度和大量参数，难以应用到临床。

本文致力于设计一个高效的分割网络具备更少的计算开销和参数量，更快的推理时间同时保持良好的性能。因此本文提出UNeXT，对将算法从实验室迁移到临床至关重要。
近期就MLP的网络也被证实可以实现计算机视觉任务，特别是MLP-Mixer是一个基于MLP的网络模型达到了基于Transformer的效果，同时计算成本更低。
因此受此启发本文提出了UNeXT是一个基于卷积和MLP的网络，仍然遵循UNet的5层架构设计，但是每个block的结构进行了改进。UNeXT包含两个stage，卷积阶段和MLP阶段，在网络的最开始和最后使用较少的卷积模块。在瓶颈部分则是使用新提出的Tokenized MLP模块来提取特征同时计算成本比较低。
TokMLP会首先将卷积特征映射到一个抽象的标记中然后使用mlp来学习有意义的信息进行分割，本文还在MLP中引入了移位操作来提取不同轴向为位移对应的局部信息。因为特征化的维度比较低因此可以显著减少计算量和参数量，同时保证了良好的分割性能。
在ISIC皮肤病变数据集和Busi胸部超声数据集上进行了评估，比其他通用分割框架取得了更优异的性能。并且参数量减少了72倍，计算成本下降了68倍，推理时间加速了10倍。

本文的工作总结如下：


（1）提出的UNeXT是第一个基于卷积MLP的医学图像分割框架
;

（2）提出一种新的标记MLP的模块可以进行轴向位移来高效学习隐空间的特征表述;


（3）本文成功的提升了医学图像分割性能并且具有更少的参数量、更快的推理时间和更低的计算复杂度。

Section II UNeXT

Network Design 
UNeXT是一个编解码网络，包含两个stage:





(1)卷积阶段





（2）Tokenized MLP阶段





输入图像首先经过编码器，前三个模块是卷积模块，后面两个是MLP模块。解码器与之镜像。encoder部分每一级分辨率会减半，decoder每一级分辨率加倍。同样使用了skip connection
通道数分别是[32,64,128,160,256],这些参数比原始UNet及其变体都要少，这也是进行的第一个修改，主要为了减少计算量和参数量。






Convolutional stage

每一个卷积模块经过conv->BN->ReLU，卷积核大小3x3，池化选用的最大池化，decoder上采样则用的双线性插值，没有用转置卷积，因为转置卷积也会引入网络参数。



Shifted MLP


在移位MLP中首先会标记卷积特征的通道轴在进行序列化，这样有助于MLP只关注卷积特征的某些位置，从而引入局部性，这有点类似Swin Transformer的基于窗口的注意力。TokMLP包含两个MLP层，本文会在宽度和高度轴上分别进行移位，这有点类似轴向注意力。


本文会将特征分成h份然后根据特定的轴对他们进行唯一，这样有助于沿着轴创建随机窗口引入局部性信息。
Tokenized MLP stage
在 TokMLP模块中首先会将feature进行shift然后进行映射编码成token。因此首选进行3x3卷积，并将通道数变为E，E就是嵌入维度也是token的数目。随后token会送入shifted MLP width,这一步的超参数是MLP的隐层维度H，这一步使用的是DWConv，处于以下两点原因：

（1）有助于编码特征的位置信息。有研究表明MLP中的卷积层足以对位置信息进行编码而且实际上比标准的位置编码效果更好，当测试和训练数据分辨率不相同时ViT中的位置编码通常会带来性能降低。

（2）DWConv参数量更少从而可以加速推理。
随后本文使用GELU激活。
随后送入shifted MLP Height,负责将维度从H变为O，这一步还引入了残差连接。

后续经过LN送入下一模块。LN比BN效果更好，因为是沿着token进行归一化。
整个MLP的计算流程表示如下:

T代表token,H代表高度，W代表宽度。


需要注意的是所有的计算基于嵌入维度H，H远远小于特征图的维度，本文中H=768，这种标记化MLP的方法有助于编码有意义的特征信息，但却不会引入过多的参数量和计算量。

Section III Experiments and Results

Datasets

ISIC2018 2594张图 512 x 512


BUSI 647张 256x256



损失函数：交叉熵损失函数 + dice loss

对比网络：

CNN:UNet,UNet++，ResUNet


Transformer:TransUNet，MedT


评价指标：F1分数，IoU,参数量、GFLOPs、推理时间


Table 1展示了对比结果，UNeXT取得了最优性能，第二名是TransUNet，但是UNeXT因为不包含任何注意力模块因此参数量很少，仅有0.57GFLOPs,而TransUNet是38.52GFLOPs，UNet是55.84.UNeXT也是最轻量级的网络，仅包含1.58M,而TransUNet是105.32M。

本文还对比了在CPU上的运行时间，因为许多临床诊断都没有GPU提供算力。本文取测试10张256图像的平均推理时间作为结果，和Swin-UNet对比发现Swin UNet参数量为41.35M，计算量为11.46 FGLOPs


Fig 4展示了不同网络的分割性能、参数量、推理时间、计算量，可以看到UNeXT均是最优最适合临床的。Fig 5则可视化了不同网络的分割结果，可以看到UNeXT的分割结果最佳。

Section IV Discussion

消融实验

Table 2展示了消融实验的对比结果，本文从UNet开始逐渐减少filter数目来降低参数量和计算复杂度。可以看到精度下降了；随后继续减少网络深度只搭建3层，也就是仅包含UNeXT的Conv STAGE，这一步参数量显著减少，但是精度也掉了4个点。随后引入MLP block可以看到参数量有所增加但是性能却显著提升。
然后在一步步增加DWConv实现的位置嵌入，获得了分割性能的进一步提升。
然后增加MLP中的shift操作，这一步不会增加计算量和参数量，但是性能进一步提升，尤其在两个轴都进行shift会获得最佳性能。
Analysis on number fof channels

UNeXT的通道数是整个网络主要的参数量来源，也决定着最终性能。Table 3展示了不同通道数下的参数量、推理时间、计算量和分割性能。



Difference from MLP-Mixer

MLP-Mixer用的全部是MLP层来进行图像识别。UNeXT则是一个卷积和MLP的网络结构.MLP-Mixer聚焦于channel混合和token混合来学习特征表达，与此相对应本文使用的是卷积提取的特征，LP部分通过shift操作建模特征表述。本文还测试了将MLP-Mixer作为encoder，传统卷积作为decoder，分割性能并不是最优的，并且拥有11M的参数量。

Section V Conclusion

本文提出一种新的深度网络UNeXT进行医学图像分割，更适合临床应用。UNeXT是一个基于卷积和MLP的分割框架，本文提出了tokenized MLP以及shift操作来有效建模特征表述同时更加高效，引入最小的计算复杂度和参数量。本文在多个医学图像分割数据集上验证了UNeXT具有最优的分割性能、最短的推理时间、最少的计算量和参数量。