Transformer 網絡推動了諸多自然語言處理任務的進步，而近期 transformer 開始在計算機視覺領域嶄露頭角。例如，DETR 将目标檢測視為一個直接集預測問題，并使用 transformer 編碼器 - 解碼器架構來解決它；IPT 利用 transformer 在單個模型中處理多個底層視覺任務。與現有主流 CNN 模型（如 ResNet）相比，這些基于 transformer 的模型在視覺任務上也顯示出了良好的性能。
谷歌 ViT（Vision Transformer）模型是一個用于視覺任務的純 transformer 經典技術方案。它将輸入圖檔切分為若幹個圖像塊（patch），然後将 patch 用向量來表示，用 transformer 來處理圖像 patch 序列，最終的輸出做圖像識别。但是 ViT 的缺點也十分明顯，它将圖像切塊輸入 Transformer，圖像塊拉直成向量進行處理，是以，圖像塊内部結構資訊被破壞，忽略了圖像的特有性質。
圖 1：谷歌 ViT 網絡架構。
在這篇論文中，來自華為諾亞實驗室的研究者提出一種用于基于結構嵌套的 Transformer 結構，被稱為 Transformer-iN-Transformer (TNT) 架構。同樣地，TNT 将圖像切塊，構成 Patch 序列。不過，TNT 不把 Patch 拉直為向量，而是将 Patch 看作像素（組）的序列。
論文連結：
https://arxiv.org/pdf/2103.00112.pdf具體而言，新提出的 TNT block 使用一個外 Transformer block 來對 patch 之間的關系進行模組化，用一個内 Transformer block 來對像素之間的關系進行模組化。通過 TNT 結構，研究者既保留了 patch 層面的資訊提取，又做到了像素層面的資訊提取，進而能夠顯著提升模型對局部結構的模組化能力，提升模型的識别效果。
在 ImageNet 基準測試和下遊任務上的實驗均表明了該方法在精度和計算複雜度方面的優越性。例如， TNT-S 僅用 5.2B FLOPs 就達到了 81.3% 的 ImageNet top-1 正确率，這比計算量相近的 DeiT 高出了 1.5%。
方法
圖像預處理
圖像預處理主要是将 2D 圖像轉化為 transformer 能夠處理的 1D 序列。這裡将圖像轉化成 patch embedding 序列和 pixel embedding 序列。圖像首先被均勻切分成若幹個 patch，每個 patch 通過 im2col 操作轉化成像素向量序列，像素向量通過線性層映射為 pixel embedding。而 patch embedding（包括一個 class token）是一組初始化為零的向量。具體地，對于一張圖像，研究者将其均勻切分為 n 個 patch：
其中是 patch 的尺寸。
Pixel embedding 生成：對于每個 patch，進一步通過 pytorch unfold 操作将其轉化成 m 個像素向量，然後用一個全連接配接層将 m 個像素向量映射為 m 個 pixel embedding：
其中
和
，c 是 pixel embedding 的長度。N 個 patch 就有 n 個 pixel embedding 組：
Patch embedding 生成：初始化 n+1 個 patch embedding 來存儲模型的特征，它們都初始化為零：
其中第一個 patch embedding 又叫 class token。
Position encoding：對每個 patch embedding 加一個 patch position encoding：
對每個 pixel embedding 加一個 pixel position encoding：
兩種 Position encoding 在訓練過程中都是可學習的參數。
圖 2：位置編碼。
Transformer in Transformer 架構
TNT 網絡主要由若幹個 TNT block 堆疊構成，這裡首先介紹 TNT block。TNT block 有 2 個輸入，一個是 pixel embedding，一個是 patch embedding。對應地， TNT block 包含 2 個标準的 transformer block。
如下圖 3 所示，研究者隻展示了一個 patch 對應的 TNT block，其他 patch 是一樣的操作。首先，該 patch 對應的 m 個 pixel embedding 輸入到内 transformer block 進行特征處理，輸出處理過的 m 個 pixel embedding。Patch embedding 輸入到外 transformer block 進行特征處理。其中，這 m 個 pixel embedding 拼接起來構成一個長向量，通過一個全連接配接層映射到 patch embedding 所在的空間，加到 patch embedding 上。最終，TNT block 輸出處理過後的 pixel embedding 和 patch embedding。
圖 3：Transformer in Transformer 架構。
通過堆疊 L 個 TNT block，構成了 TNT 網絡結構，如下表 1 所示，其中 depth 是 block 個數，#heads 是 Multi-head attention 的頭個數。
表 1：TNT 網絡結構參數。
實驗
ImageNet 實驗
研究者在 ImageNet 2012 資料集上訓練和驗證 TNT 模型。從下表 2 可以看出，在純 transformer 的模型中，TNT 優于所有其他的純 transformer 模型。TNT-S 達到 81.3% 的 top-1 精度，比基線模型 DeiT-S 高 1.5%，這表明引入 TNT 架構有利于在 patch 中保留局部結構資訊。通過添加 SE 子產品，進一步改進 TNT-S 模型，得到 81.6% 的 top-1 精度。與 CNNs 相比，TNT 的性能優于廣泛使用的 ResNet 和 RegNet。不過，所有基于 transformer 的模型仍然低于使用特殊 depthwise 卷積的 EfficientNet，是以如何使用純 transformer 打敗 EfficientNet 仍然是一個挑戰。
表 2：TNT 與其他 SOTA 模型在 ImageNet 資料集上的對比。
在精度和 FLOPS、參數量的 trade-off 上，TNT 同樣優于純 transformer 模型 DeiT 和 ViT，并超越了 ResNet 和 RegNet 代表的 CNN 模型。具體表現如下圖 4 所示：
圖 4：TNT 與其他 SOTA 模型在精度、FLOPS 和參數量名額上的變化曲線。
特征圖可視化
研究者将學習到的 DeiT 和 TNT 特征可視化，以進一步探究該方法的工作機制。為了更好地可視化，輸入圖像的大小被調整為 1024x1024。此外，根據空間位置對 patch embedding 進行重排，形成特征圖。第 1、6 和 12 個 block 的特征圖如下圖 5(a) 所示，其中每個塊随機抽取 12 個特征圖。與 DeiT 相比，TNT 能更好地保留局部資訊。
研究者還使用 t-SNE 對輸出特征進行可視化（圖 5(b)）。由此可見，TNT 的特征比 DeiT 的特征更為多樣，所包含的資訊也更為豐富。這要歸功于内部 transformer block 的引入，能夠模組化局部特征。
圖 5：DeiT 和 TNT 特征圖可視化。
遷移學習實驗
為了證明 TNT 具有很強的泛化能力，研究者在 ImageNet 上訓練的 TNT-S、TNT-B 模型遷移到其他資料集。更具體地說，他們在 4 個圖像分類資料集上評估 TNT 模型，包括 CIFAR-10、CIFAR-100、Oxford IIIT Pets 和 Oxford 102 Flowers。所有模型微調的圖像分辨率為 384x384。
下表 3 對比了 TNT 與 ViT、DeiT 和其他網絡的遷移學習結果。研究者發現，TNT 在大多數資料集上都優于 DeiT，這表明在獲得更好的特征時，對像素級關系進行模組化具有優越性。
表 3：TNT 在下遊任務的表現。
總結
該研究提出了一種用于視覺任務的 transformer in transformer（TNT）網絡結構。TNT 将圖像均勻分割為圖像塊序列，并将每個圖像塊視為像素序列。本文還提出了一種 TNT block，其中外 transformer block 用于處理 patch embedding，内 transformer block 用于模組化像素嵌入之間的關系。線上性層投影後，将像素嵌入資訊加入到圖像塊嵌入向量中。通過堆疊 TNT block，建構全新 TNT 架構。與傳統的視覺 transformer（ViT）相比，TNT 能更好地儲存和模組化局部資訊，用于視覺識别。在 ImageNet 和下遊任務上的大量實驗都證明了所提出的 TNT 架構的優越性。