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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

圖1LIT的整體架構(gòu)如圖1所示。設(shè)為輸入的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使用標(biāo)準(zhǔn)Transformer Block來編碼處理更長的依賴關(guān)系。在每個Stage，將輸入特征圖縮放為，其中和分別表示第s個Stage的塊大小和隱藏維度。對于最后2個Stage，在每個Transformer Block中設(shè)置個Self-Attention heads。

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

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

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

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

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

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

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

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

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

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