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

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

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

所提出的 LIT 在圖像識別任務(wù)（包括圖像分類(lèi)、對象檢測和實(shí)例分割）上取得了可觀(guān)的性能，可作為許多視覺(jué)任務(wù)的強大支柱。

開(kāi)源地址：https://github.com/zhuang-group/LIT

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

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

在本文中提出了一種Less attention Vision Transformer（LIT）來(lái)解決上述HVT問(wèn)題。具體來(lái)說(shuō)，建議在早期Stage使用MLP層來(lái)捕獲局部信息，同時(shí)引入具有足夠數量的Head來(lái)處理的MSA層后期的長(cháng)期依賴(lài)關(guān)系。

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

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

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

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

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

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

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

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

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

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

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

在這里通過(guò)考慮卷積層、FC層和MSA層之間的關(guān)系來(lái)證明在前期Stage應用純MLP Block的合理性。

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

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

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

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

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

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

本文僅做學(xué)術(shù)分享，如有侵權，請聯(lián)系刪文。