論文提出了T2T-ViT模型，引入tokens-to-token（T2T）模塊有效地融合圖像的結構信息，同時借鑑CNN結果設計了deep-narrow的ViT主幹網絡，增強特徵的豐富性。在ImageNet上從零訓練時，T2T-ViT取得了優於ResNets的性能MobileNets性能相當

來源：曉飛的算法工程筆記公衆號

論文: Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet

論文地址：https://arxiv.org/abs/2101.11986
論文代碼：https://github.com/yitu-opensource/T2T-ViT

Introduction

儘管ViT證明了純Transformer架構對於視覺任務很有前景，但在中型數據集（例如ImageNet）上從零訓練時，其性能仍然不如大小類似的CNN網絡（例如 ResNets）。

論文認爲這種性能差距源於ViT的兩個主要限制：

簡單地對輸入圖像分割成14x14或16x16的token序列使得ViT無法對圖像的局部結構（如邊緣和線條）建模，需要更多的訓練樣本（如JFT-300M用於預訓練）才能與CNN有相似的性能。
ViT的主幹網絡沒有像CNN那樣爲視覺任務進行精心設計，包含了大量的冗餘結構，特徵豐富程度有限，模型訓練困難。

爲了驗證，論文對ViTL/16和ResNet50學習到的特徵進行可視化對比。如圖2所示，ResNet逐層捕獲所需的局部結構信息（邊緣、線條、紋理等），而ViT特徵的結構信息建模不佳，所有注意力塊都捕獲全局關係（例如，整隻狗）。這表明，ViT將圖像拆分爲具有固定長度的token時忽略了局部結構。此外，論文發現ViT中的許多通道的值爲零，這意味着ViT的主幹網絡不如ResNets高效。如果訓練樣本不足，則只能提供特徵的豐富度有限。

基於上面的觀察，論文設計了一個新的Vision Transformer模型來克服上述限制：

提出了一種漸進式的token生成模塊Tokens-to-Token，通過transformer層提取特徵並將相鄰的token聚合爲一個token，代替ViT中將圖像簡單分割爲token的行爲。該模塊能夠迭代地對周圍toekn的局部結構信息進行建模並減少token序列長度。
爲了設計高效的Vision Transformer主幹網絡，提高特徵豐富度，論文從CNN中借用一些結構設計ViT主幹網絡。論文發現，通道數較少、層數較多的“deep-narrow”架構設計能夠顯著減少ViT模型的大小和MAC（Multi-Adds），而性能幾乎沒有下降。這表明CNN的架構優化可以借鑑到Vision Transformer主幹的設計。

基於T2T模塊和deep-narrow主幹架構，論文設計了Tokens-to-Token Vision Transformer (T2T-ViT)。對比原生的ViT，在ImageNet上從零開始訓練的性能有顯着的提高，與CNN網絡相當甚至更好。

總體言之，論文的貢獻有三方面：

通過精心設計的Transformer架構（T2T模塊和高效主幹網絡）證明，Vision Transformer可以無需JFT-300M上的預訓練，在ImageNet上以不同的複雜度勝過CNN。
爲ViT開發了一種新穎的漸進式token生成策略T2T模塊，更好地融合圖像結構信息，優於ViT的簡單token生成方法。
驗證CNN的架構優化可以用於ViT的主幹網絡設計，提高特徵豐富度並減少冗餘。通過大量實驗，deep-narrow的架構設計最適合ViT。

Tokens-to-Token ViT

爲了克服ViT的簡單token生成和低效主幹網絡的侷限性，論文提出了Tokens-to-Token Vision Transformer(T2T-ViT)，可以逐步將圖像轉換爲token並且主幹網絡更高效。因此，T2T-ViT由兩個主要組件組成：

一個多層的Tokens-to-Token（T2T）模塊，用於對圖像的局部結構信息進行建模並逐漸減少token數量。
一個高效的T2T-ViT主幹網絡，用於對T2T模塊生成的token提取全局注意力關係。在探索了幾種基於CNN的架構設計後，論文采用了一種deep-narrow結構來減少冗餘並提高特徵豐富度。

Tokens-to-Token: Progressive Tokenization

Token-to-Token(T2T)模塊主要爲了克服ViT中簡單token生成的限制，逐步將圖像結構化爲token以及對局部結構信息進行建模，並且可以迭代地減少token數量。每個T2T操作都包含兩個步驟：Re-structurization和Soft Split(SS)。

Re-structurization

如圖 3 所示，給定token序列\(T\)，先通過自注意模塊（T2T Transformer）進行變換：

其中MSA爲具有層歸一化的多頭自注意操作，MLP是標準Transformer中具有層歸一化的多層感知器。MSA輸出的\(T^{'}\)將被重塑爲空間維度上的圖像：

Reshape表示將\(T^{'}\in \mathbb{R}^{l\times c}\)重新組織爲\(I\in \mathbb{R}^{h\times w\times c}\)，其中\(l\)是\(T^{'}\)的長度，h、w、c 分別是高度、寬度和通道數，並且\(l=h\times w\)。

Soft Split

如圖3所示，在獲得重構圖像\(I\)後，使用Soft Split來建模局部結構信息並減少token的長度。爲了避免信息丟失，將圖像拆分爲重疊的分割區域，每個區域都與周圍的區域相關。這樣就建立了一個先驗，即相鄰分割區域生成的token之間應該有更強的相關性。隨後將每個分割區域中的token拼接爲一個token，從周圍的像素或token中聚合局部信息。

進行Soft Split時，每個分割區域的大小爲\(k\times k\)，區域重疊爲\(s\)，圖像邊界填充爲\(p\)，其中\(k-s\)類似於卷積操作中的步長。對於重建圖像\(I\in \mathbb{R}^{h\times w\times c}\)，Soft Split後輸出的token \(T_{o}\)的長度爲：

每個分割區域的大小爲\(k\times k\times c\)，將所有分割區域展平後得到token序列\(T_{o}\in \mathbb{R}^{l_{o}\times ck^2}\)。在Soft Split之後，輸出token可進行下一輪T2T操作。

T2T module

通過反覆進行Re-structurization和Soft Split，T2T模塊可以逐步減少token的長度以及變換圖像的空間結構。T2T模塊的迭代過程可以表述爲：

對於輸入圖像\(I_{0}\)，先應用Soft Split將其拆分爲token序列\(T_{1} = SS(I_{0})\)。在最後一次迭代之後，T2T模塊的輸出固定長度的token序列\(T_{f}\)。因此，T2T-ViT 的主幹網絡可以在\(T_{f}\)上建模全局關係。

此外，由於T2T模塊中的token長度大於ViT中的一般設置（16 × 16），MAC和內存使用量都很大。爲了解決這個問題，將T2T層的通道維度設置爲較小的值（32或64）來減少 MAC，也可以採用高效的Transformer層變種，例如 Performer層，從而在有限的GPU內存下減少內存使用。

T2T-ViT Backbone

由於ViT主幹網絡中許多通道是無效的，論文打算爲T2T-ViT重新設計一個高效的主幹網絡，減少冗餘並提高特徵豐富度。論文借鑑了CNN的一些設計，探索不同的ViT架構設計。由於每個Transformer層都具有ResNets的短路連接，可以參考DenseNet增加特徵複用和特徵豐富程度，或者參考Wide-ResNets和ResNeXt調整通道維度和head數。

論文在ViT上探索了以下五種CNN的架構設計：

Dense connection as DenseNet。
Deep-narrow vs. shallow-wide structure as in Wide-ResNets。
Channel attention as Squeeze-an-Excitation(SE) Networks。
More split heads in multi-head attention layer as ResNeXt。
Ghost operations as GhostNet。

論文對以上結構移植進行了實驗，有以下兩點發現：

採用deep-narrow結構，減小通道尺寸可以減少通道冗餘，增加層深度可以提高特徵豐富度。不僅模型大小和MAC都減小了，性能還得到了提高。
SE模塊的通道注意力也能提升ViT，但不如deep-narrow結構有效。

基於這些發現，論文爲T2T-ViT主幹網絡設計了一個 deep-narrow的架構，具有較小的通道數和隱藏維度\(d\)，但層數\(b\)更多。對於T2T模塊輸出的固定長度的token序列\(T_{f}\)，爲其添加一個class token，然後加入Sinusoidal Position Embedding(PE)，最後與ViT一樣進行分類：

其中，\(E\)是Sinusoidal Position Embedding，LN是層歸一化，fc是用於分類的全連接層，\(y\)是輸出預測。

T2T-ViT Architecture

T2T-ViT包含兩部分：Tokens-to-Token(T2T)模塊和T2T-ViT主幹網絡。T2T模塊有多種設計選擇，論文設置\(n = 2\)，T2T模塊中有\(n+1=3\)次Soft Split和\(n=2\)次Re-structurization。三次Soft Split的分區區域設置爲\(P = [7, 3, 3]\)，重疊區域設置爲\(S=[3, 1, 1]\)，可以將\(224\times 224\)的輸入圖片壓縮爲\(14\times 14\)的token序列。

T2T-ViT主幹網絡從T2T模塊中取固定長度token序列作爲輸入，基於deep-narrow架構設計，中間特徵維度(256-512)和MLP大小(512-1536)比ViT小很多。例如，T2T-ViT-14的主幹網絡中有14個Transofmer層，中間特徵維度爲384，而ViT-B/16有12個Transformer層，中間特徵維度爲768，參數量和MACs是T2T-ViT-14的3倍。

爲了方便與ResNet進行比較，論文設計了三個的T2T-ViT模型：T2T-ViT-14、T2T-ViT-19 和 T2T-ViT-24，參數量分別與ResNet50、ResNet101和ResNet152相當。而爲了與MobileNets等小型模型進行比較，論文設計了兩個lite模型：T2T-ViT-7、T2TViT-12，其模型大小與MibileNetV1和MibileNetV2相當。兩個lite TiT-ViT沒有使用特殊設計或技巧，只是簡單地降低了層深度、中間特徵維度以及MLP比例。