本文探討了普通視覺Transformer（ViT）用于語義分割的能力，并提出了SegViT。以前基于ViT的分割網絡通常從ViT的輸出中學習像素級表示。不同的是，本文利用基本的組件注意力機制生成語義分割的Mask。

具體來說，作者提出了Attention-to-Mask（ATM）模塊，其中一組可學習 class tokens和空間特征映射之間的相似性映射被轉移到Segmentation masks。

實驗表明，本文提出的使用ATM模塊的SegViT優于在ADE20K數據集上使用普通ViT主干的同類算法，并在COCO-Stuff-10K和PASCAL上下文數據集上實現了最新的性能。

此外，為了降低ViT主干的計算成本，提出了基于查詢的下采樣（QD）和基于查詢的上采樣（QU）來構建收縮結構。使用建議的收縮結構，該模型可以節省多達40%的計算，同時保持競爭力。

1、簡介

語義分割是計算機視覺中的一項密集預測任務，需要對輸入圖像進行像素級分類。全卷積網絡（FCN）在最近最先進的方法中得到了廣泛的應用。這種模式包括作為編碼器/主干的深度卷積神經網絡和提供密集預測的面向分段的解碼器。通常將1×1卷積層應用于代表性特征圖以獲得像素級預測。為了獲得更高的性能，以前的工作側重于豐富上下文信息或融合多尺度信息。然而，由于感受野有限，在FCN中很難明確建模空間位置之間的相關性。

最近，利用空間注意力機制的視覺Transformer（ViT）被引入計算機視覺領域。與典型的基于卷積的主干不同，ViT具有一個簡單的非層次結構，可以始終保持特征圖的分辨率。缺少下采樣過程（不包括圖像token化）給使用ViT主干進行語義分割任務的體系結構帶來了差異?；赩iT主干的各種語義分割方法由于從預訓練的主干中學習到的強大表示，已經取得了良好的性能。然而，注意力機制的潛力尚未得到充分發掘。

與以前的逐像素分類范式不同，本文考慮學習一個有意義的 class tokens，然后找到與之相似性較高的局部patch。為了實現這個目標，作者提出了Attention-to-Mask（ATM）模塊。更具體地說，使用了一個transformer block，它將可學習 class tokens作為query，并將空間特征映射作為key和value進行傳輸。點積運算符計算query和key之間的相似性映射。

作者鼓勵屬于同一類別的區域為相應類別（即特定類別tokens）生成更大的相似性值。圖1顯示了特征與“Table”和“Chair”標記之間的相似性映射。通過簡單地應用Sigmoid操作，可以將相似性映射轉移到mask。同時，在設計了一個典型的transformer block后，還將Softmax操作應用于相似性圖，以獲得交叉注意力圖。然后，“Table”和“Chair”標記像在任何常規transformer 解碼器中一樣，通過以交叉注意力圖作為權重的值的加權和進行更新。由于Mask是定期仔細計算的副產品，因此在操作過程中涉及的計算可以忽略不計。

在這個高效的ATM模塊的基礎上，提出了一種新的具有普通ViT結構的語義分割范式，稱為SegViT。在該范式中，在不同的層上使用多個ATM模塊，通過將不同層的輸出相加得到最終的Segmentation masks。SegViT的性能優于基于ViT的同類算法，計算成本更低。然而，與以前使用分層網絡作為編碼器的編碼器-解碼器結構相比，作為編碼器的ViT主干通常更重。為了進一步降低計算成本，采用了一種收縮結構，由基于查詢的下采樣（QD）和基于查詢的上采樣（QU）組成。QD可插入ViT主干以將分辨率降低一半，QU與主干平行以恢復分辨率。收縮結構與作為解碼器的ATM模塊一起可以減少多達40%的計算量，同時保持具有競爭力的性能。

主要貢獻總結如下：

提出了一個注意力掩碼（ATM）解碼器模塊，該模塊對于語義分割有效且高效。首次利用注意力圖中的空間信息為每個類別生成Mask預測，這可以作為語義分割的一種新范式。

設法以級聯方式將ATM解碼器模塊應用于普通、非分層的ViT主干，并設計了一種結構，即SegViT，它在競爭對手ADE20K數據集上實現了55.2%的mIoU，這是使用ViT主干的方法中最好、最輕的。作者還在PASCAL Context數據集（65.3%mIoU）和COCO-Stuff-10K數據集（50.3%mIoO）上對本文的方法進行了基準測試，并實現了新的SOTA。

進一步探索了ViT主干的架構，并制定了適用于主干的收縮結構，以降低總體計算成本，同時仍保持競爭力。這減輕了ViT主干的缺點，與層次結構的對等物相比，這些主干通常計算量更大。在ADE20K數據集上的收縮版SegViT達到了mIoU 55.1%，計算成本為373.5 GFLOP，與原始SegViT（637.9 GFLOP）相比降低了約40%。

2、本文方法

2.1、編碼器

給定輸入圖像，一個普通視覺Transformer主干將其reshape為一系列tokens ，其中，P是patch size，C是通道數。添加與大小相同的可學習位置嵌入以獲取位置信息。然后，將 token sequence 應用于m個transformer layers以獲得輸出。將每個層的輸出token 定義為。通常，一個transformer layer由一個multi-head self-attention block和一個point-wise multilayer perceptron block組成，layer norm介于其間，然后再添加一個residual connection。transformer layers重復堆疊若干次。對于像ViT這樣的普通視覺transformer ，沒有涉及其他模塊，并且對于每個層，token 的數量不會改變。

2.2、解碼器

（1）Mask-to-Attention (ATM)

Cross Attention可以描述為兩個 token sequence之間的映射。將2個 token sequence定義為，長度N等于類別數和。首先，對它們中的每一個應用線性變換以形成query（Q）、key（K）和value（V），如等式（1）所示。

在query和key之間計算相似度圖。根據scaled dot-product attention機制，相似度圖和注意力圖通過以下公式計算：

其中是縮放因子，等于key的尺寸。相似度圖S(Q,K)的形狀由兩個token sequence N和L的長度確定。然后，注意力機制通過加權和V更新G，其中分配給總和的權重是沿維度L應用Softmax的相似度圖。

Dot-product attention使用Softmax函數專門將注意力集中在最相似的token 上。然而，作者認為，除了產生最大相似性的token 之外，其他token 也有意義?；谶@種直覺，作者設計了一個輕量級模塊，可以更直接地生成語義預測。更具體地說，將G指定為 segmentation 任務的class embeddings，將指定為ViT主干的第層的輸出。將semantic mask與G中的每個token 配對，以表示每個類的語義預測。Mask的計算為：

Mask的形狀為N×L，可進一步reshape為N×H/P×W/P。ATM機制的結構如圖2右側所示。Mask是Cross Attention的中間輸出。ATM模塊的最終輸出token 用于分類。對輸出類token 應用線性變換，然后進行Softmax激活，以獲得類概率預測。

請注意這里添加了一個“無對象”類別，以防圖像不包含某些類。在推理過程中，輸出由類概率和Mask組之間的點積產生。

像ViT這樣的普通骨干沒有具有不同規模特征的多個階段。因此，用于合并具有多個比例的特征的結構（如FPN）不適用。然而，除最后一層之外的特征包含豐富的底層語義信息，并有利于性能。

作者設計了一種結構，它可以利用來自ViT不同層的特征映射，與ATM解碼器SegViT進行壓縮。在這項研究中，還找到了一種在不犧牲性能的情況下壓縮ViT主干網計算成本的方法。這個建議的Shrunk版本的SegViT使用基于查詢的下采樣（QD）模塊和基于查詢的上采樣（QU）模塊來壓縮ViT主干，并整體降低計算成本。

（2）The SegViT structure

如圖2所示，ATM解碼器將N個token作為類嵌入，并將另一個token序列作為基礎來計算key和value，以便ATM模塊生成mask。ATM的輸出是N個更新token和對應于每個類token的N個mask。首先使用隨機初始化的可學習token作為類嵌入，并將ViT主干最后一層的輸出作為基礎。為了利用多層信息，第一個ATM解碼器的輸出被用作下一個ATM解碼的類嵌入，ViT主干網的另一層的輸出作為基礎。這個過程再重復一次，這樣可以得到3組token和mask。每層的損失函數形式上可以表示為，

在每組中，輸出令牌由上述分類損失（Lcls）監督，mask按順序求和，并由mask損失（）監督，該mask損失是focal loss和 dice loss的線性組合，分別乘以超參數和，如DETR中所示。然后將所有三組的損失加在一起。還有進一步的實驗表明，這種設計是有益和有效的。

（3）The Shrunk structure

眾所周知，像ViT這樣的普通transformer主干比具有類似性能的同類transformer具有更大的計算成本。本文作者提出了一種使用基于查詢的下采樣（QD）和上采樣（QU）的收縮結構。由于注意力模塊輸出的形狀由query的形狀決定，可以在query轉換之前應用下采樣來實現QD，或者在交叉注意力期間插入新的query tokens來實現QU，提供更大的靈活性來保存（恢復）重要區域。

更具體地說，在QD層中，使用最近的采樣來減少query tokens的數量，同時保持key和value token的大小。當通過transformer層時，這些值通過query tokens和key tokens之間的注意力度映射進行加權和求和。這是非線性下采樣，將更加關注重要區域。在QU層中，使用了一個transformer decoder結構，并根據期望的輸出分辨率將新的可學習token初始化為query。

如圖3所示，設計了一個單層作為基線的SegViT結構（a）。首先嘗試一種簡單的方法（b），即在主干的1/3深度（例如，具有24層的主干的第8層）應用一次QD，以將層輸出的分辨率從1/16降到1/32，從而降低總體計算成本。由于QD過程涉及信息丟失，性能如預期般下降。

為了補償原始“收縮”版本中的信息損失，進一步應用兩個與主干并行的QU層。這是建議的Shrunk版本（c）。第一個QU層從主干的低層接收1/16分辨率的特征。然后將其輸出用作query ，以與主干最后一層分辨率為1/32的下采樣特征進行交叉注意力。該QU結構的輸出形狀為1/16分辨率。

直接減少query tokens的數量必然會損害最終性能。然而，使用設計的QU層和ATM模塊，Shrunk結構能夠減少40%的總計算成本，同時在性能上仍然具有競爭力。

3、實驗

3.1、消融實驗

（1）ATM module的影響

（2）使用不同的層作為SegViT的輸入

（3）ATM Decoder

（4）SegViT on hierarchical backbones

（5）QD module

（6）Ablation of the components in Shrunk structure

3.2、SOTA對比

（1）ADE20K

（2）COCO-Stuff-10K

（3）PASCAL-Context