Emerging Properties in Self-Supervised Vision Transformers

作者: Mathilde Caron, Hugo Touvron, Ishan Misra, Herv´e Jegou, Julien Mairal, Piotr Bojanowski, Armand Joulin

Jiun Da Lin
7 min readMay 31, 2021

Papaer 連結: https://arxiv.org/abs/2104.14294
Code: https://github.com/facebookresearch/dino

簡介

Attention maps from multiple heads.

這篇 Paper 是Facebook AI Research(FAIR)所提出的,他們將這方法命名為DINO,主要是在探討若是 Self-supervised learning(SSL) 應用在 Vision Transformer(Vit) 上是否會產生與 Convolutional network(convnets)不同的效果,而在實驗結果當中確實可以看到透過 DINO attention block 可視化效果上看來,使用 SSL 的方式來訓練 Vit 確實能夠得到更佳的語意切割效果,厲害的是這篇論文所提出的方法完全不需要給定 Ground truth boundary 便能夠自最後一個 Self-attention block 自行生成,換句話說就是 不需要給定任何 Label 就能做 Knowledge distillation。

而在透過 SSL 訓練後得到的 Features 也能直接透過 KNN 直接做分類,訓練時完全不需要使用任何 Finetuning、Linear classifier、Data augmentation 便能在 ImageNet 上達到 78.3% Top-1 accuracy。

方法介紹

(圖1) Self-distillation with no labels

DINO 在實做上與一般 SSL 一樣也是採用類似 Joint embedding 的架構,輸入 x 會隨機生成不同的 Image 給 Teacher network 以及 Student network,兩個 Network 具有相同的架構但彼此間具有各自獨立的權重並非 Siamese networks (share weight),與其他 Teacher-Student 不同的地方在於在這篇 Paper 當中完全不需要對 Teacher 做 Pretrain。

Sg 全名為 Stop gradient,因為兩邊算的 Loss 是共用的,所以在 Propagate gradients 時只藉由 Student 做計算,而 Centering 則是用來避免 Output 被特定 Dimension 主宰, Student 不做 Centering 的原因是因為 SSL 在訓練時需要避免兩個 Ouput collapse,因為如果兩個都一樣就無法做學習,過去SSL要避免這種問題出現的方式有 Contrastive loss、Clustering Constraints、Batch normalization等。

另外在 Softmax 使用上也並非使用傳統的 Softmax,而是 Temperature softmax,Temperature 縮寫為下方公式當中的 T 變數,加入這個 T 變數的用意在於讓差異過大的 Distribution 能夠縮小差距,在這篇論文當中 Teacher temperature 被設為 0.04,而 Student temperature被設為 0.1,讓 Teacher 比 Student 小的原因是因為 x1 與 x2 當初在做 Random crop 時隨機依 0.4~1 的 Scale 生成 2 個 Global crop 並指定為 Teacher 的輸入 x2,而 Student 則包含那 2 個 Global crop 之外又依 0.05~0.4 的 Scale 額外多出了 8 個 Local crop,也因為上述的設定,導致 Student 在預測時較不穩定,導致 Output 端容易產生較不確定的結果,所以在 Teacher 端的 Softmax 結果因為較準確所以除以一個較小的數字讓預測結果更高機率的落在某一類別來讓 Student 參考,這個概念也可以想像成現實生活中 Teacher 扮演的角色通常是傳遞 "專業且正確"的知識,而 Student 則是學習這些 "專業且正確" 的知識同時去探索假設不一定正確的 "小細節" 之後再跟 Teacher 詢問正確性。

Teacher: T=0.04,x2 = 2 global ROI
Student: T=0.1,x1 = 2 global+8 local ROI

Temperature softmax:

Exponential moving average:

上方(圖3)的 EMA 箭頭代表的是 Exponential moving average,lambda 則設為 0.996,筆者認為這邊所傳遞的意思比較像是 Teacher 在自我學習,同時會向 Student 學習一些較細微的東西,只是這個比例並不高。

實驗結果

在 DINO 用 SSL 訓練完之後能夠將模型部署在多種任務上,其中包括 Image retrival、Copy detection、Object segmentation、Transfer-learning等等,因此是一個蠻泛用的架構,對於有做類似研究的人十分值得嘗試看看,因為實驗眾多,下方就列了一個主要的實驗,(表1)當中主要是比較搭配不同的架構訓練效果上的差異,RN50 代表的是 ResNet-50,Vit 以及 Deit 後面的/8、/16代表的是 Vision transformer 的 Patch size,並提供了對應的 Image thoughput 給大家參考,從實驗上來看可以發現 Linear 的分數普遍都比 K-NN 來得好,且使用 Vision transformer 做為 Backbone 會比使用 ResNet 來的高。

(表1) Linear and k-NN classification on ImageNet,

最後我自己也做了一個實驗,將模型訓練在假臉資料集上頭,模型確實完全不需要給定任何 Label 就能夠學習到臉部特徵,如下方(圖2)、(圖3) 所示,透過這種 SSL 方式先讓模型自我學習後更能夠讓抽取出來 Features 有較好的空間域分布特性,在往後的各項任務上都能較輕鬆的應付。

(圖2) 各層 Attention block 下的 heatmap,第一列為 mask,第二列為 mask 在 threshold = 0.6 之下套在輸入圖片的結果
(圖3) 最後一個 Attention block 下的 heatmap,第一列為 mask,第二列為 mask 在 threshold = 0.6 之下套在輸入圖片的結果

--

--