ViT+Swin+TNT论文阅读

ViT+Swin+TNT论文阅读

能力不足，理解有限，请多指点。
记录一下阅读的三篇论文，三个应用于视觉方面的Transformer变体模型。

---

---

一、ViT(Visual Transformer)

Visual Transformer，Dosovitskiy A, Beyer L, Kolesnikov A, et al. An image is worth 16x16 words: Transformers for image recognition at scale[J]. arXiv preprint arXiv:2010.11929, 2020.

ViT模型架构如上图所示。在ViT中，先将作为输入数据的图像$x\in R^{H\times W\times C}$重构为二维补丁序列$x_p\in R^{N\times P^{2}C}$，其中$(P,P)$是每个图像补丁的分辨率，$N=HW/P^2$是图像补丁的个数。
分成N块补丁后，补丁经过线性映射，所有向量的维度都变成D维，向量还需包含补丁信息(Patch embedding)和位置信息(Position embedding)。在进入Transformer Encoder前，每个向量都需插入一个可学习的分类头(Extra learnable class embedding)，这些分类头随机初始化后在训练时通过学习来调整参数。
ViT使用的Transformer Encoder与NLP中的Encoder没有很大的变化，此处模型由L层构成，其中使用Layer Normalization进行归一化、使用GeLU作为MLP的非线性激活函数。

公式(1)的$x_p^i$均为分类头，第一个$x_{class}$作为不包含任何信息的分类头附加的补丁信息为空，E为Patch Embedding，$E_{pos}$是Position Embedding。公式(4)的结果y是Transformer Encoder的输出结果，也是MLP Head的输入，最后得到分类结果。
ViT这篇论文主要做的工作是将NLP领域中的Transformer模型应用于视觉领域上，对于CV和NLP这两者不同的输入，ViT将有三个维度的图像降维再升维得到了能够用于Transformer的二维向量，一定程度上解决了输入不同的问题。但论文只提到了分类问题上的应用，还需要讨论能否普遍应用于CV领域的其他问题上。

二、Swin Transformer

Liu Z, Lin Y, Cao Y, et al. Swin transformer: Hierarchical vision transformer using shifted windows[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2021: 10012-10022.

Swin这个名字取自Shifted的S+Windows的win

Swin的模型架构如上图所示。输入图像$x\in R^{H\times W\times 3}$首先经过补丁分割(Patch Partition)下采样至四倍得到大小为$\frac{H}{4}\times \frac{W}{4}\times 48$（即$\frac{H}{4}\times \frac{W}{4}\times 3\times 16$）。而后由Linear Embedding将维度调整至C维，将$\frac{H}{4}\times \frac{W}{4}\times C$输入到Stage 1的Swin Transformer Block中。
Swin Transformer Block由两个类似Transformer Encoder的结构组成，W-MSA是常规的窗口内多头自注意力模块，SW-MSA是带有滑动窗口的多头自注意力模块。x2和x6都是偶数倍也是由于该两层连续结构，现在先将该模型的Transformer Block当作一个黑盒。
Stage 1输出的特征图尺寸仍为$\frac{H}{4}\times \frac{W}{4}\times C$，作为输入在Patch Merging中。

Patch Merging如(a)图所示，进行下采样操作，缩小特征图的H和W，并将通道维度翻倍，得到多尺度特征图。正常来说，维度应该呈4C倍增长，但由于使用了全连接层使得维度以2C倍增长。作者给出的(a)图中，Swin是以下采样4倍、8倍、16倍、32倍来得到层次特征图，而ViT一直都以16倍来操作，Swin这样做的好处是能得到更高维的特征信息，更适用于一些分割检测任务。

考虑到如果用图像作为输入数据，全局计算复杂度回呈二次型，这并不适用于处理密集预测任务和处理高分辨率图像的视觉问题。
因此，作者提出在互不重叠的窗口里做自注意力计算，Swin Transformer Block用基于滑动窗口的(S)W-MSA模块来代替常规的自注意力机制。

先看这4个公式，${\hat{z}}^l$是(S)W-MSA模块的输出特征图，$z^l$是MLP的输出特征图，将前一个模块得到的$z^{l-1}$作为输入，经过一个Block后得到$z^l$，再将$z^l$作为输入数据放进连续的Block中，得到$z^{l+1}$，输入和输出的维度不变。
这么看起来与常规的自注意力机制没有什么区别，那么基于窗口的自注意力机制有什么不同呢？

Swin会在第$l$层将图像划分成大小为$M\times M$的窗口，并在每个窗口里进行自注意力计算。接着在下一层Block即$l+1$层中，将窗口向右和向下移动$\lfloor \frac{M}{2}\rfloor$位，得到新窗口后再计算新窗口内的相关性（即计算自注意力）。这样基于滑动窗口的自注意力机制不仅降低了计算复杂度，还将相邻窗口之间的信息连接起来，窗口之间互相通信，并且从架构图可以看到做了多次这样的移位操作，其实足够计算出整张图像的相邻信息。
但这种方法还是存在一些问题，像上面的示意图那样，窗口数量从4个变成了9个，而且有些窗口大小不一致，计算量也增至2.25倍。
为了解决这个问题，作者提出一种将元素向左上角循环移位的方法（Cyclic Shift）。

循环移位可以看作是先向上移动$\lfloor \frac{M}{2}\rfloor$位，再向左移动$\lfloor \frac{M}{2}\rfloor$位。移位后的窗口由于存在不相邻的子窗口，而不相邻的子窗口做相关性计算是没有意义的，因此需要遮挡掉不需要的计算结果。
Swin的操作是先全部窗口进行相关性计算后，利用Masked MSA对不相邻子窗口的结果$Q{K}^T$加上-100。因为经过归一化后数值一般都会变得很小，-100对这些小数值来说是一个更小的负数，在进行Softmax后会忽略掉（变为0），这就实现了遮挡不需要的计算结果。
结束后还需要将移位的子窗口归位，保持各元素之间相对位置不变。
下面简单比划一下为什么Masked这样设置：




首先假设窗口大小M=7，故分别向上和向左移动3位，这里取序号为6和4的子窗口计算。将窗口内的各元素拉直成C维，得到的Q和K也是49xC维，计算$Q{K}^T$，结果中有序号6和序号6子窗口内元素计算、有序号4和序号4的元素计算、还有序号6和序号4元素的计算，而序号6和序号4窗口的计算是无意义的，因此我们加上一个掩码矩阵，对不需要的计算结果加上-100。
其他窗口的计算也类似。




解释一下计算复杂度，对于常规的多头自注意力机制计算，输入图像与矩阵内积得到QKV，这三次内积的计算复杂度为$3hw{C}^2$，$Q{K}^T$的计算复杂度为${(hw)}^{2}C$，忽略softmax的计算，Attention这一步的计算复杂度为${(hw)}^{2}C$，最后多头映射的计算复杂度为$hw{C}^2$。
故常规的多头自注意力机制计算复杂度$=4hw{C}^2+2{(hw)}^{2}C$。
对于(S)W-MSA，在单个窗口内计算，$h=M,w=M$，故单个窗口的计算复杂度
$=3M^2{C}^2+{(M^2)}^{2}C+{(M^2)}^{2}C+M^2{C}^2=4M^2{C}^2+2{(M^2)}^{2}C$，
又因由$\frac{h}{M}\times \frac{w}{m}$个窗口，故总体计算复杂度为
$=(\frac{h}{M}\times \frac{w}{m})\times [4M^2{C}^2+2{(M^2)}^{2}C]=4hw{C}^2+2M^{2}hwC$

三、Transformer in Transformer(TNT)

Han K, Xiao A, Wu E, et al. Transformer in transformer[J]. Advances in Neural Information Processing Systems, 2021, 34.

1. 首先将输入的图像统一分割成n块补丁$X=[X^1,X^2,...,X^n]\in R^{n\times P\times P\times 3}$，$X^i$为句子向量(sentence embedding)；
2. 每块句子补丁又分成m块子补丁$X^i\to [x^{i,1},x^{i,2},...,x^{i,m}],x^{i,j}\in R^{S\times S\times 3}$，$x^{i,j}$为词向量(word embedding)；
3. 句子向量线性映射得到$Z_0=[Z_{c}lass,Z_0^1,Z_0^2,...,Z_0^n]\times R^{(n+1)\times d}$；
4. 词向量线性映射得到$Y^i=[y^{i,1},y^{i,2},...,y^{i,m}],y^{i,j}=FC(Vec(x^{i,j}))$，词向量的集合$Y_l=[Y_l^1,Y_l^2,...,Y_l^m]$；
5. Inner Transformer Block处理词向量的数据流，$Y{‘}_l^i=Y_{l-1}^i+MSA(LN(Y_{l-1}^i)),Y_L^i=Y{‘}_l^i+MLP(LN(Y{‘}_l^i))$；
6. 词向量与句子向量相加，$Z_{l-1}^i=Z_{l-1}^i+FC(Vec(Y_l^i))$；
7. Outer Transformer Block处理句子向量的数据流，$Z{‘}_l=Z_{l-1}+MSA(LN(Z_{l-1})),Z_l=Z{‘}_l+MLP(LN(Z{‘}_L))$；
8. TNT整个过程可以表示为：$Y_l,Z_l=TNT(Y_{l-1},Z_{l-1})$。

TNT虽然也有考虑到图像的局部信息，但与Swin不同的是它没有考虑到与相邻窗口联系。