Rolling-Unet：重振MLP对于医学图像分割高效提取长距离依赖的能力

pub: 2024 AAAI
[ paper] [ code]

Introduction

1 由于卷积运算固有的局域性，纯粹的CNN架构很难学习到清晰的全局（global）和远程语义（remote semantic）信息。
2 CNN和Transformer结合不能很好地平衡性能和计算成本。
3 除了将CNN和Transformer结合起来，是否还有其他方法可以同时拥有局部信息（local information）和远程依赖（long-distance dependencies）?

● 本文的主要贡献：

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

Related Work

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

Swin - unet是第一个纯基于Transformer的Ushaped架构，UCTransNet引入了基于transformer的模块来替代U-Net中的跳过连接。尽管这些工作都采用了融合全局和局部特征的策略来增强模型的能力，但仍然不能满足医学图像精确分割的需求。

● 图像任务的MLP范式

MLP- mixer 是用于视觉的深度MLP网络的先驱，AS-MLP 采用两个平行分支进行水平和垂直移动。这些作品仅仅具有局部接受场，放弃了纯MLP模型捕捉全局特征的原始动机。在医学图像领域，基于MLP的分割模型还较少。

Method

●整体架构

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

● R-MLP Module

给定空间分辨率为H×W，通道号为C的特征矩阵X∈H×W×C，其中hi(i∈[1,H])表示高度指数，wj (j∈[1,W])表示宽度指数，ck(k∈[1,C])表示通道指数，我们沿同一方向对特征矩阵中各通道层的特征映射进行Rolling（滚动）操作，如图2所示(以宽度方向为例)。Rolling操作包括两个步骤:移动和裁剪。首先，以通道索引为Ck的特征映射，移动步长为k。然后，以通道索引为C0的特征映射为参考，将其他特征映射的多余部分裁剪为缺失部分。最后，我们在每个空间位置索引(hi, wj)上执行一个权重共享的通道投影来编码长距离依赖。在图2中，对于所有通道，原始特征矩阵在固定的空间索引(hi, wj)下只有一个宽度wj特征。在宽度方向上施加Rolling操作后，不同的通道具有不同的宽度特征。当C≥W时，我们可以对整个图像的宽度特征进行编码，可以理解为全局的、单向的、线性的接受场。当C < W时，这个线性感受野是非全局的。同样，R-MLP也可以在高度方向上捕获长距离依赖。众所周知，MLP对输入的位置信息非常敏感。R-MLP执行移动和裁剪特征映射的循环操作，使每个通道上的位置索引顺序不固定。这初步降低了R-MLP对位置的敏感性。其次，通过权值共享，所有信道投影共享一组参数，进一步降低了灵敏度。

● OR-MLP and DOR-MLP

R-MLP可以沿宽度或高度方向对远程依赖进行编码。我们怎样才能捕捉到其他方向上的长距离依赖?先在宽度方向上应用R-MLP，再在高度方向上应用R-MLP，相当于在两个正交的方向上对特征图进行同步移位操作，得到对角的感受场。如式(1)所示

对于输入X，我们首先沿一个方向MLP1 R应用R-MLP，然后沿垂直方向MLP2 R连接另一个R-MLP。我们在两者之间使用GELU激活函数，然后与输入X添加残差连接。这就形成了正交滚动mlp (OR-MLP)模块，如图1所示。
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 1 OR，然后并行化另一个OR-MLP2 OR。我们沿着通道维度连接它们的输出并应用LayerNorm。然后，我们使用通道混合(CM) (Tolstikhin等2021)融合特征并将通道减少回到C。最后，我们添加与输入x的残余连接。这形成了双正交滚动mlp (DOR-MLP)模块，如图1所示。

● Lo2 Block and Feature Incentive Block

DOR-MLP模块捕获二维空间中沿四个方向的全局线性远程依赖关系，但缺乏局部上下文信息（local context information）。我们认为，更好地整合局部信息和全局依赖关系对于提高下游任务的结果至关重要。深度可分离卷积(DSC)是一种自然选择(Chollet 2017)。由于其参数少，计算成本低，与DOR-MLP兼容。众所周知，MLP-Mixer中的Channel-mixing、ViT中的MLP以及本文中的R-MLP都等价于CNN中的标准1×1卷积，允许不同通道之间的特征交互。
R-MLP中的Rolling操作不涉及任何parameters和FLOPs，因此R-MLP的参数为0 (C2)， FLOPs为0 (C2)。进一步推导出OR-MLP的参数和FLOPs分别为O(2C 2)和O(2HW C2)， DOR-MLP的参数和FLOPs分别为O(62c)和O(6HW C2)。如图1所示，我们将DOR-MLP与DSC并行化，然后沿着通道维度连接它们的输出，最后使用通道混合来融合特征并将通道恢复到c。这形成了Long-Local (Lo2)块，参见公式(3)。在DSC中，我们使用3×3卷积核。因此，我们可以得出Lo2块的参数为O(9c2)， FLOPs为O(9HW C2)。这与标准的3×3卷积具有相同的级别。

我们在编码器的第四层和瓶颈层采用了特征激励块。它本质上是一个卷积块，主要用于对特征和通道的变化进行编码。由于后续Lo2模块进行MLP，我们采用了GELU激活函数和LayerNorm，延续了之前的一系列MLP工作。在解码器的第四层，Feature Incentive块由卷积块、RELU激活函数和BatchNorm组成，后续网络按照CNN的一系列习惯进行卷积运算。

Experiments

● Datasets

在四个具有不同特征、数据大小和图像分辨率的数据集上评估了我们的方法:国际皮肤成像协作(ISIC 2018)、乳腺超声图像(BUSI)、腺体分割数据集(GlaS)和CHASEDB1。

● Implementation Details

我们使用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次。

● Comparison with State-of-the-Art Method

在图4中，我们可视化了分割结果与Ground Truth之间的差异图，以突出显示差异。紫色表示分割过度，黄色表示分割不足。由于篇幅限制，我们省略了MedT的结果。在BUSI和ISIC 2018的图像中，我们可以看到分割目标缺乏明确的边界。在分割结果中，除Rolling-Unet之外的其他方法产生了大量的分割不足或过度分割区域。这说明RollingUnet在目标轮廓的提取上是很好的。GlaS中目标边界复杂，只有rollling - unet才能获得接近Ground Truth的分割结果。

CHASEDB1数据集的可视化结果如图5所示。几乎所有的方法都能正确分割粗血管，而细微的区别在于蓝框内的细血管。rollling - unet考虑了图像的长距离依赖特征，提高了细血管的分割效果。

● Ablation Studies

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

Conclusion:

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

有任何问题欢迎联系：438285719@qq.com 共同学习交流