ECCV 2022 | k-means Mask Transformer

語言: CN / TW / HK

本站內容均來自興趣收集,如不慎侵害的您的相關權益,請留言告知,我們將盡快刪除.謝謝.

前言目前,大多數現有的基於transformer的視覺模型只是借用了自然語言處理的思想,忽略了語言和圖像之間的關鍵差異,特別是空間扁平像素特徵的巨大序列長度。這阻礙了在像素特徵和對象查詢之間交叉注意的學習。

在本文中,作者重新思考像素和對象查詢之間的關係,並提出將交叉注意學習重新定義為一個聚類過程。受傳統k-means聚類算法的啟發,開發了一種用於分割任務的k-means Mask Xformer (kMaX-DeepLab),它不僅改進了最先進的技術,而且具有簡單優雅的設計。

kMaX-DeepLab在COCO val set和Cityscapes val set 上實現了新的SOTA的性能,無需測試時間增強或外部數據集。

QQ交流羣: 444129970。羣內有大佬負責解答大家的日常學習、科研、代碼問題。

歡迎關注公眾號 CV技術指南 ,專注於計算機視覺的技術總結、最新技術跟蹤、經典論文解讀、CV招聘信息。目前公眾號正在徵稿中,可以獲取對應的稿費哦。

QQ交流羣: 444129970。羣內有大佬負責解答大家的日常學習、科研、代碼問題。

論文: arxiv.org/pdf/2207.04…

代碼: github.com/google-rese…

創新思路

基於transformer的端到端框架已經成功地應用於多種計算機視覺任務,尤其是交叉注意模塊的transformer解碼器。然而,幕後的工作機制仍不清楚。交叉注意起源於自然語言處理社區,最初是為語言問題設計的,如神經機器翻譯,其中輸入序列和輸出序列共享相似的短長度。

當涉及到某些視覺問題時,這種隱含的假設就會產生問題,即在對象查詢和空間扁平像素特徵之間執行交叉注意。具體來説,通常使用少量的對象查詢(如128次查詢),而輸入圖像可以包含數千個像素,用於檢測和分割的視覺任務。

每次對象查詢在交叉注意學習過程中都需要學習在豐富的像素中突出最具區別性的特徵,這導致訓練收斂緩慢,性能較差。

在這項工作中, 作者觀察到交叉注意方案實際上與傳統的k-means聚類具有很強的相似性,將對象查詢視為具有可學習嵌入向量的聚類中心。

因此,提出了新的k-means Mask Xformer (kMaX-DeepLab),它重新考慮像素特徵和對象查詢之間的關係,並從k-means聚類的角度重新設計交叉注意。

方法

基於掩碼 transformer 的分割框架

transformer已經有效地部署到分割任務中。在不損失一般性的情況下,在下面的問題公式中考慮全視域分割,它可以很容易地推廣到其他分割任務。全景分割的目的是將圖像I分割為一組不重疊且有關聯語義標籤的掩碼:

從DETR和MaX-DeepLab開始,全視域分割的方法轉移到一個新的端到端範式,其中預測直接匹配ground truth的格式與N個掩碼及其語義類:

基於N個對象查詢,預測N個掩碼,通過一個transformer解碼器從像素特徵中聚合信息,由自我注意和交叉注意模塊組成。通過多個transformer解碼器更新的對象查詢被用作掩碼嵌入向量,該向量將與像素特徵相乘,得到由N個掩碼組成的最終預測Z:

交叉注意與 k-means 聚類的關係

雖然基於transformer的分割框架成功地以端到端方式連接對象查詢和掩碼預測,但關鍵問題是如何將對象查詢從可學習的嵌入轉換為有意義的掩碼嵌入向量。

Cross-attention

交叉注意模塊用於聚合關聯像素特徵以更新對象查詢。

如式(4)所示,在更新對象查詢時,對圖像分辨率(HW)應用softmax函數,通常在數千像素範圍內進行分割任務。

k – means 聚類

在式(4)中,交叉注意計算對象查詢和像素之間的親和力,沿着圖像分辨率操作轉換為注意圖。然後使用注意圖來檢索關聯的像素特徵,以更新對象查詢。整個過程實際上與經典的k-means聚類算法相似,其工作原理如下:

比較式(4)、式(5)和式(6),k-means聚類算法是無參數的,因此查詢、鍵和值都不需要線性投影。集羣中心上的更新不是以殘留方式進行的。最重要的是,k-means在將親和度轉換為注意圖(即檢索和更新特徵的權重)時,採用了一個集羣級的argmax(即沿着集羣維度操作的argmax),而不是空間級的softmax。

從聚類的角度來看,圖像分割相當於將像素分組到不同的聚類中,其中每個聚類對應一個預測掩碼。然而,交叉注意機制也試圖將像素分組到不同的對象查詢中,相反,它使用了不同於k-means中集羣argmax的空間wise softmax操作。

鑑於k-means的成功,假設在像素聚類方面,聚類argmax比空間softmax更適合,因為聚類argmax執行硬分配,並有效地將操作目標從數千個像素減少到只有幾個聚類中心,加快訓練收斂,從而獲得更好的性能。

k-means Mask Transformer

k – means cross-attention

提出的k-means交叉注意以類似於k-means聚類的方式重新定義交叉注意:

比較Eq.(4)和Eq.(7),空間上的softmax現在被集羣上的argmax所取代。如圖1所示,通過這種簡單而有效的改變,可以將一個典型的transformer解碼器轉換為kMaX解碼器。

圖1:將一個典型的transformer解碼器轉換為kMaX解碼器

與原來的交叉注意不同,本文提出的k-means交叉注意採用了不同的操作(即按簇argmax)來計算注意圖,並且不需要多頭機制。然而,對於聚類而言,argmax是一種難以聚合像素特徵以進行聚類中心更新的賦值方法,它不是一種可微操作,這給訓練帶來了挑戰。

在Eq.(7)中,像素特徵和簇中心之間的親和對數直接對應於分割掩碼的softmax對數,因為簇中心旨在將親和相似的像素聚集在一起,形成預測的分割掩碼。這個公式使每個kMaX解碼器添加深度監督,以便訓練k-means交叉注意模塊中的參數。

Meta architecture

圖2展示了kMaX- deeplab的元架構,它包含三個主要部分:像素編碼器、增強像素解碼器和kMaX解碼器。像素編碼器通過CNN或transformer主幹提取像素特徵,而增強像素解碼器負責恢復特徵圖分辨率,並通過transformer編碼器或軸向注意增強像素特徵。

最後,kMaX解碼器從k-means聚類的角度將對象查詢(即簇中心)轉換為掩碼嵌入向量。

圖2 :k-means Mask Transformer元架構

Model instantiation

基於MaX-DeepLab和官方代碼庫構建kMaX。將整個模型分為兩條路徑:像素路徑和簇路徑,分別負責提取像素特徵和簇中心。圖3詳細描述了kMaX-DeepLab的實例化,其中有兩個示例主幹。

圖3:以ResNet-50和MaX-S為骨幹的kMaX-DeepLab示意圖

Pixel path

Pixel path包括像素編碼器和增強像素解碼器。像素編碼器是imagenet預訓練的骨幹,如ResNet, MaX-S(即具有軸向注意的ResNet-50)和ConvNeXt。增強像素解碼器由幾個axial attention塊和bottleneck 塊組成。

Cluster path

Cluster path共包含6個kMaX解碼器,均勻分佈在不同空間分辨率的特徵圖中。具體來説,分別對輸出步長32、16和8處的像素特徵部署了兩個kMaX解碼器。

Loss functions

訓練損失函數大多遵循MaX-DeepLab。採用相同的PQ-style的損失、輔助語義損失、mask-id交叉熵損失和像素實例識別損失。

實驗

表1:COCO val set results

圖4:kMaX- deeplab像素簇分配在每個kMaX解碼器階段的可視化,以及最終的全光預測。

表3:不同方法對像素-簇交互的消融實驗

表4:kMaX解碼器數量的消融實驗

結論

在本文中,提出了一個新的端到端框架,稱為k-means Mask Transformer (kMaX-DeepLab),用於分割任務。

kMaX-DeepLab從聚類的角度重新思考像素特徵和對象查詢之間的關係。因此,它簡化了掩碼transformer模型,用提出的單頭k-means聚類代替了多頭交叉注意。

通過建立傳統k-means聚類算法和交叉注意之間的聯繫,為分割任務定製了基於transformer的模型。