Shunted Self-Attention | 源于 PvT又高于PvT，解決小目標(biāo)問(wèn)題的ViT方法

以下文章來(lái)源于集智書童 ，作者ChaucerG

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

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

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

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

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

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

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

圖 2

在圖 2 中展示了標(biāo)準(zhǔn) Self-Attention（來(lái)自 ViT）、下采樣輔助注意力（來(lái)自 PVT）和提議的 SSA 之間的定性比較。當(dāng)對(duì)相同大小的特征圖應(yīng)用不同的注意力時(shí)，ViT 可以捕獲到細(xì)粒度的小目標(biāo)，但計(jì)算成本非常高（圖2（a））；PVT 降低了計(jì)算成本，但它的注意力僅限于到粗粒度的較大目標(biāo)（圖2（b））。

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

圖 3

在圖 3 中顯示了注意力圖的視覺比較，以突出 SSA 的學(xué)習(xí)尺度自適應(yīng)注意力。

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

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

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

• 在包括分類、目標(biāo)檢測(cè)和分割在內(nèi)的各種研究中評(píng)估了Shunted Transformer。實(shí)驗(yàn)結(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個(gè)主要區(qū)別：

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

2. 它通過(guò)增強(qiáng)跨Token交互來(lái)增強(qiáng)在逐點(diǎn)前饋層中提取局部信息的能力。

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

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

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

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

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

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

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

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