Transformer论文阅读（三）：UNETR: Transformers for 3D Medical Image Segmentation

Abstract

• 介绍背景
  • 近年来， 具有收缩和扩张路径的（编码器和解码器的U形结构）全卷积神经网络在各种医学图像分割应用中已经展现了突出的地位。[U-net统治医学图像处理]
  • 在这些结构中，编码器是学习全局上下文表示的基石（具有不可或缺的角色），而这些全局上下文表示在解码器中进一步用来执行语义输出预测。
  • 尽管他们成功，作为FCNNS的主要构造块卷积的局部性限制了在这些网络中学习长程空间依赖的能力。
• 介绍目前情况以及启发
  • 受到最近Transformer在自然语言处理NLP中学习长程依赖上取得的成功， 我们将3D医学图像分割任务修改为了序列到序列的预测任务。
• 干了啥工作，提出了啥模型
  • 特别地，我们引入了一个新的结构， 被称为UNEt TRansformers (UNETR),，这是一个使用纯Transformer作为编码器来学习输入体素地序列表示，有效地捕获了全局多尺度信息。
  • 在不同分辨率上，transformer编码器通过跳跃连接直接与解码器连接，最终计算最终地语义分割输出。
• 实验情况如何
  • 我们已经使用MSD数据集广泛验证了我们所提出算法在不同成像模式（MR，CT）上对体积脑肿瘤和脾脏分割任务地性能，我们的结果一致证明了是有利的基准。

1 Introduction

• 介绍医学图像分割的发展
  • 医学图像分割是 众多临床诊断方法的基石，它通常是解剖结构量化分析的第一步。【我为啥要干分割，分割的重要作用】
  • 自从深度学习的出现，FCNNS特别是编码解码结构在多种医学语义分割任务中取得了最优的结果。
  • 在典型的U-Net结构中，编码器通过逐步下采样所提取的特征，负责学习全局上下文表示。同时解码器上采样所提取的特征表示到输入分辨率执行像素级/体素级语义预测。
  • 初次之外，跳跃连接在不同分辨率上合并了编码器和解码器的输出，因此恢复了在下采样过程中丢失的空间信息。
• 介绍基于CNN方法的缺点【单靠CNN解决不了的那种】
  • 虽然基于FCNN的方法具有强大的表示学习能力，由于局部感受野，它们在学习长程依赖的性能受到限制。
  • 因此，这种捕获多尺度上下文信息的缺陷导致了在具有可变形状和尺度的结构的次优分割。【例如不同大小的脑损伤】
  • 多种努力尝试去使用空洞卷积层减轻这个问题的影响，
  • 然而由于CNN的局部性，它们的感受野仍然限制到一个很小的区域。
• Transformer可以解决CNN的问题，所以Transformer in vision
  • 在NLP领域，基于Transformer的模型已经在多项任务取得了最优的基准
  • Transformer 中的自注意力机制使得它可以动态关注词序列中的重要特征，学习长程依赖。
  • 这种意图最近已经被扩展至计算机视觉领域，通过引入ViT的形式。
  • 在ViT种，图像被表示为块嵌入的序列，可以用于图像分类任务中直接预测类别标签。
• 介绍本文工作，基于上述情况提出了什么模型，解决了什么问题。
  • 在本文种，我们计划去使用Transformer 的能力进行体素医学图像分割，并计划引入了一个新的结构叫UNTER。
  • 特别地，我们将3D分割任务修改为了1D地序列到序列地预测问题，使用一个纯Transformer作为编码器从嵌入地输入块种来学习上下文信息。
  • 从Transformer编码器种提取的表示在多个分辨率上通过跳跃连接与解码器合并，最终预测分割输出。
  • 我们在MSD数据集种的脑肿瘤和脾脏分割任务上广泛的验证了我们的UNETR的有效性，在我们的验证集上实验证明了相比于其他模型，我们的方法优秀的性能。
  • 据我所知，我们是第一个提出完全基于Transformer编码器用来进行体素医学图像分割（也就是3D医学图像分割）
  • 考虑到医学图像种体素数据的流行，以及它们在分割种的广泛使用。我们相信我们的UNETR为新类型的基于Transformer的分割模型开辟了新道路。它可以被用在各种应用中。

2 Related Work

• 介绍CNN ,Transformer， CNN-Tranns
  • 基于CNN的分割网络
    • 自从开创性的U-Net引入以来，基于CNN的网络已经在多种2D和3D医学图像分割任务中取得了最优结果。
    • 尽管它们成功，这些网络存在的限制是学习全局上下文和长程空间依赖上存在低性能。这几乎影响了这些挑战性任务的分割性能。
  • 视觉Transformer
    • 视觉Transformer在各种计算机视觉任务中越来越受欢迎，ViT通过纯Transformer的大规模预训练和微调证明了最先进的性能在图像分类数据集上。
    • 在目标检测中，端到端的基于Transformer模型在很多基准上已经展现突出。
    • 最近，一些努力已经探索了在2D图像分割任务中探索使用基于Transformer模型的可能性。
      • chen提出了一个使用Transformer作为Unet的Bottleneck中一层的2D方法，使用该方法进行多器官分割。
      • 除此之外，Zhang等人提出了在单独流中使用CNN和Transformer，然后合并它们的输出。
      • Valanarasu等人提出一个基于Transformer的轴向注意力机制用于2D医学图像分割。
• 我们的模型与这些工作存在三个关键不同点：
  • 1·我们的UNETR是专为3D分割量身定制的，可以直接使用体素数据。
  • 2.我们的UNETR使用transformer作为分割网络主要的编码器，并将其通过跳跃连接直接连接至解码器。而不是使用其作为分割网络的注意层。
  • 3.我们的UNETR不依赖backboneCNN生存输入序列，直接使用标记的块。

3 Methodology

3.1 Architecture

• 介绍网络的主要结构

  • 我们已经概述了提出模型的整体架构如图1
  • UNETR使用由Transformer堆栈组成的收缩苦扩展模式作为编码器，该编码器通过跳跃连接连接到解码器。
  • 我们开始首先描述我们Transformer编码器的工作机制，如在NLP中常用的，Transformer在输入的1D序列上操作。
    • 在我们的UNETR中，我们创建$3\mathrm{D}$输入体素$\mathrm{x}\in {\mathbb{R}}^{H\times W\times D\times C}$$1\mathrm{D}$ 序列，通过将该3D体素划分为扁平的均匀的不重叠的块。 ${\mathbf{x}}_v\in {\mathbb{R}}^{L\times C\times N^3}$
    • 这；里 $(N,N,N)$ 代表每一个块的维度，其中 $L=(H\times W\times D)/N^3$ 是序列的长度。
  • 随后，我们使用一个线性层来投影扁平的块到一个K维的编码空间内，该空间在整个Transformer中保持恒定。
  • 除此之外，为了保存提取到的块的空间特征，我们添加了1D可学习的位置嵌入 ${\mathbf{E}}_{pos}\in {\mathbb{R}}^{L\times D}$到投影块编码中 $\mathbf{E}\in {\mathbb{R}}^{L^2\times C\times K}$ ，公式如下（通过）：
    $${\mathbf{z}}_0=\left[{\mathbf{x}}_v^1\mathbf{E};{\mathbf{x}}_v^2\mathbf{E};\dots ;{\mathbf{x}}_v^L\mathbf{E}\right]+{\mathbf{E}}_{pos}$$
  • 在嵌入层后，我们使用Transformer块的堆栈，这个Transformer块由多头注意力和多层感知机子层组成，
    $$\begin{array}{rl}{\mathbf{z}}_i^{\prime }=& \mathrm{MSA}\left(\mathrm{Norm}\left({\mathbf{z}}_{i-1}\right)\right)+{\mathbf{z}}_{i-1}\\ {\mathbf{z}}_i& =\mathrm{MLP}\left(\mathrm{Norm}\left({\mathbf{z}}_i^{\prime }\right)\right)+{\mathbf{z}}_i^{\prime }\end{array}$$
  • 这里Norm代表层归一化，MLP（多层感知机）由两个具有GELU激活函数的线性层组成。i是中间块的标识符，其范围是在我们当前设置的总块1-12。
  • MSA块是由n个平行的自注意头组成（SA).自注意头（SA)块是一个参数化的函数，用于学习输入序列z中两个元素及其q，k表示之间的相似系。因此，SA的输出计算如下：$$\mathrm{S}\mathrm{A}(\mathbf{z})=\mathrm{Softmax}\left(\frac{q{k}^{†}}{\sqrt{C_h}}\right)v$$
    这里 $\mathbf{v}$ 代表输入序列的值， $C_h=C/n$ 是缩放因子，初次之外，MSA的输出定义如下：
    $$\mathrm{MSA}(\mathbf{z})=\left[{\mathrm{S}\mathrm{A}}_1(z);{\mathrm{S}\mathrm{A}}_2(z);\dots ;{\mathrm{S}\mathrm{A}}_n(z)\right]{\mathbf{W}}_{msa}$$
    这里${\mathbf{W}}_{msa}$代表不同头的可学习权重矩阵。

• 介绍网络结构思想的来源

  • 受到UNet形结构的启发，来自编码器中多个分辨率的特征被合并到解码器，我们从Transformer中提取了序列表示 z i ( i ∈ { 3 , 6 , 9 , 12 } ) \mathbf{z}_{i}(i \in\{3,6,9,12\}) zi​(i∈{3,6,9,12}), 大小为$\frac{H\times W\times D}{N^3}\times C$, 然后将其转为 $\frac{H}{N}\times \frac{W}{N}\times \frac{D}{N}\times C$的tensor。
  • 在我们定义中的表示是在嵌入空间中，如果它已经被调整为Transformer的输出，并且具有的特征大小为C（Transformer 的嵌入大小。【不太懂所需要表达的意思】
  • 因此，我们通过利用连续的3x3x3卷积层，将Reshaped Tensor从嵌入空间从嵌入空间投影到输入空间中（有关详细信息，参见图1）。

• 介绍一下bottleneck

  • 在我们编码器的bottleneck中（transformer的最后一层输出） ，我们使用一个反卷积来变换特征如，以此增加了它的分辨率（增加2倍），然后我们串联调整大小后的特征图和前面Transformer输出的特征图。
  • 将它们送入连续的3x3x3的卷积层，使用反卷积上采样输出。这个过程反复执行直到所有的其它子层回到原始的输入分辨率，并在最后输出的时候将其送入带有softnax激活函数的1x1x1的卷积层来生成像素级语义预测。

3.2 Loss Function

• 我们的损失函数是结合了Dice和交叉熵损失。以体素级的方式计算:$$\mathcal{L}=1-\frac{2}{J}\sum _{j=1}^J\frac{{\sum }_{i=1}^I{G}_{i,j}{Y}_{i,j}}{{\sum }_{i=1}^I{G}_{i,j}^2+{\sum }_{i=1}^I{Y}_{i,j}^2}-\frac{1}{I}\sum _{i=1}^I\sum _{j=1}^J{G}_{i,j}\log Y_{i,j}$$
  这里 $I$ 是体素的数量; $J$ 是种类数; $Y_{i,j}$和 $G_{i,j}$ 分别代表在体素i上类别j的输出的概率值和one-hot编码的标签。

4 Experiements

4.1 Datasets

• 介绍本文使用的数据集
  *

4.2 Implementation Details

• 介绍论文实现细节

4.3 Quantitative Evaluations

4.4 Qualitative Results

5 Ablation

6 Conclusion

• 总结本文干了啥
  • 在本文中，我们引入了一个新的基于Transformer的结构，该结构用于医学图像中的体素语义分割 ，并提出将其重构为1D序列到序列的预测。
  • 我们提出使用一个纯Transformer编码器来增加模型学习长程依赖的计算能力，可以在多个尺度上有效的捕获全局上下文表示。
• 简述通过实验发现了啥
  • 我们在MSD数据集中的体素脑肿瘤分割和脾脏分割任务上验证了我们的UNETR的有效性，该数据集是由CT和MRI图像模式组成。
  • 我们的基准证明了在这些任务上具有优秀的性能。
  • 我们提出的UNETR层是用于医学图像分割中基于Transformer模型的新种类
  • 虽然UNETR初始设计用于3D应用，可以直接扩展至2D应用，未来将有更多的探索。