Medical Transformer 论文阅读笔记

• 这是arxiv上的一篇文章，文章作者：

---

• 目前用卷积神经网络来做图像分割虽然也能取得一些好效果，然而卷积神经网络无法对长依赖进行建模，虽然后续提出了一些弥补的trick如注意力机制等，仍然无法完全解决这个问题。而transformer以能够建模长依赖著称，几个月前更是有篇用transformer做图像分类的文章横空出世，在大量数据集训练的条件下取得了比卷积神经网络更好的准确率。
• 文章提出了用transformer来做医学图像分割。要解决的问题是，transformer在图像任务上相比卷积神经网络需要更大的数据集来训练，而医学图像处理的一个难题就是数据不足，数据集不够大。
• 文章主要贡献是两点，一是在现有的注意力机制上增加了一个机制，一个是提出了一种训练机制（其实很难算得上是一种“训练机制”）。
• 首先开头便贴了一张图，试图说明因为建立了长依赖所以能够降低假阳：
  
  然而我也不知道为什么降低假阳一定是长依赖的功劳，只能说可能是精心挑了张对比效果较好的图片。
• 后面introduction介绍了一些transformer用于cv和segmention的论文我才知道，自从第一个吃螃蟹的出来后，短短几个月已经这么多工作关注在这上面了。文章介绍完现有工作后来了句：
  
  大概意思是说本文只用了transformer机制在unet结构的encoder部分的自注意力机制上，并且不像其它transformer用于cv的方法一样依赖于大数据集预训练的权重，本方法不需要预训练。
• 由于医学图像数据集没有足够的图片数量来训练transformer的位置编码，文章提出一种门限的位置敏感的轴向注意力机制来限制进入位置编码的信息量。
• 网络整体结构如下，为两分支的U-shape结构，其中只在两个分支的encoder部分用到了transformer，也就是下图红色部分：
  
• ViT提出的时候，transformer的每个token会对所有的每个token都计算注意力，所以是$(hw{)}^2$次计算，这太复杂了，于是有一篇文章（Ho, J., Kalchbrenner, N., Weissenborn, D., Salimans, T.: Axial attention in multidimensional transformers. arXiv preprint arXiv:1912.12180 (2019)）提出了轴向注意力（axial attention），每个token只对同一行和同一列的token计算注意力，本文用的就是轴向注意力；
• 还有，transformer其实是不具有位置表达能力的，为了添加位置表达能力，需要加一个position embedding，就是用一个onehot的位置向量，经过一个全连接的embedding，产生位置编码，这个全连接是可训练的。这个位置编码原先只加在Q K V的Q上，而这篇文章（Wang, H., Zhu, Y., Green, B., Adam, H., Yuille, A., Chen, L.C.: Axialdeeplab: Stand-alone axial-attention for panoptic segmentation. arXiv preprint arXiv:2003.07853 (2020)）提出了要把它放在Q K V三个上面。
• 原来的公式是这样的：
  
  加上轴向注意力和多个位置编码的trick后现在变成这样（w方向的注意力，h方向和这个类似）：
  
• 然而上述的trick需要大量数据集进行训练，小量的数据不足以训练QKV的三个position embedding，而医学数据集多数情况下就是少量的。在这种情况下，不准确的position embedding会给网络准确率带来负面影响，为此文章提出了个方法用来控制这个影响的程度，修改上述公式如下：
  
  这里三个G都是可学习的参数，当数据集不足以使得网络预测准确的position embedding时，网络的G会小一点，反之会大一点，因此起到一个所谓的“Gated”的作用。其实看到这里稍微有点失望，我原以为G可以是网络通过某种方法预测的，对不同样本有不同的G，不过这样也合理吧应该，因为position embedding只要相对位置一样，对不同的样本应该是一样的，毕竟只是位置信息，没有包含语义信息。不过前面铺垫了一堆，结果不过是在前人的方法加了两个可学习的权重。这就是文章的第一个贡献点。
• 然后是关于两个branch的，这里文章有一句话我一开始不能理解：

However, the patch-wise training alone is not sufficient for the tasks like medical image segmentation.It is highly likely that the segmentation mask will be larger than the patch size.

• 后来想明白了，文章说的意思是，transformer做图像分割可以用patch-wise的方式去做，也就是说把一张完整图片切割成多个patch，每个patch和这个patch对应的mask作为一个样本，用来训练transformer，这样又快，然而问题在于，一张图片的一个病灶可能比一个patch大，这样的话这个patch看起来就会很奇怪，因为被病灶充满了。
• 于是文章关于两个branch的做法是：global branch不做特殊处理，就是整张完整特征图，进行两次transformer block后送进decoder；而local branch切分成4x4个patch，每个patch单独送transformer block 前向传播，patch和patch之间没有任何联系，最后再把这4x4个patch的feature map拼回去。用深的网络来处理local信息，用浅的网络来处理global信息。这就是文章的第二个贡献点，感觉第二个贡献点称为"training strategy"不太恰当。
• 至此文章两个贡献点讲完，文章给的启示主要是有两点
  • 最主要一点是local和global两个信息都利用上的方法，用简单的浅的网络去对整张图做处理，作为一个分支，第二个分支可以对每个patch分别做处理，然后把特征图拼起来。
  • 第二点是如果对一个trick产生的结果不太信任，可以用一个可学习的权重去weight它。