摘要

问题：图像有着很复杂且较多的细节，以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模块中，输入 $\in \mathbb{R}^{N \times d}$ 会被线性变换到三个部分，queries $XW_{Q} \in \mathbb{R}^{N \times d_k}$ , keys $W_{K} \in \mathbb{R}^{N \times d_k}$ , 和 values $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操作可以表示为：
$\operatorname{Attention}(Q, K, V)=\operatorname{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层之间，用于特征变换和非线性激活：
$C(\sigma(F C(X))), \quad F C(X)=X W+b$

LN (Layer Normalization)

在transformer中，LN是稳定训练和快速收敛的关键，LN被用在每个样本 $\in \mathbb{R}^{d}$ 上：
$N(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} \rightarrow\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的大小， $\cdots, m$ 。将visual words映射到word embedding：
$Y^{i}=\left[y^{i, 1}, y^{i, 2}, \cdots, y^{i, m}\right], \quad y^{i, j}=F C\left(\operatorname{Vec}\left(x^{i, j}\right)\right),$
其中 $y^{i, j} \in \mathbb{R}^{c}$ 是第 $j$ 个word embedding， $c$ 是word embedding的维度， $\operatorname{Vec}(\cdot)$ 是vectorization操作。

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

对于word embedding，利用inner transformer block $T_{\text {in}}$ 来找到visual words之间的关系：
$\begin{aligned} Y_{l}^{\prime i} &=Y_{l-1}^{i}+\operatorname{MSA}\left(\operatorname{LN}\left(Y_{l-1}^{i}\right)\right) \\ Y_{l}^{i} &=Y_{l}^{\prime i}+\operatorname{MLP}\left(\operatorname{LN}\left(Y_{l}^{\prime i}\right)\right) . \end{aligned}$
其中 $\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}+F C\left(\operatorname{Vec}\left(Y_{l}^{i}\right)\right),$
通过上式相加的操作后sentence embedding的表示通过word-level features进行了增强，然后使用标准的outer transformer block处理上述的sentence embedding：
$\begin{aligned} \mathcal{Z}^{\prime}{ }_{l} &=\mathcal{Z}_{l-1}+\operatorname{MSA}\left(\operatorname{LN}\left(\mathcal{Z}_{l-1}\right)\right), \\ \mathcal{Z}_{l} &=\mathcal{Z}_{l}^{\prime}+\operatorname{MLP}\left(\operatorname{LN}\left(\mathcal{Z}^{\prime}{ }_{l}\right)\right) \end{aligned}$
该outer transformer block $T_{\text {out }}$ 被用于建立sentence embeddings之间的关系。

总之，TNT block的输入和输出同时包含了word embeddings 和 sentence embeddings:
$\mathcal{Y}_{l}, \mathcal{Z}_{l}=T N T\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。

对于sentence：
$\mathcal{Z}_{0} \leftarrow \mathcal{Z}_{0}+E_{\text {sentence }}$ ,
其中 $E_{\text {sentence }} \in \mathbb{R}^{(n+1) \times d}$ 是 sentence 的position encodings。
对于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 是 $d\left(d_{k}+d_{v}\right)+n^{2}\left(d_{k}+d_{v}\right)$ ，multi-layer perceptron的复杂度是 $2 n d_{v} r d_{v}$ ，其中 $r$ 是MLP中隐藏层中的dimension expansion ratio。总之，标准transformer block的FLOPs为：
$\text { FLOPs}_{T}=2 n d\left(d_{k}+d_{v}\right)+n^{2}\left(d_{k}+d_{v}\right)+2 n d d r .$
由于 $r$ 通常设置为4，input, key (query)和 value 的维度通常设置的是一样的，FLOPs可以简化为：
$\mathrm{FLOPs}_{T}=2 n d(6 d+n) .$
参数的数量为：
$\text { Params }_{T}=12 d d .$

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