以下文章来源于CVHub ，作者派派星
Title: A Robust Volumetric Transformer for Accurate 3D Tumor Segmentation
PDF: https://arxiv.org/pdf/2111.13300
Code: https://github.com/himashi92/VT-UNet
导读
本文提出了一种基于 UNet 架构的 Transformer 网络（VT-UNet）用于3D医学图像分割。3D医学图像分割一直以来是一项极具挑战性的任务，大部分现有的3D分割网络由于局部与全局上下文信息捕获不足，而影响分割精度；且3D分割网络计算效率低下阻碍了模型的工业落地部署进程。
为此，本文提出了一种新的网络VT-UNet用于医学图像分割，其基于Transformer进行网络构建，并采用经典的UNet（Encoder-Decoder）作为整体的网络架构：其中Encoder借助强大的自注意力机制高效的捕获局部与全局上下文信息，而Decoder则采用一种并行的自注意力和交叉注意力机制来捕获细节，并进行边界细化。借鉴SWin Transformer窗口化的思路，这两种注意力机制采用了两类窗口机制来降低计算复杂度。该网络凭借高计算效率，高精度和高鲁棒性，在医学分割十项全能（MSD）的脑肿瘤分割（BraTS）竞赛任务上表现SOTA！代码已开源，欢迎小伙伴尝试！
引言
本文所提出的VT-UNet网络其实也是受Transformer的启发。相比于CNN网络，Transformer网络在扩展性和鲁棒性上表现更佳，且其内在的注意力机制可以灵活的捕获局部和全局上下文信息， 这对于3D图像分割的精度至关重要。过去已经有一些将Transformer应用于3D医学图像分割的研究， 且这类方法的分割精度均要优于基于CNN的方法，但这些方法普遍将一个三维体积划分为二维切片并将该2D切片作为输入进行处理，因此可能会丢失重要的体积信息（eg，2D缺乏depth维度的信息）；当然也有一些混合（CNN与Transformer结合）方法可以保留3D体积信息的完整性；但纯用Transformer进行3D分割网络的构建，并保留3D体积信息的完整性作为网络的输入，是还未曾探索过的领域。
而本文首次提出了这一领域的3D分割网络模型设计，大量实验证明：VT-UNet不仅实现了更好的分割性能，还对数据伪影具有更强的鲁棒性。
The model size vs Dice Similarity Coefficient (DSC) is shown in this plot
虽然Transformer模型具有高度动态和异常灵活的感受野，且能捕获长远距离的信息并进行交互，但为3D体积分割设计的Transformer架构仍然是一项具有挑战性的任务，这主要是因为：
封装一个体素信息并捕捉体积序列中任意位置的连接并不简单，与使用基于Transformer的2D图像分割方法相比：体积中每个切片的数据都需连接到三个视图（View），丢弃其中任何一个都可能会对分割结果产生不利影响
保留体积中的空间信息是一项艰巨的任务，即使对于2D图像，在将图像切成patch并将patch投射到token中的ViT引入了Tokens-to-token，ViT中显示的局部结构线索可能会丢失。因此同时有效地编码局部线索和捕捉体积多个轴上的全局信息交互是一项具有挑战性的任务
由于自注意力机制的高计算量，像3D体积输入的大尺寸所带来的高昂计算也是基于Transformer的3D分割模型需要仔细考虑的因素。
因此，本文提出了VT-UNet模型来有效的解决这些挑战：在Encoder中开发了两种类型的Transformer模块，直接对3D体积进行分层式处理，以捕获局部和全局上下文信息；在Decoder中，引入并行的自注意力和交叉注意力，以创建Decoder的query和Encoder的key之间的桥梁。通过交叉注意力和自注意力的并行处理，模型能够在解码过程中保留全局上下文，这对于分割任务非常重要；同时作者还在每个Transformer块中计算注意力时应用了相对位置编码，并在解码过程中注入从序列Token的傅里叶特征位置提取的补充信息，以增强特征表达。
本文的主要贡献包括：
从序列到序列的视角重新构建了体积肿瘤分割，并提出了一种基于UNet架构的体积Transformer模型，适用于多模态医学图像分割
设计了一个编码器块和一个解码器块，使其能够同时捕获局部和全局上下文特征。同时在解码器中引入了并行处理的自注意力和交叉注意力来保留全局上下文，并提出了一种凸组合方法和傅里叶位置编码来增强特征表达。
通过对以上注意力采用类似于Swin Transformer窗口机制的设计，有效地减少了模型参数，并保持了较低的FLOPs，相比现有方法以更高的计算效率取得了更好的体积分割结果
大量实验表明：VT-UNet在体积分割中实现了SOTA结果，同时具有更强的鲁棒性。
方法
图1. llustrates VT-UNet Architecture
如上图1展示了VT-UNet的网络架构图，假设输入是一个大小为的三维体积，输出是一个大小为的三维体积，代表分割类别数。VT-UNet的Encoder由一个带有Linear Embedding Layer的3D Patch Partitioning，和带有两个连续VT Encoder Block的3D Patch Merging组成。
3D Patch Partitioning
如何处理一组token序列：VT-UNet模型的第一个模块接受一个维度的医学图像（如MRI），通过将这个3D体积分割成非重叠的3D块（见上图2（b）），创建了一组tokens。partitioning kernel的大小为，所以一个3D体积可以用个tokens来描述。3D Patch Partitioning后连接一个linear embedding，将每个维度为的token映射到一个C维度的向量。根据实验，、和的值会分别设置为4、4和72。
VT Encoder Block
在ViT中，由于tokens的构造方式，它们携带着重要的空间信息。在ViT中进行SA窗口化操作的重要性已经在几项最近的研究中得到证明，其中最显著的是在Swin Transformer中。我们提出了在VT编码器块（VT-Enc-Blks）中针对体积数据进行三维窗口化操作，遵循与Swin Transformer类似的设计原则。特别地，我们提出了两种类型的窗口化操作，即常规窗口化和移位窗口化，为简单起见分别用VT-W-MSA和VT-SW-MSA表示。图2b展示了VT-W-MSA和VT-SW-MSA的设计细节。VT-W-MSA和VT-SW-MSA都使用具有窗口化操作的注意层，紧随其后的是带有高斯误差线性单元（GELU）非线性的两层多层感知器（MLP）。在每个MSA和MLP之前都应用了层归一化（LN），并在每个模块之后应用了残差连接。窗口化使我们能够在对tokens之间的长距离依赖进行建模时注入归纳偏差。在VT-W-MSA和VT-SW-MSA中，窗口化操作有助于表示学习中的tokens之间的注意力。在VT-W-MSA中，我们将体积均匀地分成较小的不重叠窗口，如图2（b）所示。由于VT-W-MSA中相邻窗口中的tokens不能相互看到，因此我们在VT-SW-MSA中利用了移位窗口（参见图2（b）最右侧），它将VT-W-MSA相邻窗口中的tokens连接起来。VT-Enc-Blk的实现可用如下公式表示：
其中和分别表示块的VT-W-MSA模块和MLP模块的输出特征。
3D Patch Merging
3D Patch Merging是指在VT-UNet的编码器中使用3D块合并，以生成特征层次结构。拥有这样的层次结构对于在密集预测任务中生成更精细的输出细节至关重要。 在每个VT-Enc-Blk之后，我们以非重叠的方式沿空间轴合并相邻的标记以产生新的标记。在这样做时，我们首先将每组相邻标记的特征连接起来。然后，通过线性映射将结果向量投影到一个空间中，该空间中标记的通道维度加倍。这样做的好处不仅限于特征层次结构。SA的计算复杂度在标记数量的平方级别上，因此,块合并每个VT-Enc-Blk之后将VT-UNet的FLOPs计数降低了16倍。VT-UNet模型的FLOPs仅为其完全体积的CNN对应物的6.7％，而其性能相当（实际上略好）！请注意，块合并不在bottleneck阶段使用。
VT解码器：在由VT-Enc-Blk和3D Patch Expanding层组成的bottleneck层之后，VT解码器从连续的VT解码器块（VT-Dec-Blks），3D补丁扩展层和分类器开始，以产生最终预测。VT-Enc-Blk和VT-Dec-Blk之间存在一些基本设计差异，详见下面部分
3D Patch Expanding
3D Patch Expanding被用来在某种程度上消除块合并的影响。换句话说，为了以与输入相同的空间分辨率构建输出，我们需要在解码器中创建新的tokens。为了进行讨论，我们考虑bottleneck block后的Patch Expanding（参见图2的中间部分）。Patch Expanding的输入标记维数为8C。在Patch Expanding中，我们首先使用线性映射将输入tokens的维数增加一倍。经过reshape后，我们可以从维数为2×8C的结果向量中获得维数为4C的2×2 token。因此，我们将沿着空间轴reshape它，因此对于D/4×H/32×W/32×8C，我们创建D/4×H/16×W/16×4C的token，这样可以达到与输入相同的空间分辨率的效果。
VT Decoder Block
在这里，作者通过引入一个混合形式的自注意力机制(SA)进一步提高模型的预测质量。具体而言：在解码器侧，每个VT-Dec-Blk模块都接收其上一个VT-Dec-Blk模块生成的tokens，以及与VT-UNet中位于同一级别的VT-Enc-Blk模块中的key(KE)和value(VE) tokens。VT-Enc-Blk模块包括两个带有正常和移位窗口操作的SA块。VT-Dec-Blk模块则包括四个分组到SA模块和交叉注意力(CA)模块中的SA块。这两个模块使用正常和移位窗口操作来注入更多的归纳偏差。具体操作可用如下公式表示：
其中和分别表示解码器模块的右分支和左分支。SA的右分支根据上公式(3)对前一个VT-Dec-Blk模块生成的token进行操作。在这里，作者通过下标强调了从解码器流出的信息。CA的左分支使用解码器生成的查询(query)以及从VT-Enc-Blk模块中获得的相同级别的key和value。这里的想法是使用由values标识的编码器生成的基础以及键，以从编码器获取的空间信息获益。这些模块也使用正常和移位窗口操作来注入更多的归纳偏差。
值得注意的是，具有相同窗口操作的SA的值和键应该组合起来，因此在图2(c)中采用了交叉连接的形式。作者还指出，一些读者可能会问为什么要考虑来自编码器的values和key。作者研究了其他可能性，如：使用编码器生成query和key，而由解码器生成value。但是，经验证明上式(3)中描述的形式能够提供更好和更稳健的结果，因此考虑来自编码器的values和key是VT-UNet的选择。
Fusion Module
CA模块和MSA模块都生成了一些token，这些token需要被结合起来，然后传递给下一个VT-Dec-Blk模块。结合的方式是使用一个线性函数进行加权平均，其中α控制了来自CA和MSA的token对于的贡献程度。可通过下面公式（4）计算得到：
其中 代表来自CA的token， 代表来自MSA的token，FPE是一个函数，用于对位置信息进行编码。同时，在这里作者选择了一个简单的加权方案，即。
Classifier Layer
在解码器的最后一个3D patch expanding层之后，还引入了一个分类器层，其中包括一个3D卷积层，用于将深度为C的特征映射到K个分割类别。此外，作者还提到了对于SA模块计算复杂度的说明：SA模块中的计算复杂度由获取Q，K，V的计算，和计算以及将Softmax（一个矩阵）的输出应用于V得到结果的计算所共同决定。这些计算的总和为，其中和分别是维度和token数量。通过使用窗口化，SA的计算负载将减少至，其中我们假设token被分组成个窗口，SA在每个窗口内运行。 对于从体积数据生成的token，，因此窗口化不仅有助于模型捕获更具判别力的特征，还有助于减少计算负载。
实验
egmentation Results on MSD BraTS Dataset & Ablation Study
ualitative Segmentation Results
总结
本文提出了一种新颖的基于自注意力机制的3D语义分割模型 VT-UNet。该模型使用3D卷积神经网络构建编码器和解码器，其中编码器通过逐步下采样操作将输入数据压缩为低分辨率表示，而解码器则将低分辨率表示映射回原始分辨率。解码器和编码器之间使用横向连接，从而引入丰富的空间信息，提高分割精度。作者还提出了一种混合的自注意力机制，将解码器中的自注意力和跨注意力结合起来，以在保留空间信息的同时嵌入丰富的语义信息；这两种注意力机制都采用了窗口机制来降低模型计算复杂度，并提高计算效率。最后，为了将深层特征映射到K个分割类别，在解码器中额外引入了一个分类器层。由于VT-UNet模型具有高效的计算能力和强大的鲁棒性，所以可以处理大型的3D医学图像数据，并在实际应用中提高模型推理速度和准确性。
转自：“arXiv每日学术速递”微信公众号
如有侵权，请联系本站删除！
