語義分割新正規化!StructToken:對per-pixel分類正規化的重新思考


MLNLP 

機器學習演算法與自然語言處理 

)社群是國內外知名自然語言處理社群,受眾覆蓋國內外NLP碩博生、高校老師以及企業研究人員。


社群的願景 是促進國內外自然語言處理,機器學習學術界、產業界和廣大愛好者之間的交流,特別是初學者同學們的進步。

本文轉載自 | 極市平臺
作者 | Lart
論文連結:https://arxiv.org/abs/2203.12612
本文是對語義分割傳統編解碼逐畫素分類正規化的一種思考和改進。
之前語義分割的工作將其視作一種逐畫素分類任務,主流正規化是編解碼結構,透過編碼器和解碼器學習逐畫素表徵後,透過對每個畫素單獨分類到不同的類別中從而獲得預測的語義掩碼結果。這篇文章選擇了另一種策略,即將結構資訊作為先驗直接構造語義掩碼然後逐步細化,而不再是按照逐畫素分類的正規化。
具體來看,對於給定的輸入影像,模型中可學習的結構token會和影像表徵進行互動,從而推理出最終的語義掩碼。這一思路和最初的ViT結構中的cls token的行為頗為類似。考慮到這份工作的實現是基於ViT-L,所以可以很直觀的推想出,StructToken的思路其實是將Transformer原本的形式向語義分割這樣的密集預測任務的一種“直接”遷移,這其中並沒有像其他工作那樣,過多受到目標任務中原始的卷積神經網路設計正規化的影響。
所以值得思考的幾點可以由此提出:
  • 本文定義的結構資訊是什麼?
  • 提出的設計是怎樣表達出這些結構資訊的?
  • 如何驗證這些設計帶來的提升與所謂的結構資訊有關?
1
『相關工作』
現有的語義分割領域已經出現了大量的工作,但是目前出現的工作中基本上都可以被歸類為逐畫素分類正規化,差別主要在於分類引數是否是動態的:
  • 靜態逐畫素分類:延續著以分割領域早期基於卷積神經網路的經典工作的正規化,例如FCN。他們主要利用上下文語義資訊的增強和多尺度特徵的融合,從而獲得更有效的影像特徵表示。並利用獨立的靜態分類器(典型如1×1卷積)實現逐畫素的語義類別預測。然而這類工作專注於提升逐畫素特徵的表達能力,卻並未在模型設計中考慮影像中的結構資訊。
  • 動態逐畫素分類:最近的工作中開始引入動態結構的思想。除了模型本身特徵處理結構之外,分類器也開始轉變為動態形式。論文列出的典型的工作有Segmenter[Segmenter: Transformer for semantic segmentation]、MaskFormer[Per-pixel classification is not all you need for semantic segmentation]、Mask2Former[Masked-attention mask transformer for universal image segmentation]和K-Net[K-Net: Towards unified image segmentation]。他們主要是使用了一系列與語義類別相關的可學習的token,與影像自身的特徵進行互動,從而實現最終mask的預測。這從形式上來看,可以認為是一種動態分類的過程。這些方法同時由沒有完全拋棄上面提到的經典範式,整體上獲得了更好的表現。但是從本文的角度來看,這類方法仍然沒有拋開逐畫素分類的正規化:)。
這些工作從整體上來看, 都是在學習針對每一類的線性判別函式,要麼是靜態的卷積,要麼動態的矩陣乘法運算。這會作用在逐畫素的特徵表示上,從而來為其賦予一個最相關的語義類別。
作者們認為,按照人識別物體的過程,先是捕獲語義類別的結構資訊(形狀等),然後關注於內部細節。想要分割影像中不同語義類別的區域,通常先根據結構生成一個粗略的mask,之後在調整mask的細節。現有的兩種逐畫素分類正規化並沒有充分的體現這一過程,而更多的是,直接在模型倒數第二層的特徵圖上分類畫素從而獲得得分圖。這一特性鼓勵網路最佳化單一畫素的表徵,而忽略甚至破壞了最重要的結構特徵。
本文中作者們提出了一種結構先驗正規化來解決這一問題,直接從結構token中構造得分圖來分割影像,然後逐漸細化。
2
『主要內容』
文中主要研究瞭如何從特徵圖中根據結構token提取有用的資訊。提出的結構整體遵循這樣的過程:
  1. 提取特徵:使用Transformer骨幹網路,例如ViT,提取特徵圖F,大小為[C,H/16,W/16]。
  2. 構造結構token:隨機初始化可學習的結構token S,大小為[K,N],K為資料集類別數量,N為patch數量,即[H/16,W/16]。
  3. 資訊互動:使用互動結構來處理S。捕獲特徵圖中的結構資訊,並根據學習到的先驗為每一類構建粗略的mask。
  4. 特徵細化:獨立的FFN用於結構token的細化,並處理特徵圖。
  5. 級聯處理:堆疊多個基礎單元(包括互動和細化)來重複處理特徵。
  6. 預測結果:尾部使用兩個卷積層和跳過連結構成的卷積塊來細化最終構建的分割mask並得到最終的結果。
這些步驟中,互動式結構的設計是本文的核心。文中主要探索了三種互動式結構。其中包含兩種動態結構和一種靜態結構。
  • 動態結構:基於Attention的思路,但是計算相關的token並非是空間patch,而是基於通道,即S中的類別token和F中的特徵通道之間的互動。
    • 第一種CSE基於Cross-Attention正規化,經過線性變換,S生成Query,F生成Key和Value,送入Cross-Attention。這裡得到的結果與S形狀一致。按照圖示,這裡也有個拆分操作,但是論文並未明說具體如何實現。
    • 第二種SSE基於Self-Attention正規化,S和F沿通道拼接後經過線性變換得到Query、Key和Value,並送入Self-Attention。結果會被按照通道的原始比例進行拆分。
  • 靜態結構:直接使用1×1卷積處理SSE模組中的相似性注意力的計算。卷積結果即為最終對應於拆分之前的結果。這一過程使用1×1卷積直接混合不同的輸入通道的資訊,實現了類似於SSE的過程。
上面結構中在執行Attention操作之前,S和F會被送入投影層處理,雖然是針對通道的Attention處理,但是這裡的投影層使用的是1×1卷積+3×3深度卷積+1×1卷積的形式,仍然是空間維度共享的操作。
這些模組的兩個輸出都會各自接一個FFN。這裡的FFN使用的是FC+3×3分組卷積+FC的結構。即可以細化區域性特徵,也可以看作是一種隱式位置編碼。
3
『實驗結果』

對比實驗

文中在三個主要的語義分割資料集上進行了驗證。本文的方法是構建在ViT的不同變體之上的,也因此沒有使用那些金字塔架構形式的多尺度特徵。
從這裡的實驗中可以發現一個趨勢,語義分割方法使用的backbone越來越大,從早期的的Res101,到現在的ViT-L、MiT-B5、Swin-L。預訓練權重甚至都開始使用ImageNet21K上的了。不知道這樣的潮流是否真的有意義。

消融實驗

對提出的結構中的不同元件進行了消融實驗。這裡的baseline模型基於ViT,僅隔著一個CSE模組和FFN模組。這裡的FFN沒有使用分組卷積,另外這裡不對Query、Key和Value的投影層進行消融實驗,因為作者們覺得如果替換成常規的全連線成,會導致無法支援多尺度推理。因為為了保持attention操作本身的原始性,僅對輸入轉置來實現通道attention而非手動修改投影層的情況下,此時的投影層就成了空間上的全連線了。
這裡還對提出的解碼塊堆疊數量進行了實驗,最終作者們考慮效能與計算複雜度的平衡,就選擇了4。實際上實驗中反映出來,更多的塊會帶來更好的效能表現。
但是作者們並沒有討論這些伴隨而來的計算量和引數量對於效能的影響。
為了驗證提出的結構token保留結構資訊的能力,作者們構建了一個逐畫素分類正規化的對等參考,backbone提取的特徵會先將通道數量調整到類別數(類似於本文提出的結構token那樣),每個通道認為對應一個類別。之後透過四個殘差塊來進行處理,最終使用1×1卷積生成最終的得分圖。來自每個殘差塊的輸出會被用來與本文模型中每個互動塊的結構token輸出進行視覺化對比。下圖中,不同的行組對應著不同的類別。可以看到,儘管輸出的得分圖很類似,但是結構token在中間的輸出卻展現出了更清晰的目標形狀、輪廓等結構資訊。而且隨著多個塊的處理,這些目標資訊更加清晰(典型如第9行)。
4
『總結與思考』
回答開頭的問題:
  • 本文定義的結構資訊是什麼?文章反覆在強調的結構其實直觀上可以理解為反映目標資訊的形狀和外觀。本文提出的結構Token在多次堆疊的處理單元的輸出中都明顯的凸顯出了特定類別的目標,確實實現了“粗略預測”的效果。
  • 提出的設計是怎樣表達出這些結構資訊的?基於通道的互動方式,使得結構token可以對影像特徵不同的通道進行自適應的組合與強化。雙流中各自的FFN實現了獨立的空間變換和通道整合,這保證了各自學習過程的差異性和多樣性。這種互動方式保留了影像空間結構資訊的獨立性和完整性。同時由於真值的監督,目標類別對應的結構token經過最佳化,會愈發具有與真值接近的空間上的外觀,也即論文中圖3中所展示的那樣。
    • 這是為什麼呢?我覺得這是因為通道注意力的使用的結果。基於通道之間的相似性計算的attention運算中,會為影像特徵中對應空間位置啟用更加明顯(即與結構token對應類別通道更加相似)的通道賦予更大的比重,這樣才會讓損失越來越小。
  • 如何驗證這些設計帶來的提升與所謂的結構資訊有關?作者對此並未進行探討,或許可以構造這樣一個實驗:在目前這種在最後單一監督的形式中,後續處理單元中結構token各個類別通道的視覺化結果非常趨於真值了。那若是使用深監督策略,直接對論文中提供的逐畫素分類正規化的卷積模型,對這些位置的特徵進行額外監督,進一步強化這些特徵對於這些目標區域的分割效果。如果效能進一步提升,則說明這樣的結構資訊的強化是有必要的。
技術交流群邀請函
△長按新增小助手
掃描二維碼新增小助手微信
請備註:姓名-學校/公司-研究方向
(如:小張-哈工大-對話系統)
即可申請加入自然語言處理/Pytorch等技術交流群

關於我們

MLNLP社群  機器學習演算法與自然語言處理 ) 是由國內外自然語言處理學者聯合構建的民間學術社群,目前已經發展為國內外知名自然語言處理社群,旗下包括  萬人頂會交流群、AI臻選匯、AI英才匯  以及  AI學術匯  等知名品牌,旨在促進機器學習,自然語言處理學術界、產業界和廣大愛好者之間的進步。
社群可以為相關從業者的深造、就業及研究等方面提供開放交流平臺。歡迎大家關注和加入我們。

相關文章