NeurIPS 2021 _ MST- 用于Transformer視覺表征的Masked自監督解讀

MST: Masked Self-Supervised Transformer for Visual Representation

Part 1 問題和挑戰

Yann LeCun曾說“如果人工智能是一個蛋糕，則蛋糕的主要成分就是無監督習”。這句話反應了無監督學習在深度學習領域扮演著重要的作用。當前普遍的方法集中在如何設計有效的代理任務來對無標注的輸入數據學習到好的視覺表征。在計算機視覺上，目前比較流行的直接有效的方法是對比學習，將訓練數據的每個實例當做單一的分類。基于這個實力判別，很多自監督方法在分類任務上得到了有效的提升。他們成功彌補了自監督方法和監督方法的代溝。然而，這項任務仍然具有挑戰：

a. 掩碼語言模型在自然語言領域得到了廣泛的應用。圖像是高維特征，多噪聲且相比于文本形式復雜。在視覺領域中，圖像的主要信息會被隨機分到不同的token中，如果這些token被隨機masked掉，將會導致很差的表現。這個隨機掩碼語言模型容易掩蓋圖像的關鍵區域的token，這樣會導致誤判且不適合直接應用于自監督視覺Transformers。

b. 很多自監督方法是利用全局特征學習圖像級別預測，對于像素級別預測優化不足。當前自監督學習方法也許對圖像分類任務過度擬合，對下游密集任務預測表現效果不好。

Part 2 方法介紹

針對以上提出的問題，我們提出掩碼Transformer自監督學習方法，如下圖所示。MST創造性的引入了注意力特征圖引導掩碼策略并利用掩碼特征來恢復全局圖像特征任務。我們將介紹如何利用注意力特征引導掩碼策略幫助掩碼語言模型應用到視覺領域。最后我們將介紹網絡的結構和實驗細節。

圖1 MST整體全流程圖

1. 自監督方法網絡結構

我們根據多重裁剪在多種通用的數據增強下，為每個圖像生成多個視圖。這樣操作可以獲得兩個標準分辨率的裁剪圖片  和  (表示全局視圖)和N個低分辨率裁剪樣本（表示局部視圖）。如圖1所示，整體方法是由兩個編碼來進行編碼，教師網絡  和學生網絡  ，參數分別是  和  。兩個編碼器都是Transformer的backbone和projection head組成。教師網絡的編碼器參數  是由學生網絡的編碼器參數  來動態平均更新。更新公式如公式1所示：

其中：m是動量系數。

給定一個固定的教師網絡  ，學生網絡  通過最小化交叉熵損失如公式2所示來學習參數  ：

2. 視覺掩碼模型的掩碼token策略

圖2 MST的注意力引導面具策略的插圖。與原始隨機掩碼相比，通過保留圖像中的關鍵區域來改進。從左到右描述圖像: (a)輸入圖像，(b)通過自注意模塊獲得的注意力圖，(c)可能導致關鍵特征丟失的隨機掩碼策略，(d) MST的注意力引導掩碼策略，僅遮蓋非必要區域。事實上，掩碼策略是掩蔽Token。

隨機掩碼策略: 受到自然語言處理領域掩碼語言模型策略的啟發，我們應用這個隨機掩碼策略到自監督學習。根據公式3所示，重要區域的token和非重要區域的token有同樣的概率可能被掩碼。如圖2所示，我們可以觀察到隨機掩碼策略會消除重要區域的tokens，導致難以區分輸入圖像的語義信息。這個隨機采樣策略會抑制輸入圖像的重要區域，不利于網絡的識別能力。這個策略不適合直接應用到自監督視覺Transformer，如果掩碼策略未正確調制，則整體性能會惡化。

其中：m代表的掩碼區域，p是掩碼的概率（默認為0.15），prob是隨機產生的概率值。

注意力引導掩碼策略: 我們提出注意力引導掩碼策略動態控制掩碼區域并降低重要區域被掩碼的概率。同時，這個方法并不會增加額外的計算時間，整體方法結構偽代碼所示。如公式4所示，我們按照升序對每張圖像的不同patch的注意力特征進行排序，并將排序后的注意力特征值作為閾值，則低于的區域作為掩碼的候選區域。學生分支模型收到不同patches的重要性，根據概率生成掩碼。

其中：Attn代表注意力特征。

遵循BERT方法，掩碼區域填充有可學習的掩碼嵌入 [MASK]。如公式5所示，注意力引導掩碼策略保證高分的patch不被掩碼。

注意力引導掩碼策略對于預訓練模型有兩點益處：

a. 模型利用語義信息來學習到不同patches之間的關系。使得模型既保留了圖像全局語義信息，同時在圖像的局部細節有關注。

b. 我們的策略可以避免掩蓋關鍵區域，同時用可學習的掩碼特征，使得模型專注于關鍵區域。

3. 視覺Transformer掩碼解碼方法

在掩碼語言模型中，利用非掩碼區域特征來預測掩碼token。區別于原始的掩碼語言模型，我們的方法利用非掩碼的區域特征來恢復重建原始的輸入圖片。為了使網絡能夠在視覺任務上實現像素級恢復，從而增強圖像的像素級信息特征和細粒度空間結構的能力。利用卷積的偏置性特征，重建任務利用卷積神經網絡作為解碼器，使用卷積層和上采樣操作交替堆疊。如下公式6展示了重建損失函數：

其中：x代表輸入圖片，g代表解碼器，代表學生分支的編碼器，  代表學生分支的參數，  代表解碼器參數。

全體損失函數如公式7所示：

其中：λ代表權重系數。

Part 3 實驗結果

我們使用不同的transformer結構在ImageNet benchmark上訓練了預訓練模型，然后驗證他們的遷移能力在下游任務上例如目標檢測和語義分割，以及不同掩碼參數對模型的影響。

1. ImageNet Benchmark的對比

如表1所示為我們的方法與當前主要自監督算法比較。所有的這些方法擁有相同的backbone用于公平比較。我們方法的300-epoch模型實現76.9% 在線性評估的top-1的精度。我們的方法在相同的訓練epochs下超越當時最好的自監督方法DINO大約1.7%，甚至接近DINO更長訓練策略的表現效果（77.0%的800epoch）。需要強調的是，我們的算法緩解了自監督學習對極長訓練時間的需求，并且能夠在僅100個epochs的情況下獲得不錯的效果（75.0%）。

MST是通用的方法可以應用到任何基于Transformer結構的自監督方法。這里我們使用流行的Swin-T作為示例，它具有與DeiT-S相似的參數數量。使用相同的訓練epochs，MST的性能比MoBY高出1.8%，這是一種為Swin-T精心設計的自監督學習方法。Swin-T與DeiT-S共享相同的超參數，它仍然可以通過進一步調整來改進。

表1 ImageNet上流行的自監督學習方法的比較

2. 目標檢測和實例分割下游任務

在表2中，我們顯示了通過不同的自我監督方法和監督訓練學習的表示的表現。為了公平比較，所有這些方法都預先訓練了100個epochs。我們觀察到，我們的方法以42.7%的bbox mAP和38.8%的掩模mAP達到了最佳效果。它比ImageNet監督模型高出1.2%和0.5%，MoBY結果在同一時期分別比ImageNet監督模型高出1.2%和0.5%。結果表明，MST不僅在圖像分類任務上表現良好，而且在下游密集預測任務中表現良好。因此它具有很強的遷移能力。

表2 在 MS COCO 上微調的對象檢測和實例分段結果

如表所示，它說明了監督方法、DINO和我們的方法在該評估上的比較。我們的方法實現了最高的mIoU74.7%和mAcc82.35%。它的表現優于監督結果（+2.71%mIoU和+2.05%mAcc）和DINO預訓練結果（+1.08% mIoU和 +1.03%mAcc）。我們的模型也適用于語義分割任務的遷移。

表3 在城市景觀上微調的語義分割結果

3. 不同Mask策略的影響

表4顯示了不同掩碼策略的影響。我們使用隨機掩碼策略，注意力引導的掩碼策略和無掩碼來訓練DeiT-S。為了公平比較，所有方法都以相同的概率p進行掩碼。可以觀察到隨機掩碼策略的性能下降。此策略可能會抑制識別圖像能力（從 73.1 到 63.2）。隨機掩碼策略可能會破壞原始圖像關鍵區域的tokens，這些tokens對于識別對象可能是必不可少的。被屏蔽的輸入可能包含不完整甚至誤導性的信息。相反，我們的注意力引導掩碼策略的表現穩步提高（從73.1到73.7）。基本區域大多被保留下來，這可能是我們假設的有力證據。

表4 不同掩模策略（DeiT-S）的線性探針結果

4. 不同掩碼超參數的影響

表5驗證了注意力引導掩碼策略下不同掩碼超參數的性能。我們按升序對每個圖像的不同patch的注意力圖進行排序，并將前1/num個patch拆分為被掩碼的候選。刪除這些候選可以迫使網絡從相鄰的patch中學習局部特征，從而在不破壞語義的情況下增強對局部上下文進行建模的能力。這些候選根據概率p進行屏蔽。ImageNet上線性評估的Top-1精度如表下表所示。當num設置為8時，任何選擇p都可以得到一個可靠的結果，這表明最后的1/8個patch作為掩碼候選是相對安全的。

表5 基于注意力的掩碼策略的超參數設置

5. 與 BERT 的區別

在表6中，我們使用純掩碼語言模型和DeiT-S在100個epochs下進行實驗，結果在相同的實驗配置下約為40%。然后我們進一步調整其學習率和其他超參數，最佳結果僅為61%，遠低于DINO的10.6%（DINO結果為71.6%），也比監督結果低7.7%（監督結果為68.7%）。表明純掩碼語言模型方法可能不適合計算機視覺任務。此外，我們用對比損失+ BERT解決方案（即DINO+純掩碼語言模型）進行實驗，線性結果為71.9%。我們的方法比其結果高出2.0%（73.9%）。結果證明我們的方法比原來的結果方法更好。同時，我們進一步進行了實驗，僅將[mask] token替換為純掩碼語言模型策略，線性結果為73.5%，這也落后于我們的結果。這些結果充分展示了MLM對計算機視覺的更好設置，并進一步突出了我們論文的技術貢獻。

表6與 BERT 的區別

Part 4 結語

本文探討了當前視覺自我監督學習的兩個問題，即缺乏局部信息提取和空間信息丟失。為了克服上述問題，我們提出了一種基于Transformer的新型自監督學習方法，稱為MST。 MST利用注意力引導的掩碼策略來捕獲patch之間的局部關系，同時保留全局語義信息。需要注意的是，注意力引導的掩碼策略是基于從教師模型中提取的Multi-Head自注意特征圖，不會造成額外的計算成本。此外，在注意力引導掩模策略下，進一步使用全局圖像解碼器來恢復圖像的空間信息，這對于密集的預測任務至關重要。該方法在多個下游視覺任務中表現出良好的通用性和可擴展性。

Reference:

 [1] Devlin, J., Chang, M.W., Lee, K., Toutanova, K.N.: Bert: Pre-training of deep bidirectional transformers for language understanding. In: Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers). pp. 4171–4186 (2018)

[2] Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., Joulin, A.: Emerging properties in self-supervised vision transformers. arXiv: Computer Vision and Pattern Recognition

[3] Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., Guo, B.: Swin transformer: Hierarchical vision transformer using shifted windows. arXiv preprint arXiv:2103.14030 (2021)

[4]Xie, Z., Lin, Y., Yao, Z., Zhang, Z., Dai, Q., Cao, Y., Hu, H.: Self-supervised learning with swin transformers. arXiv preprint arXiv:2105.04553 (2021)