本文是華為諾亞方舟實驗在Transformer方面的又一次探索，針對現(xiàn)有Transformer存在打破圖像塊的結(jié)構(gòu)信息的問題，提出了一種新穎的同時進(jìn)行patch與pixel表達(dá)建模的TNT模塊，它包含用于塊嵌入建模的Outer Transformer 模塊與像素嵌入建模的Inner Transformer模塊，通過這種方式使得TNT可以同時提取全局與局部結(jié)構(gòu)信息。在ImageNet數(shù)據(jù)集上，TNT-S模型以81.3%的top1精度超過了DeiT-S的的79.8%；TNT-B以82.8%的top1精度超過了DeiT-B的81.8%的top1精度。

Abstract

Transformer是一種自注意力機(jī)制神經(jīng)網(wǎng)絡(luò)，最早興起于NLP領(lǐng)域。近來，純transformer模型已被提出并用于CV的各個領(lǐng)域，比如用于low-level問題的IPT，detection的DETR，classification的ViT，segmentation的SETR等等。然而這些Visual Transformer通過將圖像視作塊序列而忽視了它們最本質(zhì)的結(jié)構(gòu)信息。

針對上述問題，我們提出了一種新穎的Transformer iN Transformer(TNT)模型用于對patch與pixel層面特征建模。在每個TNT模塊中，outer transformer block用于處理塊嵌入，而inner transformer block用于處理像素嵌入的局部特征，像素級特征通過線性變換投影到塊嵌入空間并與塊嵌入相加。通過堆疊TNT模塊，我們構(gòu)建了TNT模塊用于圖像識別。

我們在ImageNet與下游任務(wù)上驗證了所提TNT架構(gòu)的優(yōu)越性，比如，在相似計算復(fù)雜度下，TNT在ImageNet上取得了81.3%的top1精度，以1.5%優(yōu)于DeiT。

Method

接下來，我們將重點描述本文所提TNT架構(gòu)并對其復(fù)雜度進(jìn)行分析。在正式介紹之前，我們先對transformer的一些基本概念進(jìn)行簡單介紹。

Preliminaries

Transformer的基本概念包含MSA(Multi-head Self-Attention)、MLP(Multi-Layer Perceptron)以及LN（Layer Normalization）等。

MSA 在自注意力模塊中，輸入將被線性變換為三部分，即queries, keys, values。其中n表示序列長度，分別表示輸入、queries以及values的維度。此時自注意力機(jī)制可以描述如下：

最后，通過一個線性層生成最終的輸出。而多頭自注意力會將queries、keys、values拆分h次分別實施上述注意力機(jī)制，最后將每個頭的輸出concat并線性投影得到最后的輸出。

MLP MLP是位于自注意力之間的一個特征變換模塊，起定義如下：

其中表示激活函數(shù)，常用GELU，其他參數(shù)則是全連接層的weight與bias，不再贅述。

LN LN 是確保transformer穩(wěn)定訓(xùn)練與快速手鏈的關(guān)鍵部分，起定義如下：

其中，分別表示特征的均值與標(biāo)準(zhǔn)差，o表示點乘操作，為可學(xué)習(xí)變換參數(shù)。

Transformer in Transformer

給定2D圖像，我們將其均勻的拆分為n塊，其中p表示每個圖像塊的大小。ViT一文采用了標(biāo)準(zhǔn)transformer處理塊序列，打破了塊間的局部結(jié)構(gòu)關(guān)系，可參考下圖a。

相反，本文提出了Transformer-iN-Transformer結(jié)構(gòu)同時學(xué)習(xí)圖像的全局與局部信息。在每個TNT模塊中，每個塊通過unfold進(jìn)一步變換到目標(biāo)尺寸，結(jié)合線性投影，塊序列變?yōu)椋?/p>

其中， c表示通道數(shù)量。具體來說，我們將每個塊視作像素嵌入信息：

其中，。

在TNT內(nèi)部，我們具有兩個數(shù)據(jù)流，一個用于跨塊操作，一個用于塊內(nèi)像素操作。對于像素嵌入，我們采用transformer模塊探索像素之間的相關(guān)性：

其中表示層索引，L表示總共層數(shù)。所有塊張量變換為。它可以視作inner transformer block，表示為，該過程構(gòu)建了任意兩個像素之間的相關(guān)性。

在塊層面，我們創(chuàng)建了塊嵌入內(nèi)存以保存塊特征，其中表示類信息，切初始化為0。在每一層，塊張量通過線性投影變換到塊嵌入空間并與塊嵌入相加：

其中表示flatten操作。然后我們采用標(biāo)準(zhǔn)transformer模塊對塊嵌入進(jìn)行變換：

該輸出即為outer transformer block'，它用于建模塊嵌入之間的相關(guān)性。

總而言之，TNT的輸入與輸出包含像素嵌入與塊嵌入，因此TNT可以表示為：

通過堆疊L次TNT模塊，我們即可構(gòu)建一個Transformer-in-Transformer網(wǎng)絡(luò)，最后類別token作為圖像特征表達(dá)，全連接層用于分類。

除了內(nèi)容/特征信息外，空間信息也是圖像識別非常重要的因素。對于塊嵌入與像素嵌入來說，我們同時添加了位置編碼信息，見上圖。這里采用標(biāo)準(zhǔn)1D可學(xué)習(xí)位置編碼信息，具體來說，每個塊被賦予一個位置編碼：

通過這種方式，塊位置編碼可以更好的保持全局空間結(jié)構(gòu)信息，而像素位置編碼可以保持局部相對位置關(guān)系。

Complexity Analysis

對于標(biāo)準(zhǔn)transformer而言，它包含兩部分：MSA與MLP。MSA的FLOPs如下：

而MLP的FLOPs則為。所以，標(biāo)準(zhǔn)transformer的整體FLOPs如下：

一般來說，所以FLOPs可以簡化為，而參數(shù)量則是

本文所提TNT則包含三部分：inner transformer block, outer transformer block與線性層。的計算復(fù)雜度分別為，線性層的FLOPS則是。因此TNT的總體FLOPs則表示如下：

類似的TNT的參數(shù)量表示如下：

盡管TNT添加了兩個額外的成分，但FLOPs提升很小。TNT的Flops大約是標(biāo)準(zhǔn)模塊的1.09x，參數(shù)量大概是1.08x。通過小幅的參數(shù)量與計算量提升，所提TNT模塊可以有效的建模局部結(jié)構(gòu)信息并取得精度-復(fù)雜度的均衡。

Network Architecture

在最終網(wǎng)絡(luò)結(jié)構(gòu)配置方面，我們延續(xù)了ViT與DeiT的配置方式。塊大小為，unfold塊大小。下表給出了TNT網(wǎng)絡(luò)的不同大小的配置信息，它們分別包含23.8M核65.6M參數(shù)量，對應(yīng)的FLOPs分別為5.2B與14.1B(注：輸入圖像尺寸為)。

Operational Optimizations 此外，啟發(fā)與SE，我們進(jìn)行tansformer的通道注意力機(jī)制探索。我們首先對所有patch/pixel嵌入進(jìn)行平均，然后采用兩層MLP計算注意力，所的注意力與所有嵌入相乘。SE模塊僅僅帶來非常少的參數(shù)量，但有助于進(jìn)行通道層面的特征增強(qiáng)。

Experiments

為驗證所提方案的有效性，我們在ImageNet以及其他下游數(shù)據(jù)上進(jìn)行了對比分析，相關(guān)數(shù)據(jù)信息如下所示。

訓(xùn)練超參方面的配置信息如下所示。

我們先來看一下TNT、CNN以及其他Transformer在ImageNet上的性能對比，結(jié)果見下表。

從上表可以看到：

• 所提TNT模型優(yōu)于其他所有Transformer模塊，TNT-S取得了81.3%top-1精度并以1.5%指標(biāo)優(yōu)于DeiT-S；通過添加SE模塊，其性能可以進(jìn)一步提升到81.6%top-1。
• 相比CNN模型，TNT優(yōu)于廣泛采用的ResNet與RegNet。

最后，我們再來看一下在下游任務(wù)的遷移效果，結(jié)果見效果。注：所有模型在分辨率進(jìn)行了微調(diào)。

從上表可以看到：

• 在遷移學(xué)習(xí)方面，TNT取得了比DeiT更優(yōu)的效果；
• 通過更高分辨率的微調(diào)，TNT-B取得了83.9%的top-1精度。

全文到此結(jié)束，更多消融實驗與分析建議各位同學(xué)查看原文。