VISION TRANSFORMER WITH PROGRESSIVE TOKENIZATION FOR CT METAL ARTIFACT REDUCTION

基于渐进式标记的CT金属伪影减少的VISION TRANSFORMER

论文链接：https://ieeexplore.ieee.org/document/10095579

代码链接：https://github.com/boolean1024/MARformer

ABSTRACT

高质量的计算机断层扫描(CT)在临床诊断中起着至关重要的作用，但金属植入物的存在会在CT图像上引入严重的金属伪影，阻碍医生的决策。以往许多金属伪影还原(MAR)的研究都是基于卷积神经网络(CNN)。最近，Transformer在计算机视觉方面展示了惊人的潜力。此外，基于Transformer的方法也被用于CT图像去噪。然而，这些方法在MAR中很少被探索。为了填补这一空白，我们提出了，据我们所知，第一个基于Transformer的MAR架构。我们的方法依赖于一个标准的视觉Transformer(ViT)。此外，我们利用渐进式标记化来避免ViT的简单标记化，这导致无法对局部解剖信息进行建模。此外，为了促进token之间的交互，我们利用了Swin Transformer的循环移位cyclic shift。最后，许多实验结果表明，基于Transformer的技术在一定程度上优于基于CNN的技术。

1. INTRODUCTION

计算机断层扫描(CT)已广泛应用于临床诊断。但是，如果患者携带金属植入物，如牙科填充物和髋关节假体，则CT图像会引入严重的金属伪影，由于图像上缺乏解剖结构，不利于医学分析。更重要的是，一些患者不可避免地要植入金属材料。为了克服这些伪影，近年来许多研究者致力于减少金属伪影(MAR)[1]。

传统的方法可分为物理效果校正、弦图补全和迭代重建三大类。物理效应校正试图修正物理努力，如光子不足伪影[2]和线束硬化[3]，但它由于存在高原子序数的金属不能达到预期的目标。弦图补全，也称为投影域的插值，将金属影响的投影视为空白，并用估计值代替它们。它的典型方法——线性插值(Linear Interpolation, LI)[4]用其周围未受影响的数据填补了空白，但由于插值值与未受影响的值不一致，LI可能会在重建图像中引入二次伪影。[5]迭代重建基于约束优化，如最小二乘法和最大似然法，逐步恢复CT图像。[6][7]然而，基于迭代的方法有共同的缺点，它们需要手工制作正则化器，并且通常很耗时，这导致难以适应复杂的MAR任务或临床应用。

随着深度学习(Deep Learning, DL)的兴起，一些计算机视觉任务正在蓬勃发展。由于深度学习的普及，诸如卷积神经网络(Convolutional Neural Network, CNN)[8]等基于深度学习的方法层出不穷，促使研究人员深入研究更多新颖的人工智能解决方案。尽管缺乏可解释性，但深度学习是数据驱动的，具有优越的特征提取能力，即可以从未知数据中自动学习最优特征。与传统的以投影域为主的MAR方法不同，基于dl的方法可以同时处理投影域和图像域的数据，也可以同时处理投影域和图像域的数据。Zhang等人提出了一种通用的开放网络——CNNMAR，合并了多种方法来提高MAR的性能[9]。U-Net[10]是语义分割中最流行的编解码网络之一，被广泛用于MAR[11] [12]，但这种方法可能会丢失一些有用的信息。此外，为了解决临床难以获得配对CT图像的问题，Liao等人提出了一种无监督学习方法——Artifact Disentanglement Network(ADN)，将组织和伪影分开[13]。尽管技术丰富，CNN仍然垄断了MAR任务，无法提取全局特征。

在transformer[14]这一通常应用于自然语言处理（NLP）的基于注意力的方法的推动下，越来越多的基于transformer的方法在计算机视觉任务中蓬勃发展。Vision Transformer(ViT)[15]直接利用标准的transformer编码器进行图像分类，无需额外修改，与最先进的CNN相比，效果更好。token到token视觉转换器(T2T-ViT)[16]使用渐进式token化来代替简单的token化。Swin Transformer[17]是一种分层transformer，其表示是用移位窗口计算的，它试图作为计算机视觉的通用主干。Wang等人将他们的工作扩展到这些基于transformer的想法。Wang等人提出了一种无卷积的token2token扩张视觉transformer，称为CTformer，用于低剂量CT图像去噪[18]。即使基于transformer的方法已经广泛应用于许多计算机视觉任务，它可以提取全局特征并对远程依赖关系进行建模，但它们很少在MAR任务中得到利用。

为了填补这一空白，我们结合上述基于transformer的策略发明了MARformer。CTformer与MARformer非常相似，但我们的工作进一步展示了对金属伪影的去除效果，而不是对噪声的去除效果。这两篇工作虽然相似，但它们毕竟属于不同的领域，即无论是在几何上还是在分布上，光子噪声与金属制品都是不一样的。MARformer建立在一个标准的ViT上，利用渐进式标记化来取代ViT的简单标记化。简单的标记化方法只包含一个标记化过程，无法对重要的局部解剖信息进行建模。为了在token到token过程中增加token之间的交互，将循环移位合并到MARformer中。无论是定量评价还是定性评价，许多实验结果都表明MARformer在某种程度上优于基于CNN的方法。

2. METHOD

2.1. 概述

图1和图2(a)说明了我们提出的基于残差U-Net的架构。从编码器到解码器在同一级别上有快捷方式。编码器包括一个token to token (T2T)过程[16]，该过程逐步对输入进行标记，并对重要的局部解剖信息进行建模。在这些token之间缺乏交互的情况下，我们参考Swin transformer[17]的循环移位方法，增加彼此之间的通信。此外，我们使用扩展展开来扩大感受野并捕获多尺度上下文。而解码器则采用与编码器各对应单元相反的操作，逐步对输入进行去符号化。这个U-Net的瓶颈是一个标准的transformer。

2.2. 残差U-Net

我们利用残差U-Net作为我们网络架构的骨干。确切地说，我们在残差网络上添加了从编码器到解码器的每个相同级别的连接[19]。这样可以避免梯度消失，并从编码器的特征图中保留更多的解剖结构细节。简而言之，残差U-Net将使网络更容易训练，并有助于获得更好的结果。

2.3. Transformer模块

图2(b)描述了我们架构中的一个标准transformer[15]。transformer包括多头自注意(MSA)和多层感知器(MLP)的交替层，MSA和MLP之前使用的层范数(LN)， MSA和MLP之后的剩余连接。具体来说，transformer完全依赖于注意力机制，这种机制不仅具有并行计算能力，而且具有最短的最大路径长度。在MSA中，将一系列token$T_0\in R^{b\times n\times d_0}$ (b为批大小，n为每批token的个数，$d_0$为token的维数)线性映射到三个张量: 查询$Q\in R^{b\times n\times d}$，键$K\in R^{b\times n\times d}$，值$V\in R^{b\times n\times d}$ (d为Q, K, V的维数)，此过程记为
$$\{\begin{array}{l}Q=T_0{W}^q\\ K=T_0{W}^k\\ V=T_0{W}^v\end{array}\text{\hspace{1em}(1)}$$
其中$Wq、Wk、Wv\in R^{d_0\times d}$为可学习参数。我们选择缩放点积注意力作为评分函数，因为当查询和键的向量长度相同时，它的计算效率更高。MSA的输出计算为
$$MSA(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d}})V\text{\hspace{1em}(2)}$$
当$T_0$通过该transformer块时，输出token$T_{tb}$表示为
$$\{\begin{array}{l}{T}_0^{{}^{\prime }}=MSA(LN(T_0))+T_0\\ T_{tb}=MLP(LN(T_0^{{}^{\prime }}))+T_0^{{}^{\prime }}\end{array}\text{\hspace{1em}(3)}$$

2.4. Tokens-to-Token模块

图3显示了与解码器相反的编码器的具体token到token过程。ViT的标记化比较简单，只包括一个标记化过程，这导致了周围标记之间缺乏依赖性，因此无法对重要的局部结构信息进行建模。为了克服这一障碍，我们采用了T2T-ViT[16]的思想，即按照将相邻token聚合为一个token的方式，逐步对输入进行token化。让我们来详细说明一下。

通过transformer块后，将token $T_{tb}\in R^{b\times n\times d}$变换为$T_i$，然后$T_i$继续重塑为$T_i^{{}^{\prime }}$。流程定义为
$$\{\begin{array}{l}{T}_i=T_{tb}^T,T_i\in {\tilde{R}}^{b\times d\times n}\\ T_i^{{}^{\prime }}=reshape(T_i),T_i^{{}^{\prime }}\in R^{b\times d\times h\times w}\end{array}\text{\hspace{1em}(4)}$$
式中$h=w=\sqrt{n}$ (h和w分别为特征映射的高度和权值)。

为了扩大这些token之间的依赖性，我们尝试了Swin Transformer[17]的循环移位。在执行循环移位之后，token被转换成$T_i^{{}^{\prime \prime }}$。该过程表示为
$$T_i^{{}^{\prime \prime }}=cycleShift(T_i^{{}^{\prime }}),T_i^{{}^{\prime \prime }}\in R^{b\times d\times h\times w}\text{\hspace{1em}(5)}$$
我们使用卷积运算来代替ViT中的patch split和linear projection[15]。此外，为了扩展感受野和捕获多尺度上下文，我们借用了在某种程度上类似于扩展卷积的扩展展开操作，而不是卷积或对后续token$T_i^{{}^{\prime \prime }}$的展开。在完成扩张展开后，token被转换成$T_i^{{}^{\prime \prime \prime }}$。该过程表示为
$$T_i^{{}^{\prime \prime \prime }}=dilatedUnfold(T_i^{{}^{\prime \prime }}),T_i^{{}^{\prime \prime \prime }}\in R^{b\times n_i\times d_i}\text{\hspace{1em}(6)}$$
其中$n_i$和$d_i$分别是每个批处理中的新token数量和新token的维度。

2.5. 目标函数

该模型在监督设置中训练，学习从带有金属伪影的CT图像$x$到其配对的干净CT图像$y$的映射。我们通过优化均方误差(MSE)来训练模型。整体损失函数L表示为
$$L=||f(W,b;x)-y|{|}_2\text{\hspace{1em}(7)}$$
其中f为神经网络，W和b为参数。

3. EXPERIMENTS

3.1. 数据集和实现细节

由于在监督环境下进行训练以及难以获得配对的临床CT图像，我们不得不通过在干净的CT图像中掺入模拟金属伪影来伪造配对数据集。我们选择了用于病灶检测的大规模CT数据集DeepLesion[20]作为我们的初步数据集。

根据CNNMAR的仿真方法[9]，我们使用Matlab脚本从初步数据集中随机选取无伪影的CT图像，模拟线束硬化效应、泊松噪声和金属局部体积效应，合成受伪影影响的CT图像。所有的图像都被调整为256 × 256。最后，我们选择了4200对，其中4000对用于训练和验证，剩下的200对用于测试。

我们采用结构相似指数(SSIM)和峰值信噪比(PSNR)进行定量评价。SSIM是一种流行的评估两幅图像之间相似性的方法，反映了人们对图像质量的判断。SSIM的取值范围为0 ~ 1，值越高表示两个图像之间的相似性越高。PSNR是评价图像质量的客观标准，值越高表示图像质量越好，即重建图像与原始图像的一致性越高。

模型在Ubuntu 20.04.4 LTS上进行训练和测试，GPU为Nvidia Tesla T4, CPU为Intel Xeon® Silver 4309Y。此外，我们利用PyTorch实现了我们提出的模型，并利用CUDA加速了训练过程。还考虑了一些必要的数据扩充。此外，我们应用对学习率不敏感的Adam算法作为更新器，以最小化损失函数，并将学习率初始化为1.0×10−5，这在1.0×10–6的预定衰减率下逐步减弱。最后，我们训练模型4000次，批大小设置为16。

3.2. 实验结果

我们将所提出的方法与其他类似规模的方法进行了比较。基线包括LI[4]和NMAR[5]等传统方法，以及基于DL的方法，包括基于监督学习的技术，如CNNMAR[9]、U-Net[12]，以及基于无监督学习的技术，如ADN[13]。除了传统的方法外，大多数都依赖于CNN。

对于定量评价，表1可以证明我们基于transformer的方法比基于CNN的方法更有潜力。的确，无论在SSIM还是PSNR上，我们的方法都具有更好的性能。

如图4所示，我们将它们的结果可视化以进行定性评估。

表2是DeepLesion中消融研究的定量结果。通过比较（a）和（b），以及（a）与（c），我们可以发现残差连接和循环移位确实在MAR中部分起作用。此外，当残余连接和循环移位都存在时，无论在SSIM还是PSNR上，我们的模型的性能都有显著提高。

4. CONCLUSION

在本文中，我们提出了一种基于transformer的mar方法。特别是，我们采用残差U-Net作为整体架构，并在标准ViT上找到了我们的方法。但是ViT的标记化过于简单，无法对重要的局部解剖信息进行建模，这就要求我们采用渐进式标记法来代替简单的标记法。为了解决token到token过程中token之间交互稀缺的问题，我们将循环移位用于我们的方法。许多实验结果表明，在某种意义上，我们基于transformer的技术比基于CNN的技术具有更大的潜力。