【TNT】《Transformer in Transformer》

NeurIPS-2021

code：

• https://github.com/huawei-noah/CV-Backbones
• https://gitee.com/mindspore/models/tree/master/research/cv/TNT

---

文章目录

• [1、Background and Motivation](#1、Background and Motivation)
• [2、Related Work](#2、Related Work)
• [3、Advantages / Contributions](#3、Advantages / Contributions)
• 4、Method
• • [4.1、Transformer in Transformer](#4.1、Transformer in Transformer)
  • [4.2、Complexity Analysis](#4.2、Complexity Analysis)
  • [4.3、Network Architecture](#4.3、Network Architecture)
• 5、Experiments
• • [5.1、Datasets and Metrics](#5.1、Datasets and Metrics)
  • [5.2、TNT on ImageNet](#5.2、TNT on ImageNet)
  • [5.3、Ablation Studies](#5.3、Ablation Studies)
  • 5.4、Visualization
  • [5.5、Transfer Learning](#5.5、Transfer Learning)
• [6、Conclusion（own） / Future work](#6、Conclusion（own） / Future work)

---

1、Background and Motivation

Background

传统视觉 Transformer（如 ViT）将图像划分成若干个较大的 Patch，并把每个 Patch 当作一个 Token，通过 Transformer 建模不同 Patch 之间的关系。

这种方法能够捕获图像的全局信息，但一个 Patch 内部的像素和局部区域会被压缩成一个整体表示。

---

Motivation

作者认为：

一个大 Patch 内部仍然包含许多有意义的局部结构，直接把它压缩成一个 Token，可能会丢失细粒度信息。

例如在人脸图像中，一个 Patch 内可能同时包含眼睛、眉毛和额头等区域；这些小区域之间的关系不能被简单忽略。

因此，TNT 采用两级划分：

1. 先把图像划分为较大的 Patch，称为 Visual Sentence；
2. 再把每个 Patch 划分为更小的 Sub-patch，称为 Visual Word。

随后使用：

• Inner Transformer：建模同一个 Patch 内部 Word 之间的局部关系；
• Outer Transformer：建模不同 Patch 之间的全局关系。

---

一句话概括

ViT 主要关注 Patch 之间的全局关系，而 TNT 在此基础上进一步挖掘 Patch 内部的局部结构，从而同时获得更细致的局部信息和完整的全局信息。

---

2、Related Work

1. ViT

   将图像划分为多个 Patch，并把 Patch 当作 Token，利用自注意力建模不同区域之间的关系，证明了 Transformer 可以用于图像识别。

2. DeiT

   通过更高效的数据训练和知识蒸馏，降低了 ViT 对大规模数据的依赖，提高了训练效率和 ImageNet 性能。

3. IPT

   将 Transformer 用于去噪、超分辨率等低层视觉任务，说明 Transformer 不仅适合分类，也能处理其他图像任务。

4. 其他视觉任务

   Transformer 逐渐被应用到图像分类、目标检测和语义分割等任务中。

---

存在的问题

现有方法大多沿用 ViT 的做法：

图像 → 大 Patch → Transformer \text{图像}\rightarrow\text{大 Patch}\rightarrow\text{Transformer} 图像→大 Patch→Transformer

这种方式能够学习不同 Patch 之间的全局关系，但会忽略一个大 Patch 内部更细小区域之间的局部关系。

TNT 在 Patch 内部进一步划分更小的 Sub-patch：

• 大 Patch：Visual Sentence；
• 小 Patch：Visual Word；
• Inner Transformer：建模 Word 之间的局部关系；
• Outer Transformer：建模 Sentence 之间的全局关系。

一句话概括：

以往视觉 Transformer 主要研究"Patch 与 Patch 之间的关系"，TNT 进一步研究"Patch 内部小区域之间的关系"，从而弥补局部细节建模不足的问题。

---

3、Advantages / Contributions

1. 提出 TNT 双层 Transformer 结构
   作者提出 Transformer-in-Transformer（TNT）：

• Inner Transformer：处理 Patch 内部的小区域，提取局部细节；
• Outer Transformer：处理不同 Patch，学习全局关系。
  因此，模型能够同时建模：
  局部信息 + 全局信息 \text{局部信息}+\text{全局信息} 局部信息+全局信息

2. 保留了 Patch 内部的细粒度结构
   普通 ViT 将一个 Patch 直接压缩成一个 Token，可能丢失内部结构。

TNT 将 Patch 进一步划分为多个 Word，使模型能够识别：

• 边缘；
• 纹理；
• 小目标；
• 局部区域之间的关系。

3. 额外计算开销较小
   不同 Patch 之间共享 Inner Transformer 参数，并且 Word embedding 维度 c c c 小于 Sentence embedding 维度 d d d。

因此，在增加局部建模能力的同时：

• FLOPs 增加较少；
• 参数量增加较少；
• 计算效率仍然较高。

4. 性能和效率之间取得更好平衡
   实验表明，TNT 在 ImageNet 以及分类、检测、分割等下游任务上具有较强表现，在准确率和计算量之间取得了较好的折中。
5. 系统分析了 Word 划分方式
   作者还研究了不同 Word 尺寸和数量对模型性能、计算量的影响，为 TNT 的结构设计提供了实验依据。

一句话总结

TNT 的核心贡献是：在 ViT 的 Patch 级 Transformer 内部嵌入一个 Word 级 Transformer，用较小的额外代价补充局部细节，从而同时获得更充分的局部信息和全局信息。

---

4、Method

4.1、Transformer in Transformer

m sub-patches

j-th visual word of the i-th visual sentence

这里的 l l l 表示 Transformer 的层数/第几个 Block

Image → Patch → Word embedding Y → Inner Transformer → Word aggregation → Sentence embedding Z → Outer Transformer \text{Image} \rightarrow \text{Patch} \rightarrow \text{Word embedding }Y \rightarrow \text{Inner Transformer} \rightarrow \text{Word aggregation} \rightarrow \text{Sentence embedding }Z \rightarrow \text{Outer Transformer} Image→Patch→Word embedding Y→Inner Transformer→Word aggregation→Sentence embedding Z→Outer Transformer

the word position encodings which are shared across sentences

4.2、Complexity Analysis

第一项：MSA 中的线性投影

---

第二项：MSA 中的注意力交互

---

第三项：MLP 计算量

MLP 通常由两个全连接层组成：

d → d r → d d\rightarrow d_{r}\rightarrow d d→dr→d

其中 d r d_{r} dr 是 MLP 中间隐藏层的维度。

假设输入为：

X ∈ R n × d X\in\mathbb{R}^{n\times d} X∈Rn×d

第一层全连接：

X W 1 XW_1 XW1

其中：

W 1 ∈ R d × d r W_1\in\mathbb{R}^{d\times d_{r}} W1∈Rd×dr

计算量：

F L O P s M L P 1 = n d d r FLOPs_{MLP1}=ndd_{r} FLOPsMLP1=nddr

第二层全连接：

H W 2 HW_2 HW2

其中：

W 2 ∈ R d r × d W_2\in\mathbb{R}^{d_{r}\times d} W2∈Rdr×d

计算量：

F L O P s M L P 2 = n d r d FLOPs_{MLP2}=nd_{r}d FLOPsMLP2=ndrd

两层相加：

F L O P s M L P = 2 n d d r \boxed{FLOPs_{MLP}=2ndd_{r}} FLOPsMLP=2nddr

这里的 2 就是因为 MLP 有两个全连接层。

---

论文进一步假设：

• d k = d v = d d_k=d_v=d dk=dv=d
• MLP 扩展比例为 4，即 d t r = 4 d d_{tr}=4d dtr=4d

代入公式（14） 得到公式 （15）

F L O P s T = 2 n d ( 2 d ) + n 2 ( 2 d ) + 2 n d ( 4 d ) = 4 n d 2 + 2 n 2 d + 8 n d 2 = 12 n d 2 + 2 n 2 d = 2 n d ( 6 d + n ) \begin{aligned} FLOPs_T &=2nd(2d)+n^2(2d)+2nd(4d)\\ &=4nd^2+2n^2d+8nd^2\\ &=12nd^2+2n^2d\\ &=\boxed{2nd(6d+n)} \end{aligned} FLOPsT=2nd(2d)+n2(2d)+2nd(4d)=4nd2+2n2d+8nd2=12nd2+2n2d=2nd(6d+n)

---

参数量

---

作者的方法，计算量

vs 标准 transformer

---

第一项：Inner Transformer

F L O P s I n n e r = 2 n m c ( 6 c + m ) \boxed{FLOPs_{Inner}=2nmc(6c+m)} FLOPsInner=2nmc(6c+m)

第二项：Linear Layer（Word-to-Sentence）

第三项：Outer Transformer

直接代入标准 Transformer 的 FLOPs 公式：

F L O P s O u t e r = 2 n d ( 6 d + n ) \boxed{FLOPs_{Outer}=2nd(6d+n)} FLOPsOuter=2nd(6d+n)

它负责建模不同 Patch 之间的全局关系。

参数量

4.3、Network Architecture

patch size is set as 16×16

sub-patches is set as m = 4 · 4 = 16 by default

---

5、Experiments

分类增广策略同 DeiT，COCO 和 ADE20K 的增广策略同 PVT

---

5.1、Datasets and Metrics

分类 top1，检测 AP，分割 mIoU

5.2、TNT on ImageNet

it is yet a challenge of how to beat EfficientNet using pure transformer

inference speed

TNT 结构比较高效，may need fewer blocks to learn its representation.

所以为了最求性能，可以把部分 TNT block 替换为传统 transformer block

5.3、Ablation Studies

Effect of position encodings

Number of heads

head width of 64

Number of visual words

When we change m m m, the embedding dimension c c c also changes correspondingly to control the FLOPs

5.4、Visualization

Visualization of Feature Maps

the features of TNT are more diverse and contain richer information than those of DeiT

Visualization of Attention Maps

5.5、Transfer Learning

Pure Transformer Image Classification

Pure Transformer Object Detection

Pure Transformer Semantic Segmentation

---

6、Conclusion（own） / Future work

• transformer in transformer（TNT）

• 介绍了 transformer 的计算量和参数量

• visual word embedding and visual sentence embedding

• inner self-attention and an outer self-attention

• build TNT architecture by stacking the TNT blocks

• 

• TNT 通过"局部 Transformer + 全局 Transformer"，弥补了 ViT 只关注 Patch 之间关系、容易忽略 Patch 内部细节的问题。

• 

• Word Position Embedding 也是在所有 Patch 之间共享的

---

更多论文解读，请参考 【Paper Reading】