A4-Unet: Deformable Multi-Scale Attention Network for Brain Tumor Segmentation【医学图像分割】

一、论文信息

1.1、中文名称

Title：A4-Unet：用于脑肿瘤分割的可变形多尺度注意力网络

1.2、论文关键词

脑肿瘤分割、卷积神经网络，通道注意力、空间注意力、Swin Transformer

1.3、核心概述

脑肿瘤分割模型面临MRI（磁共振成像）的复杂性和可变性挑战，包括不规则的形状和不明确的边界，导致噪声，错误分类和分割不完整，从而限制了模型的准确性。为了解决这个问题，本文以UNet模型为参考在编码器中掺入了可变形的大核注意力（DLKA），在瓶颈中采用了具有跨通道注意的Swin空间金字塔池（SSPP），并在解码器中引入了带有离散余弦变换（DCT）正交性的组合注意模块（CAM），在跳跃连接中添加了注意门（AG），设计出了一种新型的脑肿瘤分割模型A4-Unet。

二、摘要

2.1、背景

近年来，脑肿瘤分割模型有助于医学诊断。

2.2、挑战

然后这些脑肿瘤分割模型面临MRI的复杂性和可变性挑战，包括不规则的形状和不明确的边界，导致噪声，错误分类和分割不完整，从而限制了准确性。

2.3、提出新方法

为了解决这些问题，本文遵守出色的卷积神经网络（CNNS）设计范式，并提出了一个名为A4-Unet的新颖网络。在A4-Unet中，编码器中掺入了可变形的大核注意（DLKA），从而改善了多尺度肿瘤的捕获。在瓶颈中采用了具有跨通道注意的Swin空间金字塔池（SSPP），以进一步研究图像和信道关系中的长距离依赖性。为了提高准确性，引入了带有离散余弦变换（DCT）正交性的组合注意模块（CAM），用于通道加权和卷积元件乘积，以用于解码器中的空间加权。在跳过连接中添加了注意门（AG），以突出前景，同时抑制无关紧要的背景信息。

2.4、贡献

对拟议的A4-Unet网络进行了三个权威性MRI脑肿瘤基准和专有数据集的评估，并在Brats 2020数据集中获得了94.4％的骰子分数，从而建立了多个新的最新基准测试。该代码可详见：https://github.com/WendyWAAAAANG/A4-Unet

三、引言

3.1、引出背景

图1：BraTS 2020数据集的样本可视化

由脑细胞异常生长引起的脑肿瘤对人类健康构成重大威胁，使早期诊断和治疗至关重要。 MRI作为一种非侵入性成像技术，提供了软组织病变的清晰可视化，广泛用于诊断和治疗脑肿瘤，如图1所示。当前的医疗图像分割方法主要依赖于U形CNN。

3.2、引出挑战

尽管进行了广泛的研究，但由于MRI图像，不明确的边界以及不规则的肿瘤形状和纹理，脑肿瘤分割仍然具有挑战性。传统的CNN模型难以适应这些不规则性，未能汇总语义信息并弥补空间信息损失。这导致噪声，错误分类，分割不完整，图像特征提取有限和准确性提高有限。

3.3、目前研究概况

从以前成功的语义细分研究中得出，Guo等人确定了表I所示的三个关键特征，即应具有良好的CNN分割模型。本文将这些关键点纳入了脑肿瘤图像分割特征，并将其汇总为如下：

1）强有力的编码器的使用：大脑图像通常包含复杂的结构，例如脑组织，血管和心室，而肿瘤通常表现出各种形状和大小。一个稳健的编码器对于捕捉和表示这些复杂的高级语义特征，以及准确分割这些结构是必需的。

2）融合多尺度信息：大脑中各种组织结构中的肿瘤可能表现出显着的大小，形状和分布差异。通过融合多尺度信息，模型可以更好地捕获图像中的细节和全局上下文，从而增强了细分模型对各种结构的理解。

3）注意机制的整合：MRI图像具有多个渠道，每个渠道都提供不同的信息。通道注意机制可帮助模型确定特定任务的关键通道。空间注意机制有助于该模型关注特定位置，以捕获局部结构细节，从而提高了分割精度。

表Ⅰ：语义分割的三个关键特征

3.4、贡献总结

受Guo等人的启发，本文重新审视了CNN设计原理以开发A4-Unet，这是一种整合了四个高级组件的脑肿瘤分割结构，即可探索的大核心（DLKA），SWIN增强型体现的空间空间金字塔（SSPP），联合注意模块，联合注意模块（CAM）和注意门（AG） - 每个增强性能。本文的主要创新是：

通过合并大型变量卷积，编码器可以更好地捕获具有低复杂性的多尺度信息。
可以通过使用SWIN空间金字塔池（SSPP）（SSPP）和瓶颈层中的卷积通道注意来提取长距离依赖性内图像内图像内图像和关系间的关系。
在解码器中，本文利用离散余弦变换（DCT）的正交性来计算通道注意力重量，然后进行跳过连接以补充细节细节。此外，本文利用简单的卷积元素乘法来引起空间注意，从而提高了模型的概括性能。

四、相关工作

4.1、骨干网络

4.1.1、基于CNN的架构

基于CNN的方法对像素块进行分类以捕获本地和全局特征。 DenseNet堆叠深层以维护多尺度特征，而基于UNet的扩展则来自完全卷积神经网络（FCN）的启发，解决了各种分割挑战。 SegNeXt通过多尺度卷积注意（MSCA）模块的增强了卷积结构。然而尽管有效地保留了低级信息，但CNN模型很难捕获高级信息，限制了它们的性能。

4.1.2、基于Transformer的网络

基于Transformer的网络使用注意力机制为图像部分分配重要性权重，这类网络在视觉任务上已经显示出了令人印象深刻的结果，最初是由视觉Transformer（ViT）取得的成功。SegFormer和Swin Transformer这样的变体使用分层Transformer编码器来提取多尺度特征，并使用简单的解码器进行分割。然而它们在检测纹理和边缘等高分辨率细节方面存在困难，限制了它们在密集视觉任务中的有效性。

4.1.3、融合CNN和Transformer

结合CNN和Transformer的混合架构利用两者的优势来克服局限性。 TransAttUnet集成了Transformer和U-Net，通过注意块和多尺度跳过连接捕获全局上下文信息，实现特征图的语义一致性。BoTNet使用CNN将输入图像处理成标记化的特征映射，然后使用变压器捕获远程依赖关系。在我们的研究中，A4-Unet结合了一个鲁棒的卷积编码器和变压器引导模块，以实现令人信服的分割性能。

4.2、注意力机制

注意力机制根据输入特征动态调整权重。如挤压激励网络（SE-Net）通道注意力，为每个通道分配不同的权重，而频率通道注意网络（FcaNet）使用离散余弦变换来关注低频通道信息。

空间注意通过创建权重掩模来增强重要区域，如卷积块注意模块（CBAM）所示，该模块将池化和连接相结合以获得统一的特征描述。本文的模型使用CBAM的轻量化设计将通道和空间注意力集成在一起，以强调重要区域并抑制无关信息，捕捉跨通道关系和空间细节以进行精确检测。

4.3、调整感受野

4.3.1、空洞卷积

空洞卷积最初出现在二进制小波变换技术中，这是一种被广泛认可的信号处理技术。深度网络降低了最终特征图的分辨率，这是由于池化层、步幅操作等的累积影响。Yu和Koltun等人提出了一种创新的方法来克服这一缺陷，同时寻求更广泛的信息谱。

4.3.2、可变形卷积

卷积神经网络（CNNs）的固定感受野限制了它们处理大规模几何变换的能力，这使得高级语义提取变得具有挑战性。受到多尺度可变形部件模型和空间Transformer模块的启发，可变形卷积通过引入2D偏移量到采样位置来解决这一问题，允许灵活的网格变形。本文采用可变形卷积来增强感受野的灵活性，以实现更好的目标分割。

4.4、多尺度上下文信息

4.4.1、空洞空间金字塔池化

在语义分割中聚合多尺度上下文信息对于准确的像素级分类至关重要。空洞卷积扩大了感受野，而不会改变输出尺寸。在SPP层的基础上，ASPP融合多个尺度从而捕获图像上下文。这激发了本文的模块从病变图像中提取丰富、全面的信息。

4.4.2、多尺度Transformer

尽管CNN有效地使用了多尺度特征表示，但该潜力尚未在Vit中进行全面探索。 Crossvit等人引入了带有交叉注意的双支流Transformer，MVIT等人将多尺度特征金字塔嵌入到Transformer中。受这些模型的启发，本文提出了一个基于层次Swin Transformer体系结构的双分支编码器。

五、具体方法

5.1、整体架构

本文的A4-Une是由三个主要组件组成的编码器-解码器体系结构，如图2所示。DLKA用于增强特征提取，SSPP用于多尺度交互，CAM用于注意力机制。编码器使用DLKA，SSPP在瓶颈处处理多尺度特征，解码器通过门控和混合注意力机制在四个上采样阶段聚合特征，从而优化脑肿瘤分割。

图2：A4-Unet 的整体架构

5.2、编码器

图3：DLKA结构图

为了构建一个健壮的编码器，我们将图3中的可变形大内核注意力（DLKA）块集成到下采样过程中。 DLKA包括一个可变形卷积模块（DConv）和一个大卷积核（LK）。

DConv对于增强边缘、纹理和形状等低级特征细节非常理想，特别是对于具有不规则尺寸和各种纹理的医疗目标。 DConv包括一个二维卷积，一个可变形卷积与可调采样网格使用偏移，非线性的激活函数，和偏移场计算。由Azad提出，标准卷积层产生偏移量，指导可变形卷积层的采样位置。 DConv模块公式如下：

其中， $Conv_{DC}$ 和 $Conv_{DW}$ 分别为可变形卷积和深度扩张卷积，F为输入特征。

另一方面，尽管CNN在捕获局部特征和低级信息方面做得很好，但它们是以忽略全局背景为代价的。 Guo等人提出的LK可以通过扩大接受野来克服这一限制。它提供了一个与自我注意机制相似的接受场，但参数更少。 LK的结构包含一个深度卷积、一个扩展卷积和一个1 × 1卷积。深度卷积（depth-wise convolution， KDW）和扩展卷积（expanded convolution， KDC）的核大小可以计算如下：

其中d为膨胀率，K为核大小。

总之，DLKA集成到编码器中，在下采样过程中提供远程依赖关系，并在上采样过程中通过跳过连接与特征映射连接，从而补偿低级特征细节。

5.3、多尺度融合

图4:SSPP和跨上下文注意模块的实现

解决医学图像中不规则尺寸和形状的挑战需要引入多尺度交互和增强空间表征。先前的方法使用了多尺度交互和更深层的网络，但多尺度信息仍然是碎片化的。

本文通过修改瓶颈层来解决这个问题，包括Swin空间金字塔池（SSPP）和如图4所示的跨上下文注意模块。这种方法集成了具有不同窗口大小的Swin Transformer块，提供了丰富的上下文信息。

Swin空间金字塔池:在DeepLab V3+中，Chen等人引入了Atrous空间金字塔池（ASPP）模块，该模块动态选择不同大小的卷积块来处理不同的目标尺度。这种方法可以防止大型目标被分散，并在不改变网络结构的情况下保持远距离依赖关系。

受Azad等人的SSPP的启发，本文用Swin transformer替换了四个扩展卷积，以更好地捕获远程依赖关系。提取的特征被合并并提供给跨上下文注意模块,这增强了模型在不同范围内捕获上下文依赖关系的能力。

交叉上下文注意力：ASPP通过深度可分离卷积连接特征映射，它不捕获通道依赖关系。为了解决这个问题，Azad在SSPP特征融合后引入了跨上下文注意。假设每个SSPP层都有token(P)，嵌入维数(C)为 $(z_m^{P \times C})$ ，表示不同尺度的对象。本文通过连接这些特征创建了一个多尺度表示 $z_{\text{all}}^{P \times MC} = [z_1 \| z_2 ... \| z_M]$ 。然后，尺度关注模块强调每个特征映射的贡献，使用全局表示和MLP层生成缩放系数 $(w_{\text{scale}})$ ，增强上下文依赖性：

其中 $W_{1}$ 和 $W_{2}$ 为可学习MLP参数，δ为ReLU函数，σ为Sigmoid函数，GAP为全局平均池化。

在第二个注意层面，跨上下文注意学习缩放参数，通过计算它们的权重图来增强信息标记，使用相同的策略：

5.4、卷积注意力模块

图5：组合注意力模块

本文通过集成新颖的卷积注意模块和有效抑制不必要信息的频率特征来构建解码器。此外，本文引入了带有注意门控融合的跳跃连接，有助于抑制不相关区域和突出特征。

如图5所示，本文的解码器包括一个用于特征上采样的vanilla块，一个用于级联特征融合的注意门（AG）和一个用于特征映射增强的组合注意模块（CAM）。本文为编码器的四个金字塔层使用四个CAM块，为跳过连接使用四个AGs。通过将前一层的上采样特征与使用AG的跳过连接特征相结合来整合多尺度特征。然后，CAM模块通过频率通道和空间注意（SA）增强像素分组并抑制背景信息。最后，Dconv将融合的特征传播到上层。

1）组合注意力模块：

通道注意力：为了提高CAM中的通道注意精度，本文用Salman等人的正交通道注意（OCA）取代了基于卷积的通道注意。 OrthoNet的信道注意通过使用离散余弦变换（DCT）来保留低频信息，解决了全局平均池化（GAP）的局限性。如图6所示，OCA的结构包括在适当的尺寸内选择合适的过滤器，并使用Gram-Schmidt处理。这种结构增强了神经网络的特征表示。

图6：通道注意力在CAM中的实现

空间注意力：空间注意通过调整对局部结构的注意来帮助模型适应空间变异性，提高泛化能力。如图7所示，对于大小为 $H*W$ 的输入特征F中的每个特征点，沿通道轴的最大值和平均值记为 $F_{\text{max}} \in \mathbb{R}^{1 \times H \times W}$ 和 $F_{\text{avg}} \in \mathbb{R}^{1 \times H \times W}$ ，并连接成一个 $2*H*W$ 张量。这个张量经过卷积来创建一个空间注意力地图，突出或抑制特定的位置。

图7:空间注意力在CAM中的实现

2)注意力门：本文将注意力门纳入跳跃连接过程。图8说明了注意力门单元的体系结构。设 $x_{l}$ 表示第 $l$ 层的特征图，对于每个像素 $i$ ，使用一个门控信号 $g_{i}$ 向量在更大的尺度上识别焦点区域。注意系数为 $\alpha$ ，取值范围为0 ~ 1，选择相关的特征响应，抑制不相关的特征细节。将 $x_{l}$ 和α逐元素相乘得到 $x_{output}$ ，计算如下：

根据公式，通过加性注意推导出门控系数 $\alpha$ 。考虑到医学图像涉及多个语义类的复杂性，本文引入了多维关注系数来集中目标区域。多维注意系数的计算包括以下几个方面：

其中 $W_{x}$ ， $W_{g}$ 为偏置， $\sigma(x) \ = \ \frac{a}{1 + e^{-x}}$ 为Sigmoid函数， $\sigma(x) \, = \, \max(0, x)$ 为ReLU函数。对于门控信号向量 $g_{i}$ ，我们采用 $1 \times 1$ 通道卷积（在式中表示为 $\Psi$ ）作为特征映射 $x_{l}$ 上的线性变换。

图8：注意力门

六、实验过程

6.1、数据集

表Ⅱ：数据集详情

6.2、评价指标

6.2.1、骰子相似性系数

6.2.2、平均交并比

6.2.3、豪斯托夫距离

6.3、实验设置

所有实验均在Pytorch 2.0.1中实现，并在具有24 GB内存的单个GEFORCE GTX 4090 GPU上进行了训练。本文使用AdamW优化器和SoftMax激活函数使用标准的后传播。训练的批次大小为16，初始学习率为1E-5，并且运行30个轮次。总训练时间因数据集大小而异：Brats 2019大约20小时，Brats 2020的30小时和Brats 2021的50小时。

6.4、消融实验

表Ⅲ：骰子分数的消融实验

表Ⅳ：A4-Unet在不同数据集上的实验结果

表Ⅴ：Brats数据集上的性能比较

七、总结

在本文中，作者提出了A4-Unet，这是一种脑肿瘤分割网络，引入了可变形大核卷积（DLKA）、Swin空间金字塔池化（SSPP）和注意力机制，同时保持了相对较低的网络复杂度。该方法能够实现高效的多尺度特征提取，捕捉长程依赖关系，并整合高层和低层语义信息。本文在三个数据集上的对比实验表明，A4-Unet显著优于多种先进模型，在分割性能上树立了新的基准。值得注意的是，本文的模型在Dice Score和mIoU指标上取得了显著提升。