恒源云(GpuShare)_医学图像分割:MT-UNet

我们社区有新的技术分享小伙伴啦🎉🎉🎉
热烈欢迎👏
作为一名合格的搬运工，我必须做点事情表达我的喜悦之情：搬运～搬运～立即搬运～

文章来源 | 恒源云社区

原文地址 | 新的混合Transformer模块(MTM)

原文作者 | 咚咚

---

摘要

存在问题	虽然U-Net在医学图像分割方面取得了巨大的成功，但它缺乏对长期依赖关系进行显式建模的能力。视觉Transformer由于其固有的通过自注意（SA）捕捉长程相关性的能力，近年来成为一种可替代的分割结构。
存在问题	然而，Transformer通常依赖于大规模的预训练，具有较高的计算复杂度。此外，SA只能在单个样本中建模self-affinities，忽略了整个数据集的潜在相关性
论文方法	提出了一种新的混合Transformer模块(MTM)，用于同时进行inter-affinities学习和intra-affinities学习。MTM首先通过局部-全局高斯加权自注意(LGG-SA)有效地计算窗口内部affinities。然后，通过外部注意挖掘数据样本之间的联系。利用MTM算法，构造了一种用于医学图像分割的MT-UNet模型

Method

如图1所示。该网络基于编码器-解码器结构

1. 为了降低计算成本，MTMs只对空间大小较小的深层使用，
2. 浅层仍然使用经典的卷积运算。这是因为浅层主要关注局部信息，包含更多高分辨率的细节。

MTM

如图2所示。MTM主要由LGG-SA和EA组成。

LGG-SA用于对不同粒度的短期和长期依赖进行建模，而EA用于挖掘样本间的相关性。

该模块是为了替代原来的Transformer编码器，以提高其在视觉任务上的性能和降低时间复杂度

LGG-SA（Local-Global Gaussian-Weighted Self-Attention）

传统的SA模块对所有tokens赋予相同的关注度，而LGG -SA则不同，利用local-global自注意力和高斯mask使其可以更专注于邻近区域。实验证明，该方法可以提高模型的性能，节省计算资源。该模块的详细设计如图3所示

local-global自注意力

在计算机视觉中，邻近区域之间的相关性往往比遥远区域之间的相关性更重要，在计算注意图时，不需要为更远的区域花费相同的代价。

因此，提出local-global自注意力。

1. 上图stage1中的每个局部窗口中含有四个token，local SA计算每个窗口内的内在affinities。
2. 每个窗口中的token被aggregate聚合为一个全局token ，表示窗口的主要信息。对于聚合函数，轻量级动态卷积(Lightweight Dynamic convolution, LDConv)的性能最好。
3. 在得到下采样的整个特征图后，可以以更少的开销执行global SA（上图stage2）。

其中$X\in R^{H\times W\times }$