[医疗 AI ]
3D TransUNet:通过 Vision Transformer 推进医学图像分割’
论文地址 - https://arxiv.org/pdf/2310.07781
0. 摘要
医学图像分割在推进医疗保健系统的疾病诊断和治疗计划中起着至关重要的作用。U 形架构,俗称 U-Net,已被证明在各种医学图像分割任务中非常成功。然而,U-Net 基于卷积的操作本身限制了其有效建模远程依赖关系的能力。为了解决这些限制,研究人员转向了以其全局自我注意机制而闻名的 Transformer 作为替代架构。一个流行的网络是我们以前的 TransUNet,它利用 Transformers 的自我关注来补充 U-Net 的本地化信息与全局背景。在本文中,我们通过构建最先进的 nnU-Net 架构,将 2D TransUNet 架构扩展到 3D 网络,并充分探索 Transformers 在编码器和解码器设计中的潜力。我们介绍了两个关键组件:1) 一个 Transformer 编码器,它对来自卷积神经网络 (CNN) 特征图的图像补丁进行标记,从而能够提取全局上下文,以及 2) 一个 Transformer 解码器,它通过利用候选提案和 U-Net 特征之间的交叉注意力来自适应地优化候选区域。我们的调查表明,不同的医疗任务受益于不同的建筑设计。Transformer 编码器在多器官分割中表现出色,其中器官之间的关系至关重要。另一方面,Transformer 解码器被证明更有利于处理小且具有挑战性的分割目标,例如肿瘤分割。广泛的实验展示了集成基于 Transformer 的编码器和解码器集成到 U 形医疗图像分割架构中。TransUNet 在各种医疗应用中优于竞争对手,包括多器官分割、胰腺肿瘤分割和肝血管分割。它明显超越了 BrasTS2021 挑战赛中的顶级解决方案。代码和模型可在 https://github.com/Beckschen/ 3D-TransUNet 上获得。
1. 引言
卷积神经网络 (CNN),尤其是全卷积网络 (FCN) [1],在医学图像分割领域已经崛起。在他们的各种迭代中,U-Net 模型 [2] 的特点是其对称编码器-解码器设计,并通过跳过连接进行了增强以改善细节保留,是许多研究人员的首选。基于这种方法,在各种医学成像任务中取得了显著进展。这些进步包括磁共振 (MR) 成像中的心脏分割 [3]、使用计算机断层扫描 (CT) 扫描 [4]-[7] 的器官描绘以及结肠镜检查记录中的息肉分割 [8]。
尽管 CNN 具有无与伦比的表示能力,但由于卷积运算的固有局部性,它们在建模远程关系时经常步履蹒跚。这种限制在患者间质地、形状和大小差异较大的病例中尤为明显。认识到这一局限性,研究界越来越被 Transformers 所吸引,由于它们在捕捉全球背景方面的天生能力,这些模型完全建立在注意力机制之上 [9]。在医学图像分割领域,我们之前与 TransUNet [10] 的合作证明了变压器的潜力。然而,我们研究中的一个关键观察表明,简单地用 Transformer 替换 CNN 编码器可能会导致次优结果。Transformer 将输入作为 1D 序列处理,并优先考虑全局上下文建模,从而无意中产生低分辨率的特征。直接对此类特征进行上采样无法重新引入丢失的粒度。相比之下,结合 CNN 和 Transformer 编码器的混合方法似乎更有前途。它有效地利用了 CNN 的高分辨率空间细节,同时还受益于 Transformers 提供的全局环境。
在这项研究中,我们将原始的 2D TransUNet 架构扩展到 3D 配置,更深入地研究了 Transformer 在编码和解码过程中的战略性整合。这一飞跃植根于 nnU-Net 框架的实力,其愿景是超越其既定标准。我们的 3D TransUNet 通过两个主要机制展开:首先,Transformer 编码器将 CNN 特征图中的图像块标记化,允许将全局自聚焦特征与从编码路径跳过的高分辨率 CNN 特征无缝融合,以实现精确定位。其次,Transformer Decoder 将传统的每像素分割重新定义为掩码分类,将预测候选者构建为可学习的查询。具体来说,这些查询通过协同交叉注意力与局部多尺度 CNN 特征来逐步完善。此外,我们在 Transformer 解码器中引入了粗到细的注意力细化,对于每个分割类,使用专注于预测前景的注意力机制精心细化初始候选集,确保每个迭代细化为后续为后续设定新标准,最终不断提高分割精度。
通过将 Transformer 集成到类似 U-Net 架构的编码器和解码器组件中,我们证明了我们的设计允许框架保留 Transformer 的优势,同时增强医学图像分割。有趣的是,多器官分割在很大程度上依赖于全局上下文信息(例如不同腹部器官之间的相互作用),倾向于使用 Transformer 编码器设计。相反,像小目标分割这样的任务,如肿瘤检测,通常从 Transformer 解码器设计中受益更多。我们广泛的实验表明,与各种医学图像分割任务中的竞争方法相比,我们的方法具有卓越的性能。总而言之,我们的贡献可以总结如下:
我们引入了一个以 Transformer 为中心的医学图像分割框架,将自我注意整合到序列到序列预测上下文中,适用于 2D 和 3D 医学图像分割任务。我们彻底研究了将视觉变压器集成到 U 形分割架构的编码器和解码器中的影响,为定制设计提供见解,以应对不同的医学图像分割挑战。我们在各种医学图像分割任务上取得了最先进的结果,并发布了我们的代码库以鼓励进一步探索将 Transformer 应用于医疗应用。网络结构:
2. 相关工作
将 CNN 与自我注意机制相结合
各种研究试图通过基于特征图对所有像素的全局交互进行建模,将自我注意机制集成到 CNN 中。例如,Wang 等人设计了一个非局部算子,它可以插入多个中间卷积层 [11]。基于编码器解码器 U 形架构,Schlemper 等人 [12] 提出Additive Attention Gate 模块,这些模块被集成到 Skip-Connections 中。与这些方法不同,我们使用 Transformers 在我们的方法中嵌入全局自我注意。
Transformers
Transformers 最早由 [9] 提出用于机器翻译,并在许多 NLP 任务中建立了最先进的技术。为了使 Transformer 也适用于计算机视觉任务,我们进行了多项修改。例如,Parmar et al. [13] 仅将自我注意应用于每个查询像素的局部邻域,而不是全局。Child 等人 [14] 提出了稀疏变压器,它采用可扩展的近似值来表示全局自我注意。最近,Vision Transformer (ViT) [15] 通过直接将具有全局自注意的 Transformer 应用于全尺寸图像,实现了 ImageNet 分类的现状。据我们所知,我们最初于 2021 年提出的 TransUNet 是第一个基于 Transformer 的医学图像分割框架,它建立在非常成功的 ViT 之上。基于 TransUNet,nnformer [16] 通过将卷积与自我注意交织来改进该方法。在另一项发展中,CoTR [17] 在 Transformer 编码器中提供了更有效的自注意力。另一种 Transformer 架构是 Swin Transformer [18],这是一种分层视觉转换器,它使用移位窗口来捕获本地和全局信息,以实现高效和可扩展的视觉处理。随后的模型,如 SwinUNet [19] 和 SwinUETR [20] 已经被开发用于医学图像分割。
用于分割的掩码分类
DETR [21] 是第一个使用 Transformer 作为解码器的工作,具有用于对象检测的可学习对象查询。在变压器 [22]–[27] 最新进展的背景下,出现了一种被称为掩模变压器的新型变体。此变体通过采用一组查询嵌入来表示对象及其关联的掩码,从而引入了分段预测。Wang et al. [23] 首先开发了带有内存嵌入的掩码 Transformer,Cheng et al. [24] 以 DETR [21] 的方式进一步表述了查询更新。掩码转换器的核心是解码器,它负责将对象查询作为输入进行处理,并逐步将其转换为掩码嵌入向量 [24]–[27]。此过程使模型能够有效地处理分割任务并产生准确的结果。
3. 方法
给定 3D 医学图像(例如 CT/MR 扫描)x ∈ R D×H×W×C,空间分辨率为 D×H ×W 和 C 通道数。我们的目标是预测大小为 D×H ×W 的相应像素级标签图。最常见的方法是直接训练 CNN(例如 U-Net)首先将图像编码为高级特征表示,然后将其解码回完整的空间分辨率。我们的方法与传统方法不同,它采用 Transformer 彻底探索了标准 U 形分割架构的编码器和解码器阶段使用的注意力机制。在第 III-A 节中,我们深入探讨了 Transformers 的直接应用,用于对分割图像块的特征表示进行编码。在此之后,在 Section III-B 中,我们详细阐述了基于查询的 Transformer 的实现,它用作我们的解码器。然后,第 III-C 节介绍了 TransUNet 的详细架构。在 Section III-B 中,我们详细阐述了基于查询的 Transformer 的实现,它用作我们的解码器。然后,第 III-C 节介绍了 TransUNet 的详细架构。
A. Transformer 作为编码器
图像序列化
按照 [15],我们首先通过将输入 x 重塑为一系列扁平化的 3D 补丁(patches)来执行分词化 {x p i ∈ R P 3 ·C |i = 1, …, N},其中每个色块的大小为 P × P × P 和 N = DHW / (PPP) 是图像色块的数量(即输入序列长度)。
补丁嵌入 我们使用可训练线性投影将矢量化补丁 X 映射到潜在的 D 维嵌入空间。为了对 patch 空间信息进行编码,我们学习了特定的位置嵌入,这些 embedding 被添加到 patch 嵌入中以保留位置信息,如下所示:
其中 E ∈ R (P 3 ·C)×denc 是补丁嵌入投影,E pos ∈ R N×denc 表示位置嵌入。
每个 Transformer 层由多头自注意力 (MSA) 和多层感知器 (MLP) 模块组成(方程 (2)(3))。因此,第 l 层的输出可以写成如下:
其中 LN(·) 表示层归一化运算符,zl 是编码的图像表示。
B. Transformer 作为解码器
1) 粗略的候选估计:传统方法(如 U-Net)主要将医学图像分割视为每像素的分类任务。在这个范式中,每个像素被分类为可能的 K 类别之一,通常是通过训练具有每个像素交叉熵(或负对数似然)损失的分割模型来实现的。
在本文中,我们的方法没有考虑单个像素,而是将医学图像分割视为掩码分类问题。我们引入了 “organ query” 的概念,这是一个表示图像中每个器官的 ddec 维特征向量。使用一组预定义的 N 个器官查询,我们对包含 K 个分割类的图像的目标是将图像分离为 N 个不同的候选区域。随后,我们的目标是为每个区域分配相应的器官标签。值得注意的是,正如先前的研究所示,N 的值不必与类的数量保持一致 [22]。事实上,我们有意将 N 设置为明显大于 K,以最大限度地降低假阴性的风险。假设对象查询的维度是 ddec,粗略预测的分割图可以通过初始器官查询 P0 ∈ R N×ddec 和 U-Net 最后一个块特征 F ∈ R D×H×W×ddec 的嵌入之间的点积来计算:
其中 g(·) 是 sigmoid 激活,后跟阈值设置为 0.5 的硬阈值操作。如方程 (4) 所示,P 可以看作是标准 U-Net 分割头中的 1×1×1 卷积核
2) Transformer 解码器:
图 1 说明了我们定制的 3D Transformer 解码器,专门用于优化器官查询,从而增强粗略预测 Z 0。类似于 Transformer 编码器中的结构(详见第 III-A 节),每一层的自注意力机制(即 MSA 块)将使 Transformer 解码器能够全面参与图像特征并捕获器官查询的相互关系。认识到中间 CNN 特征的丰富定位,这补充了 Transformer 的全局图像上下文,我们通过积分优化了每个解码器层中的器官查询。
我们的策略包括 CNN 解码器和 Transformer 解码器的并发训练。在第 t 层中,精炼的器官查询表示为 Pt ∈ R N×ddec 。同时,一个中间的 U-Net 特征被映射到一个表示为 F 的 ddec 维特征空间,以促进交叉注意力计算。值得注意的是,当上采样块的数量与 Transformer 解码器层一致时,多尺度 CNN 特征可以投影到特征空间 F ∈ R (DtHtWt)×ddec 中,其中 Dt、Ht 和 Wt 指定第 t 个上采样块处特征图的空间维度。移至第 t+1- 层,使用 cross-attention 更新器官查询,如下所示:
其中,第 t 个查询特征经过线性投影,以使用权重矩阵 wq ∈ R ddec×dq 形成对下一层的查询。U-Net 特征 F 同样使用参数权重矩阵 wk ∈ R ddec×dk 和 wv ∈ R ddec×dv 转换为键和值。请注意,根据以前的研究 [25],残差路径用于更新 P。接下来,我们将介绍如何合并从粗到细的注意力细化,以进一步提高分割结果的准确性。
3) 粗到细的注意力优化:
粗细细化在医学图像分割中的价值是公认的,特别是对于小目标分割 [4],[28],[29]。该技术从初始阶段就采用粗略掩码来指导后续的优化。在这里,为了在 Transformer 解码器中集成一个无缝的粗到精细的细化过程,我们集成了一个掩码注意力模块 [25]。此增强功能旨在根据之前对每个类别的粗略预测,将交叉注意力置于前景区域内,以减少背景噪声并更好地关注感兴趣区域。这种改进的注意力图迭代地帮助了后续的、更精细的细分阶段。
具体来说,我们首先将器官查询和粗略级掩码预测分别设置为 P0 和 Z 0(基于方程 (4)),然后开始迭代细化过程。在第 t 次迭代中,使用当前器官查询特征 Pt 和粗略预测 Z t ,我们计算掩蔽的 crossattention,从而为后续迭代优化 Pt+1。这计算将现有的粗略预测 Z t 合并到亲和矩阵中,如 方程 5 中所述
其中 i, j, s 是坐标索引。此公式将交叉注意力机制限制为仅关注前台,从而使所有其他区域无效。通过迭代更新 organ 查询和相应的掩码预测,我们的 Transformer 解码器系统地优化了多次迭代的分割结果。算法 1 中概述了此迭代过程的详细说明。细化周期一直持续到迭代计数 t 达到最大阈值 T,这相当于 Transformer 解码器中的层数。
精细分段解码
在最后一次迭代之后,更新后的器官查询 PT 可以按照方程 (4) 通过具有 U-Net 最后一个块特征 F 的点积解码回最终的精细二元化分割图 Z T。为了将每个二值掩码与一个语义类相关联,我们进一步使用权重矩阵为 wf c ∈ R d×K 的线性层,将精炼的器官嵌入 PT 投影到输出类 logits O ∈ R N×K。形式上,我们有:
其中 k 是标签索引。与优化预测掩码 Z T 关联的最终类标签是 yˆ ∈ R N 。
C TransUNet
我们的 3D TransUNet 建立在最先进的 nnUNet 架构之上,旨在超越其既定标准。我们在图 1 中说明了拟议的 TransUNet 的整体框架。我们使用三种不同的配置来实例化我们的 3D TransUNet:
1) 仅编码器
在 CNN 首次用作特征提取器以生成输入特征图的地方,采用 CNN-Transformer 混合编码器。补丁嵌入应用于特征补丁,而不是 Raw 图像。对于解码阶段,我们使用标准的 U-Net 解码器。我们选择这种设计是因为 1) 它允许我们在解码路径中利用中间高分辨率 CNN 特征图;2) 我们发现混合 CNN-Transformer 编码器的性能比简单地使用纯 Transformer 作为编码器要好。仅编码器模型将使用由像素交叉熵损失和骰子损失组成的混合分割损失进行训练。
2) 仅解码器:
在这种配置中,我们在编码阶段使用传统的 CNN 编码器。至于解码阶段,我们在分割模型中使用 CNN-Transformer 混合解码器。最初,器官查询 P 设置为零。在被 Transformer 解码器处理之前,它们按照方程 (1) 通过可学习的位置嵌入进行增强。然后,如前所述,在第 III-B 节中,P 将根据 U-Net 特征逐渐细化,并被解码回全分辨率分割图。我们遵循以前的工作 [21]、[23] ,用匈牙利匹配损失来训练网络,以更新整个解码层的器官查询。此损失旨在匹配预测和真实区段之间的对。它结合了每个分段预测的像素级分类损失和二进制掩码损失:
其中像素级分类损失 Lce 和 Ldice 分别表示二进制交叉熵损失和 dice 损失 [30]。分类损失 Lcls 由每个候选区域的交叉熵损失实例化。λ0 和 λ1 是用于平衡每像素分割损失和掩码分类损失的超参数.
我们还采用深度监督,将训练损失应用于 TransUNet 解码器每个阶段的输出。
3) 编码器 + 解码器:在这里,我们将 Transformer 编码器和 Transformer 解码器集成到 3D nnUNet 模型中。然后类似于纯解码器模型,这里我们也使用 Hungarian matching loss 来训练整个网络。
最后,我们想指出的是,我们的方法虽然建立在 3D nnU-Net 之上,但可以通过简单地切换主干模型并将所有操作简化回 2D 来轻松修改以适应 2D 任务
IV. 实验和讨论
A. 数据集和评估
Synapse 多器官分割数据集 1 .我们在 MICCAI 2015 多图谱腹部标记挑战赛中使用了 30 次腹部 CT 扫描,总共有 3779 张轴向对比增强腹部临床 CT 图像。
每个 CT 体积由 85 ∼ 198 个 512 × 512 像素的切片组成,体素空间分辨率为 ([0.54 ∼ 0.54] ×[0.98 ∼ 0.98] × [2.5 ∼ 5.0])mm3 .根据 [31],我们报告了 8 个腹部器官(主动脉、胆囊、脾脏、左肾、右肾、肝脏、胰腺、脾脏、胃)的平均 DSC 和平均 Hausdorff 距离 (HD),随机拆分 18 个训练案例(2212 个轴向切片)和 12 个案例进行验证,遵循 [10] 中的拆分设置。
BraTS2021 脑肿瘤分割挑战赛
.BrasTS2021 Challenge 是最新、规模最大的脑肿瘤分割数据集。向参与者提供 1251 例多参数 MRI 扫描并带有分割标签。MRI 扫描有 4 种造影剂:天然 T1 加权图像、造影剂后 T1 加权 (T1Gd)、T2 加权和 T2 液体衰减反转恢复 (T2-FLAIR)。注释由 1 到 4 名评估员手动执行,并得到经验丰富的神经放射科医生的最终批准。标记包括 GD 增强肿瘤 (ET) 、瘤周水肿/侵袭组织 (ED) 和坏死肿瘤核心 (NCR) 的区域。所有 MRI 扫描均通过共同配准到同一解剖模板、插值至各向同性 1mm3 分辨率和颅骨剥离进行预处理。所有 MRI 扫描和相关标签的图像大小为 240×240×155。在我们的实验中,我们应用了 5 倍交叉验证,并使用与 BraTS2021 挑战中的 1 号解决方案 [32] 相同的数据分割
医疗分割十项全能 (MSD)
HepaticVessel3 .MSD HepaticVessel 是医疗分割十项全能 [33] 的一项任务,包括 443 次门静脉期 CT 扫描,这些扫描来自各种原发性和转移性肝肿瘤患者。相应的目标 ROI 是肝脏内的血管和肿瘤。这个数据集是选择该基因是由于与异质性肿瘤相邻的肝血管的管状和连接性。我们应用 5 倍交叉验证来评估此数据集上的方法。
大规模胰腺质量数据集。我们的静脉期 2930 CT 扫描数据集是从一家高容量的美国医院收集的。据我们所知,它是规模最大的胰腺肿瘤 CT 数据集之一。胰腺导管腺癌 (PDAC) 是所有胰腺异常中最优先的,5 年生存率约为 10%,是最常见的类型(约占所有胰腺癌的 90%)。标签包括 Pancreas、PDAC 和 Cyst。该数据集随机分为 2123 次 CT 扫描的训练和 807 次 CT 扫描的测试数据集。模型验证在训练集的子集上进行。训练集包括 1017 个 PDAC、462 个囊肿和 644 个正常胰腺。测试集包括 506 个 PDAC、271 个囊肿、300 个正常胰腺。评估指标包括骰子分数、敏感性和特异性,遵循 [34] 中的标准。
BraTS2023 脑转移数据集。BraTS-MET 数据集 [35] 在开发用于检测和分割脑转移的先进算法方面发挥着至关重要的作用,目的是确保无缝集成到临床实践中。该数据集包括按照标准临床方案从多个机构收集的各种未经治疗的脑转移 mpMRI 扫描 [36]–[38]。值得一提的是,我们的工作专门利用了数据集的 BraTS2023-MET 版本。目前,它由 238 个案例组成,我们在分析中采用了默认的 5 折交叉验证方法
B. 实施细节
训练
在我们所有的 3D 实验中,我们都遵循 nnUNet 规定的数据增强程序,以增强我们训练数据集的多样性。为了促进有效的训练,我们使用 1 个 Nvidia RTX 8000 GPU 采用 2 个批量。表 I 中详细介绍了我们的实现细节,包括架构超参数、训练配置和为各种数据集定制的数据增强技术等关键方面。我们试验了 1 层和 12 层 ViT 来实现 Transformer 编码器。具体来说,12 层 ViT 模型在 ImageNet21k 数据集 [39] 上进行了预训练,并增加了 LayerScale [40]。潜在维度 denc 和 ddec 分别设置为 768 和 192。为了计算匈牙利匹配损失,λ0 和 λ1 设置为 0.7 和 0.3。与 nnU-Net 开创性的自适应架构设计相一致,我们的 TransUNet 表现出根据其处理的数据特性量身定制的适应性。与 nnU-Net 类似,我们的骨干架构本质上是针对不同医学数据集的自适应性的,特别是在确定关键因素方面,例如下采样层的数量和每个阶段的通道分配。详情见表 I。值得注意的是,我们的 TransUNet 保持与 2D 图像数据的兼容性,只需最少的架构修改即可适应各种成像场景.
测试
给定 CT/MR 扫描,我们以滑动窗口方式进行推理。通过利用所有补丁的聚合,P 我们为位置 (i, j, s) 的体素分配一个概率向量:N n=1(Z T n,ijs) ∈ R K,后跟一个 argmax 以获得硬预测。
C. 三种配置的分析研究
我们进行了各种消融研究,以彻底评估 TransUNet 框架下提出的三种配置,即仅编码器、仅解码器和编码器+Decoder.
为了评估 Transformer 编码器对 CNN 编码器的有效性,以及对解码器的有效性,我们对 3D TransUNet 模型进行了全面的实验。此评估包括仅编码器和仅解码器配置,并将它们与基线 nnU-Net 进行比较跨多器官分割和肝血管肿瘤分割任务。表 II 和表 III 总结了 3D TransUNet 上的这些比较结果。
对于多器官分割,虽然仅解码器设计表现出适度的性能增强(87.63% 对 87.33%),但仅编码器配置,尤其是当使用使用 ImageNet 的预训练权重初始化的 12 层 ViT 编码器时,Dice 分数显着提高了 0.8%,达到 88.11%。当采用 2D U-Net 主干网时,这一优势变得更加明显,在 224×224 和 512×512 的分辨率下,性能分别提高了 2.80% 和 3.02%,如表 VIII 所示。
至于血管肿瘤分割,仅编码器设计的性能改进虽然存在,但仍然相对微妙。1 层和 12 层 ViT 编码器都产生了相当的结果(66.30% 和 66.35%),略高于基线 nnU-Net 的 66.04% 分数。相比之下,仅解码器配置表现出显着的增量,记录了 1.63% 的增益(67.67% 对 66.04%)
总而言之,我们的结果表明,仅编码器设计在多器官分割中蓬勃发展,而仅解码器配置更擅长血管肿瘤分割。这种区别与 Transformer 编码器的固有优势相一致,后者捕获全局上下文信息,例如各种目标之间的复杂关系,使其对多器官任务特别有效。相反,Transformer 解码器专为优化小而硬的目标而量身定制,非常适合肿瘤分割。然而,Transformer 编码器和解码器(编码器 + 解码器)的组合方法并不能为多器官或肝血管分割提供进一步的增强。原因可能归因于它们的优势可能会重叠或抵消,从而阻止在这些特定任务中产生协同效益。因此,我们随后的实验将主要采用仅编码器模型来分割多个对象,而仅使用解码器模型来针对更小和更具挑战性的肿瘤或病变进行分割任务。
用于 Transformer 解码的深度分析
1) 器官/肿瘤查询的数量
为了成功训练,可学习查询的数量必须至少等于类的数量(即,每个类必须至少有一个查询)。但是,当我们在分割过程中改变查询数量时,我们观察到具有仅解码器配置的 3D TransUNet 的性能在很大程度上不受此参数的影响。表 IV 列出了这些发现的详细摘要。
2) 用于更新查询的多尺度 CNN 特征
我们的 Transformer 解码器的一个决定性特征是它集成了来自 CNN 解码器的多尺度特征,这些特征具有丰富的定位细节。这些功能在通过交叉注意力与本地化多尺度 CNN 表示的协同作用逐步完善可学习的查询方面发挥着关键作用。如表 V 所示,我们的实验包括 Decoder-only 的配置。与 t 相比,始终观察到的性能增强基线 Transformer 解码器(其中分割掩码是通过直接使用学习查询的点积和最后一层 CNN 特征来计算的)强调了在查询更新过程中整合多尺度 CNN 特征不可或缺的性质。
3) Transformer 解码器中的粗细到精细化
在每个 Transformer 解码器层中,粗细到细细化使用当前迭代中预测的掩码来限制前景区域内的交叉注意力,从而在下一次迭代中细化器官查询。为了证明这种策略的有效性,我们选择了血管肿瘤分割作为代表性案例研究。之所以选择这种选择,是因为 Transformer 解码器在分割小目标(如肿瘤或病变)方面表现出的熟练程度。如表 V 所示,粗细化到细化(掩蔽交叉注意力)的整合始终产生增强的结果。为了更直观地理解,我们在图 2 中提供了一个定性示例,阐明了这种关注如何为复杂的目标细化掩码。从第一次到第三次迭代,肿瘤的分割质量得到了显著提高。
E. 泛化到其他肿瘤/病变数据集
我们评估了我们的 3D TransUNet 在不同成像模式和任务中的泛化性,包括从 MRI 数据中分割脑转移瘤 (BraTS2023)、胰腺肿瘤 (PDAC) 和囊肿分割。结果详见表 VI 和表 VII。我们的研究结果证实,仅编码器变体提供了适度的增强,而仅解码器配置大大提高了肿瘤/病变分割的性能。例如,在脑转移性病变分割中,仅编码器模型提供了边际增强,而仅解码模型有助于2.9%.这种趋势在 PDAC 和囊性病变分割中更为明显。具体来说,仅编码器模型将 PDAC 细分从 56.94% 提高到 58.38%,而仅解码器模型将其进一步推高至 62.66%。在囊性病变的情况下,仅编码器模型将 nnUNet 的性能从 56.88% 提高到 57.98%,而仅解码模型将其进一步提高到 61.04%。因此,平均 Dice、Sensitivity 和 Specificity 的聚合指标都见证了仅解码器模型的显着改进。
F. 与最先进的比较
我们在表 VIII 和表 IX 中将我们的 TransUNet 与以前的 2D 和 3D 最先进的多器官分割和肝血管肿瘤分割方法进行了比较。凭借基于 U-Net 架构构建的 2D 版本和基于 3D nnU-Net 框架的 3D 变体,我们的 TransUNet 始终优于其他最先进的方法,强调了其在各种 U-Net 框架中的有效性。如上所述,利用 Transformer Encoder 捕获全局器官关系的能力,我们使用 Encoderonly 设计进行多器官分割。相反,给定Transformer Decoder 在优化小目标方面的实力,我们选择了仅 Decoder 设置进行肿瘤分割。具体来说,我们将 TransUNet 与一系列方法进行了比较,包括:1) 2D 技术,如 UNet [2]、DeepLabv3+ [42] 和 UNet++ [43],辅以注意力增强的 CNN 方法,如 AttnUNet [12] 和 Pyramid Attn [41],在 224 × 224 和 512 × 512 的分辨率下进行评估;2) V-Net [30]、DARR [31] 和 3D nnU-Net [44] 等 3D 方法,以及包括 CoTR [17]、nnformer [16]、VT-UNet [45] 和 Swin UNETR [20] 等尖端的以 Transformer 为中心的策略。正如表 VIII 中的结果所证实的那样,TransUNet 不仅超越了传统的基于 CNN 的自我注意模型,而且优于许多最先进的面向 Transformer 的技术。例如,当与最近最先进的基于 Transformer 的方法(如 CoTr 和 nnformer)进行基准测试时,我们的 TransUNet 在具有挑战性的胆囊分割任务中实现了约 10% 的 Dice 分数提高约10%,整体分割提高了约 3%.
值得注意的是,如表 X 所示,我们的 3D TransUNet 超越了 BraTS2021 挑战赛中排名第一的解决方案 nnU-Net-Large [32],强调了我们提出的方法的稳健性和有效性.
V. 结论
虽然 U-Net 已经成功,但它在处理长期依赖关系方面的局限性促使人们探索 Transformer 作为一种替代架构,我们之前开发的 TransUNet 就是一个例子,利用了 U-Net 和 Transformers 的综合优势。我们的研究将 TransUNet 架构扩展到 3D 网络,建立在 nnU-Net 的基础上。利用 Transformers 在编码器和解码器设计方面的能力,我们推出了Transformer 编码器,对 CNN 特征图补丁进行标记化,有助于更丰富地提取全局上下文;2) Transformer 解码器,旨在自适应地优化分割区域,利用候选提案和 U-Net 特征之间的交叉注意力机制。我们的调查进一步阐明了医疗分割任务具有不同的架构偏好,并提供了对 tailor 的见解设计配置以满足这些特定要求。实证结果表明,我们的 3D TransUNet 在各种医疗分割任务中的卓越性能。值得注意的是,我们的方法在多器官、胰腺肿瘤、肝血管和脑转移瘤分割方面表现出色。此外,我们在 BraTS2021 挑战赛中取得了显著的成果。
相关的图标: