PMTrans：Pyramid Medical Transformer

最近Transformer风头很劲，自己也在关注其在Medical Image上的应用。

论文概要

这一篇Pyramid Medical Transformer(PMTrans)提出的是一个整合了多尺度注意力和金字塔结构卷积神经网络特征提取的方法。文章总结了该方法的三大显著特点：

1. 模型采用了金字塔结构，包含四个分支:3个Transformer（多尺度）分支和一个CNN分支；
2. 在短距离（short range）transformer中采用了局部门控轴向注意力（local gated axial attention），在中距离和长距离分支采用全局门控轴向注意力（global gated axial attention）；
3. 提出了一种合成和标签指导机制（fusion and guidance scheme），用以3个不同尺度transformers和CNN分支输出的多尺度特征图。

结构概览

方法

Gated axial attention

Axial Attention计算的是近似的自注意力（approximate self-attention）。假设输入图像或特征图为$X$，则$X$自注意力$Y$为：
$Q=W_{q}X，K=W_{k}X，V=W_{b}X，$
$Y=\text{softmax}(\frac{Q{K}^T}{\sqrt{d_k}})V,$
$Q,K和V$分别是query，key和value矩阵，$d_k$是keys和queires的维度。引入$d_k$作为分母，是为了抵御高维keys和queires点积造成的梯度过小的问题。

文中提到的axial attention是为了解决self-attention计算成本过高而进行的一种近似。它仅将self-attention应用于高度和宽度方向。
假设输入特征图$X\in {\mathbb{R}}^{H\times W\times C_{in}}$,沿宽度轴向的self-attention机制可通过以下公式计算：
$y_{ij}={\sum }_{w=1}^W\text{softmax}(q_{ij}^T{k}_{iw}+q_{ij}^T{r}_{iw}^q+k_{iw}^T{r}_{iw}^k)(v_{iw}+r_{iw}^v)$
其中像素在$i$和$j$是宽度和高度轴的位置；$r_{iw}^q$,$r_{iw}^k$he $r_{iw}^v$是queries, keys, 和values的相对位置编码。

Gated axial attention其实就是对相对位置编码增加权重。如果模型无法从小的医学数据集中学习到准确的位置编码门控将降低这些编码的影响。
$y_{ij}={\sum }_{w=1}^W\text{softmax}(q_{ij}^T{k}_{iw}+G_Q{q}_{ij}^T{r}_{iw}^q+G_K{k}_{iw}^T{r}_{iw}^k)(G_{V1}{v}_{iw}+G_{V2}{r}_{iw}^v)$
三个transformer有不同的receptive field,但都使用相同的门控轴向注意力机制，使得多尺度信息能够被高效地捕捉，也避免了将图像分割成块面临的问题。

Fusion and deep supervision

除了3个transformer，网络中还有一个CNN分支。通过attention U-net中提出的金字塔融合机制将不同维度的CNN特征图逐步与transformer的特征进行融合。此外，还通过增加auxiliary classifier
的方式增加deep supervision。

实验结果

数据集

Gland segmentation(GLAS)
MoNuSeg dataset

SOTA比较

Quantitative comparison
Qualitative comparison