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

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

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

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

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

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

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

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

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

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

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

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

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

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

• 其次，提出了一個(gè)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，因此每個(gè)Patch的初始特征維度為4×4×3=48。接下來(lái)，利用線性嵌入層將每個(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使用標(biāo)準(zhǔn)Transformer Block來(lái)編碼處理更長(zhǎng)的依賴(lài)關(guān)系。在每個(gè)Stage，將輸入特征圖縮放為，其中和分別表示第s個(gè)Stage的塊大小和隱藏維度。對(duì)于最后2個(gè)Stage，在每個(gè)Transformer Block中設(shè)置個(gè)Self-Attention heads。

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

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

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

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

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

表 2結(jié)果如表 2 所示。總的來(lái)說(shuō)，LIT在使用更少的 FLOP（3.6G 與 3.8G）的同時(shí)，將PVT-S的Top-1精度提高了0.6%。對(duì)于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)勢(shì)。

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

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

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