[Transformer]UTNet:A Hybrid Transformer Architecture for Medical Image Segmentation

Paper
Code

Abstract

Transformer已被广泛用于NLP领域的诸多任务，但是将Transformer应用于图像领域还有很大的探索空间。本文提出一种简洁高效的混合Transformer框架-UTNet，将self-attention与CNN结合在一起用于提升医学图像分割的性能。
UTNet在encoder和decoder部分均引入了self-attention模块，以最小的开销捕获不同尺度的长程依赖关系。
为此，本文提出了一种有效的结合了相对位置信息编码的自注意力机制，可以将自注意力的时间复杂度从O(n^2)降到O(n)；并提出了一种新的自注意力解码器，可以从skip connection中恢复encoder中的细粒度特征。

本文的意义在于解决了目前Transformer应用的困境：即需要从大量的训练数据中学习归纳偏执。UTNet中的混合层（hybrid layer）允许将Transformer初始化为CNN，不需要预训练。
本文还评估了UTNet在多标签多来源的心脏MRI分割图像上的性能，UTNet展现了优越的分割性能和鲁棒性，也展现了推广到其他医学图像分割任务的巨大前景。

Section I Introduction

CNN因为其优异的特征提取能力，被广泛的应用于计算机视觉任务；目前基于卷积操作搭建的encoder-decoder结构非常适合解决位置敏感类任务，如语义分割。借助卷积操作可以捕获相邻像素之间的纹理信息、局部特征；再通过堆叠不同层次提取的局部特征，逐步扩大感受野，获得更高层次的全局特征。但这种方法存在两方面的局限性：一是卷积只能提取临近像素之间的信息，无法对全局关联进行有效的建模；二是卷积核的尺寸和大小往往是固定的，无法根据输入内容进行调整。
Transformer已经被广泛用于NLP领域，其self-attention机制可以有效的捕获长程依赖关系。自注意力机制会计算像素之间的交互关系，聚合之后作为输出，可以根据输入内容动态聚合相关特征，也可以有效的捕获长程关联。
已经有研究初步表明，self-attention在分割、检测、重建等领域均十分有效。
尽管将Transformer应用于视觉任务十分具有前景，但是具体的应用和实现还存在一些挑战。

首先，self-aatention的计算时间复杂度为O(n^2)，n是序列长度，因此会有较大的训练和推理计算开销。
一些前人的工作尝试简化self-attention的时间复杂度，但还有很大的优化空间。
受限于self-attention的复杂度，在ViT中只能将图像切分为patch作为输入序列；或者采用CNN提取到的已经降为过的feature map作为输入。
但是对于医学图像分割任务，细节信息往往对精确分割更为重要，因为大多数误分割区域都位于边缘；其次，由于Transformer没有引入任何归纳偏执，导致其在小规模数据集上表现不佳，一般都需要现在大型数据集上进行预训练再进行迁移。但是即使在ImageNet预训练过，Transformer的表现仍然稍逊于ResNet，遑论医学图像数据集这种数据集规模更小的情况。

本文提出了混合Transformer的U型网络-UTNet，充分将self-attention和卷积结合在一起，利用卷积层提取足够的局部特征，避免Transformer对大规模数据集的预训练需求；同时借助self-attention捕获长程、全局特征。

网络设计遵循标准的UNet结构，但是将每层卷积模块中的最后一个卷积操作替换为Transformer结构。为了进一步提升分割效果，本文还在较高分辨率的feature map使用self-attention来提取细节之间的长程依赖，并且将self-attention的计算时间和空间复杂度从O(n^2)降低至O(n)；此外还引入了相对位置编码来学习医学图像中内容-位置的关联。
通过在多器官分割、心室分割任务中的实验表明，UTNet均表现出了优异的分割性能和鲁棒性；UNet这种设计也有望推广到其他医学图像分割任务。

Section II Method

Part 1 Self-Attention

Transformer基于多头注意力模块建立的（MHSA,Multi-head self-attention），多头使得模型可以使用不同子空间提取的注意力，最终多头的输出结果会级联后再送入前馈网络FFN层。



本文设置head=4
对于输入X(CWH)首先经过1x1卷积获得映射后的Q,K,V，然后拆分成不同head。
注意力计算公式：


计算出的注意力经softmax处理后称之为：
上下文聚合矩阵，或者叫相似性矩阵，表明了每一个q与所有key的匹配/相似程度；这一相似程度随后会作为权重 乘以value的值，这样就计算出了注意力，并且基于的感受野是全局的感受野，是将所有的输入都考虑其中的。

这种基于self-attention计算得到的上下文聚合矩阵会随着输入内容动态调整，可以更好的进行特征聚合；但是n*d的点积运算 时间复杂度是O(n^2)，因为n作为序列长度一般远大于维度d。

Part 2 Efficient Self-Attention Mechanism

图像是高度结构化的数据，在局部区域的大部分像素都具有相似的特征，所以在计算所有像素点的成对注意力是存在大量冗余的，因此计算也十分低效。

已经有研究通过理论分析证明，对于一些较长序列，其self-attention的都是低秩矩阵，表明其绝大多数信息都集中在其奇异值上。本文受此启发，提出了一种更高效的self-attention计算方法，详情参见Fig 2.





核心思想是分别对key和value进行映射，并且计算时使用的是K,V的低维近似，这样就可以减少计算量，将时间复杂度降到O(nkd)。




值得注意的是：低维近似可以是任何降采样操作，如均值池化、最大池化、带步长的卷积等，本文则是使用1x1卷积接双线性插值来对特征图进行降采样。
Fig 2分别展示了在Transformer encoder和decoder中使用的多头自注意力模块（MHSA），他们流程相似，但是decoder中输入有两个，一个是通过skip connection连接过来的encoder的特征图，一个则是decoder前一层的低分辨率的feature map。

Part 3 Relative Positional Encoding

原始的self-attention模块完全将位置信息丢弃，使得模型无法对结构信息进行建模；最初使用的正弦位置嵌入方式也并不具有像CNN的平移不变性。本文使用2D的相对位置编码方式，在进行softmax之前为每一个像素点额外的加入高度和宽度信息：
与efficient self-attention计算相似，相对宽度和高度也是取低维映射后的结果。





## Part 4 Network Architecture


Fig 1展示了UTNet的网络结构，
本文致力于将卷积与self-attention的优势结合在一起，这样一方面借助卷积可以学习归纳偏置，避免在大规模数据集上预训练Transformer；另一方面还可以利用Transformer对捕获全局特征。
。通过本文提出的efficient self-attention和相对位置编码，使得本文可以借助Transformer有效的聚集不同尺度的全局上下文信息。
因为误分割经常发生在ROI区域的边界，因此高分辨率的上下文信息就对精确分割至关重要。

本文设计的焦点就在于self-attention模块如何更高效的计算大尺寸的feature map。本文并没有简单的在CNN提取的feature map上计算self-attention，而是在encoder-decoder每一级别上使用Transformer来收集不同尺度的长程依赖关系。 但原始输入并没有使用Transformer处理，因为在较浅层次使用Transformer不仅没什么额外的作用还会增加计算成本。
一种可能的原因是浅层feature map更关注细节纹理，全局性信息没那么强。


Fig 1(b),©则是residual block和Transformer block的具体结构。细节：残差连接中使用的是pre-activation。




Section III Experiment

本文将UTNet在多器官分割、心室分割任务上进行了实验。
Table 1展示了UTNet与其他网络对比的结果。

UNet-医学图像分割基准网络之一

ResUNet-以残差连接作为基本模块的UNet

CBAM-两路卷积分别提取通道和空间注意力


Dual Attention-分别通过两个网络提取空间和通道注意力


可以看到UTNet在所有分类器官中的性能均是最好的；借助残差连接ResUNet性能略优于UNet；CBAM进一步提升了性能但仍受限于有限的感受野。Dual Attention几乎与ResNet方法相同，但是计算复杂度为2倍，因为引入了两个注意力提取网络，但仍然无法借助高分辨率的特征修复边界处的误分割现象；而UTNet网络参数量更少，却可以更好的捕获细节和全局信息。

消融实验


Fig 3(a)显示在不同位置使用self-atention对性能的影响。比如 34 表示在下采样level3和level4处施加self-attention；随着level越深，self-attention可以收集更多细粒度特征从而提升性能；但是当在原始分辨率图像级也加入self-attention时曲线达到饱和。本文认为是因为较浅层次更关注于局部问题，这个时候聚合全局上下文是无法提供太过全局信息的。



Fig 3(b)则展示了在计算self-attention时不同降维方式以及不同降维尺度对性能的影响，可以看到插值略优于最大池化。



Fig 3©则展示了Transformer不同部位，encoder,decoder,相对位置编码等对性能的影响，可以看到Transformer encoder 和decoder结合才是性能最佳的一种组合，同时相对位置编码对性能也有显著影响。




此外，本文还在计算复杂度和空间复杂度上将UTNet与DualAttention进行了对比，使用25625616的输入图像，batch size相同，来对比二者的推理时间和内存需求。




GPU内存 3.8GB vs 36.9GB 推理时间：0.146s vs 0.243s,显然UTNet更具优势。











鲁棒性分析






Table 2展示了不同网络结构的泛化性能，即测试和训练使用不同的数据集，UTNet仍然展现了优异的性能。这一结果可以归因于UTNet在不同层次feature map上均使用了self-attention来聚合全局信息、内容-位置之间的关联信息，这样赋予了UTNet更好的关注全局上下文的能力，而不仅仅是关注局部纹理。







Fig 4进一步展现了UTNet在边界上的分割结果是最一致的，其他方法都不能很好的捕获边界特征，特别是在切割RV和MYO区域时。









Section IV Conclusion

本文提出的UTNet有效的将CNN与Transformer中的self-attention机制结合在一起，用于提升医学图像分割的性能。这种混合框架无需Transformer在大规模数据集上进行预训练，其中self-attention可以有效地捕获不同层次的长程信息。我们相信这种设计将有助于设计更加丰富的Transformer模型，使其更加适合于医学图像分割任务；此外UTNet中优异的处理长序列的能力为迁移到其他下游任务开辟了可能性。