【導(dǎo)讀】近日，谷歌團(tuán)隊(duì)推出了一項(xiàng)新Transformer，可用于優(yōu)化全景分割方案，還登上了CVPR 2022。

最近，谷歌AI團(tuán)隊(duì)受Transformer和DETR的啟發(fā)提出了一種使用Mask Transformer進(jìn)行全景分割的端到端解決方案。

全稱是end-to-end solution for panoptic segmentation with mask transformers，主要用于生成分割MaskTransformer架構(gòu)的擴(kuò)展。

該解決方案采用像素路徑（由卷積神經(jīng)網(wǎng)絡(luò)或視覺Transformer組成）提取像素特征，內(nèi)存路徑（由Transformer解碼器模塊組成）提取內(nèi)存特征，以及雙路徑Transformer用于像素特征和內(nèi)存之間的交互特征。

然而，利用交叉注意力的雙路徑Transformer最初是為語言任務(wù)設(shè)計(jì)的，它的輸入序列由幾百個(gè)單詞構(gòu)成。

而對(duì)視覺任務(wù)尤其是分割問題來說，其輸入序列由數(shù)萬個(gè)像素組成，這不僅表明輸入規(guī)模的幅度要大得多，而且與語言單詞相比也代表了較低級(jí)別的嵌入。

全景分割是一個(gè)計(jì)算機(jī)視覺問題，它是現(xiàn)在許多應(yīng)用程序的核心任務(wù)。

它分為語義分割和實(shí)例分割兩部分。

語義分割就比如為圖像中的每個(gè)像素分配語義標(biāo)簽，例如「人」和「天空」。

而實(shí)例分割僅識(shí)別和分割圖中的可數(shù)對(duì)象，如「行人」和「汽車」，并進(jìn)一步將其劃分為幾個(gè)子任務(wù)。

每個(gè)子任務(wù)單獨(dú)處理，并應(yīng)用額外的模塊來合并每個(gè)子任務(wù)階段的結(jié)果。

這個(gè)過程不僅復(fù)雜，而且在處理子任務(wù)和整合不同子任務(wù)結(jié)果時(shí)還會(huì)引入許多人工設(shè)計(jì)的先驗(yàn)。

在 CVPR 2022 上發(fā)表的「CMT-DeepLab: Clustering Mask Transformers for Panoptic Segmentation」中，文章提出從聚類的角度重新解讀并且重新設(shè)計(jì)交叉注意力cross attention（也就是將相同語義標(biāo)簽的像素分在同一組），從而更好地適應(yīng)視覺任務(wù)。

CMT-DeepLab 建立在先前最先進(jìn)的方法 MaX-DeepLab 之上，并采用像素聚類方法來執(zhí)行交叉注意，從而產(chǎn)生更密集和合理的注意圖。

kMaX-DeepLab 進(jìn)一步重新設(shè)計(jì)了交叉注意力，使其更像一個(gè) k-means 聚類算法，對(duì)激活函數(shù)進(jìn)行了簡(jiǎn)單的更改。

結(jié)構(gòu)總覽

研究人員將從聚類的角度進(jìn)行重新解釋，而不是直接將交叉注意力應(yīng)用于視覺任務(wù)而不進(jìn)行修改。

具體來說，他們注意到Mask Transformer 對(duì)象查詢可以被認(rèn)為是集群中心（旨在對(duì)具有相同語義標(biāo)簽的像素進(jìn)行分組）。

交叉注意力的過程類似于 k-means 聚類算法，（1）將像素分配給聚類中心的迭代過程，其中可以將多個(gè)像素分配給單個(gè)聚類中心，而某些聚類中心可能沒有分配的像素，以及（2）通過平均分配給同一聚類中心的像素來更新聚類中心，如果沒有分配像素，則不會(huì)更新聚類中心）。

在CMT-DeepLab和kMaX-DeepLab中，我們從聚類的角度重新制定了交叉注意力，其中包括迭代聚類分配和聚類更新步驟

鑒于 k-means聚類算法的流行，在CMT-DeepLab中，他們重新設(shè)計(jì)了交叉注意力，以便空間方面的softmax操作（即沿圖像空間分辨率應(yīng)用的 softmax 操作），實(shí)際上將聚類中心分配給相反，像素是沿集群中心應(yīng)用的。

在 kMaX-DeepLab 中，我們進(jìn)一步將空間方式的 softmax 簡(jiǎn)化為集群方式的 argmax（即沿集群中心應(yīng)用 argmax 操作）。

他們注意到 argmax 操作與 k-means 聚類算法中使用的硬分配（即一個(gè)像素僅分配給一個(gè)簇）相同。

從聚類的角度重新構(gòu)建MaskTransformer的交叉注意力，顯著提高了分割性能，并簡(jiǎn)化了復(fù)雜的Masktransformer管道，使其更具可解釋性。

首先，使用編碼器-解碼器結(jié)構(gòu)從輸入圖像中提取像素特征。然后，使用一組聚類中心對(duì)像素進(jìn)行分組，這些像素會(huì)根據(jù)聚類分配進(jìn)一步更新。最后，迭代執(zhí)行聚類分配和更新步驟，而最后一個(gè)分配可直接用作分割預(yù)測(cè)。

為了將典型的MaskTransformer解碼器（由交叉注意力、多頭自注意力和前饋網(wǎng)絡(luò)組成）轉(zhuǎn)換為上文提出的k-means交叉注意力，只需將空間方式的softmax替換為集群方式最大參數(shù)。

本次提出的 kMaX-DeepLab 的元架構(gòu)由三個(gè)組件組成：像素編碼器、增強(qiáng)像素解碼器和 kMaX 解碼器。

像素編碼器是任何網(wǎng)絡(luò)主干，用于提取圖像特征。

增強(qiáng)的像素解碼器包括用于增強(qiáng)像素特征的Transformer編碼器，以及用于生成更高分辨率特征的上采樣層。

一系列 kMaX 解碼器將集群中心轉(zhuǎn)換為 (1) Mask嵌入向量，其與像素特征相乘以生成預(yù)測(cè)Mask，以及 (2) 每個(gè)Mask的類預(yù)測(cè)。

kMaX-DeepLab 的元架構(gòu)

研究結(jié)果

最后，研究小組在兩個(gè)最具挑戰(zhàn)性的全景分割數(shù)據(jù)集 COCO 和 Cityscapes 上使用全景質(zhì)量 (PQ) 度量來評(píng)估 CMT-DeepLab 和 kMaX-DeepLab，并對(duì)比 MaX-DeepLab 和其他最先進(jìn)的方法。

其中CMT-DeepLab 實(shí)現(xiàn)了顯著的性能提升，而 kMaX-DeepLab 不僅簡(jiǎn)化了修改，還進(jìn)一步提升了，COCO val set 上的 PQ 為 58.0%，PQ 為 68.4%，44.0% Mask平均精度（Mask AP），Cityscapes 驗(yàn)證集上的 83.5% 平均交集比聯(lián)合（mIoU），沒有測(cè)試時(shí)間增強(qiáng)或使用外部數(shù)據(jù)集。

從聚類的角度設(shè)計(jì)，kMaX-DeepLab 不僅具有更高的性能，而且還可以更合理地可視化注意力圖以了解其工作機(jī)制。

在下面的示例中，kMaX-DeepLab 迭代地執(zhí)行聚類分配和更新，從而逐漸提高M(jìn)ask質(zhì)量。

kMaX-DeepLab 的注意力圖可以直接可視化為全景分割，讓模型工作機(jī)制更合理

結(jié)論

本次研究展示了一種更好地設(shè)計(jì)視覺任務(wù)中的MaskTransformer的方法。

通過簡(jiǎn)單的修改，CMT-DeepLab 和 kMaX-DeepLab 重新構(gòu)建了交叉注意力，使其更像一種聚類算法。

因此，所提出的模型在COCO 和 Cityscapes數(shù)據(jù)集上實(shí)現(xiàn)了最先進(jìn)的性能。

研究團(tuán)隊(duì)表示，他們希望 DeepLab2 庫中 kMaX-DeepLab 的開源版本有助于未來對(duì)專用于視覺Transformer架構(gòu)設(shè)計(jì)的研究。

審核編輯 ：李倩