Pyramid Vision Transformer 论文笔记

Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions

PVT

论文链接: https://arxiv.org/abs/2102.12122

PVT v2 论文笔记链接

一、 Problem Statement

提出一个类似于CNN中的pyramid 结构，使得Transformer结构适用于多种dense prediction task。

二、 Direction

1. 采用4x4大小的image patch作为输入
2. 引入一个progressive shrinking pyramid，随着网络加深，减少transformer序列长度
3. 采用一个spatial-reduction attention(SRA)层，当学习high-resolution特征的时候，进一步减少资源消耗

三、 Method

先来看看整体的结构:

1. Feature Pyramid for Transformer

PVT使用的是progressive shrinking strategy来控制feature maps的大小。假设patch size在$i_{th}$ stage为$P_i$。在stage $i$的最开始阶段，把输入特征图平均分为$\frac{H_{i-1}{W}_{i-1}}{P_i^2}$patches，因此特征图的大小为$F_{i-1}\in {\mathbb{R}}^{H_{i-1}\times W_{i-1}\times C_{i-1}}$。然后把每个path flatten，变成$C_i$维度的embedding。通过线性变化之后，embedded patches的大小为$\frac{H_{i-1}}{P_i}\times \frac{W_{i-1}}{P_i}\times C_i$。这样的话，高度和宽度都是$P_i$倍小于输入。所以，通过改变patch size的大小，可以灵活地调节每个阶段feature maps的大小。

2. Spatial-reduction Attention

Transformer包括两个部分:attention layer和feed-forward layer。作者提出了spatial-reduction attention(SRA) layer 来取代传统的multi-head attention(MHA) layer。

SRA层也是接收Q, K, V作为输入。区别是在SRA在attention operation之前，降低了K和V的空间尺度，极大地减少了计算和存储资源消耗。

$$\begin{gathered} \text{SRA}(Q,K,V)=\text{Concat}({\text{head}}_0,...,{\text{head}}_{N_i})W^o \\ {\text{head}}_j=\text{Attention}(Q{W}_j^Q,SR(K)W_j^K,SR(V)W_j^V) \end{gathered}$$
其中$W_j^Q\in {\mathbb{R}}^{C_i\times d_{head}},W_j^K\in {\mathbb{R}}^{C_i\times d_{head}},W_j^V\in {\mathbb{R}}^{C_i\times d_{head}},W^o\in {\mathbb{R}}^{C_i\times C_i}$ 都是线性变换的参数。$N_i$是attention layer中head的数量。因此，每个head的维度$d_{head}$就是$\frac{C_i}{N_i}$。$SR(\cdot )$是降低spatial dimension的操作:
$$SR(x)=\text{Norm}(\text{Reshape}(x,R_i)W^S)$$

其中$x\in {\mathbb{R}}^{H_i{W}_i\times C_i}$代表输入，$R_i$表示reduction ratio。$\text{Reshape}(x,R_i)$表示把输入序列$x$ reshap成大小为$\frac{H_i{W}_i}{R_i^2}\times (R_i^2{C}_i)$的序列。$W^S\in {\mathbb{R}}^{(R_i^2{C}_i)\times C_i}$是线性变化，用来减少输入序列的维度到$C_i$。$\text{Norm}(\cdot )$表示layer normalization。Attention操作为:
$$\text{Attention}(q,k,v)=\text{Softmax}(\frac{q{k}^T}{\sqrt{d_{head}}})v$$

通过对比，计算量会比传统的MHA少$R_i^2$倍。

3. PVT 结构

其中$R_i$是SRA中的reduction ratio。$N_i$是SRA中head的数量。$E_i$是feed-forward层中的expansion ratio。

四、 Conclusion

纯粹的transformer backbone，用于下游的一些任务，比如目标检测，语义分割等等。不需要dense anchors和NMS操作，较为简洁，表现SOTA。

Reference

1. https://zhuanlan.zhihu.com/p/80226180