【寫在前面】

最近的視覺Transformer（ViT）模型在各種計算機視覺任務中取得了令人鼓舞的結果，這得益於其通過自注意力建模圖像塊或token的長期依賴性的能力。然而，這些模型通常指定每個層內每個token特徵的類似感受野。這種約束不可避免地限制了每個自注意力層捕捉多尺度特徵的能力，從而導致處理具有不同尺度的多個對象的圖像的性能下降。為了解決這個問題，作者提出了一種新的通用策略，稱為分流自注意力（SSA），該策略允許VIT在每個注意力層的混合尺度上對注意力進行建模。SSA的關鍵思想是將異質感受野大小注入token：在計算自注意力矩陣之前，它選擇性地合併token以表示較大的對象特徵，同時保留某些token以保留細粒度特徵。這種新的合併方案使自注意力能夠學習不同大小對象之間的關係，同時減少了token數和計算成本。跨各種任務的大量實驗證明了SSA的優越性。具體來説，基於SSA的Transformer達到了84.0%的Top-1精度，優於ImageNet上最先進的Focal Transformer，模型尺寸和計算成本僅為其一半，在相似的參數和計算成本下，在COCO上超過了Focal Transformer 1.3 mAP，在ADE20K上超過了2.9 mIOU。

1. 論文和代碼地址

Shunted Self-Attention via Multi-Scale Token Aggregation

論文地址：http://arxiv.org/abs/2111.15193

代碼地址：http://github.com/oliverrensu/shunted-transformer

2. Motivation

最近的視覺Transformer（ViT）模型在各種計算機視覺任務中表現出了卓越的性能。與專注於局部建模的卷積神經網絡不同，ViTs將輸入圖像劃分為一系列patch，並通過全局自注意力逐步更新token特徵。自注意力可以有效地模擬token的長期依賴性，並通過聚合來自其他token的信息來逐步擴大其感受野的大小，這在很大程度上解釋了VIT的成功。

然而，自注意力機制也帶來了昂貴的內存消耗成本，即輸入token數量的平方比。因此，最先進的Transformer模型採用了各種降採樣策略來減少特徵大小和內存消耗。一些方法努力計算高分辨率特徵的自注意力，並通過將token與token的空間縮減合併來降低成本。然而，這些方法傾向於在一個自注意力層中合併過多的token，從而導致來自小對象和背景噪聲的token的混合。這種行為反過來會降低模型捕獲小對象的效率。

此外，以前的Transformer模型在很大程度上忽略了注意力層內場景對象的多尺度性質，使得它們在涉及不同大小對象的野生場景中變得脆弱。從技術上講，這種無能歸因於其潛在的注意機制：現有方法僅依賴token的靜態感受野和一個注意層內的統一信息粒度，因此無法同時捕獲不同尺度的特徵。

為了解決這一侷限性，作者引入了一種新的通用自注意力方案，稱為分流自注意力（SSA），該方案明確允許同一層內的自注意力頭分別考慮粗粒度和細粒度特徵。與以前合併過多token或捕捉小對象失敗的方法不同，SSA有效地在同一層的不同注意頭上同時對不同規模的對象建模，使其具有良好的計算效率和保留細粒度細節的能力。

作者在上圖中展示了自注意力（來自ViT）、下采樣輔助注意力（來自PVT）和SSA之間的定性比較。當對相同大小的特徵映射應用不同的注意時，ViT捕捉細粒度的小對象，但具有計算成本極高（上圖（a））；PVT降低了計算成本，但其注意力僅限於粗粒度較大的對象（上圖（b））。相比之下，提出的SSA保持了較輕的計算負載，但同時考慮了混合尺度注意（上圖（c））。有效地，SSA不僅精確地關注粗粒度的大對象（如沙發），而且還關注細粒度的小對象（如燈光和風扇），甚至一些位於角落的對象，這些對象不幸被PVT忽略。作者還在下圖中展示了注意力圖的視覺比較，以突出SSA的學習尺度自適應注意力。

SSA的多尺度注意力機制是通過將多個注意力頭分成幾個組來實現的。每個組都有一個專用的注意力粒度。對於細粒度組，SSA學習聚合少量token並保留更多局部細節。對於剩餘的粗粒度頭組，SSA學習聚合大量token，從而降低計算成本，同時保持捕獲大型對象的能力。多粒度組共同學習多粒度信息，使模型能夠有效地建模多尺度對象。

如上圖所示，作者展示了從堆疊多個基於SSA的塊獲得的分流Transformer模型的性能。在ImageNet上，本文的分流Transformer優於最先進的聚焦Transformer，同時將模型尺寸減半。當縮小到微小尺寸時，分流Transformer實現了與DeiT Small相似的性能，但只有50%的參數。對於對象檢測、實例分割和語義分割，在模型大小相似的COCO和ADE20K上，分流Transformer始終優於聚焦Transformer。

本文的貢獻如下：

作者提出了分流自注意力（SSA），它通過多尺度token聚合在一個自注意力層內統一多尺度特徵提取。本文的SSA自適應地合併大對象上的token以提高計算效率，並保留小對象的token。

基於SSA，作者構建了分流Transformer，能夠有效捕獲多尺度對象，尤其是小型和遠程孤立對象。

作者評估了提出的分流Transformer的各種研究，包括分類，目標檢測和分割。實驗結果表明，在相似的模型尺寸下，本文的分流Transformer始終優於以前的視覺Transformer。

3. 方法

本文提出的分流Transformer的整體架構如上圖所示。它建立在新型分流自注意力（SSA）塊的基礎上。本文的SSA塊與ViT中的傳統自注意力塊有兩個主要區別：1）SSA為每個自注意力層引入了分流注意力機制，以捕獲多粒度信息和更好地建模不同大小的對象，尤其是小對象；2） 它通過增強跨token交互，增強了在逐點前饋層提取局部信息的能力。此外，本文的分流Transformer部署了一種新的patch嵌入方法，用於為第一個注意力塊獲得更好的輸入特徵圖。在下文中，作者將逐一闡述這些新穎之處。

3.1. Shunted Transformer Block

在提出的分流Transformer的第i個階段，有$L_i$個Transformer塊。每個transformer塊包含一個自注意力層和一個前饋層。為了減少處理高分辨率特徵圖時的計算成本，PVT引入了空間歸約注意力（spatial-reduction attention，SRA）來取代原來的多頭自注意力（multi-head self attention，MSA）。然而，SRA傾向於在一個自注意力層中聚合太多token，並且僅在單個尺度上提供token特性。這些侷限性阻礙了模型捕捉多尺度對象尤其是小尺寸對象的能力。因此，作者在一個自注意力層內並行引入了具有學習多粒度的分流自注意力。

3.1.1 Shunted Self-Attention

輸入序列$F \in \mathbb{R}^{h \times w \times c}$首先投影到查詢（Q）、鍵（K）和值（V）張量中。然後，多頭自注意力採用H個獨立注意頭並行計算自注意力。為了降低計算成本，作者遵循PVT並減少K和V的長度，而不是像在Swin Transformer中那樣將{Q，K，V}分割為多個區域。

如上圖所示，本文的SSA不同於PVT的SRA，因為K，V的長度在同一個自注意力層的注意頭之間不相同。相反，長度在不同的頭部不同，用於捕獲不同粒度的信息。這提供了多尺度token聚合（MTA）。具體而言，對於由i索引的不同頭部，將鍵K和值V下采樣到不同大小：

$$ \begin{aligned} Q_{i} &=X W_{i}^{Q} \ K_{i}, V_{i} &=M T A\left(X, r_{i}\right) W_{i}^{K}, M T A\left(X, r_{i}\right) W_{i}^{V} \ V_{i} &=V_{i}+L E\left(V_{i}\right) \end{aligned} $$

這裏，$\operatorname{MAT}\left(; ; r_{i}\right)$是第i個頭中的多尺度token聚合層，下采樣率為$r_i$。在實踐中，作者使用卷積核和步長為$r_i$的卷積核進行卷積層。$W_{i}^{Q}, W_{i}^{K}, W_{i}^{V}$是第i個頭中線性投影的參數。在注意力頭的一層中有變體$r_{i}$。因此，鍵和值可以在自注意力中捕捉不同的尺度。$L E(\cdot)$是MTA的局部增強分量，用於深度卷積的V值。與空間縮減相比，保留了更多細粒度和低層次的細節。

然後通過以下公式計算分流的自注意力：

$$ h_{i}=\operatorname{Softmax}\left(\frac{Q_{i} K_{i}^{\top}}{\sqrt{d_{h}}}\right) V_{i} $$

其中$d_{h}$是尺寸。多虧了多尺度鍵和值，本文分流的自注意力在捕捉多尺度對象時更強大。計算成本的降低可能取決於r的值，因此，可以很好地定義模型和r，以權衡計算成本和模型性能。當r變大時，K，V中合併了更多token，並且K，V的長度較短，因此，計算成本較低，但仍保持捕獲大型對象的能力。相反，當r變小時，保留了更多細節，但帶來了更多計算成本。在一個自注意力層中集成各種r使其能夠捕獲多粒度特徵。

3.1.2 Detail-specific Feedforward Layers

在傳統的前饋層中，全連接層是逐點的，無法學習交叉token信息。在這裏，作者的目標是通過指定前饋層中的細節來補充局部信息。如上圖所示，作者通過在前饋層的兩個完全連接的層之間添加數據特定層來補充前饋層中的局部細節：

$$ \begin{aligned} x^{\prime} &=F C\left(x ; \theta_{1}\right) \ x^{\prime \prime} &=F C\left(\sigma\left(x^{\prime}+D S\left(x^{\prime} ; \theta\right)\right) ; \theta_{2}\right) \end{aligned} $$

其中$D S(\cdot ; \theta)$是具有參數θ的細節特定層，在實踐中通過深度卷積實現。

3.2. Patch Embedding

Transformer首先設計用於處理順序數據。如何將圖像映射到序列對於模型的性能很重要。ViT直接將輸入圖像分割為16×16非重疊patch。最近的一項研究發現，在patch嵌入中使用卷積可以提供更高質量的token序列，並有助於transformer比傳統的大步非重疊patch嵌入“看得更好”。因此，一些文獻使用7×7卷積進行重疊patch嵌入。

在本文的模型中，作者根據模型大小採用不同的重疊卷積層。作者將步長為2且零填充的7×7卷積層作為patch嵌入的第一層，並根據模型大小增加步長為1的額外3×3卷積層。最後，使用步長為2的非重疊投影層生成大小為$\frac{H}{4} \times \frac{W}{4}$的輸入序列。

3.3. Architecture Details and Variants

給定一個大小為H×W×3的輸入圖像，作者採用上述patch嵌入方案獲得長度為$\frac{H}{4} \times \frac{W}{4}$、token維數為C的信息量更大的token序列。根據之前的設計，本文的模型中有四個階段，每個階段包含幾個分流Transformer塊。在每個階段，每個塊輸出相同大小的特徵圖。作者採用帶步長2（線性嵌入）的卷積層來連接不同的階段，在進入下一階段之前，特徵圖的大小將減半，但維數將加倍。因此，有每個階段輸出的特徵映射$F_{1}, F_{2}, F_{3}, F_{4}$，$F_{i}$的大小為$\frac{H}{2^{i+1}} \times \frac{W}{2^{i+1}} \times\left(C \times 2^{i-1}\right)$。

作者提出了本文模型的三種不同配置，以便在類似參數和條件下進行公平比較。如上表所示，head和$N_{i}$表示一個塊中的頭數和一個階段中的塊數。變體僅來自不同階段的層數。具體來説，每個塊中的頭數設置為2,4,8,16。patch嵌入中的卷積範圍為1到3。

4.實驗

上表展示了本文方法在ImageNet-1K上的實驗結果，，可以看出，本文的方法在各個模型大小上都能達到SOTA結果。

上表展示了用Mask R-CNN進行目標檢測和語義分割時，各個模型的性能，可以看出，本文的方法在性能上具有明顯優勢。

上表展示了用RetinaNet進行目標檢測和語義分割時，各個模型的性能。

結果如上表所示。本文的分流Transfomer在所有框架中都具有較高的性能和較少的參數，優於以前最先進的Transfomer。

作者還以SegFormer為框架，在Segformers中比較了本文的主幹與MiT的主幹。結果見上表。

上表展示了不同patch embedding的實驗結果。

作者提出了一種新的token聚合函數，用於合併多尺度對象的令牌，同時保留全局和局部信息。從上表中可以看出，本文的新token聚合函數具有與卷積空間歸約相似的計算，但獲得了更多改進。

上表展示了本文提出的FFN的性能改進結果。

5. 總結

在本文中，作者提出了一種新的分流自注意力（SSA）方案來明確説明多尺度特徵。與以往只關注一個注意力層中靜態特徵圖的工作不同，作者維護了多尺度特徵圖，這些特徵圖關注一個自注意力層中的多尺度對象。大量實驗表明，本文的模型作為各種下游任務的主幹是有效的。具體來説，該模型優於先前的Transformers，並在分類、檢測和分割任務上實現了最先進的結果。

---

已建立深度學習公眾號——FightingCV，歡迎大家關注！！！

ICCV、CVPR、NeurIPS、ICML論文解析彙總：http://github.com/xmu-xiaoma666/FightingCV-Paper-Reading

面向小白的Attention、重參數、MLP、卷積核心代碼學習：http://github.com/xmu-xiaoma666/External-Attention-pytorch

加入交流羣，請添加小助手wx：FightngCV666

本文由mdnice多平台發佈