在基于圖像級別標注的弱監(jiān)督語義分割問題中，CNN分類模型中卷積操作的局部信息感知會導致不完全的語義區(qū)域激活。為解決這個問題，本文引入視覺Transformer結構，并探索了適合視覺Transformer的初始偽標簽生成方法。同時，受視覺Transformer中學習到的自注意力與圖像中的語義Affinity的一致性啟發(fā)，本文提出了一個Affinity from Attention（AFA）模塊，從Transformer的注意力矩陣中學習高質(zhì)量的語義Affinity信息，用于對初始偽標簽進行改善。為了進一步補充偽標簽的局部細節(jié)信息，同時保證端到端訓練的效率，本文基于像素自適應卷積設計了一個高效的處理模塊。本文提出的方法在兩個數(shù)據(jù)集上超過了當前的端到端方法，以及部分多階段方法。

本工作是由京東探索研究院、武漢大學、悉尼大學聯(lián)合完成，已經(jīng)被CVPR 2022接收。

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

論文鏈接：https://arxiv.org/abs/2203.02664

代碼鏈接：https://github.com/rulixiang/afa

01 研究背景及動機

弱監(jiān)督語義分割（WSSS）的目標在于基于僅有弱標注（如image-level類別標簽）的數(shù)據(jù)訓練模型生成像素級的預測。目前主流的弱監(jiān)督語義分割方法通常首先訓練分類模型，基于類別激活圖（CAM）或其變種生成初始偽標簽；然后對偽標簽進行細化作為監(jiān)督信息訓練一個獨立的語義分割網(wǎng)絡作為最終模型。其中，分類模型一般基于卷積網(wǎng)絡，無法準確感知圖像中的全局特征關聯(lián)，導致初始偽標簽通常只覆蓋語義物體中最具有判別性的部分，顯著影響了最終語義分割的效果。

本文引入視覺Transformer結構來解決這個問題。首先，視覺Transformer中的自注意力機制天然保證了全局特征感知，解決了CNN的局部性缺陷，因而能提高初始偽標簽的準確性。同時如圖1所示，自注意力圖和像素之間的semantic affinity也存在天然的一致性，而語義affinity可以被進一步用于對初始偽標簽進行細化。然而，如圖2所示，由于自然學習到的自注意力矩陣缺乏監(jiān)督信息，直接將原始自注意力矩陣作為affinity信息對偽標簽進行細化并不能取得令人滿意的效果。

我們提出了Affinity from Attention（AFA）模塊，以初始偽標簽作為監(jiān)督信息指導自注意力的訓練，從而學習到高質(zhì)量的affinity信息，用于偽標簽的改善。

02 方法介紹

如圖3所示，本文使用Transformer作為基礎編碼器。初始偽標簽由CAM方法生成，然后使用提出的像素自適應細化模塊（PAR）進行修正。在affinity學習模塊（AFA）中，從Transformer中的多頭自注意力（MHSA）生成語義affinity預測。AFA使用基于細化初始偽標簽的偽affinity標簽進行監(jiān)督。然后，利用學習到的affinity通過隨機游走方法對初始偽標簽進行改善。在經(jīng)過PAR的進一步修正后，作為分割分支的監(jiān)督信息。

1、 視覺Transformer網(wǎng)絡

如上所示, 本文的方法使用視覺Transformer作為基礎網(wǎng)絡。在視覺Transformer網(wǎng)絡中, 對于輸入圖像, 首先將其劃分成 個patch, 然后經(jīng)過線性層得到相同數(shù)量的patch token。在每個block中, Transformer使用多頭自注意力 (MHSA）學習token之間的特征關聯(lián)。具體而言, 用 分別表示注意力中的query, key和value, 其中 表示MHSA的序號, 對應注意力矩陣和輸出分別計算為:

多頭注意力的輸出會進行拼接，然后輸入到FFN層作為Transformer block的輸出。

2、初始偽標簽

考慮到生成偽標簽的簡潔性和效率，我們采用CAM作為初始的偽標簽，其計算方法為使用分類層的權重對特征圖進行加權，使用ReLu函數(shù)消除負激活值后進行歸一化，再選取合適的背景閾值得到初始的偽標簽。

3、AFA模塊

我們注意到Transformer中的MHSA和語義affinity之間的一致性。然而，由于在訓練過程中沒有對MHSA施加明確的約束，因此在MHSA中學習到的affinity通常是不準確的，這意味著直接應用MHSA來細化初始標簽并不能很好得到令人滿意的效果。因此，我們提出了AFA模塊學習準確的affinity信息。

假設一個Transformer block中的多頭注意力表示為，在AFA中我們直接用簡單的MLP層對其進行映射輸出affinity矩陣的預測。本質(zhì)上，自注意力結構是一種像素間有向的建模方式（即自注意力矩陣是不對稱的），而語義affinity的關聯(lián)應當是無向的，因此我們簡單地將和其轉置進行相加來達到這種轉換：

為了學習到可靠的affinity矩陣 , 關鍵的一步是生成可靠的偽affinity標簽對其進行監(jiān)督。為了得到可靠的偽標簽, 我們選取兩個背景閾值將初始偽標簽分為可靠的前景、背景和不確定區(qū)域, 根據(jù)可靠區(qū)域的標簽關系得到可靠的偽affinity標簽 ?。然后計算affinity損失為：

根據(jù)上式，一方面網(wǎng)絡能從MHSA中學習到可靠的affinity關系；另一方面，由于affinity是MHSA的線性組合，上式又能夠保證注意力矩陣中更好的特征交互。在得到可靠的affinity矩陣之后，我們通過隨機游走算法對初始的偽標簽進行修正：

通過隨機游走算法，可以對初始偽標簽中的高affinity區(qū)域進行激活，并抑制低affinity的錯誤激活區(qū)域，從而使得偽標簽更好的貼合圖像中的語義邊界。

4、 PAR模塊

如圖3所示，偽affinity標簽是根據(jù)偽標簽生成的。然而，初始偽標簽通常是粗糙的并且局部不一致，即具有相似low-level圖像信息的相鄰像素可能不具有相同的語義。為保證局部一致性，之前有工作采用Dense CRF來細化初始偽標簽。然而，CRF在端到端框架中并不是一個有益的選擇，因為它顯著降低了訓練效率。受PAMR的啟發(fā)，它利用像素自適應卷積提取局部像素的RGB信息以對偽標簽進行細化，我們結合 RGB 和空間信息來定義低級的pairwise affinity并構建我們的像素自適應細化模塊（PAR）。首先，我們定義任意兩個像素的RGB信息和位置信息的kernel分別為：

根據(jù)定義的位置和RGB信息kernel函數(shù)計算像素之間的low-level affinity值為：

與CRF中計算全局像素的low-level affinity不同，我們根據(jù)像素的8個鄰居像素進行計算，該過程可以使用33的像素自適應卷積實現(xiàn)，從而能夠高效的插入到端到端的訓練框架中。同時為了擴大感受野，使用多個空洞卷積進行多個鄰居像素信息的提取。得到像素之間的low-level affinity之后，通過多次迭代進行偽標簽的修正：

5、網(wǎng)絡訓練損失

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

03 實驗結果

我們在弱監(jiān)督語義分割領域常用的VOC 2012和COCO 2014數(shù)據(jù)集上對我們的方法進行了實現(xiàn)。在CNN中，分類層前的池化方法對CAM生成效果有顯著的影響，我們實驗了對transformer結構的合適池化方式，并且發(fā)現(xiàn)：和CNN中全局平均池化（GAP）效果更好不同，對于transformer結構，全局最大池化（GMP）能夠得到更好的類別激活圖。

圖4和表2中展示了AFA模塊中的多頭注意力和語義affinity可視化圖，以及在不同設置下的偽標簽的定量和定性結果。我們注意到AFA損失以及隨后的隨機游走傳播過程都能夠有效的提高偽標簽的效果，以及直接使用MHSA作為affinity信息并不能帶來效果的提升。下圖中展示了AFA的affinity預測層不同head 的注意力的權重，說明不同head的attention的差別，也證明了將MHSA直接作為affinity信息并不是合理的做法。

在VOC和COCO數(shù)據(jù)集上的實驗都顯示AFA超過了最近的端到端語義分割方法，并且同時超過了部分多階段的方法以及一些引入了額外顯著性監(jiān)督的方法。

04 結論

本文中，我們探索了視覺transformer在弱監(jiān)督語義分割任務上的合適做法，證實了transformer結構能夠比CNN生成更好的偽標簽。受自注意力圖和像素的affinity關系的啟發(fā)，我們提出了從注意力中學習語義affinity，從而同時對初始偽標簽進行改善和監(jiān)督注意力的訓練保證更好的特征交互。我們設計了一個像素自適應細化模塊，在不影響端到端訓練效率的同時修正偽標簽的局部細節(jié)信息。提出的方法在兩個數(shù)據(jù)集上取得了當前最佳的效果。

參考文獻

[1] ? Lixiang Ru, Yibing Zhan, Baosheng Yu and Bo Du, Learning Affinity from Attention: End-to-End Weakly-Supervised Semantic Segmentation with Transformers, CVPR 2022.

[2] Ahn, Jiwoon and Kwak, Suha, Learning Pixel-Level Semantic Affinity with Image-Level Supervision for Weakly Supervised Semantic Segmentation, CVPR 2018.

[3] Nikita Araslanov and Stefan Roth, Single-stage Semantic Segmentation from Image Labels, CVPR 2020.

[4] Bolei Zhou, Aditya Khosla, Agata Lapedriza, Aude Oliva, Antonio Torralba, Learning deep features for discriminative localization, CVPR 2016.

[5] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, Neil Houlsby, An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, ICLR 2021.

公眾號后臺回復“CVPR 2022”獲取論文合集打包下載～