【点云处理之论文狂读扩展版2】—— Transformer in Transformer

摘要

• 问题：图像有着很复杂且较多的细节，以patch为单位进行划分还远不能在不同的尺度和位置上发掘目标的特征
• 解决方法：局部patch中的attention对于构建高性能的视觉transformer也是很重要的，因此构建了一个新的架构 Transformer iN Transformer (TNT)
• 技术细节：
  (1)将局部patch(16×16)看作为visual sentences，然后进一步将他们划分为更小的patch(4×4)，称为visual words
  (2)每个word的注意力将用给定的visual sentences中的其他wordd来计算，计算成本可以忽略不计
  (3)words和sentences的特征会被聚合以提高表示能力
• 代码详见https://github.com/huawei-noah/CV-Backbones

1.引言

• Self-attention ——> Non-local networks
• 与NLP任务不同，在CV任务中输入图像和ground truth label之间还存在着semantic gap ——> ViT
• 首先将局部patch看作为visual sentences，然后进一步将他们划分为更小的patch，称为visual words
• 然后在网络中嵌入一个sub-transformer挖掘visual words的特征和细节
• 每个visual sentences中，visual words间的特征和attention都使用shared network共同和计算，复杂度和参数可忽略不计
• words的特征将会和对应的visual sentences进行聚合
• class token可继续用于下游任务

2.方法

2.1 Preliminaries

MSA (Multi-head Self-Attention)

在Self-attention模块中，输入$X\in {\mathbb{R}}^{N\times d}$会被线性变换到三个部分，queries $Q=X{W}_Q\in {\mathbb{R}}^{N\times d_k}$, keys $K=X{W}_K\in {\mathbb{R}}^{N\times d_k}$, 和 values $V=X{W}_V\in {\mathbb{R}}^{N\times d_v}$，其中$N$是序列长度，$d,d_k,d_v$分别是输入、queries (keys) 和 values的维度。$W_Q\in {\mathbb{R}}^{d\times d_k},W_K\in {\mathbb{R}}^{d\times d_k},W_V\in {\mathbb{R}}^{d\times d_v}$是要学习的参数。The scaled dot-product attention操作可以表示为：
$$\mathrm{Attention}(Q,K,V)=\mathrm{Softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$$
最后，线性层用于产生输出。Multi-head Self-Attention便是将queries, keys 和 values分成$h$个部分，并行执行attention函数，然后将每个head的输出进行拼接，映射到最终的输出。

MLP (Multi-Layer Perceptron)

MLP被用在Self-attention层之间，用于特征变换和非线性激活：
$$MLP(X)=FC(\sigma (FC(X))),FC(X)=XW+b$$

LN (Layer Normalization)

在transformer中，LN是稳定训练和快速收敛的关键，LN被用在每个样本$x\in {\mathbb{R}}^d$ 上：
$$LN(x)=\frac{x-\mu }{\delta }\circ \gamma +\beta$$

其中$\mu \in \mathbb{R},\delta \in \mathbb{R}$分别是特征的平均值和标准差，$\circ$表示元素间的乘法，$\gamma \in {\mathbb{R}}^d,\beta \in {\mathbb{R}}^d$是应该学习的仿射变换参数。

2.2Transformer in Transformer

给定一张2D图像，均匀将其分为$n$个patch$\mathcal{X}=\left[X^1,X^2,\cdots ,X^n\right]\in {\mathbb{R}}^{n\times p\times p\times 3}$，其中$(p,p)$是每个patch的像素。TnT将这些块看作是visual sentences，每个patch进一步被分为$m$个visual words：

$$X^i\to \left[x^{i,1},x^{i,2},\cdots ,x^{i,m}\right],$$

其中$x^{i,j}\in {\mathbb{R}}^{s\times s\times 3}$表示第$i$个visual sentence的第$j$个visual word， $(s,s)$ 是sub-patches的大小，$j=1,2,\cdots ,m$。将visual words映射到word embedding：
$$Y^i=\left[y^{i,1},y^{i,2},\cdots ,y^{i,m}\right],y^{i,j}=FC\left(\mathrm{Vec}\left(x^{i,j}\right)\right),$$
其中$y^{i,j}\in {\mathbb{R}}^c$是第$j$个word embedding，$c$ 是word embedding的维度，$\mathrm{Vec}(\cdot )$是vectorization操作。

TnT中，有两个分支，一个处理visual sentence，另外一个处理visual sentence里的visual words。

• 对于word embedding，利用inner transformer block $T_{\text{in}}$来找到visual words之间的关系：
  $$\begin{array}{rl}{Y}_l^{\prime i}& =Y_{l-1}^i+\mathrm{MSA}\left(\mathrm{LN}\left(Y_{l-1}^i\right)\right)\\ Y_l^i& =Y_l^{\prime i}+\mathrm{MLP}\left(\mathrm{LN}\left(Y_l^{\prime i}\right)\right).\end{array}$$
  其中$l=1,2,\cdots ,L$表示第$l$个block，$L$是blocks的总数。输出结果表示为${\mathcal{Y}}_l=\left[Y_l^1,Y_l^2,\cdots ,Y_l^n\right]$。

• 在处理sentence时，先将sentence转换为 sentence embedding ${\mathcal{Z}}_0=\left[Z_{\text{class }},Z_0^1,Z_0^2,\cdots ,Z_0^n\right]\in {\mathbb{R}}^{(n+1)\times d}$， 其中 $Z_{\text{class }}$ 是class token，初始值设置为0。 在每一层中，word embeddings的序列会通过线性映射的方式转换到sentence embedding的空间域，并与sentence embedding进行相加：
  $$Z_{l-1}^i=Z_{l-1}^i+FC\left(\mathrm{Vec}\left(Y_l^i\right)\right),$$
  通过上式相加的操作后sentence embedding的表示通过word-level features进行了增强，然后使用标准的outer transformer block处理上述的sentence embedding：
  $$\begin{array}{rl}{\mathcal{Z}}^{\prime }{}_l& ={\mathcal{Z}}_{l-1}+\mathrm{MSA}\left(\mathrm{LN}\left({\mathcal{Z}}_{l-1}\right)\right),\\ {\mathcal{Z}}_l& ={\mathcal{Z}}_l^{\prime }+\mathrm{MLP}\left(\mathrm{LN}\left({\mathcal{Z}}^{\prime }{}_l\right)\right)\end{array}$$
  该outer transformer block $T_{\text{out }}$被用于建立sentence embeddings之间的关系。

总之，TNT block的输入和输出同时包含了word embeddings 和 sentence embeddings:
$${\mathcal{Y}}_l,{\mathcal{Z}}_l=TNT\left({\mathcal{Y}}_{l-1},{\mathcal{Z}}_{l-1}\right).$$

在TNT block中，inner transformer block被用于构建visual words之间的关系，提取局部特征。outer transformer block捕获sentences序列之间的固有信息。通过将TNT block堆叠 $L$ 次，我们构建了transformerin-transformer网络。

Position encoding

空间信息很重要，在sentence embeddings and word embeddings中都加了对应的position encodings，在本文中使用了标准的1D position encodings。

1. 对于sentence：
   ${\mathcal{Z}}_0\leftarrow {\mathcal{Z}}_0+E_{\text{sentence }}$,
   其中$E_{\text{sentence }}\in {\mathbb{R}}^{(n+1)\times d}$是 sentence 的position encodings。
2. 对于sentence中的word，便是将word position encoding与word position encoding相加：
   $$Y_0^i\leftarrow Y_0^i+E_{\text{word }},i=1,2,\cdots ,n$$
   其中$E_{\text{word }}\in {\mathbb{R}}^{m\times c}$是word的position encodings，该position encodings在sentence中是共享权值的。

这样，sentence position encoding可以保留全局空间信息，word position encoding能够用于保留局部相对信息。

2.3 Complexity Analysis

标准的transformer block包含了两个部分， the multi-head self-attention 和 multi-layer perceptron。multi-head self-attention的FLOPs 是$2nd\left(d_k+d_v\right)+n^2\left(d_k+d_v\right)$，multi-layer perceptron的复杂度是$2n{d}_{v}r{d}_v$ ，其中$r$是MLP中隐藏层中的dimension expansion ratio。总之，标准transformer block的FLOPs为：
$${\text{ FLOPs}}_T=2nd\left(d_k+d_v\right)+n^2\left(d_k+d_v\right)+2nddr.$$
由于$r$通常设置为4，input, key (query)和 value 的维度通常设置的是一样的，FLOPs可以简化为：
$${\mathrm{F}\mathrm{L}\mathrm{O}\mathrm{P}\mathrm{s}}_T=2nd(6d+n).$$
参数的数量为：
$${\text{ Params }}_T=12dd.$$

本文提出的transformer block包含了3个部分，an inner transformer block $T_{in}$，an outer transformer block $T_{\text{out }}$ and a linear layer。$T_{\text{in }}$ 和 $T_{\text{out }}$的FLOPs分别是$2nmc(6c+m)$ 和 $2nd(6d+n)$，linear layer的FLOPs是$nmcd$。总之TnT block的FLOPs为：
$${\mathrm{F}\mathrm{L}\mathrm{O}\mathrm{P}\mathrm{s}}_{TNT}=2nmc(6c+m)+nmcd+2nd(6d+n).$$
参数为：
$${\mathrm{P}\mathrm{a}\mathrm{r}\mathrm{a}\mathrm{m}\mathrm{s}}_{TNT}=12cc+mcd+12dd.$$
尽管在TnT中加了额外两项，但是FLOPs增加的数量很小，因为$c\ll d$ ， $\mathcal{O}(m)\approx \mathcal{O}(n)$ 。以较小的牺牲换取较高的性能，很nice。

2.4 Network Architecture

默认情况下，patch的大小为16×16，sub-patch的大小为$m=4\cdot 4=16$

TnT有三种变体：

• TNT-Ti 6.1M
• TNT-S 23.8M
• TNT-B 65.6M

3. 实验

下次再写

生词

• granularity n.粒度
• excavate v. 挖掘