Shunted Self-Attention | 源于 PvT又高于PvT，解決小目標問題的ViT方法

以下文章來源于集智書童 ，作者ChaucerG

最近的 Vision Transformer (ViT) 模型在各種計算機視覺任務(wù)中都展示了不錯的性能，這要歸功于其通過Self-Attention對圖像塊或Token的遠程依賴關(guān)系進行建模的能力。然而，這些模型通常指定每一層內(nèi)每個Token特征的相似感受野。這種約束不可避免地限制了每個Self-Attention層捕獲多尺度特征的能力，從而導致在處理具有不同尺度的多個對象的圖像時性能下降。

為了解決這個問題，本文提出了一種新穎的通用策略，稱為Shunted Self-Attention(SSA)，它允許 ViT 在每個注意力層的混合尺度上對注意力進行建模。SSA 的關(guān)鍵思想是將異構(gòu)感受野大小注入到Token中：在計算Self-Attention矩陣之前，它選擇性地合并Token以表示更大的對象特征，同時保留某些Token以保留細粒度的特征。這種新穎的合并方案使Self-Attention能夠?qū)W習不同大小目標之間的關(guān)系，同時減少Token數(shù)量和計算成本。

各種任務(wù)的廣泛實驗證明了 SSA 的優(yōu)越性。具體來說，基于 SSA 的 Transformer 實現(xiàn)了 84.0% 的 Top-1 準確率，并在 ImageNet 上以只有一半的模型大小和計算成本超過了最先進的 Focal Transformer，并且在 COCO 上以相似的參數(shù)和計算成本下超過了 Focal Transformer 1.3 mAP，ADE20K 超越了 2.9 mIOU。

最近的Vision Transformer (ViT) 模型在各種計算機視覺任務(wù)（例如分類、目標檢測、語義分割和視頻動作識別）中展示了卓越的性能。與專注于局部建模的卷積神經(jīng)網(wǎng)絡(luò)不同，ViT 將輸入圖像劃分為一系列Patch（Token），并通過全局 Self-Attention 逐步更新Token特征。Self-Attention 可以有效地模擬Token的長期依賴關(guān)系，并通過聚合來自其他Token的信息逐漸擴大其感受野的大小，這在很大程度上是 ViT 成功的原因。

然而，Self-Attention 機制也帶來了昂貴的內(nèi)存消耗成本。因此，最先進的 Transformer 模型具有采用各種下采樣策略來減少特征大小和內(nèi)存消耗。例如，ViT 在第一層進行 16×16 的下采樣投影，并在生成的粗粒度和單尺度特征圖上計算 Self-Attention；因此，所產(chǎn)生的特征信息丟失不可避免地會降低模型性能。其他方法努力計算高分辨率特征的 Self-Attention，并通過將Token與Token的空間縮減合并來降低成本。然而，這些方法往往會在一個Self-Attention中合并太多的Token，從而導致來自小物體和背景噪聲的Token混合在一起。反過來，這種行為也使模型在捕獲小目標方面的效率降低。

此外，先前的 Transformer 模型在很大程度上忽略了注意力層內(nèi)場景目標的多尺度性質(zhì)，使得它們在涉及不同大小目標的野外場景中變得脆難以檢測。從技術(shù)上講，這種低效率歸因于它們潛在的注意力機制：現(xiàn)有方法僅依賴于Token的靜態(tài)感受野和一個注意力層內(nèi)的統(tǒng)一信息粒度，因此無法同時捕獲不同尺度的特征。

為了解決這個限制，作者引入了一種新穎的通用Self-Attention方案，稱為Shunted Self-Attention(SSA)，它明確允許同一層內(nèi)的Self-Attention Head分別解釋粗粒度和細粒度特征。與合并太多Token或捕獲失敗的先前方法不同小物體，SSA 有效地在同一層內(nèi)的不同注意力頭上同時對各種尺度的物體進行建模，使其具有良好的計算效率以及保留細粒度細節(jié)的能力。

圖 2

在圖 2 中展示了標準 Self-Attention（來自 ViT）、下采樣輔助注意力（來自 PVT）和提議的 SSA 之間的定性比較。當對相同大小的特征圖應用不同的注意力時，ViT 可以捕獲到細粒度的小目標，但計算成本非常高（圖2（a））；PVT 降低了計算成本，但它的注意力僅限于到粗粒度的較大目標（圖2（b））。

相比之下，所提出的 SSA 保持了比較小的計算成本，但同時考慮了混合尺度的注意力（圖2（c））。實際上，SSA 不僅可以精確地處理粗粒度的大物體（例如沙發(fā)），還可以精確地處理細粒度的小物體（例如燈和風扇），甚至一些位于角落的物體，不幸的是 PVT 錯過了這些物體。

圖 3

在圖 3 中顯示了注意力圖的視覺比較，以突出 SSA 的學習尺度自適應注意力。

SSA 的多尺度注意力機制是通過將多個注意力頭分成幾組來實現(xiàn)的。每個組占一個專門的注意力粒度。對于細粒度的組，SSA 會聚合很少的Token并保留更多的局部細節(jié)。對于剩余的粗粒度頭部組，SSA 學習聚合大量的Token，從而降低計算成本，同時保留捕獲大目標的能力。多粒度組共同學習多粒度信息，使模型能夠有效地建模多尺度目標。

• 提出了Shunted Self-Attention (SSA)，它通過多尺度Token聚合在一個Self-Attention層內(nèi)統(tǒng)一多尺度特征提取。SSA 自適應地合并大目標上的Token以提高計算效率，并保留小目標的Token。

• 基于 SSA 構(gòu)建了Shunted Transformer，它能夠有效地捕獲多尺度物體，尤其是小型和遠程孤立物體。

• 在包括分類、目標檢測和分割在內(nèi)的各種研究中評估了Shunted Transformer。實驗結(jié)果表明，Shunted Transformer在相似的模型尺寸下始終優(yōu)于以前的Vision Transformer。

圖 4 Shunted Transformer & SSA

本文提出的Shunted Transformer的整體架構(gòu)如圖 4 所示。它建立在Shunted Self-Attention (SSA) 塊之上。SSA 模塊與 ViT 中的傳統(tǒng)Self-Attention模塊有2個主要區(qū)別：

1. SSA 為每個self-attention層引入了分流注意力機制，以捕獲多粒度信息，更好地建模不同大小的目標，尤其是小目標；

2. 它通過增強跨Token交互來增強在逐點前饋層中提取局部信息的能力。

此外，還在的Shunted Transformer部署了一種新的Patch Embedding方法，以便為第一個注意塊獲得更好的輸入特征圖。

在所提出的Shunted Transformer的第 i 個Stage中，有  個Transformer Block。每個Transformer Block都包含一個Self-Attention層和一個前饋層。為了降低處理高分辨率特征圖時的計算成本，PVT 引入了Spatial Reduction Attention（SRA）來代替原來的Multi-Head Self-Attention（MSA）。然而，SRA 傾向于在一個Self-Attention層中合并Token，并且只提供單一尺度的Token特征。這些限制阻礙了模型捕捉多尺度物體的能力，尤其是小尺寸物體。因此，本文引入了Shunted Self-Attention，可以在一個Self-Attention層內(nèi)并行學習多粒度特征。

首先將輸入序列  投影到 query(Q)、key(K) 和 value(V) 張量中。然后Multi-Head Self-Attention采用H個獨立的注意力Head并行計算Self-Attention。為了降低計算成本，遵循 PVT 的思想減少 K 和 V 的長度，而不是像 Swin Transformer 中那樣將 {Q, K, V} 拆分為區(qū)域。圖5如圖 5 所示， SSA 與 PVT 的 SRA 不同之處在于，K、V 的長度在同一Self-Attention層的注意力Head上并不相同。取而代之的是，不同Head的長度不同，以捕獲不同的粒度信息。這給出了Multi-scale Token Aggregation (MTA)。具體來說，對于由 i 索引的不同Head，Key K 和Value V 被下采樣到不同的大小：這里的  是第 i 個Head中的多尺度Token聚合層，下采樣率為 。在實踐中，采用kernel-size和stride為  的卷積層來實現(xiàn)下采樣。、、是第i個Head的線性投影的參數(shù)。在一層中存在變體  注意頭。因此，key 和 value 可以在 Self-Attention 中捕獲不同的尺度。LE(·) 是 MTA 通過深度卷積對Value V 進行的局部增強分量。與Spatial Reduction Attention相比，保留了更多細粒度和低層次的細節(jié)。那么shunted self-attention計算如下：其中  是維度。由于多尺度 key 和value，shunted self-attention在捕獲多尺度目標方面更加強大。而計算成本的降低取決于 r 的值，因此，可以很好地定義模型和 r 來權(quán)衡計算成本和模型性能。當r變大時，K,V中更多的token被合并，K,V的長度更短，因此計算成本低，但仍然保留了捕獲大目標的能力。相反，當 r 變小時，保留更多細節(jié)但帶來更多計算成本。將各種 r 集成到一個Self-Attention層中，使其能夠捕獲多粒度特征。

在傳統(tǒng)的前饋層中，全連接層是逐點的，無法學習到交叉token信息。在這里通過指定前饋層中的細節(jié)來補充局部信息。圖 6如圖 6 所示，通過在前饋層的2個全連接層之間添加數(shù)據(jù)特定層來補充前饋層中的局部細節(jié)：其中 DS(·; θ) 是具有參數(shù) θ 的細節(jié)特定層，在實踐中通過深度卷積實現(xiàn)。

Transformer 最初是為處理順序數(shù)據(jù)而設(shè)計的。如何將圖像映射到序列對于模型的性能很重要。ViT 直接將輸入圖像分割成 16 × 16 的非重疊Patches。最近的一項研究發(fā)現(xiàn)，在Patch Embedding中使用卷積可提供更高質(zhì)量的Token序列，并幫助Transformer比傳統(tǒng)的大步幅非重疊Patch Embedding“看得更清楚”。因此，一些作品像使用 7×7 卷積一樣進行重疊塊Embedding。在Shunted Transformer中，根據(jù)模型大小采用不同的卷積層重疊。將步長為 2 和零填充的 7×7 卷積層作為Patch Embedding的第1層，并根據(jù)模型大小添加額外的stride為 1 的 3×3 卷積層。最后，使用stride為 2 的非重疊投影層生成大小為 H/4×W/4 的輸入序列。

給定尺寸為 H×W×3 的輸入圖像，采用上述Patch Embedding方案來獲得長度為 H/4×W/4 和Token維度 C 的更多信息Token序列。按照之前的設(shè)計，模型中有4個階段，每個階段都包含幾個Shunted Transformer blocks。在每個階段，每個Block輸出相同大小的特征圖。采用步長為 2（線性嵌入）的卷積層來連接不同的階段，特征圖的大小在進入下一個階段之前將減半，但維度將加倍。因此，有每個階段輸出的4個特征圖 、、、， 的大小為 。表 1這里提出了3種不同的模型配置，以便在相似的參數(shù)和計算成本下進行公平比較。如表 1 所示，head 和  表示一個Block中的Head數(shù)和一個Stage中的Block數(shù)。變體僅來自不同階段的層數(shù)。具體來說，每個block的head數(shù)設(shè)置為2,4,8,16。Patch Embedding中的卷積范圍為 1 到 3。

[1].Shunted Self-Attention via Multi-Scale Token Aggregation

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