大道至簡 | 設計 ViT 到底怎么配置Self-Attention才是最合理的？

Transformer已成為深度學習中的主要架構之一，尤其是作為計算機視覺中卷積神經網絡 (CNN) 的強大替代品。然而，由于Self-Attention在長序列表示上的二次復雜性，特別是對于高分辨率密集預測任務，先前工作中的Transformer訓練和推理可能非常昂貴。為此，我們提出了一種新穎的少注意力視覺Transformer（LIT），基于Transformers中早期的Self-Attention仍然專注于局部模式，并在最近的分層視覺Transformers中帶來較小的好處。

具體來說，提出了一個分層 Transformer，使用純多層感知器（MLP）在早期階段對豐富的局部模式進行編碼，同時應用Self-Attention模塊在更深層捕獲更長的依賴關系。此外，進一步提出了一種 Learned Deformable Token Merging Module，以非均勻方式自適應地融合信息Patch。

所提出的 LIT 在圖像識別任務（包括圖像分類、對象檢測和實例分割）上取得了可觀的性能，可作為許多視覺任務的強大支柱。

開源地址：https://github.com/zhuang-group/LIT

Transformers在自然語言處理(NLP)和最近在計算機視覺(CV)領域取得了長足的進步。受CNN中金字塔設計的啟發，最近的分層視覺轉換器(HVT)將Transformer Block劃分為多個階段，并隨著網絡的深入逐漸縮小特征圖。然而，早期的高分辨率特征圖導致了很長的token序列，由于Self-Attention的二次復雜度，帶來了巨大的計算成本和內存消耗。例如，大小為56×56×96 的特征圖在一個Multi-Head Self-Attention（MSA）中需要2.0G FLOPs，而ResNet-18的整個模型只需要1.8G FLOPs。如此巨大的計算成本使得將Transformer應用到廣泛的計算機視覺任務中變得非常困難。

在HVT的早期階段，已經做出了一些努力來降低計算成本。例如，一些工作減少了MSA層中Self-Attention Head的數量或進一步減少了Transformer Block的數量。另一行工作建議通過啟發式近似來權衡MSA的準確性和效率，例如空間縮減注意力(SRA)和基于Shift Window Multi-Head Self-Attention(SW-MSA)。也有研究當特征圖的分辨率相當大時使用卷積層。然而，早期采用Self-Attention層對最終性能有多大貢獻仍不清楚。

在本文中提出了一種Less attention Vision Transformer（LIT）來解決上述HVT問題。具體來說，建議在早期Stage使用MLP層來捕獲局部信息，同時引入具有足夠數量的Head來處理的MSA層后期的長期依賴關系。

作者的動機來自2個方面。首先，先前的CNN和Transformer相關研究表明，模型淺層會專注于局部信息，而深層傾向于捕獲高級語義或全局關系，這就產生了在早期Stage是否有必要使用Self-Attention的問題。其次，從理論的角度來看，一個具有足夠Head的Self-Attention應用于圖像可以表達任何卷積層。

然而，MSA層中較少的Head理論上阻礙了逼近具有大kernel-size的卷積層的能力，其中極端情況與1×1卷積一樣具有表現力，可以被視為獨立應用于每個像素的標準FC層。雖然最近的HVT在早期階段采用很少的Head來提供金字塔表示，但作者認為這不是最優的，因為這樣的配置會引入高計算和內存成本，但帶來的收益卻很小。

需要強調的是，通過在前期Stage利用MLP Block可以避免Self-Attention在高分辨率特征圖上產生的巨大計算成本和內存占用。此外，在后期Stage應用Self-Attention來捕獲遠程依賴關系。綜合實驗結果表明，經過如此簡單的架構設計后可以帶來模型性能和效率之間的最佳平衡點。

此外，最近的HVT要么采用標準卷積層，要么采用線性投影層來合并鄰近的Token，旨在控制特征圖的規模。然而，考慮到并非每個像素對輸出單元的貢獻都相同，這種方法阻礙了Vision Transformer對幾何變換建模的表示能力。

為此，作者提出了一個受Deformable Convolutions啟發的Deformable Token Merging(DTM) 模塊，在該模塊中，學習了一個偏移網格以自適應地增加空間采樣位置，進而合并來自特征圖中子窗口的相鄰Patch。通過這種方式可以獲得更多信息的下采樣Token用于后續處理。

• 首先，確定了早期MSA層在最近的HVT中的微小貢獻，并在早期Stage提出了一個簡單的HVT結構，其中包含純MLP Block;

• 其次，提出了一個Deformable Token Merging模塊，以自適應地合并更多信息Patch以提供分層表示，并具有增強建模能力;

• 最后，進行了廣泛的實驗，以表明所提出的LIT在計算復雜度和內存消耗方面的有效性。

圖1LIT的整體架構如圖1所示。設為輸入的RGB圖像，其中H和W分別表示高度和寬度。首先將I拆分為不重疊的Patch，Patch大小為4×4，因此每個Patch的初始特征維度為4×4×3=48。接下來，利用線性嵌入層將每個Patch投影到維度，用作以下過程的初始輸入。整個模型分為4個Stage。讓s∈[1,2,3,4] 作為Stage的索引，在每個Stage使用個Block，其中前2個Stage僅使用MLP Block來編碼局部信息，后兩個Stage使用標準Transformer Block來編碼處理更長的依賴關系。在每個Stage，將輸入特征圖縮放為，其中和分別表示第s個Stage的塊大小和隱藏維度。對于最后2個Stage，在每個Transformer Block中設置個Self-Attention heads。

如圖1所示，LIT采用2種類型的模塊:MLP模塊和Transformer模塊。在前期Stage應用MLP Block。具體來說，MLP Block建立在MLP之上，該MLP由2個FC層組成，中間有GELU。對于第s階段的每個MLP，使用的擴展比。具體來說，第1個FC層將token的維度從擴展到，另一個FC層將維度縮減回。形式上，令為第s階段的輸入，l為塊的索引，MLP Block可以表示為：其中LN表示層歸一化。在最后階段，ViT中描述的Transformer Block包含一個MSA層和一個MLP，可以表示為：使用這種架構有2個主要優點：

• 首先，避免了在早期Stage由長序列引入的巨大計算成本和內存占用；
• 其次，與最近使用子窗口縮小注意力圖或減少注意力圖的工作不同在key和value矩陣的空間維度上，在最后2個Stage保留標準MSA層，以保持LIT處理遠程依賴關系的能力。

在這里通過考慮卷積層、FC層和MSA層之間的關系來證明在前期Stage應用純MLP Block的合理性。

建議背誦 | 字節用4大準則教你設計一個擁有CNN的速度，Transformer精度的模型！首先，從回顧標準卷積層開始。令為輸入特征圖，令為卷積權重，其中K為Kernel-size，和分別為輸入和輸出通道維度。為簡單起見，省略了偏置項，并用表示，其中(i,j)表示像素索引，。給定K×K個采樣位置的卷積核，像素p的輸出可以表示為：其中是采樣索引到預先指定的偏移量ΔK的雙射映射。例如，令ΔK={(?1,?1),(?1,0),...,(0,1),(1,1)}為3×3的kernel，膨脹率為1，則g(0)=(-1,-1)表示第1個采樣偏移量。當K=1時，權重張量W等價于一個矩陣，使得$。在這種情況下，Eq(4)可以表示FC層，像素p的輸出定義為：最后，讓是MSA層中的Head數，$是第h個Head的可學習參數。在特定的相對位置編碼方案下，Cordonnier等人證明了像素p處的MSA層的輸出可以表示為：其中是Head到像素位移的雙射映射。在這種情況下，Eq(6)可以看作是Kernel-size為 的卷積層近似。從Eqs(4)-(6)觀察到，雖然具有足夠數量的Head的MSA層能夠逼近任何卷積層，但理論上更少的Head限制了這種逼近的能力。作為極端情況，具有一個Head的MSA層是只能逼近FC層。請注意，MSA層在實踐中肯定不等同于卷積層。然而，d'Ascoli 等人觀察到，早期的MSA層可以在訓練時學會卷積的類似表征。考慮最近的HVT在早期Stage采用很少的Head，這種卷積行為可能會限制在小的感受野內。圖 3在圖3中，在可視化中顯示PVT-S中的早期MSA層確實只關注query像素周圍的一小塊區域，而刪除它們會導致性能輕微下降，但會顯著降低模型復雜性。這證明了在前2個stage應用純MLP Block是合理的。

以前關于 HVT 的工作依賴于Patch Merge來實現金字塔特征表示。然而，從規則網格合并Patch，并忽略了并非每個Patch對輸出單元的貢獻相同的事實。受可變形卷積的啟發，提出了一個Deformable Token Merging模塊來學習偏移網格，以自適應地采樣更多信息塊。形式上，可變形卷積被表述為：與Eq(4)中的標準卷積運算相比，DC為每個預先指定的偏移量g(k)學習一個偏移量Δg(k)。學習Δg(k) 需要一個單獨的卷積層，它也應用于輸入特征圖X。為了以自適應方式合并Patch，在DTM模塊中采用一個DC層，可以表示為：其中 BN 表示批量歸一化，并使用了 GELU。在消融研究中，當物體的尺度和形狀發生變化時，DTM 中的采樣位置會進行自適應調整，從而受益于學習到的偏移量。另請注意，與Baseline中的常規網格采樣相比，輕量級DTM引入了可忽略的FLOP和參數，因此使其成為最近HVT的即插即用模塊。

表 2結果如表 2 所示。總的來說，LIT在使用更少的 FLOP（3.6G 與 3.8G）的同時，將PVT-S的Top-1精度提高了0.6%。對于Swin-Ti將 FLOPs 減少了0.4G，同時實現了同等性能。還值得注意的是，PVT-S 和 Swin-Ti 的參數總數都減少了。整體性能證明了所提出架構的有效性，這也強調了早期MSA在PVT和Swin中的微小優勢。

表 3結果如表 3 所示。對于這2個模型，DTM 引入了可忽略的 FLOP 和參數，同時在 Top-1 精度方面分別將 PVT-S 和 Swin-Ti 提高了 0.7% 和 0.3%。圖 2此外，在圖 2 中可視化了學習到的偏移量。如圖所示，與之前的統一Patch Merge策略不同，之前的統一Patch Merge策略將采樣位置限制在綠色矩形內，DTM 根據對象的比例和形狀（例如，考拉腿、貓尾巴）。這再次強調了LIT適應各種幾何變換的能力。

結果如表4所示。首先，在用標準MSA層替換PVT-S中的SRA層后，觀察到 Top-1 準確度提高了1.1%，而FLOPs幾乎翻了一番。這表明PVT在性能和效率之間進行了權衡。接下來，通過在前2個階段逐步去除MSA層，Top-1準確率僅分別下降了0.1%、0.5%。這意味著 PVT 早期階段的self-attention層對最終性能的貢獻低于預期，并且它們的性能并不比純MLP層好多少。這可以歸因于淺層更多地關注編碼局部信息的事實。然而，可以觀察到在最后2個階段移除self-attention時性能大幅下降。結果表明，self-attention層在后期發揮著重要作用，捕獲長距離依賴對于表現良好的分層視覺Transformer至關重要。圖 3為了更好地理解這一現象，在不移除任何MSA層的情況下可視化PVT-S的注意力概率，如圖3所示。首先，第1階段的注意力圖顯示query像素幾乎不注意其他位置。在第2階段，query像素的感受野略有擴大，但與第1階段相似。考慮到PVT-S在第1階段只有1個head，在第2階段只有2個head，這有力地支持了作者的假設：MSA層中過少的head會導致較小的感受野，此時self-attention幾乎等同于FC層。此外，從最后2個階段的注意力圖中觀察到相對較大的感受野。由于大的感受野通常有助于對更長的依賴關系進行建模，這解釋了在最后2個階段移除MSA層后表4中的巨大性能下降。

[1].Less is More: Pay Less Attention in Vision Transformers

本文僅做學術分享，如有侵權，請聯系刪文。