MIXED TRANSFORMER U-NET FOR MEDICAL IMAGE SEGMENTATION（医学图像分割）

目录

ABSTRACT

1. INTRODUCTION

2.METHODS

        2.1 网络整体结构

        2.2 Mixed Transformer Module（MTM模块）

        2.3 Local-Global Gaussian-Weighted Self-Attention

 Local-Global Self-Attention

 Gaussian-Weighted Axial Attention

 External Attention

---

ABSTRACT

unet缺乏建立长程依赖关系的能力，VIT虽然具有通过自我注意(SA)捕捉长程相关性的能力，却需要大规模的预训练，计算复杂度很高。并且，SA只能在单个样本中建模，忽略了整个数据集的潜在相关性。因此提出了混合transformer—Mixed Transformer Module (MTM)模块以用于样本内和样本间的学习。

MTM首先使用本文精心设计的局部-全局高斯权重自注意力模块（LGG-SA）来计算自相似性，然后会通过外部注意力模块学习样本间关系。基于MTM模块本文构建了一个名为MT-UNet的U型网路，用于精确的医学图像分割。在两个公开数据集上的结果显示超过了其他方法。

---

1. INTRODUCTION

尽管U-Net在医学图像处理中占据主导地位，但是由于卷积运算的固有局限性，U-Net缺乏建模远程相关性的能力；Transformer是一个基于注意力的模型，最初用于序列对序列的预测，自注意(SA)是变压器的关键部件，可以对所有输入标记之间的相关性进行建模，从而使Transformer能够处理长期依赖关系。但它们往往严重依赖大规模的预训练，给方法的使用带来了不便，另外，SA的计算复杂度是二次的，这可能会降低对医学图像等高维数据的处理速度，SA也有忽略样本间相关性的局限性。

故文中提出重新设计SA以获得更好的局部感知和更低的计算成本，然后将其与外部注意(EA)集成，以同时管理内部和内部的相关性。由于在大多数视觉任务中，邻近区域之间的视觉依赖关系通常比那些遥远区域之间的要强，因此在细粒度局部上下文执行局部SA，在粗粒度全局上下文执行全局SA。在计算全局注意力映射时，使用轴向注意力减少计算量，并进一步引入可学习的高斯矩阵来增强附近tokens的权重。

总结：

（1）提出MTM模块可以同时学习样本内和样本间的关系

（2）提出LGG-SA可以在细粒度局部和粗粒度全局上下文依次执行SA，还引入一个可学习的高斯矩阵来强调每次查询的附近区域

（3）提出的MT-UNet是用于医学图像的分割框架，在两个数据集验证了其有效性

---

2.METHODS

        2.1 网络整体结构

 该网络主要运用了编解码器解构，在解码时，使用跳过连接来保持低级特征。由图可知，只有在空间大小较小且较深的层次才使用MTM模块。通过使用卷积，我们还可以引入一些结构先验信息到模型中，这对于尺寸相对较小的医学图像数据集是有帮助的。需要注意的是，对于所有的Transformer模块，都采用了2步卷积/反卷积核来实现下采样/上采样以及信道扩展/压缩。

        2.2 Mixed Transformer Module（MTM模块）

 MTM由LGG-SA和EA组成，LGG-SA用于建模不同粒度的依赖关系，包括局部和全局的；EA则计算样本之间的相关性。该模块是为了替代原来Transformer的SA。

        2.3 Local-Global Gaussian-Weighted Self-Attention

LGG-SA采用局部全局策略和高斯掩码：

 Local-Global Self-Attention

在计算机视觉中，邻近区域之间的相关性往往比那些遥远区域之间的相关性更重要。所以提出了Local-Global Self-Attention。局部注意力计算每个窗口内的自亲和力，然后每个窗口中的token被聚合为一个全局token，来表示窗口的主要信息。
聚合函数使用了步长卷积、最大池化等，其计算流程如下：

（LSA：local-self-attention；LDConv：动态卷积；GSA：全局操作） 

 Gaussian-Weighted Axial Attention

文中提出了gasian-weighted-Axial-Attention (GWAA)用于GSA。GW AA通过一个可学习的高斯矩阵增强了每个查询对附近token的感知，同时由于轴向注意降低了时间复杂度。

 

 External Attention

所有样本共享两个记忆单元Mk和Mv，描述了整个数据集中最重要的信息。