自監(jiān)督表征學習近兩年十分火熱。機器學習界的三位泰斗 Geoffroy Hinton、 Yann Lecun 、 Yoshua Bengio 一致認為自監(jiān)督學習有望使 AI 產(chǎn)生類人的推理能力，其中 Hinton 與 Lecun 近兩年也在 ICML / NeurIPS 中發(fā)表自監(jiān)督表征學習的研究工作。

可以說在機器學習頂會發(fā)表自監(jiān)督表征學習的文章，是與 Hinton 和 Lecun 站到了同一賽道上。而今年的 NeurIPS 2021，Lecun 剛發(fā)推感嘆他與另外一位 CV 泰斗 Jean Ponce 的自監(jiān)督投稿 VICReg 被拒掉了，可見在機器學習領(lǐng)域，自監(jiān)督學習的競爭激烈程度。另外一方面，最近熱門的 Transformer 給計算機視覺算法帶來了全面的升級。那么 Transformer 跟自監(jiān)督表征學習在一起會迸發(fā)出怎樣的火花？

來自港大、騰訊 AI Lab、牛津大學的學者在 NeurIPS 2021 發(fā)表的文章會帶來一個啟發(fā)性的答案。

該研究受現(xiàn)有自監(jiān)督表征學習架構(gòu) BYOL 的啟示，結(jié)合前沿的 Transformer ，提出利用 Transfomer 來提升 CNN 注意力的自監(jiān)督表征學習算法。本文將現(xiàn)有的架構(gòu)歸為 C-stream，另提出 T-stream。在 CNN 骨干網(wǎng)絡(luò)的輸出并行接入 T-stream。將 Transformer 置于 T-stream 中提升 CNN 輸出的注意力，并以此結(jié)果來監(jiān)督 CNN 自身的輸出，從而達到提升 CNN 骨干網(wǎng)絡(luò)注意力的效果。在現(xiàn)有的標準數(shù)據(jù)集中，也進一步提升了 CNN 骨干網(wǎng)絡(luò)在下游識別任務(wù)的各類性能。

• 論文地址：https://arxiv.org/pdf/2110.05340.pdf
• Github 地址：https://github.com/ChongjianGE/CARE

背景和啟示：自監(jiān)督表征學習與樣本對比

基于圖像內(nèi)容的自監(jiān)督表征學習目標為訓練一個普適的視覺特征編碼器(encoder backbone)。在給定神經(jīng)網(wǎng)絡(luò)架構(gòu)（如 VGG, ResNet）的情況下，擺脫數(shù)據(jù)標注依賴構(gòu)建自監(jiān)督的過程，進行從零開始的初始化訓練(pretext training)。將訓練好的網(wǎng)絡(luò)認為類似于用 ImageNet 標注數(shù)據(jù)預訓練的結(jié)果，后續(xù)將該網(wǎng)絡(luò)進行一系列下游識別任務(wù)的遷移(downstream finetuning)，包括圖像分類，物體檢測和分割。由于在初始化訓練中未使用數(shù)據(jù)標簽做約束，預訓練的網(wǎng)絡(luò)表征并不針對下游具體的某個識別任務(wù)，從而具備普適的視覺表征能力。其潛在的應(yīng)用在于利用海量的互聯(lián)網(wǎng)數(shù)據(jù)，擺脫人工標注的、依賴自適應(yīng)學習神經(jīng)網(wǎng)絡(luò)的視覺表征能力，從而能夠受益于一系列的下游識別任務(wù)。自監(jiān)督學習在下游識別任務(wù)中可以媲美有監(jiān)督學習。

在自監(jiān)督表征學習的研究中，對比學習 (contrastive learning) 為常用的方法。給定一批未標注的數(shù)據(jù)，以當前一個數(shù)據(jù)為正樣本，其余數(shù)據(jù)為負樣本。對比學習通過這樣的方式，構(gòu)建正負樣本及其增廣的數(shù)據(jù)來確定損失函數(shù)從而訓練網(wǎng)絡(luò)。其中一個輸入數(shù)據(jù)通過兩路網(wǎng)絡(luò)形成兩個不同的 view，進行后續(xù)的樣本對比。在處理海量數(shù)據(jù)的過程中，有效的從數(shù)據(jù)中構(gòu)建樣本和防止模型坍塌成為了熱門研究方向。從 MoCo[a]的隊列設(shè)計及網(wǎng)絡(luò)動量更新開始，一系列的研究工作應(yīng)運而生。這里介紹幾個代表性的工作，為簡潔起見，算法框架圖中的映射器 (projector) 沒有畫出：

SimCLR

上圖為 Hinton 團隊的 SimCLR[b]算法框架，其采用 large batch 的設(shè)計，擺脫隊列存儲的依賴，直接對正負樣本進行對比構(gòu)造損失來更新網(wǎng)絡(luò)。

BYOL

上圖為 DeepMind 團隊的 BYOL[c]算法框架，其單純利用當前樣本進行自身的多種數(shù)據(jù)增廣進行對比。同時引入 projector 來擺脫上下游任務(wù)對于網(wǎng)絡(luò)的影響。在更新網(wǎng)絡(luò)時也采用了動量更新的方式防止模型的坍塌。

MoCo V3

上圖為 Kaiming He 團隊的 MoCo V3[d]，其將 transformer 做為學習編碼器（encoder backbone），利用現(xiàn)有的自監(jiān)督學習框架進行樣本對比學習。同時也是將 vision transformer (ViT[e])作為編碼器引入自監(jiān)督學習中的工作之一。

本文方法

與 MoCo V3 的出發(fā)點不同，本文的工作旨在利用 transformer 的自注意力機制來提升 CNN 編碼器的性能。其中將原有的 CNN 框架歸為 C-stream，然后提出包含 Transformer 的 T-stream。兩個 stream 同時接收 CNN 編碼器的輸出，然后用 T-stream 監(jiān)督 C-stream。相比于 Yann Lecun 團隊的 Barlow Twins[f]利用協(xié)方差矩陣衡量兩個 view 的冗余，本文引入可學習的 transformer 能夠自適應(yīng)的在不同網(wǎng)絡(luò)訓練狀態(tài)下提供注意力的輔助，從而更有效的提升 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)。

至此，算法框架的前向過程已經(jīng)設(shè)計完成。后續(xù)通過對f_1(x)、f_2(x)、f_3(x)、f_4(x)進行針對性的損失函數(shù)設(shè)計進行反向傳播的學習。本算法在反向傳播過程中，僅更新 C-stream 以及 T-stream 的其中一路，而對應(yīng)的另外一路則利用動量更新 (momentum update) 的手段進行。具體形式在后續(xù)介紹。

網(wǎng)絡(luò)架構(gòu)設(shè)計：本算法旨在通過自監(jiān)督學習框架的搭建，利用自定義的輔助任務(wù)來學習一個能夠有效提取圖像特征的 CNN 編碼器。本算法對任意的 CNN 編碼器均具有一定的適用性，因此在 CNN 編碼器的選取上有著很好的靈活性。例如，ResNet50，ResNet101 以及 ResNet152 皆可以作為本算法的 CNN 編碼器。Transformer 的結(jié)構(gòu)如下圖所示：

Transformer 結(jié)構(gòu)示意圖

該 Transformer 以 CNN 編碼器輸出為輸入，并輸出f_1(x)、f_2(x)、f_3(x)、f_4(x)更具空間專注度性質(zhì)的特征。本算法所設(shè)計的 Transformer 主要包括 4 個串行的模塊，其中單個模塊如上圖所示。單個模塊主要包含由一個 1x1 的卷積層，一個多頭自注意力層（Multi-head Self-attention, MHSA）[g]以及額外一個 1x1 卷積層組成。其中 MHSA 層可以很好地學習到具有空間專注度性質(zhì)的特征。此外，映射器 projector 和預測器 predictor 的主要結(jié)構(gòu)為多層感知器(Multi-layer perceptron)。兩者皆包含兩個全線性連接層(fully connected layers)，一個激活層以及一個標準化層(batch normalization)。

本文設(shè)計的損失函數(shù)基于流程框架中的四個輸出f_1(x)、f_2(x)、f_3(x)、f_4(x)。其中本文用

表示 C-stream 的損失項，用

表示 T-stream 的損失項。其具體形式如下：

此外，本文用 T-stream 的輸出來監(jiān)督 C-stream 的輸出。這個約束用

表示，具體形式如下：

該約束表明在自監(jiān)督學習中，C-stream 的輸出會與 T-stream 的輸出盡量相似。所以最終整體的損失函數(shù)可以由如下表示：

在計算整體損失后，本算法只后向傳播梯度更新 C-stream 和 T-stream 的上支。其下路分支主要通過動量更新的方式來更新參數(shù)。所謂動量更新指的主要是利用當前 C-Stream 和 T-Stream 的上路分支的參數(shù)，以及其之前時刻的參數(shù)信息，來動量更新其下路分支的網(wǎng)絡(luò)參數(shù)。在網(wǎng)絡(luò)訓練結(jié)束后，只保留 CNN encoder1 作為目標編碼器。該編碼器隨后會用做下游識別任務(wù)的 backbone 網(wǎng)絡(luò)。

可視化分析：CNN 編碼器的注意力展示

在自監(jiān)督訓練結(jié)束后，本文對 CNN 編碼器進行特征響應(yīng)的可視化展示，從而觀察編碼器在訓練后對視覺內(nèi)容關(guān)注程度的變化。本文對同樣的編碼器進行兩種自監(jiān)督策略，一種是只使用 C-stream 的結(jié)構(gòu)進行訓練，一種是使用全部結(jié)構(gòu)進行訓練。本文對這兩種訓練策略下的同樣的編碼器進行可視化展示，如下圖所示：

CNN 編碼器的注意力可視化展示

從圖中可以看到，第一行為輸入圖像，第二行為單純利用 C-stream 結(jié)構(gòu)進行訓練的編碼器的注意力，第三行為利用本文提出 CARE 結(jié)構(gòu)進行訓練的編碼器的注意力。通過觀察注意力在圖像上面的分布和強度可以看出，本文提出的 CARE 算法訓練的編碼器對圖像中的物體更敏感，注意力更強烈。

實驗結(jié)果

在實驗過程中，本文從多個方面驗證提出算法的有效性。包括上游訓練好的模型在線性分類問題中的性能、在半監(jiān)督數(shù)據(jù)中的效果，以及在下游物體檢測和分割任務(wù)中的性能。在骨干網(wǎng)絡(luò)模型選擇方面，本文選取了 CNN 通用的 ResNet 系列模型進行訓練。對比模型為 ResNet 系列和 Transformer 結(jié)構(gòu)。驗證的方式為利用不同的自監(jiān)督學習算法在多種模型上進行各類任務(wù)不同訓練階段中的性能驗證。在本文算法的訓練過程中，使用 8 卡 V100 算力即可進行模型訓練的收斂。在當前海量算力的視覺自監(jiān)督表征學習任務(wù)下相對算力友好。

上游任務(wù)線性分類的比較。在固定 ResNet-50 為骨干網(wǎng)絡(luò)情況下，針對不同的自監(jiān)督學習算法進行訓練，展示在不同訓練階段的線性分類效果。如下圖所示，本文提出的 CARE(CNN attention revitalization)方法取得的優(yōu)異的表現(xiàn)。

此外，本文也與 BYOL 方法在 ResNet 不同骨干網(wǎng)絡(luò)、不同訓練階段的性能進行對比，如下圖所示。本文的 CARE 方法在不同骨干網(wǎng)絡(luò)下性能更佳。

以上為相同骨干網(wǎng)絡(luò)、不同學習算法的對比。本文同時也對比了 Transformer 的骨干網(wǎng)絡(luò)以及現(xiàn)有的學習算法。效果如下圖所示，跟 Transformer 結(jié)構(gòu)相比，本文利用 ResNet 網(wǎng)絡(luò)，在參數(shù)量相近的情況下，取得了更好的結(jié)果。

更多的實驗對比，以及 CARE 算法的 Ablation Study 詳見文章中的實驗章節(jié)。

總結(jié)

綜上，本文提出了一個利用 Transformer 結(jié)構(gòu)來輔助 CNN 網(wǎng)絡(luò)訓練的視覺自監(jiān)督表征學習框架。其核心貢獻在于利用一種網(wǎng)絡(luò)結(jié)構(gòu)的特性（即 Transformer 的注意力提升特性），在訓練中監(jiān)督目標網(wǎng)絡(luò)（即 CNN 骨干網(wǎng)絡(luò)），從而使得網(wǎng)絡(luò)特性能夠得到遷移并提升目標網(wǎng)絡(luò)性能的效果。在視覺識別的各類任務(wù)中也得到了充分驗證。本文的框架對自監(jiān)督表征學習具有很強的啟示意義，現(xiàn)有網(wǎng)絡(luò)結(jié)構(gòu)設(shè)計繁多，功能各異。如何利用這些網(wǎng)絡(luò)獨有的特點，進而集成在一個網(wǎng)絡(luò)中達到渾然一體的目標，也是后續(xù)自監(jiān)督表征學習可探索的重要方向。