Transformer论文阅读（一）：CoTr: Efficiently Bridging CNN and Transformer for 3D Medical Image Segmentation_综合

CoTr: Efficiently Bridging CNN and Transformer for 3D Medical Image Segmentation 【有效桥接CNN与Transformer进行3D医学图像分割】

1 Introduction
2 Materials
3 Methods
- 3.1 CNN-encoder
- 3.2 DeTrans-encoder
- - 3.2.1 Input-to-sequence Transformation
  - 3.2.2 MS-DMSA Layer【多尺度可变形自注意力机制】
  - 3.2.3 DeTrans Layer
- 3.3 Decoder
- 3.4 Implementation details
4 Results
- 4.1 Comparing to models with only Transformer encoder.
- 4.2 Comparing to models with only CNN encoder.
- 4.3 Comparing to models with hybrid CNN-Transformer encoder.
- 4.4 Computational Complexity.
5 Discussion on Hyper-parameter Settings
6 Conclusion

论文地址：https://arxiv.org/abs/2103.03024v1
代码：https://github.com/YtongXie/CoTr

介绍大背景，为啥要研究Transformer
- 卷积神经网络（CNNs)已经成为了当今3D医学图像分割中的事实上的标准。然而，由于卷积的归纳偏置的局部性和权重共享，这些网络中使用的卷积操作不可避免地在建模远程依赖方面存在局限性。【CNN具有一些难以处理的问题，但是这些个问题恰好能被Transformer解决】【权重共享与局部性是CNN的优势也是劣势】
介绍Transformer存在的问题
- 尽管 Transformer 就是为了解决这个问题而诞生的，但它在处理高分辨率 3D 特征图时遇到了极端的计算和空间复杂性。【在这一点上MLP与Transformer几乎等同，因此为了提高Transformer的效率，大多数策略在考虑从CNN的发展及其策略来启发改进视觉Transformer。所以有了本文的工作】
引出本文研究，为了解决上面的问题
- 在本文中，我们提出了一个新的框架，该框架将卷积神经网络和Transformer有效的连接起来（CoTr），用于精确的3D医学图像分割。【CNN与Transformer融合结构更好】
- 在此框架下，CNN被构造来提取特征表示，并构造一个有效的可变形的Transformer(DeTrans）来对提取的特征图进行长程依赖建模。【不同尺度上获取的信息进行编码输入到一个Transformer，对不同尺度的信息进行concat】
- 与同等对待所有图像位置的Transformer 不同，我们的 DeTrans 通过引入可变形的自注意力机制只关注一小部分关键位置。【本文所特殊的地方】
- 因此，DeTrans的计算和空间复杂性已经得到大大降低，使得处理多尺度和高分辨率特征图成为可能，这对图像分割来说至关重要。【与传统的Transformer in vision或融合模型在视觉上的表现不同，我们是可变形的Transformer，重点关注也不一样】
实验与实验结果
- 我们对涵盖11个主要人体器官的数据集Multi-Atlas Labeling Beyond the Cranial Vault (BCV)进行了广泛的评估。
- 实验结果表明我们的CoTr在3D多器官分割任务上比其他基于CNN、基于Transformer和混合的方法带来了可观的性能提升。代码开源：https://github.com/YtongXie/CoTr。

1 Introduction

介绍为啥干这个工作
- 图像分割是医学图像分析中一个长期存在的挑战。自从引入U-Net以来，全卷积神经网络（CNNs）已经成为解决此类任务的主要方法。
- 尽管他们普遍存在，由于归纳偏置的局部性和权重共享，CNNs仍然受制于有限的感受野难以去捕获长程依赖。
介绍前人在这个方面的努力
- 很多努力都致力于扩大CNN的感受野，从而提升上下文建模能力。【CNN本身存在的问题，前人也在CNN领域内开展了一些研究，但是终归是治标不治本，提升有限，比如空洞卷积等策略】
  - yu提出一种具有可以调节扩张率的空洞卷积，在语义分割中表现出优异的性能。【空洞卷积】
  - 更直接地，Peng设计大的卷积核来捕获丰富的全局上下文信息。
  - Zhao等人在多特征尺度上采用金字塔池化的方法来聚合多尺度的全局信息。
  - Wang提出了非局部算子，通常嵌入在编码器的末尾来捕获远程依赖。【non-local 算子，这个也是一个研究热点】
- 虽然在一定程度上提升了上下文建模能力，但是受困于CNN架构，这些方法仍然不可避免地存在有限感受野地问题。
什么是Transformer架构，以及其优点
- Transformer是一个序列到序列地预测框架，由于其有着强大的长程依赖建模能力，在自然语言处理和机器翻译方面有着良好的记录。
- Transformer中的自注意力机制可以根据输入内容动态的调整感受野，因此在建模长程依赖上优于卷积运算。
Transformer in Vision的进展
- 最近，Transformer以及被认为是一个替代架构，在诸多计算机视觉任务中取得了竞争性的性能，比如图像识别，语义分割/实例分割，目标检测，低级视觉和图像生成。【自VIT引爆cv圈，transformer in vision 成为了热点研究方向】
- 一个典型的例子是ViT，其在识别任务上性能优于基于CNN架构的ResNet，但是代价是使用300M的数据来训练。由于大规模的数据集不总是可用，最近研究尝试去结合CNN和Transformer成为一个混合模型。【目前发展主线有两条：纯Transformer结构的性能优化，CNN与Transformer的联合优化，互相借鉴优势】
  - Carion使用CNN来提取图像特征，Transformer进一步处理提取的特征。
  - Chen设计了TransUnet，其中CNN和Transformer以级联的方式组合起来，为2D医学图像分割制作了一个强大的编码器。虽然TransUNet的设计美好性能好，由于自住意的存在，优化该模型具有挑战性。
    - 首先，它需要非常长的训练时间才能将注意力集中在显着位置，尤其是在 3D 场景中，最初将注意力均匀地投射到每个像素上。
    - 其次，由于其高计算复杂度， Transformer很难处理多尺度和高分辨率的特征图，这在图像分割中起着至关重要的作用。【多尺度与高分辨率是图像中常用的操作，但是Transformer在这个上面却难以处理。】
本文提出的新结构来解决以上问题
- 在本文中，我们为3D医学图像分割提出了一个混合架构，该架构有效的桥接卷积神经网络（CNN）和Transformer（CoTr) 。 CoTr具有一个编码解码结构。
  - 在编码层，一个简单的CNN结构被用来提取特征图，Transformer被用来捕获全局依赖（长程依赖），如下图
  - 受到一些工作的启发，我们向Transformer引入了可变形的自注意力机制。这种注意力机制只将注意力集中在一小部分关键采样点上，从而大大降低了 Transformer 的计算和空间复杂度。【这里的工作就是一个比较有名的baseline：Deformable Transformer】
  - 因此，Transformer 可以处理 CNN 生成的多尺度特征图，并保留丰富的高分辨率信息进行分割。
文章的主要贡献有三点
- 1.我们第一个来探索Transformer 的3D医学图像分割，特别是在计算和空间上的效率方法。【transformer优化技巧：可变形注意力】
- 2.我们引入了可变形自注意力机制来减少Transformer的复杂性，因此使得我们的CoTr可以使用多尺度特征来建模长程依赖。
- 3 我们的CoTr在3D多器官分割任务上优于基于CNN的、基于Transformer、以及混合模型。

2 Materials

Multi-Atlas Labeling Beyond the Cranial Vault (BCV) 数据集被用在这项研究中。
- 包含30个标签的CT扫描用来自动分割11个腹部器官，包括spleen (Sp), kidney (Ki), gallbladder (Gb),
  esophagus (Es), liver (Li), stomach (St), aorta (Ao), inferior vena cava (IVC),portal vein and splenic vein (PSV), pancreas (Pa), and adrenal gland (AG).
- 脾(Sp)、肾(Ki)、胆囊(Gb)、食管(ES)、肝(LI)、胃(ST)、主动脉(Ao)、下腔静脉(IVC)、门静脉和脾静脉(PSV)、胰腺(Pa)和肾上腺(AG)。

3 Methods

介绍CoTr
- CoTr旨在学习更加有效的图像分割的表达，通过桥接CNN与Transformer。
- 如图2所示，其包含了一个CNN编码器来特征提取，一个用于长程依赖建模的可变形Transformer编码层，一个用于分割的解码器组成。
- 我们将深入研究每个模块的细节。

3.1 CNN-encoder

介绍CNN编码器
- CNN编码器 $FCNN(?)\mathcal{F}^{C N N}(\cdot)$ 包含Conv-IN-ReLU的块（block）和三个stages的3D残差模块（residual block)
  - Conv-IN-ReLU块包含3D卷积层，然后是实例归一化（Instance Norm）和修正（Rectified ）线性单元激活函数（ReLU）。
  - 三个Stages的3D残差块的数量分别是3，3，2.【这里的3，3，2 应该是实验效果比较好的组合】
CNN编码器的公式：
- 给定的输入图像x具有的长度为H，宽度为W，深度为D（切片数量）。则由 $FC∣NN(?)\mathcal{F}^{C \mid N N}(\cdot)$ 产生的特征如可以正式表示为：
  ${fl}l=1L=FlCNN(x;Θ)∈RC×D2l×H2l+1×W2l+1\left\{\boldsymbol{f}_{l}\right\}_{l=1}^{L}=\mathcal{F}_{l}^{C N N}(x ; \Theta) \in \mathbb{R}^{C \times \frac{D}{2^{l}} \times \frac{H}{2^{l+1}} \times \frac{W}{2^{l+1}}}$
  其中 $L$ 表示特征层的数量，【常见1，2，3，4，每增加一个层级，特征图缩小一倍】 $Θ\Theta$ 代表CNN编码层的参数， $C$ 代表通道数
  【这里这个公式比较奇怪的是，为啥在HW维度下采样的次数比D多了一次，可能具体要看代码，目前知识受限，暂时还没明白具体缘由。】

3.2 DeTrans-encoder

介绍DeTrans相比CNN和Transformer的优点
- 由于卷积算子内在的局部性，基于CNN的编码层不能够有效的捕获长程依赖。
- 因此，我们提出了DeTrans-encoder，该编码器引入了多尺度可变形的自注意力机制（MS-DMSA）进行有效的长程上下文建模。
- DeTrans-encoder是由输入到序列层和 $L_{D}$ 堆叠可变形Transformer层组成。

3.2.1 Input-to-sequence Transformation

紧接上文介绍DeTrans-decoder中的重要组成之一：将CNN得到的特征图转换为Seq。
- 考虑到Transformer以序列到序列的方式处理信息，首先我们将由CNN编码器 ${fl}l=1L\left\{\boldsymbol{f}_{l}\right\}_{l=1}^{L}$ 产生的特征图展开为1维的序列。不幸的是，注意特征的操作会导致丢失对图像分割至关重要的空间信息。【单纯的展开会丢失空间信息，所以需要位置embedding】
- 为了解决这个问题。我们在展开的 ${fl}l=1L\left\{\boldsymbol{f}_{l}\right\}_{l=1}^{L}$ 中增加了3D位置编码序列 ${pl}l=1L\left\{\boldsymbol{p}_{l}\right\}_{l=1}^{L}$ 。
  - 针对此研究，我们使用不同频率的正弦和余弦函数[19]来计算每个维度pos的位置坐标，如下所示：
    ${PE#(pos,2k)=sin?(pos?v)PE#(pos,2k+1)=cos?(pos?v)\left\{\begin{array}{c}P E_{\#}(p o s, 2 k)=\sin (p o s \cdot v) \\ P E_{\#}(p o s, 2 k+1)=\cos (p o s \cdot v)\end{array}\right.$
    这里的 $#∈{D,H,W}\# \in\{D, H, W\}$ 表示三个维度中的每一个维度， $10000^{2 k / \frac{C}{3}}$ .对于每一个特征层 $l$ , 我们拼接 $P E_{D}, P E_{H}$ , and $P E_{W}$ 作为3D位置编码 $pl\boldsymbol{p}_{l}$ ，并将其与前面展开得到的 $fl\boldsymbol{f}_{l}$ 进行元素求和得到 DeTrans-encoder的输入序列。【固定形位置编码】

3.2.2 MS-DMSA Layer【多尺度可变形自注意力机制】

介绍所谓的关注重点区域的可变注意力这个组件【多尺度可变形自注意力机制】
- 在Transformewr的架构中，自住意层可以看到特征图中所有可能的位置，所以其具有收敛速度慢和高计算复杂性的缺点，进而导致难以处理多尺度特征图。【不是所有位置都有必要，如果这一步是一种可学习的操作，可能性能更好】
- 为了弥补这个问题，我们设计MS-DMSA层仅关注参考位置周围的一小组键采样位置，而不是所有位置。【多尺度可变形自注意力机制层】
介绍这一模块具体怎么设计的
- $zq∈RC\boldsymbol{z}_{q} \in \mathbb{R}^{C}$ 是查询矩阵 $q$ 的特征表示【此处应该是涉及QKV】， $p^q∈[0,1]3\hat{\boldsymbol{p}}_{q} \in[0,1]^{3}$ 是参考点的归一化后的3D坐标
- 给定的多尺度特征图 ${fl}l=1L\left\{\boldsymbol{f}_{l}\right\}_{l=1}^{L}$ 是从CNN编码器的后 $L$ stages 提取的,可以用公式计算第i个注意头的未来表示为。如下：
  $i=∑lL∑kKΛ(zq)ilqk?Ψ(fl)(σl(p^q)+Δpilqk)\text { head }_{i}=\sum_{l}^{L} \sum_{k}^{K} \Lambda\left(\boldsymbol{z}_{q}\right)_{i l q k} \cdot \Psi\left(\boldsymbol{f}_{l}\right)\left(\sigma_{l}\left(\hat{\boldsymbol{p}}_{q}\right)+\Delta_{\boldsymbol{p}_{i l q k}}\right)$
  这里 $K$ 是采样关键点的数量,
  $∈[0,1]\Lambda\left(\boldsymbol{z}_{q}\right)_{\text {ilqk }} \in[0,1]$ 是注意力权重,
  $Δpilqk∈R3\Delta_{\boldsymbol{p}_{i l q k}} \in \mathbb{R}^{3}$ 是第l个特征级别的第k个采样点的采样偏移量【类似神经网络中的b】,
  $σl(?)\sigma_{l}(\cdot)$ 重新缩放 $p^q\hat{\boldsymbol{p}}_{q}$ 到第 $l$ -th 级特征
  遵循前人的设定， $\Lambda\left(\boldsymbol{z}_{q}\right)_{\text {ilqk }}$ and $Δpilqk\Delta_{\boldsymbol{p}_{i l q k}}$ 通过查询特征 $zq\boldsymbol{z}_{q}$ 上的线性投影获得. 因此，MS-DMSA层可以公式化为：
  $2,…,head?H))\mathrm{MS}-\operatorname{DMSA}\left(\boldsymbol{z}_{q},\left\{\boldsymbol{f}_{l}\right\}_{l=1}^{L}\right)=\Phi\left(\text { Concat }\left(\text { head }_{1}, \text { head }_{2}, \ldots, \operatorname{head}_{H}\right)\right)$
  这里 $H$ 是注意头的数量, $Φ(?)\Phi(\cdot)$ 是线性投影层，用于加权和聚合所有头的特征表示

3.2.3 DeTrans Layer

最终汇总介绍DeTrans Layer与DeTrans编码器
- DeTrans-Layer是由MS-DMSA层和前馈网络构成，每个层后面跟着层归一化。
- 跳跃连接策略被用在每一个子层来避免梯度消失
- DeTrans编码器由重复堆叠DeTrans层构成

3.3 Decoder

介绍解码器
* DeTrans-encoder的输出序列根据每一个尺度的大小被重塑成为特征图。【这里应该就是所谓的seq to map，与前面的map to seq 和位置嵌入是一个相反的操作】
* 解码器，一个纯CNN结构，使用转置卷积渐进地将特征图上采样至输入的分辨率（D x H x W），然后使用3D残差块来细化上采样特征图
* 除此之外，为了更好的分割，还增加了编码器和解码器之间的跳跃连接，以保持更多的低层细节。
* 我们使用深度监督策略，通过在不同尺度的解码器的输出端添加辅助损失函数来实现【深度监督策略？？？？】
* 我们模型的损失函数是Dice损失和交叉熵损失的和，关于网络更多的细节在附录A中。

3.4 Implementation details

介绍一下具体的操作细节，比如影像的预处理
- 遵循前人的操作，我们首先使用[-958，327]的范围截断每个扫描的HU值以过滤不相关的区域，然后通过减去82.92和除以136.97来归一化截断的体素值。【这里的操作应该是减均值除标准差？所以这两个数值应该分别是均值和标准差】
- 我们将BCV数据集分为两个部分：21个扫描用于训练，9个扫描用于测试，随机选择6个训练扫描作为验证集，该验证集仅仅用于选择CoTr的超参数。测试集中的最终结果是由所有训练扫描训练得到的模型获得。
介绍训练策略
- 在训练策略中，我们从CT扫描中随机裁剪出48x192x192大小的子体积作为输入，为了减轻有限数据的过拟合问题，我们采用了实时数据增强。
- 包括随机旋转、缩放、翻转、添加高斯白噪声、高斯模糊、调整亮度和对比度、模拟低分辨率和伽马变换，以使训练集多样化。
- 由于Instance Norm（实例归一化）好处，我们采用了一个小的batchsize：2的微batch训练策略【这里的IN BN LN GN 的选用策略，需要具体研究一下】
- 为了权衡训练时间成本和性能回报之间的平衡，CoTr被训练了1000个时期，每个时期包含250次迭代。
- 我们采用动量为0.99、初始学习率为0.01的随机梯度下降算法作为优化器
- 我们在MS-DMSA和前馈神经网络中设置了隐藏大小分别为384和1536，根据经验设置超参数 LD = 6、H = 6 和 K = 4。
- 此外，我们用小的CNN编码器形成了CoTr的两个变体，分别表示为 $CoTr?\mathrm{CoTr}^{*}$ 和 $CoTr?\mathrm{CoTr}^{\dagger}$
  - 在 $CoTr?\mathrm{CoTr}^{*}$ 中，每一个CNN编码器地Stage中仅仅有一个3D残差块
  - 在 $CoTr?\mathrm{CoTr}^{\dagger}$ 中，每一个CNN编码器地Stage中地残差块地数量为2.
介绍测试模块
- 在测试模块中，我们使用滑动窗口机制，其中窗口大小等于训练块大小。
- 此外，为了提高分割的稳健性，还采用了高斯重要性加权[12]和沿所有轴翻转的测试实时增强（TTA)
- 为了定量评估分割结果，我们计算了骰子相关性分数(Dice)度量，该度量度量预测与其基础事实之间的重叠。【用Dice指标来衡量】

4 Results

4.1 Comparing to models with only Transformer encoder.

对比实验，先与纯Transformer编码器的模型对比
- 我们首先针对最先进的SETR的两个变种评估了我们的CoTr，该SETR使用了随机初始化和预训练模型ViT-B/16作为编码器。
- 我们同样也对比了一种移除CNN编码器的CoTr变体（CoTr w/o CNN-encoder).）为了确保公正的比较，所有型号都使用相同的解码器。其中这些模型的分割结果如表1所示：
我们可以从中得到三个结论：
- 1. 虽然Transformer架构不限制输入图像的形式，但是在2D自然图像中的预训练模型ViT-B/16不能在3D医学图像上很好的工作（work well），次优的性能可能是由于2D自然图像和3D医学图像之间的领域偏差（domain shift）。
- 2 CoTr w/o CNN-encoder【没有CNN编码器】有大概22M的参数，优于有100M参数的SETR模型，我们相信轻量化的Transformer模型对于医学图像分割任务更友好，因为医学图像分割任务往往只有较小的训练数据集。
- 3 我们的 $CoTr?\mathrm{CoTr}^{*}$ 明显优于CoTr w/o CNN-encoder，具有可比较的参数【大致意思是参数差不多，但是性能更优】，在超过11个器官提升了Dice达到4%，它表明混合CNN-Transformer编码器在医学图像分割中的比纯Transformer编码器上具有明显的优势

4.2 Comparing to models with only CNN encoder.

既然是一个混合架构，那么必须跟单一架构都要比较一下
- 然后，我们将COTR与去除DeTrans编码器的COTR的变体（CoTr w/o DeTrans）和三个基于CNN的上下文建模方法进行比较
  - 这三个基于CNN的上下文建模方法是：空洞空间金字塔汇集（ASPP）模块，金字塔解析（PP）模块和非局部（non-local）模块。
- 为了公平比较，我们使用了相同的CNN编码器和解码器，分别使用ASPP , PP, Non-local模块来取代我们的DeTrans编码器。，结果如上面的表1，结果表明我们的CoTr分割性能比（CoTr w/o DeTrans）高，在所有器官上高1.4%。
- 这证实了我们的CoTr使用了混合的CNN-transformer编码器比纯CNN编码器在医学图像分割中学习有效表示能力更强。
- 此外，与这些上下文建模方法相比，我们的变压器架构有助于更准确的分割。

4.3 Comparing to models with hybrid CNN-Transformer encoder.

既然单一模型比完了，还是要与同是混合模型的其他模型进行比较
- 我们同样将我们的CoTr与其他混合CNN-Transformer结构进行了对比，例如Trans-UNet
- 为了直接处理3D图像，我们像CoTr一样使用3D CNN编码器和解码器扩展了原始的2DTransUNet成为3D TransUNet。我们将3D TransUNet中Transformer的一些头和层设置的与我们的CoTr相同。结果表明我们的CoTr在所有器官的分割上都强于3D TransUNet,特别是对于胆囊和胰腺分割.
- 甚至一些小的CNN编码器 $CoTr?\mathrm{CoTr}^{\dagger}$ 都取得了比TransUNet更优的性能在七个器官的分割上。
介绍取得卓越性能的原因
- 卓越的性能归功于CoTr中的可变形机制，该机制减少了计算和空间复杂度，使得处理高分辨率和多尺度的特征图成为可能。

4.4 Computational Complexity.

最后在强调一下本文的重点：计算复杂度优化
- 提出的CoTr是使用一块2080ti训练的，使用了Pytorch，大概花了两天的时间进行训练
- 每一个小体积48 x 192 x 192 的分割时间小于30ms【介绍一下分割速度】

5 Discussion on Hyper-parameter Settings

讨论一下研究过程中发现的一些东西
- 在Detrans-encoder中，有三个超参数，例如K，H， $L_{D}$ ，其分别表示采样关键点，头部和堆叠DeTrans 层的数量。要研究它们的设置对分割的影响，我们将K设置为1,2和4，H为2,4和6，并将 $L_{D}$ 为2,4和6. 如图3所示。
- 我们绘制了K，H， $L_{D}$ 对应值在验证集上获得的所有器官的平均Dice。
  - 它表明，增加K，H， $L_{D}$ 的数量可以提高分割性能。
  - 为了证明从多尺度策略引起的性能增益，我们还试图从最后一级的单级特征映射训练CoTr【类似TransUNet】。图3（d）表明，使用多尺度特征映射而不是单尺度特征映射可以有效的提升平均Dice达到1.2%。

6 Conclusion

在本文中，我们提出了一个CNN Transformer的混合模型CoTr，利用该模型进行3D医学图像分割。
- 在这个模型中，我们设计了可变形的Transformer，采用可变形的自注意力机制减小模型在多尺度和高分辨率特征图上建立长程依赖时的计算和空间复杂度。
- 对比实验在BCV数据集上进行，我们的CoTR在基于CNN和基于Transformer的模型中的卓越性能表明：
  - 通过结合CNN和Transformer的优势，提出的CoTr实现了保持低层特征细节和建立长程依赖的平衡，作为一个强壮的baseline，我们的CoTR可以在未来被扩展至解决其他结构（例如，脑结构和种类分割）。