ICCV2021 | Tokens-to-Token ViT:在ImageNet上從零訓練Vision Transformer
前言?本文介紹一種新的tokens-to-token Vision Transformer(T2T-ViT)，T2T-ViT將原始ViT的參數(shù)數(shù)量和MAC減少了一半，同時在ImageNet上從頭開始訓練時實現(xiàn)了3.0%以上的改進。通過直接在ImageNet上進行訓練，它的性能也優(yōu)于ResNet，達到了與MobileNet相當?shù)男阅堋?/p>
論文：Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet
?Background
Vision Transformer(ViT)是第一個可以直接應用于圖像分類的全Transformer模型。具體地說，ViT將每個圖像分割成固定長度的14×14或16×16塊(也稱為tokens)；然后ViT應用Transformer層對這些tokens之間的全局關系進行建模以進行分類。
盡管ViT證明了全Transformer架構在視覺任務中很有前途，但在中型數(shù)據(jù)集(例如ImageNet)上從頭開始訓練時，其性能仍遜于類似大小的CNN對等架構(例如ResNets)。
論文假設，這種性能差距源于ViT的兩個主要局限性：
1)通過硬分裂對輸入圖像進行簡單的tokens化，使得ViT無法對圖像的邊緣和線條等局部結(jié)構進行建模，因此它需要比CNN多得多的訓練樣本(如JFT-300M用于預訓練)才能獲得類似的性能；
2)ViT的注意力骨干沒有很好地像用于視覺任務的CNN那樣的設計，如ViT具有冗余性和特征豐富度有限的缺點，導致模型訓練困難。
為了驗證論文的假設，論文進行了一項初步研究，通過圖2中的可視化來調(diào)查ViTL/16和ResNet5的獲知特征的差異。論文觀察ResNet的功能，捕捉所需的局部結(jié)構(邊、線、紋理等)。從底層(Cv1)逐漸向中間層(Cv25)遞增。
然而，ViT的特點卻截然不同：結(jié)構信息建模較差，而全局關系(如整條狗)被所有的注意塊捕獲。這些觀察結(jié)果表明，當直接將圖像分割成固定長度的tokens時，原始 ViT忽略了局部結(jié)構。此外，論文發(fā)現(xiàn)ViT中的許多通道都是零值(在圖2中以紅色突出顯示)，這意味著ViT的主干不如ResNet高效，并且在訓練樣本不足的情況下提供有限的特征豐富度。
圖2.在ImageNet上訓練的ResNet50、ViT-L/16和論文提出的T2T-VIT-24的功能可視化。綠色框突出顯示學習的低級結(jié)構特征，如邊和線；紅色框突出顯示值為零或過大的無效要素地圖。注意：這里為ViT和T2T-ViT可視化的特征圖不是attention圖，而是從tokens重塑的圖像特征。
創(chuàng)新思路
論文決意設計一種新的full-Transformer視覺模型來克服上述限制。
1)與ViT中使用的樸素tokens化不同，論文提出了一種漸進式tokens化模塊，將相鄰tokens聚合為一個tokens(稱為tokens-to-token模塊)，該模塊可以對周圍tokens的局部結(jié)構信息進行建模，并迭代地減少tokens的長度。具體地說，在每個tokens-to-token(T2T)步驟中，transformer層輸出的tokens被重構為圖像(restructurization)，然后圖像被分割成重疊(soft split)的tokens，最后周圍的tokens通過flatten分割的patches被聚集在一起。因此，來自周圍patches的局部結(jié)構被嵌入要輸入到下一transformer層的tokens中。通過迭代進行T2T，將局部結(jié)構聚合成tokens，并通過聚合過程減少tokens的長度。
2)為了尋找高效的Vision Transformer主干，論文借鑒了CNN的一些架構設計來構建Transformer層，以提高功能的豐富性，論文發(fā)現(xiàn)ViT中通道較少但層數(shù)較多的“深度窄”架構設計在同等型號和MAC(Multi-Adds)的情況下性能要好得多。具體地說，論文研究了寬ResNet(淺寬VS深窄結(jié)構)、DenseNet(密集連接)、ResneXt結(jié)構、Ghost操作和通道注意。論文發(fā)現(xiàn)其中，深窄結(jié)構對于ViT是最有效和最有效的，在幾乎不降低性能的情況下顯著地減少了參數(shù)數(shù)目和MACs。這也表明CNNs的體系結(jié)構工程可以為Vision Transformer的骨干設計提供幫助。
基于T2T模塊和深度窄骨干網(wǎng)架構，論文開發(fā)了tokens-to-token Vision Transformer(T2T-ViT)，它在ImageNet上從頭開始訓練時顯著提高了性能，而且比普通ViT更輕便。
?
?Methods
T2T-ViT由兩個主要部分組成(圖4)：
1)一個層次化的“Tokens-to-Token模塊”(T2T模塊)，用于對圖像的局部結(jié)構信息進行建模，并逐步減少tokens的長度；
2)一個有效的“T2T-ViT骨干”，用于從T2T模塊中提取對tokens的全局關注關系。
在研究了幾種基于CNN的體系結(jié)構設計后，對主干采用深窄結(jié)構，以減少冗余度，提高特征豐富性。
圖4.T2T-ViT的整體網(wǎng)絡架構。在T2T模塊中，首先將輸入圖像soft split為patches，然后將其展開為token T0序列。在T2T模塊中，token的長度逐漸減小(在這里使用兩次迭代和輸出Tf)。然后，T2T-VIT主干將固定token作為輸入并輸出預測。兩個T2T塊與圖3相同，PE為位置嵌入。
?
Tokens-to-Token
Tokens-to-Token(T2T)模塊旨在克服ViT中簡單tokens化的限制。它將圖像逐步結(jié)構化為表征，并對局部結(jié)構信息進行建模，這樣可以迭代地減少表征的長度。每個T2T流程有兩個步驟：重組和Soft Split(SS)(圖3)。
圖3.T2T流程圖解。
經(jīng)過變換和reshape后，tokens Ti被重構為圖像Ii，然后重疊split為tokens Ti+1。具體地說，如粉色面板中所示，輸入Ii的四個tokens(1、2、4、5)被串聯(lián)以形成一個tokens 在Ti+1。T2T transformer可以是普通的transformer?層或有限GPU存儲器中的像Performer層這樣的其他高效transformer。
在進行soft split時，每個塊的大小為k×k，在圖像上疊加s個，其中k?類似于卷積運算中的步長。因此，對于重建圖像I_∈_rh×w×c，soft split后的輸出tokens的長度為
每個分割patches的大小為k×k×c。將空間維度上的所有patches展平，以To表示。在soft split之后，為下一個T2T過程饋送輸出tokens。
通過迭代進行上述重構和soft split，T2T模塊可以逐步減少tokens的長度，并轉(zhuǎn)換圖像的空間結(jié)構。T2T模塊中的迭代過程可以表示為
對于輸入image I0，首先應用soft split將其分割為tokens：T1=SS(I0)。在最終迭代后，T2T模塊的輸出tokens Tf具有固定的長度，因此T2T-ViT的主干可以對Tf上的全局關系進行建模。
?
T2T-ViT?Backbone
論文探索了不同的VIT體系結(jié)構設計，并借鑒了CNN的一些設計，以提高骨干網(wǎng)的效率，增強學習特征的豐富性。由于每個transformer層都有跳躍連接，一個簡單的想法是采用如DenseNet的密集連接來增加連通性和特征豐富性，或者采用Wide-ResNets或ResNeXt結(jié)構來改變VIT主干中的通道尺寸和頭數(shù)。
論文探討了從CNN到VIT的五種架構設計：
密集連接如DenseNet；
深-窄與淺-寬結(jié)構如寬ResNet]；
通道注意如擠壓-激勵(SE)網(wǎng)絡；
多頭注意層中更多的分頭如ResNeXt；
Ghost操作如Ghost Net。
實驗發(fā)現(xiàn)：1)采用簡單降低通道維數(shù)的深窄結(jié)構來減少通道中的冗余，增加層深來提高VIT中的特征豐富度，模型尺寸和MACs都有所減小，但性能有所提高；2)SE塊的通道關注度也提高了VIT，但效果不如深窄結(jié)構。
?基于這些發(fā)現(xiàn)，論文為T2T-VIT骨干網(wǎng)設計了一種深窄結(jié)構。具體地說，它具有較小的通道數(shù)和隱藏維度d，但具有更多的層b。對于T2T模塊最后一層定長的Token，論文在其上拼接一個類Token，然后添加正弦位置嵌入(PE)，與VIT一樣進行分類：
T2T-ViT Architecture
T2T-VIT的結(jié)構細節(jié)。T2T-VIT-14/19/24的型號尺寸與ResNet50/101/152相當。T2T-VIT-7/12的型號大小與MobileNetV1/V2相當。對于T2T transformer 層，在有限的GPU內(nèi)存下，論文采用了T2T-VITT-14的transformer層和T2T-VIT-14的Performer層。對于VIT，‘S’表示小，‘B’表示基本，‘L’表示大?！甐IT-S/16’是原始VIT-B/16的變體，具有更小的MLP大小和層深。
?Conclusion
如圖1所示，論文的215M參數(shù)和5.2G MACS的T2T-ViT在ImageNet上可以達到81.5%的TOP-1準確率，遠遠高于ViT的48.6M參數(shù)和10.1G MACs的TOP-1準確率(78.1%)。這一結(jié)果也高于流行的類似大小的CNN，如具有25.5M參數(shù)的ResNet50(76%-79%)。此外，論文還通過簡單地采用更少的層來設計T2T-ViT的精簡變體，取得了與MobileNets(圖1)相當?shù)慕Y(jié)果。
T2T-VIT與VIT在ImageNet上從頭訓練的比較
將CNN中的一些常用設計移植到VIT&T2T-VIT中，包括DenseNet、Wide-ResNet、SE模塊、ResNeXt、Ghost操作。相同的顏色表示相應的遷移。所有模型都是在ImageNet上從頭開始訓練的。
長按掃描下方二維碼添加小助手。
可以一起討論遇到的問題
聲明：轉(zhuǎn)載請說明出處
掃描下方二維碼關注【集智書童】公眾號，獲取更多實踐項目源碼和論文解讀，非常期待你我的相遇，讓我們以夢為馬，砥礪前行！
