TransBTS: Multimodal Brain Tumor Segmentation Using Transformer

1. Abstract

• 研究背景介绍，引出本文工作
  • Transformer可以从使用自注意力机制进行全局信息建模中受益，近来来已经在自然语言处理和2D图像分类方面取得了成功。【研究背景，介绍Transformer的优点：全局信息处理的好】
  • 然而，局部和全局特征对于密集的预测任务都至关重要（不可或缺），特别是3D医学图像分割任务。【3D图像分割即关注局部也关注全局，这就是本文研究的出发点，即不仅仅需要Transformer这种全局信息，同时也需要得到CNN那种的局部信息，所以进行了改进】
• 介绍本文工作
  • 本文，我们首次探索了将3DCNN中的Transformer进行MRI的脑部肿瘤分割，提出了一个基于编码解码结构新的网络TransBTS。【3DCNN提取局部信息，Transformer获取全局信息】
    • 为了捕获局部的3D上下文信息，编码器首先使用3DCNN来提取三维空间特征图。
    • 同时，对用于映射令牌的特征图进行了精心的改革，这些令牌被馈送到Transformer中进行了全局特征建模。【tokens是Transformer中一个重要的概念】
    • 解码器利用Transformer嵌入的特征并执行渐进式上采样以预测详细的分割图。
• 实验结果
  • BraTS2019数据集上的实验结果表明，TransBTS优于基于3D MRI扫描的脑肿瘤分割的最新方法。

Introduction

• 介绍研究背景，胶质瘤
  • 神经胶质瘤是最常见的恶性程度不同的恶性脑肿瘤。
  • 在核磁共振成像（MRI图像）上自动精准地分割这些恶性肿瘤对临床诊断中具有重要意义。
• 介绍CNN家族网络
  • 卷积神经网络CNN已经在各种视觉任务中取得了巨大的成功，比如分类，分割和目标检测。
    • 分割家族
      • 全卷积网络FCN首次实现了端到端的语义分割，且取得了令人印象深刻的结果
      • U-Net使用跳过连接的对称的编码-解码器结构来改善细节的保持，变成了医学图像分割中的主流架构。许多U-Net网络的变体例如U-Net++，Res-UNet进一步提升了图像分割的性能。
    • 缺陷
      • 尽管基于CNN的方法具有出色的表示能力，但是由于卷积核感受野的限制，所以难以去建立一个明确的长程依赖关系【即难以获取全局信息，这里如果将HDC和相应的变种融入进来会不会有改善】
      • 这种卷积运算的限制对学习全局语义信息增加了挑战，这种全局语义信息对于像分割这种密集预测任务来说至关重要。
• 介绍Attention机制
  • 受到自然语言处理NLP中注意力机制的启发，现有的研究通过将注意力机制与CNN模型进行融合来克服这个限制。
  • 具体融合产生的网络
    • 非局部神经网络基于自注意力机制设计了一个即插即用的非局部算子， 该算子可以捕获特征图中的长距离依赖关系，但是会消耗大量内存和计算成本。
    • Schlemper提出了一个注意力门模型，它可以以最小的计算开销集成到标准的CNN模型中，同时增加模型的灵敏度和预测精度。
    • Transformer 设计用于seq2seq任务中的远程依赖性进行建模，并捕获序列中任意位置之间的关系。这个结构仅仅基于自住意而提出，同时完全消除了卷积。与之前的基于CNN的方法不同的是，Transformer不仅在建模全局上下文方面功能强大，而且在大规模预训练的情况下也可以在下游任务上取得出色结果。
• 介绍Transformer in vision的工作进展
  • 近期，基于Transformer的框架已经在各种计算机视觉任务中取得最优的性能。 Vision Transformer（ViT）将图像分割为小块，并将这些小块之间的相互性建模为序列，从而在图像分类上获得令人满意的结果
  • DeiT进一步介绍了一种知识蒸馏的方法来训练Transformer。
  • DETR在Transformer的帮助下将目标检测视为一项预测任务。
  • TransUNet是一项并行工作，它使用VIT进行医学图像分割。我们将在2.4节中介绍我们的方法与TransUNet之间的不同。（elaborate）
• 研究动机 （Research Motivation ）
  • Transformer的成功主要体现在图像分类任务中，而对于分割这种密集型预测任务，局部信息和全局信息都同样重要，正如（某一篇论文）指出，当直接将图像分割为小块作为Transformer的tokens的时候局部信息被忽略了。
  • 然而，对于超出2D的医学三维数据（3D MRI scans），连续切片之间的局部特征建模（即深度维）也是三维分割的关键。
  • 问：如何设计出一种能够提高网络性能的神经网络，利用高表现力的Transformer 在体数据的空间和深度维度上对局部和全局特征进行建模
• 总结
  • 本文中，我们首次提出了尝试去探索在3D CNN中的Transformer用于3D MRI脑部肿瘤分割（TransBTS)
    细节
    • 提出的TransBTS建立在一个编码解码结构，网络的编码器首先使用3DCNN提取三维空间特征，然后同时对输入3D图像进行下采样，从而产生可以有效的捕获到局部3D上下文信息的紧凑的三维特征图。
    • 每一个特征图被调整为一个向量（例如 token）被输入到Transformer进行全局特征建模。
    • 3D CNN的解码器从Transformers中提取特征嵌入，并进行渐进上采样以预测全分辨率分割图。
  • 实验结果
    • 在BraTS2019数据集上进行实验， 证明TransBTS在3D MRI扫描上实现了比最先进的脑肿瘤分割方法更优越的性能。
    • 我们也进行了综合消融研究，阐明在3DCNN中加入Transformer的结构工程来释放这两个模型的力量。
    • 我们希望TransBTS可以作为一个强力的3D 基线（baseline）来促进未来在三维分割上的研究。

2. Method

2.1 overall Architecture of TransBTS

• 下图给出了所提出的TransBTS结构：
  
  • 给定一个MRI 扫描$X\in {\mathbb{R}}^{C\times H\times W\times D}$ 其空间分辨率为$H\times W$, 深度维度为D（切片数量），通道C（模态数量）。
  • 我们首先使用3D CNN来生成紧凑的特征图捕获空间和深度信息，然后利用Transformer编码器在全局空间中建立长距离依赖模型。
  • 然后对上采样层和卷积层进行多次叠加，逐步得到高分辨率的分割结果。
  • TransaBTS的网络细节如下：我们将细致描绘TransBTS的组件。

2.2 Network Encoder

• 因为Transformer的计算复杂度与tokens数目（即序列长度）呈平方关系，因此直接将输入图像转为序列作为Transformer的输入是不切实际的。

  • 因此 ViT将图像分割固定大小的块（16x16为一个patches），然后将每一个块转为一个token，这样讲序列的长度减少了256倍（16的平方）。
  • 对于3D 的三维数据，直接像ViT那样进行直接token化，就是讲数据分割为3D的块，然而这种简单的策略使得Transformer无法跨空间和深度维度对图像局部上下文信息进行建模以实现三维分割。

• 解决策略

  • 因此，未来解决这个挑战，我们的解决方法是堆叠3x3x3的卷积块，并使用步长为2的卷积进行下采样，逐渐将输入图像编码为低分辨率的具有高层特征表示的$F\in {\mathbb{R}}^{K\times \frac{H}{8}\times \frac{W}{8}\times \frac{D}{8}}(K=128)$，这只有输入尺寸H,W,D的1/8。
  • 在这种方法中，丰富的局部3D上下文特征被有效的嵌入到了F中，然后F被输入到Transformer的编码器来进一步学习全局感受野有关的长程关系。

• Transformer编码器的特征嵌入：

  • Feature Embedding of Transformer Encoder 对于给定的特征图F，未来确保每个体积的综合表示。我们使用线性投影（一个3x3x3的卷积层）将通道维度从K=128增加到d=512。Transformer层期待一个序列作为输入，因此，我们将空间和深维度压缩称为一个维度，结果为$d\times N\left(N=\frac{H}{8}\times \frac{W}{8}\times \frac{D}{8}\right)$的特征图f，这个被视为一个N个d维的tokens。
  • 为了编码分割任务中至关重要的局部信息，我们引入了一个可学习的位置嵌入，将他们通过直接相加与特征图f进行融合，产生的特征嵌入如下：$$z_0=f+PE=W\times F+PE$$
    这里 $W\in {\mathbb{R}}^{d\times K}$是一个线性投影矩阵, $PE\in {\mathbb{R}}^{d\times N}$ 表示位置嵌入, $z_0\in {\mathbb{R}}^{d\times N}$ 指的是特征嵌入。

• Transformers 层

  • Transformer Layers. Transformer 的编码器由L个Transformer层组成。每个层都有一个标准的结构，即 由一个多头注意力块 【Multi-Head Attention (MHA) block】和一个前馈网络 【Feed Forward Network (FFN)】组成。第 $?$ -th $(?\in [1,2,\dots ,L])$ 个Transformer 层可以通过下面的公式计算：
    $$\begin{array}{c}{z}_{?}^{\prime }=MHA\left(LN\left(z_{??1}\right)\right)+z_{??1}\\ z_{?}=FFN\left(LN\left(z_{?}^{\prime }\right)\right)+z_{?}^{\prime }\end{array}$$
    这里 $LN(?)$ 表示层归一化 and $z_{?}$是第$?$ -th Transformer 层的输出.

2.3 Network Decoder

• 为了在原始的3D图像空间上$(H\times W\times D)$生成分割结果，我们使用3D CNN解码结构来进行特征上采样和像素级分割。
• Feature Mapping 特征映射：
  • 为了适应3D CNN解码器的输入维度，我们设计了一个特征映射模块来讲序列数据投影回一个标准的4D特征图。特别的，Transformer$z_L\in {\mathbb{R}}^{d\times N}$的输出序列首先被规整为$d\times \frac{H}{8}\times \frac{W}{8}\times \frac{D}{8}$。同时为了降低解码器的计算复杂度，卷积块被用来减少通道维度从d到K。【前面从K升到d，这里从d降到K,对称式设计】。通过这些操作，得到了特征图$Z\in {\mathbb{R}}^{K\times \frac{H}{8}\times \frac{W}{8}\times \frac{D}{8}}$，其具有与在特征编码部分的F一样的维度。
• 连续特征上采样 Progressive Feature Upsampling 。在特征映射后，级联的上采样操作和卷积块被用来将Z逐渐的恢复到一个全分辨率的分割结果$R\in {\mathbb{R}}^{H\times W\times D}$。除此之位，使用跳过连接将编码特征和解码器附件进行融合，通过级联具有更丰富的空间细节的更精细的分割掩膜。

2.4 Discussion

• 一个近期的工作TransUNet是使用Transformer实现医学图像分割。TransUnet和TransBTS的关键区别如下：
  • （1）TransUNet是一个2D网络 ，以逐片的方式处理每个3D医学图像。但是，我们的TransBTS是基于3DCNN并一次处理所有图像切片，从而可以更好的利用切片之间的连续信息。即TransUNet仅仅关注标记化图像块之间的空间相关性，但是我们的方法同时可以对切片/深度维度和空间维度中的长程依赖关系进行建模，以实现三维分割。
  • （2）由于TransUNet采用了ViT结构，因此它依赖于大规模图像数据集上预训练的ViT模型，相比之下，我们的TransBTS具有灵活的网络设计，并且在任务特定的数据集上从头开始训练，而无需依赖预训练的权重。

3 Experiments

• 数据和评估指标（Data and Evaluation Metric）
  • 实验中使用的3D MRI数据集由Brain Tumor Segmentation （BraTS）2019挑战赛提供。训练集由355个病例组成，验证集由125个病例组成。每个样本由脑部MRI扫描的四个模态构成，即，原始T1加权（T1），对比读T1加权（T1CE)，T2加权（T2)和流体衰减反转恢复（FLAIR),每个模态具有的体积为240x240x155，且已经对其到了同一空间。
  • 标签具有四个类别：背景（0),坏死性和非增强性肿瘤（标记1），肿瘤周围水肿（2）和GD增强性肿瘤（4）。通过Dice系数和Hausdirff（95%）指标来测量分割的准确性，以增强区域（ET，1），肿瘤核心区域（TC，1，4）以及整个肿瘤区域（WT,1,2,4)。
• 实现细节（Implementation Details）
  • 提出的TransBTS使用Pytorch实现，在4块3090GPU下训练（每一块具有24GB显存），可从头开始使用6000个epochs（batchsize为8）
  • 我们使用Adam优化器来训练模型，使用多元学习率策略将初始学习率被设置为0.0002。其中，初始速率在每次迭代中均以0.9的幂进行衰减。
  • 使用了下面的数据增强方法：（1）从240x240x155随机裁剪到128x128x128体素。(2)随机镜以0.5的概率横穿轴向，冠状，矢状平面。（3）[-0.1,0.1]之间的随机强度变化和[0.9,1.1]之间的比例变化【预测是数据不够，使用数据增强】。
  • softmax Dice损失用于训练网络，L2 Norm也用于模型正则化，权重衰减率为10-5。
  • 在测试阶段，我们利用测试时间增强（TTA)来进一步提高提出的TransBTS的性能。
    

3.1 Main Results

• 我们首先在BraTS 2019训练集上进行五折交叉验证评估，这是依据许多现有工作的常规设定。
• 我们的TransBTS实现在ET,WT,TC的平均Dice系数分数为78.92%，90.23%，81.19%。
• 我们同时在BraTS2019验证集上进行了实验，比较了TransBTS与现有最先进（SOTA)的3D方法，量化结果如表1。
  
• TransBTS 在ET,WT,TC上分别取得了78.93%，90.00%，81.94%的最好Dice系数分数，大大优于现有的最优方法，即在ET,TC,WT分别高1.83%，1.40%，0.64%。
• 就Hausdorff距离度量而言，分割也取得了相当大的进步。与普通的3D U-Net相比，我们的TransBTS在两个指标上都显示出极大的优势，并且有明显的改进。这清楚的揭示了利用Transformer建模全局关系的好处。
• 对于定性分析，我们还显示出各种方法的脑肿瘤分割结果的视觉比较，包括3D U-Net，V-Net，和我们的TransBTS。由于验证集的基础事实的不可靠，因此我们对所有方法的训练设置了5折交叉验证评估，从图二可以看出，TransBTS可以更准确的描述脑肿瘤，通过在每个体积之间建模的长距离依赖性来更准确地描述脑肿瘤，并产生更好地分割掩膜。

3.2 Ablation Study

• 我们进行了广泛的消融实验研究，来验证TransBTS的有效性，并根据BraTS2019训练集上的交叉验证评估其设计选择的合理性。

  • （1）我们调查了Transformer令牌的序列长度（N）的影响，该长度受网络编码器中3D CNN的总体跨度（OS）的控制。
  • （2）我们以各种模型比例（即深度（L）和嵌入尺寸（d））探索Transformer
  • （3）我们还分析了跳过连接的不同位置的影响。

• Sequence length N

  • 表2给出了Transformer各种序列长度的消融研究。 在降采样后，第一行（OS = 16）和第二行（OS = 8）都将特征图的每个体积重塑为特征向量。
  • 值得注意的是，通过将OS从16调整为8，增加令牌的长度会导致性能的显着提高。 ET，TC和WT的Dice得分分别达到（1：73％，1：02％和1：02％）。 由于内存限制，在将OS设置为4之后，我们无法直接将每个体积重塑为特征向量。 因此，我们进行了一些修改，以将序列长度保持为4096，这将每个2x2x2patch展开为特征向量，然后再传递给Transformer。 我们发现，尽管OS从8下降到4，而没有本质上增加序列长度，但是性能却没有提高甚至变差。
    

• Transformer Scale

  • 特征嵌入尺寸（d）和Transformer层数（深度L）这两个超参数主要决定了transformer的比例。 我们进行消融研究，以验证Transformer规模对细分效果的影响。 为了提高效率，我们只训练每个模型配置1000个纪元。 如表3所示，d = 512和L = 4的网络获得了ET和WT的最佳分数。 增大嵌入尺寸（d）不一定会导致性能提高（L = 4，d：512对768），但会带来额外的计算成本。 我们还观察到，就性能和复杂性而言，L = 4是“transformer”的“最佳点”。
  • 

• Position of Skip-connections(SC)

  • 为了提升模型的表达能力，我们进一步探究了跳过连接的位置。消融实验结果被列在了表4。

• 如果将跳过连接附加到第三个Transformer层，则更像是在不补偿空间细节损失的情况下，从相邻层进行聚合。遵循传统的从U-Net跳过连接的设计，即连接到3D Conv层，重要的ET和TC获得了可观的收益（3.96%，1.23%），这要归功于低级空间细节信息的恢复。

4 Conclusion

• 本文中，我们提出了一个新颖的分割框架，该框架有效的将Transformer融入3D CNN中，以进行MRI中的多模态脑部肿瘤分割。
• 最终架构TransBTS，不仅继承了3D CNN对局部上下文信息建模的优点，同时利用了Transformer学习全局语义相关。
• 在BraTS2019数据集上的实验结果验证了提出的TransBTS的有效性，在未来的工作中，我们将探究Transformer中的计算和存储有效的注意力机制，以开发以效率为中心的体积分割模型。