論文閱讀筆記：Momentum Contrast for Unsupervised Visual Representation Learning

Momentum Contrast for Unsupervised Visual Representation Learning

論文地址

期待代碼開源

Summary

• 提出了Momentum Contrast（MoCo）的無監督視覺表示學習方法；
• 把對比學習比作查字典過程，把字典當作隊列，引入動量更新；
• 對end-to-end、Memory Bank、MoCo三種對比學習方式進行了比較；
• 在ImageNet和Instagram 數據集上進行大規模的訓練和對比實驗，並將訓練後的特徵遷移至下游任務進行實驗。

Problem Statement

• 無監督表示學習在NLP領域中取得巨大的成功如GPT、BERT，但有監督預訓練仍是計算機視覺主流方法。

Methods

1.Contrastive Learning as Dictionary Look-up

Contrastive learning 即對比學習，可以將其視爲訓練一個編碼器進行字典查詢的任務。

假設有一個編碼的查詢q以及一組編碼的樣本$\lbrace k_0,k_1,k_2,... \rbrace$（字典中的鍵值），假設q和一個字典裏單獨的鍵值$k_+$匹配，Contrastive loss 即對比損失，它的取值低表示q和鍵值$k_+$（positive key）相似而和字典中其它所有鍵值（negative keys）不相似。

定義一種對比損失函數InfoNCE，形式爲：

$\mathcal{L}_q = -log \frac {exp(q·k_+/\tau)}{\sum_{i=0}^K exp(q·k_i/\tau)}$

$\tau$用來控制concentration level of distribution。

對比損失作爲無監督的目標函數用來訓練編碼器網絡來表示查詢（queries）和鍵值（keys），

查詢表示爲$q=f_q(x^q)$，$f_q$爲一個編碼器網絡，$x^q$爲查詢樣本，同樣有$k=f_k(x^k)$，輸入的具體形式由特定的任務決定。

2.Momentum Contrast

從上述分析來看，對比學習是一種在高維連續輸入（如圖片）中建立離散字典的方法，字典是動態的，鍵值是隨機採樣得到的，並且key encoder在訓練中進行更新。假設好的特徵可以通過包含大量negative樣本的字典中學習而來，並且key encoder能夠在更新中儘可能保持一致，基於這種思想作者提出了MoCo算法。

Dictionary as a queue. 方法的核心是將詞典保持爲數據樣本隊列。這樣可以重新利用當前mini-batch中已編碼的鍵值。同時隊列能夠將字典大小和mini-batch大小進行解耦，字典大小可以遠遠大於mini-batch的大小，可被當作超參數。由於mini-batch遵循先進先出的準則，字典總是表示一個所有數據的子集。

Momentum update. 使用隊列可以擴充字典的大小，但是對鍵值編碼器key encoder進行反向傳播變得更難（梯度會在隊列中的所有數據進行傳播）。而簡單地將query encoder $f_q$直接複製給key encoder $f_k$，這樣快速地改變key encoder會破壞鍵值表示的一致性。於是作者提出動量更新方法：

$\theta_k \leftarrow m\theta_k + (1-m)\theta_q , m\in[0,1)$

只有$\theta_q$通過反向傳播更新，$\theta_k$的變換更加平滑，這樣一來，儘管隊列中的鍵值被不同的編碼器進行編碼，但是這些編碼器的差別很小，在實驗中，大的動量（例如0.999）往往效果好於小的動量（例如0.9），意味着緩慢變化的key encoder是利用好隊列的關鍵所在。

Relations to previous mechanisms.

以上三種方法的不同之處在於對鍵值的保持方式以及鍵值編碼器的更新方法的不同。

a方法，字典大小和mini-batch大小相同，受限於GPU顯存，對大的mini-batch進行優化也是挑戰，有些pretexts進行了一些調整，能夠使用更大的字典，但是這樣不方便進行遷移使用。

b方法，Memory Bank包含數據集所有數據的特徵表示，從Memory Bank中採樣數據不需要進行反向傳播，所以能支持比較大的字典，然而一個樣本的特徵表示只在它出現時纔在Memory Bank更新，因此具有更少的一致性，而且它的更新只是進行特徵表示的更新，不涉及encoder。

3.Pretext Task

將一對查詢query和以及鍵值key組成樣本對，如果它們出自同一圖像，那麼是正樣本對，否則爲負樣本對。查詢和鍵值分別編碼自$f_q$和$f_k$。在隨機數據增強下從同一圖像中任意提取兩個"view"構建正樣本對，負樣本取自隊列。

Technical details. 使用ResNet作爲編碼器，最後一層輸出爲128D向量，即查詢query和鍵值key的表示。

Shuffling BN. 在實驗中發現Batch Norm會阻止模型學到良好的特徵表示。模型似乎會欺騙pretext task並容易找到低損失的解決方案。可能是因爲由BN導致的intra-batch communication among samples泄露了信息。

作者通過Shuffling BN來解決該問題。在訓練時使用多個GPU，在每個GPU上分別進行BN（常規操作），對於鍵值編碼器$f_k$，在當前mini-batch中打亂樣本的順序，再把它們送到GPU上分別進行BN，然後再恢復樣本的順序；對於查詢編碼器$f_q$，不改變樣本的順序。這能夠保證用於計算查詢和其正鍵值的批統計值出自兩個不同的子集。

Experiment

ImageNet-1M (IN-1M)、Instagram-1B (IG-1B)10億圖片數據集。

訓練：使用ResNet-50，SGD優化器，weight_decay=0.0001，momentum=0.9

對於（1N-1M）：mini-batch size=256，8GPUs，初始學習率爲0.03，訓練200epochs，在120~160epoch時將學習率乘以0.1，花費72小時;

對於（1G-1B）：mini-batch size=1024，64GPUs，初始學習率爲0.12，指數衰減（每62.5k iterations乘以0.9）訓練1.25M iterations，花費6天。

上圖爲end-to-end、memory bank、MoCo三種對比損失方法在ImageNet線性分類評價下的對比結果。

上圖爲ImageNet上，MoCo和其它方法在線性分類評價下的對比結果。

上圖爲PASCAL VOC trainval07+12上進行微調的目標檢測結果。

上圖爲end-to-end、memory bank、MoCo三種對比損失方法在PASCAL VOC目標檢測的結果。

上圖爲MoCo與之前方法在PASCAL VOC trainval2007上微調的目標檢測結果對比。

上圖爲在COCO上微調的目標檢測和實例分割結果。

上圖爲MoCo和ImageNet上有監督預訓練並微調的任務的對比。

相關工作

• Pretext tasks. The term “pretext” implies that the task being solved is not of genuine interest, but is solved only for the true purpose of learning a good data representation.

  術語“pretext”表示要解決的任務不是真正意義上的，而是僅出於學習良好數據表示的真正目的。

• loss functions. 常被用來研究pretext task的獨立性，包括但不限於

  • Contrastive losses [1] measure the similarities of sample pairs in a representation space

  • Adversarial losses [2] measure the difference between probability distributions，GAN和NCE有着相關聯繫

部分論文引用情況

• Dimensionality reduction by learning an invariant mapping. In CVPR, 2006.[1] （Contrastive Learning）

• Generative adversarial nets. In NIPS, 2014.[2]

• Unsupervised feature learning via non-parametric instance discrimination. In CVPR, 2018. [3]

  …

Notes

• 本文受[3]的啓發挺多的。

• [3]的作者在Improving Generalization via Scalable Neighborhood Component Analysis這篇文章裏提到了動量更新Memory Bank的方法。

• 作者在文章實驗部分"Shuffling BN"中寫道"batch norm prevents the model from learning good representations"。