超越 U-Net,跨尺度交互模块的引入,TransDAE 重新定义自注意力机制以增强图像分割 !

图像处理机器学习关系型数据库

点击下方卡片,关注 「AI视界引擎」 公众号

( 添加时备注:方向+学校/公司+昵称/姓名 )

picture.image

在医疗保健领域,医学图像分割对于准确的疾病诊断和制定有效的治疗策略至关重要。早期检测可以显著帮助管理和控制疾病,可能防止其进展。

机器学习,尤其是深度卷积神经网络,已成为解决分割挑战的有前途的方法。

传统方法如U-Net使用编码块进行局部表示建模,解码块以揭示语义关系。然而,这些模型通常难以处理具有显著纹理和形状变化的多种尺度的物体,且它们往往捕捉不到输入数据中的长程依赖性。

为了解决这些问题,作者提出了TransDAE:一种新颖的方法,重新设想了自注意力机制,包括整个特征空间中的空间和通道级关联,同时保持计算效率。

此外,TransDAE通过添加跨尺度交互模块增强了 Short-Cut 路径,促进了特征重用并提高了局部化准确性。

值得注意的是,TransDAE在Synaps多器官数据集上超过了现有最先进的性能,而无需依赖预训练权重。

I Introduction

早期疾病诊断在医疗健康领域至关重要,因为它能帮助检测疾病严重程度和扩散的初期阶段 [1]。医学图像分割是计算机辅助疾病诊断(CAD)、治疗规划以及外科预评估的关键组件。分割过程涉及通过像素级分类对目标器官和组织的形状和体积进行划分 [2]。传统的手工标注劳动密集、耗时且容易出错 [3]。因此,医学图像分割的自动化已成为研究热点,以减轻这种负担。近年来,深度学习在计算机辅助诊断(CAD)中的应用潜力得到了广泛探索,因为它在不同医学领域中得到了成功应用。

多年来,卷积神经网络(CNNs)和全卷积网络(FCNs)在医学图像分割中占主导地位,主要是因为它们可以通过卷积操作学习层次特征。在这些模型中,U-Net 作为一种特别有效的模型,是由于其具有 U 字形结构和对称的编码和解码路径 [4]。这种结构,通过跳接连接,促进低级和高级特征的融合,提高上下文建模并产生准确的分割结果。Res-UNet [5],Dense-UNet [6],U-Net++ [7],UNet3+ [8] 等变体通过解决特定限制,如更深层的空间信息丢失,进一步提高了性能。

然而,尽管取得了这些进展,基于 CNN 的模型还是存在固有的局限性。卷积操作的局部性质意味着这些模型往往难以捕捉医学图像中的长时间依赖和多尺度的变化,这对分割复杂的解剖结构至关重要。尽管 atrous convolution [9] 和 pyramid pooling [10] 等方法旨在捕捉更大的上下文信息,但处理全球关系仍具有挑战性。因此,CNN 架构引入了注意力机制来提高其关注重要区域的能力,如 Attention U-Net [11] 及其变体 [12]。

近年来,Transformer 架构的引入提供了一种新的解决这些挑战的方法。最初为了自然语言处理(NLP)而发展的 Transformer,通过自注意力机制捕捉长时间的依赖关系。这个能力使其在计算机视觉任务中的应用,包括医学图像分割,表现出惊人的结果。然而,ViT [14] 和类似的 Transformer 模型也存在一些挑战。自注意力的二次计算复杂度和大量训练数据的要求使得这些模型在处理高分辨率医学图像时变得不切实际。此外,尽管 Transformers 在建模全局依赖关系方面表现出色,但往往缺乏捕捉局部细节的能力,这对于准确的分割至关重要 [15]。

为了克服这些局限性,作者提出了一种名为 TransDAE 的分层 Transformer 模型,专门用于医学图像分割。作者的方法通过在整个特征空间集成空间和通道关联来重新想象自注意力机制。这种双注意力机制使模型在保持计算效率的同时,更有效地捕捉局部和全局依赖性。此外,作者还引入了一个 Inter-Scale Interaction Module (ISIM),以增强跳接路径,提高特征利用率并改进定位精度。这个模块在确保模型处理医学图像的多尺度性质方面发挥关键作用,这也是 CNN 和 Transformers 面临的挑战。在这项工作中,作者可以总结以下贡献:

  • 作者提出了一种同时捕捉空间和通道依赖性的双注意力机制,解决了许多现有方法主要集中在一或者另一方面的局限性。
  • 作者引入了高效的自注意力和增强的自注意机制,以降低计算复杂性,同时有效地模拟低 Level 和高 Level 依赖性,使作者的模型能够扩展到高分辨率医学图像。
  • 作者强调了跳接连接在连接编码器和解码器组件中的重要性,并集成了一个 Inter-Scale Interaction Module (ISIM),以增强特征重用并提高定位准确性。
  • 作者集成了一个大型 Kernel 注意模块,进一步增强通过跳接传递的信息,放大低 Level 局部化信息的有效性,并产生更强大、更有效的网络。

II Literature Review

在 II 文献综述部分,作者将对相关研究进行全面梳理,以深入了解目前在该领域的研究动态和发展趋势。在此过程中,作者将从以下几个方面展开讨论:

  1. 现有研究的分类与概述:首先,作者将对已有的研究进行分类,以便更好地理解和组织相关内容。同时,作者将对每个类别下的研究成果进行简要概述,以帮助读者快速了解研究的重点。
  2. 研究方法与理论基础:研究方法是任何研究领域中不可或缺的部分,它决定了研究者能够达到的研究深度和广度。在文献综述中,作者将对已有的研究方法进行详细讨论,并对各种方法的优缺点进行分析。此外,作者还将探讨一些主要的研究理论基础,并阐述它们在相关领域的应用。
  3. 研究结果与分析:通过对相关文献的阅读和分析,作者将梳理出已有的研究结果,并对这些结果进行总结和归纳。在此基础上,作者将对现有研究的不足之处进行深入剖析,并提出可能的改进措施。
  4. 未来研究方向:在讨论了现有研究的基础上,作者将对领域内的未来研究趋势进行预测,并提出一些建议性的研究方向。这将为新入行的研究者提供一些启示,同时也为已在该领域取得一定成果的研究者提供新的视角。

综上所述, II 文献综述部分的目标是对研究领域内的研究进行全面梳理,找出研究中的不足之处,提出未来研究方向,以期为学术界贡献一份力量。

Cnn-Based Segmentation Networks

近年来,深度学习方法在图像分割任务中取得了主导地位,取代了基于手工特征的机器学习方法。CNN已成为各种医学图像分割任务的首选,主要归功于U-Net的成功[16]。U-Net的U型对称结构包括编码器与解码器之间的 Short-Cut ,实现从编码器网络的更高分辨率特征图与上采样特征的拼接,提供更准确的表达。U-Net的成功激发了研究行人调整其架构,通过各种策略提高性能,如Res-UNet、Dense-UNet、U-Net++、UNet3+等。3D U-Net[17]被提出作为U-Net的增强版,专门针对三维体积分割设计。

Oktay等人[11]在U-Net的 Short-Cut 中引入了注意力门,强调关注关键目标而忽略无关目标的重要性。Alryalat等人[12]在U-Net Short-Cut 中采用了双注意力策略,使用通道注意力和空间注意力使网络专注于更多的代表性通道,并在图像中识别最具有代表性的空间区域。周等人[7]开发了U-Net++,并证明使用嵌套和密集 Short-Cut 将编码器特征图注入解码器,而不是直接获取,可以提高网络性能。然而,由于卷积操作的有限感受野大小,CNN方法主要捕获局部依赖性,难以表示长期依赖性。虽然不同网络块的维数尺寸各不相同,但操作仍然局限于局部信息,而无法捕捉全局上下文。卷积操作的局部性和共享权重特性使其难以捕捉全局上下文。

为克服CNN网络的局限性,近年来出现了各种方法。 Yu等人[9]试图在不降低分辨率的情况下扩大感受野大小以捕捉全局上下文,采用了 dilation rate 的空洞卷积。Zhao等人[10]在不同的特征尺度上使用金字塔池来模拟全局信息。 Wang等人[18]提出了一种非局部网络,通过计算输入特征映射内所有特征的加权和来捕获长期依赖性。一些研究[12, 19]发现自注意力模块在长期依赖性建模中具有潜力。虽然这些努力试图缓解CNN的缺陷,但由于这些应用数据中存在强长期依赖性,它们仍无法完全满足临床应用要求。

Transformers

Transformer方法在自然语言处理领域的成功,尤其是在高依赖性的词汇之间存在的情况下,激励研究行人利用这些模型在图像分割和识别任务中利用其远程依赖能力。ViT [14]作为一个基础方法,将Transformer方法引入机器视觉领域,并超越了传统的CNN基础架构。这种方法将输入图像划分为称为patch的块,并在每个窗口内的每个patch中嵌入其位置,这样网络可以考虑块之间的空间依赖性。这些block然后被输入到Transformer编码器中,该编码器采用多头自注意力模块,然后是一个多层感知机进行分类。

为了提高这种新颖方法的性能,已经提出了几种增强版本的ViT,包括Swin Transformer [20]、LeViT [21]和Twins [22]。考虑到这些模型的复杂性,Swin Transformer [20]试图通过将图像patch划分为窗口并将Transformer在窗口内的每个patch内唯一应用来减少模型参数的数量。还建议一个额外的步骤,即允许相邻窗口之间的交互,这是基于CNN的基本原则:移动窗口然后重新应用Transformer模块。

虽然基于Transformer的方法在各种领域都取得了巨大成功,但他们也存在一些局限性。一个明显的缺陷是它们在捕获局部信息表示方面的弱项。与CNN不同,由于卷积操作的固有性质,Transformers主要擅长捕获长期依赖性。这种局限性可能导致在处理数据的关键时刻理解信息性能不佳。因此,已经变成必须探索具有CNN和Transformer优点的混合模型,以克服这些限制并提高整体性能。

Transformer方法在自然语言处理领域的成功,尤其是在词汇之间高度依赖的情况下,鼓励研究行人利用这些模型的远程依赖能力进行图像分割和识别任务。ViT [14]作为一个基础方法,将Transformer方法引入机器视觉领域,并超越了传统的CNN基础架构。该方法将输入图像划分为名为patch块,并在每个窗口内的每个patch中嵌入其位置,使得网络可以考虑block之间的空间依赖性。这些block接着被输入到Transformer编码器中,该编码器采用多头自注意力模块,然后是一个多层感知机进行分类。

为了提高这种新颖方法的性能,已经提出了几种增强版本的ViT,包括Swin Transformer [20]、LeViT [21]和Twins [22]。考虑到这些模型的复杂性,Swin Transformer [20]试图通过将图像patch划分窗口并只在窗口内的每个patch内应用Transformer来减少模型参数的数量。还建议一个额外的步骤,即允许相邻窗口之间交互,这是基于CNN的基本原则:移动窗口然后重新应用Transformer模块。

虽然基于Transformer的方法在各种领域都取得了巨大成功,但他们也存在一些局限性。一个明显的缺陷是它们在捕获局部信息表示方面的弱项。与CNN不同,由于卷积操作的固有性质,Transformers主要擅长捕获长期依赖性。这种局限性可能导致在处理数据的关键时刻理解信息性能不佳。因此,已经变成必须探索具有CNN和Transformer优点的混合模型,以克服这些限制并提高整体性能。

Hybrid CNN-transformer approaches

近年来,医学图像分割领域取得了显著进展,主要通过将Transformer架构的优势与CNN的优势相结合来实现。具体来说,就是将Transformer模块引入CNN网络的编码组件中,使其能够同时捕捉局部信息并有效建模长程依赖关系。在这一点上,TransUNet[23]是一种开创性的方法,它利用ResNet-50 Backbone 网络生成低分辨率特征图,然后使用ViT模型进行编码。编码后的特征通过级联上采样层进行上采样的过程得到最终的分割图。然而,将纯Transformer模型与CNN模型集成在一起,可以提高网络复杂度多达八倍。为了解决这一挑战,Cao等人[24]提出了Swin-UNet,它采用固定窗口(类似于Swin-Transformer的方法)计算注意力。作为一种额外的特性,Swin-UNet包括一种填充相邻特征图以生成更高分辨率特征图的层,该层在upsampling过程中 Reshape 。此外,Wu等人[25]将Transformer模块集成到编码层,用双编码器(包含CNN和Transformer分支)替换单个支路编码器。研究行人还设计了一个功能适应模块(FAM)和一种记忆高效的解码器来克服这些分支和解码组件的计算效率低下问题。类似地,Azad等人[26]通过引入“Context Bridge”解决了传统CNN方法的限制。此特性将U-Net的局部表示能力与Transformer模型相结合,克服了在处理长程依赖关系和多样物方面的问题。此外,他们用一种“高效自注意力”策略替换了标准的注意力机制,这样在保持性能的同时简化了结构。

基于CNN的方法在捕获局部信息方面表现出色,但在建模长程依赖关系(尤其是医学图像分析中很重要)方面存在困难。相反,基于Transformer的方法在长程依赖关系表示方面具有优势,但缺乏捕捉局部信息的机制。因此,作者的研究目标是以保持可接受的网络复杂度为基础,开发一种同时结合两者优点的模型。作者提出了一个双注意力模块来处理空间输入特征和通道上下文,利用Wang等人[27]的效率自注意力方法和增强自注意力模块[28]以降低复杂度。作者的重新设计的Transformer块被集成到类似于U-Net的架构中,突显了 Short-Cut 对改进性能和准确特征重建设的的重要性。通过集成一个大核方法,作者增强了信息传递,增加了低级局部定位信息的有效性,并通过更好的编码器-解码器通信最终加强了模型的整体性能。

III Proposed Method

图1展示了作者提出的模型的概述,这是一个具有U-Net类似结构的层次Transformer模型,利用了局部和全局特征表示以及增强的跳跃连接模块。给定一个输入图像,其空间维度为和通道,模型使用[24, 28]中的填充嵌入模块获得大小为的交叉patch Token 。令

picture.image

Dual Attention Transformer Block

将双注意力机制集成到作者的模型中的动机来自意识到在医学图像分割任务中,通道和空间注意力都是不可或缺的。准确的分割结果依赖于有效地表示特征张量。通道注意力使得模型可以选择性地关注最具信息量的表示,从而促进对医学图像内在结构的更深入理解。相反,空间注意力强调特征之间的空间关系,允许模型捕捉图像中各个区域之间的关键上下文信息和依赖关系。通过集成双注意力机制,作者的模型有效地结合了通道和空间注意力的优势,最终提高了在医学图像分割任务中的性能。这种方法使作者开发了一种更强大、更高效的网络,能够有效地表示特征张量,从而导致了分割结果的改进。图2直观地说明了双注意力机制的工作方式。图2说明了通道和空间注意力组件如何共同工作,以提高模型的分割能力。值得注意的是,作者的设计是将注意力机制逐个应用,而不是并行应用,从而提高了性能。

picture.image

为了在不牺牲双注意力机制相关复杂性优势的同时利用其优势,作者使用高效的注意力模块用于通道注意力,使用增强的 Transformer 块用于空间注意力。标准自注意力机制的局限性在于其二次计算复杂度为,如公式(1)所示。这限制了架构在处理高分辨率医学图像时的适用性。

在公式(1)中,、和分别表示 Query 、键和值向量,而表示嵌入维度。采用高效的注意力机制,作者可以降低计算开销,同时不牺牲通道注意力方法提供的受益。这使作者能够更有效地处理特征图,并在医学图像分割任务中实现增强性能。此外,高效的注意力机制确保模型保持可扩展性,从而使其适用于更广泛的应用场景和数据集。

与点积注意力不同,Efficient Attention在初始阶段不计算点之间的配对相似度。相反,键被表示为个注意力图,其中表示输入特征的位置。这些全局注意力图反映了输入特征的语义方面,而不是输入位置的相似性。这种重新排列大大降低了注意力机制的计算复杂性,同时保持了高的表示能力。在作者的设置中,内存复杂度为,计算复杂度为(典型的设置为,),因此作者的结构采用Efficient Attention来捕捉输入特征图的通道重要性。

为了降低空间注意力模块的复杂性,作者遵循黄等人[28]的策略,这是一种适用于高分辨率特征图的空间减少自注意力。由黄等人提出的这种策略,使用空间减少比例R允许空间分辨率降低,从而实现有效的自注意力。方程(3)说明了这种减少策略的数学公式。

论文原文: 如图中的方程所示,首先,K和V被 Reshape 为新的形状 。然后,使用线性投影W,通道深度恢复到C。这些操作将自注意力复杂度降低到,这在处理高分辨率特征图上是计算上可行的。为了实现空间减少,可以采用卷积或池化等技术。

Inter-Scale Interaction Module

注意力机制作为动态选择器,能够强调不同尺度上的相关特征并忽略冗余特征,这依赖于输入特征。这个机制的一个重要副产品是注意力图,它像聚光灯一样,高亮不同尺度上各个特征的相对重要性。这个聚光灯有助于解读不同特征之间的关系。

在分析建立特征之间关系的各种方法时,出现两种主要策略,它们分别解决不同尺度的特征。

第一种策略通常被称为“自注意力机制”[14; 30; 31]。虽然这种机制擅长理解长时间依赖,但在多个尺度上存在局限:

  • 它无意识地将图像作为一维序列处理,忽略了其固有的二维结构。
  • 它的计算需求巨大,特别是对于分辨率高的图像,其二次复杂度 especially 麻烦。
  • 尽管它在空间适应性方面表现出色,但它无法充分适应不同尺度和通道。

相比之下,第二种策略利用了大规模核卷积(large-kernel convolutions)的能力[32; 33; 34]。这些卷积天然就擅长在多个尺度上工作,分辨出特征重要性并生成注意力图。然而,这种方法不是没有挑战。主要问题是引入大规模核卷积会增大计算开销并增加参数数量。

作者从Guo等人提出的视觉注意力网络[35]中获得灵感,提出了一种创新的混合策略:自注意力和大规模核卷积。这种混合策略能够解决不同尺度间的交互问题。通过分解大规模核卷积操作,作者希望能更细致地理解不同尺度上的长时间依赖关系。如图3所示,大规模核卷积可以分解为三个主要部分,分别针对不同的尺度:

picture.image

  • 空间局部卷积通过深度卷积。
  • 空间长时间卷积通过深度膨胀卷积。
  • 通道卷积,由紧凑的卷积实现。

进一步观察,卷积可以分解为三个子部分,分别针对不同的尺度和通道:

  • 由卷积实现了通道上的 。
  • 将 卷积分解为三个子部分,分别针对不同的尺度和通道:
  • 空间局部卷积通过深度卷积。
  • 空间长时间卷积通过深度膨胀卷积。
  • 通道卷积,由紧凑的卷积实现。

作者的跨尺度交互策略避免了辅助归一化函数,如Sigmoid和softmax。作者认为注意力方法的实质并不在于归一化注意力图,而在于根据输入特征在不同尺度上的可适应性输出。通过和谐地集成卷积和自注意力,作者的方法是全面的,考虑了局部上下文、广泛的感受野、线性复杂度和跨尺度的动态性。

由于深度神经网络中的不同通道通常对应着神经网络中的独特物体,这种适配性深度膨胀卷积、扩大的深度卷积和最后一个是卷积,在视觉任务中变得必不可少。如图3所示,跨尺度交互模块的详细结构和复杂性。

IV Experimental Results

本节详细介绍了训练过程、作者在实验评估中使用的指标以及实验结果的详细分析。

训练过程

在本研究中,作者在配备NVIDIA Tesla V100 GPU的PyTorch上实现所提出的 method,不使用数据增强,批量大小为24。在400个epoch中,所有模型以的学习率和的衰减率进行训练。模型的权重初始化使用标准正态分布,从开始就稳定,且确保重量波动较小。此外,如果在训练过程中的十连续epoch中,验证性能未发生变化,训练过程停止。在训练和验证数据集上,优化算法逐渐减小损失值,并在训练过程中最终收敛到最优解。训练过程中没有出现不稳定性。

Dataset

所提出的算法在Synapse多器官分割数据集[36]上进行了评估。数据集 [36]包含37个器官的临床CT图像,总共3779个轴向对比增强的腹部CT扫描。解释者对每个实例中的13个器官进行了标注,包括脾脏、右侧肾脏、左侧肾脏、胆囊、食管、肝脏、胃、主动脉、下腔静脉、门静脉、脾静脉、胰腺、左侧肾上腺和右侧肾上腺。每个CT扫描都通过对比增强获得,其体积范围在85到198个切片、512x512像素之间。

Quantitative and qualitative results

表1展示了作者提出的方法与几种基准方法的比较分析。这其中包括作者的初步模型( Baseline )和一些性能优异的最新架构。

picture.image

为了全面了解作者方法的有效性,作者评估了三个不同的 Baseline :Baseline (Baseline): 这是作者方法的基础,排除了双注意力(dual attention)和ISIM(Integrated Splanchnic Systematics and Representations)的增强。相反,它只在每个 Transformer 块中使用一个高效的注意力模型。提出的方法(没有ISIM): 方法的演变版本,这个版本结合了通道(channel)和空间(spatial)注意力。正如作者描述的那样,这个结合了两种注意力机制的效率。提出的方法(Proposed Method): 作者的整体方法,利用所有特征,包括ISIM。

作者模型的连续增强无疑提升了它的性能。引入双注意力和然后是ISIM,这在实践中验证了作者策略解决医学图像分割挑战的有效性。

与最先进的模型进行全面比较,作者的方法无疑具有优势。特别的是,作者的提出的Dice相似系数(DSC)令人印象深刻地达到了82.16%,超过了像HiFormer这样的强大竞争者,其值为80.39%。作者的模型的一个显著特点是,它可以更准确地划分更细微的解剖结构。ISIM的集成显著增强了这一点。这在胆囊的分割中得到体现,作者的技术产生了71.48%的得分,超越了像TransDeepLab的69.16%。同样,由于其体积复杂,胰腺的分割一直是一个传统的困难。然而,作者的方法取得了65.02%的得分,超越了像FFUNet-trans的64.17%。

对于分割更明显的器官,作者的模型仍然无与伦比。左肾(L)和右肾(R)分别获得了85.45%和80.85%的分数。值得注意的是,肝的分割,作者的方法获得了94.85%的分数,几乎等于HiFormer的94.61%。此外,在分割脾脏时,作者的模型以91.62%的分数略微超过了作者自己的 Baseline (91.05%)。

Ablation Study

作者技术的一个重要特征是其能够灵活捕捉长程依赖性。与其他模型相比,对于更大的器官(如肝脏)的优越预测能力证明了这一点。该模型在预测范围内充分容纳这些长程依赖性的能力非常重要。

此外,作者观察到对于较小的器官(如主动脉),U-Net模型往往在其他基于Transformer的方法中表现优越。这突显了在预测较小实体时局部特征表示的重要作用,以及将此信息整合到预测矩阵中的迫切需求。

为了进一步强调模型利用远程信息的能力,作者有必要注意到作者的方法在分割小和大器官方面的熟练程度。这需要较大的感受野大小以实现目标预测的精度。作者在图4中通过分别针对两种器官类型的类激活图进一步阐述了这一点,揭示了作者的模型在识别局部模式方面增强了能力,从而实现了精细的 segmentation。

picture.image

V Conclusion

在本研究中,作者提出并评估了一种用于医学图像分割的新架构,该架构以高效和增强的注意力机制为基础,并结合了 ISIM 的独特功能。通过对模型进行分步 Baseline 评估的有序方法无疑明确了每个组件的贡献,特别强调 ISIM 在提升整体性能方面的变革性作用。

不仅超越了作者的基本模型,作者还提出了一个与顶级当代架构竞争并与许多实例超过他们性能的方法。考虑到其显著的准确性和效率,作者的模型具有重要的临床价值,将自己定位为医疗保健从业者在诊断和治疗工作中的宝贵助手。

这个完美地结合了开创性研究和实际结果的方法不仅突出了作者的方法的重要性,而且为未来的医学图像创新设置了一条有前景的轨迹。

点击上方卡片,关注 「AI视界引擎」 公众号

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
字节跳动 XR 技术的探索与实践
火山引擎开发者社区技术大讲堂第二期邀请到了火山引擎 XR 技术负责人和火山引擎创作 CV 技术负责人,为大家分享字节跳动积累的前沿视觉技术及内外部的应用实践,揭秘现代炫酷的视觉效果背后的技术实现。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论