Revitalizing MLP’s Ability to Efficiently Extract Long-Distance Dependencies for Medical Image

Rolling-Unet: 重振MLP有效提取医学图像分割远程依赖关系的能力

论文链接：https://ojs.aaai.org/index.php/AAAI/article/view/28173

项目链接：https://github.com/Jiaoyang45/Rolling-Unet

Abstract

基于深度学习网络的医学图像分割方法主要分为CNN和Transformer。然而，CNN很难捕获长距离依赖关系，而Transformer的计算复杂度高，局部特征学习能力差。为了有效地提取和融合局部特征和远程依赖关系，本文提出了一种结合MLP的CNN模型Rolling-Unet。具体来说，我们提出了核心R-MLP模块，该模块负责学习整个图像在单一方向上的长距离依赖关系。通过对不同方向的R-MLP模块进行控制和组合，形成OR-MLP和DOR-MLP模块，以捕获多方向的远程依赖关系。此外，在不增加计算负担的情况下，提出了Lo2块对本地上下文信息和远程依赖关系进行编码。Lo2块具有与3×3卷积相同的参数大小和计算复杂度。在四个公共数据集上的实验结果表明，Rolling-Unet的性能优于当前的方法。

Introduction

随着计算机技术和人工智能的飞速发展，卷积神经网络(CNN)强大的建模能力得到了广泛的研究。基于深度学习的分割算法也被引入到医学图像中。U-Net (Ronneberger, Fischer, and Brox 2015)是医学图像分割领域最著名的网络架构之一，它是一个全卷积分割网络。U-Net的编码器和解码器是对称的，形成一个u型段，并通过跳过连接融合来自不同阶段的特征图。U-Net可以适应较小的训练集，输出更准确的分割结果。这一优势使得U-Net取得了巨大的成功并得到了广泛的应用。遵循这一技术路线，诸如UNet++ (Zhou等人2018)，AttUNet (Oktay等人2018)，3D U-Net (C°C°ek等人2016)和V-Net (Milletari, Navab和Ahmadi 2016)已被开发用于各种医学成像模式的图像和体积分割。尽管这些方法能很好地执行，尽管这些方法表现良好，但由于卷积运算的固有局部性，纯CNN架构很难学习清晰的全局和远程语义信息(Chen et al 2021)。

为了克服CNN的局限性，受到Transformer在自然语言处理(NLP)领域取得巨大成功的启发，研究人员试图将Transformer引入视觉领域(Carion et al 2020)。视觉转换器(Vision Transformer, ViT) (Dosovitskiy et al . 2020)完全基于多头自注意机制，使网络能够捕获远程依赖并编码形状表示。然而，它需要大量的训练数据才能达到良好的性能。此外，它具有很高的计算复杂度，这使得网络无法支持高分辨率输入(Azad et al 2022)。Swin Transformer (Liu et al . 2021)减少了计算量，但代价是其窗口之间没有信息交互，从而导致更小的感受野。与CNN模型相比，纯Transformer模型在捕获局部表示方面也表现不佳(Chen et al . 2021)。鉴于CNN和Transformer的特点，一些方法尝试将CNN和Transformer结合起来(Chen et al . 2021;Valanarasu等2021;Wang et al . 2022)进一步增强网络的能力。但是这些方法仍然不能很好地平衡性能和计算成本。

多层感知器(Multilayer perceptron, MLP)或全连接(fully connected, FC)是最早的神经网络类型，它由多个线性层和非线性激活叠加在一起组成(Rosenblatt 1957)。理论上，MLP是一个通用逼近器(Pinkus 1999)。但是，MLP计算量大，在数据不足时容易出现过拟合。此外，输入平坦化限制了输入分辨率。由于当时硬件和可用数据集的限制，MLP的发展并不顺利。2021年，MLP-Mixer (Tolstikhin et al 2021)重新焕发了MLP的活力。它主要由两个模块组成：Token-Mixing MLP和Channel-Mixing MLP，这两个模块在没有卷积和注意的情况下实现了竞争性能。MLP具有较小的归纳偏差，在大型数据集上，纯MLP架构可以更好地提取全局语义信息。但这也使得它在小数据集上表现不佳。为了获得更好的性能，引入了局部偏差(Hou et al . 2022;Tang et al . 2022;Yu et al . 2022;Lian et al . 2021)。但他们忽视了全局层面。

如何更有效地捕获和融合局部特征和远距离依赖关系是实现医学图像精确分割的关键。在本文中，我们重新思考了这个话题:除了将CNN和Transformer结合起来，是否还有其他方法可以同时拥有本地信息和远程依赖?答案是肯定的。本文将CNN与MLP相结合，提出了一种医学图像分割网络Rolling-Unet。其核心是灵活的Rolling-MLP (R-MLP)模块，可以在整个图像的单一方向上捕获线性距离依赖。通过连接两个垂直的R-MLP模块，形成正交Rolling-MLP (OR-MLP)模块，该模块可以捕获多个方向的远程依赖关系。我们采用U-Net的U型框架，包括编解码器结构、瓶颈层和跳跃连接，以保持良好的空间细节。在编解码器的第4层和瓶颈层，我们用Feature Incentive block和Long-Local (Lo2) block代替了原来的卷积block。Feature Incentive block对特征进行编码，控制特征输出的尺寸和形状。Lo2块由双正交Rolling-MLP (DOR-MLP包含两个互补的OR-MLP)模块和深度可分离卷积(DSC)模块组成，既捕获图像的局部上下文信息，又捕获图像的远程依赖关系。大量的实验表明，我们的方法优于现有的最佳方法。本文的主要贡献有：

1. 提出了一种新的远程依赖捕获方法，构建了R-MLP模块。
2. 在1的基础上，构建OR-MLP和DORMLP模块，可以获得更多方向上的远程依赖关系。
3. 在2的基础上，提出Lo2区块。它同时提取本地上下文信息和远程依赖关系，而不增加计算负担。Lo2块具有与3×3卷积相同的参数和计算级别。
4. 在3的基础上，构建了不同参数尺度的Rolling-Unet网络。在4个数据集上，Rolling-Unet的所有尺度都超过了现有方法，充分验证了我们方法的有效性。

Related Work

用于医学图像分割的CNN和Transformer

受U-Net的启发，UNet++（Zhou et al 2018）在模型中加入了一组密集的跳跃连接，以缓解特征融合的语义差距。随后的几项工作利用了注意力机制、图像金字塔和残差结构等技术（Oktay等人2018；Jha等人2020）（Jha等人2019），进一步提高了基于CNN的模型的性能。DconnNet (Yang and Farsiu 2023)是一种基于cnn的最先进的模型，它利用从共享潜在空间中提取的方向性特征来丰富整体数据表示。在医学图像领域，也出现了纯基于transformer的分割范式：MISSFormer (Huang et al . 2021)， DAE-Former (Azad et al . 2022)， Swin-Unet (Cao et al .2023)。Swin-Unet是第一个纯基于Transformer的U形架构，它采用Swin Transformer来增强特征表示。鉴于CNN和Transformer各自的缺点，人们提出了各种整合这两种范式的作品。MedT (Valanarasu等2021)设计了一个门控轴向注意模型，解决了医学图像中数据样本有限的问题。UCTransNet (Wang et al . 2022)引入了基于transformer的模块来替代U-Net中的跳跃连接。尽管这些工作都采用了融合全局和局部特征的策略来增强模型的能力，但仍然不能满足医学图像精确分割的需求。

图像任务的MLP范式

MLP-Mixer (Tolstikhin et al 2021)是用于视觉的深度MLP网络的先驱。由于其在小数据集上的性能较差，后来的工作试图在MLP中加入局部先验。Vision Permutator (ViP）(Hou et al.2022)用沿高度和宽度维度的线性投影对特征表示进行编码。稀疏MLP (Tang et al . 2022)遵循类似的策略，除了它直接沿着图像的高度和宽度进行映射。然而，这种设计缺乏灵活性，因为它的参数和计算开销与图像大小有关，这限制了输入图像的大小。S2MLP (Yu et al 2022)设计了一个空间移位模块，它将不同的令牌特征对齐到同一通道。AS-MLP (Lian et al . 2021)采用两个平行分支进行水平和垂直移动。然而，这些作品仅仅具有局部感受野，放弃了纯MLP模型捕捉全局特征的原始动机。在医学图像领域，据我们所知，基于MLP的分割模型很少。UNext (Valanarasu and Patel 2022)推出了一种轻量级模型，该模型采用轴向移位模块，但仍然只能捕获短距离线性接受场。PHNet (Lin等2023)是一种三维分割网络，它提出了一个多层排列感知器模块，该模块通过保留位置信息来增强原始MLP。

Method

架构概述

图1展示了拟议的Rolling-Unet的整体架构，它遵循U-Net设计。它由编码器-解码器、瓶颈层和跳跃连接组成。编解码器有下采样和上采样四个阶段，分别由最大池化和双线性插值实现。编码器-解码器的前三层分别包含两个标准的3×3卷积块。第四层和瓶颈层采用Feature Incentive block来处理特征通道的压缩和扩展，Lo2块来捕获图像的局部上下文和远程依赖关系。跳跃连接通过加法融合了相同尺度的特征。下面将详细描述每个模块。

R-MLP模块

给定空间分辨率为$H\times W$，通道数为$C$的特征矩阵$X{\in }^{H\times W\times C}$，其中$h_i(i\in [1,H])$表示高度指数，$w_j(j\in [1,W])$表示宽度指数，$c_k(k\in [1,C])$表示通道指数，我们沿同一方向对特征矩阵中各通道层的特征映射进行Rolling操作，如图2所示(以宽度方向为例)。Rolling操作包括两个步骤：移动和裁剪。首先，以通道索引为$c_k$的特征映射移动步长为$k$。然后，以通道索引为0的特征映射为参考，将其他特征映射的多余部分裁剪为缺失部分。最后，我们在每个空间位置索引$(h_i,w_j)$上执行一个权重共享的通道投影来编码长距离依赖。在图2中，对于所有通道，原始特征矩阵在固定的空间索引$(h_i,w_j)$下只有一个宽度$w_j$特征。在宽度方向上施加Rolling操作后，不同的通道具有不同的宽度特征。当$C\ge W$时，我们可以对整个图像的宽度特征进行编码，可以理解为全局的、单向的、线性的感受野。当C < W时，这个线性感受野是非全局的。同样，R-MLP也可以在高度方向上捕获长距离依赖。

众所周知，MLP对输入的位置信息非常敏感。R-MLP执行移动和裁剪特征映射的循环操作，使每个通道上的位置索引顺序不固定。这初步降低了R-MLP对位置的敏感性。其次，通过权值共享，所有信道投影共享一组参数，进一步降低了灵敏度

OR-MLP和DOR-MLP

R-MLP可以沿宽度或高度方向对远程依赖进行编码。我们怎样才能捕捉到其他方向上的长距离依赖？先在宽度方向上应用R-MLP，再在高度方向上应用R-MLP，相当于在两个正交的方向上对特征图进行同步移位操作，得到对角的感受野。如公式(1)所示，对于输入X，我们首先沿一个方向$ML{P}_R^1$应用R-MLP，然后沿垂直方向$ML{P}_R^2$连接另一个R-MLP。我们在两者之间使用GELU激活函数，然后与输入X添加残差连接。这就形成了正交Rolling-MLP (OR-MLP)模块，如图1所示。
$$ML{P}_{OR}(X)=(ML{P}_R^2(GELU(ML{P}_R^1(X))))+X\text{\hspace{1em}(1)}$$
R-MLP是一种具有巨大潜力的高度灵活的模块。移位步长$k$的符号决定了编码顺序。当单独使用R-MLP时，颠倒顺序不影响线性感受野提取。然而，在使用OR-MLP时，$k$的符号是至关重要的。对于宽度方向，给定正$k$值，表示从左向右移动(LR)，负$k$值表示从右向左移动(RL)。对于高度方向，给定正$k$值，表示从上到下移动(TB)，给定负$k$值表示从下到上移动(BT)。如图3所示，我们考虑两个互补的OR-MLP模块。

第一种方法首先沿LR方向应用R-MLP，然后依次沿TB方向应用。第二种方法首先沿BT方向应用R-MLP，然后依次沿LR方向应用。通过并行化这两个OR-MLP，我们捕获了四个方向上的远程依赖关系：宽度、高度、正对角线和负对角线！如公式(2)所示，对于输入$X$，我们首先应用OR-MLP $ML{P}_{OR}^1$，然后并行化另一个OR-MLP $ML{P}_{OR}^2$。我们沿着通道维度连接它们的输出并应用LayerNorm然后，我们使用通道混合(CM) (Tolstikhin等2021)融合特征并将通道减少回$c$。最后，我们添加与输入x的残余连接。这形成了双正交Rolling-MLP (DOR-MLP)模块，如图1所示。
$$\begin{array}{rl}ML{P}_{DOR}(X)=& CM(LN(Concat[ML{P}_{OR}^1(X),\\ & ML{P}_{OR}^2(X)]))+X\end{array}\text{\hspace{1em}(2)}$$

Lo2块和Feature Incentive block

DOR-MLP模块捕获二维空间中沿四个方向的全局线性远程依赖关系，但缺乏局部上下文信息。我们认为，更好地整合本地信息和全球依赖关系对于提高绩效至关重要。深度可分离卷积(DSC)是一种自然选择(Chollet 2017)。由于其参数少，计算成本低，与DOR-MLP兼容。众所周知，MLP-Mixer中的Channel-mix、ViT中的MLP以及本文中的R-MLP都等价于CNN中的标准1×1卷积，允许不同通道之间的特征交互。

R-MLP中的Rolling操作不涉及任何参数和FLOPs，因此R-MLP的参数为O(C²)， FLOPs为O(HWC²)。进一步推导出OR-MLP的参数和FLOPs分别为O(2C²)和O(2HWC²)， DOR-MLP的参数和FLOPs分别为O(6C²)和O(6HWC²)。如图1所示，我们将DOR-MLP与DSC并行化，然后沿着通道维度连接它们的输出，最后使用通道混合来融合特征并将通道恢复到$c$。这形成了Long-Local (Lo2)块，参见公式(3)。在DSC中，我们使用3×3卷积核。因此，我们可以得出Lo2块的参数为O(9C²)， FLOPs为O(9HWC^2)。这与标准的3×3卷积具有相同的级别。
$$Lo2(X)=CM(Concat[ML{P}_{DOR}(X),DSC(X)])\text{\hspace{1em}(3)}$$
我们在编码器的第四层和瓶颈层采用了特征激励块。它本质上是一个卷积块，主要用于对特征和通道号的变化进行编码。由于后续Lo2区块主要进行MLP，我们采用了GELU激活函数和LayerNorm，延续了之前的一系列MLP工作。在解码器的第四层，Feature Incentive块由卷积块、RELU激活函数和BatchNorm组成，后续网络按照CNN的一系列习惯进行卷积运算。

Experiments

数据集

我们在四个具有不同特征、数据大小和图像分辨率的数据集上评估了我们的方法:国际皮肤成像协作(ISIC 2018)、乳腺超声图像(BUSI)、腺体分割数据集(GlaS)和CHASEDB1。ISIC 2018数据集包含相机获取的皮肤图像和相应的皮肤病变分割图。我们只使用了ISIC 2018数据集的训练集，该数据集包含2594张图像。该数据集的难点在于分割目标的边界往往比较模糊，随着图像尺寸的增大，分割目标的边界更加模糊。因此，我们将图像调整为256×256和512×512两种分辨率，分别进行实验。BUSI数据集由正常、良性和恶性乳腺癌的超声图像和相应的分割图组成。它与ISIC 2018数据集有类似的问题，但它们具有不同的病变类型和成像方法。我们使用了647张乳腺良性和恶性肿瘤的超声图像，大小调整为256×256。GlaS数据集包含165张图像，我们将其大小调整为512×512。CHASEDB1数据集是一个具有28张999×960分辨率图像的血管分割数据集。为了保留细血管的细节，我们将图像大小调整为960×960。

实现细节

我们使用Pytorch在NVIDIA A6000 GPU上实现了Rolling-Unet。对于ISIC 2018, BUSI和GlaS数据集，批量大小设置为8，学习率为0.0001 (Valanarasu和Patel 2022)。对于CHASEDB1数据集，将批大小设置为4，学习率为0.001 (Tomar et al 2022)。我们使用Adam优化器来训练模型，并使用最小学习率为0.00001的余弦退火学习率调度器。损失函数是二元交叉熵(BCE)和骰子损失的组合。我们将每个数据集随机分成80%的训练子集和20%的验证子集。考虑到BUSI、GlaS和CHASEDB1数据集的数据量有限，我们重复了这一过程三次，并报告了结果的平均值和标准差。为了公平地评估网络的能力，所有的实验都没有使用任何预训练权值和后处理方法，只使用了两种简单的在线数据增强：随机旋转和翻转。我们一共训练了400个epoch。

与最先进方法的比较

我们将Rolling-Unet与其他最先进的方法进行了比较，包括基于cnn的方法：U-Net (Ronneberger, Fischer, and Brox 2015)， UNet++ (Zhou et al .2018)， Att- Unet (Oktay等人2018)，DconnNet (Yang和Farsiu 2023)；基于Transformer的方法：UCTransNet (Wang et al . 2022)， MedT (Valanarasu et al . 2021)；基于MLP的方法：UNeXt (Valanarasu and Patel 2022)。由于内存限制，MedT未能在GlaS, ISIC 2018(图像大小= 512)和CHASEDB1数据集上产生结果。类似地，UNet++在CHASEDB1数据集上也没有产生结果。为了充分证明Rolling-Unet的有效性，我们训练了不同大小的Rolling-Unet。图1中通道号C = 16 /32/64时，分别命名为Rolling-Unet (S) / Rolling-Unet (M) / Rolling-Unet (L)。我们采用Intersection over Union (IoU)、F1评分和95% Hausdorff Distance (HD95)作为评价指标。

BUSI和GlaS的评价结果见表1。ISIC 2018的结果如表2和表3所示。CHASEDB1的结果如表4所示。

我们的方法在所有数据集上都优于所有其他方法。特别是在BUSI和ISIC 2018上，Rolling-Unet取得了显著的优势。在这两个数据集中，许多目标边界模糊，难以从背景中区分出来。Rolling-Unet更有效地提取了远程依赖项，提高了分割性能。在ISIC 2018上改变图像尺寸的实验进一步验证了这一结论。当图像尺寸增大时，只有RollingUnet和UNeXt保持了相似的性能，而其他方法都有不同程度的下降。对于ISIC 2018中Rolling-Unet (X)的指标低于Rolling-Unet (S)的现象，我们有两个假设。一是训练的波动，需要取多个结果的平均值来减少影响。二是该数据集的语义信息相对简单，网络参数较多容易出现过拟合，从而降低性能。最近的轻量化模型(Valanarasu and Patel 2022; Ruan et al 2023; Cheng et al 2023)也从侧面反映了这一点。在后续的工作中，我们将通过更多的实验来解释这一现象。

在GlaS和CHASEDB1上，没有一种方法取得明显的优势，但Rolling-Unet仍然是最好的，标准差较小。GlaS中的图像密度大，微小的细胞和组织；分割目标和背景通常具有相似的纹理、颜色和形状。在CHASEDB1数据集中，较厚的血管对所有方法来说都不是困难的，分割的困难就在这些较薄的血管上，如图5所示。这些问题需要更强大的方法来解决。

各模型的参数数量见表1。我们将参数量小于2M的模型定义为一级模型，将参数量大于20M的模型定义为二级模型(只有Rolling-Unet (L)在2-20M之间)。在四个数据集上，我们的方法在一级和二级模型上都是最好的，证明了方法的有效性。

在图4中，我们可视化了分割结果与Ground Truth之间的差异图，以突出显示差异。紫色表示分割过度，黄色表示分割不足。由于篇幅限制，我们省略了MedT的结果。在BUSI和ISIC 2018的图像中，我们可以看到分割目标缺乏明确的边界。在分割结果中，除Rolling-Unet之外的其他方法产生了大量的分割不足或过度分割区域。这说明Rolling-Unet在目标轮廓的提取上是很好的。GlaS中目标边界复杂，只有Rolling-Unet才能获得接近Ground Truth的分割结果。CHASEDB1数据集的可视化结果如图5所示。几乎所有的方法都能正确分割粗血管，而细微的区别在于蓝框内的细血管。Rolling-Unet考虑了图像的长距离依赖特征，提高了细血管的分割效果。

消融研究

为了研究各种因素对模型性能的影响，我们在ISIC 2018数据集(图像大小= 512)上进行了消融研究。具体描述如下。

Lo2 Block由DOR-MLP和DSC模块并行组成。前者负责捕获远距离依赖关系，后者负责提取局部上下文信息。为了确保DOR-MLP和DSC的组合是最优的，并探讨它们各自的贡献，实验结果如表5所示。无论是否存在DSC模块，R-MLP、OR-MLP和DOR-MLP的性能都会逐渐提高。这证明了所提出的模块在捕获远程依赖关系方面的有效性，并验证了从多个方向提取远程依赖关系的思想。当与DSC模块结合使用时，可以进一步提高性能。因此，必须融合远程依赖项和本地上下文信息。

为了排除参数和FLOPs增加带来的性能提升，我们将Rolling-Unet中的R-MLP替换为常规MLP。这使得模型在保持参数和FLOPs一致的同时失去了捕获远距离依赖关系的能力。如表6所示，性能显著下降。这个结果是意料之中的，因为没有能力捕获远程依赖关系的Rolling-Unet与原始的U-Net具有相似的网络结构。

进一步，我们探索了DOR-MLP与DSC的结合。系列1意味着先执行DOR-MLP，然后执行DSC。系列2意味着先执行DSC，然后执行DOR-MLP。并行是指DSC和DOR-MLP并行连接，两个支路并行执行，最后通过通道混频实现特性集成。结果如表7所示。系列1和系列2差别不大，最好是并联。这证明:提取局部特征和远程依赖关系的顺序并不重要，最好在同时提取后进行融合。

Conclusion

本文提出的Rolling-Unet模型可以在不增加计算成本的情况下捕获远程依赖关系，并且优于现有的方法。值得注意的是，来自多个方向的远程依赖不是全局感受野，严格意义上讲，它们仍然是MLP的妥协。然而，R-MLP是一个非常灵活的模块。结合起来，还可以捕捉到大尺度的区域特征，甚至是全局特征。在今后的工作中，我们将在这方面进行探索。我们还将研究其在三维医学图像分割以及其他图像任务中的潜力。