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

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

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

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

1、簡介

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

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

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

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

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

主要貢獻總結(jié)如下：

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

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

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

2、本文方法

2.1、編碼器

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

2.2、解碼器

（1）Mask-to-Attention (ATM)

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

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

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

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

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

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

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

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

（2）The SegViT structure

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

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

（3）The Shrunk structure

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

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

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

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

直接減少query tokens的數(shù)量必然會損害最終性能。然而，使用設(shè)計的QU層和ATM模塊，Shrunk結(jié)構(gòu)能夠減少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