PaperWeekly - CVPR 2022 | 從自注意力中學習語義Affinity，用於端到端弱監督語義分割－鑽石舞台

©PaperWeekly 原創 ·作者 | 邵琪鈞研究方向 | 計算機視覺、深度學習

論文標題：

Learning Affinity from Attention: End-to-End Weakly-Supervised SemanticSegmentation with Transformers

論文地址：

https://arxiv.org/abs/2203.02664

代碼地址：

https://github.com/rulixiang/afa

推薦一個自己的深度學習項目：主要用於論文復現，包含圖像識別、目標檢測、圖像分割、自監督等領域。每個項目可以獨立運行，並且有對應論文解讀。歡迎 star：

https://github.com/KKKSQJ/DeepLearning

前言

本文是一篇基於弱監督，並且使用Transformer 來進行端到端的語義分割算法。目前基於 image-level 的弱監督語義分割基礎都是建立一個良好的類激活圖（CAM），而後利用像素間的關係進行建模（pixel-level 的 affinity），從而在 CAM 基礎上生成偽標籤，再對偽標籤進行細化，作為語義分割模型的監督信息。目前的方法主要基於卷積神經網絡（CNN），由於 CNN 無法很好的挖掘全局信息，所以，往往導致對象區域不完整。如下圖（a）所示。

為了解決這個問題，本文引入自然集成全局信息的 Transformer，為端到端的弱監督語義分割（WSSS）生成完整的初始偽標籤。基於 Transformer 的自注意力與語義 affinity 之間的內在一致性，本文提出一個從注意力中學習 affinity 的模塊此外（affinity from Attention，AFA）。為了有效的獲得可靠的 affinity 標籤來監督 AFA 並保證偽標籤的局部一致性，作者設計了一個像素自適應細化模塊（PAR），該模塊融合低級圖像外觀信息來細化標籤。實驗證明，該方法在 PASCAL VOC 2012 和 MS COCO 2014 數據集上分別實現了 66.0% 和 38.9% 的 mIoU，顯著優於近期的端到端方法和幾個多階段算法。

相關介紹

弱監督語義分割（WSSS）：目前的基於 image-level 的弱監督語義分割基礎都是建立於一個良好的類激活圖（CAM），而後利用像素間的關係進行建模（pixel-level 的 affinity），從而在一個 CAM 的基礎上得到一些良好的結果，進而為語義分割生成偽標籤。

先前的一些方法：

1. 雙階段方法：第一階段訓練一個分類網絡，利用 CAM 來產生初始偽標籤。偽標籤經過細化之後，作為第二階段的監督信息，去訓練一個獨立的分割模型。（該方法流程複雜化，並且效率低下）

2. 端到端方法：目前大多數端到端的方法基於 CNN，利用 CNN 來產生 CAM，進而生成偽標籤。然而 CNN 無法很好的獲取全局信息，從而導致無法獲取完整的對象區域，從而影響生成的偽標籤的質量。

視覺中的 Transformer：Vision Transformer 在論文 VIt 中被提出，將 Transformer 應用到圖像中，並取得了顯著的成功。由於 Transformer 的自注意力機制，它能很好的獲取圖像的全局信息。

動機&貢獻

3.1 動機

處理弱監督語義分割大致分為兩類方法：雙階段方法和端到端方法。其中，雙階段方法流程複雜，效率低下。有必要提出一個端到端的方法。

端到端方法大都基於 CNN，但是 CNN 無法很好的獲取全局信息，從而導致無法獲取完整的目標信息（如下圖（a）所示），從而導致生成的偽標籤效果不佳。

Transforerm 的固有優點是可以捕獲全局信息，但是，Transformer 自然（不加約束）產生的 CAM 是粗糙的（如上圖（c）所示），需要對其進一步細化。

3.2 貢獻

本文提出一個端到端的基於 Transformer 的弱監督語義分割框架。據我們所知，這是第一個在 WSSS 領域探索 transformer 的工作。

作者利用 Transformer 的固有優點，設計了affinity from attention 模塊（AFA），它旨在從 MHSA 中學習可靠的語義 affinity，並利用學習到的 affinity 去修正初始偽標籤。

本文提出一個高效的像素自適應細化模塊（PAR），該模塊融合了局部像素的 RGB 和位置信息，用於標籤細化。

方法

4.1 Transformer Backbone

如圖 3 所示，該框架使用 Transformer 作為 backbone。輸入一張圖片，首先劃分為個 patchs，然後經過線性層得到相同數量的 patch tokens。在每一個 Transformer block 中，多頭注意力（multi-head self-attention (MHSA) ）將會捕獲全局特徵。具體來說，對於第個 head，patch tokens 經過 MLP 被映射成 queries ，keys ，values ，其中表示特徵維度。在的基礎之上，自注意力矩陣和輸出的計算如公式 1 所示：

最終每個 Transformer block 的輸出是將送入前向傳播網絡（feed-forward layers (FFN)）得到。

例如：。其中表示 concatenation，由 Layer Norm 和 MLP 組成。

通過堆疊多個 Transformer block，backbone 產生 feature maps。

4.2 CAM Genaration

使用類激活圖（class activation maps (CAM) ）作為初始偽標籤。

對於從 backbone 提取的 feature maps ，對於一個類別，它的類激活圖將會根據它的第 c 個類別權重（由分類層產生類別權重）和 feature maps中第 c 個通道的特徵相乘生成。具體如公式 2 所示：

其中，RELU 函數用於去除負激活。隨後，經過 Min-Max 歸一化到 [0,1] 區間。用 backbone score 來區域前景和背景區域。

4.3 Affinity from Attention

如上圖所示，Transformer 中的 MHSA 和語義 affinity 之間存在一致性。這啟示我們使用 MHSA 去發現目標區域。但是，在訓練過程中，不受額外約束的自注意矩陣學習到的 affinity 是粗糙並且不精確的。這意味着直接應用 MHSA 作為 affinity 去細化偽標籤是不行的，如圖 2（c）所示。

因此，提出 Affinity from Attention module（AFA）去解決這個問題。

假設 Transformer block 中的 MHSA 用表示。其中，是 flattened spatial size。是 attention heads 的數量。在 AFA 模塊中，直接使用 MLP 來編碼 MHSA 以此來產生語義 affinity。

值得注意的是：自注意力機制是一個定向圖模型，而 affinity 矩陣應該是對稱，因為共享相同語義的節點應該是平等的。

為了執行上述變換，作者將 S 和 S 的轉置相加。如公式 3 所示：

表示預測的語義 affinity。

Pseudo Affinity Label Generation：為了學到良好的語義 affinity ，一個關鍵的步驟是獲取一個可信賴的偽 affinity 標籤作為監督信息。

如上圖所示，Transformer 提取的特徵通過 CAM 生成初始偽標籤，初始偽標籤經過 PAR 模塊進行細化得到細化後的偽標籤（refined pseudo labels）。（PAR 模塊再下一節進行介紹）

作者使用兩個 background scores 去過濾 refined pseudo label，得到可靠的前景，背景以及不確定區域。

給定 CAM ，偽標籤的生成如公式 4 所示：

其中，0 和 255 代表背景和 ignore 區域的索引。表示提取的類別索引。

隨後，pseudo affinity label 由生成。具體來說，對於，如果像素和像素共享相同語義，則設置它們的 affinity 為 positive。反之，設置 affinity 為 negative。如果像素來自於 ignore 區域，則它們的 affinity 也被忽略。除此之外，只考慮同一個局部窗口下的像素和像素，並且忽視遠處像素對的 affinity。

Affinity Loss：生成的偽 affinity 標籤去監督網絡預測的 affinity 。affinity loss 如公式 5 所示：

其中，分別表示在進行採樣的 positive 和 negative 集合。分別表示的數量。

該損失函數強迫網絡從 MHSA 中學習高置信度的語義 affinity 依賴，並且有益於自注意去發現完整的目標區域。

Propagation with Affinity：學習到的語義 affinity 用於修正初始的 CAM，作者利用 random walk 完成這一步驟。具體來說，對於學習到的語義 affinity 矩陣，其對應的語義變換矩陣 T 如公式 6 所示：

其中是一個超參數，用於忽略中平凡的 affinity 值。是一個對角矩陣，用於將歸一化。

對於初始 CAM ，rankdom walk 算法如公式 7 所示：

其中，表示將矢量化。這種傳播過程，擴散了具有高 affinity 的語義區域，抑制了錯誤激活的區域，從而使激活圖更好地與語義邊界對齊。

4.4 Pixel-Adaptive Refinement

如圖 3 所示，偽 affinity 標籤由初始偽標籤生成。但是初始偽標籤是粗糙的並且局部不一致。比如：具有相似圖像外觀的相鄰像素不會共享同一個語義。

為了確保局部的一致性，一些方法使用 dense CRF 去細化初始偽標籤，但是 dense CRF 不是一個端到端的好選擇，因為它的訓練效率低下。

本文採用像素自適應卷積（ pixel-adaptive convolution）來提取局部 RGB 信息進行細化，即集合 RGB 和空間信息定義低級的 pairwise affinity 並且構建像素自適應模塊（ Pixel-Adaptive Refinement module (PAR).）

具體來說，給定一張圖片，對於像素位置 (i,j),(k,l)，其 RGB 和空間對的 kernel 如公式 8 所定義：

其中表示 RGB 信息。表示像素 (i,j) 的空間位置。表示 RGB 的標準差。表示位置的標準差。

然後將經過 softmax 進行歸一化，再進行加權，得到作為 PAR 模塊的 affinity kernel。如公式 9 所示。

通過這個 affinity kernel，我們同時細化初始的 CAM 以及傳播的 CAM。這個細化操作在多個迭代過程被構建。對於 CAM ，在第次迭代，有：

與 CRF 中計算全局像素的 low-level affinity 不同，作者根據像素的 8 個最近鄰像素進行計算。該過程可以使用的像素自適應卷積實現，從而能夠高效的插入到端到端的訓練框架中。同時為了擴大感受野，使用多個空洞卷積進行多個最近鄰像素信息的提取。得到像素之間的 low-level affinity 之後，通過多次迭代進行偽標籤的修正。

4.5 Network Training

如圖 3 所示，訓練損失函數包括分類損失、分割損失、正則化損失和前文介紹的 affinity 損失。其中分類損失採用常見的多分類軟間隔損失，分割損失採用常見的交叉熵損失：

實驗

數據集：

PASCAL VOC2012

MS COCO2014

實驗細節：

Transformer backbone：Mix Transformer(MiT)（Segformer）

優化器：AdamW

學習率：

學習率衰減策略：polynomial

數據增強：隨機縮放 [0.5,2.0]、隨機翻轉、隨機裁剪 [512,512]

batch size：8

迭代次數：20,000

: (0.55,0.35)

在PAR，空洞卷積的空洞率：[1,2,4,8,12,24]

: (0.3,0.3,0.01)

5.1 消融實驗

不同 tok-k 池化對 CAM 的影響。

top-k 池化指的是去每一個 feature map 的 top k% 的值進行池化操作。gmp：global max-pooling（top-100%），gap：global average-pooling（top-1）。對於 transformer based 的分類網絡來說，使用 gmp 可以獲得更好的性能。

不同模塊帶來的性能提升