機器之心報道
機器之心編輯部
來自港大、騰訊 AI Lab、牛津大學的研究者用 Transformer 振興 CNN 注意力。
自監督表徵學習近兩年十分火熱。機器學習界的三位泰斗 Geoffroy Hinton、 Yann Lecun 、 Yoshua Bengio 一致認為自監督學習有望使 AI 產生類人的推理能力,其中 Hinton 與 Lecun 近兩年也在 ICML / NeurIPS 中發表自監督表徵學習的研究工作。
可以說在機器學習頂會發表自監督表徵學習的文章,是與 Hinton 和 Lecun 站到了同一賽道上。而今年的 NeurIPS 2021,Lecun 剛發推感嘆他與另外一位 CV 泰斗 Jean Ponce 的自監督投稿 VICReg 被拒掉了,可見在機器學習領域,自監督學習的競爭激烈程度。另外一方面,最近熱門的 Transformer 給計算機視覺演算法帶來了全面的升級。那麼 Transformer 跟自監督表徵學習在一起會迸發出怎樣的火花?
來自港大、騰訊 AI Lab、牛津大學的學者在 NeurIPS 2021 發表的文章會帶來一個啟發性的答案。
該研究受現有自監督表徵學習架構 BYOL 的啟示,結合前沿的 Transformer ,提出利用 Transfomer 來提升 CNN 注意力的自監督表徵學習演算法。本文將現有的架構歸為 C-stream,另提出 T-stream。在 CNN 骨幹網路的輸出並行接入 T-stream。將 Transformer 置於 T-stream 中提升 CNN 輸出的注意力,並以此結果來監督 CNN 自身的輸出,從而達到提升 CNN 骨幹網路注意力的效果。在現有的標準資料集中,也進一步提升了 CNN 骨幹網路在下游識別任務的各類效能。
- 論文地址:https://arxiv.org/pdf/2110.05340.pdf
- Github 地址:https://github.com/ChongjianGE/CARE
背景和啟示:自監督表徵學習與樣本對比
基於影象內容的自監督表徵學習目標為訓練一個普適的視覺特徵編碼器(encoder backbone)。在給定神經網路架構(如 VGG, ResNet)的情況下,擺脫資料標註依賴構建自監督的過程,進行從零開始的初始化訓練(pretext training)。將訓練好的網路認為類似於用 ImageNet 標註資料預訓練的結果,後續將該網路進行一系列下游識別任務的遷移(downstream finetuning),包括影象分類,物體檢測和分割。由於在初始化訓練中未使用資料標籤做約束,預訓練的網路表徵並不針對下游具體的某個識別任務,從而具備普適的視覺表徵能力。其潛在的應用在於利用海量的網際網路資料,擺脫人工標註的、依賴自適應學習神經網路的視覺表徵能力,從而能夠受益於一系列的下游識別任務。自監督學習在下游識別任務中可以媲美有監督學習。
在自監督表徵學習的研究中,對比學習 (contrastive learning) 為常用的方法。給定一批未標註的資料,以當前一個數據為正樣本,其餘資料為負樣本。對比學習透過這樣的方式,構建正負樣本及其增廣的資料來確定損失函式從而訓練網路。其中一個輸入資料透過兩路網路形成兩個不同的 view,進行後續的樣本對比。在處理海量資料的過程中,有效的從資料中構建樣本和防止模型坍塌成為了熱門研究方向。從 MoCo[a]的佇列設計及網路動量更新開始,一系列的研究工作應運而生。這裡介紹幾個代表性的工作,為簡潔起見,演算法框架圖中的對映器 (projector) 沒有畫出:
SimCLR
上圖為 Hinton 團隊的 SimCLR[b]演算法框架,其採用 large batch 的設計,擺脫佇列儲存的依賴,直接對正負樣本進行對比構造損失來更新網路。
BYOL
上圖為 DeepMind 團隊的 BYOL[c]演算法框架,其單純利用當前樣本進行自身的多種資料增廣進行對比。同時引入 projector 來擺脫上下游任務對於網路的影響。在更新網路時也採用了動量更新的方式防止模型的坍塌。
MoCo V3
上圖為 Kaiming He 團隊的 MoCo V3[d],其將 transformer 做為學習編碼器(encoder backbone),利用現有的自監督學習框架進行樣本對比學習。同時也是將 vision transformer (ViT[e])作為編碼器引入自監督學習中的工作之一。
本文方法
與 MoCo V3 的出發點不同,本文的工作旨在利用 transformer 的自注意力機制來提升 CNN 編碼器的效能。其中將原有的 CNN 框架歸為 C-stream,然後提出包含 Transformer 的 T-stream。兩個 stream 同時接收 CNN 編碼器的輸出,然後用 T-stream 監督 C-stream。相比於 Yann Lecun 團隊的 Barlow Twins[f]利用協方差矩陣衡量兩個 view 的冗餘,本文引入可學習的 transformer 能夠自適應的在不同網路訓練狀態下提供注意力的輔助,從而更有效的提升 CNN 編碼器的視覺表徵能力。
本文的方法:CARE (CNN Attention REvitalization)
Proposed pipeline
本文提出的演算法流程圖如上所示。首先將輸入影象x進行兩次不同的預處理得到兩個正樣本x_1、x_2。然後,用 C-stream 的兩個 CNN 編碼器分別提取x_1、x_2的特徵,其中將一路 CNN 提取的特徵輸入對映器 projector1 和預測器 predictor1 得到高維特徵f_1(x),同時將另一路 CNN 提取的特徵僅輸入動量更新的對映器 (momentum projector1) 得到高維特徵f_2(x)。此外,雙路 CNN 提取的這兩組特徵也會被同時輸入到 T-stream。其中一路的 Transformer1 提取具有空間注意力的特徵,並將此特徵輸入到對映器 projector2 和預測器 predictor2 得到高維特徵f_3(x)。另一路動量更新的 Transformer 同樣提取 CNN 特徵並輸入動量更新的對映器 momentum projector2 得到高維特徵f_4(x)。
至此,演算法框架的前向過程已經設計完成。後續透過對f_1(x)、f_2(x)、f_3(x)、f_4(x)進行針對性的損失函式設計進行反向傳播的學習。本演算法在反向傳播過程中,僅更新 C-stream 以及 T-stream 的其中一路,而對應的另外一路則利用動量更新 (momentum update) 的手段進行。具體形式在後續介紹。
網路架構設計:本演算法旨在透過自監督學習框架的搭建,利用自定義的輔助任務來學習一個能夠有效提取影象特徵的 CNN 編碼器。本演算法對任意的 CNN 編碼器均具有一定的適用性,因此在 CNN 編碼器的選取上有著很好的靈活性。例如,ResNet50,ResNet101 以及 ResNet152 皆可以作為本演算法的 CNN 編碼器。Transformer 的結構如下圖所示:
Transformer 結構示意圖
該 Transformer 以 CNN 編碼器輸出為輸入,並輸出f_1(x)、f_2(x)、f_3(x)、f_4(x)更具空間專注度性質的特徵。本演算法所設計的 Transformer 主要包括 4 個序列的模組,其中單個模組如上圖所示。單個模組主要包含由一個 1x1 的卷積層,一個多頭自注意力層(Multi-head Self-attention, MHSA)[g]以及額外一個 1x1 卷積層組成。其中 MHSA 層可以很好地學習到具有空間專注度性質的特徵。此外,對映器 projector 和預測器 predictor 的主要結構為多層感知器(Multi-layer perceptron)。兩者皆包含兩個全線性連線層(fully connected layers),一個啟用層以及一個標準化層(batch normalization)。
本文設計的損失函式基於流程框架中的四個輸出f_1(x)、f_2(x)、f_3(x)、f_4(x)。其中本文用
表示 C-stream 的損失項,用
表示 T-stream 的損失項。其具體形式如下:
此外,本文用 T-stream 的輸出來監督 C-stream 的輸出。這個約束用
表示,具體形式如下:
該約束表明在自監督學習中,C-stream 的輸出會與 T-stream 的輸出儘量相似。所以最終整體的損失函式可以由如下表示:
在計算整體損失後,本演算法只後向傳播梯度更新 C-stream 和 T-stream 的上支。其下路分支主要透過動量更新的方式來更新引數。所謂動量更新指的主要是利用當前 C-Stream 和 T-Stream 的上路分支的引數,以及其之前時刻的引數資訊,來動量更新其下路分支的網路引數。在網路訓練結束後,只保留 CNN encoder1 作為目標編碼器。該編碼器隨後會用做下游識別任務的 backbone 網路。
視覺化分析:CNN 編碼器的注意力展示
在自監督訓練結束後,本文對 CNN 編碼器進行特徵響應的視覺化展示,從而觀察編碼器在訓練後對視覺內容關注程度的變化。本文對同樣的編碼器進行兩種自監督策略,一種是隻使用 C-stream 的結構進行訓練,一種是使用全部結構進行訓練。本文對這兩種訓練策略下的同樣的編碼器進行視覺化展示,如下圖所示:
CNN 編碼器的注意力視覺化展示
從圖中可以看到,第一行為輸入影象,第二行為單純利用 C-stream 結構進行訓練的編碼器的注意力,第三行為利用本文提出 CARE 結構進行訓練的編碼器的注意力。透過觀察注意力在影象上面的分佈和強度可以看出,本文提出的 CARE 演算法訓練的編碼器對影象中的物體更敏感,注意力更強烈。
實驗結果
在實驗過程中,本文從多個方面驗證提出演算法的有效性。包括上游訓練好的模型線上性分類問題中的效能、在半監督資料中的效果,以及在下游物體檢測和分割任務中的效能。在骨幹網路模型選擇方面,本文選取了 CNN 通用的 ResNet 系列模型進行訓練。對比模型為 ResNet 系列和 Transformer 結構。驗證的方式為利用不同的自監督學習演算法在多種模型上進行各類任務不同訓練階段中的效能驗證。在本文演算法的訓練過程中,使用 8 卡 V100 算力即可進行模型訓練的收斂。在當前海量算力的視覺自監督表徵學習任務下相對算力友好。
上游任務線性分類的比較。在固定 ResNet-50 為骨幹網路情況下,針對不同的自監督學習演算法進行訓練,展示在不同訓練階段的線性分類效果。如下圖所示,本文提出的 CARE(CNN attention revitalization)方法取得的優異的表現。
此外,本文也與 BYOL 方法在 ResNet 不同骨幹網路、不同訓練階段的效能進行對比,如下圖所示。本文的 CARE 方法在不同骨幹網路下效能更佳。
以上為相同骨幹網路、不同學習演算法的對比。本文同時也對比了 Transformer 的骨幹網路以及現有的學習演算法。效果如下圖所示,跟 Transformer 結構相比,本文利用 ResNet 網路,在引數量相近的情況下,取得了更好的結果。
更多的實驗對比,以及 CARE 演算法的 Ablation Study 詳見文章中的實驗章節。
總結
綜上,本文提出了一個利用 Transformer 結構來輔助 CNN 網路訓練的視覺自監督表徵學習框架。其核心貢獻在於利用一種網路結構的特性(即 Transformer 的注意力提升特性),在訓練中監督目標網路(即 CNN 骨幹網路),從而使得網路特效能夠得到遷移並提升目標網路效能的效果。在視覺識別的各類任務中也得到了充分驗證。本文的框架對自監督表徵學習具有很強的啟示意義,現有網路結構設計繁多,功能各異。如何利用這些網路獨有的特點,進而整合在一個網路中達到渾然一體的目標,也是後續自監督表徵學習可探索的重要方向。
參考文獻:
[a]. Momentum contrast for unsupervised visual representation learning. Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. CVPR 2020.
[b]. A Simple Framework for Contrastive Learning of Visual Representations. Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. ICML 2020.
[c]. Bootstrap Your Own Latent: A New Approach to Self-Supervised Learning. Grill et al. NIPS 2020.
[d]. An Empirical Study of Training Self-Supervised Vision Transformers. Xinlei Chen, Saining Xie, and Kaiming He. ICCV 2021.
[e]. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. Dosovitskiy et al. ICLR 2021.
[f]. Barlow Twins: Self-Supervised Learning via Redundancy Reduction. Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. ICML 2021.
[g]. Bottleneck Transformers for Visual Recognition. Aravind Srinivas, Tsung-Yi Lin, Niki Parmar, Jonathon Shlens, Pieter Abbeel, Ashish Vaswani. CVPR 2021.