A Robust Volumetric Transformer for Accurate3D Tumor Segmentation
摘要Abstract
- 介绍自己的工作
- 我们提出了一个用于体素分割的Transformer结构,这是一项具有挑战性的任务,需要在编码局部和全局空间线索的时候保持复杂度平衡,并沿体素的所有轴保留信息。
- 提出的编码器的受益于注意力机制,能够同时编码局部和全局线索,而解码器采用并行的自我和交叉注意公式来捕获精细细节以进行边界细化
- 介绍提出工作的效果或者成就
- 根据经验,我们表明,所提出的设计选择产生了一个计算效率高的模型,在医学分割十项全能 (MSD) 脑肿瘤分割 (BraTS) 任务上具有竞争性和有希望的结果。
- 我们进一步表明,我们的模型学习的表示对数据伪影具有鲁棒性
1.引言 Introduction
- 介绍一下近期Transformer结构在医学影像上的进展
- 受计算机视觉中基于Transformer的模型的强大经验结果、它们有希望的泛化和鲁棒性特征 以及它们对远程交互建模的灵活性的启发,我们提出了一种用于分割3D 医学影像模式(MRI,CT),称为VT-UNet。
- 介绍早期Transformer in Vision(medical Image)的 缺点及缺陷。
- 早期为 3D 医学扫描开发基于Transformer的分割模型的努力已被证明优于最先进的 CNN 同行然而,这些方法将 3D 体积分成 2D 切片并将 2D 切片作为输入处理。
- 因此,对于封装切片间依赖关系至关重要的大量且潜在的关键体积信息会丢失。
- 而一些混合方法(使用卷积模块和Transformer层)保持了3D体积的完整无缺,纯Transformer结构的设计,能够在输入时保持完整的体积数据,在文献中尚未探索。
- 我们的工作朝着这个方向迈出了第一步,并提出了一个模型,该模型不仅实现了更好的分割性能,而且还展示了对数据伪影的更好鲁棒性。、
- 介绍我们工作的意义
- Transformer 模型具有高度动态和灵活的感受野,并且能够捕获远程交互,但是为体积分割设计基于 Transformer 的 UNet 架构仍然是一项具有挑战性的任务。 原因如下:
- 封装体素信息并捕获体积序列中任意位置之间的连接并不简单。与基于 Transformer 的 2D 图像分割方法 [4] 相比,体积的每个切片中的数据都连接到三个视图,丢弃其中任何一个都可能是有害的。
- 在一个体积中保存空间信息是一项艰巨的任务。即使对于2D图像,在视觉转换器(ViT)中引入的将图像分成小块并将小块投影到标记中时,局部结构线索也会丢失,如标记到标记ViT所示[28]。因此,有效地编码局部线索,同时沿着体积的多个轴捕获全局交互是一项具有挑战性的任务。
- 由于自注意的二次方复杂性,以及三维体积张量输入的大尺寸,设计一个基于Transformer的分割模型,在计算上是高效的,这需要仔细的设计考虑。
- 我们提出的VT-UNet模型通过提出一些模块有效地解决了上述设计挑战。在我们基于UNet的架构中,我们开发了两种类型的Transformer块。
- 介绍我们模型的细节
- 首先,我们的编码器中的块,直接在三维体积上工作,以分层的方式,共同捕捉局部和全局信息,在精神上类似于Swin Transformer块。
- 其次,对于解码器,我们在扩展路径中引入了并行的交叉注意和自我注意,这在解码器的查询和编码器的键和值之间建立了一座桥梁。通过这种交叉注意和自我注意的并行化,我们的目标是在解码过程中保留完整的全局背景,这对分割任务很重要。
- 由于VT-UNet不受卷积影响,在解码过程中结合了两个模块的注意力输出,因此序列的顺序对于获得准确的预测非常重要。
- 受[24]的启发,除了在每个Transformer块中计算注意力时应用相对位置编码外,我们还增强了解码过程,并注入了从序列中标记的傅里叶特征位置中提取的补充信息。
- 总结一下本文的主要贡献
- 总结,我们的主要贡献如下:
- 我们从序列到序列的角度重新阐述了肿瘤的体积分割,并提出了一个UNet形状的体积变换器,用于多模式的医学图像分割。
- 我们设计了一个具有两个连续自我注意层的编码器模块,以联合捕获局部和全局上下文线索。此外,我们设计了一个解码器模块,它支持基于并行(移位)窗口的自我和交叉注意。这种并行化使用查询的一个共享投影,并独立计算交叉和自我关注。为了进一步增强我们的解码功能,我们提出了一种凸组合方法以及傅立叶位置编码。
- 结合我们提出的设计选择,与现有的方法相比,我们大大限制了模型参数,同时保持较低的FLOPs(见图1)。
- 我们进行了广泛的评估,并表明我们的设计达到了最先进的体积分割结果,同时对数据伪影有更强的鲁棒性。
- 总结,我们的主要贡献如下:
2 方法论 Methodology
- 介绍文章提出模型的具体设计情况
- 介绍一下一些公式的定义
- 我们分别用黑体小写x和黑体大写X来表示向量和矩阵。让 X = { x 1 , x 2 , ⋯ , x τ } , x i ∈ R C \mathbb{X}=\left\{\mathrm{x}_{1}, \mathrm{x}_{2}, \cdots, \mathrm{x}_{\tau}\right\}, \mathrm{x}_{i} \in \mathbb{R}^{C} X={x1,x2,⋯,xτ},xi∈RC 是一个代表感兴趣的信号的序列(例如,一个MRI卷),我们称每一个 x i \mathbf{x}_{i} xi 为token. 我们假设原始形式的记tokens可能不是定义跨度的最佳选择。
- 因此,在自我注意(SA)中,我们通过从输入tokens学习线性映射来定义跨度,这一点我们可以用
R
τ
×
C
v
∋
V
=
X
W
V
\mathbb{R}^{\tau \times C_{v}} \ni \mathbf{V}=\mathbf{X W}_{\mathbf{V}}
Rτ×Cv∋V=XWV,这里我们堆叠tokens
x
i
s
\mathbf{x}_{i} \mathrm{~s}
xi s作为
X
\mathbf{X}
X 的行(i.e.,
X
=
[
x
1
∣
x
2
∣
⋯
x
τ
]
⊤
\mathbf{X}=\left[\mathrm{x}_{1}\left|\mathbf{x}_{2}\right| \cdots \mathbf{x}_{\tau}\right]^{\top}
X=[x1∣x2∣⋯xτ]⊤ ) 我们略加修改了子注意力机制在我们的任务中,其形式如下:
S A ( Q , K , V ) = SoftMax ( Q K ⊤ / C + B ) V , \mathrm{SA}(\mathbf{Q}, \mathbf{K}, \mathbf{V})=\operatorname{SoftMax}\left(\mathbf{Q K}^{\top} / \sqrt{C}+\mathbf{B}\right) \mathbf{V} \text {, } SA(Q,K,V)=SoftMax(QK⊤/C+B)V, - 其中 R τ × τ ∋ B \mathbb{R}^{\tau \times \tau} \ni \mathrm{B} Rτ×τ∋B 是可寻来你的,并且作为卷中不同标记的相对位置偏差,其中 V = X W V , K = X W K \mathbf{V}=\mathbf{X} \mathbf{W}_{\mathbf{V}}, \mathbf{K}=\mathbf{X} \mathbf{W}_{\mathbf{K}} V=XWV,K=XWK, Q = X W Q \mathbf{Q}=\mathbf{X} \mathbf{W}_{\mathbf{Q}} Q=XWQ. 在实践中,多次并行计算多个注意力头的子注意力称为多头注意力,例如,公式1是基于体积Transformer窗口的多头自注意力和基于体积Transformer变换窗的多头自注意的基本构建单元。
- VT-UNet的概述
- . 图2显示了所提出的体积Transformer网络或者简称为VT-UNet的概念图。 我们模型输入的3D卷大小为 D × H × W × C D \times H \times W \times C D×H×W×C. 输出是一为个 D × H × W × K D \times H \times W \times K D×H×W×K维度的卷, 代表存在或者不存在体素级类别标签,(K是类别的数量)下面我们将条论VT-UNet的结构形式,并详细解释我们设计背后的功能和原理。
- 介绍VT编码器
- . VT编码器由具有线性嵌入层的3D patch 分割层组成,3Dpatch 合并层后有两个连续的VT编码模块。
- 介绍什么是3D Patch分割(块):
- 基于Transformer的模型以一连串的tokens工作,VT-UNet的第一个模块接受一个 D × H × W × C D \times H \times W \times C D×H×W×C 维度的医学卷,例如, MRI)。通过将3D卷分为不重叠的3D patches来创建一系列的tokens。如图2b,. 分割核的大小为 P × M × M P \times M \times M P×M×M, resulting 结果用 τ = ⌊ D / P ⌋ × ⌊ H / M ⌋ × ⌊ W / M ⌋ \tau=\lfloor D / P\rfloor \times\lfloor H / M\rfloor \times\lfloor W / M\rfloor τ=⌊D/P⌋×⌊H/M⌋×⌊W/M⌋ tokens来描述卷. 3D patch分割后跟着一个线性的嵌入层来将维度为 P × M × M P \times M \times M P×M×M的标记映射为一个 C C C 维向量. T,根据我们的实验,M,P,C的典型值分别为4,4,72.
- 介绍什么是VT编码块
- 在VIT中,由于它们的构造方式,tokens可以带有明显的空间信息。最近的一些研究表明,在ViT中通过窗口化进行SA的重要性,最值得注意的是在SwinTransformer中。遵循在Swin Transformer设计上相似的原则,尽管是针对体积数据(3D),我们提出了3D窗口操作在VT编码模块中(VT-Enc-Blks)。特别的,我们提出了两种类型的窗口,称为普通窗口和移位窗口。为简单起见,我们分别用VT-W-MSA和VT-SW-MSA表示,图2 提供了VT-W-MSA和VT-SW-MSA的设计细节,图2(b)说明了窗口操作。
- VT-W-MSA和VT-SW-MSA使用了具有窗口的注意力层,其后接着2层的多层感知机,中间是高斯误差线性单元非线性(GELU).在每个MSA和MLP之前都会使用层归一化(LN),在每个模块之后都会应用残差链接。窗口化使我们能够在建模标记之间的长距离依赖关系时注入归纳偏见。
- 在VT-W-MSA和VT-SW-MSA中我们将体积均匀地分成更小的非重叠窗口,如图 2 (b) 所示。我们将体积均匀地分成更小的非重叠窗口,如图 2 (b) 所示。由于相邻窗口中的令牌不能用 VT-W-MSA 互相看到,我们利用 VT-SW-MSA 中的移位窗口(参见最右边的面板图 2(b)),它桥接 VT 的相邻窗口中的令牌-W-MSA。窗口化由Swin Transformer实现,可以理解为体素数据地泛化。
- 请注意,我们工作中的窗口操作类似于 [15],它将窗口的好处从图像扩展到视频。将所有东西放在一起,VT-Enc-Blk 实现了以下功能: z ^ l = V T − W − M S A ( LN ( z l − 1 ) ) + z l − 1 , z ^ l + 1 = V T − S W − M S A ( LN ( z l ) ) + z l , z l = MLP ( LN ( z ^ l ) ) + z ^ l , z l + 1 = MLP ( LN ( z ^ l + 1 ) ) + z ^ l + 1 , \begin{aligned} &\hat{\mathrm{z}}^{l}=\mathrm{VT}-\mathrm{W}-\mathrm{MSA}\left(\operatorname{LN}\left(\mathrm{z}^{l-1}\right)\right)+\mathrm{z}^{l-1}, \quad \hat{\mathrm{z}}^{l+1}=\mathrm{VT}-\mathrm{SW}-\mathrm{MSA}\left(\operatorname{LN}\left(\mathrm{z}^{l}\right)\right)+\mathrm{z}^{l} \text {, }\\ &\mathbf{z}^{l}=\operatorname{MLP}\left(\operatorname{LN}\left(\hat{\mathbf{z}}^{l}\right)\right)+\hat{\mathbf{z}}^{l} \text {, }\quad \mathrm{z}^{l+1}=\operatorname{MLP}\left(\operatorname{LN}\left(\hat{\mathrm{z}}^{l+1}\right)\right)+\hat{\mathrm{z}}^{l+1} \text {, } \end{aligned} z^l=VT−W−MSA(LN(zl−1))+zl−1,z^l+1=VT−SW−MSA(LN(zl))+zl, zl=MLP(LN(z^l))+z^l, zl+1=MLP(LN(z^l+1))+z^l+1,
- 这里 z ^ l \hat{\mathbf{z}}^{l} z^l 和 z l \mathbf{z}^{l} zl 分别代表块l中VT-W-MSA和MLP模块输出地特征。
- 介绍3D patch Merging
- 我们利用 3D 补丁合并块在 VT-UNet 的编码器中生成特征层次结构,拥有这样的层次结构对于在密集预测任务的输出中生成更精细的细节至关重要.【这个东西应该类似于CNN中地池化操作】
- 在每一个VT-Enc-Blk后,我们以一种非重叠地方式沿着空间轴合并相邻地tokens来产生新的tokens。这样做时,我们首先连接每组 2x2 个相邻标记的特征,结果向量通过线性映射投影到令牌的通道维度加倍的空间。
- 补丁合并的好处不仅限于特征层次结构,SA 的计算复杂度是令牌数量的二次方。因此,在每个 VT-Enc-Blk 之后,补丁合并将 VT-UNet 的 FLOP 计数减少了 16 倍。为了给读者一个更好的主意,我们将在§Sec中讨论。如图 4 所示,微型 VT-UNet 模型与其完全体积的 CNN ,相比仅使用 6.7% 的 FLOP,同时实现了相似的性能(确实略好一些)!请注意,在瓶颈阶段不使用补丁合并块。[为何会出现这种差异]
- 介绍VT解码器
- 在由 VT-Enc-Blk 和 3D 补丁扩展层组成的瓶颈层之后,VT 解码器从连续的 VT 解码器块 (VT-Dec-Blks)、3D 补丁扩展层和最后的分类器开始生成最终预测。 VT-Enc-Blk 和 VT-Dec-Blk 之间存在一些基本的设计差异,我们将在下面讨论。
- 3D Patch Expanding【这个似乎咋iSwin-Unet中见到过】
- 此功能用于以某种方式恢复补丁合并的效果。 换句话说,为了构建与输入具有相同空间分辨率的输出,我们需要在解码器中创建新的标记。为了便于讨论,考虑在瓶颈层之后扩展的补丁(参见图 2 的中间部分)。
- 补丁扩展的输入标记的维数为 8C。在补丁扩展中,我们首先使用线性映射将输入标记的维度增加两倍。在重塑之后,我们可以从维度为 2x8C 的结果向量中获得维度为 4C 的 2x2 个标记。这样,我们将沿空间轴重塑,因此对于 D/4xH/32 xW/32 x8C,我们创建 D/4 x H/16 x W/16 x 4C 令牌
- 介绍VT 解码模块
- 同样是一个纯Transformer构成的解码器
3 Related work
- 介绍视觉Transformer的发展历程
- 视觉Transformer已经针对不同的计算机视觉任务表现出卓越的经验结果和有前途的特点。例如,相比CNN,它们较少偏向纹理 ,并表现出更好的泛化性和鲁棒性。
- 介绍Transformer用于图像分割
- Transformer 最近被用于研究图像分割,TransUNet是第一个基于 Transformer 的医学图像分割方法它采用了 UNet 结构,并用 ViT [8] 替换了瓶颈层,其中补丁嵌入应用于从 CNN 编码器生成的特征图(其中输入是 3D 卷的 2D 切片)。
- 与这些混合方法(同时使用卷积和自注意力)不同,Cao 等人。 提出了 Swin-UNet,一个纯粹的基于变压器的医学图像分割网络。它继承了 swin-transformer 块 并显示出比 TransUNet 更好的分割结果。
- TransUnet [6] 的 3D 版本,称为 TransBTS [25],具有 CNN 编码器-解码器设计和 Transformer 作为瓶颈层
- 周等人。 [30] 提出了具有基于 3D Swin Transformer 的块的 nnFormer 作为编码器和解码器,具有交错的卷积干。
- 有人提出了一种使用变压器作为编码器并通过跳跃连接将中间编码器输出直接连接到解码器的模型。
- SOTA方法的编解码器结构比较如图2(d)所示。
- Transformer 模型的优点
- 上述用于 3D 医学图像分割的基于Transformer的方法已经显示出它们的前景,与它们的 CNN 对应物相比,它们获得了更好的性能。
- 另一方面,我们提出的模型完整地处理体积数据,从而完全编码切片之间的交互。
- 此外,我们提出的模型完全基于 Transformer 构建,并在编码器-解码器设计中引入横向连接以执行 CA 和 SA。这些设计元素有助于实现更好的分割性能,以及增强的鲁棒性
4 实验
- 介绍实现细节
-
实现细节
- 我们使用了来自MSD BraTS任务的484个MRI扫描,遵循前任的设定,我们将484个扫描分别分为80%的训练集,15%的验证集,5%测试集。
- 我们使用Pytorch,具有单卡的A40,Swin-T在ImageNet-22K的与训练权重被用来初始化模型。
- 对于训练,我们使用AdamW优化器,学学习率为1e-4,1000epochs ,batch_size = 4 ,我们使用了旋转,添加噪声,模糊,和添加gamma作为数据增强技术。
-
分析实验结果
-
表1比较了VT-UNet和近期的基于Transformer方法和最先进的基于CNN的方法。
-
我们通过改变模型训练中使用的嵌入维度的数量引入了VT-UNet的变体,我们的变体如下:(a)small VT-UNet-S C = 48;(2) Base VT-UNet-B C=72;我们遵循别人使用的评价策略,使用DICE相似系数,HD距离作为评价指标,分别计为三类计算他们;(1)增强肿瘤ET,(2)肿瘤核心TC,(3)全肿瘤。
-
我们的量化结果如表1,表明VT-UNet在DSC和HD上取得最优秀的整体性能,表4表表明在不可见病人数据上的量化结果。我们可以观察到我们的模型可以准确地分割结构并描绘肿瘤的边界。
-
我们相信,捕获相邻切片之间的远程依赖在我们的模型性能中扮演着重要角色。
-
我们在表 3 中的实证结果揭示了在 VT-Dec-Blks中引入并行 CA 和 SA 以及 FPE 和凸组合的重要性。我们可以注意到所有这些组件都有助于模型的性能
-
-
鲁棒性分析
-
患者的运动和采集条件等因素会给 MRI 带来噪音。在这里,我们通过在推理时将伪影综合引入 MR 图像来研究 VT-UNet 的鲁棒性。包括(1)运动伪影,(2)重影伪影,(3)尖峰伪影(人字形伪影)。表2比较累我们模型和nnFormer的鲁棒性。
-
结果表明VT-UNet在存在这些干扰因素的情况下,表现得更加可靠。
-
我们的研究结果与现有的RGB图像的工作是一致的,其中,基于Transformer的模型对遮挡[17]、自然和对抗性扰动[21,17]表现出更好的鲁棒性,这是因为其高度动态和灵活的接受场。
-
-
5 总结
- 总结本文的工作
- 本文提出了一种用于医学图像分割的体积变换器网络,该网络在处理大型 3D 体积方面具有计算效率,并学习对伪影具有鲁棒性的表示。
- 我们的结果表明,所提出的模型在体积分割中实现了对现有状态方法的一致改进。我们相信我们的工作可以帮助更好的临床诊断和治疗计划。