MIA 2024 | VSmTrans:一种用于医学图像分割的融合自注意力和卷积的混合范式

科技   2024-10-19 11:53   中国香港  

点击下方“ReadingPapers”卡片,每天获取顶刊论文解读

论文信息

题目:VSmTrans: A hybrid paradigm integrating self-attention and convolution for 3D medical image segmentation

VSmTrans:一种融合自注意力和卷积的混合范式用于3D医学图像分割

作者:Tiange Liu, Qingze Bai, Drew A. Torigian, Yubing Tong, Jayaram K. Udupa

源码链接:https://github.com/qingze-bai/VSmTrans

论文创新点


  1. 提出一种新颖的混合Transformer骨干网络:作者提出了一种名为Variable-Shape Mixed Transformer (VSmTrans)的混合范式,用于3D医学图像分割。这种混合范式能够紧密整合自注意力卷积,以利用这两种范式的优势。

    1. 设计了有效的自注意力机制VSW-MSA:作者设计了一种名为Variable-Shape Window Multi-head Self-attention (VSW-MSA)的新型自注意力机制。这种机制可以快速扩展接受域,并在全局和局部信息收集之间实现良好的平衡,而无需额外的计算成本。

    1. 混合Transformer模块的创新设计:在新的混合模块中,CNN不仅仅是引入归纳偏置的独立路径,而是将并行卷积增强模块嵌入到Transformer模块中,可以享受大接受域和强大归纳偏置的好处。

    1. 在多个公共医学图像数据集上的广泛实验:作者在AMOS CT数据集和BraTS2021 MRI数据集上进行了广泛的实验,验证了所提出方法的有效性。实验结果表明,该方法在性能上具有竞争力,甚至超过了其他一些最先进的方法。

    1. 消融实验验证了混合机制的有效性:通过一系列消融实验,作者验证了所提出的混合机制能够充分利用自注意力卷积模块,有效平衡大接受域与局部归纳偏见,从而实现准确的分割结果,尤其是在物体边界上。

    关键词图像分割、3D医学图像、Transformer、卷积

    摘要

    目的:近期,视觉Transformer因其出色的全局表征学习能力,在性能上与CNNs不相上下。然而,在将它们应用于3D图像分割时存在两个主要挑战:i) 由于3D医学图像的庞大尺寸,由于巨大的计算成本,很难捕获全面的全局信息。ii) Transformer中局部归纳偏置的不足影响了分割细节特征的能力,例如模糊和微妙定义的边界。因此,要将视觉Transformer机制应用于医学图像分割领域,需要充分克服上述挑战。方法:作者提出了一种名为Variable-Shape Mixed Transformer(VSmTrans)的混合范式,它整合了自注意力和卷积,并能够享受自注意力机制带来的复杂关系自由学习的益处以及卷积带来的局部先验知识。具体来说,作者设计了一种Variable-Shape自注意力机制,它可以在不增加额外计算成本的情况下快速扩展接受域,并在全局意识和局部细节之间实现良好的平衡。此外,平行卷积范式引入了强大的局部归纳偏置,以促进挖掘细节的能力。同时,一对可学习的参数可以自动调整上述两种范式的的重要性。作者在两种公共医学图像数据集上进行了广泛的实验,这些数据集具有不同的模态:AMOS CT数据集和BraTS2021 MRI数据集。结果:我们的方法在这些数据集上实现了88.3%和89.7%的最佳平均Dice分数,这优于以前的基于Swin Transformer和基于CNN的架构。作者还进行了一系列消融实验,以验证所提出的混合机制及其组成部分的效率,并探索VSmTrans中那些关键参数的有效性。结论:所提出的用于3D医学图像分割的混合Transformer骨干网络可以紧密整合自注意力和卷积,以利用这两种范式的优势。实验结果证明了我们方法的优越性,与其他最先进的方法相比。混合范式似乎最适合医学图像分割领域。消融实验还表明,所提出的混合机制可以有效平衡大接受域和局部归纳偏置,从而实现高精度的分割结果,特别是在捕获细节方面。

    方法

    3.1. 总体架构

    在本文中,我们设计了一个U形编码器-解码器架构,如许多其他基于Transformer的方法(Dosovitskiy等人,2020;Hatamizadeh等人,2022)所示,以验证新提出的变压器,称为Variable-Shape Mixed Transformer(VSmTrans)。如图2所示,该网络使用VSmTrans块作为编码器,解码器主要由常规卷积块组成。输出的每个编码器层还通过跳跃连接传输到解码器。具体来说,对于输入的CT补丁,尺寸为H × W × D,我们使用不重叠的卷积获得尺寸为H/2 × W/2 × D/2的特征图。线性嵌入层应用于将每个特征图投影到C个特征通道。然后,将投影的特征输入到VSmTrans中,它由四个阶段组成。在每对相邻阶段之间,有一个补丁合并层,用于降低分辨率并为特征图的通道翻倍。ResBlock(He等人,2016)用作解码器的主要成分,以上采样特征图,直到它们达到原始分辨率。
    作为编码器的主要组成部分,VSmTrans是基于作者新设计的可变形状混合窗口多头自注意力(VSmW-MSA)(第3.3节)。在VSmW-MSA中,作者将并行卷积嵌入到所提出的可变形状混合窗口多头自注意力中,以增加归纳偏见。值得注意的是,每个阶段由两个块组成,我们在第二块中使用移位窗口划分方法(Lee等人,2022;Liu等人,2021)以进一步扩大接受域。因此,两个连续的VSmTrans块可以总结如下:
    其中VSmSW-MSA是具有移位窗口的VSmW-MSA。多层感知器(MLP)和层归一化(LN)分别是多层感知器和层归一化。分别表示在同一给定层中VSmW-MSA模块和MLP模块的输出。

    3.2. 可变形状窗口多头自注意力

    视觉Transformer具有来自自注意力机制的高复杂性的固有特征。这导致了在计算复杂性和性能之间寻找平衡的挑战。大多数现有方法通过堆叠许多具有正方形注意力窗口的Transformer块(Xie等人,2021b;Zhou等人,2021;Tang等人,2022)或应用各种条纹注意力窗口(Ho等人,2019;Dong等人,2022)来解决这个问题。然而,正方形注意力窗口无法迅速扩展接受域,并可能在多层堆叠过程中丢失部分注意力特征,而条纹注意力窗口则更倾向于捕获全局信息而不是局部特征。我们设计了可变形状窗口多头自注意力(VSW-MSA),以扩展接受域并考虑局部细节。关键的洞察是,应用基于多个不同形状窗口的自注意力可以迅速扩展接受域,而无需额外的计算负担。在这些窗口中,包括立方体和条纹窗口,以实现全局和局部信息之间的良好平衡。我们的实验(第4.4节)表明,这种方法实现了强大的形状感知。
    如图3所示,输入特征图将被平均分成四组自注意力窗口。这四组分别对应条纹轴向、条纹矢状、条纹冠状和局部窗口。具体来说,假设是特征图空间,其中包含C个通道,每个通道的特征图尺寸为。然后输入特征图集,使得每个组具有个通道。在每个组中,特征图将在多头机制下进一步排列为K个头。我们共同使用三个条纹窗口形状,以实现大规模接受域,并使用立方体窗口形状来学习局部信息。在自注意力计算之后,对所有组的输出应用连接,以捕获局部和全局信息。正式地,自注意力路径可以定义为:
    其中表示相应自注意力操作的输出。分别表示第i组的查询、键和值的投影矩阵。
    为了进一步捕获全局信息,类似于先前的工作(Liu等人,2021),我们引入了移位窗口划分方法,为每个组的自注意力建立跨窗口连接。每个窗口区域将在两个连续块之间移动像素,忽略无效计算。

    3.3. 整合自注意力和卷积

    由于自注意力机制不包含对交互局部性的内置归纳先验,许多先前的工作(Dosovitskiy等人2020; Guo等人,2021)使用位置编码来学习位置关系。然而,仅通过位置编码学习归纳偏见信息需要大量的训练数据,这对于医学图像分析领域来说是非常具有挑战性的。卷积自然具有内置的归纳偏见,如局部性和平移不变性,这可以有效学习目标对象的结构信息和详细纹理。因此,将自注意力和卷积整合到一个范式中是一个有前景的解决方案,它不需要在大型数据集上进行预训练,尤其是在医学图像分析领域。因此,我们提出了一个新颖的混合模块,称为可变形状混合窗口多头自注意力(VSmW-MSA),它整合了VSW-MSA模块和轻量级卷积模块,以利用和统一两者的优势。
    在图1中,展示了两个代表性的Transformer块的结构和我们的结构,以显示差异。具体来说,内部位置编码机制(Chu等人,2021;Liu等人,2021)将位置信息纳入每个自注意力窗口中(图1(a))。在提出的方法中,我们在自注意力操作旁边添加了一个额外的轻量级卷积模块(图1(c))。VSmW-MSA的详细组件也在图4中显示,其中卷积路径在投影值特征图上工作,而不是在自注意力窗口内,以便可以强有力地提取特征图的结构信息。轻量级卷积模块由一个3×3×3卷积、一个1×1×1卷积和一个LN层组成,两个卷积之间有一个GELU激活函数(Hendrycks和Gimpel,2016)以增加非线性。卷积路径可以表示为:
    如图4所示,我们提出的混合模块整合了可变形状窗口多头自注意力和轻量级卷积模块。其中,自注意力路径将投影特征图分成四个组自注意力窗口,以建立长期依赖性。卷积路径可以有效地引入结构和位置信息。这两个并行路径共享投影特征图。同时,考虑到自注意力和卷积在不同阶段可能扮演不同的角色,自注意力和卷积路径的权重由两个可学习的参数α和β控制,可以表示为:

    3.4. 损失函数

    我们使用Dice损失和交叉熵损失的组合作为优化目标。正式地,损失可以定义为:
    其中分别代表类别的预测和真实情况,分别表示Dice损失和交叉熵损失的权衡权重值。

    实验和结果

    声明

    本文内容为论文学习收获分享,受限于知识能力,本文对原文的理解可能存在偏差,最终内容以原论文为准。本文信息旨在传播和学术交流,其内容由作者负责,不代表本号观点。文中作品文字、图片等如涉及内容、版权和其他问题,请及时与我们联系,我们将在第一时间回复并处理。

    小白学视觉
    哈工大在读博士的公众号,《OpenCV 4快速入门》的作者,面向初学者介绍计算机视觉基础知识、OpenCV使用、SLAM技术,深度学习等内容。
     最新文章