ECCV 2022 | k-means Mask Transformer

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

前言目前，大多數現有的基於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的模型。