2023 港科大医学图像分割新作 | PHNet: 当MLP与CNN巧妙 ...

余生不扰

本文首发至微信公众号：CVHub，不得以任何形式转载或售卖，仅供学习，违者必究！

Title: A Permutable Hybrid Network for Volumetric Medical Image Segmentation
PDF: https://arxiv.org/pdf/2303.13111
Code: coming soon...

导读

近年来，VIT（Vision Transformer）在3D医学图像分割中取得了实质性进展，与此同时，多层感知器MLP（Multi-Layer Perceptron）网络由于其与VIT具有相当性的结果而重新受到研究人员的青睐。本文提出了一种用于3D医学图像分割的混合网络PHNet，它利用了卷积神经网络CNN和MLP的优点，通过利用2D和3D CNN提取局部信息，解决了3D volume 数据内在的各向同性问题；此外，本文还提出了一种高效的多层置换感知器模块MLPP，它通过保留位置信息来增强原始的MLP，并获得长程距离依赖。通过在两个公共数据集COVID-19-20和Synapse上的实验结果表明，本文所提出的PHNet方法优于最先进的SOTA方法。
引言

卷积神经网络的发展极大推动了计算机辅助诊断CAD（Computer-aided diagnosis）性能，尤其是在医学图像分割方向。其中最受欢迎的医学图像分割结构当属U-Net，它以编码-解码结构和经典的跳跃连接不仅捕获了图像丰富的语义信息，还极大保持了医学图像的空间细节信息。继U-Net之后，多个优秀的变种相继出现，比如ResUNet，Y-Net和N-Net等。



图1. Illustration of different architectures

随后，基于注意力机制的Transformer在自然语言处理领域大火。以计算机视觉领域为代表的Transformer网络VIT架构频频霸榜，在各类计算机视觉任务上都展现出了比CNN架构更先进的性能。鉴于Transformer在自然图像识别任务方面的显着进展，许多研究人员已经研究了各种神经网络在医学图像分割方面的有效性。例如，TransUNet提出在U-Net体系结构的瓶颈处使用Transformer进行全局信息通信。同样，UNETR 和CoTr 设计了一个层次化的Transformer和CNN架构进行融合也取得了不错的精度。
Transformer大法固然好，但其内在的自注意力机制产生的巨大计算量一直是一个诟病（这一点极大限制了该架构在工业界的推广应用），尤其是在3D医学图像数据中。
因此，基于多层感知器MLP的方法在研究界重新引起了兴趣，因为它已经展示出了与CNN和Transformer架构相当的性能，而不需要沉重的自我注意机制。例如，MLP-Mixer 通过一系列MLP可以实现特征间的信息交流，以捕获输入数据中的长远距离依赖。然而，MLP在3D医学图像分割方面的有效性仍然缺乏研究。
为此，本文将CNN和MLP相结合，提出了一种新的混合网络PHNet，以实现准确的3D医学图像分割。PHNet采用编码器-解码器结构，其中编码器利用2.5D CNN结构，可以利用医学图像固有的等向性，并通过捕获不同方向上体积医学图像的变化信息密度来避免浅层损失信息。论文进一步提出了MLPP，一种可以在计算效率高的情况下保持位置信息并集成全局相互依赖性的多层排列感知器模块。为了提高计算效率，引入了令牌token组操作，可以高效地在令牌级别上聚合特征映射，从而减少所需的计算量。本文首次提出将CNN和MLP进行结合，并应用于3D医学图像分割。通过在两个公开数据集上的评估实验（COVID-19肺CT病灶分割挑战赛2020（COVID-19-20）和Synapse多器官分割），证明了PHNet表现均优于当前最先进的SOTA方法，并超越了MICCAI Covid-19-20挑战赛的冠军！
方法




图2. Overview of PHNet and Detailed architecture of MLPP

如上图2所示PHNet总体思路非常简单：采用编码器-解码器架构，编码器由两个主要组件组成：2.5D卷积模块和多层置换感知器（MLPP）模块。2.5D卷积阶段提取局部特征，输出的特征图传递给MLPP模块以学习全局特征。解码器对分层特征进行处理以进行最终预测。
2.5D Convolution

本文通过引入卷积层来提取局部特征，这是基于先前关于医学图像分析偏差的研究[1]和3D体积医学图像的各向异性本质而得出的结论。3D体积图像如CT和MRI扫描常常由于其厚切片扫描而受到各向异性问题的影响，导致高面内IP分辨率和低面外TP分辨率，这种差异在COVID-19-20中特别明显，其IP分辨率平均为0.74mm，而TP分辨率仅为5mm。为了解决这一问题，本文使用2D卷积块来捕获IP信息，直到特征在三个轴面（横断、冠状和矢状）上形成近似均匀的分辨率，然后使用3D卷积块来处理体素信息。每个编码器层由两个残差卷积块组成，每个块包括两个顺序的Conv-IN-ReLU操作。
Multi-Layer Permute Perceptron (MLPP)

尽管卷积神经网络（CNN）通过深层堆叠卷积层能够建模长距离依赖关系，但研究表明：基于多层感知器MLP的网络在学习全局上下文方面更有效。 因此本文设计了MLPP模块（如上图2b所示）用以提取深层的全局信息。MLPP模块按顺序分解平面内IP特征和垂直方向TP特征的训练。作者分别将这两个块称为IP-MLP和TP-MLP。为了实现跨轴令牌之间的通信，作者还在IP-MLP中提出了一个辅助注意力分支，称为AA-MLP。具体细节如下：
IP-MLP

常见的基于多层感知机MLP的方法将输入特征图直接展平成一维向量，这会导致卷积特征中的空间信息丢失。为了解决这个问题，作者提出了一种称为轴向分解的方法，在水平轴、垂直轴和通道轴上分别对输入特征进行单独处理，从而在编码某个轴向信息时保留其他轴向的精确位置信息。
为了平衡远距离交互和计算成本，并减轻图像分辨率敏感性问题，作者还提出了一种称为token分割的操作，将特征向量分割成多个token，这些token可以被后续的全连接层高效处理。作者以水平轴为例，将输入特征沿水平方向分成不重叠的段，然后将每个段再分成多个不重叠的通道组，其中每个组有g = C/L个通道。接下来，作者将每个分段展平并使用一个全连接层将R_{Lg}映射到R_{Lg}，从而得到一个新的特征向量Y_i^{k}。最后，将所有分段重新排列，得到Y_H、Y_W和Y_C三个特征向量，它们分别代表沿垂直，水平和通道方向的编码信息。这些特征向量经过逐元素求和后，被送入一个新的全连接层，得到最终的输出Y_{IP}，其中W \in R^{C×C}是全连接层的权重矩阵。
AA-MLP

本文提出的IP-MLP模块有两个限制，可能会损害分割性能。首先，轴向分解截断了不在相同水平或垂直位置的令牌之间的直接交互。其次，与vanilla MLP相比，令牌分割操作的局部接收字段较小。为了解决这些限制，作者设计了一个辅助分支，以实现轴内令牌通信，并通过轻量级但有效的类MLP架构充当注意力函数。具体来说，给定输入的特征图 X\in \mathbb{R}^{H\times W\times C}，将X划分为非重叠的窗口，窗口大小设置为L，因此得到X_i \in \mathbb{R}^{L\times L}，其中i \in {1,\dots,HWC/L^2}。然后应用一个FC矩阵W \in \mathbb{R}^{L^2\times L^2}来转换每个窗口，并得到Y_i \in \mathbb{R}^{L\times L}。通过将所有窗口排列回原始尺寸，最终得到注意力图YA\in \mathbb{R}^{H\times W\times C}。最后，通过Y_{IP}和Y_A的残差注意力得到IP-MLP的特征映射F_{IP}


其中 \odot 表示逐元素相乘。
TP-MLP

通过IP-MLP获取了平面内的信息后，接下来使用TP-MLP来捕捉长期的沿平面垂直方向的特征。同样地，对于输入的特征映射F_{IP} \in \mathbb{R}^{H \times W \times D \times C}，我们首先沿深度维度将X=F_{IP}划分为长度为L的不重叠的片段。这样我们获得了X_i \in \mathbb{R}^{L \times C}，其中i \in {1,...,\frac{HWD}{L}}。接下来，我们将X沿通道维度分为几个不重叠的组，每组有g=C/L个通道，然后得到X_{ki} \in \mathbb{R}^{L \times g}，其中k \in {1,...,C/g}。然后，我们将每个段展平并通过全连接层将\mathbb{R}^{L \times g}映射到\mathbb{R}^{L \times g}，得到Y_{i}^k。最后，我们将所有Y_{i}^k \in \mathbb{R}^{L \times g}的段重新排列到原始维度，并输出\mathrm{F_{TP}} \in \mathbb{R}^{H \times W \times D \times C}。
Decoder

本文的解码器采用纯CNN架构，使用转置卷积逐渐上采样特征映射以匹配输入图像分辨率。在上采样过程之后，使用残差卷积块来细化特征映射。为了进一步提高分割准确性，作者在编码器和解码器之间包括跳跃连接，允许保留low-level的细节信息。
实验结果




Synapse数据集上不同方法的定性可视化效果，为了更好地展示细节，显著改进的区域被放大



COVID-19-20数据集上不同SOTA方法对比



Synapse数据集上不同SOTA方法对比



消融实验

（a）作者对不同结构的组合进行了性能比较，包括在浅层和深层使用 Conv、Attention 和 MLP 的不同组合。其中Conv 和 MLP 使用了与 PHNet 相同的模块，Attention 使用了 Swin Transformer block 并将窗口大小设置为分段长度以进行公平比较。结果显示，在浅层使用 Conv 和在深层使用 MLP 的组合可以获得最佳性能，这与作者的论点相符，即 Conv 擅长提取局部特征，而 MLP 更有效地模拟长程依赖。
（b）作者还比较了不同 MLP 设计的性能，包括 MLP-Mixer（M）、ShiftMLP（S） 和 WaveMLP（W）。作者仅将 PHNet（P）中的 MLPP 模块替换为其他模块以进行公平比较。结果表明，所提出的结构分别在 Dice 指标上获得了 2.15％、1.85％ 和 0.45％ 的改进，证明了本文设计的良好性能。
（c）作者还研究了不同分段长度 L 对性能的影响。具体地，将分段长度设置为宽度（W）的不同比例，即 1、1/2、1/3 和 1/4。这有利于不同大小的感兴趣区域（ROI）。结果显示，当 L = 1/2 W 时，性能最佳。
（d）最后，作者研究了 MLP 层数的影响。结果表明，在 MLP 层数为 2 时，性能最佳。
结论

本文提出了一种新的可置换混合网络PHNet，用于3D体积医学图像分割任务。PHNet集成了2D CNN、3D CNN和MLP，以实现局部和全局特征的提取和集成。此外，本文还引入了一个可置换的MLP块，以解决空间信息丢失问题并最小化计算负担。通过在两个公共数据集上进行的广泛实验证明，所提出的方法优于现有的SOTA方法。结合不同任务的特点，考虑精度与计算负载问题，系统地研究CNN、Transformer和MLP之间的效率、有效性和交互作用，会是一个不错的idea。
<hr/>CVHub是一家专注于计算机视觉领域的高质量知识分享平台，全站技术文章原创率常年高达99%，每日为您呈献全方位、多领域、有深度的前沿AI论文解决及配套的行业级应用解决方案，提供科研 | 技术 | 就业一站式服务，涵盖有监督/半监督/无监督/自监督的各类2D/3D的检测/分类/分割/跟踪/姿态/超分/重建等全栈领域以及最新的AIGC等生成式模型。欢迎关注微信公众号CVHub或添加小编好友:cv_huber，备注“知乎”，参与实时的学术&技术互动交流，领取CV学习大礼包，及时订阅最新的国内外大厂校招&社招资讯！

夏曰阳光

这里面提到的MLP部分的方法，在提取全局特征方面比Transformer效果好，敢不敢拿这个用来替换[思考]

晓珂侃世界

可以一试