[医疗 AI ] 3D TransUNet：通过 Vision Transformer 推进医学图像分割

[医疗 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 应用于医疗应用。

网络结构：