作者丨派派星
來源丨CVHub
編輯丨極市平台
Title: Integrally Pre-Trained Transformer Pyramid Networks
Paper: https://arxiv.org/pdf/2211.12735.pdf
Code: https://github.com/sunsmarterjie/iTPN
導讀
自 ViT 提出以後，Transformer 在計算機視覺領域逐漸衍生出兩個重要分支，一個分支是以 Vision Transformer 為代表的給為 Transformer 主幹網絡，而另一個分支便是以 MAE 和 BEiT 為代表的 掩碼圖像重建(Masked Image Modelling, MIM)技術。通過結合這兩項技術，極大的促進了包括分類、檢測和分割等下遊任務的發展。
今天為大家介紹的是一篇與 MIM 和 Vision Transformer 相關的一篇工作，其旨在解決上遊預訓練和下遊微調之間的遷移差距。以代表性工作 MAE 和 BEiT 為例，其使用的主幹網絡均為樸素的 ViT 模型。盡管 SimMIM、ConvMAE以及GreenMIM 等模型應用了分層結構，但本質上僅會作用到 Backbone 上，而不會影響到 Neck，即特征金字塔。這會導緻一個問題，僅當我們直接應用于下遊任務時，如果你是采用 Linear probing 的方式，雖然不會破壞預訓練的特征提取器，但由于整個優化過程從随機初始化的 Neck 層開始，作者認為這并不能確定當機的 Backbone 與 Neck 能夠很好的“搭配合作”。
是以，本文設計了一種簡潔有效的預訓練架構來緩解這種現象。考慮到部分讀者可能對這方面不太熟悉，本篇文章将會先引入相關的概念和必要的背景基礎知識，最後再詳細介紹此項工作。
背景
定義
掩碼圖像重建是一種利用圖像掩碼去執行計算機視覺任務的技術。相比于基于對比學習等方法，基于 MIM 的方法能提供更有競争力的結果。說到這，那就讓我們先簡單回顧下對比學習的知識吧。
對比學習
對比學習(Contrastive Learning)是一種通過在相同資料的扭曲視圖(distorted views)上提取不變特征來學習執行個體級判别表示。例如凱明的MoCo[1]和Hinton的SimCLR[2]便是早期的兩項代表性工作，它們采用不同的機制引入負樣本與正樣本進行對比：
随後，DeepMind 團隊提出了 BYOL[3] 進一步解決了對負樣本的依賴(避免表示崩潰)。當然，同期也有不少其他類似的工作，如凱明的SimSiam[4]便探索了孿生神經網絡表征學習的崩潰解。
除了這種組成一對對的表征學習方式外，SwAV[5]提倡對資料進行線上聚類，同時加強同一圖像的多重增強視圖之間的一緻性：
SwAV
最後便是一些緻力于将對比學習應用于提升特定下有任務的方法，如同時發表在 CVPR & ICCV 2021 上的三篇代表性文章：Detco[6]、ReSim[7]以及CAST[8]：
自凱明的 MIM 提出以後，此類對比學習方法的熱度也逐漸下降了，後面也就一些結合 Transformer 去做的工作如 DINO[9] 和 MoCov3[10]。
掩碼圖像重建
掩碼圖像重建是自監督任務的一種形式，它能夠很好的解決對資料的依賴。相比于 CV 領域，NLP 領域一直都處于領先地位，無論是基于 GPT 中的自回歸語言模組化或者 BERT 中的掩碼自編碼方案，本質上都是基于删除-預測的機制，這些方法很容易推廣到 LLM 上。
然而，正如凱明在 MAE 論文中指出，掩碼自編碼器的概念是一種更通用的去噪自編碼器(Denoising Auto Encoder, DAE)，它是自然的，也适用于計算機視覺。DAE 可以算是介于 AE 與 VAE 之間的産物。衆所周知，自編碼器(Auto-Encoder)是一種非常典型的網絡架構，它允許在沒有标注的情況下進行表征學習，由 Hinton 等人于 1993 年提出。随後過了十幾年，人們便提出了将噪聲“強制”引入到學習的表征上，其通過對潛在特征疊加高斯噪聲進而形成“損壞”的信号以作為網絡的輸入來重建未校正的輸入信号，這便是變分自編碼器(Variational auto-encoder, VAE)。
BERT 的提出帶火了 NLP 領域。而對于 CV 領域而言，筆者早期接觸的自監督學習主要有圖像修複、圖像着色、圖像拼圖等形式，強烈推薦大家去看下 Amit Chaudhary 的自監督學習筆記：https://amitness.com/2020/02/illustrated-self-supervised-learning/。随着 ViT 的引入，便湧現出了許多基于掩碼預測的自監督學習方法，下面帶大家快速過一遍：
iGPT較早提出在給定一系列像素作為輸入的情況下預測後續像素值。
MaskFeat應用 HOG 作為預測目标，而非 RGB 像素值。
SimMIM中采用線性層作為解碼器
BEiT和MAE算是兩個同期的工作，其一個亮點便是基于 ViT 模型重建缺失塊。
則是提出了一種與架構無關的 MIM 架構，将 Transformer 和 CNN 相容起來。
方法
Framework
整體的架構圖如上所示。首先，讓我們先看下圖中左側内容，其顯示了傳統的預訓練範式。需要注意的是，這裡我們先不區分具體的微調任務(如分類、檢測和分割)，假設它們都共享相同的 Backbone，而無需 Neck 和 Head (參考MAE和BEiT等)。那麼，問題來了，這種架構容易引起兩個問題：
- 骨幹網絡的參數并未針對多級特征提取進行專門的優化，即缺乏一種合理的機制來高效的提取和融合多尺度特征；
- 微調階段的優化若從随機初始化的 Neck 和 Head 開始，那麼這會顯著減慢訓練過程且不容易獲得更優的效果。
是以，iTPNs 提倡将重建和識别過程統一起來，以最大限度的減輕這種現象，如圖中右半部分所示。
Unifying Reconstruction and Recognition
具體地，iTPNs 基于 HiViT 和 特征金字塔建構了一個全新的架構。其中，HiViT 通過以下方式完成進一步的簡化：
- 采用通道多層感覺器(Channel-wise MultiLayer Perceptron, C-MLP)替換 shiftedwindow attentions；
- 将感受野為 的 stage 剔除掉，而是直接在 的 stage 直接計算全局注意力；
如此一來，該主幹網絡就不用像 SimMIM 方法一樣需要輸入全圖才能跑，直接節省 30%–50 % 的計算成本。下表展示了具體的參數比對：
此外，為了更好的聯合優化骨幹網絡(HiViT)和頸部(Feature Pyramid)，iTPNs 采用了以下兩個技術細節：
首先，作者通過将特征金字塔插入預訓練階段(用于重建)并在微調階段複用訓練好的權重(用于識别)以此來統一上遊和下遊頸部特征。
其次，為了更好地預訓練特征金字塔，本文提出了一種新穎的掩碼特征模組化(Masked Feature Modeling, MFM)任務為特征金字塔提供多階段監督，該任務通過将原始圖像輸入一個 moving-averaged backbone 計算出中間特征，同時使用特征金字塔的每一層輸出來重建中間目标。
總的來說，MFM 可以了解為 MIM 的一種補充方法，其更好的提高了重建和識别的準确性。此外，MFM 還可以适應從預訓練教師模型(本文應用了CLIP)“吸收知識”以獲得最佳的性能。
實驗
ImageNet-1K classification
可以看出，iTPN 顯示出優于現有方法的顯着優勢，無論是僅使用像素監督還是利用來自預訓練教師的知識（括号中内容為教師模型的名稱）。
ImageNet-1K classification
此外，iTPN 在幾個重要基準測試中的識别準确率均超過了之前的 SOTA。
COCO and ADE20K
不僅在圖像分類，在下遊的目标檢測和圖像分割任務下表現也很不錯。
Linear probing
正如我們前面所提到的，傳統 MIM 方法使用 Linear probing 微調時效果并不好，這在很多 MIM 相關的方法中也經常被提及。确實，與 Fine-tune 相比，這種方式對預訓練骨幹更加敏感。不過，從上表可以明顯的看出，iTPN 仍然能獲得不錯的精度，例如在 CLIP 監督下，超越同等配置的 MVP 将近兩個百分點。
總結
本文提出了一個用于預訓練 HiViT 的完整架構，其核心貢獻在于利用特征金字塔統一重建和識别任務，進而最大限度地減少預訓練和微調任務之間的遷移差距。此外，為了更好的優化特征金字塔，本文提出了一種掩碼特征模組化任務，旨在補充掩碼圖像模組化能力。最後，預訓練的 iTPN 在一些主流的視覺識别任務中報告了卓越的識别能力。