论文解析[13] MIXED TRANSFORMER U-NET FOR MEDICAL IMAGE SEGMENTATION
发表年份：2021
论文地址：https://arxiv.org/abs/2111.04734
代码地址：https://github.com/dootmaan/mt-unet
文章目录
- 摘要
- 2 方法
- 2.1 总体结构
- 2.2 Mixed Transformer Module
- 2.3 Local-Global Gaussian-Weighted Self-Attention
- 2.3.1 Local-Global Self-Attention
- 2.3.2 Gaussian-Weighted Axial Attention
- 2.4 External Attention
- 3 实验
- 4 结论
摘要
由于Vision Transformer天生可以通过Self-Attention（SA）来提取长距离的相关性，经常被用于分割结构。然而，Transformer通常依赖于大尺度的预训练和很高的计算复杂度。并且，SA只能对单个样本的自相关性进行建模，忽略了整个数据集的潜在相关性。
为了解决这个问题，我们提出了一个全新的Transformer模块MTM，可以同时进行样本之间的和内部的相关性学习。
MTM首先计算自相关性通过Local-Global Gaussian-Weighted Self-Attention（LGG-SA）。通过External Attention（EA）来挖掘数据样本之间的内部连接。
通过使用MTM，我们建造了一个U型模型Mixed Transformer U-Net (MT-UNet) 用于精确的医学图像分割。
2 方法
2.1 总体结构
基于一个解码器-编码器结构，当解码时使用跳跃连接来保持低水平的特征。
MTM仅用于空间尺度更小的深层来减少计算开销，上层仍然使用经典的卷积操作。因为我们想要在初始层关注局部的关系，因为它们包含更多的高分辨率细节。
2.2 Mixed Transformer Module
MTM包括LGG-SA和EA。
LGG-SA用于建模不同粒度的短距离和长距离依赖。EA用于利用样本之间的联系。
这个模块用于替代原始的Transformer编码器，由于它在视觉任务上更好的的表现和更低的时间复杂度。
2.3 Local-Global Gaussian-Weighted Self-Attention
LGG-SA体现了聚焦计算的想法。不像传统的SA对所有token使用相同的注意，LGG-SA可以更多的关注于临近的区域，因为local-global的策略和高斯掩码。实验证明，它可以提升模型表现，也节约了计算资源。
2.3.1 Local-Global Self-Attention
SA旨在捕获输入序列的所有实体之间的相互联系。为了实现目标，SA引入了三个矩阵，key（K）、query（Q）和value（V）。这三个矩阵是输入X的线性变换。
Local SA计算每个窗口中的自相关性
尝试的传播函数中，Lightweight Dynamic Convolution (LDConv)是表现最好的。
2.3.2 Gaussian-Weighted Axial Attention
不像LSA使用原始的SA，我们提出了Gaussian-Weighted Axial Attention (GWAA)用于GSA。
S ( q , K ) S(q,K) S(q,K) 表示q和K之间的相似度，高斯权重为 e − D i , j 2 2 σ 2 e^{-\frac{D^2_{i,j}}{2σ^2}} e−2σ2Di,j2，输出结果为：
因为想让方差是可学习的，可以写成：
可以简单使用w来表示 D i , j 2 D^2_{i,j} Di,j2之前的系数， w D i , j 2 w D^2_{i,j} wDi,j2也充当相对位置偏置，通过它可以在MTM中强调位置信息。它提高了显式提供相对关系的模型性能，这是普通绝对位置嵌入无法做到的。
2.4 External Attention
EA首先被提出解决SA不能挖掘不同样本之间的关系。不像SA使用每个样本自己的线性变换来计算注意力分数，EA中每个样本共享两个注意力单元 M K M_K MK 和 M V M_V MV，用于描绘整个数据集中最关键的信息。
在我们的设计中，一个附加的线性映射是用于扩大矩阵Q的通道，进而提升这个模块的表示学习能力。
3 实验
4 结论
我们提出了一个有效的名为MT-UNet的ViT用于医学图像分割。模型的特点在于MTM，由于LGG-SA和EA，能够同时学习内部和外部的关联。提出的模型有更低的时间复杂度，在我们的实验中超过了sota的ViT。