点击下方卡片,关注 「AI视界引擎」 公众号
准确的医学图像分割对于临床量化、疾病诊断、治疗计划等应用至关重要。卷积基和Transformer基的u型架构在各种医学图像分割任务中取得了显著的成功。前者可以有效地学习图像的局部信息,而需要更多的图像特定归纳偏见,这些偏见与卷积操作有关。后者可以有效地捕捉不同特征尺度上的长程依赖性,利用自注意力,然而,当序列长度增加时,它通常会遇到计算和内存需求的挑战。
为了解决这个问题,作者通过在精心设计的u型架构中集成这两种范式的优点,提出了一种混合但有效的CNN-Transformer网络,名为BRAU-Net++,用于准确的医学图像分割任务。具体来说,BRAU-Net++使用双级路由注意力作为核心构建块来设计作者的u型编码器-解码器结构,其中编码器和解码器都是分层构建的,以便同时学习全局语义信息并减少计算复杂性。此外,该网络通过结合通道-空间注意力(采用卷积操作)重构跳过连接,旨在最小化局部空间信息损失并放大多尺度特征的全球维度交互。
在三个公共基准数据集上的广泛实验表明,作者提出的算法在其他所有状态最先进的方法中超过了 Baseline :BRAU-Net,在几乎所有评估指标上。在Synapse多器官分割、ISIC-2018挑战和CVC-ClinicDB上的平均Dice-相似系数(DSC)分别为82.47、90.10和92.94,以及ISIC-2018挑战和CVC-ClinicDB上的平均mIoU分别为84.01和88.17。代码将在GitHub上提供。
1 Introduction
准确且鲁棒的医学图像分割在计算机辅助诊断系统中起着至关重要的作用,尤其是在图像引导的临床手术、疾病诊断、治疗计划和临床量化方面。医学图像分割通常认为与自然图像分割基本相同,并且其相应技术通常源于后者。这两个社区共同之处在于,他们都将提取图像的准确感兴趣区域(ROIs)作为研究目标,无论是手动还是全自动地进行。得益于深度学习技术,自然图像视觉分割任务已经取得了令人瞩目的性能。
然而,与自然图像分割不同,医学图像分割需要更准确的ROIs分割结果,例如异常和器官,以快速识别病变边界并准确评估病变程度。这是因为临床实践中,医学图像的微小分割错误可能导致临床设置中的用户体验不佳,并增加随后的计算机辅助诊断。此外,在各种成像模式中手动勾画病变及其边界需要大量的努力,这既耗时又实际,并且分割结果可能受到临床医生的喜好和专业知识的影响。因此,作者认为开发智能和鲁棒的技术以有效地和准确地分割医学图像中的病变区域或器官至关重要。
随着深度学习的不断发展以及广泛而具有前景的应用,许多依赖卷积操作的医学图像分割方法已经提出,用于在医学图像中分割特定的目标物体。在这些方法中,像U-Net和全卷积网络(FCN)这样的u型编码器-解码器架构已经成为医学图像分割的主导方法。后续出现了许多变体,例如U-Net++,U-Net 3+,Attention U-Net,3D U-Net,和V-Net ,用于各种医学影像模式(如心电图分割、多器官分割和息肉分割)的图像和体积分割,并在广泛的医学应用中取得了卓越的成功。
这些基于CNN的方法的优秀性能表明CNN具有强大的学习语义信息的能力。但是,由于卷积操作固有的局部性,它们在明确捕捉长程依赖性方面往往存在局限性。一些研究试图通过使用空洞卷积层,自注意力机制,和图像金字塔来解决这个问题。然而,这些方法在提高建模长程依赖性的能力方面并没有显著改善。
最近,受到自然语言处理(NLP)领域中transformer的巨大成功启发,许多研究试图将transformer应用于视觉领域。这些工作在各种视觉任务上都取得了一致的改进,表明视觉transformer在视觉领域具有巨大的潜力。在这些工作中,一个热门的话题是如何通过改进核心构建块(即注意力)来提高模型性能。
作为视觉transformer的核心构建块,注意力是一种强大的工具来捕捉长程依赖性。然而,原始注意力是一种全注意力机制,它计算所有空间位置之间的配对标记相关性,因此具有很高的计算复杂性和占用大量的内存空间。为了缓解这个问题,一些工作尝试将稀疏注意力应用于视觉transformer,其中每个 Query 标记只关注关Key和Value标记的一部分,而不是整个序列。
为此,已经探索了一些手工制作的稀疏模式,例如在局部窗口中限制注意力,膨胀窗口,或轴向条纹。在医学图像视觉领域,许多研究也将transformer引入医学图像分割任务,如nnFormer, UTNet, TransUNet, TransCeption, HiFormer, Focal-UNet,和MISSFormer。然而,据作者所知,较少的工作考虑引入稀疏性到这个领域,其中代表性的工作涉及Swin-Unet和Gated Axial UNet(MedT)。但这些稀疏注意力机制是将或选择稀疏模式以手工方式。因此,这些模式是 Query 无关的。也就是说,它们由所有 Query 共享。将动态和 Query 感知的稀疏性应用于医学图像分割仍然是一个 largely 未探索的领域。
以上提到的所有问题都激励作者探索一种全自动的高级分割算法,该算法可以依靠医学图像的性质产生有效的分割结果,从而使更多的图像引导医学应用受益。
最近,受到BiFormer在将稀疏注意力应用于视觉transformer方面的成功启发,作者提出了BRAU-Net++,以利用transformer的力量进行医学图像分割。据作者所知,BRAU-Net++是首先将动态稀疏注意力集成到CNN-Transformer架构的混合模型。BRAU-Net++也来自BRAU-Net,它使用BiFormer块构建了一个u型纯transformer网络结构,带有跳过的连接用于公共的symphysis-fetal头分割。
与Swin-Unet和BRAU-Net相似,网络结构的主要组成部分包括编码器、Neck、解码器和跳过连接。编码器、Neck和解码器都基于BiFormer的核心构建块:双级路由注意力,它有效地模拟了长程依赖性并节省了计算和内存。同时,受到全局注意力机制的启发,作者通过结合通道-空间注意力(通过卷积操作实现)重新设计了跳过连接,以最小化局部空间信息损失并放大多尺度特征的全球维度交互。此外,所提出的架构利用了深度卷积来隐式编码位置信息。
在三个公开的医学图像数据集上进行了大量的实验:Synapse多器官分割,ISIC-2018挑战,和CVC-ClinicDB,表明所提出的方法取得了有前途的表现和鲁棒的一般化能力。
主要的贡献可以总结如下:
- 引入了一个u型混合CNN-Transformer网络,该网络使用双级路由注意力作为核心构建块来设计编码器-解码器结构,其中编码器和解码器都是分层的,这样可以有效地学习局部-全局语义信息,同时减少计算复杂性。
- 重新设计了传统的跳跃连接,使用通道-空间注意力机制,并提出了Skipp Connection with Channel-Spatial Attention (SCCSA),旨在在通道和空间方面增强跨维度交互,并补偿下采样造成的空间信息损失。
- 在三个常用的数据集上验证了BRAU-Net++的有效性:Synapse多器官分割、ISIC-2018挑战和CVC-ClinicDB数据集。结果表明,在几乎所有评估指标上,所提出的BRAUNet++优于其他最先进的(SOTA)方法。
II Related Work
U-Shaped Architecture
Ii-A1 CNN-Based U-Shaped Architecture for Medical Image Segmentation
这个范式的关键技术包括U-Net和FCN,以及后续的变体,其中一些被引入到2D或3D医学图像分割领域。由于U型结构的简单性和优越性能,各种类似于U-Net的方法,如U-Net++,UNet 3+,和DCSAU-Net在2D医学图像分割领域不断涌现。
此外,还有其他方法被引入到3D医学图像分割领域,如3D-Unet和V-Net。这种方法采用一系列卷积池化操作来设计其编码器和解码器。由于其强大的表示能力,该范式在广泛的医学应用中取得了巨大的成功。
Ii-A2 Transformer-Based U-Shaped Architecture for Medical Image Segmentation
原始的Transformer架构最初是针对机器翻译任务提出的,已成为自然语言处理(NLP)问题的实际标准。后续的工作试图将Transformer应用于计算机视觉。最近,研究行人试图开发纯Transformer或混合Transformer以执行医学图像分割。在[35]中,提出了一种纯Transformer,即Swin-Unet,用于医学图像分割,其中来自原始图像的标记化块而不是CNN特征图被输入到架构中以进行局部-全局语义特征学习。
在[1]中,提出了一种CNN-Transformer混合模型,TransUNet,利用CNN的高分辨率空间信息以及Transformer编码的全局上下文来实现卓越的分割性能。类似于TransUNet,UNETR和Swin UNETR在编码器中采用Transformer,并利用卷积解码器生成分割图。这些工作使用全注意力或静态稀疏注意力计算配对标记的亲和力。与这些方法不同,引入了动态稀疏注意力来选择最相关的标记,网络输入来自原始图像的标记化块。因此,由于分辨率较低,信息不会丢失。同时,作者将卷积操作应用于跳过连接,以增强多尺度特征的全局维度交互。
Sparse Attention Mechanism
稀疏连接模式已经引入来解决原始注意力机制的计算和内存复杂性。稀疏注意力在视觉Transformer中越来越受欢迎。在Swin Transformer中,注意力被限制在非重叠的局部窗口上,并引入了移位窗口操作以方便相邻窗口之间的跨窗口通信。因此,这种注意力是基于局部窗口的手工设计。随后的研究还引入了各种手工设计的稀疏模式,例如膨胀窗口或交叉形状窗口。
最近,基于动态标记稀疏度的有效视觉Transformer取得了巨大的成功。在[51]中,通过分层剪枝来动态选择要传递到下一层的标记数量,从而实现了推理加速。在[25, 24]中,他们分别提出了四叉树注意力和双级路由注意力,以实现从粗粒度到细粒度的 Query 自适应稀疏度。差异在于,双级路由注意力旨在定位最相关的几个关Key-Value对,而四叉树注意力构建了一个标记金字塔并从不同粒度 Level 组装信息。在本工作中,作者试图使用BiFormer块作为基本单元来构建一个u型编码器-解码器架构,并使用SCCSA模块进行医学图像分割。
Channel-Spatial Attention
在计算机视觉中,关注机制的研究取得了很大的进展,其中通道注意力和空间注意力是两个重要方向。通道注意力关注CNN中的通道信息。例如,SENet可以自适应地校准CNN的通道特征响应,从而增强网络的判别能力。
另一方面,空间注意力关注相关的空间区域。例如,STN可以自动地转换空间中的各种变形数据,并捕获重要的区域特征。在这些个体成功的基础上,CBAM将通道注意力和空间注意力以 ConCat 的方式结合,共同捕获通道和空间位置之间的复杂依赖关系。受到全局注意力机制的启发,作者使用通道-空间注意力来重新设计跳过连接,以增强通道-空间维度的交互作用,并补偿由于下采样导致的空间信息损失。
III Method
在这一节中首先简要总结一下B-Level Routing Attention(BRA)。然后,描述了所提出的BRAU-Net++的整体架构。最后,介绍了BiFormer块和Skip Connection Channel-Spatial Attention module(SCCSA)。
Preliminaries: Bi-Level Routing Attention
双级路由注意力(BRA)是一种动态的、 Query 感知的稀疏注意力机制,其核心思想是在粗糙粒度的区域 Level 过滤掉最不相关的关Key-Value对,只保留大多数相关路由区域的一小部分,以实现细粒度的Token-Token注意力。
与其他手工制作的静态稀疏注意力机制相比,BRA更容易建模长程依赖性。这一点与原始注意力相似。但是,BRA的复杂度为,而原始注意力的复杂度为。
Iii-A1 Region Partition and Linear Projection
将一个2D输入特征图分成个非重叠区域,每个区域的特征维度可以得到。接下来,基于产生的特征图,可以通过线性投影得到 Query 、键、值。
其中,分别是 Query 、键、值对应的线性投影权重矩阵。
Iii-A2 Region-to-Region Routing
该过程首先分别计算每个区域的和的平均值,得到区域级 Query 和键,。接下来,通过矩阵乘法在和转置的之间计算区域-区域邻接矩阵,。最后,关键步骤是通过路由索引矩阵,仅保留每个 Query 区域的前个最相关区域,,使用行向量顶运算符:topKRIndex()。区域-区域路由可以表示为:
Iii-A3 Token-to-Token Attention
由于路由区域可能在整个特征图上空间分散,需要将路由区域中的Key和Value张量聚合。然后,在这些Key-Value张量上应用细粒度的Token-Token注意力。
这个过程如图1所示,可以表示为:
其中,是收集到的Key和Value张量。函数LCE()使用深度卷积进行参数化。
Architecture Overview
BRAU-Net++的整体架构如图2(a)所示。BRAU-Net++包括编码器、解码器、Neck和SCCSA模块。对于编码器,给定一个大小为的输入医学图像,将医学图像分成重叠的 Patch ,每个 Patch 的特征维数为(定义为C)的任意维度。通过 Patch 嵌入将转换后的 Patch 标记传递给多个BiFormer块和 Patch 合并层以生成层次特征表示。具体而言, Patch 合并用于降低特征图的分辨率并增加维度,而BiFormer块用于学习特征表示。对于Neck,特征图的分辨率和维度保持不变。
受到U-Net和Swin-Unet的启发,作者设计了一个对称的基于Transformer的解码器,由BiFormer块和 Patch 扩展层组成。 Patch 扩展层负责上采样和降低维度。通过SCCSA模块将提取的上下文特征与编码器中的多尺度特征融合,以补充下采样造成的空间信息损失并增强全局维度交互。最后,使用 Patch 扩展层进行上采样以恢复特征图的原分辨率,然后使用线性投影层生成像素级分割预测。
BiFormer Block
结构模块的核心是双级路由注意力(BRA)。如图3所示,BiFormer块包括一个3×3的深度卷积,2个LayerNorm(LN)层,一个BRA模块,3个残差连接和2层MLP,其中扩展比 = 3。3×3的深度卷积可以隐式编码相对位置信息。BiFormer块可以表示为:
其中,,和分别表示第个块的深度卷积、BRA模块和MLP模块的输出。
Encoder
编码器采用分层的金字塔结构进行构建。具体而言,第一阶段包括一个由两个3×3卷积层组成的 Patch 嵌入层和一个由3×3卷积层组成的 Patch 合并层,用于降低输入空间分辨率并增加通道数。
如图2所示,具有分辨率和C通道的标记化输入被输入到第一阶段的连续两个BiFormer块中,以进行表示学习。第二阶段和第三阶段的标记化输入也以类似的方式处理。 Patch 合并层进行2×下采样,将标记数减半,并将特征维度增加2×。
Decoder
与编码器类似,解码器也基于BiFormer块构建。受到Swin-Unet的启发,采用 Patch 扩展层来上采样解码器中提取的深度特征。 Patch 扩展层主要用于将特征图重新形状为更高分辨率的特征图,即增加分辨率2倍,并将特征维度减半。最后, Patch 扩展层进行4倍上采样,输出分辨率为的特征图,用于像素级分割预测。
Skip Connection Channel-Spatial Attention (SCCSA)
通道和空间注意力的结合可以增强模型捕捉更广泛的上下文特征的能力,与使用单一注意力机制相比。受到[39]的启发,考虑在跳过连接中应用顺序通道-空间注意力机制,并提出跳过连接通道-空间注意力,简称SCCSA。SCCSA模块可以有效地补偿下采样引起的空间信息损失,增强解码器中每个层的多尺度特征的全局维度交互,从而在生成输出Mask的同时恢复细粒度细节。
如图2(b)所示,SCCSA模块包括通道注意子模块和空间注意子模块。具体而言,首先通过将编码器和解码器的输出进行拼接得到。然后,通道注意子模块利用多层感知机(MLP)结构,降低比率为 = 4,放大跨维度通道-空间依赖性。使用两个7×7卷积层来聚焦于具有相同降低比率的空间信息。给定输入特征图,中间状态和输出的定义如下:
其中,和分别是通道和空间注意子模块的输出;和分别表示元素乘法和 sigmoid 激活函数。
Loss Function
在训练过程中,对于Synapse数据集,采用混合损失来解决类别不平衡问题,将 dice 损失和交叉熵损失相结合。对于 ISIC-2018 和 CVC-ClinicDB 数据集,仅使用 dice 损失来优化模型。dice 损失(),交叉熵损失()和混合损失()的定义如下:
其中,表示像素数量,和分别表示类别的真实标签和生成的概率。是类别数量,且所有类别的权重和为1。是一个平衡因子,用于平衡和的影响。在作者的研究中,和分别设置为和0.6。BRAU-Net++的训练过程可以概括为算法1。
4 Experimental Settings
Datasets
作者在三个公开可用的医学图像分割数据集上训练和测试了提出的BRAU-Net++:Synapse多器官分割,ISIC-2018挑战,和CVC-ClinicDB。数据集的详细划分已在表1中给出。
所有数据集都与临床诊断相关,其分割结果对患者的治疗至关重要,并包括图像及其相应的真实Mask。选择不同成像模式数据集的主要原因是为了评估所提出方法性能和鲁棒性。
4.1.1 Synapse Multi-Organ Segmentation Dataset
自动多器官分割在腹部计算机断层扫描(CT)上可以支持临床诊断、治疗计划和治疗流程。实验中使用的数据集包括MICCAI 2015多部位CT标注挑战中30个腹部的CT扫描,每个CT体积包含3,779个轴向CT图像,每个图像大小为512×512像素, Voxel 空间分辨率为[0.54-0.54][0.98-0.98][2.5-5.0]mm³。遵循[1]和[35],训练集和测试集分别包含18个样本(包含2,212个轴向切片)和12个样本。
4.1.2 ISIC-2018 Challenge Dataset
本文中的数据集指的是用于ISIC-2018挑战中病变分割任务的训练集,其中包含2,594张皮肤镜像图像及其真实分割标注。为评估模型性能,进行了五倍交叉验证,并选择最佳模型进行推理。
4.2.3 CVC-ClinicDB Dataset
CVC-ClinicDB数据集通常用于息肉分割任务。它也是MICCAI 2015自动息肉检测挑战的训练数据集。该数据集包含612张图像,这些图像随机分为490张训练图像、61张验证图像和61张测试图像。
Evaluation Metrics
为了评估所提出的BRAU-Net++的性能,采用了平均Dice相似系数(DSC)和平均Hausdorff距离(HD)作为评估指标,用于评估BRAU-Net++在8个腹部器官(主动脉、胆囊、脾脏、左肾、右肾、肝脏、胰腺、脾脏和胃)上的性能,其中仅在单个器官的评估中使用DSC。
此外,平均交点与交点(mIoU)、DSC、精确度、召回率等也被用于评估在ISIC-2018挑战和CVC-ClinicDB数据集上的模型性能。形式上,预测可以分为真阳性(TP)、假阳性(FP)、真阴性(TN)和假阴性(FN),然后计算DSC、IoU、精确度、召回率等指标。
HD可以描述为:
其中,和分别表示真实Mask和预测分割图,表示点和之间的欧几里得距离。
Implementation Details
在配备24GB内存的NVIDIA 3090显卡上训练BRAU-Net++模型及其各种变体。使用Python 3.10和PyTorch 2.0实现。
在训练期间,将模型在上述三个数据集(Synapse多器官分割、ISIC-2018挑战和CVC-ClinicDB)上进行初始化和微调,使用在ImageNet-1K上预训练的BiFormer权重,并考虑空间,仅在Synapse多器官分割数据集上从零开始训练所提出的模型。在这些产生的模型上,进行了一系列的ablation研究,以分析每个组件的贡献。
对于Synapse多器官分割数据集,将所有图像调整为224224的分辨率,并使用随机梯度下降算法进行400个epoch的训练,批量大小为24,学习率为0.05,动量为0.9,权重衰减为1e-4。对于ISIC-2018挑战和CVC-ClinicDB数据集,将所有图像调整为256256的分辨率,并使用Adam优化器进行200个epoch的训练,批量大小为16。使用CosineAnnealingLR计划,初始学习率为5e-4。数据增强方法包括水平翻转、垂直翻转、旋转和剪裁,概率为0.25,以增强数据多样性。
还有一些超参数也是经验性地设置的。例如,区域划分因子根据224224和256256的分辨率分别设置为7和8。从阶段1到阶段7的的数目分别为2、4、8、、8、4和2,其中表示使用全注意力。
5 Experimental Results
在本节中将详细比较所提出的BRAU-Net++与其他最先进的(SOTA)方法,包括基于CNN的方法、基于Transformer的方法以及同时基于CNN和Transformer的混合方法,在Synapse多器官分割、ISIC-2018挑战和CVC-ClinicDB数据集上进行比较。同时,将Synapse多器官分割数据集作为一个示例,在该数据集上进行广泛的本实验分析,以分析BRAU-Net++中每个组成部分的影响。
Comparison on Synapse Multi-Organ Segmentation
正如上面提到的,自动多器官腹部CT分割在提高临床工作流程的效率,包括疾病诊断、预后分析和治疗计划方面起着重要作用。因此,选择这个数据集来评估各种方法的表现。在Synapse多器官腹部CT分割数据集上,与以前的最先进方法(SOTA)在DSC和HD方面的比较如表2所示,其中粗体表示最佳结果。
结果是在作者实验设置下根据公开发布的代码复制的,其他结果直接来自各自发表的论文。作者的BRAU-Net++在两个评估指标上都明显优于基于CNN的方法和作者自己的 Baseline BRAU-Net,这表明深度混合CNN-Transformer模型可能能够建模全局关系和局部表示。
与当前的基于Transformer的方法:TransUNet和Swin-Unet相比,作者的BRAU-Net++在DSC方面分别提高了4.49%和3.34%,在HD方面分别降低了12.62mm和2.48mm。这表明使用双级路由注意力作为核心构建块来设计u型编码器-解码器结构可能有助于有效地学习全局语义信息。具体而言,BRAU-Net++在大多数器官的分割上稳定地优于其他方法,特别是对于左肾和肝脏的分割。
从表2中可以看出,BRAU-Net++在DSC值上达到82.47%,与其他方法相比具有更高的重叠,表明预测的分割图与地面真相Mask之间的差异最小。作者还在HD方面达到19.07mm,与其他方法相比具有相对较低的值,分别获得最佳结果14.7mm和次佳结果18.20mm。与HiFormer和MISSFormer相比,BRAU-Net++在HD上的提高仅为0.87mm,但比HiFormer明显提高了4.37mm,这表明BRAU-Net++学习目标边缘信息的能力可能不如HiFormer。
总的来说,表2表明除HiFormer和MISSFormer之外,BRAU-Net++在先前的研究中取得了显著的改进,例如DSC性能提高范围为0.51%至12.2%,HD性能提高范围为1.59mm至20.63mm。因此,作者认为BRAU-Net++仍有潜力获得相对更好的分割结果。
另外,从表2中可以看出,BRAU-Net++的参数数量约为50.76M,其中SCCSA模块产生了约19.36M的参数。但是,带有SCCSA模块的BRAU-Net++在DSC方面的性能比不带SCCSA模块的BRAU-Net++略有提高。在HD方面也有类似的观察。参数数量的有效性将在随后的部分中进行讨论。
一些不同方法在Synapse数据集上的定性结果如图4所示。从图4中可以看出,BRAU-Net++生成了对胆囊、左肾和胰腺的平滑分割图,这表明双级路由注意力可能擅长捕捉小型目标的特征,而BRAU-Net++可以更好地学习全局和长程语义信息,从而产生更好的分割结果。
在ISIC-2018挑战数据集上,BRAU-Net++在mIoU、DSC和Accuracy方面取得了最佳性能,在Precision和Recall方面取得了次佳结果。
如图5(左)所示,与最近发表的DCSAU-Net和BRAU-Net相比,BRAU-Net++在mIoU方面取得了1.84%和1.2%的改进。此外,BRAU-Net++在召回率方面达到了0.9224,这在临床应用中更有利。从以上分析和图5(左)可以看出,BRAU-Net++在ISIC-2018挑战数据集上对其他方法的边界分割预测更好,分割Mask的轮廓更接近 GT 值。
Comparison on CVC-ClinicDB
在息肉有可能转变为结直肠癌之前进行早期检测可以提高生存率。这对临床实践具有重要意义。因此,作者在实验中选择了这个数据集。
表4中呈现了BRAU-Net++在mIoU、DSC、Precision和Recall方面的定量结果。BRAU-Net++在mIoU、DSC、Precision和Recall方面取得了最佳结果,分别比第二好的方法提高了1.99%、1.27%、2.12%和1.03%。定性结果如图5(右)所示。可以看出,BRAU-Net++生成的息肉Mask与 GT 值的边界和形状非常接近。
Ablation Study
在本节中对上述三个数据集进行了广泛的消融研究,以彻底评估BRAU-Net++中涉及的每个组件的有效性。具体而言,作者对SCCSA模块、跳过连接数量和值、输入大小和划分因子、模型尺度和预训练权重的影响进行了消融研究。
4.4.1 Effectiveness of SCCSA Module
SCCSA模块是BRAU-Net++的重要组成部分。它使用通道-空间注意力来增强通道和空间方面的跨维度交互,并有助于生成更精确的分割Mask。
表2显示了BRAU-Net++在Synapse上的结果。与BRAU-Net++无SCCSA相比,BRAU-Net++在DSC方面提高了0.91%,在HD方面减少了0.39mm。这种微小的改进是以引入大量参数为代价的:它为这个模型带来了巨大的参数数量。
一个主要原因可能是多尺度CNN特征与由分层 Transformer 结构学习的全局语义特征的组合并不能显著提高分割任务。至于确切的原因,作者打算将其留为未来的工作,进一步探索和分析。ISIC-2018挑战和CVC-ClinicDB数据集的分割结果如表5所示。
可以看出,将SCCSA模块添加到BRAU-Net++模型中可以在几乎所有评估指标上实现最佳结果。例如,SCCSA可以在与mIoU指标相比分别提高ISIC-2018挑战和CVC-ClinicDB的结果0.6%和0.9%。此外,作者还计算了参数数量、浮点运算(FLOPs)和每秒帧数(FPS),以进一步研究该模块的有效性。作者可以观察到,SCCSA在两个数据集上并没有显著降低FPS,特别是对于CVC-ClinicDB。
4.4.2 Effectiveness of the Number of Skip Connections
已经观察到,u型网络的跳过连接可以帮助通过恢复低级空间信息来改善更精细的分割细节。这个消融主要是为了探索不同的跳过连接数量对BRAU-Net++性能提升的影响。这个实验在Synapse数据集上进行。在1/4、1/8和1/16分辨率尺度的位置添加跳过连接,通过在不同位置的连接组合,可以将跳过连接的数量改变为0、1、2和3,其中"0"表示不添加跳过连接。其他添加的连接及其相应的在平均DSC和HD指标上的分割性能如表6所示。
可以观察到,随着跳过连接数目的增加,分割性能逐渐提高,而在所有位置的1/4、1/8和1/16分辨率尺度上添加跳过连接时,最佳的平均DSC和HD被实现。因此,采用这种配置为BRAU-Net++增强学习精确低级细节的能力。这可能是BRAU-Net++能够捕捉小目标特征的主要原因。
4.4.3 Effectiveness of Input Resolution and Partition Factor
进行这项消融实验的主要目的是测试输入分辨率对模型性能的影响。在Synapse数据集上对128128、224224和256256分辨率尺度进行了三组实验,并将结果报告在表7中。遵循[24],将划分因子设置为特征图大小的除数以避免填充,具有不同输入分辨率的图像应采用不同的划分因子。
因此,作者将上述三种分辨率的对应划分因子设置为 = 4、 = 7和 = 8。可以看出,保持 Patch 大小相同(例如,32)并逐渐增加分辨率尺度,即增加标记序列长度,可以导致模型性能的一致提高。这与常识相符,即分辨率更高的图像包含更多的语义信息,从而提高性能。然而,这需要付出更大的计算成本。因此,考虑到计算成本,为了与其他方法进行公平比较,所有实验都基于默认的224224分辨率作为输入。
V-B4 Effectiveness of the Number of Top-
类似于[24],随着下一阶段的routed区域大小逐渐减小,相应地增加以保持合理的token数量以进行注意力。在Synapse数据集上对每个阶段中top-的数量进行了消融研究,并将结果列在表8中,其中列出了网络每个阶段中token的数量。
可以看出,在编码器接近顶部阶段增加token数量似乎可以提高分割性能。这可能是因为网络的接近顶部块可以捕获低级信息,例如边缘或纹理,这对于分割任务至关重要。此外,盲目增加token以进行注意力的数量可能会损害性能,这表明显式稀疏约束可以作为正则化来提高模型的泛化能力。
V-B5 Effectiveness of Model Scale and Pre-trained Weights
研究了网络深度的效果。同样地,众所周知,基于Transformer的模型性能受到模型预训练的影响。因此,考虑从零开始训练和预训练两种不同的BRAU-Net++模型规模的消融研究。BRAU-Net++的两种不同模型规模分别称为微型和基础模型。它们的配置和Synapse数据集上的结果列在表9中。
可以看出,基础模型获得了更优的结果。特别是在HD评估指标上,基础模型的结果比微型模型提高了14.77mm。这表明基础模型可以实现更好的边缘预测。因此,采用基础模型进行医学图像分割。考虑到计算性能,在所有实验中采用"基础"模型。
VI Discussion
在这项工作中证明了动态和 Query 感知稀疏注意力在降低计算复杂度和提高模型性能方面的有效性。为了更直观地说明稀疏注意力在医学图像分割任务上的作用,作者按照[24]的方法,可视化了路由区域和与 Query 标记相关的注意力响应。
采用了路由索引和注意力得分,这些值来自编码器第三阶段的最后块。也就是说,这些值来自分辨率的特征图,而可视化结果以原始分辨率呈现。在Synapse多器官分割、ISIC-2018挑战和CVC-ClinicDB数据集上的结果如图6所示。
可以看出,稀疏注意力的类型可以有效地找到语义上最相关的区域,这表明动态稀疏注意力计算机制对于医学图像的计算和选择稀疏模式是有效的。然而,探索其他更有效的稀疏模式计算方法仍然必要,并且是作者未来工作的重点。
作者进行了一系列消融研究来评估BRAU-Net++中每个相关组件的贡献,其中作者提出了SCCSA模块来增强来自编码器阶段和来自解码器阶段的特征在通道和空间方面的跨维度交互。实验结果显示,在几乎所有评估指标下,这种微小的改进都带来了巨大的参数数量。这是作者的工作的一项缺点。作者认为主要原因可能是多尺度CNN特征与由分层 Transformer 结构学习的全局语义特征的组合不能显著提高分割任务。在未来的工作中,作者将关注如何有效地解决这个问题。
三个不同的成像模式数据集:Synapse多器官分割、ISIC-2018挑战和CVC-ClinicDB,被故意选择作为基准。选择这个的主要原因是为了评估所提出方法在多模态医学图像分割任务上的性能和鲁棒性。大量的实验揭示了BRAU-Net++在多模态医学图像分割任务上的广泛适用性。
VII Conclusion
在本文中,作者提出了一种精心设计的u型混合CNN-Transformer架构BRAU-Net++,该架构利用动态稀疏注意力代替全注意力或静态手工稀疏注意力,可以有效地学习局部-全局语义信息,同时降低计算复杂度。此外,作者还提出了一种新的模块:跳过连接通道-空间注意力(SCCSA),以集成多尺度特征,以补偿空间信息的损失并增强跨维度交互。
实验结果显示,BRAU-Net++可以在Synapse多器官分割、ISIC-2018挑战和CVC-ClinicDB数据集上的几乎所有评估指标上实现最先进的性能,尤其是在捕捉小目标特征方面表现出色。对于未来的工作,作者将专注于设计更复杂和通用的架构,用于多模态医学图像分割任务。
参考
[1]. BRAU-Net++: U-Shaped Hybrid CNN-Transformer Network for Medical Image Segmentation
点击上方卡片,关注 「AI视界引擎」 公众号