Mixed Transformer U-Net For Medical Image Segmentation
Transformer中的自注意力机制的缺陷:需要大规模预训练;具有二次计算复杂度,降低对医学图像等高维数据的处理速度;忽略了样本间的相关性。
本文重新设计自注意力,然后将其与外部注意力集成。由于在大多数视觉任务中,邻近区域之间的视觉依赖关系通常比那些遥远区域之间的要强,因此在细粒度局部上下文执行局部SA,在粗粒度全局上下文执行全局SA。在计算全局注意力映射时,使用轴向注意力减少计算量,并进一步引入可学习的高斯矩阵来增强附近tokens的权重。
网络整体架构:
为了降低计算成本,MTMs只对空间大小较小的深层使用,而上层仍然使用经典的卷积运算。通过使用卷积引入一些先验信息到模型中,对于尺寸较小的医学图像数据集是很有帮助的。
Mixed Transformer Module(MTM):
MTM包含Local-Global Gaussian-Weighted Self-Attention(LGG-SA)和External Attention(EA).LGG-SA对不同粒度的长短程依赖关系进行建模,EA用于发掘样本之间的相关性。
Local-Global Gaussian-Weighted Self-Attention
LGG-SA采用局部全局策略和高斯掩码:
Local-Global Self-Attention
在计算机视觉中,邻近区域之间的相关性往往比远处区域之间的相关性更重要,在计算注意图时,不需要为更远的区域花费相同的代价。因此提出Local-Global Self-Attention。局部注意力计算每个窗口内的自亲和力,然后每个窗口中的token被聚合为一个全局token,来表示窗口的主要信息。
Gaussian-Weighted Axial Attention
与LSA使用原始SA不同,GSA使用Gaussian-Weighted Axial Attention,GWAA通过一个可学习的高斯矩阵增强了每个查询对附近token的感知,同时由于轴向注意,时间复杂度较低
External Attention
所有样本共享两个记忆单元
M
K
M_{K}
MK,
M
V
M_{V}
MV,描述了整个数据集中最重要的信息。