U-Net Transformer：用于医学图像分割的自我和交叉注意力模块

Abstract-摘要

对于复杂和低对比度的解剖结构，医学图像分割仍然特别具有挑战性。本文提出的一种U-Transformer网络，它将Transformer中的self-attention和Cross attention融合进了UNet，这样克服了UNet无法建模长程关系和空间依赖的缺点，从而提升对关键上下文的分割。本文集合了两种注意力机制：自注意力机制建立编码器之间的全局交互，交叉注意力被添加到跳跃连接当中。通过过滤非语义特征在UNet Decoder中恢复空间信息。

Introduction

尽管FCN，Unet系列都取得了不错的性能，但是仍然具有局限性，这种网络结构处理有局部二义性或者低对比度器官时表现不佳。

本文提出的U-Transformer，借助Transformer的特性来帮助建模长程依赖关系和空间依赖关系，U-Transformer保留了原始的U型架构提取归纳偏置，此外引入了两种注意力机制来改善网络的决策能力。

首先。self-attention在编码器的结尾用来显式的建模语义之间的全局交互关系，其次在跳跃连接中加入了Cross attention从而过滤到非语义特征，从而更好的协助U-Net恢复空间信息。

The U-Transformer结构

U型网络的编码器结构可以有效的提取全局上下文信息进行复杂的医学图像分割任务，本文的U-Transformer将多头Transformer中的模块引入UNet，通过两种注意力机制（多头自注意力和多头交叉注意力）来提取长程关系和空间依赖性。其网络结构如下图：

如图所示，多头注意力机制（MHSA）加在编码器最后，其感受野是整张图像1，而原始的UNet的感受野只是局部（蓝色区域），多头交叉注意机制用在解码器中的跳跃连接之后，主要将高层次语义更丰富的特征图谱与来自跳跃连接的高分辨率图结合起来。

Part-1自注意力机制

多头自注意力模块用于提取图像中的长程结构信息，因此本文在UNet做完编码器部分后在底部加入了多头注意力机制，这个多头注意力机制的主要作用是将特征图中的每个元素建立关联，对应的感受野是整个图像，因此当前像素点的分类取决于所有输入的像素，注意力的计算公式如下：

自注意力机制的计算涉及到了Q,K,V三个矩阵

MHSA模块

MHSA模块的具体计算过程如下图：

首先特征图+位置编码后reshape到一定的维度，Wq，Wk，Wv是训练的3个矩阵，通过这三个矩阵得到自注意力模块的计算结果，而Positional embedding对于医学图像分割十分重要，因为不同的组织结构分别在图像中不同的固定位置，加入Position信息之后就可以捕获器官之间的绝对信息和相对信息。

Q,K经过Softmax后的结果为注意力系数矩阵A，A乘以V得到最后的注意力输出。

MHCA模块

 交叉注意力负责将输入图像的每个像素之间建立关联，MHCA的核心思想就是想办法滤除跳跃连接中无关或噪声区域，突出显示相关区域。

其结构如下图所示：

MHCA的作用类似于一个门控函数，输入分别是跳跃连接过来的结果S和上一层的特征图Y处理后的结果。Y嵌入后的结果作为Q,K，S输入的结果作为V，最后跳跃连接的输出S是经过Y加权处理后的结果，计算出来的注意力权重会被rescale到0-1之间，最终交叉注意力的计算结果Z作为一个Fliter，再与S做点积，其中权值较小的元素代表噪声或者不相关的区域，可以被去除。经过这样处理的精简版的S，再将这一精简版的结果S再与Y做级联。

总结

本文提出的U-Transformer借助Transformer中的自注意力和交叉注意力有效的增加了UNet的性能，可以有效的捕获长程依赖关系，尤其对复杂器官或较小器官具有很好的分割性能。